PyTorch 的 KV Cache 管理模式
1. 初始化阶段
通过 register_buffer 在模型初始化时预分配固定空间：
register_buffer("k_cache", torch.zeros(N_LAYERS, 1, MAX_SEQ_LEN, D_MODEL))
register_buffer("v_cache", torch.zeros(N_LAYERS, 1, MAX_SEQ_LEN, D_MODEL))
特点：
- KV cache 是模型的固有组成部分，生命周期与
nn.Module绑定 - 需要提前确定
MAX_SEQ_LEN，超出则溢出 - 数据常驻内存（GPU 或 CPU），不随模型序列化
2. Forward 阶段
每层生成新的 KV 后，通过 copy_ 追加到预分配空间：
new_k = ... # 当前 step 的 key
self.k_cache[layer_idx].copy_(torch.stack(new_k), non_blocking=True)
特点：
- 原地更新，避免额外分配
- 手动管理写入位置和偏移量
- 需要开发者保证并发安全和边界检查
Core AI 的 KV Cache 管理模式
1. Export 阶段：state_names 机制
在 torch.export 时声明 state_names，Converter 遍历 FX Graph 时会自动识别这些状态变量，并为它们生成对应的 placeholder 节点。
效果： KV cache 不再是模型的内部状态，而是被提升为图的输入/输出接口。
2. 推理阶段：显式输入/输出
导出的推理函数中，KV cache 作为独立的输入参数传入：
inputs + kv_cache → executable() → outputs + new_kv_cache
关键差异：
- 权限完全交给开发者：KV cache 的创建、加载、复用、销毁全由调用方控制
- 正确性也由开发者保证：框架不再隐式管理 KV cache 的生命周期
- 可序列化：KV cache 作为 NDArray 可以独立保存/传输，不绑定特定设备内存
两种模式的本质对比
| 维度 | PyTorch (register_buffer) | Core AI (state_names) |
|---|---|---|
| 所有权 | 模型持有，生命周期绑定 | 开发者持有，独立于模型 |
| 内存位置 | 固定（GPU/CPU） | 灵活（NDArray 可在不同 backing 间迁移） |
| 序列化 | 不随模型导出 | 可作为 NDArray 独立序列化 |
| 动态扩展 | 受限于 MAX_SEQ_LEN | 每次调用独立分配，无预分配限制 |
| 正确性责任 | 框架 + 开发者 | 完全由开发者承担 |
| 多请求复用 | 困难（模型状态互斥） | 自然支持（KV cache 是独立对象） |
对 vLLM 的启示
PyTorch 模式的局限： KV cache 与模型强耦合，导致 vLLM 必须自己实现一套外部的 paged attention 来绕过这个限制。
Core AI 模式的优点： 将 KV cache 抽象为独立的输入/输出流，天然支持：
- 多请求间的 KV cache 共享（同前缀场景）
- KV cache 的序列化与持久化
- 灵活的内存调度策略（CPU offload、分页交换等）
感悟
抛开一些目标上的区别，core ai的设计扩展性更好