JSSA makes token-level sparse decoding gather-free, reaching up to full-KV throughput at 128K context
Synopsis
Long-context decoding is bottlenecked by repeatedly reading a growing KV cache, and sparse attention reads only part of it, yet token-level selection has not reliably delivered both accuracy and lower latency. JSSA scores candidate tokens in a shared low-dimensional query–key subspace and pairs it with a gather-free indexed-attention kernel that keeps a distinct selected token set for each KV head. Across GQA, hybrid sliding-window/global attention and MLA models it stays closest to full-KV accuracy, and on a B200 at 128K context its decoding throughput reaches up to that of full-KV attention.
Interpretation
Long-context decoding can now skip most KV-cache reads at token granularity instead of falling back to block- or page-level selection. Practical sparse methods have largely operated at coarse granularity, and token-level selection, though finer, had not reliably translated into either higher accuracy or lower latency. On LongBench at a 256-token KV budget, JSSA stays within 0.67 points of full-KV on LLaMA-3.1-8B-Instruct, Qwen3-14B and Ministral-3-8B-Instruct-2512, and no single competing baseline is consistently strongest across the three architectures.
The low-dimensional scoring space holds queries and keys together, so the compressed ranking stays close to the original attention scores. Objectives that preserve queries or keys independently do not guarantee preservation of the bilinear scores used for ranking, and an arbitrary low-rank projection alone does not guarantee more accurate ranking. At default projection ranks on LLaMA-3.1-8B-Instruct and Gemma4-26B-A4B-It, JSSA's oracle top-k recall exceeds Loki by 6.25 to 35.37 percentage points and FASA by 22.44 to 35.25 points.
Selected tokens are read in place through scattered indices rather than staged into a contiguous buffer. A common implementation gathers the selected key-value entries into a contiguous buffer before calling a dense kernel, and that step offsets the benefit of token-level sparsity. On LLaMA-3.1-8B-Instruct at 128K context and batch size 16, gathering is comparable to the reduced-attention computation itself and accounts for more than half of the residual latency; with gather-free execution the attention-stage speedup over full-KV roughly doubles at batch sizes 4 and 16.
Perspective
Evaluation covers standard GQA, hybrid sliding-window/global GQA and MLA architectures; long-context understanding and language modeling use LLaMA-3.1-8B-Instruct, Qwen3-14B and Ministral-3-8B-Instruct-2512, reasoning and throughput use Qwen3-8B, Gemma4-26B-A4B-It and GLM-4.7-Flash, and system experiments run on a single NVIDIA B200. For Gemma4 only the global-attention layers are sparsified while local layers keep native sliding-window attention; for MLA models selection operates over the shared latent cache and executes through FlashMLA's sparse kernels.
The marginal bound controls the exact score error but establishes neither its tightness nor global optimality of the spectral solution, and direct refinement of the exact objective gives no measurable selection benefit in the tested settings. Calibration length matters more than corpus: bases calibrated at 32K degrade to perplexity 16.05–18.31 when evaluated at 128K, versus 11.31–11.48 for 128K-calibrated bases. At roughly 700B scale the reasoning-accuracy differences between JSSA and native DSA are small, and the authors read them as comparable accuracy rather than a consistent advantage for either method.