惯性聚合 高效追踪和阅读你感兴趣的博客、新闻、科技资讯
阅读原文 在惯性聚合中打开

推荐订阅源

腾讯CDC
T
Threatpost
Threat Intelligence Blog | Flashpoint
Threat Intelligence Blog | Flashpoint
cs.CL updates on arXiv.org
cs.CL updates on arXiv.org
T
Tenable Blog
AWS News Blog
AWS News Blog
Know Your Adversary
Know Your Adversary
TaoSecurity Blog
TaoSecurity Blog
P
Palo Alto Networks Blog
Spread Privacy
Spread Privacy
I
Intezer
Security Latest
Security Latest
The Last Watchdog
The Last Watchdog
Google DeepMind News
Google DeepMind News
Help Net Security
Help Net Security
Cyberwarzone
Cyberwarzone
N
News and Events Feed by Topic
O
OpenAI News
A
Arctic Wolf
S
Secure Thoughts
Attack and Defense Labs
Attack and Defense Labs
N
News and Events Feed by Topic
M
MIT News - Artificial intelligence
F
Full Disclosure
P
Privacy International News Feed
The GitHub Blog
The GitHub Blog
T
Troy Hunt's Blog
C
CXSECURITY Database RSS Feed - CXSecurity.com
H
Hacker News: Front Page
aimingoo的专栏
aimingoo的专栏
S
Security @ Cisco Blogs
H
Hackread – Cybersecurity News, Data Breaches, AI and More
Apple Machine Learning Research
Apple Machine Learning Research
Engineering at Meta
Engineering at Meta
Cloudbric
Cloudbric
大猫的无限游戏
大猫的无限游戏
Google Online Security Blog
Google Online Security Blog
Recent Announcements
Recent Announcements
H
Help Net Security
量子位
V
V2EX
美团技术团队
G
Google Developers Blog
www.infosecurity-magazine.com
www.infosecurity-magazine.com
S
Schneier on Security
V2EX - 技术
V2EX - 技术
D
Docker
博客园 - 【当耐特】
Project Zero
Project Zero
博客园 - 司徒正美

博客园 - marsggbo

