ScoutAttention: Efficient KV Cache Offloading via Layer-Ahead CPU Pre-computation for LLM Inference

发表时间: 2026-03 · arXiv:2603.27138 (DAC 2026)

原文: https://arxiv.org/abs/2603.27138

Qiuyang Zhang, Kai Zhou, Ding Tang, Kai Lu, Jiguang Wan (Huazhong University of Science and Technology, Wuhan, China); Cheng Li, Zhenyu Yang (Huawei Technologies, Hefei, China); Peng Xu (Zhejiang Lab, Hangzhou, China)

速读

一句话结论 提出了一种名为 ScoutAttention 的 GPU-CPU 协同 KV Cache 卸载框架,通过提前一层的 CPU 预计算和异步周期性回取机制,在精度损失不到 2.4% 的情况下,将长上下文大模型推理的解码吞吐量提升至现有卸载方法的 2.1 倍。

要解决什么问题 在长上下文大模型推理的解码(Decode)阶段,KV Cache 的显存占用极大地限制了批处理大小(Batch Size),导致吞吐量低下。为了打破显存容量瓶颈,现有的做法通常将部分 KV Cache 卸载到内存(DRAM)中,但这引入了两个新的致命卡点。第一条路线是基于回取(Recall)的方法,例如 InfiniGen,它在需要时将 Token 从内存预取回 GPU。然而,GPU 的 HBM 带宽与 PCIe 传输带宽之间存在巨大鸿沟,缓慢的 I/O 导致 GPU 在高达 61% 的时间内处于空闲等待状态。第二条路线是 CPU-GPU 协同注意力(Co-attention)方法,例如 HGCA,它直接在 CPU 上并行计算被卸载的 Token,从而避开 I/O 瓶颈。但由于 GPU 的注意力计算速度大约是 CPU 的 20 倍,这种并行机制又将瓶颈转移到了 CPU 算力上,导致 GPU 仍有 57% 的时间在等待 CPU 计算完成。这两种路线都无法真正打满 GPU 的利用率。

怎么做的 核心思路是构建一个 GPU-CPU 协同的块级稀疏注意力机制,并通过跨层流水线彻底掩盖 CPU 的计算延迟。该方法由三个关键部件构成。首先是 GPU-CPU 协同块级稀疏注意力。系统将 KV Cache 分块,仅在 GPU 上保留极度压缩的块摘要(Block Digest)和少量最重要的块,其余卸载到内存。计算时,GPU 通过查询向量与块摘要的内积选出 Top-k 个重要块。此时,驻留在 GPU 上的块由 GPU 计算,未驻留的块交由 CPU 进行近数据计算,最后在 GPU 上合并结果。由于相邻 Token 关注的重要块高度重合,CPU 只需要处理极少量的增量块。其次是提前一层的 CPU 预计算(Layer-Ahead CPU Pre-computation),这是绕开 CPU 算力瓶颈的核心。传统的并行计算是 CPU 和 GPU 同时计算同一层,而 ScoutAttention 让 CPU 提前一层开工。当 GPU 正在处理第 $i$ 层时,利用残差连接下相邻层输入高度相似的特性,直接用第 $i$ 层的输入 $X^i$ 乘上第 $i+1$ 层的投影矩阵 $W_Q^{i+1}$,近似预测出下一层的查询向量:

$$Q_{\mathrm{pred}}^{i+1} = W_Q^{i+1} X^i$$

利用这个预测值,系统提前找出第 $i+1$ 层需要 CPU 处理的块集合,并立即异步启动 CPU 计算。这样,CPU 获得了一整层(包含注意力、前馈网络等)的超长执行窗口,彻底消除了 GPU 的等待时间。最后是异步周期性 KV Cache 回取。随着解码步数增加,重要块的分布会发生偏移,导致 CPU 需要处理的未驻留块越来越多。为此,系统会定期将内存中的重要块重新拉回 GPU。为了不阻塞推理,这个 I/O 传输动作被设计为在某一层注意力计算完成后异步触发,利用直到下一解码步才需要该数据的 20 毫秒以上的时间差,完美隐藏了 PCIe 传输代价。

