Prism Attention: Proposal-Refined Index Sharing Mechanism for Efficient LLMs Inference

Zhenxu Tian; Kebin Liu; Zhengwu Yang; Yi Su; Qingqing Dang; Yanlin Sha; Kaipeng Deng; Dianhai Yu

cs.CLOct 16, 2026

Prism Attention: Proposal-Refined Index Sharing Mechanism for Efficient LLMs Inference

Authors: Zhenxu Tian, Kebin Liu, Zhengwu Yang, Yi Su, Qingqing Dang, Yanlin Sha, Kaipeng Deng, Dianhai Yu,

Organizations: Soochow University · Baidu Inc.

Abstract

Long-context decoding in large language models (LLMs) is increasingly bottlenecked by the memory and computation required to attend over an ever-growing Key-Value (KV) cache at each autoregressive step. Although token-level sparse attention reduces this burden by restricting computation to a small subset of relevant past tokens, its practical gains are often offset by repeatedly recomputing layer-wise sparse indices as the context grows. A common remedy is to share sparse indices across layers by computing them only at designated anchor layers, but such static reuse ignores the layer-wise evolution of attention patterns, causing the reused indices to progressively deviate from the layer-specific optima and incur substantial attention loss. By revisiting cross-layer sparsity from the perspective of attention dynamics, we find that inter-layer salient token shift is highly localized: although the exact top- $κ$ indices differ across layers, the core high-attention tokens remain stable, with changes concentrated near the expanded boundary of previously selected regions. Motivated by this, we propose Prism Attention, a proposal-refined sparse attention framework. It replaces rigid cross-layer index reuse with a coarse-to-fine strategy, preserving layer-wise adaptability while retaining the efficiency benefits of cross-layer sharing. Extensive experiments show that Prism Attention achieves state-of-the-art performance among mainstream sparse attention methods on both reasoning and long-context benchmarks, while attaining a $3.4\times$ end-to-end decoding speedup over FlashAttention at 128K context length.

Figures & tables

Figure 1: Cross-layer overlap ratio across varying distances of different layer offset.

Panel of Figure 2 — Figure 2: Localized evolution of sparse indices across layers within proposal regions. For each target (refined) layer, we categorize its own top-128 sparse indices according to their ranks in the corresponding reference (proposal) layer: “Prop. Top-128”, “Prop. Top-128 $\sim$ 512”, and “Outside Prop. Top-512” denote the ratios of target-layer top-128 indices falling in the reference-layer top-128 set, ranks 128 $\sim$ 511, and outside the top-512 set, respectively. The above two panels share the same legend. ( 2(a) ) tracks this distribution from proposal layers $L_{0}$ , $L_{9}$ , $L_{18}$ , and $L_{27}$ to their subsequent refined layers, showing that a slightly expanded proposal pool ( $\rho=4$ , i.e., $4\times$ the base size) consistently captures a significant majority of layer-wise sparse-index variations. ( 2(b) ) measures, for each adjacent pair $L_{\ell}\!\rightarrow\!L_{\ell+1}$ , how the top-128 indices of $L_{\ell+1}$ fall into the proposal regions of $L_{\ell}$ . This demonstrates that adjacent-layer shifts are heavily concentrated within the expanded proposal window, significantly reducing the risk of missing salient tokens while avoiding the prohibitive cost of full-context search. A more comprehensive analysis is provided in Appendix B .

Panel of Figure 3 — Figure 3: Effect of expanded proposal regions on cross-layer sparse-index coverage. We conduct the analysis on Qwen3-4B using decoding traces collected from AIME24, and report the average cross-layer coverage of refined-layer top-128 sparse indices by the reference-layer proposal set. Left: the reference proposal is restricted to its top-128 indices, corresponding to rigid top- $\kappa$ index reuse by the refined layer. Right: the reference proposal is expanded from top-128 to top-256 ( $\rho=2$ ), doubling the candidate pool used for refinement while keeping the refined-layer sparse support fixed at top-128. Compared with rigid top-128 reuse, the expanded proposal region preserves substantially higher coverage across layers, showing that many refined-layer salient tokens missed by rigid reuse are still captured within a moderately enlarged candidate window.

