Joint Subspace Sparse Attention for Efficient Long-Context Decoding

Jiwon Song; Yoongon Kim; Jae-Joon Kim

cs.CLOct 13, 2026

Joint Subspace Sparse Attention for Efficient Long-Context Decoding

Authors: Jiwon Song, Yoongon Kim, Jae-Joon Kim

Organizations: Seoul National University

Abstract

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.

Figures & tables

Figure 1: JSSA at one decoding step for a single KV head. The mean of its associated query heads scores a projected-key cache in a shared low-dimensional subspace to select the top- $B$ token positions. The query heads share this index set, and gather-free indexed attention accesses the corresponding original full-dimensional keys and values.

Figure 2: Low-dimensional scoring and gather overhead on LLaMA-3.1-8B-Instruct. (a) LongBench accuracy versus channel budget

r

for FASA and Loki. (b) Gather-then-attend latency at 128K context, normalized to dense attention, with a breakdown of the residual sparse-attention cost.

Model	FKV	QUEST	SqueezedAttn.	SnapKV	FASA	Loki	JSSA
LLaMA-3.1-8B-Instruct	48.35	46.94	47.95	43.53	46.74	45.65	49.16
Qwen3-14B	50.49	48.44	48.58	45.27	47.93	49.54	49.82
Ministral-3-8B-Instruct-2512	47.90	46.48	47.91	44.82	46.04	48.07	47.79

Table 1: Average LongBench performance at a KV budget of 256 tokens. FASA, Loki, and JSSA use 32-dimensional scoring representations. The best and second-best sparse results for each model are shown in bold and underlined, respectively.

Figure 3: RULER accuracy across context lengths. Average accuracy under a fixed

1.5625\%

KV budget, corresponding to 512, 1,024, and 2,048 selected tokens at 32K, 64K, and 128K context, respectively. JSSA remains close to full-KV across models, while competing methods exhibit stronger model- and length-dependent degradation.

Model	Context	FKV	QUEST	FASA	Loki	JSSA
LLaMA-3.1-8B-Instruct	32K	11.07	35.10	23.12	50.15	11.27
	64K	10.14	21.68	14.76	41.45	10.14
	128K	10.38	25.50	56.59	56.29	11.49
Qwen3-14B	32K	11.24	16.38	16.89	14.89	11.49
	64K	10.93	14.51	19.34	13.49	11.01
	128K	10.41	13.29	15.05	14.61	10.53

Table 2: PG-19 perplexity across context lengths under the same fixed-density KV budgets used for RULER. Lower is better. The best and second-best sparse results are shown in bold and underlined, respectively.

Model	Method	AIME 2025	MATH-500	GPQA-Diamond
Qwen3-8B	FKV	69.58 (1.00 $\times$ )	94.53 (1.00 $\times$ )	57.95 (1.00 $\times$ )
	TriAttention	20.00 (1.75 $\times$ )	86.20 (1.70 $\times$ )	40.40 (4.13 $\times$ )
	QUEST	62.92 (1.29 $\times$ )	93.50 (1.21 $\times$ )	58.14 (1.03 $\times$ )
	SeerAttention-R	42.08 (1.54 $\times$ )	91.80 (1.35 $\times$ )	55.93 (1.43 $\times$ )
	FASA	60.42 (1.30 $\times$ )	93.58 (1.17 $\times$ )	57.45 (1.06 $\times$ )
	JSSA	68.33 (1.07 $\times$ )	94.23 (1.09 $\times$ )	57.39 (1.02 $\times$ )

Table 3: Reasoning performance at a KV budget of 1,024 tokens. Each entry reports pass@1 accuracy, with generation length relative to FKV shown in parentheses. Higher accuracy and lower length ratio are better. Bold indicates the best sparse accuracy for each model and benchmark.

Figure 4: End-to-end decoding throughput. Throughput on a single NVIDIA B200 GPU at 32K and 128K context with a KV budget of

B=2{,}048

tokens. Annotations report JSSA’s maximum speedup over FKV for each model and context setting.

Model	Method	AIME 2025	GPQA-Diamond
DeepSeek-V3.1-Terminus	FKV	86.67	81.57
DeepSeek-V3.1-Terminus	JSSA	87.50	80.30
DeepSeek-V3.2-Exp	Native DSA	87.50	78.54
GLM-5.2	JSSA	83.33	90.91
GLM-5.2	Native DSA	82.50	90.91

