RouteSparse: Input-Conditional Pattern Routing for Budgeted Long-Context Prefilling

2026-08-29

AI Summary

论文提出RouteSparse方法，通过输入条件路由选择稀疏模式以解决长上下文预填充的高成本问题，提高了效率并减少了性能下降。

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.

Problem

• long-context prefilling
• dynamic sparse attention
• pattern routing

Innovation

• Dynamic Sparse Attention
• Input-Conditional Routing
• Latency-Aware Router
• Constrained Risk Minimization
• Attention-Output Error Certificate