Figure 4: An overview of Prism Attention. Algorithm 1 presents the per-step decoding procedure.

LLMs	Methods	AIME24	AIME25	GPQA	MATH500	Avg.
Qwen3-4B	Full Attention	71.25	66.41	56.19	93.93	71.95
	Quest	46.67	39.79	24.74	73.40	46.15
	Seer-R	69.32	58.59	55.40	94.10	69.35
	LessIsMore	71.48	62.87	56.23	94.12	71.18
	IndexCache	70.99	61.30	54.29	94.05	70.16
	Prism	72.60	64.27	54.83	94.00	71.43

Table 1: Accuracy comparison on reasoning benchmarks using Qwen3-4B and Qwen3-8B.

	Methods	2WikiMQA	TriviaQA	HotpotQA	TREC	Qasper	PCount	MF-en	PRe	GovReport	Avg.
	Full Attention	62.68	89.69	69.79	45.00	29.35	22.26	49.16	99.06	28.66	55.07
Qwen3-4B	Seer-R	59.07	90.42	69.59	48.00	19.19	24.36	50.49	100.00	28.19	54.37
	LessIsMore	62.57	82.62	70.73	41.00	21.43	16.67	49.83	98.00	28.42	52.36
	IndexCache	57.73	87.28	70.64	37.00	26.92	22.12	49.16	100.00	28.54	53.27
	Prism	60.36	89.76	70.16	52.00	23.89	20.35	49.75	100.00	28.57	54.98
	Full Attention	64.77	89.02	70.45	31.33	31.52	24.95	50.54	100.00	28.66	54.58

Table 2: Performance comparison of different methods across various LLMs on LongBench.

Figure 5: Accuracy-Efficiency Pareto Comparison.

Panel of Figure 6 — Figure 6: ( 6(a) ) Kernel latency breakdown at BSZ=4 across different context lengths. ( 6(b) ) The speedup of the average per-layer kernel latency over FlashAttention across context lengths and batch sizes.

Figure 7: End-to-end latency comparison across different context lengths on Qwen3-4B at BSZ=16.

Per-Layer Indices	Reuse Indices	Quantization	Group	Calibration	Proposal-Refined	Qwen3-4B
Per-Layer Indices	Reuse Indices	Quantization	Group	Calibration	Proposal-Refined	AIME24	AIME25
$\checkmark$						71.35	63.87
$\checkmark$		$\checkmark$				72.08	63.26
	$\checkmark$	$\checkmark$	4			65.68	50.05
	$\checkmark$	$\checkmark$	4	$\checkmark$		72.19	59.22
	$\checkmark$	$\checkmark$	4		$\checkmark$	71.56	62.60

Table 3: Ablation study of index construction and proposal-refined sharing on Qwen3-4B. Group

x

denotes that one in every

x

layers acts as a proposal layer (i.e., a proposal-to-refined ratio of

1{:}(x-1)

). All configurations use a sparse budget of

\kappa=1024

; in proposal-refined configurations, each refined layer selects its final

\kappa

tokens from a candidate pool of

n=\rho\kappa

with

\rho=4

Figure 8: Sensitivity to proposal expansion factor

\rho

and sparse budget

\kappa

on AIME25.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Cross-layer index/attention-mass distribution (Left) and attention score recovered by proposal expansion (Right) across model architectures. Left figure : within each pair of bars, the solid bar denotes the index overlap ratio (the share of the target/refined layer’s 128 indices) and the hatched bar denotes the attention-mass ratio (the share of the total attention score those indices carry in the target layer). Both are normalized over the target layer’s own Top-128, so the horizontal axis is a percentage and each bar sums to 100%. Right figure : the legend groups samples by the original attention coverage, i.e., the share of the target layer’s Top-128 attention score that is already covered by reusing only the proposal layer’s Top-128 indices. The vertical axis reports the additionally recovered attention score (in %) obtained by expanding the proposal region from Top-128 to Top-512, measured against the same target-layer Top-128 attention total, so both quantities are expressed on a common scale. The recovered attention decreases consistently as the original coverage grows: when the proposal layer’s Top-128 already captures most of the target layer’s attention, little is left for expansion to recover, whereas in the low-coverage regime—exactly the cases where directly reusing the proposal indices would be most lossy—the expanded region recovers the largest amount of attention score. Note: for conciseness, we abbreviate DeepSeek-R1-Distill-Llama-8B as DSR1-8B and GLM-4.7-Flash as GLM4.7 in both figures above.

