RedKnot - 小红书开源的长文本推理加速引擎
来源:互联网
时间:2026-07-01 15:58:06
RedKnot是什么?
但凡做过大模型长文本推理的,想必都知道那个让人头疼的KV Cache膨胀问题。说起来,这几乎成了长上下文场景下绕不过去的坎儿。RedKnot这个方案,是小红书开源的一把新钥匙,专门用来给长文本LLM推理加速。它的思路很有意思:不把KV Cache当作一个整体去处理,而是按注意力头的维度把它拆解开,离线给每个头打上标签——哪些是需要全局重算的“全局头”,哪些是可以局部复用的“局部头”。然后搭配一个叫SegPagedAttention的内存管理机制,实现按头分页存储,再配合一个能处理变长序列的内核,从根本上避免了attn_mask带来的性能惩罚。不仅如此,它还引入了一个稀疏FFN的设计,只对注意力分数高的token做稠密计算,直击短上下文场景下的FFN瓶颈。
核心功能一览
- :离线给每个注意力头贴上“全局”或“局部”的标签,这个分类在请求间非常稳定,在线运行时完全不需要额外开销。
头分类稀疏
- :就算相同的文档片段出现在非前缀位置,也能复用之前算好的KV Cache,打破了传统前缀匹配的枷锁。
位置无关的KV复用
- :只对注意力得分最高的top-k token做稠密FFN计算,其他的token直接走残差路径,短上下文下的FFN瓶颈就这么被削掉了。
稀疏FFN加速
- :KV Cache不再是一个稠密的张量,而是被改成了按段分页存储,物理上只保留每个头真正需要的token。
SegPagedAttention存储
- :在模型的浅层采用比较保守的策略——局部注意力加稠密FFN,保护早期的残差流;到了深层再启用全局头重算和稀疏FFN,在精度和效率之间找到平衡点。
弹性稀疏恢复
- :通过PROFILE、BUILDSTATE、SELECTVISIBLESTATE、EXECUTE这四个适配器接口,统一支持GQA、MoE、MLA等各种注意力架构。
架构无关的运行时
技术原理深度拆解
- :传统系统把KV Cache看成个稠密的四维张量,RedKnot却敏锐地发现,不同注意力头关心的上下文范围和重要性差异很大,所以沿着头这个维度把缓存系统拆开了。
按头解耦KV Cache
- :token级的恢复策略需要把所有头认为重要的token取并集,结果就是集合迅速膨胀。RedKnot改为按头恢复,只对少数全局头重算,局部头直接复用,避开了精度和延迟的两难困境。
头级稀疏 vs. Token级稀疏
- :短上下文(2–8K)下,FFN的计算占了首Token生成时间(TTFT)的57–62%,单做注意力优化根本够不着。RedKnot用注意力信号来挑选重要token执行FFN,这个优化和KV稀疏是乘法叠加的效果。
FFN与注意力的正交优化
- :稠密布局加上attn_mask,会让FlashAttention的快速路径失效,带来4.9–7.6倍的性能惩罚。SegPagedAttention按头分页存储,配合融合的varlen内核,整个过程不需要构造mask。
分页存储消除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,在线运行时直接查表,不增加任何推理延迟。
离线画像零在线开销
- :只对大约12–15%的全局头重算,85–88%的局部头直接复用,避免了token级并集膨胀导致的重算量失控,同时也消除了级联误差传播。
头级恢复替代token级恢复
- :只对高注意力分数的top-k token做稠密FFN,其余走残差路径,直接削掉了短上下文下占TTFT一半以上的FFN瓶颈,与注意力优化形成乘法叠加收益。
稀疏FFN正交加速
项目地址
- :https://github.com/rednote-machine-learning/RedKnot
GitHub仓库
- :https://arxiv.org/pdf/2606.06256
arXiv技术论文
与同类方案的对比
| 对比维度 | RedKnot | CacheBlend |
|---|---|---|
核心定位 |
按头解耦的KV Cache管理系统 | 位置无关的KV Cache混合复用系统 |
稀疏粒度 |
按注意力头(head) |
按Token |
KV复用范围 |
任意位置(位置无关PIC) | 任意位置(非前缀片段) |
在线开销 |
零 |
有(在线选择token子集并混合) |
FFN优化 |
稀疏FFN |
无 |
存储布局 |
按头分页 |
稠密张量 + attn_mask,逻辑稀疏 |
内核效率 |
全程无attn_mask,FlashAttention快速路径 | 构造attn_mask,SDPA慢路径,4.9–7.6× 内核惩罚 |
浅层恢复 |
浅层保守用局部注意力+稠密FFN,保护残差流 | 浅层token并集膨胀,需重算大量token |
精度表现 |
通常 ≥ 稠密基线F1的95%,长文本可反超 | 浅层误差大,精度依赖重算token比例 |
哪些场景能派上用场?
- :需要把数万检索片段拼进prompt,位置无关的KV复用加上头级稀疏恢复,能把长文本预填充延迟从数十秒降到数秒。
RAG长文档问答
- :连续调用几十次工具,上下文越积越多,稀疏FFN直接削减短片段下占首Token生成时间过半的FFN瓶颈。
编程Agent多轮工具调用
- :把用户记忆、工具输出和历史状态都纳入长上下文,按头复用可以让单卡并发从4路提升到30路以上。
长会话记忆系统
- :多个Agent动态交换、重排上下文片段,位置无关KV复用打破了“必须前缀匹配”的限制,避免重复预填充。
多Agent协作框架
- :局部头直接复用近期分页KV,全局头按需重算,在128K上下文下也能保持低首字延迟和高生成稳定性。
实时流式长文本生成