ASPLOS'26 | SpeContext 让小模型先圈出重要 KV Cache，长上下文推理吞吐最高提升 24.89 倍

August 12, 2026

2026   ·   LLM   KV Cache   长上下文推理   稀疏注意力   论文解读  

ASPLOS’26 | SpeContext 让小模型先圈出重要 KV Cache，长上下文推理吞吐最高提升 24.89 倍

原文：SpeContext: Enabling Efficient Long-context Reasoning with Speculative Context Sparsity in LLMs

1. 前言：长推理真正麻烦的不是 prompt，而是越想越长

现在很多 reasoning model 会先生成很长的思维链，再给出答案。对模型能力来说，这叫 test-time scaling：不继续扩大模型参数，而是在推理时多花一些计算，让模型一步步思考。

问题是，模型每多“想”一个 token，都要把这个 token 的 Key 和 Value 状态追加到 KV Cache 里。KV Cache 可以理解成模型对历史内容做过的中间笔记，有了它，生成下一个 token 时就不用把之前的文字重新计算一遍；代价是这本笔记会随着上下文长度线性增长。

以 Llama3.1-8B 为例，论文给出的数据是：32K 上下文仅 KV Cache 就要约 4GB。放到 24GB 的 RTX 4090 上，模型权重已经占去大头，再同时服务多个长请求，显存很快就不够了。

更麻烦的是，长上下文推理和“读一篇很长的文档再回答一个短问题”不是一回事：

  • 长上下文输入：prompt 很长，但输出可能只有几十个 token，KV Cache 的主体在 prefill 阶段一次性产生。
  • 长上下文推理：prompt 可能不长，模型却会连续生成几千甚至几万个 token，KV Cache 在 decode 过程中不断长大。

很多已有方案只在 prefill 结束后整理一次 prompt 的 KV Cache，之后新生成的 KV 全部保留。拿它们直接跑长推理，相当于只收拾了一次房间，后面东西继续往里堆，优化效果很快就没了。

SpeContext 想解决的就是这个问题：在长推理过程中，不再让大模型每层自己临时翻遍所有历史 KV，而是让一个小模型提前猜出“下一步最可能用到哪些历史 token”，把对应 KV 预取到 GPU。

这个思路和 speculative decoding 有点像，但猜的东西不同：

  • speculative decoding 猜的是“接下来会生成什么 token”
  • SpeContext 猜的是“接下来会关注哪些历史 token”

所以论文把它叫 speculative context sparsity，可以理解成“投机式地预测上下文稀疏性”。

2. 为什么 KV Cache 可以只读一部分

标准 attention 会让当前 query 和所有历史 key 做点积，再经过 softmax 得到注意力权重：

\[\text{AttentionWeight}=\text{softmax}\left(\frac{QK^\top}{\sqrt d}\right)\]

虽然公式把所有历史 token 都算了一遍，但最后的权重通常很不均匀：少数 token 拿走了大部分注意力，大量 token 的权重接近 0。

比如模型正在回答“太平洋为什么是世界上最大的海洋”，它可能重点关注前文里的“面积”“太平洋”“最大”等词，而不是平均关注每一个逗号和虚词。

这就给了 sparse attention（稀疏注意力） 一个机会：每次只取 Top-K 个最重要的历史 token 参与 attention。这里的 K 通常也叫 KV budget，表示这一轮、每个 attention head 最多允许读取多少个历史位置。budget 越大越接近 full attention，准确率通常越好，但显存读取和计算也越多。

如下图左边是普通 full attention，右边是只读取被选中 KV 的 sparse attention。下方两张散点图展示了速度和准确率之间的 Pareto frontier：越靠右表示吞吐越高，越靠上表示精度越好。

KV Cache 稀疏选择的收益

麻烦在于：你怎么知道该选哪 K 个 token？

Quest 会把 KV 切成 page，用每页的最大最小值估算重要性；ClusterKV 先做聚类，用聚类中心代替一组 KV；ShadowKV 则把 key 量化以后再计算近似分数。它们本质上都在做一件事：先把历史 KV 压成便宜的索引，再用当前 query 检索候选。

如下图右半部分列出了三条典型路线：paging、clustering 和 quantization。它们比遍历原始 KV 便宜，但都需要预处理。

现有 KV Cache 优化路线

这些方法处理长 prompt 很合适，因为 prompt 的 KV 在 prefill 后基本不再变化，预处理一次可以反复使用。但到了长推理，每一轮都会新增 KV，如果持续分页、聚类或者量化，预处理本身就会变成新的瓶颈。