LLM agent 为什么不稳?问题可能不在模型,在 harness Meta-Harness:让 LLM 自己搜索最优 harness,模型不动,性能白涨 从零理解 ASR:音频基础、Qwen3-ASR 架构,以及离线 vs 流式推理原理 arXiv'26 | LLM Agents 让群体信念变得可编程:当 AI 开始系统性操控舆论 巴西「主权大模型」翻车:当模型可以随便融合,怎么证明它偷了你的权重? 进阶篇 | 不靠人工设计,让遗传算法自己进化出 SOTA 的 LLM 剪枝指标 说人话:一文搞懂现在火热的 LLM agent 自进化原理 ICML'26 | Transformer 真的需要三个投影矩阵吗?Q-K=V 让 KV Cache 直接砍半 UPenn & Meta | NF-CoT:当 LLM 的思维链不再是文字,而是连续概率流 arXiv'26 | 为什么你的多智能体系统越加 agent 越慢?DeLM 用去中心化解了这个矛盾 arXiv'26 | Self-Harness:让 Agent 自己改自己的 harness,pass rate 最高翻倍 Anthropic | 当 AI 开始造自己:递归自我改进离我们有多远? arXiv'26 | Mirage:把世界模型的 3D 记忆搬进 Latent Space,快 10 倍还省 55 倍显存 arXiv'26 | FlashMemory-DeepSeek-V4:用 13.5% 的显存干 100% 的活,超长上下文推理的 less is more MoE 压缩新思路:别删专家、别合并专家,把它们"重映射"就够了 MoE 变 Dense:剪枝+蒸馏能救内存瓶颈吗? ReMoE:只动 Router 就让 MoE 推理快 2 倍?这才是端侧 MoE 部署该有的姿势 LLM 时代,还有人搞 AutoML 吗?有,而且变得更难了 arXiv'26 | Frontier:LLM 推理仿真器,端到端误差从 51.7% 降到 2.6% arXiv'26 | ELF:Flow Matching 生成文字,用 10 倍少的数据全面超越主流 Diffusion LM LLMRouterBench:当所有 routing 方法被拉到同一起跑线,结果有些尴尬 LLM Agent Memory 全景拆解:从 RAG 到 KV Cache 到参数写入,100+ 篇工作的方法演进与真实取舍 EuroSys'26 | TokenFlow:让 LLM 流式输出真正「流」起来 AAAI'23 | NAS-LID:用「局部内在维度」给超网做体检,省 86% 显存 ICLR'26 Workshop Spotlight | Lang-PINN:让 LLM 多智能体帮你从自然语言一键搭建物理信息神经网络 KDD'25 | BurstGPT:我们收集了 1031 万条 Azure OpenAI 真实 trace,LLM 推理系统没你想的那么稳 KBS'21 | 我写的 AutoML 综述被引 2700+ 次,今天来聊聊这篇文章的来龙去脉 EuroSys'26 | PARD 提前丢掉注定超时的请求,goodput 最高提升 176% EuroSys'26 | MFS 把整个 model family 融进一套嵌套模型,KVCache 跨 tier 直接共享 EuroSys'26 | LLMFolder 用常量折叠把 FFN 参数砍 80%,精度反超剪枝方法 65% EuroSys'26 | KUNSERVE 把冗余参数副本临时让给 KVCache,P99 TTFT 最快降 72× EuroSys'26 | IBP 用无损 bit 压缩缓解 PCIe 瓶颈,GNN/DLRM/LLM 推理都能用 FAST'26 | SolidAttention 把 SSD 搬进 LLM 推理,笔记本也能跑 128k 上下文 LLM 推理启动慢?华为用一个「可编程 Page Cache」把模型加载砍了 79% 延迟降47%!FineMoE如何用「细粒度」打破MoE推理的显存-延迟死局 训练一个「会管技能库」的 AI——SkillOS 让 agent 真正越用越强 MoE 训练通信瓶颈有救了?DySHARP 直接在交换机里做计算,干掉 50% 冗余流量 2026-05-12-2508_06526 把 Dense LLM 变成 MoE 还能推理提速?NeurIPS 2024 Read-ME 做到了 KV Cache 也能「语义共享」?SemShareKV 用 LSH 做到了 写完 Markdown 还要手动排版?我写了个 VS Code 插件一键搞定微信公众号、知乎、小红书 多 Agent 协作不需要说「人话」?LatentMAS 让 LLM 在隐空间里直接协作 让不同 LLM 之间共享 KV Cache?DroidSpeak 是怎么做到的 RouteMark: 基于路由行为指纹的模型合并知识产权归属 | A Fingerprint for IP Attribution in Routing-based Model Merging Lang-PINN: 从自然语言到物理信息神经网络的多智能体框架 | From Language to PINNs via a Multi-Agent Framework Ghost in the Cloud: 地理分布式大模型训练的安全隐患 | Your Geo-distributed LLM Training is Easily Manipulated GM-Skip: 基于度量引导的 Transformer 块跳过策略加速视觉语言模型 | Metric-Guided Transformer Block Skipping for Efficient VLMs ExpertFlow: 基于预测性专家缓存与令牌调度的高效MoE推理 | Efficient MoE Inference via Predictive Expert Caching and Token Scheduling AutoHete: 面向大语言模型的自动化高效异构训练系统 | An Automatic and Efficient Heterogeneous Training System for LLMs DAC'26 | ExpertFlow:让 MoE 大模型在单卡上跑起来,内存省 93%、速度快 10 倍 Eurosys26 | FineMoE如何用「细粒度」打破MoE推理的显存-延迟死局 LoRA fine-tune吞吐量提升1.96倍!LoRAFusion如何把内存带宽浪费和pipeline bubble一起干掉 Fast26 | LLM 推理启动慢?华为用一个「可编程 Page Cache」把模型加载砍了 79% KV Cache 的两层存储到底卡在哪?FAST'26 这篇论文给出了答案 NeurIPS24 | 把Dense LLM变身MoE还提速 ICML25 | EPIC:KV Cache 复用的「编译-链接」范式(附可运行代码复现) KV Cache 复用的第三条路:FAST 2026 CacheSlide 是怎么解决 Agent 推理的位置漂移问题的 MoE 推理的内存墙,被一块多芯粒芯片打穿了? KVCOMM:让多 Agent 系统的 KV Cache 真正“通起来”,TTFT 直接砍掉 7.8 倍 NSDI26 | DroidSpeak让不同 LLM 之间共享 KV Cache TokenDance 解决多 Agent LLM 推理的 KV Cache 冗余问题 当 AI 开始学会"记住":LLM Agent 记忆系统的统一视角 【转载】ACM MM 投稿论文模板修改成投稿模式 尝试从源头理解 SVD 原理和计算 LLM 场景下的强化学习技术扫盲 解决 Overleaf 中插入 PDF 图片失败的问题:排查与修复 Tmux ctrl+B快捷键失效处理办法 对抗训练综述学习笔记 【转知乎回答】一文看懂 LLaMA 中的旋转式位置编码(Rotary Position Embedding) 二进制中为什么负数是正数取反再加一 leetcode 常见题型代码总结 Prompt-Tuning、P-Tuning和Prefix-Tuning区别和代码实现【转】 Deepspeed ZeRO系列算法原理+通信开销详解 NSCC集群使用笔记 Huggingface Transformers实现张量并行的小坑 set/get_output_embeddings Pytorch 如何使用 storage 实现参数 offload? TACC 集群使用笔记 图解 vLLM 的推理调度策略 大模型推理框架 vLLM 源码解析(二):Block 模块分配和管理 OpenAI 的视频生成大模型Sora的核心技术详解(一):Diffusion模型原理和代码详解 大模型推理框架 vLLM 源码解析(一)
说人话理解 EPIC:KV Cache 复用的「编译-链接」范式(附可运行代码复现)
marsggbo · 2026-07-23 · via 博客园 - marsggbo

