PaperlayerView exact revision on arXiv

Source-linked AI summary

Joint Subspace Sparse Attention for Efficient Long-Context Decoding

Jiwon Song, Yoongon Kim, Jae-Joon Kim

arXiv:2610.19026v1cs.CL

TL;DR

Long-context decoding repeatedly accesses a growing KV cache, while token-level sparse attention faces challenges in preserving query–key ranking and avoiding gather overhead. JSSA combines a shared query–key scoring subspace with gather-free indexed attention; across evaluated architectures it remains close to full-KV accuracy, with up to 6.43× full-KV decoding throughput at 128K on a B200 GPU.

  • Problem

    Long-context decoding repeatedly reads the growing KV cache, while token-level selection can lose query–key ranking structure in low-dimensional scoring and incur overhead gathering irregular per-KV-head selections.

    Source: §1 · p. 1

  • Method

    JSSA uses a shared low-dimensional query–key subspace for token scoring and a gather-free indexed-attention kernel that supports distinct per-KV-head selections.

    Source: §1 · pp. 1–2

  • Results

    Across GQA, hybrid local/global attention, and MLA models, JSSA remains close to full-KV accuracy and reaches up to 6.43× full-KV decoding throughput at 128K context on a B200 GPU.

    Source: §1 · p. 2

  • Takeaways & Limitations

    At approximately 700B scale, JSSA matches model-native DSA reasoning accuracy and reaches 96–99% of its decoding throughput without model training.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.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.

1 INTRODUCTION

JSSA targets two obstacles to token-level sparse decoding: preserving query–key ranking with low-dimensional scores and avoiding the overhead of gathering scattered selections. It combines a shared scoring subspace with gather-free indexed attention, and reports accuracy and throughput close to its stated full-KV and model-native comparisons.

1 INTRODUCTION

  • Token-level selection offers finer-grained sparse patterns, but low-dimensional scores can lose query–key ranking structure and irregular per-KV-head selections create gathering overhead under GQA.Forcing KV heads to share one selected set can also degrade accuracy.Source: §1 · p. 1
  • JSSA learns a shared low-dimensional query–key subspace using a closed-form spectral surrogate, while direct refinement of the exact objective gives no measurable selection benefit in tested settings.Its gather-free indexed-attention kernel preserves distinct per-KV-head token sets and accesses selected entries in the standard GQA cache layout.Source: §1 · pp. 1–2
  • At each decoding step, the mean of a KV head’s associated query heads scores projected keys to select top-B positions, which those query heads share for attention over original full-dimensional keys and values.Figure 1 depicts this scoring and execution path for one KV head.Source: Figure 1 · p. 2
  • Up to 6.43× decoding throughput over full-KV attention at 128K context on a B200 GPU, with accuracy close to full-KV across GQA, hybrid local/global attention, and MLA models.At approximately 700B scale, JSSA also matches model-native DSA reasoning accuracy without model training.Source: §1 · p. 2

2 MOTIVATION

Token-level selection requires scoring every candidate position, so its scoring dimension must be reduced without losing query–key interactions. Even when selection is sparse, gathering scattered KV entries can consume much of the remaining latency.

2.1 LOW-DIMENSIONAL TOKEN SCORING REQUIRES THE RIGHT OBJECTIVE

  • Token-level methods must score every candidate position and reduce the channel dimension; they can select original channels, as FASA does, or project onto a low-rank subspace, as Loki does.Retaining r original channels is a restricted orthonormal projection, while arbitrary projection alone does not guarantee more accurate ranking.Source: §2.1 · p. 2
  • At every evaluated rank, FASA’s discrete channel selection outperforms Loki’s key-reconstruction projection, motivating a subspace objective that preserves query–key interactions rather than key representations alone.Figure 2(a) compares LongBench accuracy against channel budget r for FASA and Loki.Source: §2.1 · p. 2, Figure 2 · p. 3