Table 4: Large-scale reasoning accuracy. JSSA is applied post hoc without model training. FKV and JSSA use DeepSeek-V3.1-Terminus, whereas native DSA uses DeepSeek-V3.2-Exp. No corresponding FKV checkpoint is available for GLM-5.2.

Figure 5: Scoring-subspace ablations. (a) LongBench accuracy versus projection rank on LLaMA-3.1-8B-Instruct at a KV budget of 256 tokens. (b) PG-19 perplexity on LLaMA-3.1-8B-Instruct under the same fixed-density KV budgets used in the main evaluation, for bases calibrated on different corpora and context lengths. Calibration transfers across datasets, but bases calibrated at 32K degrade when evaluated at 128K.

Figure 6: Effect of gather-free indexed attention ( $B=2{,}048$ tokens). (a) Attention-stage speedup over full-KV attention on LLaMA-3.1-8B-Instruct at 128K, comparing gather-then-attend with our gather-free kernel. (b) End-to-end decoding speedup of gather-free execution over gather-then-attend on Gemma4-26B-A4B-It across context lengths and batch sizes.

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Attention architecture	LongBench	RULER	PG-19	Reasoning	Throughput
LLaMA-3.1-8B-Instruct	GQA	✓	✓	✓
Qwen3-14B	GQA	✓	✓	✓
Ministral-3-8B-Instruct-2512	GQA	✓	✓	✓
Qwen3-8B	GQA				✓	✓
Gemma4-26B-A4B-It	Hybrid local/global GQA				✓	✓
GLM-4.7-Flash	MLA				✓	✓

Appendix

Table 5: Model and evaluation coverage. Checkmarks indicate the evaluations performed for each model. For Gemma4, JSSA is applied only to global-attention layers, while local sliding-window layers remain unchanged.

Model	Temperature	Top- $p$	Top- $k$	Max. new tokens
Qwen3-8B	0.6	0.95	20	40960
Gemma4-26B-A4B-It	1.0	0.95	64	40960
GLM-4.7-Flash	1.0	0.95	-	131072

Appendix

Table 6: Reasoning decoding configurations. All compared methods use identical parameters and evaluation seeds for a given model and benchmark. Each evaluation is repeated with eight independent seeds, and the reported results are averaged across runs.

Model	LLaMA-3.1-8B	Qwen3-14B	Ministral-3-8B	Gemma4-26B	GLM-4.7-Flash
Latency (s)	203	295	129	365	247

Appendix

Table 7: End-to-end JSSA calibration latency on an NVIDIA B200 GPU.

Method	Calibration Corpus	Packed 128K	Plain
JSSA	LongAlign	49.16	49.06
JSSA	Qasper	49.13	49.15
FASA	LongAlign	43.99	46.74
FASA	Qasper	44.70	47.21

Appendix

Table 8: Average LongBench accuracy under different calibration corpora and input constructions. Packed calibration uses 128K-token sequences, whereas plain calibration uses the original unmodified examples.

Basis	$B=256$	$B=2048$
Raw group-mean moments	70.05%	74.92%
Unit-trace-normalized group-mean moments	71.34%	75.94%

Appendix

Table 9: Effect of trace normalization on oracle top-

B

recall for LLaMA-3.1-8B-Instruct at

r=32

$r$	$\eta$	$L_{\mathrm{sel}}^{(0)}$	$L_{\mathrm{sel}}^{(1600)}$	Relative Change	Overlap
8	0.005	$1.752351{\times}10^{-3}$	$1.751730{\times}10^{-3}$	$-0.04\%$	0.999999
	0.05	$1.752351{\times}10^{-3}$	$1.747971{\times}10^{-3}$	$-0.25\%$	0.999958
	0.5	$1.752351{\times}10^{-3}$	$1.736314{\times}10^{-3}$	$-0.92\%$	0.998841
	5.0	$1.752351{\times}10^{-3}$	$1.724382{\times}10^{-3}$	$-1.60\%$	0.989120
32	0.005	$3.968264{\times}10^{-4}$	$3.967997{\times}10^{-4}$	$-0.01\%$	1.000000
	0.05	$3.968264{\times}10^{-4}$	$3.965928{\times}10^{-4}$	$-0.06\%$	0.999999

