PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Yanjun Ma, Juntao Li, Min Zhang

arXiv:2610.22985v1cs.CLcs.AI

TL;DR

Long-context token-level sparse attention is hindered by repeated index construction and static cross-layer reuse that misses evolving layer-specific attention patterns. Prism Attention refines expanded proposals to retain layer-wise adaptation, reporting the best average accuracy among evaluated sparse methods and a 3.4× decoding speedup at 128K context.

  • Problem

    Token-level sparse attention reduces computation over growing KV caches, but repeatedly constructing layer-wise indices is costly, and static cross-layer reuse ignores evolving attention patterns.

    Source: §1 · pp. 1–2

  • Method

    Prism Attention builds expanded candidate pools at proposal layers, lets refined layers reselect their own tokens within them, and uses quantized proposal generation to reduce construction cost.

    Source: §1 · pp. 2–3

  • Results

    Prism achieves the best average accuracy among all evaluated sparse attention methods and a 3.4× end-to-end decoding speedup over FlashAttention on 128K-token contexts.

    Source: §6 · p. 10

  • Takeaways & Limitations

    Preserving layer-wise adaptability within a shared proposal region is an effective way to combine sparse-attention accuracy with efficient long-context decoding.

    Source: §6 · p. 10

    Prism inherits a frozen IndexCache partition calibrated greedily on a held-out mathematical-reasoning corpus, while task-dependent drift leaves calibration alternatives for exploration.

    Source: F Limitation and Future Work · p. 20

Abstract

from arXiv · show

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-$\kappa$ 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.

1 Introduction

Token-level sparse attention reduces the burden of attending over growing KV caches, but index construction and static cross-layer reuse limit its practical efficiency and accuracy. Prism Attention uses proposal refinement to retain layer-wise adaptation, and reports strong benchmark performance with faster 128K decoding.

1 Introduction

  • Repeatedly computing layer-wise sparse indices requires scoring the full context and can offset the practical gains of token-level sparsity as sequences grow.At each decoding step, attention over the accumulated KV cache also scales linearly with sequence length.Source: §1 · pp. 1–2
  • Static cross-layer reuse reduces redundant computation but ignores evolving attention patterns, so overlap with anchor-layer indices falls with layer distance and more sharply at longer contexts.This can cause reused indices to miss critical tokens and degrade long-context performance.Source: §1 · p. 2
  • Inter-layer salient-token shifts are localized: core high-attention tokens remain stable while most changes occur near the boundary of selected regions.Prism builds an expanded candidate pool at proposal layers, lets refined layers select tokens within it, and uses quantized proposal scoring to reduce construction cost.Source: §1 · pp. 2–3
  • 3.4× end-to-end decoding speedup over FlashAttention at 128K context length accompanies state-of-the-art performance among mainstream sparse attention methods on reasoning and long-context benchmarks.Prism Attention combines Dynamic Index Refinement with Quantized Proposal Generation.Source: §1 · p. 2

2 Related Work

Prior sparse-attention work explores block-level and token-level selection, while cross-layer sharing addresses the cost of repeatedly constructing token indices. These approaches leave a trade-off between index-construction overhead and the accuracy of reused indices.

2.1 Sparse Attention

  • Token-level sparsity preserves finer-grained context than block-level sparsity, but recomputing token indices at every layer adds computational overhead.The cited block-level approaches target different inference phases, including prefill, decoding, or both.Source: §2.1 · p. 3

2.2 Cross-Layer Index Sharing

  • Cross-layer methods propagate top-κ indices from anchor layers, but sharing across more reuse layers lowers overhead while increasing drift from layer-specific optima.Kascade and HyLRA optimize anchor placement with dynamic programming, while IndexCache uses greedy search to mitigate error propagation.Source: §2.2 · p. 3

3 Observations

Sparse indices drift from layer-wise optima as layers move farther from an anchor, but salient-token changes are localized rather than arbitrary. Expanded proposal regions recover coverage of tokens that rigid reuse misses.

3 Observations

  • As distance from an anchor layer increases, index overlap shifts across depth, and rigid reuse deviates from the layer-wise optimum.The passage attributes this deviation to information loss during autoregressive decoding.Source: §3 · p. 3
  • High-attention tokens remain stable across adjacent layers, while new selections mostly fall between the original and expanded proposal boundaries rather than across the full context.This localized pattern suggests that an expanded candidate pool can retain high coverage for subsequent layers.Source: §3 · pp. 3–4
  • A moderately expanded proposal region preserves substantially higher cross-layer coverage than rigid top-128 reuse, with refined-layer support fixed at top-128.The comparison uses Qwen3-4B decoding traces from AIME24 and expands the proposal to top-256 (ρ = 2).Source: Figure 3 · p. 4, §3 · p. 4

