4、KV Cache 数据结构设计：张量形状与维度、批次管理与序列长度对齐、内存布局
好，咱们今天聊点实在的。KV Cache 这东西，说白了就是个缓存。但你要是真把它当成一个简单的缓存来设计，那后面推理的时候，性能会给你颜色看的。我当年第一次手写推理引擎时，就在这上面栽过跟头。
KV Cache 的数据结构设计，直接决定了你的显存占用、访存效率，甚至决定了你能不能跑起来一个长序列。咱们今天就把这三个核心问题掰开揉碎讲清楚。
4.1 张量形状与维度：K 和 V 到底长什么样？
先看最基本的。一个 Transformer 层，注意力计算需要 Key 和 Value。那 KV Cache 里存的是什么？就是每一层、每一个头的 K 和 V。
它的形状，我习惯用 [batch_size, num_heads, seq_len, head_dim] 来表示。为什么这么排？因为 PyTorch 和大多数框架的注意力实现，都默认这个布局。你想想看，batch_size 是批次里的样本数，num_heads 是多头注意力的头数，seq_len 是当前已经生成的 token 数量，head_dim 是每个头的维度。
举个例子，一个 7B 模型，通常 num_heads=32，head_dim=128。如果 batch size 是 1，生成了 1024 个 token，那这一层的 KV Cache 大小就是 [1, 32, 1024, 128]。注意，这是 K 和 V 各一份，所以实际占用要乘以 2。
关键点： 这个形状不是一成不变的。在推理过程中，seq_len 这个维度是动态增长的。每生成一个 token，就要在最后一个维度上拼接新的 K 和 V。所以，你的数据结构必须支持高效的 append 操作。
我在项目中遇到过一个问题：有人把维度顺序搞反了，变成了 [batch_size, seq_len, num_heads, head_dim]。这在某些自定义算子里也能跑，但一旦要跟 Flash Attention 这类优化库对接，就得做 transpose，性能损失很大。所以，一开始就定好规范，能省很多事。
4.2 批次管理与序列长度对齐：动态 batch 的烦恼
实际线上服务，请求是源源不断来的。有的请求生成了 10 个 token，有的已经生成了 500 个。这时候，KV Cache 的批次管理就成了一个头疼的问题。
最简单的做法是「静态 batch」：把所有请求的序列长度对齐到最长的那个。短的序列后面补 padding。这样做，代码写起来简单，但显存浪费严重。你想想看，一个只生成 10 个 token 的请求，却要占着 500 个 token 的显存位置，太亏了。
更好的做法是「动态 batch」或者叫「continuous batching」。每个请求的 KV Cache 独立管理，长度各不相同。推理引擎在调度时，把当前所有活跃请求的 K 和 V 拼成一个不规则的张量。
这里就引出了「序列长度对齐」的问题。怎么对齐？
- Padding 对齐： 简单粗暴，短序列补 0。适合请求长度差异不大的场景。
- 分页对齐（PagedAttention）： 把 KV Cache 切成固定大小的块（比如 16 个 token 一块）。每个请求按需分配块，不浪费。vLLM 就是靠这个出名的。
- 动态拼接： 每次推理前，把当前批次里所有请求的 KV Cache 按实际长度拼接成一个连续的张量。这需要高效的显存管理。
我的建议： 如果你的场景是离线批量推理，请求长度差不多，用 padding 对齐就够了，实现简单。如果是线上实时服务，请求长度差异大，一定要上分页或者动态拼接。我曾经在一个项目里用 padding 对齐，结果显存利用率不到 30%，后来改成动态拼接，直接翻了三倍吞吐。
4.3 内存布局：连续 vs 非连续
这是最底层，也是最容易忽视的问题。内存布局决定了你的 GPU 能不能把数据喂给计算单元。
连续内存布局： 所有数据在显存里是一块连续的地址空间。比如 [batch_size, num_heads, seq_len, head_dim] 这个张量，在内存里就是按行优先顺序排列的。这种布局的优点是：访存局部性好，GPU 的缓存命中率高。缺点是：动态增长时，需要整体重新分配内存，代价很大。
非连续内存布局： 数据分散在多个不连续的内存块里。比如每个请求的 KV Cache 单独分配一块显存，然后用一个指针数组把它们串起来。这种布局的优点是：动态增长灵活，每个请求可以独立扩容。缺点是：访存时可能跨块，导致缓存 miss，性能下降。
怎么选？我个人的经验是：
- 对于 prefill 阶段（一次性处理所有输入 token），用连续布局。因为这时候序列长度是固定的，没有动态增长的问题，连续布局的访存效率最高。
- 对于 decode 阶段（逐个生成 token），用非连续布局。因为每个请求的序列长度都在变，非连续布局可以避免频繁的内存拷贝。
但现实往往更复杂。很多高性能推理引擎，比如 TensorRT-LLM，会在 decode 阶段也尽量保持连续布局。怎么做？它们会预先分配一块足够大的显存，然后在这块显存里「切」出连续的空间给每个请求。这样既保证了连续性，又避免了动态分配。
注意： 非连续布局还有一个隐藏的坑——内存碎片。频繁地分配和释放 KV Cache 块，会导致显存碎片化。时间长了，明明总显存还有剩余，但就是分配不出一个连续的大块。我曾经在一个长期运行的服务上遇到过这个问题，最后不得不加了一个定期的内存整理机制。
知识体系核心逻辑
下面这张图，是我自己梳理的 KV Cache 数据结构设计决策树。你可以对照着看，遇到实际问题时，顺着这个思路走，基本不会跑偏。
嗯，这张图把咱们今天聊的三个维度串起来了。你从根节点往下走，先确定你的张量形状，再决定批次管理策略，最后选内存布局。每一步都有 trade-off，没有银弹。
最后说一句，数据结构设计没有绝对的对错，只有适合不适合。关键是你得清楚你的场景是什么，瓶颈在哪里。是显存不够？还是计算速度慢？还是动态请求太多？想清楚了，再动手。