You Only Index Once: Cross-Layer Sparse Attention with Shared Routing
About
Long-context inference in modern LLMs is increasingly constrained by decoding efficiency, especially in reasoning-heavy settings where models generate long intermediate chains of thought. Existing sparse attention methods often face a practical efficiency-quality trade-off. Structured block sparse methods typically provide stronger acceleration but incur noticeable quality loss, while token sparse methods are usually more accurate yet deliver limited end-to-end speedup because top-k routing over the full cache remains expensive. In this work, we propose cross-layer sparse attention (CLSA), which is built on top of KV-sharing architectures such as YOCO. The core idea is to share not only the KV cache across cross-decoder layers, but also the routing index. A single indexer computes token-level top-k selection once and reuses the resulting index across layers, thereby preserving the fine-grained selectivity of token sparse attention while amortizing the routing overhead. The resulting architecture improves all major inference bottlenecks jointly, including pre-filling, KV-cache storage, and long-context decoding. Experiments across short-context and long-context benchmarks show that CLSA is both accurate and efficient, achieving up to 7.6x decoding speedup and 17.1x overall throughput improvement at 128K context. These results suggest a more complete architectural solution for long-context LLMs that jointly advances model quality and inference efficiency.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Reasoning | BBH | Accuracy41.8 | 770 | |
| Commonsense Question Answering | WinoGrande | Accuracy63 | 89 | |
| Grade School Math Word Problems | GSM8K | Accuracy0.47 | 76 | |
| Reasoning | ARC Challenge | Accuracy (ARC)0.465 | 56 | |
| Reasoning | DROP | -- | 42 | |
| Commonsense QA | HellaSwag | Accuracy67.6 | 19 | |
| Decode Throughput Generation | NVIDIA B200 GPU | Throughput (tokens/s)6.74e+3 | 15 | |
| End-to-end Text Generation | NVIDIA B200 GPU Overall Generation | Throughput (tokens/s)4.92e+3 | 15 | |
| Prefill Throughput | Synthetic | Prefill Throughput (tokens/s)2.09e+4 | 15 | |
| Long-context retrieval and reasoning | RULER 16k context | S1 Accuracy100 | 10 |