4 Method

Prism Attention replaces rigid cross-layer index reuse with a coarse-to-fine scheme: proposal layers transmit expanded candidate pools, and refined layers select locally within them. INT2 key-cache compression makes proposal scoring more efficient during long-context decoding.

4 Method

  • Prism Attention transmits expanded candidate pools from proposal layers for refined layers to adapt their sparse support, while retaining the computational advantages of cross-layer reuse.The method also compresses the proposal-layer key cache to INT2 for candidate scoring.Source: §4 · p. 4

4.1 Proposal-Refined Index Sharing

  • Static reuse assumes salient-token identities remain unchanged, but layer-wise optima shift with depth and mismatch accumulates farther from the proposal layer.Prism instead forms an expanded proposal pool and lets refined layers recompute scores within that inherited region before selecting final support.Source: §4.1 · p. 5

4.2 Quantized Proposal Construction

  • Proposal construction remains bandwidth-intensive because proposal layers repeatedly access the historical key cache, so Prism uses group-wise asymmetric 2-bit quantization for candidate scoring.The proposal score aggregates query heads assigned to a shared KV group; max-reduction lets each KV head be loaded once per decoding step.Source: §4.2 · pp. 5–6

5 Experiment

Prism is evaluated on reasoning and long-context tasks against full attention and sparse-decoding baselines, with accuracy, latency, and component ablations examined. Its results show strong long-context accuracy and a favorable accuracy-efficiency trade-off, while ablations indicate that calibration and local refinement help counter the losses of direct index reuse.

5.1 Experimental Setup

  • Evaluation covers Qwen3-4B and Qwen3-8B on LongBench subsets averaging over 8192 input tokens and reasoning benchmarks AIME-24/25, MATH500, and GPQA-Diamond.Reasoning accuracy is Pass@1; maximum generation lengths are 4096 tokens for LongBench and 32768 for reasoning tasks.Source: §5.1 · p. 6

5.3 Efficiency Analysis

  • 64.43 accuracy at 3.84 ms makes Prism-g4-1K 1.84× faster than Full Attention while retaining 97.0% of its accuracy; it strictly dominates LIM-2K and Seer-R-4K.On AIME 25 with Qwen3-4B, Figure 5 compares the accuracy-efficiency trade-off across sparse budgets; IDC-g4-1K is 0.55 ms faster but 5.21 points less accurate.Source: Figure 5 · p. 7, §5.3 · p. 9

5.2 Accuracy Results

  • 54.98 on Qwen3-4B and 54.97 on Qwen3-8B are Prism’s LongBench average scores; Prism leads sparse methods, improving over IndexCache by 1.71 and over Seer-R and IndexCache by 1.33 and 2.32 points, respectively.On Qwen3-8B, Prism slightly exceeds Full Attention.Source: §5.2 · p. 7

5.3 Efficiency Analysis

  • At 128K and BSZ=16, Prism’s average per-layer kernel latency is 7.50× faster than FlashAttention; refinement latency stays about 0.011 ms as context grows.At BSZ=4, proposal construction grows linearly with context, while index refinement remains constant; the calibrated 1:3 proposal-to-refined ratio amortizes proposal cost.Source: §5.3 · p. 8, Figure 6 · p. 8

5.4 Ablation

  • On AIME25, direct cross-layer index reuse drops accuracy from 63.87 to 50.05, while calibration plus proposal-refined sharing reaches 64.43 under group-4.The same improvement over direct reuse holds under group-9, where one proposal layer is followed by eight refined layers.Source: §5.4 · p. 9, Table 3 · p. 9
  • Increasing ρ from 1 to 4 raises accuracy from 59.22 to 64.43, but ρ=8 slightly degrades it; increasing κ from 512 to 4096 raises accuracy from 57.66 to 65.26.Most of the κ-related gain is achieved by κ=1024; the reported effective balance uses ρ between 2 and 4 and κ between 1024 and 2048.Source: §5.4 · p. 9, Figure 8 · p. 9

6 Conclusion

Prism Attention replaces rigid cross-layer index reuse with coarse-to-fine proposal and refinement. The paper reports the best average accuracy among evaluated sparse methods and a 3.4× end-to-end decoding speedup over FlashAttention on 128K-token contexts.

6 Conclusion

  • Prism Attention replaces rigid cross-layer index reuse with a coarse-to-fine strategy.Proposal layers generate expanded candidate regions, and refined layers re-select top-κ tokens within them; refinement cost is independent of context length, while INT2-quantized scoring reduces proposal-construction cost.Source: §6 · p. 10

A Detailed Settings for the Observation Experiments

The appendix tests whether localized cross-layer attention shifts and proposal-expansion benefits extend across architectures, context lengths, and task types. The experiments compare target-layer top-128 indices with their ranks in a corresponding proposal layer.

