小红书技术REDtech公众号文章

提速不降智,小红书DoSA将首token延迟无损减少近一半

百万 token 输入,首字还要等多久?小红书 DoSA 保持原始评分函数下的精确 Top-K,在三款模型上将首字等待缩短约 37%~48%。

小红书技术REDtech

公众号 · 13 篇文章


导读

稀疏注意力减少了主注意力的计算,负责筛选 token 的索引器却可能成为长上下文推理的瓶颈。小红书自研 DoSA 利用分数上下界,提前排除不可能进入 Top-K 的候选,再配合 GPU 内核优化减少实际耗时。在 GLM-5.2 的测试中,平均约 97% 的查询—key 候选组合无需进行全头精算。在约 1M token 的测试中,Dots3-Note-preview、GLM-5.2 和 DeepSeek-V3.2 的首 token 延迟相对各自的 SGLang 基线分别减少约 37%、46% 和 48%。

01

注意力稀疏了,为什么还是慢?

把一本长书或整套代码交给大模型,回答往往不会马上开始。生成第一个 token 之前,模型需要先处理输入,完成预填充。这段首 token 延迟(TTFT),就是用户感受到的首字等待。上下文越长,等待越久。


DeepSeek-V3.2 提出的 DSA 为主注意力减轻了负担:先由索引器(Indexer)给候选 token 打分,选出分数最高的 K 个,再让主模型集中计算这些内容。问题在于,Indexer 自己仍要遍历候选,并通过多个索引头计算分数。


在本文约 1M 上下文的测试中,Dots3-Note-preview、GLM-5.2 和 DeepSeek-V3.2 的 Indexer 分别占预填充 GPU 时间的 81.6%、73.0% 和 90.8%。负责节省注意力计算的筛选步骤,已经成为主要耗时。


图 1|长上下文预填充中的 Indexer 成本。左:占 GPU 时间的比例,包含查询/键准备、打分和 Top-K;右:占理论矩阵 FLOPs 的比例。


层间复用、减少索引头或按块共享,可以降低筛选成本,但这些近似方法可能改变选中的 token。我们的目标是保留原始选择结果。DoSA 提前证明一部分候选不可能进入 Top-K,跳过它们的完整打分,把计算集中到仍有机会入选的候选上。

02

都做 Indexer 加速,DoSA 有何不同?

表 1|稀疏注意力的索引器加速路线。近似路线可能改变原始 Top-K 选择结果;LiteTopK 通过算子融合减少中间写回,DoSA 进一步利用数学上下界跳过不必要的全头打分。


上述层间复用、减少索引头和按块共享等路线,通过减少独立打分或改变选择粒度来节省计算,但不保证与原始逐层、逐查询的完整评分 Top-K 等价。即使下游评测没有明显下降,也不能据此认定选择过程无损。


DoSA 的本质区别,是在保持原始评分函数的前提下,用数学上下界证明哪些计算可以省掉。少量头只负责构造保守上界,不会替代最终评分;只有被证明不可能进入 Top-K 的候选才被排除,其余候选仍用原始全部头精算。

03

先判断能否入选,再决定是否精算

理解 DoSA,只需要抓住两个量:进入 Top-K 的“门槛下界”,以及每个候选的“得分上界”。当一个候选的最高可能得分,仍低于这个门槛,它就可以被安全排除。


图 2|DoSA 的筛选与精算流程:计算筛选点积,用上界筛除候选,对保留候选计算全部头,再按原始规则归约分数并写回。灰色和紫色分别表示省去的计算与写回。


例如,要选出前 3 名,已知有 3 个候选的分数都不低于 80 分,最终第 3 名就至少有 80 分。另一个候选即使按上界估计也只有 75 分,便可直接排除。做出这个判断,并不需要先找到真正的前 3 名。


DoSA 先选出一组种子,用当前查询的全部索引头重新打分。种子中第 K 大的分数,就是门槛下界,因为已经有至少 K 个候选达到这个分数。种子可来自历史入选频次,也可在冷启动时由少量锚点查询产生;它们都需要针对当前查询重新计算分数。