效果如何 实验基于 SGLang 框架搭建,采用 Qwen3-8B 评估精度,Qwen3-14B 评估性能,测试集为涵盖单文档问答、多文档问答、摘要等任务的 LongBench,上下文长度最高达 64K。对比基线包括三个:代表传统无卸载路线的 FullKV、代表 I/O 回取路线的 InfiniGen,以及代表 CPU 并行计算路线的 HGCA。在精度方面,由于预测查询向量与真实查询向量的余弦相似度极高(大于 0.93),且 CPU 仅处理少量边缘块,在 2048 的稀疏预算下,ScoutAttention 相比 FullKV 的平均精度损失仅为 2.1%。在性能方面,ScoutAttention 显著打破了显存与计算瓶颈。在 64K 输入长度下,其解码吞吐量达到了 FullKV 的 5.1 倍;在各种长度下,吞吐量最高达到现有卸载基线(InfiniGen 和 HGCA)的 2.1 倍。更关键的是,通过预计算和异步回取,ScoutAttention 将 GPU 因等待 CPU 或 I/O 而产生的空闲时间比例从基线的 57% 至 61% 暴降至仅 6%。该方法的代价与局限在于,使用预测查询向量替代真实查询向量会带来微小的精度折损(精度略低于 InfiniGen);此外,周期性回取的触发间隔需要通过离线分析来设定阈值,以在 CPU 计算开销和 I/O 传输量之间寻找平衡。

一、 主要贡献

在长上下文推理过程中,大型语言模型(LLMs)面临着严峻的GPU内存容量限制,其中KV缓存的内存消耗严重限制了解码阶段的批处理大小(Batch Size)。虽然现有的研究探索了将KV缓存卸载到DRAM的方法,但这些方法要么需要频繁的GPU-CPU数据传输,要么对CPU提出了大量的计算要求,导致系统在等待I/O操作或CPU处理完成时GPU利用率低下。

本文提出了ScoutAttention,一种新颖的KV缓存卸载框架,通过GPU-CPU协同的注意力计算来加速LLM推理。为了防止CPU计算成为系统的瓶颈,ScoutAttention引入了GPU-CPU协同的块级稀疏注意力机制,显著降低了CPU负载。与传统的并行计算方法不同,该框架采用了一种新颖的提前一层CPU预计算算法(layer-ahead CPU pre-computation),使得CPU能够提前一层启动注意力计算,并辅以异步定期召回机制以保持极低的CPU计算负载。实验结果表明,ScoutAttention在保持精度下降不超过基线 $2.4\%$ 的同时,实现了比现有卸载方法高达 $2.1\mathrm{x}$ 的加速比。

图1:不同KV缓存卸载方法的推理流水线。(a) 全注意力:长上下文下的高注意力计算。(b) InfiniGen:由于缓慢的KV缓存召回导致的流水线气泡。(c) HGCA:由于缓慢的CPU侧计算导致的流水线气泡。(d) ScoutAttention:通过提前一层的CPU预计算实现高效的推理流水线。
图1b
图1c
图1d

二、 背景知识与设计动机

LLM推理架构与阶段特征
大型语言模型(LLMs)利用深度Transformer块堆叠实现自回归序列建模【16, Attention is all you need, 2017, NIPS】。在推理过程中,该架构分为两个不同阶段:预填充(Prefill)阶段一次性处理整个输入前缀,生成并存储 $K$ 和 $V$ 向量到KV缓存中,由于工作负载高度并行,该阶段是计算密集型的;解码(Decode)阶段则自回归地生成token,重复访问不断扩展的KV缓存,因此是内存密集型的。这种资源需求的差异导致在同一硬件上共享两个阶段会产生性能干扰。预填充-解码分离架构(Prefill-Decode disaggregation)通过在独立的专用计算集群上运行这两个阶段来解决此问题,从而实现更好的资源利用率和更低的尾部延迟【24, DistServe: disaggregating prefill and decoding for goodput-optimized large language model serving, 2024, OSDI】【14, Splitwise: Efficient Generative LLM Inference Using Phase Splitting, 2024, ISCA】。

