Joint Subspace Sparse Attention for Efficient Long-Context Decoding

Joint Subspace Sparse Attention for Efficient Long-Context Decoding. By Jiwon Song and two others, from Seoul National University. Published on the thirteenth of October. Its code repository is public but slight: nine source files.

When a large language model decodes over a long context, the cost is dominated by traffic: reading the KV cache. Sparse attention cuts that traffic by reading only part of it, but practical methods mostly work at coarse granularity. Choosing individual tokens would give finer control, and it brings two difficulties of its own. Scoring tokens in a low-dimensional space may lose the query and key structure needed to rank them accurately. And gathering the scattered tokens that were selected can eat up the latency benefit the selection was meant to deliver.

Joint Subspace Sparse Attention, or JSSA, takes on both difficulties, one for selection quality and one for execution speed. On the quality side, it learns a single low-dimensional subspace shared by queries and keys, chosen by minimizing the error in how well token scores are preserved.

That objective is non-convex. The authors upper-bound it with a tractable marginal objective that has a closed-form spectral solution. They also report that directly refining the exact objective gave no measurable improvement in the settings they tested, so the tractable version is not a compromise there.

On the speed side, they built a gather-free, token-level indexed-attention kernel. It allows each KV head to make its own distinct selection, while working within the standard grouped-query attention, or GQA, layout of the KV cache. The point is to get token-level choice without paying for the gather.

Testing spans three kinds of model: standard GQA, hybrid sliding-window and global attention, and Multi-head Latent Attention, or MLA. Across them, JSSA consistently stays close to full-KV accuracy. Competing methods showed substantial inflation in generation length and failures that depended on the architecture; JSSA avoids both.

On an NVIDIA B200 GPU at a context of one hundred twenty-eight thousand tokens, JSSA reaches up to six point four three times the decoding throughput of full-KV attention. Against model-native DSA, it achieves comparable reasoning accuracy without any model training, and reaches ninety-six to ninety-nine percent of DSA's decoding throughput, even on large models at seven hundred billion parameters.

Token-level sparsity had two separate reasons to disappoint: a ranking that might be wrong, and a fetch that might be slow. JSSA answers each with its own tool. A learned shared subspace with a spectral solution handles the first. A kernel that never gathers handles the second.