Figure 10: Cross-layer index/attention-mass distribution (Left) and attention score recovered by proposal expansion (Right) across context lengths. Both follow the conventions of Figure 9 .

Figure 11: Cross-layer index/attention-mass distribution (Left) and attention score recovered by proposal expansion (Right) across task types. Both follow the conventions of Figure 9 . Note: for conciseness, we abbreviate Mathematical Reasoning, Summarization, Few-shot Learning, and Multi-Document QA as Math, Summ, FS, and MDQA, respectively, in both figures above.

Panel of Figure 12 — Figure 12: Effect of expanded proposal regions on cross-layer sparse-index coverage, with the refined-layer sparse support fixed at top-128 (Qwen3-4B; GPQA Rein et al. (2023) and MATH500 Hendrycks et al. (2020) traces).

Model	Config	Pattern
Qwen3-4B	group-4	$\mathcal{FSSSSSSFFSSSSFSFSFFSSSSSSSFSSSSSSSSF}$
Qwen3-4B	group-9	$\mathcal{FSSSSSSFSSSSSFSSSSFSSSSSSSSSSSSSSSSS}$
Qwen3-8B	group-4	$\mathcal{FSSSSSSFFSSSSFFSSSFSFFSSSSSSFSSSSSSS}$
Qwen3-8B	group-9	$\mathcal{FSSSSSSFSSSSSFFSSSSSSSSSSSSSSSSSSSSS}$

Appendix

Table 4: Searched Proposal/Refined layer assignment patterns ( F : Proposal layer, S : Refined layer).

Figure 13: Sensitivity of proposal layer assignment to calibration sample selection on Qwen3-4B.

Metric	Seed 0	Seed 1	Seed 2	Seed 3	Seed 4
AIME25	64.43	64.32	64.69	64.06	63.91

Appendix

Table 5: Performance across layer assignment patterns from random calibration seeds.

Batch Size	Component	8K	16K	32K	64K	128K
BSZ=1	FlashAttention	0.0293	0.0463	0.0679	0.1097	0.1945
	Quantization	0.0111	0.0109	0.0111	0.0109	0.0109
	Proposal Construction	0.0132	0.0182	0.0270	0.0446	0.0807
	Index Refinement	0.0108	0.0111	0.0109	0.0109	0.0109
	Sparse Decode	0.0333	0.0331	0.0333	0.0332	0.0334
BSZ=4	FlashAttention	0.0644	0.1066	0.1911	0.3617	0.6985

Appendix

Table 6: Detailed latency breakdown (ms) of Prism Attention kernels compared with FlashAttention across various batch sizes and sequence lengths.

Panel of Figure 14 — Figure 14: Component-wise latency breakdown of Prism Attention compared to the FlashAttention baseline across different sequence lengths at (a) BSZ=1 and (b) BSZ=16.

BSZ	Seq	Gate	Refine	Eff.	Gate	Reuse	Eff.	FA
		Prism			IndexCache
$1$	8K	0.0328	0.0208	0.0238 (0.61 $\times$ )	0.0279	0.0116	0.0157 (0.93 $\times$ )	0.0145
	16K	0.0491	0.0198	0.0271 (0.67 $\times$ )	0.0561	0.0116	0.0227 (0.80 $\times$ )	0.0182
	32K	0.0490	0.0196	0.0269 (0.79 $\times$ )	0.0565	0.0117	0.0229 (0.93 $\times$ )	0.0213
	64K	0.0518	0.0197	0.0278 (0.89 $\times$ )	0.0578	0.0115	0.0231 (1.07 $\times$ )	0.0247
	128K	0.0545	0.0198	0.0285 (1.50 $\times$ )	0.0603	0.0118	0.0240 (1.78 $\times$ )	0.0426

Appendix

Table 7: Attention-path kernel latency breakdown under a

2