论文总结了三个具体问题：

  1. 逐层检索太慢。 每一层都要先检索、再搬 KV、最后才能计算 attention，三者有数据依赖，只能串行。检索开销会随模型层数增长，最高占到 60% 延迟。
  2. 新生成的 KV 越积越多。 已有方法为了避免反复预处理，通常完整保留 decode 阶段的新 KV，正好错过了长推理里增长最快的部分。
  3. 静态 offloading 会出现性能悬崖。 推理开始前决定“全放 GPU”或“全放 CPU”，上下文稍微长一点导致显存不够，系统就可能突然把大量 KV 搬到 CPU，性能下降超过 80%。

这三件事决定了 SpeContext 不能只换一个检索算法，而要同时改算法、执行流水线和显存管理。

3. 核心 insight：小模型不仅会模仿答案，也会模仿“看哪里”

SpeContext 最有意思的地方，是把 knowledge distillation（知识蒸馏） 和 KV 检索连了起来。

知识蒸馏会让一个较小的 student model 学习大模型 teacher 的输出概率分布。比如大模型看到某段上下文后，认为下一个 token 是 A/B/C 的概率分别为 0.8/0.15/0.05，小模型训练的目标就是尽量给出相似分布。

论文的推理是：如果小模型和大模型最后给出的概率分布很像，它们从上下文里提取的重要信息也不应该完全不同。 否则小模型忽略了大模型真正依赖的信息，却还能长期模仿大模型输出，这件事很难成立。

作者用 mutual information（互信息，衡量两个变量共享多少信息）和 data processing inequality（数据处理不等式，信息经过处理不会凭空增加）解释了这个直觉：

  • 大模型输出和输入上下文高度相关
  • 蒸馏迫使小模型输出接近大模型
  • 那么小模型内部表示也必须保留足够多的上下文信息
  • 所以小模型 attention 关注的位置，可以作为大模型重要 token 的廉价预测

注意，这不是在说小模型每一层、每一个 head 的 attention 都和大模型完全一致。论文真正需要的是一个弱得多的条件：

小模型选出的 token，能覆盖大模型 attention 权重的大头即可。

这个差别很关键。SpeContext 不需要复制大模型的完整思考过程，只要让小模型当一个“划重点工具”。

4. 整体架构：先选位置，再边算边搬

下面这张架构图建议从上往下看。

SpeContext 系统架构

最上面是输入条件：模型结构、硬件规格和请求负载。系统先在 compilation 阶段根据这些信息计算显存切换阈值。

真正推理时，流程分成两条并行路径：

  • Lightweight Retrieval Head 读取当前输入，输出下一轮需要的 KV 位置索引
  • 原始 LLM 使用上一轮已经准备好的稀疏 KV 做 forward

中间的 Asynchronous Prefetcher 收到位置索引后，在另一条 CUDA stream 上把缺失的 KV 从 CPU 搬到 GPU。理想情况下，大模型还在计算当前 token 时，下一个 token 需要的 KV 已经在路上了。

整个循环可以压成下面几行：

# 当前 token 到来后，轻量检索头先预测下一轮会关注的位置
selected = retrieval_head(input_ids).topk(k=kv_budget)

# GPU 已有的 KV 不重复搬，只预取新增部分
to_load = selected - resident_kv_indices
async_prefetch(kv_cache[to_load])

# 大模型使用上一轮准备好的稀疏 KV 继续生成
next_token = llm.forward(input_ids, sparse_kv=resident_kv)

resident_kv_indices = selected
input_ids = next_token

这段伪代码省略了多层、多 head 和流水线 warmup，但核心关系就是：小模型负责预测位置，大模型负责真正计算；预测第 $t+1$ 步和执行第 $t$ 步重叠。

下面分别拆开三个组件。

5. C1：把 0.5B 蒸馏模型剪成只会“划重点”的 retrieval head

作者直接使用 EAGLE-3 训练好的 distilled language model（DLM，蒸馏语言模型）。完整 DLM 大约有 0.5B 参数，也有 embedding、attention、FFN、output projection 和 LM head。

如果把它完整跑一遍，额外开销大约是 20%。问题很明显：为了少读一点 KV，先额外跑一个不算小的模型，可能省下来的时间还没花出去的多。

但 SpeContext 根本不需要 DLM 生成 token，只需要它的 attention weight。所以作者把后半条链路几乎全剪了：

  • 保留 embedding
  • 保留 Query/Key projection
  • 保留用于计算 attention weight 的矩阵乘
  • 删掉 Value、attention output、FFN、output projection、LM head 和采样

最后得到约 0.03B 参数、约 60MB 的 retrieval head，参数量减少 90% 以上。

下图左上角的绿色部分是保留下来的路径，右侧灰掉的 FFN、LM Head、Logits 都不再计算。右边几组小图说明它还能适配 MHA、GQA、MQA 和 MLA 四种 attention 结构。

轻量 retrieval head 设计

