怎样在Python中使用PyTorch实现LLM的KV Cache缓存-Python教程-PHP中文网

key_cache和value_cache必须是torch.tensor且shape对齐，因为f.scaled_dot_product_attention要求输入张量stride严格匹配，shape需为(batch_size, num_heads, max_seq_len, head_dim)，否则报错；预分配+切片写入可避免cat导致的显存拷贝与碎片化，确保零拷贝高效推理。

怎样在python中使用pytorch实现llm的kv cache缓存

为什么 `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 在循环中拼接不可取；应直接用索引赋值：

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 上。

Galileo python sdk

Galileo AI 平台 Python SDK 完整参考，用于评估、监控和保护 GenAI 应用，适用于构建 Python 应用。

下载

典型写法是：

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，但比动态拼接稳定得多。