论文阅读：HeadInfer

HeadInfer 要解决的问题很具体：在一张消费级显卡（如 24GB 的 RTX 4090）上，跑百万乃至四百万 token 的长上下文推理。它的做法是把 KV cache 的卸载（offloading）粒度从整层压到单个 attention head，任意时刻只在 GPU 上保留一个 head 的 KV cache，其余全部放在 CPU 内存里，并且这个过程是数学等价的——不是量化，也不是丢 token。

背景：长上下文推理撞上显存墙#

Transformer 推理分两个阶段。prefill（预填充）阶段一次性处理整段输入 prompt，为每个 token 算出 key/value 并存进 KV cache；decoding（解码）阶段每步生成一个新 token，新 token 要对之前所有 token 的 key/value 做注意力。KV cache 的作用是缓存历史 key/value，避免每步重算，代价是它随序列长度线性增长。

上下文一长，KV cache 就成了显存的主要消耗者。论文给出的例子很有代表性：用 Llama-3-8B 在 BF16 下生成 100 万 token，总共需要约 207 GB 显存，其中：

模型权重：约 15 GB
激活（activation）：约 64 GB
KV cache：约 128 GB

KV cache 一项就吃掉 128 GB，是权重的八倍多。RTX 4090 只有 24 GB，连零头都放不下。

已有的 offloading 方法（如 FlexGen、LayerKV、Infinigen）把 KV cache 搬到 CPU 内存来突破显存限制，但论文指出它们在消费级显卡上仍然吃力，原因有两个：一是粒度太粗，即使把 KV cache 按层卸载，留在 GPU 上的那部分仍然比其他内存都大；二是这些方法只处理 KV cache，没有一并压下激活内存。两者叠加，24GB 的卡依然跑不动百万级上下文。

核心 idea：把卸载粒度降到单个 attention head#

多头注意力里，每个 head 的计算是相互独立的：head $h$ 的输出只依赖它自己的 query、key、value，各 head 算完再拼接起来。既然如此，注意力就可以按 head 逐个计算，任意时刻 GPU 上只需要驻留一个 head 的 KV cache，算完就把它换回 CPU、把下一个 head 换上来。

A_t^{(h)} = \mathrm{Softmax}\!\left(\frac{Q_t^{(h)} {K_{\text{cache}}^{(h)}}^{\top}}{\sqrt{d_k}}\right) V_{\text{cache}}^{(h)}

这样 GPU 上的 KV cache 占用被压到极小，而且因为只是改变了计算和存放 KV 的顺序，注意力的结果和标准推理逐位相同——HeadInfer 是无损的，区别于 H2O、StreamingLLM 那类靠丢弃 token 的方法，也区别于 KV 量化那类有损压缩。

对 Llama-3-8B 的 100 万 token 推理，这一招把 GPU 上的 KV cache 从 128 GB 压到 1 GB，总显存从 207 GB 压到约 17 GB（实测），相比 BF16 baseline 降低约 92%，从而让 4090 能跑起来。

三个维度的切分：sequence、layer、head#

要理解 HeadInfer 的粒度，可以把整个注意力计算看成沿三个维度展开的立方体：序列 $S$ 、层 $L$ 、头 $H$ 。不同的省内存方法，本质是在不同维度上决定”哪一部分留在 GPU、哪一部分放到别处”：

chunked-prefill（分块预填充）切的是序列维 $S$ ：prefill 时不一次性处理整段序列，而是分成小块逐块处理，降低瞬时的激活内存。
layer-wise offload（按层卸载）切的是层维 $L$ ：GPU 上只保留当前层的 KV cache，其余层放到 CPU，按需取回。
head-wise offload（按头卸载）切的是头维 $H$ ：在一层内部，GPU 上只保留一个 head（或一组 head）的 KV cache。

这三个维度是层层嵌套的——块是序列的子集，层是模型的子集，头是注意力层的子集——所以三种技术可以叠加。HeadInfer 同时用上三者：chunked-prefill 压激活，layer-wise 加 head-wise 一起把 KV cache 压到单个 head。

下图对比了 KV cache（按层 × 头展开）在三种策略下留在 GPU 上的部分：

KV cache 在 layer×head 维度上的卸载粒度

用一个变量把这件事量化：设 $\alpha$ 为留在 GPU 上的 KV cache 占总量的比例。全部驻留 GPU 时 $\alpha=1$ ；按层卸载时 GPU 只留一层， $\alpha=1/L$ ；而 HeadInfer 只留一个 head：

\alpha = \frac{H_{\text{on}}}{H_{\text{all}}} = \frac{1}{L \times H}

其中 $H_{\text{all}} = L \times H$ 是全模型的 head 总数。对 Llama-3-8B（ $L=32$ 层，GQA 下有 $H=8$ 个 KV head），按层卸载留下 $1/32$ ，按头卸载留下 $1/256$ ，后者比前者又小了 8 倍。这正是表中”layer-wise 的 GPU 端 KV cache 是 8 GB、HeadInfer 是 1 GB”的来源。此时 GPU 上的 KV cache 大小为：