注意力机制的内在稀疏性
注意力机制具有固有的稀疏性,不到 $20\%$ 的token贡献了超过 $80\%$ 的总注意力权重。基于此观察,现有研究提出了稀疏注意力算法以减少计算量【22, H2O: heavy-hitter oracle for efficient generative inference of large language models, 2023, NIPS】【15, QUEST: query-aware sparsity for efficient long-context LLM inference, 2024, ICML】【18, Efficient Streaming Language Models with Attention Sinks, 2024, ICLR】【10, SnapKV: LLM knows what you are looking for before generation, 2024, NIPS】。其中,块级稀疏性(block-wise sparsity)备受关注。它将KV缓存划分为固定大小的块,并用块摘要 $K_{\mathrm{digest}}$ 来概括每个块。在计算注意力时,根据查询和 $K_{\mathrm{digest}}$ 的点积选择前 $k$ 个(top-$k$)重要的块。不同方法生成摘要的策略不同,例如Quest【15】使用通道级最小/最大池化,MoBA【11, MoBA: Mixture of Block Attention for Long-Context LLMs, 2025, arXiv】使用均值池化,NSA【21, Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention, 2025, arXiv】结合可学习的MLP来生成摘要。

缓解长上下文内存占用的现有策略及其瓶颈
为了减轻长上下文处理带来的巨大GPU内存占用,近期研究探索了将KV缓存卸载到DRAM的方法,主要分为以下两种策略:
* 基于召回的卸载方法及其I/O瓶颈:该方法涉及在需要计算时将KV缓存重新加载回GPU内存。例如InfiniGen【8, InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management, 2024, OSDI】将大部分token卸载到DRAM,仅在GPU保留关键token,并通过预测机制提前一层发起召回I/O。然而,互连带宽构成了严重瓶颈。对于80GB HBM GPU和CPU之间通过PCIe 4x16接口通信,有效I/O带宽仅约 $800\mathrm{MB}/s$(细粒度传输时)或 $15\mathrm{GB}/s$(粗粒度传输时),远低于 $1.9\mathrm{TB}/s$ 的HBM带宽。评估表明,在批处理大小为40时,InfiniGen导致GPU在 $61\%$ 的执行时间内处于空闲状态。
* 协同注意力方法及其计算瓶颈:卸载到DRAM的token直接在CPU上计算。例如HGCA【3, HGCA: Hybrid GPU-CPU Attention for Long Context LLM Inference, 2025, arXiv】在GPU上保留 $25\%$ 的滑动窗口token,其余 $75\%$ 在CPU上计算。虽然消除了I/O瓶颈,但由于GPU在解码阶段的注意力计算速度比CPU快约20倍,HGCA中的并行执行导致在批处理大小为40时产生约 $57\%$ 的GPU空闲时间。

图2:GPU和CPU之间的I/O带宽。
图3:HGCA和InfiniGen的低GPU利用率。

三、 方法细节

ScoutAttention的设计原则与架构
ScoutAttention是一个旨在实现高效KV缓存卸载的GPU-CPU协同稀疏注意力机制。为了最大化性能,其建立在三个关键设计原则之上:1)采用GPU-CPU协同注意力,以缓解PCIe带宽瓶颈;2)引入CPU侧预计算,以隐藏CPU计算延迟;3)与块级稀疏注意力集成,以减轻CPU的计算负担。在解码阶段,ScoutAttention将大多数不重要的KV块卸载到DRAM,仅在GPU上保留块摘要和少量重要块。在注意力计算期间,采用GPU-CPU协同稀疏注意力,并引入提前一层的CPU预计算算法以防止CPU计算成为瓶颈。此外,为了纠正随着解码进行的块重要性漂移,提出了一种异步定期召回机制,从而保持较低的CPU计算负载。

图4:ScoutAttention的整体架构。
图4:ScoutAttention的整体架构。