插播:之前写的《动手学 AutoML》终于出版了,从 NAS 到超参优化都有覆盖,适合想系统入门 AutoML 的同学。好了广告结束,现在进入正题。

动手学AutoML书籍封面

说人话理解 EPIC:KV Cache 复用的「编译-链接」范式(附可运行代码复现)

原文:EPIC: Efficient Position-Independent Caching for Serving Large Language Models


1. 前言

你有没有想过,当你用 RAG 系统给 LLM 塞了 5 篇文档 + 一个问题时,这 5 篇文档的 KV cache 其实在不同请求之间是可以复用的?

毕竟文档内容没变,变的只是用户的问题。但现实是,绝大多数系统只支持前缀匹配的 KV cache 复用——也就是说,只有当两个请求的开头完全一致时,缓存才能命中。这就很尴尬了:你 RAG 检索出来的文档顺序稍微变一下,或者 system prompt 改了一个字,前面缓存的 KV 全部失效,得重算。

这也是整个行业在 KV cache 复用上的核心矛盾:你想复用的 chunk(文档/few-shot 示例)在不同请求里出现的位置不一样,但 KV 向量里偏偏编码了位置信息(RoPE),导致"位置一变,缓存作废"。

今天想和大家聊聊这篇 EPIC,它把这个问题用一个非常漂亮的类比讲明白了——位置无关代码(PIC)。学过操作系统的同学应该不陌生:动态链接库之所以能被加载到内存的任意地址,靠的就是位置无关代码。EPIC 要做的,是让 KV cache 也能实现"位置无关"。

2. 核心问题:为什么 KV Cache 不能随便拼?

先交代下背景。LLM 的 prefill 阶段要给所有 prompt token 算 KV 向量,这一步是 compute-bound 的,也是 TTFT(Time-To-First-Token)的主要瓶颈。Context caching(也叫 prompt caching)的思路很简单:重复出现的 token 序列,把它们的 KV 向量缓存下来,下次直接用

vLLM、SGLang 这些推理框架都实现了 prefix caching——但问题是,它只能复用公共前缀

位置无关缓存 vs 位置无关代码的类比

如上图,EPIC 把 KV cache 的复用类比成编译和链接:

  • Compile 步骤:每个文档 chunk 被独立送进 LLM,生成 KV 向量并缓存(类比把 C 源文件编译成 .o 文件)
  • Link 步骤:使用时把多个 chunk 的 KV 缓存拼起来,加上用户 query,重算一部分 KV 来保证精度(类比链接 .o 文件生成可执行程序)

但这里有个关键问题:每个 chunk 独立编译时,position ID 都是从 0 开始的。这意味着拼接后,第二个 chunk 的第一个 token 的位置信息是"0",但它在完整 prompt 里的实际位置可能是 1024。

RoPE 位置编码直接编进了 K 和 Q 向量,位置信息的错位导致注意力计算出错。更具体地说,每个 chunk 开头的 token 会产生"attention sink"效应——它们会过度吸引注意力,让后面的 token 无法有效关注到真正重要的内容

3. 现有方案的痛点

