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

来源：市场资讯

（来源：机器之心）

机器之心发布

稀疏注意力减少了主注意力的计算量，但负责筛选 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 等学术会议发表多篇论文，主要研究方向为大语言模型及其算法与系统协同设计。