CRISP: Cliff-awaRe Input-adaptive Sparse Prefilling with Structural-Mass-Motivated Routing

AI summary

Overview

Research area: Efficient long-context inference for large language models — specifically sparse attention during the prefill phase.

Technical level: Advanced. The paper combines an asymptotic analysis of cumulative coverage thresholds with systems-level latency benchmarks, and assumes familiarity with softmax attention, sparse attention patterns, and Jensen-Shannon Divergence.

One-sentence scope: CRISP is a training-free dynamic sparse prefill method that replaces FlexPrefill's JSD-based head routing with a structural-mass proxy and replaces cumulative coverage thresholding on Vertical-Slash heads with a sink-aware noise-floor threshold.

What This Paper Is About

The prefill phase of long-context LLM inference scales quadratically, and self-attention is the bottleneck. Existing dynamic sparse methods route each attention head at runtime either to a Vertical-Slash (VS) pattern or a Pooled-Estimation (PE) path, but they use indirect routing signals and budget rules that overlook how attention mass is actually distributed after softmax. The paper diagnoses two specific structural problems in that paradigm — an expensive, indirect routing proxy, and a coverage threshold that accumulates noise at long contexts — and proposes a fix for each.

Key Contributions

1. A structural account of VS/PE routing. The authors show empirically that FlexPrefill's JSD signal tracks the same head-level distinction that can be read directly off the attention map, and they report the routing agreement between their proxy and JSD (94.0% on Llama-3.1-8B, 88.1% on Qwen2.5-7B).

2. The C_struct structural-mass proxy. A constant-time indexed slice reduction over the already-computed proxy attention map, measuring mass only at architectural sinks (first 128 tokens) and the local recency window (last 128 tokens). It removes the pooled matmul and KL divergence that JSD routing requires.

3. Formalization of the post-softmax "mass cliff." The paper argues theoretically that strictly cumulative coverage thresholds (γ) accumulate O(n) background noise at long contexts, and identifies two failure modes: sink-only collapse and residual noise accumulation.

4. CRISP, a sink-aware thresholding method grounded in the noise floor rather than a coverage target, which the authors report achieves parity with dense attention on InfiniteBench and up to 5.30× attention speedup at 512k tokens.

Main Findings

• Routing agreement is high. C_struct reproduces JSD's routing decision on 94.0% (Llama) and 88.1% (Qwen) of measured heads, measured on InfiniteBench inputs. The top-1 mass block is a sink or recency block for 91.3% (Llama) and 98.3% (Qwen) of heads.

• Cumulative thresholds fail in two ways. At γ=0.95, high sink mass can cause selection to terminate inside the sink ("sink-only collapse"), leaving important blocks unreached; when sink mass is distributed, selection exhausts the important band and in the paper's illustration accumulates 13 noise blocks.

• γ-tuning cannot fix it. Increasing FlexPrefill's coverage from γ=0.95 to γ=0.97 degrades some benchmarks while marginally improving others, inconsistently across models and tasks, and adds substantial latency (46,369 ms vs 40,949 ms at 512k), without retrieval improvement.