在 EPIC 之前,CacheBlend 是第一个尝试解决 Position-Independent Caching (PIC) 的工作。但它有两个硬伤:

第一,复杂度太高。 CacheBlend 需要先在第一层完整重算所有 KV,通过对比 attention map 选出"变化最大"的 15% token,然后在剩余层只重算这 15%。但别忘了,第一层的完整重算就已经是 O(N²) 了——当 N 到达 35000 token 时,CacheBlend 直接 OOM。

第二,动态选择 token 的开销巨大。 CacheBlend 用的是 dynamic sparsity,需要运行时计算 attention map 来选 token,这个"选 token"本身就占了 TTFT 的 16% - 64%

TTFT 随 context 长度变化

如上图,随着 context 变长,CacheBlend-15 的 TTFT 以二次方增长,在约 35000 token 时 OOM;而 LegoLink-16 几乎是线性增长,50000 token 也没问题。

4. LegoLink:简单到离谱的核心算法

EPIC 提出的 LegoLink 算法,核心思想简单到让人拍大腿:

只重算每个 chunk 开头的 k 个 token(第一个 chunk 除外),让它们"意识到"自己不在序列起始位置。

就这么简单。k 通常设为 2~16,在论文实验中 k=2 就足以把精度损失控制在 7% 以内

但光这一句话可能还是抽象,下面我用一个具体例子把整个流程拆开讲。

假设一个 RAG 场景,用户问"Chrysan Company 的创始人是谁?",系统检索出了 3 篇文档:

Chunk 0 (系统指令, 4 tokens):  "你 是 一个 助手"       → position [0,1,2,3]
Chunk 1 (文档A, 5 tokens):     "苹果 公司 由 乔布斯 创立"  → position [0,1,2,3,4]
Chunk 2 (文档B, 5 tokens):     "Chrysan 公司 由 张三 创立" → position [0,1,2,3,4]
Query   (用户问题, 3 tokens):   "谁 创立 Chrysan"

⚠️ 注意看 position:每个 chunk 编译时都是从 0 开始的! 这就是 PIC 的关键——chunk 独立编码,不知道自己在完整 prompt 里的位置。

Step 1: Compile — 独立编码每个 chunk

每个 chunk 被独立送进模型,position 从 0 开始:

Chunk 0: tokens=["你","是","一个","助手"],   positions=[0,1,2,3]   → 生成 KV_0 (4个KV向量)
Chunk 1: tokens=["苹果","公司","由","乔布斯","创立"], positions=[0,1,2,3,4] → 生成 KV_1 (5个KV向量)
Chunk 2: tokens=["Chrysan","公司","由","张三","创立"], positions=[0,1,2,3,4] → 生成 KV_2 (5个KV向量)

每个 chunk 的 KV 向量都被缓存下来。

现在用户发来 query,我们要把 3 个 chunk 的 KV 拼起来。拼接后的完整序列长这样:

全局位置:  [0  1  2  3 | 4    5    6   7      8   | 9       10   11  12    13   | 14  15     16]
Token:    [你 是 一个 助手| 苹果 公司  由  乔布斯  创立 | Chrysan 公司  由  张三   创立  | 谁  创立  Chrysan]
                       chunk 0          chunk 1                  chunk 2             query

问题来了:chunk 1 的"苹果"这个 token 的 KV 向量里,RoPE 编码的位置是 0(编译时的位置),但它在完整 prompt 里的正确位置应该是 4。"Chrysan" 的 KV 里位置是 0,正确应该是 9

这导致了什么?每个 chunk 开头的 token 都以为自己是"位置 0"——序列的起点。在 causal attention 中,位置 0 的 token 会被后续所有 token 关注(因为它"最老"),吸收了大量注意力。这就是 attention sink

LegoLink (k=2) 的做法:

需要重算的 token(★ 标记):
位置:  [0  1  2  3 | ★4   ★5   6   7      8   | ★9      ★10  11  12    13   | ★14 ★15    ★16]
Token: [你 是 一个 助手| 苹果 公司  由  乔布斯  创立 | Chrysan 公司  由  张三   创立  | 谁  创立  Chrysan]
                       ↑ chunk1开头2个     ↑ chunk2开头2个       ↑ query全部