GPU-CPU协同的块级稀疏注意力
在现代AI服务器中,CPU的注意力计算吞吐量(约 $100\mathrm{GB}/s$)显著高于PCIe上的KV缓存传输吞吐量(约 $15\mathrm{GB}/s$)。基于此差异,ScoutAttention采用GPU-CPU协同注意力而不是KV缓存召回。与完全在CPU上执行稀疏计算的HGCA不同,该设计在CPU和GPU上共同执行协同的块级稀疏注意力。本文采用Quest【15】作为稀疏化方法,但也完全兼容DeepSeek NSA【21】等算法。具体操作流程为:ScoutAttention将不重要的KV块卸载到DRAM,GPU仅保留块摘要和固定的关键块子集。计算时,GPU首先通过计算查询与每个块摘要的点积来识别top-$k$块。接着,GPU处理驻留在其内存中的块,而未驻留在GPU上的那部分top-$k$块则由CPU计算。最后,使用FlashAttention算法【2, FLASHATTENTION: fast and memory-efficient exact attention with IOawareness, 2022, NIPS】【19, From Online Softmax to FlashAttention, 2023】在GPU上合并两台设备的中间结果。这种划分策略非常有效,因为重要块在相邻token之间表现出强烈的时间局部性(平均不到 $15\%$ 的重要块会发生变化)。由于GPU保留了前几步识别的重要块,CPU只需处理少量未驻留在GPU上的top-$k$块,从而大幅降低了计算开销。

图5:具有提前一层CPU预计算的GPU-CPU协同块级稀疏注意力工作流程。
图5:具有提前一层CPU预计算的GPU-CPU协同块级稀疏注意力工作流程。

提前一层的CPU预计算算法
为了进一步防止CPU计算成为瓶颈,提出了提前一层的CPU预计算(layer-ahead CPU pre-computation)技术,允许CPU侧的注意力计算比GPU执行提前一层开始。在GPU计算第 $i$ 层注意力时,ScoutAttention首先识别第 $i+1$ 层的top-$k$重要块,主动触发第 $i+1$ 层的CPU侧计算,使其与GPU处理第 $i$ 层并行运行。随后,GPU计算第 $i$ 层注意力,并与前一层触发的CPU侧计算结果合并。实现CPU侧预计算的关键挑战是如何在第 $i$ 层获取第 $i+1$ 层的注意力查询。受InfiniGen【8】启发,由于残差连接导致连续层输入高度相似,假设可以通过将下一层的查询投影矩阵 $W_Q^{i+1}$ 应用于当前层输入 $X^i$ 来近似 $Q^{i+1}$,得到预测查询 $Q_{\mathrm{pred}}^{i+1}$。通过在多个模型上的实验证实(如下表所示),预测查询与真实查询之间的余弦相似度始终保持在很高水平,验证了该预测查询为CPU侧预计算提供了可靠的近似。

模型 Qwen 3 8B Gemma 3 12B Llama 3.1 8B Mistral 7B GLM 4 9B
余弦相似度 0.94 0.93 0.96 0.97 0.94

表1:预测查询与真实查询之间的余弦相似度。

基于此,设计了如下预计算算法。当GPU计算第 $i$ 层时,它首先执行以下步骤:
1. 通过对当前输入 $X^i$ 应用投影矩阵 $W_Q^{i+1}$ 预测下一层查询表示 $Q_{\mathrm{pred}}^{i+1}$。
2. 计算 $Q_{\mathrm{pred}}^{i+1}$ 与下一层摘要键 $K_{\mathrm{digest}}^{i+1}$ 的点积,选择top-$k$块,记为 $B_{\mathrm{pred}}^{i+1}$。
3. 将 $B_{\mathrm{pred}}^{i+1}$ 与已在GPU内存中的块 $B_{\mathrm{gpu}}^{i+1}$ 进行比较,未在GPU上的块标记为 $B_{\mathrm{cpu}}^{i+1}$。
4. 触发 $B_{\mathrm{cpu}}^{i+1}$ 在CPU上进行异步预计算。
完成第 $i+1$ 层的设置后,GPU继续第 $i$ 层的计算:先执行GPU侧注意力产生 $A_{\mathrm{gpu}}^i$,再与上一层预计算完成的CPU侧输出 $A_{\mathrm{cpu}}^i$ 进行合并(Merge),得到最终的注意力输出 $A^i$。

