Summary (Overview)

• JSSA (Joint Subspace Sparse Attention) is a training-free, token-level sparse attention method that addresses two key challenges in long-context LLM decoding: (1) low-dimensional scoring that may lose query–key structure needed for accurate token ranking, and (2) gather overhead from irregular per-KV-head token selections.
• Key theoretical contribution: JSSA learns a shared low-dimensional query–key subspace by minimizing bilinear score error. The non-convex objective is upper-bounded by a tractable marginal objective that admits a closed-form spectral solution via eigendecomposition.
• Key systems contribution: A gather-free indexed-attention kernel that supports distinct per-KV-head selections under the standard GQA KV-cache layout, avoiding the staging buffer required by gather-then-attend approaches.
• Results: JSSA achieves up to 6.43× decoding throughput of full-KV attention at 128K context on an NVIDIA B200 GPU, while remaining within 0.67 points of full-KV on LongBench and matching model-native DSA reasoning accuracy at 700B scale without training.
• Validation breadth: Evaluated across GQA, hybrid sliding-window/global attention, and Multi-head Latent Attention (MLA) models, demonstrating robustness across architectures and scales.

---

Introduction and Theoretical Foundation

Background and Motivation

Long-context autoregressive decoding is bottlenecked by repeatedly accessing the growing KV cache. With full attention, each decoding step scores the current query against the entire context, causing memory traffic to grow linearly with sequence length. Sparse attention reduces KV access by selecting only part of the context at each decoding step.

Token-level selection offers finer-grained sparse patterns than block- or page-level methods but presents two challenges:

1. Selection quality: Objectives that preserve queries or keys independently may not preserve the bilinear scores used for token ranking.
2. Execution efficiency: Token-level selection produces irregular per-KV-head index sets in GQA; gathering them into a contiguous buffer adds overhead, and forcing all KV heads to share one selected set degrades accuracy.

Theoretical Foundation

For each layer and KV head, JSSA uses an orthonormal basis $M \in \mathbb{R}^{r \times d}$ with $MM^\top = I_r$ applied to both queries and keys. Let $P = M^\top M$ be the corresponding rank-$r$ projector. In GQA, selection uses the mean query $\bar{q} = \frac{1}{G}\sum_{g=1}^{G} q_g$ of the query heads associated with the KV head. The projected score is:

$$\hat{s} = \bar{q}^\top P k,\tag{1}$$

and the resulting score error is:

$$s - \hat{s} = \bar{q}^\top (I - P) k.\tag{2}$$

The exact selection objective, defined over independently unit-trace-normalized second-moment matrices $\widetilde{G}_{\bar{Q}}$ and $\widetilde{G}_K$ of calibration group-mean queries and keys, is:

$$\mathcal{L}_{\mathrm{sel}}(P) = \mathrm{tr}\left((I - P)\widetilde{G}_{\bar{Q}}(I - P)\widetilde{G}_K\right).\tag{3}$$

This objective is non-convex with no general closed-form solution. However, for any orthogonal projector $P$:

$$\mathcal{L}_{\mathrm{sel}}(P) \leq \frac{1}{4}\mathcal{L}_{\mathrm{marg}}(P)^2, \qquad \mathcal{L}_{\mathrm{marg}}(P) = \mathrm{tr}\left((I - P)\left(\widetilde{G}_{\bar{Q}} + \widetilde{G}_K\right)\right).\tag{4}$$

The marginal objective admits the closed-form solution:

$$M^\star = \mathrm{TopEig}_r\left(\widetilde{G}_{\bar{Q}} + \widetilde{G}_K\right)^\top,\tag{5}$$

where $\mathrm{TopEig}_r(A)$ returns the top-$r$ eigenvectors of $A$ as columns.

---

Methodology

3.1 Overview

JSSA maintains a low-dimensional projected-key cache $\mathbf{K}_{\mathrm{proj},h}$ alongside the original KV cache for each layer and KV head. At each decoding step:

1. Project the mean query $\bar{q}_{t,h}$ of the $G$ query heads associated with KV head $h$
2. Score it against $\mathbf{K}_{\mathrm{proj},h}$ and select the top-$B$ token positions
3. Compute attention from original full-dimensional queries, keys, and values using the shared per-KV-head index set

3.2 Joint Query–Key Subspace

• The subspace is chosen to minimize expected squared error of the bilinear score used for token ranking, directly preserving the query–key interaction that determines token ordering
• Direct optimization from $M^\star$ produces negligible subspace changes and no measurable improvement in oracle top-$B$ recall, so $M^\star$ is used directly
• Calibration reduces to second-moment accumulation followed by one eigendecomposition per layer and KV head

3.3 Gather-Free Indexed Attention

