You Only Index Once: Sharing Routing Across Layers for Efficient Long-Context LLMs

Paper: 2606.06467 Authors: Yutao Sun, Yanqi Zhang, Li Dong, Jianyong Wang, Furu Wei Categories: cs.CL, cs.AI, cs.LG

The Gap

Long-context LLMs face a hard efficiency-quality trade-off in sparse attention. Block sparse methods (like sliding window) give strong speedups but lose important tokens. Token-level sparse methods (top-k selection) maintain quality but remain slow because every layer must scan the full KV cache to pick its top-k tokens. The routing operation itself becomes the bottleneck — you’re repeatedly sorting through thousands of tokens just to decide which ones to attend to.

Prior work treats this as unavoidable: finer granularity means more routing cost. YOCO-style architectures share KV cache across layers but still route independently per layer. The gap: routing overhead scales linearly with layer count even when the important tokens don’t change much layer-to-layer.

Problem: Token sparse accurate but slow (routing overhead × L layers)
         |
         v
Assumption: Important tokens stable across adjacent layers
         |
         v
Method: Compute top-k index once, reuse across all layers (CLSA)
         |
         v
Evidence: 7.6x decoding speedup, <2% quality drop at 128K context
         |
         v
Conclusion: Routing overhead amortizable without sacrificing selectivity

The Increment

One sentence: Before this paper, token-level sparse attention meant paying full routing cost per layer; after, you route once and reuse the index across layers with minimal quality loss.

Core Mechanism

CLSA builds on YOCO (You Only Cache Once), which already shares KV cache across decoder layers. The insight: if layers share the same KV cache, why not share the routing decision too? CLSA adds a single indexer module that computes top-k token selection once per forward pass. This indexer outputs a binary mask indicating which tokens to keep. All subsequent decoder layers reuse this mask — no re-sorting, no re-ranking, just direct attention to the pre-selected tokens.

The indexer itself is a small transformer block trained jointly with the main model. During prefill, it sees the full sequence and learns which tokens matter for the task. During decoding, it updates the mask incrementally as new tokens arrive. The key architectural choice: the indexer’s parameters are tiny compared to decoder layers (think 1-5% of total parameters), so the overhead of one extra forward pass is negligible compared to L repeated routing operations.

Data flows like this: input → indexer → sparse mask → decoder layer 1 (uses mask) → decoder layer 2 (reuses same mask) → … → decoder layer L (still same mask) → output. The mask doesn’t change between layers. Only when a new token is generated does the indexer recompute.

Standard token sparse (per-layer routing):
    Input ---+---> Layer 1: route (expensive) + attend
             |
             +---> Layer 2: route (expensive) + attend
             |
             +---> Layer L: route (expensive) + attend

CLSA (route once, reuse):
    Input ---> Indexer: route once (cheap module)
                  |
                  v
               [mask]
                  |
         +--------+--------+
         |        |        |
         v        v        v
    Layer 1  Layer 2  Layer L
    attend   attend   attend
    (reuse)  (reuse)  (reuse)

Think of CLSA as a train station with security screening. In the standard approach (per-layer routing), every train car has its own security checkpoint — passengers get rescanned at each car even though they already passed screening. CLSA consolidates security at the station entrance (indexer). One comprehensive scan determines who boards, and all train cars (decoder layers) trust that decision. The security team (indexer) is small but thorough. The cars (decoders) just check tickets (mask) instead of rescanning luggage. Passengers (tokens) only prove their importance once, then move freely through the train.

Key Concepts

• Routing overhead: In token-level sparse attention, “routing” means scanning all cached tokens to rank them by relevance, then selecting top-k. This involves computing attention scores (or approximations) over the full cache, sorting, and indexing. At 128K context with 32 layers, you’re doing this expensive scan 32 times per generated token. CLSA’s core bet: most layers would pick similar tokens anyway, so compute it once with a dedicated module and share the result. The indexer learns during training which tokens generalize across layers.

• Cross-layer KV sharing: YOCO architecture caches key-value pairs once and shares them across all decoder layers, reducing memory from O(L × N) to O(N) where L is layers and N is sequence length. CLSA extends this: if the cache is shared, the access pattern (which tokens to attend to) can be shared too. This turns routing from a per-layer operation into a per-forward-pass operation.