2.2 GATHER OVERHEAD LIMITS TOKEN-LEVEL SPEEDUP

  • At 128K context, gathering selected KV entries remains comparable to reduced attention and exceeds half of residual latency at batch size 16, despite sparsity lowering total latency.Figure 2(b) compares gather-then-attend latency, normalized to dense attention, with the reduced sparse-attention cost.Source: §2.2 · p. 3, Figure 2 · p. 3
  • The two requirements are a low-dimensional scoring subspace that preserves token rankings and an attention kernel that directly accesses selected KV entries..Source: §2.2 · p. 3

3 JOINT SUBSPACE SPARSE ATTENTION

JSSA selects tokens using projected keys and group-mean queries, then computes attention using original full-dimensional KV entries. Its shared query–key scoring subspace and gather-free execution are designed for token-level selection across GQA and MLA.

3.1 OVERVIEW

  • JSSA maintains a projected-key cache per layer and KV head, scores it with the projected mean query of associated query heads, and uses the selected top-B positions for attention over original full-dimensional queries, keys, and values.For GQA, associated query heads share each KV head’s index set; MLA uses FlashMLA’s sparse attention kernel.Source: §3.1 · p. 3

3.2 JOINT QUERY–KEY SUBSPACE

  • JSSA applies one orthonormal basis to queries and keys, choosing the shared subspace to minimize expected bilinear score error; a tractable marginal upper bound yields a closed-form spectral solution, while direct refinement adds no measurable recall benefit in tested settings.For GQA, the objective uses the group-mean query; its calibration query and key second moments are independently unit-trace-normalized.Source: §3.2 · pp. 3–4

3.3 GATHER-FREE INDEXED ATTENTION

  • The Triton indexed-attention kernel resolves distinct per-KV-head selections in one layer launch and reads selected entries in place, avoiding both a gather buffer and separate per-head launches.In MLA, each token has one shared latent vector per layer, so selection yields a single index set.Source: §3.3 · p. 4

3.4 PROJECTED SCORING CACHE

  • The projected-key cache is scoring-only, with attention still using original full-dimensional KV; the reported RULER figure shows JSSA close to full-KV across models under a fixed 1.5625% KV budget.Table 1 covers average LongBench performance at a 256-token KV budget, with 32-dimensional scoring representations for FASA, Loki, and JSSA.Source: §3.4 · p. 4, Figure 3 · p. 5, Table 1 · p. 5

4 EXPERIMENTS

Experiments compare JSSA with full-KV and sparse baselines across GQA, hybrid local/global GQA, and MLA models; the setup also varies tasks, context lengths, and KV budgets.

4.2.1 LONGBENCH

  • At a 256-token KV budget, JSSA is the best or near-best sparse method on all three LongBench models and stays within 0.67 points of full-KV.No competing baseline is consistently strongest across architectures.Source: §4.2.1 · p. 5

4.2.2 RULER

  • Across 32K–128K contexts, JSSA stays closest to full-KV on RULER and has the highest sparse accuracy at 128K.On Ministral-3-8B-Instruct-2512, it matches full-KV at 128K and is within 2 points at 32K and 64K.Source: §4.2.2 · p. 5

4.3 LANGUAGE MODELING

  • Across three models and 32K–128K contexts, JSSA’s PG-19 perplexity stays within 1.11 points of full-KV.Competing methods are less consistent, including sharp degradation for Loki on LLaMA and a FASA spike at 128K.Source: §4.3 · p. 6, Table 2 · p. 6

4.4 REASONING

  • At a 1,024-token KV budget, JSSA offers the strongest overall reasoning accuracy–length trade-off, staying within 1.3 points of full-KV on most benchmarks.Generation-length inflation is 2–17% on Qwen3-8B and Gemma4-26B-A4B-It; on GLM-4.7-Flash, JSSA matches or exceeds full-KV accuracy on all three benchmarks.Source: §4.4 · p. 6, Table 3 · p. 6