• After selecting top-$B$ positions per KV head, attention is computed directly from scattered locations in the original full-dimensional KV cache
• A single token's key/value slice is contiguous over head dim, so each indexed read remains a burst-sized transaction
• A Triton kernel launches once per layer, resolves per-head index sets internally, and maps each query head to indices selected by its associated KV head
• This avoids both the staging buffer (gather-then-attend) and per-KV-head launches

3.4 Projected Scoring Cache

• Maintains $\mathbf{K}_{\mathrm{proj}} = \mathbf{K}M^\top$ as a scoring-only auxiliary cache
• For GQA: FP8 projected-key cache occupies at most 6.25% of original BF16 KV cache ($r = 32$ for LLaMA, Qwen, Ministral; $r = 64$ for Gemma4)
• For MLA: 128-dimensional scoring representation (64 latent + 64 RoPE dimensions), matching DSA indexer dimensionality
• FP8 throughout with negligible accuracy impact vs. BF16

---

Empirical Validation / Results

LongBench (Table 1) — KV budget of 256 tokens

JSSA achieves best or near-best sparse results on all three models, remaining within 0.67 points of full-KV.

PG-19 Perplexity (Table 2) — Fixed-density KV budgets

JSSA remains close to full-KV across all models and context lengths with a maximum gap of 1.11 perplexity points, while competing methods show inconsistent degradation (e.g., Loki degrades sharply on LLaMA, FASA spikes at 128K on LLaMA).

Reasoning Performance (Table 3) — KV budget of 1,024 tokens

JSSA limits generation-length inflation to 2–17% on Qwen3-8B and Gemma4, versus substantially longer generations for TriAttention, QUEST, and SeerAttention-R.

Decoding Throughput (Figure 4)

At 128K context, JSSA reaches up to:

• 1.51× full-KV throughput on Qwen3-8B
• 1.88× on Gemma4-26B-A4B-It
• 6.43× on GLM-4.7-Flash

Large-Scale Validation (Table 4) — ~700B parameters

JSSA matches native DSA on AIME 2025 for both DeepSeek and GLM-5.2, reaching 96.1–99.2% of native DSA's decoding throughput without continued training.

Ablation Studies

Projection rank (Figure 5a): JSSA performs best at every rank (4–32), with largest gains under tightest constraints (+5.40 and +24.48 LongBench points over FASA and Loki at $r = 4$).

Calibration robustness (Figure 5b): Insensitive to calibration corpus (LongAlign, NarrativeQA, LooGLE yield nearly identical perplexity), but calibration length matters — bases calibrated at 32K degrade sharply at 128K evaluation.

Gather-free execution (Figure 6): Roughly doubles attention-stage speedup at 128K on LLaMA-3.1-8B (25.35× → 65.26× at batch 16), and improves end-to-end throughput by 1.07×–2.22× over gather-then-attend on Gemma4.

---

Theoretical and Practical Implications

Theoretical Implications

• Score-preserving projection: The paper demonstrates that preserving query–key bilinear interactions is more important than preserving queries or keys independently for token ranking accuracy. The marginal upper bound provides a tractable surrogate that empirically matches direct optimization of the exact objective.
• Shared vs. asymmetric subspaces: JSSA constrains both queries and keys to a single shared scoring subspace, contrasting with asymmetric factorizations like KQ-SVD; the shared formulation proves effective for token selection.
• Closed-form spectral solution: The reduction of a non-convex objective to a tractable eigendecomposition makes calibration efficient (second-moment accumulation + one eigendecomposition per layer/KV head).

Practical Implications

• Architecture-agnostic: Works across standard GQA, hybrid sliding-window/global attention, and MLA models without architectural modifications or training.
• Immediate deployment: Training-free calibration means JSSA can be applied post hoc to existing models, including 700B-scale models, achieving accuracy comparable to model-native DSA (which requires continued pre-training).
• Systems design insight: Demonstrates that gather overhead can dominate residual sparse-attention cost, motivating gather-free kernel design as a first-class concern for token-level sparsity.
• Calibration guidance: Calibration should cover sequences near the maximum deployment context length for optimal transfer.

---

Conclusion

JSSA presents a training-free token-level sparse-attention method that jointly addresses selection quality and execution efficiency:

1. Selection quality: A shared query–key scoring subspace derived from a closed-form spectral objective that directly preserves bilinear score structure, with the marginal upper bound proving empirically sufficient.
2. Execution efficiency: A gather-free indexed-attention kernel supporting distinct per-KV-head selections under standard GQA KV-cache layout, eliminating staging-buffer overhead.

Across GQA, hybrid local/global attention, and MLA models, JSSA remains close to full-KV accuracy while delivering substantial throughput gains (up to 6.43× at 128K context). At approximately 700B scale, JSSA matches model-native DSA's reasoning accuracy while reaching 96–99% of its decoding throughput, without any model training.

Future directions implied by this work include: exploring tighter bounds or direct optimization of the selection objective for further gains, extending gather-free execution to additional cache layouts, and investigating calibration strategies that transfer across context lengths more robustly.