Appendix

Table 10: Exact-objective refinement from the spectral initialization on LLaMA-3.1-8B-Instruct. We sweep the learning rate over three orders of magnitude and report the exact objective before and after 1,600 optimization steps.

Model	$r$	$B$	Oracle Recall		Subspace Overlap
Model	$r$	$B$	Step 0	Step 1600	Subspace Overlap
LLaMA-3.1-8B-Instruct	8	256	48.76%	48.74%	0.999958
	8	2048	53.81%	53.81%	0.999958
	32	256	71.34%	71.36%	0.999999
	32	2048	75.94%	75.95%	0.999999
Gemma4-26B-A4B-It	64	256	79.85%	79.85%	1.000000

Appendix

Table 11: Selection stability under exact-objective refinement. Starting from

M^{\star}_{\mathrm{marg}}

, we optimize the exact objective for 1,600 steps with

\eta=0.05

. Oracle top-

B

recall and subspace overlap remain essentially unchanged. For Gemma4-26B-A4B-It, results are reported over the five global-attention layers to which JSSA is applied.

Model	$B$	JSSA	Loki	FASA
LLaMA-3.1-8B-Instruct ( $r=32$ )	256	71.34%	35.97%	48.90%
LLaMA-3.1-8B-Instruct ( $r=32$ )	2048	75.94%	51.84%	53.31%
Gemma4-26B-A4B-It ( $r=64$ )	256	79.85%	67.79%	44.60%
Gemma4-26B-A4B-It ( $r=64$ )	2048	83.33%	77.08%	49.90%

Appendix

Table 12: Oracle top-

B

recall on LLaMA-3.1-8B-Instruct and Gemma4-26B-A4B-It at their default projection ranks.

Model	Method	32K	64K	128K
LLaMA-3.1-8B-Instruct	Asymmetric (KQ-SVD)	84.25	81.77	73.23
LLaMA-3.1-8B-Instruct	JSSA	86.35	82.55	75.38
Qwen3-14B	Asymmetric (KQ-SVD)	71.46	82.42	75.78
Qwen3-14B	JSSA	88.86	88.47	82.21
Ministral-3-8B-Instruct-2512	Asymmetric (KQ-SVD)	89.71	86.77	86.28
Ministral-3-8B-Instruct-2512	JSSA	90.72	89.02	85.97

Appendix

Table 13: RULER accuracy of JSSA and an asymmetric KQ-SVD factorization adapted to token scoring under the same fixed-density KV budgets as the main evaluation.

		AIME 2025			GPQA-Diamond
Model	Method	512	1024	2048	512	1024	2048
Qwen3-8B	Asymmetric (KQ-SVD)	0.83	9.17	34.58	43.94	53.54	59.09
Qwen3-8B	JSSA	64.58	68.33	67.92	58.08	57.39	58.21
Gemma4-26B-A4B-It	Asymmetric (KQ-SVD)	26.67	77.50	85.83	55.81	75.25	81.82
Gemma4-26B-A4B-It	JSSA	88.75	86.67	88.33	80.68	81.94	81.12

Appendix

Table 14: Reasoning pass@1 accuracy of JSSA and an asymmetric KQ-SVD factorization adapted to token scoring across token budgets.

Figure 7: End-to-end decoding throughput on DeepSeek. Throughput across context lengths at batch size 4 (left) and across batch sizes at 128K context (right), measured on eight NVIDIA B200 GPUs with tensor parallelism degree eight.

Figure 8: Complete end-to-end decoding-throughput sweeps. Results across context lengths and batch sizes on a single NVIDIA B200 GPU. Annotations show JSSA’s maximum speedup over FKV in each panel.

Figure 9: Component-level sparse-attention latency at 128K context. Bars decompose JSSA and QUEST into scoring preparation, candidate scoring, top-

B

selection, and sparse attention. JSSA’s key-projection overhead is negligible, and its latency advantage over QUEST increases at larger batches.

Figure 10: Effect of projected scoring-cache precision. Average LongBench accuracy at a KV budget of 256 tokens using FP8 and BF16 projected-key caches. The original full-dimensional KV cache and final attention computation remain in BF16.