4.5 DECODING THROUGHPUT

  • At 128K, JSSA reaches 1.51×, 1.88×, and 6.43× the full-KV throughput on Qwen3-8B, Gemma4-26B-A4B-It, and GLM-4.7-Flash, respectively.The reported throughput is end-to-end decoding on one NVIDIA B200 GPU with a 2,048-token KV budget.Source: Figure 4 · p. 7, §4.5 · p. 7

4.6 VALIDATION AT LARGE SCALE

  • For pass@1 on AIME 2025 and GPQA-Diamond, JSSA is comparable to model-native DSA; on DeepSeek, its decoding throughput reaches 96.1–99.2% of native DSA across measured context lengths and batch sizes.At 128K context and batch size 8, JSSA reaches up to 3.49× full-KV throughput. Table 4 compares JSSA and FKV on DeepSeek-V3.1-Terminus with native DSA on DeepSeek-V3.2-Exp; no corresponding FKV checkpoint is available for GLM-5.2.Source: Table 4 · p. 7, §4.6 · p. 8

4.7 ABLATION STUDIES

  • Figure 5(a) shows JSSA leading FASA and Loki at every tested rank, while 32K-calibrated bases have worse 128K perplexity than 128K-calibrated bases.At rank 4, JSSA gains 5.40 and 24.48 LongBench points over FASA and Loki; 128K perplexity is 16.05–18.31 for 32K-calibrated bases versus 11.31–11.48 for 128K-calibrated bases. Calibration corpus has little effect.Source: §4.7 · p. 8, Figure 5 · p. 8
  • With the same selected tokens, gather-free execution roughly doubles attention-stage speedup at 128K on LLaMA and raises Gemma4 end-to-end throughput by 1.07×–2.22×.The Gemma4 gain grows with batch size, showing that gathering selected KV entries remains a substantial overhead after sparsification.Source: §4.7 · p. 8

5 RELATED WORK

Related work spans coarse-grained sparse attention, dynamic token selection, and low-rank attention representations. JSSA is distinguished by its token-level selection and gather-free execution for per-KV-head selections in standard GQA.

Coarse-grained sparse attention

  • Page- and block-level methods estimate importance for groups that map naturally to block-sparse kernels, while prefill methods identify input-dependent sparse patterns and some model-native designs train selection with hardware-aligned execution.Because selected groups may contain irrelevant tokens, JSSA instead selects individual tokens.Source: §5 · pp. 8–9

Token-level sparse attention and execution

  • Dynamic token-level methods differ in their scoring signals and cache policies; JSSA retains the complete cache, uses joint query–key score-preservation error, and keeps distinct selections across KV heads.Prior execution paths primarily target regular block selections or MLA’s shared latent cache; JSSA provides gather-free indexed attention for irregular per-KV-head selections in standard GQA.Source: §5 · p. 9

Projected and low-rank attention representations

  • Other low-rank methods reduce or transform representations used by attention, whereas JSSA keeps the original full-dimensional KV cache and uses projection only for an auxiliary scoring cache.KQ-SVD instead uses asymmetric query and key projections aimed at attention-score fidelity.Source: §5 · p. 9

6 CONCLUSION

JSSA is a training-free token-level sparse-attention method that combines a shared query–key scoring subspace with gather-free indexed attention. The evaluations cover GQA, hybrid local/global attention, and MLA, with model-specific handling of local layers and cache selection.

6 CONCLUSION

  • JSSA combines a closed-form spectral scoring objective with gather-free indexed attention and remains close to full-KV accuracy while delivering substantial throughput gains across GQA, hybrid local/global attention, and MLA.At approximately 700B scale, it matches native DSA’s reasoning accuracy and reaches 96–99% of its decoding throughput without model training.Source: §6 · p. 9

