大模型推理做到一定规模后,你会发现瓶颈往往不在算力,而在显存。跑并发推理时,最长序列决定了最底层的显存分配——一个 4K context 的槽位可能只为 100 token 的请求服务了几秒就被释放。KV Cache 的碎片化、padding 带来的无效计算,这两个问题叠在一起,直接把整卡利用率拖下去了。cann-recipes-infer 里的 Paged KV Cache 和 Packed Sequence,就是专门处理这两件事的。这篇笔记顺着源码把这两套机制的来龙去脉、实现要点和调参经验都过一遍,适合已经看过推理引擎或者想自己动手改 Serve 层分配的工程师,也适合刚入行、想搞懂大模型推理为什么这么吃显存的人。
1. 为什么要同时聊 Paged KV Cache 和 Packed Sequence
1.1 先算一笔 KV Cache 的显存账
很多做推理优化的人,第一课都是算 KV Cache 的体积。公式很简单:KV Cache 只需要缓存 K 矩阵和 V 矩阵,所以每个 token 的显存开销是2 × num_layers × num_heads × head_dim × dtype_size。
以 7B 规模模型为例,常见配置是 32 层、32 个头、head_dim 128,如果跑 FP16(每元素 2 字节),那么:
每个 token 的 KV 大小 = 2 × 32 × 32 × 128 × 2 = 524288 字节 = 512 KB单序列生成 2048 个 token,KV Cache 就是 1 GB。如果按最大序列长度 4096 预分配,单个槽位需要 2 GB 显存。一张 40 GB 加速卡,连续跑几个长上下文请求,显存一下子就吃紧了。而实际情况往往更糟——很多请求实际长度只有几百 token,你按最长长度给它占了一个大坑,等请求结束释放时,显存已经碎片化得没法连续分配大块空间了。
这个账算完你就明白,KV Cache 是大模型推理延迟和吞吐的隐形天花板。它不吃算力,但吃显存带宽和容量。所以 Paged KV Cache 的思路非常直接:像操作系统管内存一样,把 KV Cache 切成固定大小的“页”,按需分配,而不是一次性给足。
1.2 静态 Padding 与动态请求之间的矛盾
除了显存浪费,还有计算浪费。静态批处理时代,一个 batch 里的所有序列都要 padding 到相同长度。假设某次 batch 有 8 条请求,实际长度分别是 128、200、512、1024、2048、256、64、1024,那么统一 padding 到最大长度 2048 后,batch 参与计算的总 token 数是8 × 2048 = 16384。但如果把所有序列拼成一个长数组,实际只需要128+200+512+1024+2048+256+64+1024 = 6256个 token。
两者对比,padding 方案里有超过 60% 的 token 是无效的。更致命的是,注意力计算是和序列长度平方相关的,padding 的部分同样参与矩阵乘和 Softmax,这不仅是显存浪费,还是在算力上白扔钱。
所以 Packed Sequence 解决的是“如何让计算只发生在真实数据上”。它把所有序列拼成一个连续的长序列,用一个cu_seqlens数组记录各序列的边界,注意力算子通过边界信息构造分块对角掩码,实现一个 batch 内多个序列并行且互不干扰。两个机制配合的原因也在这里:序列一旦打散重拼,KV Cache 的位置就不再是按序列连续排布的,而是需要一套“逻辑连续、物理离散”的管理方式,这正是分页机制擅长的领域。
2. Paged KV Cache:像操作系统一样管显存
2.1 把上下文切成固定大小的块
Paged KV Cache 的核心思想和虚拟内存分页一模一样。在传统方案里,每条序列拥有从seq_start到seq_start + max_len的连续显存区间。分页方案里,我们把显存预切成若干个固定大小的 KV Block,每个 Block 能存放block_size个 token 对应的 K、V 数据。
在 CANN 这一类异构架构上,KV Block 的数据排布一般是四维张量:
形状:[num_blocks, 2, block_size, num_heads, head_dim] 含义:[物理块号, 0表示K/1表示V, 块内token位置, 注意力头, 向量维度]为什么块大小不能太大也不能太小,后面第 5 节专门讲。先记住结论:经验值在 16~32 个 token 之间,关注内部碎片和块表开销的平衡。
2.2 块表:逻辑连续到物理离散的桥梁
每条序列在逻辑上仍然是一段从 0 开始递增的 token 位置,但这些位置映射到物理块时,可能东一块西一块。这个映射关系就存在 Block Table(块表)里。
块表本质上是一个二维数组:行是序列 ID,列是该序列按顺序占用的物理块号。举个例子,序列 3 的逻辑 token 位置 0~15 落在物理块 7 上,位置 16~31 落在物理块 29 上,那么块表的block_table[3] = [7, 29]。
// 简化的块表条目示意 struct KVBlock { int32_t block_id; // 物理块号 int32_t seq_id; // 属于哪个序列,-1 表示空闲 int32_t token_used; // 当前块内已写入的 token 数 int32_t next_block; // 下一块的物理块号,形成链表 }; struct SequenceMeta { int32_t seq_id; int32_t logical_len; // 当前逻辑长度 int32_t head_block; // 逻辑首块 int32_t tail_block; // 逻辑尾块 int32_t block_count; // 已经占用的块数 };有了块表,查询一个 token 的物理位置就是两次寻址。假设block_size = 16,要查序列 s 的逻辑第 t 个 token:
block_idx = block_table[s][ t / 16 ] offset_in_block = t % 16 物理地址 = block_idx × block_size × kv_per_token + offset_in_block × kv_per_token整个计算过程非常简单,和查字典差不多——先翻到第几页,再在页内找偏移。
2.3 分配、释放与回收
Paged KV Cache 的管理动作有三个:分配、追加、释放。分配发生在序列进入推理的 prompt 阶段,根据 prompt 长度一次性申请ceil(prompt_len / block_size)个块。追加发生在 decode 阶段,每生成一个 token 就往当前块里塞,如果当前块满了就再领一个新块挂到链尾。释放发生在序列结束时,把占用的所有块依次放回空闲块链表。
std::vector<int32_t> allocate_blocks(int num_needed) { std::vector<int32_t> blocks; for (int i = 0; i < num_needed; ++i) { int32_t b = free_list_.back(); free_list_.pop_back(); blocks.push_back(b); } return blocks; } void append_token(int32_t seq_id) { SequenceMeta& meta = seq_meta_[seq_id]; KVBlock& tail = blocks_[meta.tail_block]; if (tail.token_used == block_size_) { int32_t new_block = allocate_blocks(1).back(); tail.next_block = new_block; meta.tail_block = new_block; meta.block_count++; blocks_[new_block].token_used = 1; } else { tail.token_used++; } meta.logical_len++; }上面是高度简化的示意,真实源码里还要考虑多线程并发、CUDA/Ascend 上 Host 与 Device 之间块表同步、内存对齐等问题。但核心脉络就是这样,并没有想象中复杂。
3. Packed Sequence:消掉 padding 的无效开销
3.1 数据布局:一个长数组加一个偏移数组
Packed Sequence 的实现思路,其实一句话就能说清楚:把所有序列的 token embedding 按顺序塞进同一个大数组,然后用一个累计长度数组标出每条序列的边界。
// 4 条序列,长度分别是 3、5、2、4 // packed 后的 token 总数为 14 int32_t cu_seqlens[] = {0, 3, 8, 10, 14}; // 序列 i 在 packed 数组中的起始位置 = cu_seqlens[i] // 序列 i 在 packed 数组中的结束位置 = cu_seqlens[i + 1] // total_tokens = cu_seqlens[num_seqs]和静态 padding 方案相比,这个布局把之前浪费的计算资源全部省回来了。FlashAttention 一类的融合注意力算子天然支持这种输入,它内部会根据cu_seqlens构造出一个分块对角掩码:同一个序列内部保留 causal mask,不同序列之间完全屏蔽。
3.2 注意力算子如何处理多个序列的边界
注意力算子在拿到 packed 输入和cu_seqlens后,每个线程块需要知道自己服务的是哪条序列、序列内部的范围是什么。这里有个容易忽略的细节:Transformer layer 的输入 X 是整段长数组,但归一化、投影矩阵乘都是逐 token 计算的,所以不需要特殊处理;真正需要区分边界的只有 softmax 里的掩码和 KV 索引关系。
一个非常有用且常见的技巧是:在注意力内部,把序列偏移信息转换成每个 token 的“序列内位置”。这样在计算注意力分数时,如果query 的序列内位置 < key 的序列内位置,就直接屏蔽。由于cu_seqlens是单调递增的,算子内部可以用二分查找或者预处理索引来快速定位每个 token 所属的序列。
3.3 Packed 与 Paged 如何协同工作
如果你自己实现了 Paged KV Cache,再实现 Packed Sequence,你可能会发现两者是天然互补的。
Packed 后,一个 batch 里的不同序列长度完全不同,而且由于 KV Cache 是按物理块动态分配的,序列边界和物理块边界不一定对齐。假设块大小为 16,一条长度为 20 的序列会占用 2 个块,第二个块里只有 4 个 token 的有效数据。传统连续显存方案遇到这种场景会很别扭,但分页机制里每个块自带token_used计数,天然支持这种“部分填充”状态。
配合方式是:推理循环里,Host 端维护cu_seqlens和块表,kernel 端通过这两个数据结构组合寻址。对于序列 i 的第 t 个 token,先通过cu_seqlens得到它在 packed 数组中的逻辑位置,再通过块表得到它的物理 KV 位置。两套机制一套管数据流,一套管缓存流,互不干扰又缺一不可。
4. 关键代码路径与参数配置
4.1 显存规划:究竟开多少个块
代码里第一步往往是按显存余量算出能开多少个 KV Block。这个计算不能拍脑袋,要严格按公式走:
每块字节数 = block_size × kv_per_token 最大块数 = 可用显存 / 每块字节数其中kv_per_token就是第 1 节算出来的 512 KB(7B FP16)。如果block_size = 16,一块就是 8 MB。一张 32 GB 的卡,留出 4 GB 给激活值和权重,再用 20 GB 做 KV Cache,那么初始块数大约20GB / 8MB ≈ 2560块。这个数字在代码里往往就是 Cache Manager 初始化时需要传进去的那个max_block_num。
举个例子,如果一批并发请求的总上下文长度峰值是 40K token,那么理论最多需要40000 / 16 = 2500块。所以 2560 块刚好压线,一旦超过就得等前一个序列释放。这也是为什么要配合调度器做显存预留——给点余量更稳。
4.2 Prompt 阶段:批量分配,一次性写入
Prompt 阶段的特点是输入已知,KV 可以批量计算并写入。代码逻辑一般是:
- 根据当前序列的 prompt_len 计算需要
n = ceil(prompt_len / block_size)个块 - 从 free list 中批量取出 n 个块,按顺序填入 block_table
- 执行 attention prompt kernel,将 K、V 输出直接写进块内偏移对应位置
- 更新 SequenceMeta 的 tail_block 和 token_used
这步有个容易踩的坑:KV 写入必须和块内偏移对齐。如果 prompt 是 20 个 token,块大小 16,那第二个块只写了 4 个 token。写 kernel 时很容易默认每块写满,实际应该按min(block_size, remaining_len)来控制循环边界。
4.3 Decode 阶段:逐 token 产出,动态扩容
进入 decode 循环后,每个 step 只产生一个 token。代码里最频繁执行的动作是:
对每个活动序列: 1. 如果 tail_block 的 token_used == block_size,申请新块并挂到链尾 2. 将新 token 的 K/V 写入 tail_block 的 token_used 位置 3. token_used++ 4. 执行 decode attention kernel,读取该序列所有物理块中的 KV这个循环的高频路径短,但要注意锁的粒度。如果有多个并发 decode 流,free list 的 pop 操作必须原子化。我在调试时遇到过 free list 为空但显存还有余量的情况,原因就是别处的激活内存没有及时释放,导致 KV Cache 池被卡住。
4.4 调度与释放:把块还给空闲池
序列解码完成或被抢占后,要把它的物理块依次回收到 free list。源码里回收逻辑通常是遍历 block_table 里的块号,把seq_id置为 -1,token_used清 0,然后 push 进 free list 的头部或尾部。用头部回收可以保留热块,下次分配优先复用最近释放的块——这会在 Cache 局部性上有点收益。
这里值得注意:释放不是只发生在序列结束。如果做并行采样或多轮对话,这些场景里“序列”的定义要理清。每开一次新对话,就是一个全新的序列,旧序列的 KV 全部作废,但如果有 prefix cache 机制,前缀部分其实可以保留,只有新增的 suffix 部分需要重新分配。这部分优化能再省很多显存,不过代码复杂度会上升,属于进阶玩法了。
5. 实操踩坑与调参经验
5.1 三个最常见的坑
第一个坑是 block_size 太小。把块大小设成 8 或 4 时,碎片确实少了,但块表数量暴涨,kernel 内部需要循环处理的块数变多,索引开销和循环分支压力会抵消掉显存收益。实测经验是 7B 模型用 16 比较稳,13B 以上模型可以试试 32。
第二个坑是 padding 的残留数据污染。做 Packed Sequence 时,如果你在拼接前没有把上一个 batch 的末尾数据清零,或者cu_seqlens和实际 token 数不一致,注意力会出现跨序列的诡异关联。排查办法是单步调试时打印每个序列的cu_seqlens差分,再和实际长度比对。
第三个坑是 decode 阶段 free list 用完时没有优雅降级。有些实现直接返回错误导致推理中断,这在线下测试看不出来,但上了真实负载马上崩。更稳妥的做法是预留一小块“overflow 池”,或者在调度器里对当前活跃序列数量做上限控制。
5.2 关键参数速查
| 参数 | 作用 | 建议值/原则 |
|---|---|---|
| block_size | 每个物理块的 token 数 | 16~32,显存带宽足够多用 32 |
| max_block_num | 初始化块数 | 按可用显存 / 每块字节数计算 |
| free_list 预分配比例 | 提前从显存切出的块池 | 建议留总块的 5%~10% 做余量 |
| cu_seqlens 精度 | 偏移数组元素宽度 | 用 int32,int16 在长上下文下会溢出 |
5.3 性能收益能有多大
直接给一个保守参考区间:在混合长度负载(短请求多、长请求少)下,Paged KV Cache 能把 KV 显存利用率从 20%~30% 提升到 70% 以上;Packed Sequence 在长短差距大的 batch 里能省 30%~50% 的无效计算。两者叠加,端到端吞吐提升 1.4~2 倍是常见的。当然模型越大、层数越深,KV Cache 占比越高,收益就越明显。短到极端的负载(全是单条长序列)收益就会打折,因为本身没有碎片也没有 padding 浪费。
5.4 调参时的观测手段
调这些参数,不能光靠猜。最有效的办法是在代码里加两个计数器:一个是total_allocated_blocks,一个是blocks_actually_used,每 100 个 decode step 打一次日志。如果使用率长期低于 50%,说明块开多了或者 buffer 策略太保守。另一个值得加的观测点是块表的平均链长,链条太长说明批量分配不充分,每次只领一块,碎片率会上升。
我个人调试时还喜欢加一个“块表可视化”的小工具:把每个物理块的 seq_id 和 token_used 打印成二维表格,一眼就能看出哪里碎片多、哪个序列在跳着用块。这种看似笨的办法,比看 profiling 数据直观得多,排查块表越界和释放遗漏尤其好使。
6. 结尾:一点个人体会
Paged KV Cache 和 Packed Sequence 这套组合,本质上是在用操作系统的内存管理思路解决深度学习推理的显存难题。自己动手实现一遍之后,你会对 FlashAttention 的实现原理、KV Cache 的存储布局、以及调度器怎么估算并发容量都有更具体的体感。调试这类代码最忌讳一上来就追求性能,先把块表的分配-释放-复用逻辑跑对,再加上计数器验证显存利用率和计算有效比例,性能自然就上来了。后续如果想继续深入,可以在此基础上加 prefix cache、做 chunked prefill,或者把块表搬到 device 端减少查询开销。希望这份笔记能给你的源码阅读和自研推理优化省一些时间。