CLAA: Cross-Layer Attention Aggregation for Accelerating LLM Prefill

Bradley McDanel; Steven Li; Harshit Khaitan

cs.CLFeb 17, 2026

CLAA: Cross-Layer Attention Aggregation for Accelerating LLM Prefill

Authors: Bradley McDanel, Steven Li, Harshit Khaitan

Organizations: Franklin and Marshall College · Meta Reality Labs

Abstract

The prefill stage in long-context LLM inference remains a computational bottleneck. Recent token-ranking heuristics accelerate inference by selectively processing a subset of semantically relevant tokens. However, existing methods suffer from unstable token importance estimation, often varying between layers. Evaluating token-ranking quality independently from heuristic-specific architectures is challenging. To address this, we introduce an Answer-Informed Oracle, which defines ground-truth token importance by measuring attention from generated answers back to the prompt. This oracle reveals that existing heuristics exhibit high variance across layers: rankings can degrade sharply at specific layers, a failure mode invisible to end-to-end benchmarks. The diagnosis suggests a simple fix: aggregate scores across layers rather than relying on any single one. We implement this as Cross-Layer Attention Aggregation (CLAA), which closes the gap to the oracle upper bound and reduces Time-to-First-Token (TTFT) by up to 39% compared to the Full KV Cache baseline.

Figures & tables

Figure 1: Illustration of our framework for evaluating token ranking heuristics for LLM Prefill acceleration. An Answer-Informed Oracle establishes a ground-truth token ranking by aggregating attention from the generated answer back to the prompt. This approach, which measures rank similarity between heuristic outputs and the oracle, motivates our Cross-Layer Attention Aggregation (CLAA) method that achieves higher alignment with the oracle.

Figure 2: Layer-wise token ranking performance on Llama-3.1-8B-Instruct. Spearman correlation with answer-informed oracle across LongBench tasks, comparing existing heuristics to our proposed CLAA method.

	Single-Document QA			Multi-Document QA			Summarization			Few-shot Learning			Synthetic		Code
Method	NrtvQA	Qasper	MF-en	HotpotQA	2WikiMQA	MuSiQue	GovReport	QMSum	MultiNews	TREC	TriviaQA	SAMSum	PCount	PRe	LCC	RB-P	Avg.
Keep Token Rate = 100%
FullKV	30.16	45.53	54.94	56.02	46.66	31.28	35.12	25.28	27.25	73.00	91.65	43.80	8.88	99.50	63.38	56.73	49.32
Keep Token Rate = 10%
Oracle	29.85	43.94	55.85	54.99	47.11	28.92	32.26	25.29	21.95	68.50	91.43	41.80	7.47	99.50	59.63	56.85	47.83
GemFilter	24.36	21.07	39.73	51.29	33.92	25.78	28.94	18.98	17.42	60.50	91.53	40.39	4.76	87.50	22.73	32.47	37.59

Table 1: LLaMA-3.1-8B LongBench results across varying token keep rate (%).

Figure 3: Needle-in-a-Haystack result of LLaMA-3.1-8B-Instruct with 40% token keep rate. X denotes out of memory on 80GB A100.

Keep Token Rate = 100%
Method	LLaMA-3.2-3B	LLaMA-3.1-8B	Mistral-Nemo-12B
FullKV	44.23	49.32	48.32
Keep Token Rate = 10%
Oracle	44.13	47.83	47.75
GemFilter	35.20	37.59	36.60
FastKV	42.45	46.81	45.62

Table 2: Average LongBench results across different models and token keep rates.

Method	Retr.	M-Hop	Agg.	QA	Avg.
FullKV	98.27	86.88	86.67	63.30	83.78
GemFilter	82.23	58.20	88.33	63.70	73.12
FastKV	87.69	87.16	86.67	63.60	81.28
CLAA	89.85	87.72	86.67	63.80	82.01

Table 3: RULER benchmark performance by category with 64K context on LLaMA-3.1-8B with 40% token keep rate.

Task	Method	Rewrite	Total TTFT	Speedup	Score
Qasper	Full KV	—	1084 ms	1.00x	0.487
	Selective Context	3134 ms	3421 ms	0.32x	0.186
	LLMLingua	901 ms	1141 ms	0.95x	0.129
	CLAA (20%)	—	755 ms	1.44x	0.470
