JSSA:面向高效长上下文解码的联合子空间稀疏注意力
原标题:Joint Subspace Sparse Attention for Efficient Long-Context Decoding
AI 摘要
论文提出 Joint Subspace Sparse Attention(JSSA),一种免训练的长上下文解码稀疏注意力方法。它通过学习共享的低维 query-key 子空间来最小化双线性分数误差,并设计无 gather 的 token 级索引注意力内核,支持 GQA 下每个 KV head 独立选择。在 GQA、混合滑窗/全局注意力和 MLA 模型上,JSSA 保持接近全 KV 的精度,在 B200 GPU 上 128K 上下文达到全 KV 解码吞吐,并在 700B 规模模型上达到模型原生 DSA 的推理精度和 96–99% 的解码吞吐。
正文节选
Joint Subspace Sparse Attention for Efficient Long-Context Decoding Abstract Sparse attention reduces the KV-cache traffic that dominates long-context LLM decoding, but practical methods largely operate at coarse granularity. Token-level selection offers finer control, yet presents two challenges: low-dimensional scoring may lose query–key structure needed for accurate ranking, while gathering scattered selected tokens can offset the resulting latency benefit. We introduce Joint Subspace Sparse