A.1 MODELS AND EVALUATION COVERAGE

  • LongBench, RULER, and PG-19 evaluations use LLaMA-3.1-8B-Instruct, Qwen3-14B, and Ministral-3-8B-Instruct-2512; reasoning and throughput evaluations use Qwen3-8B, Gemma4-26B-A4B-It, and GLM-4.7-Flash.JSSA is used only on Gemma4’s global-attention layers, and GLM-4.7-Flash selection operates over the shared MLA cache.Source: §A.1 · p. 13

A.2 EVALUATION PROTOCOLS

The appendix specifies evaluation and implementation choices, then develops the calibration objective and its spectral surrogate. It also reports calibration sensitivity and normalization results and clarifies what the marginal bound does—and does not—guarantee.

A.2 EVALUATION PROTOCOLS

  • Evaluations use benchmark-specific prompts and generation limits, greedy decoding for LongBench and RULER, and teacher-forced perplexity for PG-19.Reasoning evaluations use each model’s recommended sampling configuration, identical settings across methods, and eight seeds with mean pass@1 and generation length reported.Source: §A.2 · pp. 13–14, Table 6 · p. 14

A.3 JSSA CONFIGURATION

  • JSSA’s GQA selections are per KV head, while MLA selection yields one layer-level token set; Gemma4 sparsifies only global-attention layers.The projected scoring cache uses FP8; final sparse attention uses the original full-dimensional BF16 KV representations, and local sliding-window layers remain unchanged.Source: §A.3 · p. 14, Table 5 · p. 14

A.4 JSSA CALIBRATION

  • Calibration completes in 129–365 seconds across evaluated models without iterative optimization or model training.JSSA uses eight packed calibration sequences centered and cropped to 131,072 tokens, then computes a basis for each layer and KV head.Source: §A.4 · pp. 14–15, Table 7 · p. 15

A.5 CALIBRATION INPUT CONSTRUCTION

  • JSSA’s LongBench accuracy varies by at most 0.10 points across calibration choices; plain examples improve FASA accuracy by 2.75 points with LongAlign and 2.51 points with Qasper.JSSA uses packed 128K-token LongAlign sequences; FASA uses plain LongAlign examples, and Loki uses packed LongAlign calibration.Source: Table 8 · p. 15, §A.5 · p. 15

A.6 BASELINE IMPLEMENTATIONS

  • The compared baselines use different adaptations and are evaluated only where their original formulations transfer without substantial method-specific choices.For MLA, only FASA is implemented as a representative post-hoc baseline; SeerAttention-R is evaluated only on Qwen3-8B using its released trained checkpoint.Source: §A.6 · pp. 15–16

B.1 PRODUCT-OF-MARGINALS FORMULATION

  • The marginal calibration objective treats group-mean queries and keys as independent draws so it measures score error across query–key combinations compared during inference.This product-of-marginals construction does not assert empirical independence between query and key activations at corresponding sequence positions.Source: §B.1 · p. 17

B.2 TRACE NORMALIZATION

  • Trace normalization improves oracle recall by 1.29 and 1.02 percentage points at token budgets of 256 and 2,048, respectively.It prevents the larger key second moment from dominating the marginal surrogate while preserving its joint query–key character.Source: §B.2 · p. 17, Table 9 · p. 18

B.3 DERIVATION OF THE MARGINAL BOUND

  • The marginal objective has a closed-form rank-r solution, but its bound does not establish that it globally minimizes the exact objective or shares its minimizer.The derivation does not require the normalized query and key second-moment matrices to commute; their normalized commutator ratio averages 0.470 on LLaMA-3.1-8B.Source: §B.3 · p. 18

B.4 EXACT-OBJECTIVE REFINEMENT

Direct refinement of the spectral basis reduces the exact objective slightly but does not measurably improve token-selection quality in the tested settings. Across two models and token budgets, JSSA has the highest oracle top-B recall among the compared scoring methods.