接下来,为每个候选估计一个不会低于真实分数的上界。Indexer 的总分由多个头的贡献加权汇总,DoSA 保留少量重要头精确计算,再分别处理剩余的正、负权重头。


正权重头会增加分数。DoSA 复用保留头的点积结果,并为近似残差留出有数学保证的余量,避免低估这部分贡献。负权重头会扣分,因此只扣除能够保证“至少会发生”的部分;通过聚合查询向量,一次额外点积就能给出这样的下界。两部分合起来,构成候选总分的上界。


只有上界严格低于门槛,候选才会被排除;等于门槛的候选仍须保留。剩余候选继续计算全部头,最终按原始分数执行 Top-K。界估得宽松一些,只会多做精算,不会误删可能入选的候选。


种子里高分候选越多,门槛越高;保留头的贡献算得越准,上界往往越紧。两者都能提高过滤率,但准备种子、计算上界也有成本。因此,我们衡量的是完整 Indexer 的耗时,准备工作同样计入。


这里的“无损”,指保持原始评分函数下的精确 Top-K。上界只用于决定哪些计算可以跳过,最终选择仍依据原始分数。GPU 实现还加入数值保护,避免浮点计算误差破坏安全筛选条件。

04

少算之后,还要让 GPU 高效执行

减少点积后,如何让 GPU 更快完成剩下的工作?不同查询留下的候选不同,逐个处理会造成零散计算、重复加载 key 和大量中间结果写回。DoSA 将流程组织为规则筛选、候选精算和紧凑写回。


图 3|筛选、精算与写回的数据流。蓝、橙色为连续处理的数据块;灰色 MMA 块为省去的全头计算。G2S 表示全局内存到共享内存的传输,A2R 表示负责累加的硬件模块到寄存器的传输。


筛选阶段仍然扫描可见的 key,但只计算上界需要的少量信息。保留头数量相近的查询被编成规则数据块,交给 Tensor Core 做矩阵运算;数据加载与计算通过流水线重叠执行。筛选结果用紧凑的位掩码记录,留给下一阶段使用。


精算阶段利用查询之间的重叠。实现中每 4 个查询组成一组,如果某个 key 被至少两个查询保留,且对整组都可见,就进入共享流:加载一次,为整组查询打分。这会多算少量点积,但能减少重复搬运;其余候选仍按查询分别处理。


写回时,低于门槛的分数直接过滤,只保存有机会参与最终选择的记录。记录采用紧凑分数编码和精确 key ID,先用编码定位 Top-K 边界,再对边界上的候选比较原始分数,保留精确选择。


这也解释了 DoSA 与 LiteTopK 的区别。LiteTopK 融合打分与 Top-K 选择,减少中间数据搬运;DoSA 进一步跳过不必要的全头打分,把计算量也降下来。

05

百万 token 下,收益能有多大?

测试覆盖英文书籍、中文书籍和代码,输入长度从 128K 到约 1M token,主要来自 InfiniteBench,并用 Project Gutenberg 和 SQLite 源码补充。对比基线为 SGLang v0.5.18 和 LiteTopK,Top-K 均取 2,048。SGLang 基线先计算并写回完整候选分数,再执行独立的 Top-K 选择;LiteTopK 通过融合打分与选择减少中间写回;DoSA 进一步利用安全上下界跳过不必要的全头打分。


在 GLM-5.2 的完整 Indexer 微基准中,5 种上下文长度、4 种查询块大小与 3 个输入领域组成 60 种配置,计时包含 query/key 准备。DoSA 在全部配置中均快于两条基线:相对 SGLang 加速 1.12~2.74 倍,相对 LiteTopK 加速 1.17~1.69 倍。到 1M 上下文时,相对 SGLang 的加速达到 2.47~2.74 倍。


图 4|预填充完整 Indexer 相对 SGLang 的加速比。三行分别为英文、中文和代码;每行比较 4K、8K、16K、32K 查询块,覆盖 128K~1M 上下文。