A Detailed Settings for the Observation Experiments

  • The observation experiments use Qwen3-4B on AIME24, MATH500, and GPQA, averaging overlap across heads, decoding steps, samples, and valid layer pairs.One setup records 32 consecutive decoding steps around 1K and 8K context; another pools datasets at about 8K context.Source: A Detailed Settings for the… · p. 13

B Generalization of Cross-Layer Attention Locality Across Architectures, Context Lengths, and Task Types

  • The generalization analysis spans three architectures, 4K–16K sequence lengths, and mathematical reasoning, summarization, few-shot learning, and multi-document QA.It categorizes each target layer’s top-128 indices by their ranks in the corresponding reference layer.Source: B Generalization of Cross-Layer… · p. 13

B.1 Different Model Architectures

  • Across Qwen3-4B, DeepSeek-R1-Distill-Llama-8B, and GLM-4.7-Flash, localized drift and proposal expansion’s recovery of attention score hold across GQA and MLA.The comparison uses 8K-token mathematical-reasoning sequences.Source: §B.1 · p. 13, Figure 9 · p. 14

B.2 Different Sequence Lengths

  • Across Qwen3-4B sequences of 4K, 8K, and 16K tokens, localized shifts and proposal expansion’s recovery of additional attention score persist.Other analysis settings are held unchanged.Source: §B.2 · p. 14, Figure 10 · p. 14

B.3 Different Task Types

The task-type experiments find varying degrees of cross-layer index drift, while expanded proposal regions retain candidate tokens with substantial attention contributions across the evaluated tasks.

B.3 Different Task Types

  • The task analysis uses 8K-token sequences, with 10 randomly selected samples for each task category.The mathematical-reasoning samples come from AIME24, MATH500, and GPQA; the other categories come from LongBench.Source: §B.3 · p. 14
  • Across mathematical reasoning, summarization, few-shot learning, and multi-document QA, drift magnitude varies by task, but expanded proposal regions retain candidates with substantial attention contributions.The localized-shift pattern is reported for all four task categories.Source: §B.3 · p. 15, Figure 11 · p. 15
  • On GPQA and MATH500, rigid top-128 reuse yields poor overlap for distant layers, while expanding the proposal pool to top-256 preserves cross-layer coverage.The visual fixes refined-layer support at top-128.Source: Figure 12 · p. 15, §B.3 · p. 16

C Sensitivity to Calibration Sample Selection

Calibration-sample sensitivity analyses find stable proposal-layer assignments and stable downstream accuracy across assignment patterns generated with different random seeds.

C Sensitivity to Calibration Sample Selection

  • Most proposal layers are stable across calibration subsets, with layers 0, 8, 13, and 35 almost always selected.The observed assignments broadly match the searched patterns in Table 4.Source: C Sensitivity to Calibration Sample… · p. 16
  • Final accuracies remain remarkably stable across layer-assignment patterns generated from different random seeds at κ = 1024 and ρ = 4.The comparison is reported in Table 5.Source: C Sensitivity to Calibration Sample… · p. 16
  • The downstream comparison tests whether assignment-pattern variation across random seeds changes performance, using a sparse budget of κ = 1024 and expansion factor ρ = 4.The reported accuracies remain stable across those patterns.Source: C Sensitivity to Calibration Sample… · p. 16

D Complete Efficiency Results: From Kernels to End-to-End

The efficiency evaluation measures Prism Attention from isolated kernels through SGLang integration and GLM-5.2 accuracy. It reports nearly length-independent refinement cost, slower-growing integrated costs, and accuracy recovery on GLM-5.2.

D Complete Efficiency Results: From Kernels to End-to-End

E Scaling Prism Attention to GLM-5.2

F Limitation and Future Work

Prism uses an offline greedy procedure to choose its Proposal/Refined layer partition, then freezes that partition for downstream evaluations. Because cross-layer index drift is task-dependent, the authors identify calibration strategy, objective, and corpus as future-work directions.

F Limitation and Future Work

  • The Proposal/Refined layer partition is searched once using IndexCache’s offline greedy calibration procedure, then frozen for all downstream evaluations.The search uses a held-out mathematical-reasoning corpus and token-level cross-entropy objective.Source: F Limitation and Future Work · p. 20
  • The calibration recipe is not treated as a design variable, despite cross-layer index drift being markedly task-dependent.Source: F Limitation and Future Work · p. 20
  • Future calibration exploration could test stronger-than-greedy search, objectives closer to downstream accuracy, and corpora matched to target domain and length.The authors suggest these changes could yield a more favorable layer assignment at no additional inference cost.Source: F Limitation and Future Work · p. 20

Related papers

Loading 2610.22985v1…