具体重算了哪些?

  • Chunk 0:不重算(它是第一个 chunk,位置本来就从 0 开始,没有错位)
  • Chunk 1 的前 2 个 token:"苹果"(pos 4)、"公司"(pos 5) — 用正确的全局位置重新算 KV
  • Chunk 2 的前 2 个 token:"Chrysan"(pos 9)、"公司"(pos 10) — 用正确的全局位置重新算 KV
  • Query 的所有 token:"谁"(pos 14)、"创立"(pos 15)、"Chrysan"(pos 16) — query 本来就要算

一共只重算了 7 个 token(2+2+3),而不是全部 17 个。

Step 3: 重算时的 Attention 计算

重算这 7 个 token 时,每个 token 的 Q 向量用正确的全局位置编码 RoPE,然后 attend to 所有 17 个位置的 KV(其中 7 个是刚重算的,10 个是缓存的)。

以 "Chrysan"(位置 9) 为例:

  • 编译时:它的 K 向量用了 RoPE(pos=0),所有后续 token 都觉得它是"起点",疯狂给它注意力 → attention sink
  • LegoLink 重算后:它的 K 向量用了 RoPE(pos=9),后续 token 能正确判断与它的相对距离,不再过度关注 → sink 消失

这就是为什么重算 chunk 开头的少数几个 token 就能解决问题:只有这几个 token 有严重的 attention sink,后面的 token 虽然位置也有偏差,但它们不在"位置 0"上,不会成为 sink

直观对比

Naive(不重算):
  query "Chrysan" 的注意力 → 80% 被 chunk 1/2 开头 token 吸走 → 找不到"张三"

LegoLink-2(重算 k=2):
  query "Chrysan" 的注意力 → 正确分配到"张三 创立"处 → 输出"张三"✓

Full Recompute(全重算):
  query "Chrysan" 的注意力 → 和 LegoLink 几乎一样 → 输出"张三"✓
  但代价是重算了全部 17 个 token 而不是 7 个

四种 PIC 算法对比

如上图对比了 Naive(不重算)、FR(全重算)、CacheBlend(重算 15% token)和 LegoLink(只重算每个 chunk 开头 k 个 token)。右边的 attention map 非常直观:

  • Naive:每个 chunk 开头有明显的亮色竖线(attention sink),注意力被这些 token 吸走了
  • FR:attention 分布正常,chunk 开头 token 虽然仍有一定 attention(因为 BOS token 的特殊性),但不再过度集中
  • LegoLink:通过只重算开头 k 个 token,attention map 几乎和 FR 一致

4.2 为什么只重算开头 k 个 token 就够了?

回到上面的例子。chunk 1 里的"由"(编译 pos=2, 正确 pos=6) 和 "乔布斯"(编译 pos=3, 正确 pos=7) 的位置也是错的,为什么不需要重算?

原因有二:

  1. Attention sink 是位置 0 的专利。在 causal attention 中,位置 0 的 token 是唯一被所有后续 token 都能看到的,自然成为注意力的"垃圾桶"。位置 2、3 虽然也有偏差,但偏差不会导致极端的 attention 集中。

  2. Query 和 decode token 才是信息聚合者。最终回答问题的是 query 和后续生成的 token,它们的 Q 向量是用正确全局位置算的,能正确 attend to 所有 KV(包括那些位置稍有偏差但内容正确的 token)。

4.3 复杂度对比

方法 重算 token 数 时间复杂度 说明
FR N=17 O(N²) 全重算,最准但最慢
CacheBlend-15 15%×N≈3 O(15%·N²) 还是 O(N²),因为第一层要全算
LegoLink-k (chunks-1)×k+q=7 O(k·N) ≈ O(N) k 是常数,近线性
LegoLink-0 0 O(q·N) 编译时处理 sink,link 零开销

5. 系统架构

EPIC 系统架构

EPIC 的系统架构很清晰:

  • KVCompile:接收用户提交的 immutable chunk,执行标准 prefill 生成 KV 向量,存入 KVCache,返回 cache ID
  • KVLink:收到请求时,根据 cache ID 取出 KV 缓存,执行 LegoLink 重算,然后正常 decode
  • KVCache:支持 HBM / DRAM / SSD 多级存储

用户通过两个 API 交互:

  1. generate_context_cache(chunks) → 获得 cache IDs
  2. chat_completion(cache_ids, query) → 获得回复

这种 explicit caching 的设计跟 Google Gemini 和 Mooncake 类似,用户自己管理缓存的生命周期。

6. 实验结果

6.1 精度-延迟 Pareto 前沿

