RouteSparse: Input-Conditional Pattern Routing for Budgeted Long-Context Prefilling
Organizations: College of Computer Science, Chongqing University · School of Information Science and Engineering, Chongqing Jiaotong University
Abstract
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 dense prefill speed with a 0.2-point RULER drop relative to dense attention, compared with 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.
Figures & tables
| Input-conditional routing for one layer |
|---|
| 1. Compute the shared probe in Equation 3 . |
| 2. Build candidate indices and estimate coverage and uncertainty. |
| 3. Look up profiled latency for each pattern–budget pair. |
| 4. Solve Equation 8 under the layer latency target. |
| 5. Expand or fall back for candidates with . |
| 6. Group heads by kernel and execute sparse or dense attention. |
| Method | Prefill (s) | Speedup | RULER | InfiniteBench | PG-19 PPL |
|---|---|---|---|---|---|
| Dense FlashAttention-2 | 43.6 | 84.7 | 47.8 | 8.24 | |
| Static A-shape | 7.1 | 78.9 | 42.5 | 8.91 | |
| MInference, fixed | 6.0 | 83.1 | 46.5 | 8.39 | |
| MInference, larger budget | 7.4 | 84.0 | 47.0 | 8.32 | |
| RouteSparse | 6.7 | 84.5 | 47.6 | 8.27 |
| Variant | Prefill (s) | RULER | Worst-decile | Fallback (%) |
|---|---|---|---|---|
| RouteSparse full | 6.7 | 84.5 | 6.8 | |
| No input routing | 6.0 | 83.1 | 0.0 | |
| No uncertainty score | 6.3 | 83.8 | 2.1 | |
| No dense fallback | 5.9 | 83.9 | 0.0 | |
| One fixed budget | 6.1 | 84.0 | 4.3 | |
| Recent queries only | 6.4 | 84.1 | 5.5 |
Long-Context Language Model Inference
Long-context language model inference makes processing and generating over lengthy prompts faster and more memory-efficient.
See all 224 papers