Model	Selection	32K	64K	128K
LLaMA-3.1-8B-Instruct	Per-KV-head	86.35	82.55	75.38
	Head-shared soft vote	86.03	80.82	70.09
	$\Delta$	-0.32	-1.73	-5.29
Ministral-3-8B-Instruct-2512	Per-KV-head	90.72	89.02	85.97
	Head-shared soft vote	90.67	87.24	84.98
	$\Delta$	-0.05	-1.78	-0.99

Appendix

Table 15: RULER accuracy of per-KV-head selection and head-shared soft voting in GQA. The budgets are 512, 1,024, and 2,048 tokens at 32K, 64K, and 128K context, respectively.

\Delta

denotes head-shared minus per-KV-head accuracy.

	Single-Document QA			Multi-Document QA			Summarization			Few-Shot Learning			Synthetic		Code		Avg.
Method	NarrativeQA	Qasper	MultiFieldQA-en	HotpotQA	2WikiMQA	Musique	GovReport	QMSum	MultiNews	TREC	TriviaQA	SAMSum	PassageCount	PassageRetrieval-en	LCC	RepoBench-P	Avg.
LLaMA-3.1-8B-Instruct
FKV	29.67	45.58	55.09	56.47	45.07	32.23	35.32	24.94	27.17	72.50	91.48	43.58	7.93	99.50	54.72	52.28	48.35
SnapKV	25.70	31.04	51.46	54.50	42.70	26.86	22.87	23.60	22.53	55.50	91.57	41.58	8.42	99.50	53.85	44.84	43.53
QUEST	29.21	43.63	54.39	54.89	44.48	28.53	32.66	24.51	26.41	72.00	91.67	42.21	8.04	98.00	54.00	46.40	46.94
SqueezedAttn.	29.32	45.20	55.31	55.43	46.02	30.45	35.41	24.30	27.14	72.50	91.75	42.03	8.65	99.00	54.64	50.12	47.95

Appendix

Table 16: Complete task-level LongBench results at a KV budget of 256 tokens. The final column reports the average across all tasks.

LLaMA-3.1-8B-Instruct
Method	Avg.	NIAH-S1	NIAH-S2	NIAH-S3	NIAH-MK1	NIAH-MK2	NIAH-MK3	NIAH-MV	NIAH-MQ	VT	CWE	FWE	QA1	QA2
32K Context
FKV	87.95	100.00	100.00	100.00	100.00	96.00	100.00	99.00	100.00	95.20	45.20	92.00	72.00	44.00
SnapKV	65.05	100.00	100.00	4.00	100.00	88.00	0.00	85.00	92.00	90.40	3.60	62.67	76.00	44.00
QUEST	74.92	100.00	100.00	100.00	100.00	80.00	4.00	95.00	99.00	94.40	6.80	78.70	68.00	48.00
SqueezedAttn.	83.14	100.00	100.00	100.00	100.00	96.00	56.00	99.00	100.00	93.60	61.60	50.67	76.00	48.00

Appendix

Table 17: Complete task-level RULER accuracy under a fixed 1.5625% KV budget. The Avg. column reports the average across the 13 RULER tasks.

	AIME 2025			MATH-500			GPQA-Diamond
Method	512	1024	2048	512	1024	2048	512	1024	2048
FKV	69.58			94.53			57.95
TriAttention	5.83	20.00	37.50	73.80	86.20	93.20	31.82	40.40	50.13
QUEST	47.08	62.92	66.25	92.73	93.50	89.93	57.01	58.14	56.12
SeerAttention-R	26.67	42.08	55.00	85.70	91.80	94.03	54.67	55.93	58.52
FASA	30.83	60.42	67.50	87.93	93.58	94.60	56.06	57.45	59.97

Appendix

Table 18: Complete reasoning accuracy on Qwen3-8B across KV budgets. FKV is reported once because it does not depend on the sparse token budget.

	AIME 2025			MATH-500			GPQA-Diamond
Method	512	1024	2048	512	1024	2048	512	1024	2048
FKV	90.00			97.15			80.87
TriAttention	13.33	22.50	54.17	68.35	83.95	95.35	40.28	43.94	68.69
QUEST	28.75	50.42	70.42	88.20	93.65	96.33	63.07	71.40	77.84
FASA	57.92	73.75	84.17	95.25	97.08	97.13	78.91	80.56	80.87
JSSA	88.75	86.67	88.33	96.18	97.08	97.10	80.68	81.94	81.12