GovReport	Full KV	—	13505 ms	1.00x	0.377
	Selective Context	5642 ms	15391 ms	0.88x	0.327

Table 4: Comparison with semantic compression methods on LLaMA-3.1-8B. Rewrite time is the compression model overhead; Total TTFT includes both rewrite and prefill.

Figure 6: Left : Impact of the pruning layer index on LongBench accuracy. Pruning later in the model improves performance for all heuristics. Right : Impact of the cross-layer aggregation window size (

n

) at different token keep rates. Increasing the window size improves accuracy.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Parameter	Symbol	Value(s) Explored	Final Value in Main Results
Oracle	Token Keep Rate	-	{0.1, 0.2, 0.4}	Main Variable
(Ours)	Pruning Layer Index	$l_{p}$	{3, 7, 11, 15, 19}	15
FastKV	Token Keep Rate	-	{0.1, 0.2, 0.4}	Main Variable
	TSP Layer Index	$l_{\text{TSP}}$	{3, 7, 11, 15, 19}	15
	Observation Window	$W$	{8}	8
	Pooling Kernel Size	-	{7}	7

Appendix

Table 5: Hyperparameter settings for all methods evaluated on LLaMA-3.1-8B-Instruct.

Figure 7: Impact of the first compression layer index (

m

) on LongBench accuracy for CLAA. For

m>0

, layers

0

m-1

are kept uncompressed. The experiment shows that deferring compression to later layers (increasing

m

) improves accuracy, especially for higher token keep rates. These results are on Llama-3.1-8B-Instruct.

Figure 8: Detailed per-task comparison of token ranking heuristics against the Answer-Informed Oracle.

Cross-Layer Attention

See all 13 papers

Stability Implies Redundancy: Delta Attention Selective Halting for Efficient Long-Context Prefilling
Apr 20, 2026Yujie Chen, Tailai Chen, Yifeng Gao +4Self-Attention Long-Context Language Model Inference
Shanghai Jiao Tong University · University of California, San Diego · Carnegie Mellon University
UniPrefill: Universal Long-Context Prefill Acceleration via Block-wise Dynamic Sparsification
May 7, 2026Qihang Fan, Huaibo Huang, Zhiying Wu +2Long-Context Language Model Inference LLM Inference Acceleration
1MAIS&NLPR, CASIA · 2UCAS · 3WeChat, Tencent
Probe and Skip: Self-Predictive Token Skipping for Efficient Long-Context LLM Inference
Jan 19, 2026Zimeng Wu, Donghao Wang, Chaozhe Jin +2LLM Inference Acceleration Long-Context LLM Inference
School of Computer Science and Engineering, Beihang University, Beijing, China · State Key Laboratory of Virtual Reality Technology and Systems, Beihang University, Beijing, China

Apr 20, 2026J/K move · Enter open · S save

Stability Implies Redundancy: Delta Attention Selective Halting for Efficient Long-Context Prefilling

Yujie Chen, Tailai Chen, Yifeng Gao, Zoe Wanying He+3

Shanghai Jiao Tong University · University of California, San Diego · Carnegie Mellon University

Self-Attention Long-Context Language Model Inference Efficient VLM Inference LLM Inference Acceleration Attention Mechanisms

Open paper AskPDF arXiv

Prefilling computational costs pose a significant bottleneck for Large Language Models (LLMs) and Large Multimodal Models (LMMs) in long-context settings. While token pruning reduces sequence length, prior methods rely on heuristics that break compatibility with hardware-efficient kernels like FlashAttention. In this work, we observe that tokens evolve toward semantic fixing points, making further processing redundant. To this end, we introduce Delta Attention Selective Halting (DASH), a training-free policy that monitors the layer-wise update dynamics of the self-attention mechanism to selectively halt stabilized tokens. Extensive evaluation confirms that DASH generalizes across language and vision benchmarks, delivering significant prefill speedups while preserving model accuracy and hardware efficiency. Code will be released at https://github.com/verach3n/DASH.git.

Abstract

Cross-Layer Attention

Explore similar work

Stability Implies Redundancy: Delta Attention Selective Halting for Efficient Long-Context Prefilling

UniPrefill: Universal Long-Context Prefill Acceleration via Block-wise Dynamic Sparsification

Probe and Skip: Self-Predictive Token Skipping for Efficient Long-Context LLM Inference