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

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

Empirical Validation / Results

General Benchmarks (4B models)

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

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