RedKnot - 小红书开源的长文本推理加速引擎

RedKnot是什么？

但凡做过大模型长文本推理的，想必都知道那个让人头疼的KV Cache膨胀问题。说起来，这几乎成了长上下文场景下绕不过去的坎儿。RedKnot这个方案，是小红书开源的一把新钥匙，专门用来给长文本LLM推理加速。它的思路很有意思：不把KV Cache当作一个整体去处理，而是按注意力头的维度把它拆解开，离线给每个头打上标签——哪些是需要全局重算的“全局头”，哪些是可以局部复用的“局部头”。然后搭配一个叫SegPagedAttention的内存管理机制，实现按头分页存储，再配合一个能处理变长序列的内核，从根本上避免了attn_mask带来的性能惩罚。不仅如此，它还引入了一个稀疏FFN的设计，只对注意力分数高的token做稠密计算，直击短上下文场景下的FFN瓶颈。

核心功能一览

• 头分类稀疏
  ：离线给每个注意力头贴上“全局”或“局部”的标签，这个分类在请求间非常稳定，在线运行时完全不需要额外开销。

• 位置无关的KV复用
  ：就算相同的文档片段出现在非前缀位置，也能复用之前算好的KV Cache，打破了传统前缀匹配的枷锁。

• 稀疏FFN加速
  ：只对注意力得分最高的top-k token做稠密FFN计算，其他的token直接走残差路径，短上下文下的FFN瓶颈就这么被削掉了。

• SegPagedAttention存储
  ：KV Cache不再是一个稠密的张量，而是被改成了按段分页存储，物理上只保留每个头真正需要的token。

• 弹性稀疏恢复
  ：在模型的浅层采用比较保守的策略——局部注意力加稠密FFN，保护早期的残差流；到了深层再启用全局头重算和稀疏FFN，在精度和效率之间找到平衡点。

• 架构无关的运行时
  ：通过PROFILE、BUILDSTATE、SELECTVISIBLESTATE、EXECUTE这四个适配器接口，统一支持GQA、MoE、MLA等各种注意力架构。

技术原理深度拆解

• 按头解耦KV Cache
  ：传统系统把KV Cache看成个稠密的四维张量，RedKnot却敏锐地发现，不同注意力头关心的上下文范围和重要性差异很大，所以沿着头这个维度把缓存系统拆开了。

• 头级稀疏 vs. Token级稀疏
  ：token级的恢复策略需要把所有头认为重要的token取并集，结果就是集合迅速膨胀。RedKnot改为按头恢复，只对少数全局头重算，局部头直接复用，避开了精度和延迟的两难困境。

• FFN与注意力的正交优化
  ：短上下文（2–8K）下，FFN的计算占了首Token生成时间（TTFT）的57–62%，单做注意力优化根本够不着。RedKnot用注意力信号来挑选重要token执行FFN，这个优化和KV稀疏是乘法叠加的效果。

• 分页存储消除mask惩罚
  ：稠密布局加上attn_mask，会让FlashAttention的快速路径失效，带来4.9–7.6倍的性能惩罚。SegPagedAttention按头分页存储，配合融合的varlen内核，整个过程不需要构造mask。

• 分层弹性策略
  ：模型浅层局部头占比高，语义选择性也弱，所以采用保守策略防止误差扩散。深层全局头占比上升，注意力也更集中，这时做稀疏化收益最大，精度损失也最小。

怎么上手使用？

• 离线画像分类
  ：对目标模型调用PROFILE接口，通过needle-in-a-haystack测试，离线把每个(layer, head)对分成全局头或局部头，生成稳定的Head Class Map，后续请求直接拿来用，零开销。

• 预构建可复用状态
  ：用BUILDSTATE接口把常用文档片段预计算成KV Cache，并按头维度分页存入Global/Local KV Pool，实现位置无关的离线预构造。

• 动态状态选择
  ：在线请求来临时，通过SELECTVISIBLESTATE接口，根据查询语义和Head Class Map，从缓存池里挑出需要复用的局部头KV和需要重算的全局头范围。

• 融合推理执行
  ：调用EXECUTE接口，全局头走完整注意力重算并写入Online KV Cache，局部头直接复用分页KV走局部注意力，同时对低分token跳过FFN计算走残差路径。

• 服务化部署集成
  ：从GitHub仓库 https://github.com/rednote-machine-learning/RedKnot 拉取基于SGLang实现的源码，按四大适配器接口规范接入现有推理服务栈即可上线。

核心优势在哪里？

• 按头解耦的粒度对齐
  ：突破了传统token级稠密KV Cache的思维定式，把存储、计算和恢复的粒度统一到了注意力头维度，和工作负载真实的按头稀疏结构完美匹配。

• 离线画像零在线开销
  ：全局头和局部头的分类在请求间高度稳定，只需要离线做一次profiling，在线运行时直接查表，不增加任何推理延迟。

• 头级恢复替代token级恢复
  ：只对大约12–15%的全局头重算，85–88%的局部头直接复用，避免了token级并集膨胀导致的重算量失控，同时也消除了级联误差传播。

• 稀疏FFN正交加速
  ：只对高注意力分数的top-k token做稠密FFN，其余走残差路径，直接削掉了短上下文下占TTFT一半以上的FFN瓶颈，与注意力优化形成乘法叠加收益。

项目地址

• GitHub仓库
  ：https://github.com/rednote-machine-learning/RedKnot

• arXiv技术论文
  ：https://arxiv.org/pdf/2606.06256

与同类方案的对比

哪些场景能派上用场？

• RAG长文档问答
  ：需要把数万检索片段拼进prompt，位置无关的KV复用加上头级稀疏恢复，能把长文本预填充延迟从数十秒降到数秒。

• 编程Agent多轮工具调用
  ：连续调用几十次工具，上下文越积越多，稀疏FFN直接削减短片段下占首Token生成时间过半的FFN瓶颈。

• 长会话记忆系统
  ：把用户记忆、工具输出和历史状态都纳入长上下文，按头复用可以让单卡并发从4路提升到30路以上。

• 多Agent协作框架
  ：多个Agent动态交换、重排上下文片段，位置无关KV复用打破了“必须前缀匹配”的限制，避免重复预填充。

• 实时流式长文本生成
  ：局部头直接复用近期分页KV，全局头按需重算，在128K上下文下也能保持低首字延迟和高生成稳定性。