RedKnot:小红书开源的长上下文LLM推理优化工具,分头KV缓存加速长文本服务

原创 发布日期:
62

一、RedKnot 是什么

RedKnot 是小红书AI Infra团队开源、面向长上下文大模型推理服务的高性能优化扩展方案。项目基于SGLang推理引擎做注意力层扩展,兼容SGLang全部原生能力;当前GitHub仓库处于初始化阶段,仅包含README说明文档,配套技术论文同步发布于arXiv平台。

传统LLM推理引擎统一存储全部注意力头的KV缓存,长文本场景下显存占用暴增、首字延迟(TTFT)高、并发能力差。RedKnot核心创新思路是将KV缓存沿注意力头维度拆分管理,搭配头感知KV复用、分段分页注意力SegPagedAttention、分层弹性稀疏三大核心机制,在不微调模型、不损失生成精度的前提下,大幅降低长文本推理显存开销、算力消耗与响应延迟。

二、功能特色

1. 头感知差异化KV缓存复用

自动区分全局/本地/检索/稠密四类注意力头,针对不同头的上下文依赖特性分配独立缓存策略,支持离线KV片段存储+RoPE位置重定位,重复上下文仅重算必要Token,最大化复用已有KV缓存。

2. SegPagedAttention 分段分页内存管理

重构传统PagedAttention,为每个注意力头分配独立内存页与专属页表,解决多头混合存储带来的无效内存加载问题,大幅减少显存碎片、提升缓存读写效率。

3. 分层弹性稀疏计算优化

分层差异化算力调度:模型浅层全量计算保障输出精度,深层启用稀疏FFN,根据Token注意力分数过滤低贡献文本,跳过冗余前馈计算,预填充阶段算力最高削减79.5%。

4. 原生兼容SGLang完整生态

无侵入式集成SGLang,完整支持Radix前缀缓存、连续批处理、FP8/AWQ/GPTQ量化、PD分布式分离、FlashAttention等成熟推理能力,无需重构原有服务架构。

5. 显著的线上性能收益(8卡H800实测)

  • 首字生成延迟TTFT加速1.6~3.54倍;

  • 单卡并发承载能力提升4.7~7.8倍;

  • Prefill预填充算力FLOPs降低67%~79.5%;

  • 超长上下文场景收益随文本长度线性放大,生成文本质量近乎无损。

RedKnot:小红书开源的长上下文LLM推理优化工具,分头KV缓存加速长文本服务

三、技术细节

3.1 核心底层痛点

传统推理框架将KV缓存统一为稠密张量[Batch, Head, SeqLen, Dim],所有注意力头共用同一组内存页:部分头仅依赖局部上下文,却被迫加载完整序列KV,造成大量显存与算力浪费;共享页表无法区分头的有效窗口,长文本场景内存碎片严重、并发上限极低。

3.2 四大核心技术实现

(1)注意力头分类机制

按照模型层、头固有属性将KV头划分为4类,匹配专属缓存策略:

  1. Global全局头:需要读取全部上下文,页表覆盖完整序列,完整存储KV缓存;

  2. Local本地头:仅依赖窗口内局部Token,仅缓存有限长度KV;

  3. Retrieval检索头:适配RAG检索片段,支持离线KV持久化复用;

  4. Dense稠密头:通用计算头,按需动态分配缓存长度。

(2)Head-Aware KV Reuse 头感知KV复用

  • 离线预存高频重复Prompt、检索片段的KV缓存;

  • 新请求匹配历史片段时,复用已存储KV,仅对新增Token重新计算;

  • RoPE重定位算法修正复用片段位置编码,保证语义与输出精度不下降。

(3)SegPagedAttention 分段分页注意力

区别于传统PagedAttention“多头共用内存页”:

  • 每个注意力头拥有独立分页内存池与独立页表;

  • 读取KV时仅加载当前头所需内存段,无需加载全部头数据;

  • 动态回收空闲内存页,将显存利用率提升至90%以上,减少碎片。

(4)分层弹性稀疏FFN

  • 浅层(前N层):对误差敏感,执行完整FFN计算,保证基础语义表征稳定;

  • 深层(后序层):语义收敛稳定,基于注意力分数筛选高权重Token,低权重Token直接走残差通路跳过FFN计算,大幅削减预填充算力开销。

3.3 部署架构

RedKnot不作为独立推理引擎,而是SGLang上层扩展模块,架构分层:

  1. 底层:CUDA算子、FlashAttention、分页内存分配器;

  2. 中间层:SegPagedAttention头分离缓存管理器、KV复用调度器;

  3. 上层:兼容SGLang服务API、OpenAI标准接口、批量推理调度逻辑。

四、应用场景

  1. RAG检索增强生成服务
    业务高频复用知识库检索片段,大量重复上下文,KV复用能力可显著降低重复预填充算力,提升问答并发。

  2. 超长上下文对话/智能Agent
    多轮长会话、工具调用累积超长历史上下文(64K/128K Token),解决长文本TTFT延迟高、单卡并发低的瓶颈。

  3. 文档解析、代码生成、长文本摘要
    单请求输入万级Token,Prefill阶段算力消耗巨大,弹性稀疏计算可大幅缩短首字输出等待时间。

  4. 企业私有化大模型服务集群
    多用户共享系统提示词、固定业务模板,前缀KV缓存复用降低集群显存占用,减少GPU硬件采购成本。

  5. 高并发ToC线上AI对话产品
    流量波动大、请求上下文长短不一,SegPagedAttention优化内存碎片,提升高峰期承载量。

