key_cache和value_cache必须是torch.Tensor且shape对齐，因为F.scaled_dot_product_attention要求输入张量stride严格匹配，shape需为(batch_size, num_heads, max_seq_len, head_dim)，否则报错；预分配+切片写入可避免cat导致的显存拷贝与碎片化，确保零拷贝高效推理。
为什么 key_cache 和 value_cache 必须是 torch.Tensor 且 shape 要对齐
LLM 推理时 KV Cache 的核心是复用已计算的 key/value，避免重复 attention 计算。PyTorch 原生不提供开箱即用的 KV Cache 管理类，得自己维护两个缓存张量。常见错误是把 cache 当成 list 追加 torch.cat 拼接——这会触发显存拷贝和碎片化，推理延迟飙升。
正确做法是预分配固定长度的 tensor，用 cache[:, :, :seq_len, :] 切片写入。shape 必须是 (batch_size, num_heads, max_seq_len, head_dim)，否则 F.scaled_dot_product_attention 会报 RuntimeError: expected stride to be a multiple of...。
- 预分配时用
torch.empty（非zeros），省掉初始化开销 -
max_seq_len要覆盖你所有请求的最大生成长度，太小会越界，太大浪费显存 - 多 batch 场景下，不同样本当前
seq_len不同，必须按实际长度索引切片，不能统一用cache[:, :, :, :]
如何在 forward 中安全接入 KV Cache 参数
标准 LLaMA / Mistral 的 nn.Module.forward 通常只接受 input_ids 和 attention_mask。要支持 cache，必须扩展签名：加 past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None，并返回更新后的 (key_cache, value_cache)。
关键点在于：cache 输入是上一轮的输出，这一轮算完要拼上新 token 的 KV，再传给下一轮。但注意——torch.cat 在循环中拼接不可取；应直接用索引赋值：
立即学习“Python免费学习笔记（深入）”；
key_cache[:, :, cur_pos, :] = new_k value_cache[:, :, cur_pos, :] = new_v
其中 cur_pos 是当前 token 在序列中的绝对位置（从 0 开始），不是相对偏移。
- 如果模型用了 RoPE，
new_k和new_v必须在旋转后、进 cache 前计算，否则位置信息错乱 -
attention_mask需同步扩展：原 mask 是(1, seq_len)，cache 模式下要变成(1, past_len + 1)，否则 padding 位置参与 softmax - HF Transformers 库的
use_cache=True就是干这事，但底层逻辑完全一样
使用 F.scaled_dot_product_attention 时 cache 的传入方式
PyTorch 2.0+ 的 F.scaled_dot_product_attention 支持 attn_mask 和 is_causal，但它**不自动管理 KV Cache**——你得手动把 key_cache 和 value_cache 拼到当前 query 对应的 KV 上。
典型写法是：
key = torch.cat([past_key, key], dim=2) value = torch.cat([past_value, value], dim=2) attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask)
但这是低效的。更优解是让 key 和 value 直接指向 cache 的 slice：
key = key_cache[:, :, :cur_pos + 1, :] value = value_cache[:, :, :cur_pos + 1, :]
这样零拷贝，且 cache 张量生命周期由外层控制，避免中间变量滞留。
- 务必确保
key_cache和value_cache设备与query一致，否则报Expected all tensors to be on the same device - 若启用
enable_math=False, enable_flash=True，某些旧驱动下 slice 可能触发 fallback，建议实测torch.backends.cuda.flash_sdp_enabled() - cache 的 dtype 必须和模型权重一致（如
torch.bfloat16），混用会导致精度坍塌
batch 内各序列长度不同时怎么处理 cache
真实服务场景中，一个 batch 里不同请求已生成的 token 数往往不同（比如有的刚输入 prompt，有的已 decode 了 20 步）。这时不能共用一个 cur_pos 标量，而要用 torch.arange 构造 position ids，并用 torch.scatter 或逐 sample 处理。
最简方案是：为每个样本维护独立的 cache_start_pos，写入时用 key_cache[i, :, start_pos[i]:start_pos[i]+1, :] = new_k[i]。虽然略慢于全量 slice，但比动态拼接稳定得多。
- 别试图用
mask把短序列 pad 成长序列再统一操作——padding 位置的 KV 会被 softmax 加权，污染注意力分布 - HuggingFace 的
DynamicCache类本质就是封装了 per-sample 的 offset 管理，但底层仍是 tensor slice - 如果你用 vLLM 或 TensorRT-LLM，它们的 PagedAttention 已解决这个问题，但 PyTorch 原生仍需手写逻辑
cache 管理本身不难，难的是和 position embedding、RoPE、mask、device placement 这些细节咬合严丝合缝。漏掉任意一环，模型就可能静默出错或显存爆炸。
