# You Only Index Once: Cross-Layer Sparse Attention with Shared Routing

> Cross-Layer Sparse Attention shares a single top-k routing index across decoder layers, delivering up to 7.6× decoding speedup at 128K context with near-lossless quality.

- **Source:** [arXiv](https://arxiv.org/abs/2606.06467)
- **Published:** 2026-08-17
- **Permalink:** https://picx.dev/p/MmJeX9
- **Whiteboard:** https://picx.dev/p/MmJeX9/image

## Summary

## Summary

- **Core Contribution**: Proposes **Cross-Layer Sparse Attention (CLSA)**, a novel architecture that extends KV-sharing (from YOCO) to also share the *routing index* across decoder layers, amortizing the expensive top-k token selection cost.
- **Key Innovation**: A single query-aware indexer computes token-level top-k routing once, and all cross-decoder layers reuse this index, preserving the fine-grained selectivity of token-sparse attention while eliminating per-layer routing overhead.
- **Efficiency Gains**: CLSA achieves up to **7.6× decoding speedup** and **17.1× overall throughput improvement** at 128K context compared to a dense Transformer baseline.
- **Quality Preservation**: CLSA remains nearly lossless relative to dense baselines across short-context (MMLU, GSM8K, etc.) and long-context (RULER, Books, ArXiv) benchmarks, and even improves on several reasoning tasks.
- **Unified Bottleneck Solution**: Unlike methods that optimize only one aspect, CLSA jointly improves pre-filling, KV-cache storage, and decoding efficiency.

## Introduction and Theoretical Foundation

Long-context inference in modern LLMs is increasingly **decoding-bound**, especially in reasoning-heavy settings like chain-of-thought generation. As sequences grow, three major bottlenecks emerge: (1) expensive pre-filling, (2) growing KV-cache memory, and (3) slow decoding.

Existing sparse attention methods face a fundamental **efficiency-quality trade-off**:
- **Block-sparse methods** (e.g., NSA, MoBA) provide strong wall-clock speedups due to GPU-friendly structured sparsity but suffer noticeable quality loss.
- **Token-sparse methods** (e.g., DSA) are more accurate but deliver limited end-to-end speedup because top-k routing over the full cache is irregular, expensive, and recomputed per layer.

The theoretical foundation builds on **YOCO (You Only Cache Once)**, which decomposes the model into:
- A **self-decoder** that encodes input into shared hidden states and constructs a single KV cache.
- A **cross-decoder** where layers read from this shared cache via dense cross-attention.

CLSA's central insight: *when multiple layers read from the same memory, the routing decision should also be tied to that memory and shared*. This extends the sharing principle from memory to routing.

## Methodology

### Architecture Overview

The CLSA architecture builds on YOCO with two key modifications:

**1. Shared Indexer**: A lightweight single-head indexing branch computes index queries and keys from shared hidden states:
$$Q_{\mathrm{idx}} = H W_{\mathrm{idx}}^{Q}, \quad K_{\mathrm{idx}} = H W_{\mathrm{idx}}^{K}$$

Routing indices are computed once:
$$I = Q_{\mathrm{idx}} K_{\mathrm{idx}}^{\top}, \quad S_t = \mathrm{TopK}(I_t, k)$$

Each cross-decoder layer then attends only to the selected tokens:
$$O_t^{(l)} = \mathrm{Attn}(Q_t^{(l)}, K_{S_t}, V_{S_t})$$

**2. Multi-Layer Distillation**: To ensure the shared index serves all layers well, a distillation objective aggregates attention distributions across all layers and heads:

$$\bar{A} = \frac{1}{LH} \sum_{l=1}^{L} \sum_{h=1}^{H} \mathrm{softmax}\left(Q^{(l,h)} K^{(h)^{\top}}\right)$$

$$\mathcal{L}_{\mathrm{KD}} = \frac{1}{n} \sum_{t=1}^{n} \mathrm{KL}\left(\mathrm{sg}\left[\bar{A}_t\right] \| \mathrm{softmax}(I_t)\right)$$

**Training Stages**:
- **Stage 1**: Indexer warmup with frozen backbone ($\mathcal{L}_{\mathrm{stage1}} = \mathcal{L}_{\mathrm{KD}}$)
- **Stage 2**: Joint sparse adaptation ($\mathcal{L}_{\mathrm{stage2}} = \mathcal{L}_{\mathrm{LM}} + \lambda \mathcal{L}_{\mathrm{KD}}$, with $\lambda = 0.1$)

### Complexity Comparison

| Model | KV Cache Memory | Prefilling Time | Decoding Time |
|-------|----------------|-----------------|---------------|
| Transformer | $\mathcal{O}(LND)$ | $\mathcal{O}(LN^{2}D)$ | $\mathcal{O}(LND)$ |
| Hybrid TRM | $\mathcal{O}(L(\gamma N+(1-\gamma)W_{1})D)$ | $\mathcal{O}(L(\gamma N^{2}+(1-\gamma)W_{1}N)D)$ | $\mathcal{O}(L(\gamma N+(1-\gamma)W_{1})D)$ |
| YOCO (Dense) | $\mathcal{O}((N+W_{1}L)D)$ | $\mathcal{O}(\frac{L}{2}W_{1}ND)$ | $\mathcal{O}(\frac{L}{2}(N+W_{1})D)$ |
| DSA | $\mathcal{O}(LND)$ | $\mathcal{O}(LW_{2}ND+\eta LN^{2})$ | $\mathcal{O}(LW_{2}D+\eta LN)$ |
| **YOCO (CLSA)** | $\mathcal{O}((N+W_{1}L)D)$ | $\mathcal{O}(\frac{L}{2}W_{1}ND)$ | $\mathcal{O}(\frac{L}{2}(W_{1}+W_{2})D+\eta N)$ |

Key advantage: The indexer cost ($\eta N$) is paid **only once** instead of $\eta LN$ in DSA.

## Empirical Validation / Results

### General Benchmarks (4B models)

| Model | ARC-C | BBH | GSM8K | HellaSwag | HumanEval | MMLU | DROP | WinoGrande |
|-------|-------|-----|-------|-----------|-----------|------|------|------------|
| Transformer | 0.453 | 0.420 | 0.434 | 0.667 | 0.384 | 0.527 | 0.366 | 0.638 |
| YOCO (Dense) | 0.461 | 0.411 | 0.430 | 0.676 | 0.396 | 0.519 | 0.387 | 0.630 |
| **YOCO (CLSA)** | **0.465** | 0.418 | **0.470** | 0.674 | **0.396** | 0.513 | **0.391** | 0.616 |

CLSA achieves the **best scores on ARC-Challenge, GSM8K, and DROP**, and matches the best HumanEval, while remaining close to dense baselines on other tasks.

### Long-Context (RULER at 32K)

CLSA achieves the **best average score (53.1)** at 32K context, outperforming both Transformer (46.2) and dense YOCO (52.3), with gains mainly from harder multi-needle settings (MK1, MK2).

### Inference Efficiency

- **At 128K context**: ~7.6× decode throughput and ~17.1× overall throughput vs. Transformer
- **Latency analysis**: Unamortized top-k at 128K can cost as much as dense attention; CLSA's amortized top-k takes only ~0.08 ms per layer
- **Comparison with other sparse methods**: DSA is *slower* than dense Transformer at 128K due to unamortized routing; IndexCache and HySparse reduce overhead but CLSA achieves the lowest per-layer latency

### Attention Sparsity Analysis

| Domain | Attn. Coverage at 2048 tokens | CE Loss at 2048 | Dense CE Loss | Δ |
|--------|------------------------------|-----------------|---------------|-----|
| StarCoder | 84.12% | 0.5699 | 0.5703 | -0.0004 |
| Books | 76.29% | 1.7500 | 1.7446 | +0.0054 |
| ArXiv | 80.67% | 1.0844 | 1.0818 | +0.0026 |

At a 1:16 activation ratio (2048 tokens), sparse attention captures ~80% of dense attention mass while introducing **negligible loss degradation** (≤0.006).

## Theoretical and Practical Implications

1. **Routing-Memory Binding Principle**: The paper establishes that when layers share memory, they should also share routing decisions. This is validated empirically: cross-layer attention scores are similar enough that one shared index serves all layers effectively.

2. **Rethinking Sparse Attention Efficiency**: The work demonstrates that theoretical FLOP reduction does not automatically translate to wall-clock speedup. The irregular top-k operation is poorly matched to GPU Tensor Cores, making routing amortization essential for practical gains.

3. **Unified Architectural Solution**: CLSA shows that addressing all three bottlenecks (pre-fill, KV-cache, decode) jointly is more effective than optimizing one in isolation. This suggests a more complete design philosophy for long-context LLMs.

4. **Token vs. Block Sparsity**: The analysis clarifies that token-level selection is superior for long-context quality because block sparsity imposes a structure misaligned with semantic saliency (nearby tokens can have very different importance). CLSA makes token sparsity practical by amortizing its main cost.

## Conclusion

CLSA extends KV-sharing to routing-sharing, creating an architecture that jointly optimizes pre-filling, KV-cache storage, and decoding efficiency. The method preserves model quality while delivering substantial speedups (up to 7.6× decode, 17.1× overall at 128K). 

Future directions suggested by this work include:
- Applying CLSA to even longer contexts and larger model scales
- Exploring adaptive selection of active token budgets per domain or task
- Investigating whether the shared routing principle extends to other architectural components beyond attention (e.g., FFN sparsity)
- Combining CLSA with other efficiency techniques like quantization or speculative decoding for further gains

---

_Markdown view of https://picx.dev/p/MmJeX9, served by PicX — AI-generated visual whiteboard summaries of research papers._