B.4 EXACT-OBJECTIVE REFINEMENT

  • Refinement reduces the exact objective by at most 1.60% at r = 8 and 0.83% at r = 32, while overlap stays at least 0.999476 at r = 32.The learning-rate sweep spans a 1000× range; the spectral basis is used directly in practice.Source: §B.4 · p. 19, Table 10 · p. 19, Table 11 · p. 20
  • Random initialization reduces the objective substantially, but its final objective remains 3.53× and 8.97× larger than the spectral initialization at r = 8 and r = 32.The random basis has much lower overlap with the spectral subspace, while the spectral solution lies in a flat, low-objective region.Source: §B.4 · p. 19
  • Direct refinement leaves oracle top-B recall and subspace overlap essentially unchanged across the tested settings.For Gemma4-26B-A4B-It, the evaluation covers the five global-attention layers where JSSA is applied.Source: §B.4 · pp. 19–20, Table 11 · p. 20

B.5 ORACLE SELECTION RECALL

  • JSSA achieves the highest oracle recall on both models and token budgets, exceeding Loki by 6.25–35.37 points and FASA by 22.44–35.25 points.These comparisons use the methods’ default projection ranks and the paper’s main-evaluation calibration configuration.Source: Table 12 · p. 20, §B.5 · p. 20

B.6 COMPARISON WITH ASYMMETRIC QUERY–KEY FACTORIZATION

The asymmetric KQ-SVD-style scoring factorization is competitive in some settings, but its accuracy varies more across architectures and contexts than JSSA’s. Its reasoning performance degrades sharply under aggressive sparsity and recovers at larger token budgets.

B.6 COMPARISON WITH ASYMMETRIC QUERY–KEY FACTORIZATION

  • On RULER, the asymmetric variant is competitive in several settings but less consistent across architectures and context lengths, with a particularly large gap on Qwen3-14B.It closely matches JSSA on LLaMA-3.1-8B-Instruct and Ministral-3-8B-Instruct-2512.Source: Table 13 · p. 21, §B.6 · p. 21
  • The asymmetric variant recovers substantially at larger reasoning budgets and becomes competitive in several 2,048-token settings.The authors suggest it may overfit calibration-specific query–key correlations, while noting that this robustness difference needs further characterization.Source: §B.6 · pp. 21–22
  • At a 512-token budget, the asymmetric factorization scores 0.83% and 26.67% on AIME 2025, versus 64.58% and 88.75% for JSSA.On GPQA-Diamond, it trails JSSA by 14.14 and 24.87 percentage points at the same budget.Source: Table 14 · p. 21, §B.6 · p. 21

C LARGE-SCALE VALIDATION

At approximately 700B scale, post-hoc JSSA reaches 96.1–99.2% of native DSA throughput on DeepSeek while showing comparable reasoning accuracy. A separate GLM-5.2 comparison also reports closely matched reasoning scores under that model’s native sparse-attention schedule.

C LARGE-SCALE VALIDATION

  • JSSA and native DSA show comparable reasoning accuracy on AIME 2025 and GPQA-Diamond, rather than a consistent advantage for either method.The passage notes small numerical differences and potential evaluation variance.Source: C LARGE-SCALE VALIDATION · p. 23
  • Across the DeepSeek throughput sweep, JSSA reaches 96.1–99.2% of native DSA throughput.At batch size 4, throughput relative to FKV rises from 0.89× at 8K context to 2.41× at 128K; at 128K and batch size 8 it reaches 3.49×.Source: C LARGE-SCALE VALIDATION · p. 23
  • On GLM-5.2, JSSA and native DSA differ by 0.83 percentage points on AIME 2025 and have identical GPQA-Diamond scores.JSSA follows GLM-5.2’s schedule, selecting only at indexer layers and reusing indices in adjacent layers.Source: C LARGE-SCALE VALIDATION · p. 23

D DECODING THROUGHPUT DETAILS

