arXiv cs.CL· Jiwon Song, Yoongon Kim, Jae-Joon Kim·· 3 天前AI 评分34
JSSA：面向高效长上下文解码的联合子空间稀疏注意力
Joint Subspace Sparse Attention for Efficient Long-Context Decoding
AI 导读
研究者提出联合子空间稀疏注意力（JSSA），通过学习共享低维 query-key 子空间提升 token 级选择质量，并设计免 gather 的 token 级索引注意力 kernel，在 GQA、混合滑窗/全局注意力和 MLA 模型上保持接近全 KV 的精度。
正文 · AI 翻译
摘要：稀疏注意力减少了在长上下文 LLM 解码中占主导地位的 KV-cache 流量，但实用方法大多在粗粒度上运行。Token 级选择提供了更细粒度的控制，但带来了两个挑战：低维打分可能丢失准确排序所需的 query--key 结构，而收集分散的已选 token 可能抵消由此带来的延迟收益。我们提出联合子空间稀疏注意力（JSSA）来应对这两个挑战。在选择质量方面，JSSA 通过最小化 token 分数保持误差来学习一个共享的低维 query--key 子空间。我们用可处理的边缘目标对上式所得的非凸目标进行上界约束，该边缘目标具有闭式谱解；在测试设置中，直接对精确目标进行细化未带来可测量的改进。在执行效率方面，我们开发了一种无收集的 token 级索引注意力内核，在标准分组查询注意力（GQA）KV-cache 布局下支持每个 KV 头不同的选择。在标准 GQA、混合滑动窗口/全局注意力和多头潜在注意力（MLA）模型中，JSSA 始终接近全 KV 精度，同时避免了竞争方法中观察到的显著生成长度膨胀和依赖架构的失败。在 NVIDIA B200 GPU 上，JSSA 在 128K 上下文下实现了高达全 KV 注意力 $6.43\times$ 的解码吞吐量。JSSA 在无需模型训练的情况下达到了与模型原生 DSA 相当的推理精度，即使在 700B 规模的大模型上也能达到其解码吞吐量的 96--99\%。
| 学科： | 计算与语言（cs.CL） |
| 引用为： | arXiv:2610.19026 [cs.CL] |
| （或 arXiv:2610.19026v1 [cs.CL] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2610.19026 arXiv 颁发的 DOI 通过 DataCite（待注册） |
提交历史
来自：Jiwon Song [查看邮箱]
[v1]
2026年10月13日星期二 18:18:47 UTC（259 KB）
来源：arXiv cs.CL · arxiv.org