Source-linked AI summary
RouteSparse: Input-Conditional Pattern Routing for Budgeted Long-Context Prefilling
Chao Zhang, Yifan Ji, Ziyan Zhang, Kai Song, Fei Lin
TL;DR
Long-context prefilling remains costly because self-attention computes every causal query–key pair, while fixed per-head sparse routing can be brittle across inputs. RouteSparse conditionally selects sparse patterns and budgets using a shared probe, latency constraints, and dense fallback, achieving quality close to dense attention at 6.5× speed on 128K-token prefilling.
Problem
Fixed per-head sparse routing may be brittle because attention structure and suitable sparsity budgets vary across input types.
Method
RouteSparse uses a shared low-resolution attention probe to select patterns and budgets under measured latency constraints, allocating dense computation to uncertain cases.
Results
6.5× speed with quality close to dense attention on 128K-token prefilling, with larger robustness gains under domain and length shift than fixed routing.
Takeaways & Limitations
Input-conditional routing, hardware profiling, and selective fallback each contribute measurably to the quality–latency tradeoff.
Takeaways & Limitations
Performance depends on hardware and implementation, while the primary experiments use one model size and one GPU generation.
Abstract
from arXiv · showhide
Dynamic sparse attention can reduce the quadratic cost of long-context prefilling without changing model weights. MInference assigns each attention head one pattern offline and estimates that pattern's sparse indices for every prompt. This design is efficient, but it assumes that a head's preferred pattern and sparsity budget remain suitable across inputs. We introduce RouteSparse, which routes each head and prompt segment among a small library of GPU-efficient sparse patterns. A low-cost probe estimates pattern utility and uncertainty; a latency-aware router then selects a pattern and budget, while uncertain cases fall back to a denser mask. We formulate routing as constrained risk minimization, derive an attention-output error certificate from omitted probability mass, and evaluate the method on long-context retrieval, question answering, summarization, and language modeling. On Llama 3.1-8B-Instruct with 128K-token prompts, RouteSparse achieves $6.5\times$ dense prefill speed with a 0.2-point RULER drop relative to dense attention, compared with $7.3\times$ speed and a 1.6-point drop for fixed per-head routing. Ablations confirm that input-conditional routing, hardware profiling, and selective dense fallback each contribute to the quality--latency tradeoff.
1 Introduction
Long-context prefilling makes exact self-attention costly, motivating sparse methods whose prompt-dependent attention locations challenge fixed per-head designs. RouteSparse addresses this by routing patterns and budgets conditionally on each input while accounting for uncertainty and latency.
- Motivation: Long-context prefilling exposes quadratic self-attention cost, making time to first token a dominant user-visible expense.Existing efficient kernels still compute every causal query–key pair.
- Motivation: Post-training sparsity is difficult because important attention locations depend on prompt content, despite evidence that online approximation is possible.Dynamic sparse attention and query-aware key selection motivate input-dependent index estimation.
- Approach: RouteSparse makes pattern and budget selection conditional on the current input while preserving a kernel-friendly sparse pattern library.A sampled attention probe scores candidate masks and supports per-layer routing under measured latency constraints.
- Approach: The router includes selection overhead in its latency objective and sends high-uncertainty heads to larger budgets or dense fallback.This couples hardware-aware performance with uncertainty-sensitive allocation.
- Approach: RouteSparse computes an attention-output error certificate from omitted probability mass and value-vector norms.The certificate is intended to quantify the effect of sparse omission during routing.
- Evaluation: The evaluation tests quality, latency, calibration, and robustness under domain and length shift using preregistered ablations and stress tests.These experiments assess both the quality–latency tradeoff and behavior beyond calibration conditions.
2 Preliminary
Sparse attention replaces dense query–key computation with selected structured connections, but fixed offline choices can be mismatched to prompts with different information topology. RouteSparse therefore evaluates candidate patterns and budgets per input under latency and uncertainty constraints.
- 2.1 Prefill Attention: Exact causal attention requires O(n^2d) arithmetic, whereas a binary sparse mask removes unselected logits and computes only selected interactions.The sparse mask must preserve useful attention while reducing realized kernel latency.
- 2.1 Prefill Attention: MInference assigns one structured pattern to each head offline, using A-shape, vertical-slash, and block-sparse structures mapped to efficient GPU kernels.RouteSparse retains this pattern library but changes the decision granularity.
- 2.2 Why Route Per Input?: Repeated keys, locally coherent prose, and separated evidence favor different spatial patterns, so the same head can face different information topologies.This motivates routing decisions that depend on the current prompt rather than only on head identity.
- 2.2 Why Route Per Input?: Offline selection minimizes average calibration loss, but when the best pattern varies with the input it commits to the best constant decision.An input-conditional router can approximate the input-specific optimum if quality gains exceed probe and routing overhead.
- 2.2 Why Route Per Input?: The evaluation tests whether conditional routing improves quality at matched latency, becomes more robust under shift, and benefits from selective dense fallback.These correspond to hypotheses H1, H2, and H3.
3 ROUTESPARSE
ROUTESPARSE selects structured sparse patterns and budgets conditionally for each input, using a shared probe, calibrated risk estimates, and latency-aware per-layer routing. It increases budget or falls back to dense attention when sparse candidates are uncertain.
- Candidate patterns and budgets: Each head considers A-shape, vertical-slash, block-sparse, and dense masks across kernel-matched discrete budget grids.Candidate tuples are profiled on target hardware so measured kernel latency, rather than equal theoretical FLOPs, guides selection.
- Shared attention probe: A shared probe samples recent, uniformly spaced, and boundary-adjacent queries while pooling keys into blocks for inexpensive candidate scoring.The probe is shared across candidates within a head and batched across heads; its softmax is only a low-resolution indexing signal.
- Risk estimation: Probe mass coverage, entropy, local concentration, and sampling-group agreement expose diffuse or nonstationary attention where aggressive sparsity is risky.The router combines these descriptors with disagreement and an optional calibrated output-error predictor; setting β = 0 gives a training-free router.
- Error certificate: The method uses omitted attention mass to form an empirical output-error certificate, correcting estimated retained mass by a held-out one-sided residual quantile.Because probe mass is estimated rather than exact, the uncorrected bound is not a formal certificate unless estimation error is bounded.
- Measurement and calibration: End-to-end latency includes probe, index construction, routing, and regrouping rather than measuring sparse kernels alone.Calibration separately profiles candidate-kernel median and tail latency and uses dense prompts to calibrate retained-mass and output-error estimates without changing model weights.
- Latency-aware routing: For each layer, a discrete router chooses one pattern-budget option per head under a latency budget using greedy risk-per-microsecond substitutions.Routing can run concurrently with projection kernels or as a fused GPU reduction, and selected heads are grouped to reduce kernel-launch fragmentation.
4 Evaluation
The evaluation measures quality, latency, calibration, and robustness across long-context tasks and controlled shifts, using dense attention and multiple sparse-routing baselines. Ablations and stress tests isolate the contributions of conditional routing, uncertainty handling, and dense fallback.
- Experimental setup: Experiments use Llama 3.1-8B-Instruct with a 128K-token context on one A100 80GB GPU across retrieval, question answering, summarization, code, synthetic retrieval, and language modeling.RULER covers retrieval, tracking, and aggregation; InfiniteBench covers long-document tasks; PG-19 measures long-form language modeling.
- Baselines: Baselines include dense FlashAttention-2, fixed per-head MInference, latency-matched enlarged MInference, a global layer router, and static local-plusglobal attention.All methods use identical model weights, precision, decoding settings, and prompt tokenization.
- Metrics: Reported measures include end-to-end prefill latency, tail latency, memory, index-building time, routing time, kernel time, benchmark scores, perplexity, and dense-audited output error.Latency is evaluated at 32K, 64K, and 128K tokens, alongside the largest supported length.
- Comparison protocol: Paired prompts, repeated warm-up timing, bootstrap confidence intervals, and latency–quality Pareto curves support comparisons across operating points.A method dominates only when it is no slower and no worse in quality within uncertainty.
- Ablations: Ablations test input routing, uncertainty, dense fallback, fixed budgets, query selection, FLOPs-based prediction, and per-head launches.The hypotheses assess matched-latency quality, gains under domain or length shift, and whether selective fallback outperforms uniform budget expansion.
- Robustness: Stress tests vary domain composition, evidence placement, repeated-token frequency, and calibration transfer from prose to code.They report routing choices, fallback rate, certificate coverage, and largest quality drop; high fallback can preserve quality while removing speedup.
5 Results & Analysis
At 128K tokens, RouteSparse improves the quality–latency tradeoff over fixed routing, especially under shift, while selective fallback and hardware-aware execution materially shape performance.
- 1.4 RULER points are recovered by ROUTESPARSE at 0.7 seconds of additional latency over fixed MInference.
- Conditional routing is most useful in the middle latency–quality regime, where aggressive sparsity omits too much mass and near-dense budgets leave little quality to recover.
- 0.2 and 2.0 points are the quality losses for conditional routing with fallback in-domain and under combined domain-and-length shift, versus 1.6 and 6.4 for fixed routing.
- Hardware profiles preserve quality while predicted-FLOPs routing costs 0.8 seconds because theoretically cheap kernel mixtures are poorly utilized.
- Per-head launches are slower still, showing that routing overhead and sparse-kernel execution must be evaluated together.
- Selective fallback recovers 0.6 average points and 2.8 worst-decile points over no fallback at a cost of 0.8 seconds, with a measured 6.8% fallback rate.
6 Related Work
Related work spans exact-attention execution, structured and dynamic sparsity, KV-cache compression, context extension, and long-context evaluation; RouteSparse targets post-training prefill execution.
- Exact-attention systems improve tiling, parallelism, distributed processing, or KV-cache batching, serving as complementary execution and memory-management baselines.
- Structured sparse architectures reduce quadratic attention through local, global, random, or strided connectivity, while RouteSparse targets post-training prefill acceleration with deployable kernels.
- Dynamic post-training methods estimate important keys, pages, or spatial patterns online; RouteSparse builds on MInference by adding input-conditional, latency-constrained routing and fallback.
- KV-cache compression primarily addresses decode-time memory or bandwidth and does not itself remove quadratic prefill attention.
- Context-window extension changes positional treatment or distributes exact processing, whereas RouteSparse concerns execution cost for an already long-context model.
- Long-context benchmarks motivate shift tests and category-level reporting, while RouteSparse additionally measures approximation error and system latency.
7 Conclusion
RouteSparse selects structured sparse patterns per input rather than fixing one pattern per head, combining latency-constrained routing with dense computation for uncertain cases.
- ROUTESPARSE achieves quality close to dense attention at 6.5× speed on 128K-token prefilling.
- Its robustness gains are larger under domain and length shift than with fixed routing.
- Ablations show measurable contributions from hardware profiling, input-conditional routing, and selective fallback.
Limitations
The paper’s guarantees and evaluation are bounded by probe failures, shift calibration limits, hardware dependence, implementation costs, and narrow experimental scale.
- The sampled probe can miss rare but essential query–key interactions.
- Empirical quantile calibration does not provide a distribution-free guarantee under arbitrary shift, and output-error bounds do not tightly bound final generation quality.
- Performance depends on hardware and implementation, while pattern heterogeneity can reduce batching efficiency and discrete kernels increase maintenance cost.
- The primary experiments use one model size and one GPU generation, limiting broader scaling conclusions.