CLAA: Cross-Layer Attention Aggregation for Accelerating LLM Prefill
Abstract
Token-ranking heuristics accelerate the prefill bottleneck in long-context LLM inference by selectively processing semantically relevant tokens. However, current evaluation relies on end-to-end benchmarks, making it difficult to isolate the quality of the token ranking itself. To address this, we introduce an Answer-Informed Reference framework with two variants: a Model Reference (Model-Ref) that measures token importance using the model generated answer, and a Ground-Truth Reference (GT-Ref) that uses human reference answers. Using GT-Ref, we establish that as few as 10% of prompt tokens suffice to match full-context performance on LongBench across all three evaluated models. The framework also reveals that existing heuristics exhibit variance across layers, with rankings degrading sharply at specific layers. Motivated by this instability, we propose Cross-Layer Attention Aggregation (CLAA), which aggregates importance scores across consecutive layers, eliminating the layer-dependent accuracy collapse observed in single-layer methods. A meaningful gap remains between the best heuristic and GT-Ref, indicating theoretical room for improved token selection.
Lay Summary
When you give a large language model a long document and ask a question about it, the model first has to "read" the entire document before it can start answering. For inputs that span tens of pages of text, this reading step dominates the wait time, often taking far longer than producing the answer itself. Researchers have noticed that most words in a long document are irrelevant to any given question, so several recent methods try to skim. They pick out the important words first and let the model read only those. But it has been hard to tell how well these skimming methods actually work, because their performance gets tangled up with other engineering choices. We built an evaluation framework that uses the answer the model produces to figure out, in hindsight, which input words really mattered. This provides a reliable yardstick for any skimming method. Using this yardstick, we show that just 10% of the words in a long document carry enough information to fully answer typical questions, and we propose a new skimming method that selects those words more reliably than prior approaches. Our framework reveals both how much room there is to skim aggressively and how far current methods still have to go to reach that ceiling.