• Index stability assumption: The method assumes that the set of important tokens doesn’t vary wildly between adjacent layers. Empirically, this holds — early layers focus on syntax/structure, later layers on semantics, but the overlap is high. When a token is critical for layer 5, it’s usually critical for layer 6 too. The indexer learns to pick tokens that are broadly useful rather than layer-specific. The <2% quality drop in experiments validates this assumption.

Framework Shift

Before (standard token sparse):          After (CLSA):

Input (128K tokens)                      Input (128K tokens)
         |                                        |
         v                                        v
Layer 1: scan 128K, sort, pick k         Indexer: scan 128K, pick k
         |                                        |
         v                                        v
Layer 2: scan 128K, sort, pick k              [mask]
         |                                   /    |    \
         v                                  /     |     \
       ...                          Layer 1   Layer 2   Layer L
         |                          (use k)   (use k)   (use k)
         v                                   \     |     /
Layer L: scan 128K, sort, pick k             \    |    /
                                                  v
Cost: L × routing overhead              Cost: 1 × routing overhead

From per-layer independence to centralized routing, the core shift is amortizing the expensive decision across layers by assuming stability in what matters.

Expert Assessment

Problem choice: Real and well-positioned. Long-context inference is the current frontier, and the routing overhead problem is measurable and annoying in production. The gap between block sparse (fast but lossy) and token sparse (accurate but slow) has been complained about since at least 2023. This paper doesn’t invent a new problem — it solves an acknowledged one.

Method maturity: The insight is clean and the execution is straightforward. The indexer design is minimal — no exotic architectures, just a small transformer block with a learned top-k selector. The risk was that layer-specific token needs would break the sharing assumption, but the experiments show it holds. One could argue for even simpler baselines (e.g., static heuristics like recency + attention entropy), but the learned indexer probably captures task-specific patterns better. The method feels like the obvious next step once YOCO exists, which is a compliment — good ideas often look obvious in hindsight.

Experimental integrity: Strong. Baselines include both block sparse (FlashAttention-2 + sliding window) and token sparse (Quest, MInference). The 7.6x decoding speedup is measured at 128K context where the bottleneck is real. Quality is evaluated on diverse benchmarks (RULER, Needle-in-Haystack, reasoning tasks). The quality drop is small (<2%) and honestly reported. One minor concern: the indexer adds parameters and a forward pass, but the cost breakdown shows it’s <5% overhead, which is fair. The ablations are thorough — they test different indexer sizes, sparsity ratios, and context lengths.

Writing quality: The paper is clear and well-structured. The related work section properly credits YOCO and positions this as an extension. The method description could be tightened — Section 3.2 spends too many words on obvious points about KV sharing. The experimental section is dense but readable. If I were revising, I’d expand the failure case analysis (Section 4.4) — when does index sharing hurt? The paper hints at “heterogeneous layer preferences” but doesn’t drill down. A concrete example of a task where per-layer routing wins would strengthen the honest assessment.

Verdict: strong accept — Solves a real problem with a simple, effective method and thorough experiments. The speedup is substantial and the quality trade-off is acceptable for most applications.

Takeaways

• Amortize expensive decisions: If multiple components need similar information, compute it once and share. This applies beyond attention — think kernel fusion, cached intermediate representations, or shared preprocessing pipelines. The CLSA pattern: identify a repeated operation, check if the inputs/outputs are stable across repetitions, centralize if yes.

• Learned routing beats heuristics: Static rules (recency, position, keyword matching) are tempting for sparse selection, but a small learned module often outperforms them by adapting to task distribution. The indexer here is <5% overhead but captures patterns that sliding windows miss.

• Validate stability assumptions with ablations: CLSA bets on cross-layer index stability. The paper validates this with per-layer similarity metrics and quality comparisons. When proposing sharing schemes, measure whether the thing you’re sharing actually stays similar — don’t just assume.

• Speedup measurement matters: Report end-to-end latency, not just FLOPs or theoretical complexity. CLSA wins because routing overhead is wall-clock expensive. If your optimization targets a non-bottleneck, the speedup won’t show up in practice.