• Retrieval recovery. CRISP recovers up to +28.0 pp on retrieval tasks over FlexPrefill γ=0.95. Breakdown: KV Retrieval (InfiniteBench) +17.8pp on Llama (Qwen's baseline is near-zero, so the difference is reported as not meaningful); Passkey (InfiniteBench) 0.0pp on Llama and +28.0pp on Qwen; Passage Retrieval (En, LongBench) +12.50pp / +13.50pp; Passage Retrieval (Zh, LongBench) +3.16pp / +4.50pp; RULER 65K +2.65pp / +1.45pp.

• Parity with dense attention. On InfiniteBench at 131K, CRISP α=1.0 scores 48.7 on Llama (vs FlashAttention 48.6) and 28.7 on Qwen (vs FlashAttention 24.0). On RULER, CRISP α=1.0 scores 88.82 on Llama and 76.20 on Qwen (vs FlashAttention 89.02 and 75.84). On LongBench, CRISP α=1.0 scores 47.77 on Llama and 47.23 on Qwen (vs FlashAttention 48.82 and 48.80).

• One honest regression. On Llama RULER, CRISP shows a −0.40pp aggregate change vs FlexPrefill γ=0.95, which the authors attribute to RULER's aggregation tasks requiring broad mid-range coverage that the α=1.0 threshold occasionally treats as noise. MInference loses −3.08pp there.

• Speedups grow with context. At 512k tokens, CRISP reaches 5.40× (α=1.5), 5.30× (α=1.25), and 5.17× (α=1.0) over FlashAttention, versus 4.41× for FlexPrefill γ=0.95 and 3.90× for γ=0.97. Latency at 512k: 33,454 ms / 34,111 ms / 34,983 ms for CRISP, vs 40,949 ms and 46,369 ms for the FlexPrefill baselines.

• Sparse attention is not profitable at short contexts. On Llama-3.1-8B at 8K, FlexPrefill runs at 225.9 ms and CRISP at 209.9 ms — 4.4× and 4.1× the cost of dense attention. By 64k, CRISP is 1.76× faster than dense, so the crossover falls between 8K and 64k.

• The components interact. Routing-only (C_struct with γ-cumsum VS selection) hurts Llama-3.1-8B by −2.2pp on InfiniteBench and −1.4pp on RULER. Sink-aware selection alone is broadly beneficial. Full CRISP recovers the regression.

Methodology in Plain English

The researchers start from an existing state-of-the-art method (FlexPrefill) that makes two runtime decisions per attention head: which sparse pattern to use, and how many tokens to keep.

For the first decision, they note that FlexPrefill builds an extra pooled attention estimate and computes Jensen-Shannon Divergence against the per-query attention just to decide whether a head is "concentrated." Instead, since concentrated heads in sink-having architectures put their mass at predictable places, CRISP simply measures how much attention mass sits in the first 128 tokens and the last 128 tokens of the proxy attention map. That is a constant-time reduction over a matrix the method already computes. They then verify against FlexPrefill's JSD decisions head by head.

For the second decision, they argue that accumulating mass until reaching a fixed γ is the wrong rule, because after softmax the mass distribution is not flat: sinks hold a dominant share, task-relevant blocks hold moderate mass, and background tokens hold near-zero mass that shrinks as O(1/n) with sequence length. A coverage target therefore either stops inside the sink or wades into background noise. Their replacement computes the expected background mass over the blocks that are not always retained, and selects each block that exceeds a multiple α of that expectation. At α=1.0 the threshold is the mean background mass itself, which they present as calibration-free; the number of kept blocks becomes an output of the input rather than a preset.

They evaluate on Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct across InfiniteBench (131K), RULER (4K–131K), and LongBench (4K–16K), with attention-only latency measured on a single NVIDIA H100 80GB from 64k to 512k tokens. Baselines are FlashAttention (exact dense), MInference (offline patterns), and FlexPrefill at γ=0.95 and γ=0.97. CRISP uses a universal configuration: τ_proxy=0.2, γ=0.95 for the PE path only, α∈{1.0, 1.25} for the VS path, block size B=128, and k_min=1024 tokens (8 blocks) per direction.

Why This Matters

Research impact. The paper reframes a systems problem as a structural one: it argues that the failure of cumulative coverage thresholds is a consequence of the post-softmax mass hierarchy rather than a tuning artifact, and it shows that an existing routing signal can be replaced by a cheaper structural measurement that agrees with it most of the time. It also provides a structural account for the VS/PE dichotomy that prior theoretical work had not formalized. The reported result that a sparse method can match or exceed exact dense attention on InfiniteBench — which the authors attribute to dense attention aggregating over architectural sink noise — is a notable claim for the long-context efficiency literature.

Real-world applications:

• Long-document retrieval and question answering, where the reported double-digit recovery on retrieval tasks (KV retrieval, passkey, passage retrieval in English and Chinese) is directly relevant.
• Retrieval-augmented generation pipelines that assemble very long contexts from many retrieved chunks and need prefill to remain fast.
• Long-context agent and coding workloads where the prefill cost dominates per-request latency.
• Serving infrastructure providers who must decide whether sparse prefill is worth enabling — the 8K crossover finding suggests gating sparse prefilling on a context-length threshold.

Industry relevance. The method is training-free and computes exact softmax attention over the blocks it selects, so it does not require distillation or fine-tuning and can be dropped into existing inference stacks. The latency results are reported as attention-only on a single H100 80GB, isolating the operation being modified; end-to-end serving throughput is not reported.

Future Directions

• Extending beyond sink-having architectures. The authors state that C_struct is a structural measurement, not a formal entropy quantity, and that architectures suppressing attention sinks (for example via gated attention) break the correspondence. Sparse attention in the sink-free regime is described as largely unexplored, and a different structural proxy would be required.

• Hybrid heads. CRISP inherits FlexPrefill's binary VS/PE routing. The authors note that real heads can carry mixed structure, and forcing such a head onto the VS path lets the sink-aware threshold discard tokens that do not fit the pattern. Probabilistic mixing and per-head confidence-weighted decisions are left to future work.

• Scaling and decoding. Empirical verification on models larger than 8B and on contexts beyond established benchmarks is left as future work. The method addresses prefilling only; extending to token-level decoding under causal masking requires characterizing the mass cliff dynamically during autoregressive generation.

• Improving the PE path and scheduling α. The authors note that gains concentrate on the VS path, and leave full PE improvements and dynamic scheduling of α — to mitigate the precision-coverage tradeoff observed on aggregation tasks such as Llama RULER — to future treatment.

Target Audience

Researchers and engineers working on long-context LLM inference efficiency, particularly those implementing or evaluating sparse attention kernels and serving systems. The paper is most useful to readers already comfortable with softmax attention mechanics and sparse pattern selection. Practitioners deciding whether to adopt sparse prefill in production will find the latency scaling and the short-context crossover findings directly actionable, though they should note that this work evaluates 7–8B models and reports attention-only latency rather than end-to-end throughput.

Authors’ abstract

The attention prefilling phase of long-context LLM inference scales quadratically, making self-attention a severe computational bottleneck. Traditional sparse attention methods mitigate this through fixed patterns or offline profiling, but lack the flexibility to adapt to input-dependent attention structure. Recent dynamic methods address this by routing heads to sparse patterns in real-time, but rely on indirect routing proxies with overhead and budget allocation mechanisms that overlook the post-softmax mass hierarchy. We present CRISP (Cliff-awaRe Input-adaptive Sparse Prefilling), which identifies and addresses two structural challenges in this dynamic routing paradigm. First, we show that the routing decision can be read directly off the structure of the proxy attention map. We replace the Jensen-Shannon Divergence (JSD) routing with C_struct, a structural proxy that measures mass at Vertical-Slash compatible positions and reproduces JSD's routing decisions while eliminating both the pooled matmul and subsequent KL divergence overhead. Second, we formalize the post-softmax mass cliff and demonstrate theoretically that strictly cumulative coverage thresholds accumulate O(n) background noise at long contexts. CRISP navigates this via a sink-aware threshold grounded in the noise floor. Empirically, across InfiniteBench, RULER and LongBench on two model families, CRISP is the strongest sparse method overall and matches or exceeds exact dense attention on retrieval-heavy benchmarks, recovering up to +28.0 pp on retrieval tasks over baselines and achieving up to a 5.30x attention speedup at 512k tokens, driven primarily by our O(n) noise elimination during selection while preserving structural integrity.