DoSA 以分數上界略過全頭計分，1M TTFT 近減半
arXiv 2610.13990 · 小紅書 Dots Studio · 2026-09-26（v1）
DeepSeek Sparse Attention（DSA）的 indexer 原本要對每個 key 算完所有 head 的分數才能做 Top-K，DoSA 改成先用便宜的上界排除不可能入選的 key，只對剩下的少數候選算完整分數，而且選出的 Top-K 與原本完全一致。論文在 GLM-5.2 與 DeepSeek-V3.2 上，約 1M tokens 的 TTFT 比 SGLang 快 1.86 倍與 1.94 倍。程式碼放在 llm-sys/flash-indexer。
原本的問題
DSA 先用多 head indexer 替每個 key 打分，公式是 s_i = Σ_h w_h · [q_h·k_i]_+，再取 Top-K 交給稀疏注意力。注意力本身只看少數 key，但挑 key 這一步仍是一次稠密掃描。
論文量到的結果是：prompt 約 1M tokens 時，indexer 佔 GLM-5.2 prefill GPU 時間的 73.0%，佔 DeepSeek-V3.2 的 90.8%（128K 時 GLM-5.2 僅 23.6%）。也就是上下文越長，「決定看哪裡」的成本越壓過「看」本身。
既有加速法各補一塊：GVR 利用解碼的時間相關性，DeepSelect 用有界緩衝串流分數，LiteTopK 把計分與候選過濾融合以省掉分數寫回。它們都仍要對每個 key 評估全部 head。MISA 用較少 head 先篩再重排，但重排救不回被誤刪的贏家，所以不保證與原 Top-K 相同。
核心改動：可證明的拒絕條件
DoSA 借用文件檢索 MaxScore／WAND 的想法：只要某 key 的分數上界 U_i 低於 Top-K 門檻的下界 τ，它就不可能入選。門檻下界來自一組種子 key：任取至少 K 個可見 key 算出完整分數，其第 K 大的值必不高於真正的門檻。
難處在於 indexer 的權重有正有負，上界不好做。論文的處理分三塊：
- 保留的 head：直接算實際分數，保留數量
r小於總 head 數H。 - 被省略的正權重 head：投影到某個保留 head 上重用已算好的內積，殘差再用 Cauchy–Schwarz 與 ReLU 次可加性界住（式 5、7）。
- 被省略的負權重 head：把一部分負 head 聚合成單一方向，其 ReLU 值是被扣掉量的下界，因此把它減掉仍是安全的上界。丟掉負 head 只會讓上界變鬆，不會出錯。
通過上界檢查的 key 才算全部 head 的原始分數，並且只寫回不低於 τ 的分數。證書只決定「少做哪些事」，最終分數仍是原式，因此論文證明在種子門檻有效、候選不漏的條件下，Top-K 與原 indexer 一致（定理 1）。
| 既有做法（如 LiteTopK） | DoSA | |
|---|---|---|
| 每個 key 的計分 | 全部 head | 先算保留 head 與上界，僅存活者算全部 head |
| 分數寫回 | 融合後減量 | 只寫回 ≥ τ 的分數 |
| Top-K 是否與原式相同 | 相同 | 相同（有證明） |
實際效果
實驗在 NVIDIA GPU 上，Top-K 為 2,048，baseline 是 SGLang v0.5.18 與 LiteTopK。GLM-5.2 與 DeepSeek-V3.2 的 indexer head 數分別是 32 與 64。
- Indexer 微基準：60 種設定（上下文 128K 到 1M、prefill chunk 4K 到 32K、英文／中文／程式碼）全部最快。對 SGLang 為 1.12 到 2.74 倍，對 LiteTopK 為 1.17 到 1.69 倍；上下文越長越賺，英文 16K chunk 從 128K 的 1.19 倍升到 1M 的 2.58 倍。
- 端到端 TTFT（TP8，約 1M tokens）：對 SGLang 為 1.86 倍（GLM-5.2）與 1.94 倍（DeepSeek-V3.2），對 LiteTopK 為 1.26 倍與 1.50 倍。
- 剪枝率：GLM-5.2 的 21 個 indexer 層中有 20 層超過 95%，各領域平均 97.2% 到 97.5%。
- 解碼：增益小很多，256K、512K、約 1M 請求長度分別只有 1.06、1.11、1.26 倍，因為每步 query 少，能共用的 key 載入有限。
消融顯示負 head 聚合最關鍵：拿掉它，候選存活率從 2.65% 升到 95.06%，等於幾乎沒有剪到東西。只開篩選省 37.4% 延遲，只開稀疏寫回省 23.4%，兩者合用省 47.3%（1.90 倍）。
影響範圍
直接相關的是自己維護 DSA 類稀疏注意力 serving kernel 的人，例如 SGLang 或 LiteTopK 上跑長上下文 prefill 的團隊。這是無損改動，不需要重訓或改模型參數，折疊後的權重只存在於篩選計畫內。要驗證的是自己模型的 indexer：head 數、權重正負分布與選擇邊界，決定了上界能拒絕多少 key。
短上下文不划算：論文自己指出 DoSA 仍要掃過每個 key，上下文短或上界鬆時，篩選成本可能蓋過省下的計算。另外作者說明，數值保證要求對原 scorer 保守評估，且並列的 Top-K 集合未必產生相同的生成文字；結果只證明計算效益，不證明任務準確度提升。論文未公布 NVIDIA GPU 的具體型號，kernel 設計依賴 Tensor Memory（TMEM）累加器，移植到其他硬體的成本論文未說明。