精度 vs TTFT

如上图(6 个数据集 × 3 个模型),LegoLink(蓝色星星系列)在大多数场景下建立了新的 Pareto 前沿。核心发现:

  • LegoLink-2 就能把精度损失控制在 0-7%,同时 TTFT 比 CacheBlend-15 减少最多
  • CacheBlend-1 或 CacheBlend-5(重算类似数量的 token)精度会崩——因为它选错了要重算的 token
  • 增大 k 带来的精度增益递减:重算几个开头 token 就够了

6.2 系统级性能

延迟和吞吐

在异步工作负载下:

  • LegoLink-16 相比 CacheBlend-15 实现 最高 8× TTFT 降低
  • 吞吐提升 最高 7×
  • 随着 CCR(Context Cache Ratio)增加,LegoLink 的 TTFT 保持稳定,而 CacheBlend 持续波动

论文还提出了一个骚操作:LegoLink-0。在 compile 阶段给每个 chunk 前面加 4 个 dummy token(如 BOS token),编译完后把这些 dummy 的 KV 丢掉。这样 attention sink 在编译时就被"预消费"了,link 阶段完全不需要重算!虽然精度比 LegoLink-2 略差一点,但 link 开销为零

7. 代码复现

光说不练假把式。我用一个最简的 Transformer 模型复现了 LegoLink 的核心逻辑,不依赖 vLLM,只需要 PyTorch,CPU 就能跑。

代码在这里:https://github.com/marsggbo/easy-kvcache

git clone https://github.com/marsggbo/easy-kvcache.git
cd easy-kvcache
python epic/epic_legolink.py

核心实现大概 400 行,把 compile/link 两步和四种算法(Naive / FR / LegoLink / LegoLink-0)都实现了。这里挑几个关键部分讲一下。

7.1 Compile 步骤:独立编码 chunk

def compile_chunk(self, token_ids: torch.Tensor) -> CompiledChunk:
    """Compile:独立编码一个 chunk,position 从 0 开始"""
    chunk_len = token_ids.shape[1]
    # ⚠️ 关键:position IDs 从 0 开始
    positions = torch.arange(chunk_len, device=token_ids.device)
    _, kv_caches, _ = self.model(token_ids, positions)
    return CompiledChunk(token_ids=token_ids, kv_caches=kv_caches, chunk_len=chunk_len)

继续沿用上面 4.1 节的例子。假设我们有 3 个 chunk,每个 chunk 被独立编译:

输入 token_ids:           含义:
chunk 0: tensor([[10, 20, 30, 40]])       → "你 是 一个 助手"     (shape: [1, 4])
chunk 1: tensor([[50, 60, 70, 80, 90]])   → "苹果 公司 由 乔布斯 创立" (shape: [1, 5])
chunk 2: tensor([[100,110,120,130,140]])   → "Chrysan 公司 由 张三 创立" (shape: [1, 5])

对于 chunk 1,调用 compile_chunk(tensor([[50,60,70,80,90]])) 时内部变量:

chunk_len = 5
positions = tensor([0, 1, 2, 3, 4])    ← 关键:从 0 开始!不管 chunk 1 在 prompt 里实际从位置 4 开始

模型前向传播后,返回的 kv_caches 是一个 list,每个元素是一个 (K, V) tuple:

kv_caches = [
    (K_layer0, V_layer0),   # 每个 shape: [1, n_heads, 5, head_dim]  即 [batch, heads, chunk_len, dim]
    (K_layer1, V_layer1),
    (K_layer2, V_layer2),
    (K_layer3, V_layer3),
]

这些 KV 向量被存进 CompiledChunk,后续 link 时复用。注意 K 向量里已经编进了 RoPE(pos=0,1,2,3,4) 的位置信息——而正确的全局位置应该是 4,5,6,7,8。这个错位就是后续需要重算的根源

7.2 LegoLink:只重算开头 k 个 token