# 算法1:提前一层的CPU预计算
1: 输入: 第i层输入 X^i, 查询投影权重 W_Q, 块摘要 K_digest.
2: 输出: 第i层注意力输出 A^i.
3: ⊲ 触发下一层的CPU侧预计算 ⊳ 
4: Q_pred^{i+1} = W_Q^{i+1} X^i  # 获取下一层的预测查询
5: B_pred^{i+1} = TopK(Q_pred^{i+1} (K_digest^{i+1})^T)  # 预测top-k块
6: B_cpu^{i+1} = B_pred^{i+1} \ B_gpu^{i+1}  # 驻留在CPU中的块
7: spawn CPUAttn(B_cpu^{i+1})  # 异步CPU预计算
8: ⊲ 计算当前层的注意力 ⊳ 
9: Q^i = W_Q^i X^i
10: A_gpu^i = GPUATTN(Q^i, B_gpu^i)
11: ⊲ 与前一层触发的A_cpu^i合并 ⊳ 
12: A^i = MERGE(A_gpu^i, A_cpu^i)
13: return A^i

我们的预计算策略没有与GPU计算同时执行CPU操作,而是利用了整个Transformer层的处理窗口(涵盖GPU侧注意力、前馈网络和QKV投影)进行CPU侧注意力计算。例如,在80G HBM GPU上运行Qwen3-32B(4k稀疏预算),注意力计算仅需 $300\mathrm{us}$,而完整的Transformer层需要 $900\mathrm{us}$。这种提前一层的预计算为CPU提供了比并行计算方法多3倍的处理时间,有效消除了GPU空闲期。

异步定期KV缓存召回机制
随着解码的进行,重要KV缓存块的集合会逐渐发生偏移。由于CPU处理未驻留在GPU上的块,而GPU保留预填充阶段后识别的重要块,这种偏移导致CPU处理的token比例越来越大(如图6a所示,CPU计算比例呈稳步上升趋势)。为了减轻性能下降,提出了异步定期KV缓存召回机制,通过定期从DRAM召回重要块来刷新GPU驻留的上下文。基于PCIe互连是细粒度传输显著瓶颈的认识,该设计围绕两个原则:最小化召回频率,以及将数据传输操作移出关键推理路径。具体而言,如果在解码步骤 $m$ 的第 $i$ 层决定触发召回,ScoutAttention会在第 $i$ 层注意力计算完成后异步发起I/O操作。这为传输提供了充足的时间窗口(通常超过 $20\mathrm{ms}$),因为召回的块直到下一个解码步骤 $m+1$ 的第 $i$ 层才会被GPU需要。
召回间隔是通过离线分析经验性确定的。通过跟踪不同层的CPU计算比例模式,建立了逐层的召回间隔。对于每一层,确定使测量比例保持在阈值 $\beta$ 以下的最大步数。该阈值由系统的GPU/CPU计算能力比和互连带宽决定。基于经验分析,默认阈值设定为 $12\%$。如图6b所示,该机制实现了平均 $8.2\%$ 的CPU计算比例,所有层的平均召回间隔为8.7。

图6:Qwen 3 8B解码步骤中的CPU计算比例(token数/预算)。(a) 无定期召回。(b) 有定期召回。
图6:Qwen 3 8B解码步骤中的CPU计算比例(token数/预算)。(a) 无定期召回。(b) 有定期召回。