这里还有一个容易忽略的设计选择：到底是整个 batch 共用一组重要 token，还是每个 attention head 各选各的？

作者最终选择 head-level retrieval。原因也很直观：不同 head 本来就可能负责不同模式，有的关注局部语法，有的关注实体，有的关注开头指令。把所有 head 强行合并成一份 Top-K，粒度太粗。

实验也支持这个判断。在长输入、budget=1024 时：

  • head-level 选出的 token 能覆盖大模型约 0.92 到 0.97 的 attention weight
  • batch-level 只能覆盖约 0.73 到 0.83

所以 SpeContext 虽然只跑一个单层 retrieval head，但它会给不同 head 保留不同的 token 集合。

对于 GQA（Grouped-Query Attention，多组 query head 共享一个 KV head），系统先对同一组里的 attention score 做 element-wise max，再按 group 选 Top-K；对于 MQA，所有 query head 共用一组 KV；对于 DeepSeek 使用的 MLA，则选择低维 latent cache，再只对命中的部分做升维。

6. C2：异步预取不难，难的是不要每次把 K 个 KV 全搬一遍

检索位置提前算出来后，SpeContext 就能把 KV transfer 和 LLM forward 放到两条 CUDA stream 上并行。

但并行不等于免费。如果每一步都从 CPU 重新搬 K 个 KV，而 PCIe 搬运时间比一层 LLM 计算还长，最后关键路径仍然会被 I/O 卡住。

论文测了不同 KV budget 下的预取延迟。budget 越大，搬运时间越长；到 512、1024 以后，完整预取已经明显超过单层模型计算时间。

预取开销与相邻命中重合率

转机在右图：相邻两个生成 step 选出的 token，高度重合，重合率超过 80%。

这很符合直觉。模型生成一句话时，连续两个 token 面对的上下文几乎一样，关注点不会突然全部换掉。比如刚生成“太平”，下一步生成“洋”时，依赖的历史信息大概率仍然是同一批。

因此 SpeContext 提出 elastic loading：

假设上一轮 GPU 上有：

S_last = {1, 8, 21, 35}

新一轮预测需要：

S_now  = {1, 8, 21, 50}

那就没必要重新搬 4 份 KV，只要把位置 35 替换成位置 50。实际加载量从 K 变成了两个集合的差集：

to_load = S_now - S_last = {50}

论文报告这种增量更新最多能减少 90% 数据传输。

下面这张时序图对比了五种路线。最下面 SpeContext 的绿色箭头表示 KV prefetch 和 FFN/attention 计算发生重叠，红色短条只搬相邻 step 之间变化的部分。

异步预取与弹性加载

从系统角度看，SpeContext 真正的收益不是“检索头比聚类快一点”，而是提前打破了 retrieve -> load -> attention 的逐层串行依赖。

7. C3：显存不够时，不要突然全量 offload

长推理还有一个动态问题：你很难提前知道模型最终会生成 2K、8K 还是 32K token。

传统静态策略可能在推理开始时发现 KV 放得下，于是全留 GPU；等长度跨过显存上限，再一次性切到 CPU offloading。结果就是上下文只增加一点，延迟却突然恶化。

SpeContext 的做法是按层渐进 offload。

系统根据四类信息预先算出一组序列长度阈值：

  • 模型权重和层数
  • KV head 数量与 head dimension
  • GPU 可用显存
  • 并发请求数和 KV budget

目标很直接：在不 OOM 的前提下，让尽可能多的层继续把完整 KV 留在 GPU。

例如 Llama3-8B 有 32 层。开始时 32 层 KV 都在 GPU；上下文超过第一个阈值，就把第 32 层 KV 放到 CPU，只在 GPU 留一个 budget 大小的工作区；再超过下一个阈值，就继续迁移第 31 层。

可以把运行时逻辑理解成：

while generating:
    while seq_len >= threshold[num_cpu_layers]:
        offload_kv(layer=last_gpu_layer)
        num_cpu_layers += 1

    next_token = llm.generate_one_token()
    seq_len += 1

它不是什么复杂调度算法，关键是用理论内存模型提前算好阈值，把“突然全搬走”的性能悬崖改成逐层、可预测的变化。

8. 效果：精度能不能保住

稀疏 attention 最大的风险永远不是速度，而是漏掉关键上下文。

论文先在 LongBench 上测试长文本问答。这里的 F1 score 衡量生成答案与标准答案的词级重合，越高越好；Passage Count 则测试模型能否从长文档里正确计数。

在 Llama3.1-8B 上，budget 只有 512 时，SpeContext 因为只做一次全局选择，精度略低于逐层检索的 ClusterKV；当 budget 达到 1024，SpeContext 基本追平 full attention，并超过几个 sparse attention baseline。