def link_legolink(self, query_ids, k=4):
    """LegoLink:只重算每个 chunk(除第一个)开头的 k 个 token"""
    # Step 1: 确定重算位置
    recompute_indices = []
    offset = 0
    for i, chunk in enumerate(self.compiled_chunks):
        if i > 0:  # 第一个 chunk 不需要重算
            for j in range(min(k, chunk.chunk_len)):
                recompute_indices.append(offset + j)
        offset += chunk.chunk_len
    # 加上 query tokens
    query_indices = list(range(total_ctx_len, total_ctx_len + query_len))
    all_recompute_indices = recompute_indices + query_indices
    
    # Step 2: 用 **正确的全局位置** 重算这些 token
    recompute_positions = torch.tensor(all_recompute_indices)  # 全局位置!
    
    # Step 3: 逐层执行——新算的 KV 替换缓存中对应位置
    for layer_idx, layer in enumerate(self.model.layers):
        # 计算新的 Q, K, V(用正确位置的 RoPE)
        Q, K_new, V_new = compute_qkv(x, recompute_positions)
        # 在缓存 KV 的对应位置覆盖
        K_exp[:, :, recompute_indices] = K_new
        V_exp[:, :, recompute_indices] = V_new
        # Q (k' tokens) attends to all N tokens
        attn = softmax(Q @ K_exp.T / sqrt(d)) @ V_exp

还是用上面的 3 chunk + query 例子,设 k=2,一步步看每个变量:

Step 1:确定重算位置

遍历 3 个 chunk:
  i=0, chunk 0 (len=4): 跳过(第一个 chunk 不重算), offset=0 → 0+4=4
  i=1, chunk 1 (len=5): 重算前 2 个 → recompute_indices=[4, 5], offset=4 → 4+5=9
  i=2, chunk 2 (len=5): 重算前 2 个 → recompute_indices=[4, 5, 9, 10], offset=9 → 9+5=14

total_ctx_len = 4+5+5 = 14
query_len = 3
query_indices = [14, 15, 16]

all_recompute_indices = [4, 5, 9, 10, 14, 15, 16]   ← 一共 7 个 token 需要重算

对应到完整序列中的位置:

位置: [0  1  2  3 | 4    5    6   7      8   | 9       10   11  12    13   | 14  15    16  ]
      [你 是 一个 助手|苹果  公司  由  乔布斯  创立|Chrysan  公司  由  张三   创立 | 谁  创立  Chrysan]
                     ★    ★                   ★       ★                    ★   ★     ★
                   重算chunk1前2个            重算chunk2前2个               重算query全部

Step 2:构造正确位置的重算输入

# 从完整 token 序列中取出需要重算的 token
all_token_ids = [10,20,30,40, 50,60,70,80,90, 100,110,120,130,140, 150,160,170]
                                                                    ↑ query tokens
recompute_ids = tensor([[50, 60, 100, 110, 150, 160, 170]])   # shape [1, 7]
#                       pos4 pos5 pos9 pos10 pos14 pos15 pos16

recompute_positions = tensor([4, 5, 9, 10, 14, 15, 16])       # ← 全局位置!不是从 0 开始

关键区别:编译时 "苹果"(token 50) 用的 positions=0,现在重算用的 positions=4。RoPE 编码不同,生成的 K 向量也不同,attention sink 效应就会消失。

Step 3:逐层重算并替换缓存

# 拿到缓存的 KV(3 个 chunk 拼接后)
K_cached shape: [1, n_heads, 14, head_dim]    # 14 = 4+5+5,三个 chunk 的 KV 拼一起
V_cached shape: [1, n_heads, 14, head_dim]

# 扩展为完整序列长度(加上 query 的位置)
K_exp shape: [1, n_heads, 17, head_dim]       # 17 = 14 + 3 (query)
V_exp shape: [1, n_heads, 17, head_dim]

# 复制缓存到前 14 个位置
K_exp[:, :, :14, :] = K_cached

# 用重算的 KV 覆盖对应位置
K_exp[:, :, [4,5,9,10,14,15,16], :] = K_new   # 7 个新 K 向量替换原来的
V_exp[:, :, [4,5,9,10,14,15,16], :] = V_new

# 现在 K_exp 中:
# 位置 0-3:  chunk 0 的原始缓存 KV(没动,因为第一个 chunk 不需要重算)
# 位置 4,5:  chunk 1 开头 2 个 token 的 **新** KV(用正确位置 4,5 重算的)
# 位置 6-8:  chunk 1 剩余 3 个 token 的原始缓存 KV(没动)
# 位置 9,10: chunk 2 开头 2 个 token 的 **新** KV(用正确位置 9,10 重算的)
# 位置 11-13:chunk 2 剩余 3 个 token 的原始缓存 KV(没动)
# 位置 14-16:query 的 KV(首次计算)