四、 实验环境

  • 数据集:使用LongBench【1, LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding, 2024, ACL】评估长上下文理解准确率,包含Qasper、NarrativeQA、2WikiMQA、DuReader、GovReport、QMSum、SAMSum和PassageRetrieval共8个数据集,涵盖单/多文档问答、摘要和消息检索等任务,上下文长度最高达64k tokens。
  • 模型架构参数:准确率评估使用Qwen 3 8B模型;性能评估使用Qwen 3 14B模型。所有基线方法的稀疏预算固定为2048 tokens,ScoutAttention的块大小设为32。
  • 硬件配置:配备80GB HBM的GPU,36核CPU,GPU与CPU之间通过PCIe 4x16接口连接通信。
  • 软件配置:ScoutAttention基于广泛使用的推理框架SGLang【23, SGLang: efficient execution of structured language model programs, 2024, NIPS】实现。为了高效识别top-$k$块,基于FlashInfer【20, FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving, 2025, arXiv】实现了块级top-$k$选择CUDA内核。CPU侧注意力计算工作器基于IPEX (Intel Extension for PyTorch)【5】构建,并将CPU线程分组,每组处理批次中的一个序列以提高吞吐量。

五、 实验结果

实验基线包括:FullKV(标准Transformer)、InfiniGen(基于召回的KV缓存卸载方法)和HGCA(基于协同注意力的KV缓存卸载方法)。性能评估仅针对预填充-解码分离架构中的解码实例。

  • 准确率评估:如图7所示,ScoutAttention在各数据集上保持了稳健的准确率。由于使用预测查询进行CPU侧计算,与InfiniGen相比有轻微精度下降。但因预测查询与真实查询余弦相似度极高且CPU计算占比仅约 $8\%$,精度损失极小。相比FullKV,ScoutAttention在1024预算下平均精度下降仅 $2.5\%$,在2048预算下平均下降仅 $2.1\%$。
  • 不同输入长度下的整体吞吐量:如图8所示,ScoutAttention在不同序列长度下始终实现最高吞吐量。随着输入长度增加,由于FullKV严重的内存容量瓶颈,ScoutAttention的加速比稳步增长,在64k输入长度下达到FullKV的 $5.1\mathrm{x}$。对于HGCA和InfiniGen,严重的I/O和CPU计算瓶颈导致它们在8k长度下的吞吐量甚至低于FullKV;尽管在更长输入下超越了FullKV,ScoutAttention相比这两种方法仍实现了高达 $2.1\mathrm{x}$ 的加速。
  • 批处理大小和块大小的影响:如图9所示,在32k输入长度下,当批处理大小从16增加到32时,受限于CPU和I/O瓶颈,HGCA和InfiniGen表现出次线性扩展,仅获得 $1.31\mathrm{x}$ 和 $1.21\mathrm{x}$ 的加速。相比之下,ScoutAttention表现出更好的可扩展性,从批次16到32实现 $1.78\mathrm{x}$ 加速,从32到64实现 $1.48\mathrm{x}$ 加速。图10表明,增大块大小会减少块总数,从而减小用于块选择的摘要缓存大小,释放内存以支持更大的批处理大小,进而提升解码吞吐量。
  • 延迟细分分析:图11展示了端到端延迟细分(Idle代表因CPU计算依赖或PCIe数据传输导致的GPU停顿)。HGCA遭遇严重的CPU计算瓶颈,空闲时间占总延迟的 $57\%$;InfiniGen因严重的I/O瓶颈,空闲时间高达 $61\%$。ScoutAttention通过最小化CPU计算和采用CPU侧预计算,将空闲时间大幅降至仅 $6\%$。
  • 消融实验:图12量化了各项优化的性能贡献。通过提前一层的预计算(PC)隐藏CPU计算延迟,实现了 $1.39\mathrm{x}$ 的加速。同时,异步定期KV缓存召回(PR)通过降低CPU计算负载,额外提供了 $1.20\mathrm{x}$ 的加速。

图7:不同方法在LongBench上的结果。
图8:不同输入长度下各方法的解码吞吐量。
图9:不同批处理大小下的解码吞吐量。图10:不同块大小下的解码吞吐量。
图11:延迟细分。图12:消融实验。

六、 结论

本文提出了ScoutAttention,一种用于KV缓存卸载的高效GPU-CPU协同注意力机制。为了实现高性能的CPU侧计算,引入了新颖的提前一层CPU预计算算法,并配合异步定期召回机制。实验结果表明,ScoutAttention在精度仅下降 $2.1\%$ 的情况下,相比全注意力实现了高达 $5.1\mathrm{x}$ 的加速。