S_{\text{on-GPU}} = 2 \times B \times S \times D_h \times \text{sizeof(dtype)}

$D_h = D/H$ 是单个 head 的维度，前面的 2 对应 key 和 value。

方法：ping-pong memory 与异步重叠#

按 head 卸载省了显存，但引入了新的开销：每算一个 head，都要通过 PCIe 把它的 KV cache 从 CPU 取上来、算完再送回去。PCIe 带宽（论文实验环境约 25 GB/s）远低于显存带宽（约 1 TB/s），如果传输和计算串行，就会被 PCIe 拖垮。

HeadInfer 用 ping-pong memory 把传输藏进计算里。它在 GPU 上预分配两块 head 大小的缓冲区，配合两条 CUDA stream 做异步搬运：当 GPU 正在计算当前 head 时，同时用一条 stream 把下一个 head 的 KV cache 预取（H2D）到 GPU，用另一条 stream 把上一个 head 的 KV cache 写回（D2H）到 CPU。三件事并行发生，PCIe 传输被计算掩盖，不产生阻塞式等待。

ping-pong memory：预取下一 head、计算当前 head、写回上一 head 三者并行

因为注意力各 head 独立，这种乱序的预取和写回不会影响结果的正确性。只要单个 head 的计算时间足够覆盖它对应的 PCIe 传输时间，整体吞吐就不会因为卸载而下降。

chunked prefill 与 adaptive head grouping#

chunked prefill 是配套的第二件事。prefill 阶段的激活内存正比于一次处理的序列长度，一次喂进百万 token 的激活自己就会撑爆显存。HeadInfer 把输入序列切成固定大小的块（默认 10K，来源见下一节的 roofline 分析），逐块建立 KV cache，把激活内存从”正比于全序列”降到”正比于一个块”。表中 HeadInfer 的激活内存只有 0.625 GB，就是这么来的。

adaptive head grouping（自适应头分组）解决的是另一头的问题：把 head 切得越细，卸载越省内存，但 PCIe 传输次数和 kernel 启动次数也越多，短上下文时这些开销会盖过收益。HeadInfer 的做法是按上下文长度动态调整分组的粒度——上下文短时把多个 head 合成一个大组一起搬（省开销），上下文长时才把组切细（省内存）。

以 Llama-3-8B 的 8 个 KV head 为例，论文用的自适应策略大致是：上下文 ≤ 500K 时所有 head 合成 1 组；500K–1M 分 2 组；1M–2M 分 4 组；2M–4M 分 8 组（每组 1 个 head，最细粒度）。分组越细，GPU 上同时驻留的 head 越少，能支持的上下文越长。

自适应头分组：上下文越长，分组越细，可支持的上下文越长

一个值得注意的边界情形是：当所有 head 合成 1 组（Head=8/Group=1）时，一层的 KV cache 作为整体一起搬运，这在功能上就等价于”layer-wise 卸载 + chunked prefill”。所以 layer-wise 卸载可以看作 HeadInfer 在分组最粗时的特例，HeadInfer 只是把粒度进一步细化到了 head。

roofline 分析：为什么卸载拖不垮 prefill#

一个自然的担心是：把 KV cache 搬到 PCIe 上，会不会让推理整体变慢？论文用 roofline 模型回答了这个问题，结论是要分 prefill 和 decoding 两阶段看。

roofline 的核心是把算子分成两类：当计算受限于显存/PCIe 带宽（搬数据的时间盖过算的时间）时是 memory-bound；当计算受限于 GPU 峰值算力时是 compute-bound。判断依据是算术强度（arithmetic intensity，每读一字节数据能做多少次浮点运算）。

roofline：prefill 在长序列下 compute-bound，decoding 始终 memory-bound

prefill 是 compute-bound 的。prefill 的计算量随序列长度 $S$ 呈平方增长（注意力是 $S^2$ ），而数据搬运量只随 $S$ 线性增长，所以序列一长，算的时间就压过搬的时间。论文的分析显示，拐点大约在 2K：只要块长 $S \ge 10\text{K}$ ，即使 KV cache 走 PCIe 卸载，prefill 仍然停在 compute-bound 区，卸载几乎不拖慢速度。这也是 chunk size 取 10K 的原因——既压住激活内存，又保证每个块的 prefill 仍在算力受限区。
decoding 是 memory-bound 的。解码每步只生成一个 token，计算量小，瓶颈在把整个 KV cache 读一遍。此时依赖 PCIe 带宽搬运 KV cache 会显著拉低吞吐，这是 HeadInfer 在长上下文解码时很慢的根本原因。
head-wise 本身不改变 roofline 的位置。因为各 head 独立，把注意力按 head 拆开计算，其算术强度和峰值性能与标准 FlashAttention 相同，adaptive 分组进一步保住性能。也就是说，卸载带来的代价主要落在 decoding 阶段的带宽上，而不是 prefill 的算力上。

效果#

显存占用。以 Llama-3-8B、100 万 token 上下文为例，各方法的显存拆解如下（GB，“总 KV cache”列含卸载到 CPU 的部分）：

