提速不降智,小红书DoSA将首token延迟减半

提速不降智,小红书DoSA将首token延迟减半
2026年10月20日 12:02 机器之心Pro

稀疏注意力减少了主注意力的计算量,但负责筛选 token 的索引器仍可能成为长上下文推理的瓶颈。小红书 Dots 团队自研的 DoSA 利用分数上下界,在保持原始 Top-K 选择结果的前提下加速索引器,并在测试中提前过滤了约 97% 的候选。在约 1M token 的上下文长度下,Dots3-Note-preview、GLM-5.2 和 DeepSeek-V3.2 的首 token 延迟较各自的 SGLang 基线分别降低约 37%、46% 和 48%。

  • 论文标题:DoSA: Dots Sparse Attention
  • 作者单位:Dots Studio
  • 项目地址:https://github.com/llm-sys/dosa-sglang

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

大模型从发出请求到收到第一个 token 的时间称为首 token 延迟(TTFT),也就是用户感受到的「首字等待时间」。上下文越长,预填充通常越耗时。

DeepSeek-V3.2 采用的 DSA 减轻了主注意力的计算负担:先由索引器(Indexer)为候选 token 打分,选出分数最高的 K 个,再由主注意力对选中的内容进行计算。然而,索引器本身成为了主要瓶颈。在约 1M token 的上下文长度下,Dots3-Note-preview、GLM-5.2 和 DeepSeek-V3.2 的 Indexer 耗时分别占预填充时间的 81.6%、73.0% 和 90.8%。

图 1|左:Indexer 耗时占预填充 GPU 时间的比例;右:Indexer 的理论矩阵运算量占比。

02 同样是加速 Indexer,DoSA 有何不同?

表 1|不同的基模团队索引器加速路线。

层间复用、减少索引头和按块共享等方法,通过减少独立打分次数或改变选择粒度来节省计算,但不保证与原始逐层、逐查询、使用全部索引头打分的 Top-K 选择结果等价。

DoSA 的关键区别在于:保留原始评分函数,并用数学上下界证明哪些计算可以省略。少量索引头仅用于构造保守上界,不会替代最终评分。只有被证明不可能进入 Top-K 的候选才会被排除,其余候选仍使用全部索引头进行精确打分。

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

理解 DoSA,只需抓住两个量:Top-K 入选的阈值下界,以及每个候选的分数上界。如果某个候选的得分上界仍低于门槛下界,就能确定它不可能入选,从而将其安全排除。

图 2|DoSA 的筛选与精算流程:先计算筛选所需的点积,利用上界排除候选,再对保留的候选计算全部索引头的得分,最后按原始规则汇总分数并写回。灰色和紫色分别表示省去的计算与写回操作。

DoSA 首先选出一组种子候选,并使用当前查询的全部索引头重新打分。种子候选中第 K 大的分数可作为入选门槛的下界,因为已有至少 K 个候选达到这一分数。种子候选可依据历史入选频次选取,也可在冷启动时由少量锚点查询产生;无论来源如何,都需要针对当前查询重新打分。

接下来,DoSA 为每个候选构造一个不低于其真实得分的上界。Indexer 的总分由多个索引头的贡献加权汇总。DoSA 保留少量重要头进行精确计算,并分别处理其余正权重头和负权重头的贡献:对于正权重头,将其投影到保留头上,并估计残差贡献的上界;对于负权重头,则利用聚合后的扣分项下界,保守地估计其扣分幅度。两部分与保留头的精确贡献相结合,得到总分的保守上界。

只有当得分上界严格低于门槛下界时,候选才会被排除。剩余候选继续使用全部索引头进行精确打分,最终按原始分数执行 Top-K 选择。界限即使估计得较为宽松,也只会增加精算量,不会误删可能入选的候选。

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

减少点积计算后,还需要让 GPU 高效完成剩余工作。不同查询保留的候选各不相同,若逐个处理,容易产生零散计算、重复的 key 加载和大量中间结果写回。为此,DoSA 将流程组织为规则化筛选、候选精算和紧凑写回三个阶段。

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

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

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

写回时,直接过滤低于门槛下界的分数,只保留可能进入最终 Top-K 的记录。每条记录包含紧凑编码后的分数和精确的 key ID。先用编码定位 Top-K 边界,再比较边界候选的原始分数,从而保持选择结果的精确性。

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

05 百万 token 上下文中,加速效果如何?

GLM-5.2 的完整 Indexer 微基准覆盖 5 种上下文长度、4 种查询块大小和 3 个输入领域,共 60 种配置。DoSA 在全部配置中均快于 SGLang 和 LiteTopK:相对两者的加速比分别为 1.12~2.74 和 1.17~1.69。当上下文长度达到 1M token 时,相对 SGLang 的加速比为 2.47~2.74。

图 4|预填充阶段完整 Indexer 相对 SGLang 的加速比。三行分别对应英文、中文和代码输入;每行比较 4K、8K、16K 和 32K 四种查询块大小,覆盖 128K~1M token 的上下文长度。

逐层分析显示,英文、中文和代码三个领域的平均候选过滤率均约为 97%。每个领域测试的 21 个 Indexer 输出层中,都有 20 层的过滤率超过 95%。因此,绝大多数查询与 key 的组合可以跳过全头精算。

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

端到端 TTFT 更直接地反映了用户的等待时间。在约 1M token 的上下文长度下,Dots3-Note-preview、GLM-5.2 和 DeepSeek-V3.2 的首 token 延迟较各自的 SGLang 基线分别降低约 37%、46% 和 48%。

图 6|Dots3-Note-preview 的首 token 延迟(越低越好)。在约 1M token 的上下文长度下,相对 SGLang 的加速比为 1.58。

图 7|GLM-5.2 的首 token 延迟(越低越好)。在 1M token 的上下文长度下,相对 SGLang 的加速比为 1.86。

图 8|DeepSeek-V3.2 的首 token 延迟(越低越好)。在 1M token 的上下文长度下,相对 SGLang 的加速比为 1.94。

在这些测试中,随着上下文变长,Indexer 在预填充中的耗时占比升高,优化它带来的端到端收益也随之增大。不过,主注意力和前馈层等计算仍需执行,因此端到端 TTFT 的加速比低于 Indexer 本身的加速比。

作者简介

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

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

财经自媒体联盟更多自媒体作者

新浪首页 语音播报 相关新闻 返回顶部