Joint Subspace Sparse Attention for Efficient Long-Context Decoding
JSSA sparse attention stays near full-KV accuracy and raises 128K decoding throughput up to 6.43×.
Joint Subspace Sparse Attention learns a shared low-dimensional query-key subspace to rank tokens while limiting KV-cache traffic in long-context decoding. A tractable spectral objective selects the subspace, and a gather-free indexed-attention kernel supports distinct per-KV-head selections in standard grouped-query layouts. Across GQA, hybrid sliding-window/global attention, and MLA models, accuracy stays close to full-KV attention without the length inflation seen in competing methods. On an NVIDIA B200 it reaches up to 6.43× full-KV decoding throughput at 128K context and 96–99% of model-native DSA throughput at 700B scale without training.
- Shared subspace scoring preserves token ranking with a closed-form spectral solution.
- A gather-free kernel supports per-KV-head selections under GQA layouts.
- Accuracy stays near full KV across GQA, hybrid attention, and MLA.
- Up to 6.43× decoding throughput at 128K on an NVIDIA B200.
- Matches native DSA accuracy and 96–99% of its throughput at 700B scale.
Full article204 words · extracted from huggingface.co · click to collapse
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 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. 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. On an NVIDIA B200 GPU, JSSA achieves up to 6.43times 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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://huggingface.co/papers/2610.19026