方法	权重	KV cache（GPU）	激活	总显存	总 KV cache
Standard	15.08	128	64	207	128
Chunked Prefill	15.08	128	0.625	143	128
4bit KV-quant	15.08	32	64	111	32
Layer-wise Offload	15.08	8	64	87	128
HeadInfer	15.08	1	0.625	16.7	128

HeadInfer 同时压住了 GPU 端 KV cache（→1 GB）和激活（→0.625 GB），把总显存压到 16.7 GB，与实测 17 GB 吻合，是唯一能塞进 24GB 卡的方案。

最大上下文。在单张 RTX 4090（24GB）上，各模型能跑到的最大上下文（无损方法）：

方法	Llama-3-8B	Llama-2-7B	Mistral-7B	Qwen2-7B	Gemma-2-9b
Standard	25K	10K	30K	35K	10K
Chunked Prefill	30K	20K	40K	70K	10K
4-bit KV-quant	45K	30K	40K	50K	20K
Layer-wise offload	45K	60K	45K	50K	35K
HeadInfer	4096K	1024K	4096K	4200K	1300K

Llama-3-8B 从标准推理的 25K 扩到 4096K（约 160 倍），其他模型也普遍进入百万级。消融实验进一步拆解了粒度的作用：同样都配了 chunked prefill，分组从粗到细，最大上下文依次是 550K（1 组，等价 layer-wise）、1100K（2 组）、2100K（4 组）、4096K（8 组），最细粒度相比最粗有 8 倍的上下文扩展。

质量无损。在 LongBench v2 上，HeadInfer（1024K）的整体得分 30.2，在 Medium 和 Long 两档难度上最高；在 SCBench 的全部 9 项任务上都优于其他方法；在 Needle-in-a-Haystack 上与标准推理精度相同。这些结果印证了它的数学等价性——扩上下文没有牺牲准确率。

吞吐。以 Llama-3-8B 为例，100 万 token 时约 516 tokens/s（prefill）和 0.15 tokens/s（decoding）；2 万 token 时约 7210 tokens/s（prefill）和 6 tokens/s（decoding），prefill 与标准推理（7235 tokens/s）基本持平。可以看到 prefill 吞吐几乎不受卸载影响，而 decoding 吞吐较低，与 roofline 的判断一致。

扩展到 70B。借助跨 8 张 RTX 4090 的流水线并行（pipeline parallelism），HeadInfer 能让 Llama-3-70B 跑到百万级上下文：10K 块时约 950K，把块调小到 5K 后可达 1M。

局限与不足#

解码极慢。HeadInfer 解决的是”能不能跑得下”，而不是”跑得快不快”。百万 token 上下文下解码只有约 0.15 tokens/s（每个 token 要六七秒），四百万 token 的 prefill 更是要数小时级别。瓶颈是 decoding 阶段 memory-bound、依赖 PCIe 带宽，这是方法本身的结构性代价。
受 CPU 内存限制。KV cache 被搬到 CPU，最大上下文因此受制于 CPU 内存容量（论文环境约 512 GB 供 KV cache）。Llama-3、Llama-2、Mistral 的上限就是被 CPU 内存卡住的；要再往上，需要更大内存或进一步卸载到磁盘，论文把这留作后续工作。
对短上下文没有收益。ping-pong 的重叠只有在单 head 计算时间足够覆盖 PCIe 传输时才成立；短上下文下传输和 kernel 启动开销占比大，所以自适应分组会退回粗粒度，这一段区间 HeadInfer 相对标准推理没有优势。

扩展：结合 head-wise sparsity#

HeadInfer 是无损方法，但它可以和有损的稀疏注意力叠加，进一步提速。论文给出的方向是结合 duo-attention 式的 head-wise sparsity：把 head 分成两类——重要的 retrieval head 保留全部 token，streaming head 只保留最近的少量 token。由于 retrieval head 的 KV cache 大得多，HeadInfer 只把这些大的 retrieval head 卸载到 CPU，把小的 streaming head 留在 GPU。在 50% 稀疏度下，百万 token 的 prefill 从约 2054 秒降到约 1152 秒，解码从约 6.51 秒/token 降到约 3.28 秒/token，接近 2 倍加速。

论文还用一个例子说明了为什么无损这件事重要：在长文档里插入一句”HeadInfer 的幸运数字是 42”，再提问这个数字。靠丢 token 的 H2O、StreamingLLM 会因为把这句”看起来不重要”的话丢掉而答错，而 HeadInfer（以及叠加 50% 稀疏后的版本）能准确答出——因为它保留了完整的 KV cache，只是换了个地方存。

一句话总结#

HeadInfer 的贡献是把 KV cache 的卸载粒度做到了单个 attention head，配合 ping-pong 异步重叠、chunked prefill 和自适应分组，在数学等价的前提下把长上下文推理的 GPU 显存需求压到消费级显卡能承受的范围，让单张 4090 跑百万级、乃至四百万 token 成为可能。代价是解码速度和对 CPU 内存的依赖：它换来的是”跑得下”，而不是”跑得快”。