K sparse budget, measured with our Triton kernels integrated into the SGLang serving framework. Gate and Refine denote the full layer chain of an index-construction layer (Prism: proposal, IndexCache: anchor) and an index-reuse layer (Prism: refined, IndexCache: shared), respectively. Eff. is the effective cost at the

1{:}3

construction-to-reuse ratio used throughout our experiments, i.e.,

T=(T_{\mathrm{construct}}+3\,T_{\mathrm{reuse}})/4

; parenthesized values are speedups over FlashAttention (FA). All latencies in ms.

	Seq $=16$ K		Seq $=64$ K		Seq $=128$ K
Method	BSZ=1	BSZ=16	BSZ=1	BSZ=16	BSZ=1	BSZ=16
FA	2.151 (1.00 $\times$ )	4.057 (1.00 $\times$ )	2.623 (1.00 $\times$ )	8.748 (1.00 $\times$ )	3.074 (1.00 $\times$ )	15.025 (1.00 $\times$ )
Seer-R	2.520 (0.85 $\times$ )	3.101 (1.31 $\times$ )	2.580 (1.02 $\times$ )	3.208 (2.73 $\times$ )	2.659 (1.16 $\times$ )	3.483 (4.31 $\times$ )
LIM	2.399 (0.90 $\times$ )	3.465 (1.17 $\times$ )	2.517 (1.04 $\times$ )	4.166 (2.10 $\times$ )	2.667 (1.15 $\times$ )	5.147 (2.92 $\times$ )
IDC-g4	2.261 (0.95 $\times$ )	2.854 (1.42 $\times$ )	2.317 (1.13 $\times$ )	3.550 (2.46 $\times$ )	2.450 (1.26 $\times$ )	4.434 (3.39 $\times$ )
IDC-g9	2.026 (1.06 $\times$ )	2.557 (1.59 $\times$ )	2.062 (1.27 $\times$ )	2.912 (3.00 $\times$ )	2.156 (1.43 $\times$ )	3.348 (4.49 $\times$ )

Appendix

Table 8: End-to-end per-step decoding latency on Qwen3-4B with the SGLang serving framework under a

2

K sparse budget. Each cell reports the absolute latency (ms) and, in parentheses, the speedup over FlashAttention (FA).

Config	Pattern
21F57S (Default)	$\mathcal{FFFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSS}$
10F68S	$\mathcal{FFSSSSSSSSSSSSFSSSFSSSSSSSFSSSSSSSFSSSSSSSFSSSSSSSSSSSSSSSSSSSFSSSSSSSFSSSFSSS}$
10F11R57S	$\mathcal{FFRSSSRSSSRSSSFSSSFSSSRSSSFSSSRSSSFSSSRSSSFSSSRSSSRSSSRSSSRSSSFSSSRSSSFSSSFSSS}$

Appendix

Table 9: Layer-assignment patterns evaluated on GLM-5.2 under a unified sparse budget

\kappa=2048

and proposal expansion factor

\rho=4

. Note: the meanings of F, S, and R here differ from those in Table 4 , and throughout this table follow the definitions given below. F (Full) layers construct sparse indices from the full context, S (Shared) layers directly reuse the indices of their nearest preceding index-producing layer, and R (Refined) layers refine an inherited candidate pool. 21F57S is the default pattern of GLM-5.2 itself, i.e., its native trained IndexCache configuration, where each F layer builds its own

\kappa

indices and each S layer reuses these

\kappa

indices from the nearest preceding F layer. 10F68S is an IndexCache variant calibrated to fewer F layers. 10F11R57S incorporates Prism as a hybrid of F, R, and S layers: each F layer builds an expanded pool of

\rho\kappa

candidate indices, each R layer selects its own top-

\kappa

from the

\rho\kappa

pool of the nearest preceding F layer, and each S layer reuses the

\kappa

indices of its nearest preceding F or R layer.

Config	LongBenchV2_128K	Ruler_64K	Ruler_128K
21F57S	65.91 / 4696	96.03 / 3802	95.53 / 3254
10F68S	65.96 / 4972	95.68 / 3507	93.05 / 3224
10F11R57S	66.06 / 4798	96.24 / 3086	95.15 / 2752

Appendix

Table 10: Accuracy on GLM-5.2 with the SGLang framework under a unified sparse budget

