RouteRelay cuts routing costs in transformer sparse attention models
RouteRelay: Event-Triggered Cross-Layer Route Reuse for Efficient Dynamic Sparse Attention
Computation and Language
Summary
Transformers use attention to figure out which parts of input text to focus on, but this can be expensive with long inputs. Dynamic sparse attention saves work by only looking at a few key pieces for each query, but it still recomputes routing scores at every layer, which can be wasteful. The authors designed RouteRelay, which reuses routing decisions across layers whenever possible, saving computation by only recomputing when necessary. Tests show RouteRelay keeps nearly all important routes while reducing routing effort significantly.
What this means in practice
- •For machine learning engineers: Reduce computational cost when scaling transformer models to long contexts by reusing routing data across layers in sparse attention.
- •For gpu software developers: Implement efficient selective rerouting routines and memory management in GPU kernels to speed up sparse attention computation in transformers.
Authors
Bin Li, Sisi Liu, Chenyang Hu, Chaoyang Zhang, Wei Li, Hui Song
Abstract
Dynamic sparse attention reduces long-context prefill cost by routing each query chunk to a small set of key chunks at every Transformer layer. The sparse attention kernel avoids most token interactions, but the router still rebuilds a chunk--chunk score matrix layer after layer, even when the selected routes change little. We introduce RouteRelay, a router-agnostic method that reuses only route metadata across depth while continuing to compute attention with the current layer's queries, keys, and values. Anchor layers perform full routing. Intermediate layers rescore the previous top-$k$ route and a compact sentinel set of near-miss and randomly probed chunks. A query row is rerouted only when a sentinel challenges its weakest selected chunk. We give a top-$k$ stability condition, a probabilistic bound on missed challengers, and a row-selective GPU execution design. In a reproducible empirical evaluation, RouteRelay retains at least 99.99% route recall while rerouting 25.0%, 55.4%, and 78.2% of rows under low, moderate, and high cross-layer drift, respectively. Across routing scales, RouteRelay retains 100.0% recall while evaluating 38.4--51.6% of full-routing score pairs as the key-chunk count grows from 128 to 1024. Its unfused CPU execution remains slower than dense matrix multiplication, exposing row compaction and ledger updates as the main kernel-engineering targets.