# Attention: Q (7个重算token) × K_exp^T (17个位置) → [1, heads, 7, 17]
attn_scores = Q @ K_exp.transpose(-2, -1) / sqrt(head_dim)
# 加 causal mask(每个 token 只能看到自己及之前的位置)
# 最终只取 query 部分(最后 3 个)的输出 → logits
def link_legolink_zero(self, query_ids, n_dummy=4):
    """LegoLink-0:compile 时加 dummy prefix,link 时零开销"""
    for chunk in self.compiled_chunks:
        # 在 chunk 前面加 n_dummy 个 dummy token(如 BOS)
        dummy_ids = torch.full((1, n_dummy), BOS_TOKEN_ID)
        padded_ids = torch.cat([dummy_ids, chunk.token_ids], dim=1)
        # 编译带 dummy 的版本
        _, kv_caches, _ = self.model(padded_ids, positions)
        # 丢掉 dummy 的 KV,只保留原始 chunk 的 KV
        kv_caches = [(K[:,:,n_dummy:,:], V[:,:,n_dummy:,:]) for K, V in kv_caches]

这个变体的思路很巧妙,还是用 chunk 1 举例:

原始编译(标准 compile):
  输入: ["苹果","公司","由","乔布斯","创立"]
  positions: [0, 1, 2, 3, 4]
  → "苹果"在 pos=0,成为 attention sink

LegoLink-0 编译(加 dummy prefix):
  输入: ["<BOS>","<BOS>","<BOS>","<BOS>", "苹果","公司","由","乔布斯","创立"]
  positions: [0,    1,     2,     3,       4,     5,    6,   7,      8]
  → 4 个 dummy 占据了 pos 0-3,"苹果"在 pos=4
  → dummy token "吃掉"了 attention sink
  
  然后丢弃前 4 个 dummy 的 KV:
  KV 原始 shape: [1, heads, 9, dim]
  丢弃后 shape:  [1, heads, 5, dim]   ← 只保留 "苹果" 到 "创立" 的 KV

这样 chunk 1 的 "苹果" 的 KV 向量里编码的是 pos=4(不再是 0),attention sink 在编译时就被消除了。link 阶段直接 naive 拼接就行,零重算开销

代价是 compile 阶段要多算 4 个 dummy token,但 compile 只做一次、link 做很多次,所以总体上是划算的。

7.4 运行结果

跑一下就能看到效果:

Step 5: 不同 k 值对 LegoLink 精度的影响

    k      Cosine Sim     Recompute Tokens   % of Total
  -------------------------------------------------------
    0          0.9975                    0         0.0%
    1          0.9995                   10         9.6%
    2          0.9996                   12        11.5%
    4          0.9996                   16        15.4%
   32          1.0000                   72        69.2%

和论文结论完全一致:k=2 就够了,继续加大 k 收益递减。

8. 一些个人 Take

  1. 类比很精髓。用编译/链接来类比 PIC 的 compile/link,这个抽象做得很漂亮。看完论文你会觉得"这也太自然了",但自然的东西往往是最难想到的。

  2. Attention sink 是个被低估的问题。之前 StreamingLLM 那篇讲 attention sink 更多是在 decode 阶段(evict 过程中保留 sink token),EPIC 把它放到了 PIC 的场景下,而且给出了比 StreamingLLM 更细粒度的解决方案。

  3. Static sparsity > Dynamic sparsity(在这个场景下)。CacheBlend 用 dynamic sparsity 去选 token 重算,看起来更"智能",但实际上选 token 的开销本身就很大。LegoLink 预先知道该重算谁(每个 chunk 开头),反而又快又准。这给了我们一个启发:不是所有问题都需要"动态"方案,如果你对问题结构足够了解,静态方案往往更优

  4. 这篇工作对 RAG 场景特别有价值。目前 RAG 系统的瓶颈之一就是长 context 的 prefill 开销。如果 document chunk 的 KV cache 可以跨请求复用,而且不需要 prefix 完全匹配,那 TTFT 可以大幅降低。可以预见未来的推理框架会原生支持 PIC。

  5. 代码质量。官方代码基于 vLLM 实现,改动了约 2K 行 Python,主要在 attention backend 和 scheduler 两处。代码结构还是比较清晰的,感兴趣的可以 diff 一下和 vLLM 0.7.0 的差异。

欢迎评论区交流,如果觉得有帮助也欢迎 star 一下 easy-kvcache 项目,后续会继续加入更多 KV cache 优化算法的简化复现。