Appendix

Table 19: Complete reasoning accuracy on Gemma4-26B-A4B-It across KV budgets. FKV is reported once because it is independent of the sparse token budget.

	AIME 2025			MATH-500			GPQA-Diamond
Method	512	1024	2048	512	1024	2048	512	1024	2048
FKV	90.42			96.80			75.44
FASA	37.08	65.00	81.25	91.08	95.65	96.80	55.30	69.38	73.61
JSSA	86.67	91.67	91.25	96.95	96.88	96.88	75.25	77.53	78.03

Appendix

Table 20: Complete reasoning accuracy on GLM-4.7-Flash across KV budgets. FKV is reported once because it is independent of the sparse token budget.

Model	Benchmark	Method	512	1024	2048
Qwen3-8B	AIME 2025	TriAttention	2.07 $\times$	1.75 $\times$	1.54 $\times$
		QUEST	1.51 $\times$	1.29 $\times$	1.16 $\times$
		SeerAttention-R	1.76 $\times$	1.54 $\times$	1.30 $\times$
		FASA	1.71 $\times$	1.30 $\times$	1.10 $\times$
		JSSA	1.16 $\times$	1.07 $\times$	1.04 $\times$
	MATH-500	TriAttention	2.59 $\times$	1.70 $\times$	1.22 $\times$

Appendix

Table 21: Generation length relative to FKV across reasoning benchmarks and KV budgets. Lower is better.

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,084 papers

SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference
Jun 3, 2026Yaosheng Fu, Guangxuan Xiao, Xin Dong +2KV-Cache Offloading Long-Context Language Model Inference
1NVIDIA · 2Thinking Machines Lab · 3ByteDance Seed +1
Stratified Sparse Attention: Reducing the Sparse-to-Dense Gap at Extreme Sparsity
Oct 20, 2026Sahil Joshi, Prithvi Dixit, Aryan Katiyar +5Sparse Attention Long-Context Language Model Inference
Department of Computer Science, Rice University, TX, USA. · Department of Electrical Engineering and Computer Sciences, UC Berkeley, CA, USA. · Department of Computer Science, Stanford University, CA, USA. +1
CommunityKV: Efficient Long-Context Decoding via Graph Partitioning
Sep 30, 2026Joe McKenna, Anastasios Alexandridis, Nathan Susanj +1Long-Context Language Model Inference KV-Cache Management
Amazon AGI

Jun 3, 2026J/K move · Enter open · S save

SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference

Yaosheng Fu, Guangxuan Xiao, Xin Dong, Song Han+1

1NVIDIA · 2Thinking Machines Lab · 3ByteDance Seed+1

KV-Cache Offloading Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention

Open paper AskPDF arXiv

Sparse attention reduces compute and memory bandwidth for long-context LLM inference. However, two key challenges remain: (1) KV cache capacity still grows with sequence length, and offloading to CPU memory introduces a PCIe transfer bottleneck; (2) the sparse selection step itself retains

O(T^2)

complexity and can dominate attention cost at long contexts. We propose SparDA, a decoupled sparse attention architecture that introduces a fourth per-layer projection, the Forecast, alongside Query, Key, and Value. The Forecast predicts the KV blocks needed by the next layer, enabling lookahead selection that overlaps CPU-to-GPU prefetch with current-layer execution. Because Forecast is decoupled from the attention query, our GQA implementation uses one Forecast head per GQA group, reducing selection overhead versus the original multi-head selector. SparDA adds

<

0.5% parameters and trains only the Forecast projections by matching the original selector's attention distribution. On two sparse-pretrained 8B models, SparDA matches or slightly improves accuracy and delivers up to 1.25

\times

prefill speedup and 1.7

\times

decode speedup over the sparse-attention offload baseline. By enabling larger feasible batch sizes on a single GPU, SparDA further reaches up to 5.3

\times

higher decode throughput than the non-offload sparse baseline. Our source code is available at https://github.com/NVlabs/SparDA.

Abstract

LLM Inference Acceleration

Explore similar work

SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference

Stratified Sparse Attention: Reducing the Sparse-to-Dense Gap at Extreme Sparsity

CommunityKV: Efficient Long-Context Decoding via Graph Partitioning