Joint Subspace Sparse Attention for Efficient Long-Context Decoding

Computer Science > Computation and Language arXiv:2610.19026 (cs) [Submitted on 13 Oct 2026] Title: Joint Subspace Sparse Attention for Efficient Long-Context Decoding Authors: Jiwon Song, Yoongon Kim, Jae-Joon Kim.

计算机科学 > 计算与语言 arXiv:2610.19026 (cs) [提交于 2026 年 10 月 13 日] 标题：用于高效长上下文解码的联合子空间稀疏注意力 作者：Jiwon Song, Yoongon Kim, Jae-Joon Kim。

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.

摘要：稀疏注意力减少了在长上下文大语言模型（LLM）解码中占主导地位的 KV 缓存流量，但实际方法大多在粗粒度上运行。词元级选择提供了更精细的控制，但也带来了两个挑战：低维评分可能会丢失准确排序所需的查询-键结构，而收集分散的选定词元可能会抵消由此产生的延迟优势。

We introduce Joint Subspace Sparse Attention (JSSA) to address both challenges. For selection quality, JSSA learns a shared low-dimensional query—key subspace by minimizing token-score preservation error. We upper-bound the resulting non-convex objective with a tractable marginal objective that admits a closed-form spectral solution; direct refinement of the exact objective provides no measurable improvement in the tested settings.

我们引入了联合子空间稀疏注意力（JSSA）来解决这两个挑战。在选择质量方面，JSSA 通过最小化词元分数保持误差来学习共享的低维查询-键子空间。我们将由此产生的非凸目标函数通过一个可处理的边缘目标函数进行上界约束，该函数允许闭式谱解；在测试设置中，直接优化精确目标函数并未提供可测量的改进。

For execution efficiency, we develop a gather-free token-level indexed-attention kernel that supports distinct per-KV-head selections under the standard grouped-query attention (GQA) KV-cache layout. Across standard GQA, hybrid sliding-window/global attention, and Multi-head Latent Attention (MLA) models, JSSA consistently remains close to full-KV accuracy while avoiding the substantial generation-length inflation and architecture-dependent failures observed with competing methods.

在执行效率方面，我们开发了一种无需收集（gather-free）的词元级索引注意力内核，它支持在标准分组查询注意力（GQA）KV 缓存布局下进行不同的每 KV 头选择。在标准 GQA、混合滑动窗口/全局注意力和多头潜在注意力（MLA）模型中，JSSA 始终保持接近全 KV 的准确性，同时避免了竞争方法中观察到的显著生成长度膨胀和架构依赖性故障。

On an NVIDIA B200 GPU, JSSA achieves up to $6.43\times$ the decoding throughput of full-KV attention at 128K context. JSSA achieves reasoning accuracy comparable to model-native DSA without model training and reaches 96—99% of its decoding throughput even at 700B scale large models.

在 NVIDIA B200 GPU 上，JSSA 在 128K 上下文下实现了高达全 KV 注意力 6.43 倍的解码吞吐量。JSSA 在无需模型训练的情况下实现了与模型原生 DSA 相当的推理准确性，即使在 700B 规模的大模型中也能达到其解码吞吐量的 96%—99%。