LongBench 准确率

长推理则使用 LongWriter，让 GPT-4o 从相关性、准确性、连贯性、清晰度、深度和阅读体验六个维度打分。这里的分数越高越好。

论文里有个挺有意思的现象：部分设置下，SpeContext 甚至略高于 full attention。作者检查后认为，full attention 的输出存在重复，而稀疏 attention 恰好缓解了重复。

LongWriter 生成质量

不过这个结果不要过度解读。它说明 sparse attention 未必必然伤害长文本生成，不等于删 KV 能稳定提升质量；这里还受到 GPT-4o judge、采样随机性和具体模型重复倾向的影响。

9. 性能：24.89 倍到底是怎么来的

云端实验使用 80GB GPU，重点测试多请求并发。下表里的 [2k, 32k] 表示输入 2K token、输出 32K token；括号中的数字是并发请求数，绿色倍数以 HuggingFace Eager 为基准。

云端多请求吞吐

以 DeepSeek-Distill-Llama-8B 的 [2k, 32k] 为例：

  • HuggingFace Eager：27.74 tokens/s，4 个请求
  • FlashInfer：314.25 tokens/s，8 个请求
  • SpeContext：690.59 tokens/s，32 个请求

换算下来，SpeContext 相对 Eager 是 24.89 倍，相对 FlashInfer 是约 2.20 倍。

这里一定要看清 24.89 倍的来源：它不只是单请求 attention kernel 更快，更重要的是 sparse KV 大幅减少显存占用，让同一张 GPU 能同时塞进更多请求。也就是说，这个数字同时包含了单请求加速和并发容量提升。

单请求结果如下图。左边是云端，右边是把 GPU 显存限制到 4GB 的端侧环境。论文报告端侧相对 HuggingFace Eager 最高加速 10.06 倍；与同样做 offloading 的 ShadowKV 相比，正文报告最高 1.17 倍。

云端与端侧单请求吞吐

消融实验把三个组件逐步叠加：

  • C1：轻量 retrieval head + sparse attention
  • C2：异步预取 + elastic loading
  • C3：自适应显存管理

SpeContext 消融实验

图里可以看到，C1 先降低 attention 计算和显存压力，C2 解决 CPU-GPU 搬运，C3 再通过尽量多留 GPU KV 提升速度。三层优化缺一块，长推理场景都会在别处重新卡住。

10. 一些冷静的讨论

这篇工作最让我觉得有意思的，不是又设计了一个 Top-K 算法，而是它换了检索信号：

过去从 KV 本身构造索引，SpeContext 则让一个与大模型同源的小模型直接预测大模型会看哪里。

这个视角把 speculative inference 从“预测结果”扩展到了“预测计算路径”，后面可以自然延伸到 attention block 预取、MoE expert 预取，甚至不同 memory tier 之间的数据调度。

但它还不是银弹：

  1. 需要对应的 DLM。 论文直接复用 EAGLE-3 的蒸馏模型。换一个没有现成 DLM 的模型，仍然要付训练成本。
  2. 互信息分析给的是合理性，不是严格保证。 输出分布接近，并不能数学上保证每个 attention head 都关注相同 token，最终可靠性还是靠实验验证。
  3. 小 budget 下有精度代价。 budget=512 时，全局选择不如 ClusterKV 的逐层选择；SpeContext 用更少同步换来了更粗的选择粒度。
  4. fallback 还只是 future work。 作者计划在 retrieval head 的注意力不够集中时退回 full attention，但当前系统还没有这个置信度兜底。
  5. 部分加速来自框架和更高并发。 24.89 倍是相对 HuggingFace Eager，若和 FlashInfer 比，最高提升约 2.20 倍。看系统论文的数字，baseline 和并发配置必须一起看。

总结一下，SpeContext 的核心逻辑其实很简单：

既然蒸馏小模型已经学会模仿大模型的输出，那就让它顺便提前告诉系统，大模型下一步大概率会读取哪些历史信息。

有了这个提前量，KV 检索从逐层串行操作变成了 LLM 前面的轻量预测；再配合增量预取和动态显存管理，长上下文推理里不断增长的 KV Cache 才真正被控制住。

顺带扯句题外话：SpeContext 属于 LLM 推理系统优化，书里没有直接讲这套 KV Cache 预取机制。不过它“用一个便宜代理预测昂贵模型行为”的思路，和 AutoML 里用 proxy 评估架构、用小成本信号减少完整执行是一脉相承的。我们把 NAS、搜索策略、LLM 剪枝量化和后训练压缩整理进了《动手学 AutoML：从 NAS 到大语言模型优化实战》，角度不同，但目标都是少做不必要的昂贵计算。

动手学AutoML书籍封面

Flag Counter