逐层分析显示,三个领域的平均候选过滤率约为 97%,每个领域的 21 个 Indexer 输出层中,都有 20 层超过 95%。大多数查询与 key 的组合因此可以跳过全头精算。这里的过滤率衡量候选筛选效果,不等同于整个模型的计算量降幅。


图 5|GLM-5.2 各 Indexer 输出层的平均过滤率。英文、中文和代码三个领域均有 20 层超过 95%,共测试 21 层。


更接近用户体验的是端到端 TTFT。在约 1M 上下文下,结果如下,数值表示基线延迟与 DoSA 延迟之比。

三款模型采用的张量并行度与预填充块大小见表。逐请求测量,并在请求之间清空前缀缓存。加速比为基线延迟除以 DoSA 延迟,延迟降幅为 1 − 1 / 加速比。相对各自的 SGLang 基线,三款模型的首 token 等待分别减少约 37%、46% 和 48%,其中后两者接近减半。不同模型的部署配置不同,表中结果用于比较同一模型内的方法收益。


图 6|Dots3-Note-preview 的首 token 延迟(越低越好)。采用 TP4、16K 预填充块;约 1M 上下文下,相对 SGLang 加速 1.58 倍,首 token 等待减少约 37%。


图 7|GLM-5.2 的首 token 延迟(越低越好)。采用 TP8、32K 预填充块;1M 下相对 SGLang 加速 1.86 倍。


图 8|DeepSeek-V3.2 的首 token 延迟(越低越好)。采用 TP8、16K 预填充块;1M 下相对 SGLang 加速 1.94 倍。


上下文越长,Indexer 在预填充中的耗时占比越高,优化它带来的端到端收益也越大。主注意力和前馈层仍需执行,因此完整请求的加速低于 Indexer 本身。

06

为更长上下文省下计算

DoSA 在保持原始 Top-K 选择结果的前提下,减少了 Indexer 的完整打分。上下界筛选省下计算,GPU 上的数据复用与紧凑写回减少执行开销,两者共同缩短了长上下文请求的首字等待。


这项优化适用于长文档阅读、代码分析等需要输入大量上下文的场景。落到具体服务时,还需结合请求长度、并发方式和部署配置,评估延迟改善能转化为多少吞吐与成本收益。


接下来,我们计划探索头以外其他维度的界,进一步减少计算与访存;也会尝试在训练阶段引入 DoSA,让模型形成更适合安全筛选的分数分布。

///

作者简介

  • 深水|小红书 Dots 部门,北京大学博士,2026 年 7 月加入小红书。在NeurIPS、ICML等学术会议以第一作者身份发表PiSSA、TransMLA等多篇论文。在DoSA之前曾提出HISA、MISA、LISA、LatentIndex等有损加速方法。

  • 星炼|小红书 Dots 部门,中科院软件所硕士,2026 年 6 月加入小红书。在MLSys、EMNLP等学术会议以第一作者身份发表过多篇论文,主要研究方向为大语言模型和大模型的算法-系统协同设计。

///

关于团队

我们是公司核心的 LLM 预训练团队,负责下一代基座模型的架构设计与训练系统研发。

团队聚焦:

  • Scaling Law

  • 优化器与训练稳定性

  • 长上下文与推理能力扩展

  • 高效训练与推理

  • MoE 与稀疏化架构

我们关注的不只是“把模型做大”。我们更关注:什么样的结构、优化与训练机制,能够让模型更高效地产生智能。

你将获得:

  • 直接参与下一代基座模型核心架构研发

  • 接触真实大规模训练环境,接触最前沿算力环境

  • 深度参与从结构设计到训练落地的完整闭环

  • 与顶尖研究员与工程师协作

  • 理解 frontier model 背后的真实技术问题

招聘官网:

https://job.xiaohongshu.com/social/position/12293


论文与代码

论文链接:

https://arxiv.org/pdf/2610.13990

huggingface:

https://huggingface.co/papers/2610.13990

代码:

https://github.com/llm-sys/dosa-sglang


前往微信阅读全文

内容来自公众号,可前往微信查看原文。

查看作者的更多文章 →订阅更新、RSS 与阅读指南 →