五、使用方法

5.1 环境前置依赖

  1. 基础环境:Python 3.10+、CUDA 12.0+、PyTorch 2.3+;

  2. 依赖推理引擎:SGLang(最新main分支版本);

  3. 硬件:NVIDIA A100/H100/H800等支持分页注意力的GPU。

5.2 仓库拉取

git clone https://github.com/rednote-machine-learning/RedKnot.git
cd RedKnot

5.3 模块集成部署

  1. 将RedKnot注意力扩展文件拷贝至SGLang源码注意力模块目录;

  2. 编译CUDA自定义算子:python setup_cuda.py install

  3. 启动SGLang服务时启用RedKnot优化开关:

python -m sglang.launch_server \
--model your-long-context-llm \
--enable-redknot \
--redknot-head-split True \
--redknot-sparse-layer-start 18

参数说明:

  • --enable-redknot:开启RedKnot全套优化;

  • --redknot-head-split:启用SegPagedAttention分头缓存;

  • --redknot-sparse-layer-start:设置稀疏FFN起始模型层数。

5.4 业务调用

完全复用SGLang原生OpenAI兼容接口,业务代码无需修改,仅增加启动参数即可生效优化。

六、竞品对比

选取行业主流长上下文推理框架vLLM、原生SGLang与RedKnot横向对比,聚焦长文本KV缓存核心能力:

对比维度 RedKnot 原生SGLang vLLM
KV缓存管理机制 SegPagedAttention(按注意力头独立分页) RadixAttention(按Token前缀分页,多头共用内存) PagedAttention(基础Token分页,无头部拆分)
长上下文KV复用策略 头分类差异化复用+离线KV持久化 前缀树缓存,全头统一复用 基础Prefix Caching,复用粒度粗
稀疏算力优化 分层弹性稀疏FFN,深层Token级跳过 无原生稀疏FFN,需二次开发 无内置稀疏计算能力
长文本TTFT加速能力 1.6~3.54倍提升 1.2~1.8倍提升 1.0~1.4倍提升
单卡并发提升上限 最高7.8倍 最高3.2倍 最高2.5倍
底层依赖 SGLang扩展,不独立运行 独立完整推理引擎 独立完整推理引擎
适配场景 128K+超长上下文、RAG、多轮Agent 通用长短文本、结构化输出 通用高并发短文本推理
显存利用率(长文本) 90%+ 65%~75% 50%~60%

七、常见问题解答(FAQ)

Q1:RedKnot可以脱离SGLang单独部署使用吗?

A:不可以。RedKnot是基于SGLang开发的注意力层扩展优化模块,并非独立推理引擎,所有内存调度、批处理、服务接口均复用SGLang底层能力,必须搭配SGLang运行。

Q2:使用RedKnot优化后,大模型生成内容精度会下降吗?

A:几乎无损。项目通过RoPE重定位对齐复用KV的位置编码,同时浅层模型强制全量计算保障基础语义;公开论文实测,主流LLaMA、Qwen、DeepSeek长上下文模型输出准确率、一致性与稠密基线差距小于1%,部分检索场景效果小幅提升。

Q3:RedKnot适合短文本对话场景吗?

A:收益有限。RedKnot核心优化目标是64K Token以上超长上下文,短文本(2K以内)KV缓存占用本身很小,分头缓存、稀疏计算带来的性能提升不明显,短文本业务直接使用原生vLLM或SGLang即可。

Q4:目前GitHub仓库无完整代码,什么时候能获取可运行源码?

A:仓库当前仅完成初始化提交,仅上传README说明文档;团队会分阶段推送算子、调度代码、测试Demo,可持续关注GitHub仓库提交记录获取更新。

Q5:RedKnot支持INT4、FP8等量化模型吗?

A:完整支持。RedKnot仅改造KV缓存存储与计算逻辑,不干预模型权重量化流程,兼容SGLang全部量化方案(AWQ/GPTQ/FP8/INT4)。

Q6:单卡4090消费级显卡可以部署RedKnot吗?

A:可以运行,但性能收益受限。SegPagedAttention对高显存带宽、大显存A/H系列服务器GPU优化最优,4090显存容量有限,长文本场景并发提升幅度低于H100/H800。

Q7:已有的vLLM业务服务可以迁移至RedKnot吗?

A:需要迁移至SGLang基座。RedKnot无vLLM适配版本,原有vLLM业务需切换SGLang服务,再集成RedKnot优化模块,接口层无需修改。

八、相关链接

  1. GitHub仓库地址:https://github.com/rednote-machine-learning/RedKnot

  2. 论文arXiv地址:https://arxiv.org/abs/2606.06256

九、总结

RedKnot是小红书AI Infra团队针对当前长上下文大模型推理显存与算力瓶颈推出的开源优化方案,跳出传统按Token管理KV缓存的固有思路,创新性沿注意力头维度拆分缓存体系,融合头感知KV复用、SegPagedAttention分段分页内存、分层弹性稀疏三大核心技术,依托成熟SGLang推理引擎实现无侵入集成,在不改动模型权重、不损失生成质量的前提下,大幅降低超长文本场景的首字延迟、显存占用与算力消耗,显著提升单卡并发承载能力,尤其适配RAG检索、超长文档解析、多轮AI智能体等长文本业务,为企业低成本落地大规模长上下文LLM服务提供轻量化、高性能的开源优化方案。

打赏
THE END
作者头像
AI工具箱
一个喜欢收集AI工具的小萌新