\kappa=2048

. Each cell reports accuracy and, to the right of the slash, the average output length. The _

L

suffix in each dataset name denotes the maximum input length.

	Config	S-NIAH-1	S-NIAH-2	S-NIAH-3	MK-NIAH-1	MK-NIAH-2	MK-NIAH-3	MQ-NIAH	MV-NIAH	HotpotQA	SQuAD	CWE	FWE	VT	AVG
64K	21F57S	100.00	100.00	98.00	100.00	100.00	100.00	99.00	100.00	79.00	84.50	97.90	90.00	100.00	96.03
	10F68S	100.00	100.00	100.00	100.00	94.00	99.00	100.00	100.00	75.00	85.75	94.70	95.33	100.00	95.68
	10F11R57S	100.00	100.00	100.00	99.00	100.00	100.00	97.00	100.00	78.00	82.83	99.30	95.00	100.00	96.24
128K	21F57S	100.00	100.00	100.00	100.00	100.00	100.00	100.00	98.00	76.00	88.17	95.00	84.67	100.00	95.53
	10F68S	100.00	100.00	100.00	100.00	86.00	88.00	96.00	98.50	76.00	84.83	89.60	90.67	100.00	93.05
	10F11R57S	100.00	100.00	98.00	100.00	98.00	100.00	100.00	94.50	80.00	79.50	95.00	92.00	100.00	95.15

Appendix

Table 11: Per-subtask accuracy (%) on RULER (64K and 128K) with GLM-5.2 under the SGLang framework and a unified sparse budget

\kappa=2048

Long-Context Language Model Inference

Long-context language model inference makes processing and generating over lengthy prompts faster and more memory-efficient.

See all 227 papers

You Only Index Once: Cross-Layer Sparse Attention with Shared Routing
Jun 4, 2026Yutao Sun, Yanqi Zhang, Li Dong +2Long-Context Language Model Inference KV-Cache Management
Microsoft Research · Tsinghua University
CSAttention: Centroid-Scoring Attention for Accelerating LLM Inference
Mar 30, 2026Chuxu Song, Zhencan Peng, Jiuqi Wei +1LLM Inference Acceleration Long-Context Language Model Inference
Department of Computer Science, Rutgers University · Oceanbase, Ant Group.
Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference
Dec 18, 2025Dhruv Deshmukh, Saurabh Goyal, Nipun Kwatra +1Self-Attention Long-Context Language Model Inference
Microsoft Research India.

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

You Only Index Once: Cross-Layer Sparse Attention with Shared Routing

Yutao Sun, Yanqi Zhang, Li Dong, Jianyong Wang+1

Microsoft Research · Tsinghua University

Long-Context Language Model Inference KV-Cache Management LLM Inference Acceleration Sparse Attention Cross-Layer Attention

Open paper AskPDF arXiv

Long-context inference in modern LLMs is increasingly constrained by decoding efficiency, especially in reasoning-heavy settings where models generate long intermediate chains of thought. Existing sparse attention methods often face a practical efficiency-quality trade-off. Structured block sparse methods typically provide stronger acceleration but incur noticeable quality loss, while token sparse methods are usually more accurate yet deliver limited end-to-end speedup because top-k routing over the full cache remains expensive. In this work, we propose cross-layer sparse attention (CLSA), which is built on top of KV-sharing architectures such as YOCO. The core idea is to share not only the KV cache across cross-decoder layers, but also the routing index. A single indexer computes token-level top-k selection once and reuses the resulting index across layers, thereby preserving the fine-grained selectivity of token sparse attention while amortizing the routing overhead. The resulting architecture improves all major inference bottlenecks jointly, including pre-filling, KV-cache storage, and long-context decoding. Experiments across short-context and long-context benchmarks show that CLSA is both accurate and efficient, achieving up to 7.6x decoding speedup and 17.1x overall throughput improvement at 128K context. These results suggest a more complete architectural solution for long-context LLMs that jointly advances model quality and inference efficiency.

Abstract

Long-Context Language Model Inference

Explore similar work

You Only Index Once: Cross-Layer Sparse Attention with Shared Routing

CSAttention: Centroid-Scoring Attention for Accelerating LLM Inference

Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference