RouteSparse: Input-Conditional Pattern Routing for Budgeted Long-Context Prefilling
Computation and Language
Summary
The authors propose RouteSparse, a method that improves speed and efficiency of attention in large language models by dynamically selecting sparse attention patterns for each part of the input. Instead of using a fixed pattern for each attention head, RouteSparse chooses from a small set of patterns based on a quick estimate of their usefulness for the current input, sometimes falling back to denser patterns to maintain quality. They tested this on a large Llama 3.1-8B-Instruct model and achieved much faster processing (6.5 times speedup) with minimal accuracy loss, better balancing speed and performance compared to fixed routing methods. Their approach considers hardware efficiency and input characteristics to optimize results.
Authors
Chao Zhang, Yifan Ji, Ziyan Zhang, Kai Song, Fei Lin
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 $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.