Across workload sweeps, JSSA’s throughput gains generally grow with context length and batch size, though scoring and selection can offset KV-access savings at smaller workloads. The appendices report component latency, projected-cache precision, per-head selection, and task-level accuracy results.

D DECODING THROUGHPUT DETAILS

  • JSSA reaches up to 1.69×, 2.82×, and 6.99× full-KV throughput on Qwen3-8B, Gemma4-26B-A4B-It, and GLM-4.7-Flash, respectively.Its benefit generally increases with context length and batch size; the Gemma4 result sparsifies only global-attention layers, while GLM benefits from sparse FlashMLA over a shared latent cache.Source: Figure 8 · p. 24, D DECODING THROUGHPUT DETAILS · pp. 24–25

E.1 COMPONENT-LEVEL LATENCY BREAKDOWN

  • At 128K context, projecting the newly generated key accounts for at most 3.2% of JSSA pipeline latency and less than 1% at larger batches.Top-B selection dominates at small batches, while the full-context score scan grows more important as batch size increases; JSSA is faster than QUEST at every measured Qwen3-8B batch size.Source: Figure 9 · p. 26, §E.1 · p. 26

E.2 PROJECTED SCORING CACHE AND QUANTIZATION

  • FP8 and BF16 projected-key caches differ by only 0.18, 0.01, and 0.05 LongBench accuracy points on LLaMA-3.1-8B-Instruct, Qwen3-14B, and Ministral-3-8B-Instruct-2512.The FP8 scoring cache halves its memory footprint relative to BF16; final attention remains in BF16.Source: §E.2 · p. 26, Figure 10 · p. 27

E.3 SELECTION GRANULARITY IN GQA

  • At 128K context, per-KV-head selection outperforms head-shared soft voting by 5.29 points on LLaMA-3.1-8B-Instruct, 0.99 on Ministral-3-8B-Instruct-2512, and 1.32 on Qwen3-14B.At 32K, the two approaches perform nearly identically; the indexed-attention kernel supports distinct per-KV-head sets without gathering selected entries.Source: §E.3 · pp. 26–27, Table 15 · p. 27

E.4.1 LONGBENCH

  • At a 256-token KV budget, JSSA achieves the best sparse LongBench average on LLaMA-3.1-8B and Qwen3-14B and remains close to the strongest sparse baselines on Ministral-3-8B.Its task-level results are comparatively stable across the reported categories.Source: Table 16 · p. 28, §E.4.1 · p. 28

E.4.2 RULER

  • On RULER, JSSA remains close to full-KV across models and context lengths, preserving needle-in-a-haystack variants and avoiding stronger long-context degradation seen for Loki and FASA on LLaMA-3.1-8B.On Ministral-3-8B, it is within 2 points of full-KV average accuracy at 32K and 64K and matches it at 128K.Source: §E.4.2 · p. 28

E.5 COMPLETE REASONING RESULTS

JSSA remains comparatively stable across KV budgets, staying close to full-KV accuracy at 512 tokens on most model–benchmark pairs and providing the most consistent trade-off between final-answer accuracy and reasoning length.

E.5 COMPLETE REASONING RESULTS

  • Tables 18–20 report complete reasoning accuracy for Qwen3-8B, Gemma4-26B-A4B-It, and GLM-4.7-Flash across KV budgets.Full-KV results are reported once because they do not depend on the sparse token budget.Source: §E.5 · p. 29, Table 18 · p. 31, Table 19 · p. 31
  • At 512 tokens, JSSA stays close to full-KV accuracy on most model–benchmark pairs and provides the most consistent trade-off between final-answer accuracy and reasoning length.Several baselines show severe accuracy loss and generation-length inflation; increasing the budget generally narrows the gaps.Source: §E.5 · p. 29, Table 18 · p. 31, Table 19 · p. 31
  • Table 17 reports RULER accuracy at a fixed 1.5625% KV budget, with its average covering 13 tasks.Source: Table 17 · p. 30

Related papers

Loading 2610.19026v1…