小红书把 KV Cache 按「头」拆了——长文本推理的「稠密时代」可能正在结束

发布日期:2026-06-29 · 分类:产品发布/更新 · 标签:AI 推理基础设施来源:小红书技术(Redtech)公众号 arXiv 2606.06256原文链接:https://mp.weixin.qq.com/s/qRrZvL0aZzYI82djFSrLug论文地址:https://arxiv.org/abs/2606.06256事件内容6 月 29 日,小红书引擎架构部 AI Infra 团队开源了 RedKnot——一套围绕「按头拆分 KV Cache」重构的长文本推理引擎。论文同步挂在 arXiv 上,代码已经公开维护。问题起点很朴素:今天所有大模型的 KV Cache 都是「稠密张量」——形状是 B, H, L, D。同一段 token 里,所有 head 被一起存、一起搬、一起算。但小红书团队的实证发现是:KV Cache 的价值不是按 token 均匀分布的,而是强烈按头分化的。有些 head 真的需要看完整上下文(比如负责全局语义对齐的 head),有些 head 主要只看局部窗口(比如负责局部语法模式的 head)。但当前的推理系统不管这一套——它把所有 head 绑在一起处理。RedKnot 做了三件事:1. 头分类稀疏。离线对每一对「层 头」做一次分类——少数是「全局头」(12-15%,需要看完整上下文),多数是「局部头」(85-88%,只看滑动窗口)。论文在 Mistral-7B、Qwen3-32B、Llama-3.3-70B、Qwen3.5-397B、DeepSeek-V4-Flash 上验证,局部头占比稳定在 83.4-96.8%。2. 稀疏 FFN。只对注意力得分最高的 top-k token 执行稠密 FFN,其他 token 走残差恒等路径。这条特别针对 agent 场景常见的 2-8K 片段——在那个长度下,FFN 才是预填充的真正瓶颈(57-62% TTFT),不是注意力。3. SegPagedAttention 存储。把稠密布局换成「按 (层, 头) 分段」的分页 KV 存储,并用一个融合的变长注意力内核,物理上只保留每个头真正需要的 token。每个 head 都留在 FlashAttention 快速路径上,不构造 attn_mask。三件事作用在「头、存储、通道」三条正交维度上,收益相乘而不是互相争抢余量。实测数据(8 卡 H800、跨 3 个模型、6 个 QA 数据集、8K-128K 上下文):TTFT 加速DeepSeek-V4-Flash 从 16K 的 3.51× 升到 128K 的 5.16×——加速比随上下文增长而提升,这是 token 级基线做不到的;Llama-3.3-70B HotpotQA精确匹配(EM)从稠密基线的 0.60 提升到 0.80;首字 top-1 / top-10 与稠密路径的一致性达到 0.93 / 0.87;单卡并发32K 上下文从 4 提升到 31;KV 传输字节PD 分离场景下最多省 6.3×;预填充阶段算力削减 67-79.5%。精度方面,整体精度通常 ≥ 稠密 F1 的 95%,甚至在长上下文区间能「反超」——因为它抑制了低价值 token 的噪声,同时保留了高质量的注意力结构。工程实现基于 SGLang。架构无关,同一套运行时覆盖标准 GQA、Qwen3.5 这种混合注意力 MoE、DeepSeek-V4 这种 MLA 压缩注意力。深度剖析为什么「按头」这个视角这么重要?过去一年,长文本推理的优化主要走两条路:一是前缀缓存(prefix cache)——同一段对话开头的 KV 算一次复用多次,DeepSeek 的 MTP、Qwen 的 PD 分离都是这条路;二是位置无关 KV 复用(PIC)——相同文档片段出现在不同位置也可以复用 KV,CacheBlend、Epic、ProphetKV 都是这条路。但 RedKnot 的论文图 1 显示了一件尴尬的事:这些 PIC 方法「不可靠」的原因,在于它们在 token 粒度上挑选要重算的子集——但不同的 head 关注的 token 子集并不相同。要同时满足所有 head,系统只能取各头重要 token 的并集——这个并集往往覆盖了片段中很大一部分 token,复用的意义被抵消了。换句话说:「按 token」是个错粒度。从一开始,「按头」才是工作负载的真实稀疏结构。为什么 FFN 优化在短上下文 agent 场景里特别重要?RAG、coding agent、长会话系统——这些场景的 prompt 通常在 2-8K token 区间。在这个区间,注意力计算本身不是瓶颈,FFN 占 TTFT 的 57-62%。即使把注意力重算降到零,FFN 这块成本也动不了。但 FFN 又很难稀疏化——如果直接砍掉某些 token 的 FFN,会破坏残差流。所以 RedKnot 的解法是「只对注意力得分最高的 top-k token 跑稠密 FFN,其他走残差恒等」——这是一种「用注意力作为 FFN 稀疏化的路由」的设计。「稀疏化即降噪」这件事比加速本身更反直觉。论文里有一句话值得反复读:「上下文越长,注意力越集中。RedKnot 在长上下文区间的准确率反而能「反超」稠密计算——因为它在抑制低价值 token 噪声的同时保留了高质量的注意力结构。」这意味着 RedKnot 不只是一个「更快的稠密计算」,它是一种「更好的计算」——长 prompt 里的低价值 token 真的在拖精度后腿。这件事 Anthropic 和 OpenAI 的工程师在内部肯定想过,但小红书把它做成了可工程化的开源实现。和 06-28 DeepSeek DSpark 的关系。昨天的 DSpark 是投机解码层面的加速(用小模型预测、大模型验证,整体生成速度 60-85%);今天的 RedKnot 是预填充层面的稀疏化(把稠密注意力 稠密 FFN 拆开,加速首字延迟)。两条路正交——一个解决生成阶段、一个解决输入阶段;一个从模型架构角度、一个从系统调度角度。中国大模型推理基础设施的开源节奏正在系统化。值得关注的原因长文本 agent 的成本结构可能重写。coding agent、deep research agent、客服 agent 这类「一次会话跑半小时」的场景,TTFT 和并发是核心成本。如果 RedKnot 的加速比在生产环境中得到验证,这些场景的单位成本可能降到原来的 1/3 到 1/5。「按头」这个视角会成为长上下文推理的通用语言。下一个 vLLM、SGLang、TensorRT-LLM 的版本,大概率都会引入类似的头分类机制。「稀疏化即降噪」打开了一个新研究维度——长 prompt 的精度问题可能不是「加更多上下文」,而是「更精确地用上下文」。小红书的 AI Infra 团队第一次完整亮相。之前他们的开源更多在 TensorFusion 这种训练侧,RedKnot 是第一个公开披露推理引擎团队规模和研究方向的 release。风险与待观察论文实验数据基于 H800,消费级 GPU(RTX 4090/5090)和国产卡的复现效果还没看到;「离线头分类」的稳定性需要长期观察——如果某些任务在分布漂移下让头分类失效,稀疏化反而会损害精度;「稀疏 FFN 走残差恒等」这个 trick 在数学上是否真的能保持模型原有的「梯度流」,还需要更多下游任务验证;论文 4.4 节承认「DeepSeek-V4-Flash 上仍有约 5% 的精度下降在某些长尾任务上」——这个数字在生产环境的容忍度,要看具体场景。