Cross-Layer Attention: How Sharing Key-Value Tensors Across Transformer Layers Slashes KV Cache Memory

In autoregressive large language models, the primary operational ceiling for serving long sequences and high batch concurrency is the Key-Value (KV) cache. During standard generation, every transformer layer computes and stores key and value activations for every token in the sequence to prevent quadratic recomputation during subsequent autoregressive decoding steps.

While this mechanism transforms inference time complexity from $O(N^2)$ to $O(N)$ per generated token, it introduces a massive memory footprint that scales linearly with sequence length, batch size, and network depth. For modern foundation models operating on 32k to 128k context windows, the KV cache rapidly exceeds the physical memory capacity of high-bandwidth GPU memory (HBM), throttling concurrent throughput and driving up inference serving costs.

To mitigate this memory wall, prior architectural innovations focused on compressing the head and channel dimensions:

• Multi-Query Attention (MQA) collapses all key and value heads into a single shared head per layer.
• Grouped-Query Attention (GQA) partitions query heads into $G$ groups that share key-value projections.
• Multi-Head Latent Attention (MLA) compresses key-value representations into low-rank latent vectors.

However, all of these approaches leave the layer dimension entirely uncompressed: every layer in an $L$-layer model maintains an independent, dedicated KV cache.

Cross-Layer Attention (CLA), introduced by Brandon et al. (NeurIPS 2024), addresses this overlooked dimension. By sharing key-value activations across adjacent transformer layers, CLA cuts the number of unique layers in the KV cache by half or more, advancing the accuracy-memory Pareto frontier while operating seamlessly within existing inference kernels.

The Memory Scaling Problem in Standard Transformer Decoders

In a standard transformer decoder consisting of $L$ layers, each attention block projects the input hidden state $h_{l-1} \in \mathbb{R}^{S \times d_{\text{model}}}$ (where $S$ is the sequence length and $d_{\text{model}}$ is the hidden dimension) into query, key, and value representations:

Q_l = h_{l-1} W_Q^l,    where W_Q^l in R^(d_model x (H_q * d_k))
K_l = h_{l-1} W_K^l,    where W_K^l in R^(d_model x (H_kv * d_k))
V_l = h_{l-1} W_V^l,    where W_V^l in R^(d_model x (H_kv * d_v))

Here, $H_q$ is the number of query heads, $H_{kv}$ is the number of key-value heads, and $d_k, d_v$ represent the per-head projection dimensions.

During token generation, the model appends the newly computed key and value vectors to the layer's historical KV cache. The total memory consumed by the KV cache across a batch of size $B$ and sequence length $S$ is expressed as:

Memory_KV = 2 * B * S * L * H_kv * d_k * bytes_per_element

Consider a 70-billion parameter model with $L = 80$ layers, $H_{kv} = 8$ heads, and $d_k = 128$ operating in 16-bit precision (2 bytes per parameter). At a sequence length of $S = 64\text{k}$ tokens, a single batch entry ($B = 1$) requires:

Memory_KV = 2 * 1 * 65,536 * 80 * 8 * 128 * 2 bytes ≈ 21.47 GB

A batch size of just 4 concurrent requests consumes 85.9 GB of VRAM solely for key-value storage, entirely consuming the capacity of an 80 GB NVIDIA H100 GPU before accounting for model parameters or runtime activations.

Standard Transformer Attention (Layer-Wise Independent KV Cache):
Layer 1:  h_0 -> [W_Q^1, W_K^1, W_V^1] -> Attention(Q_1, K_1, V_1) -> KV Cache Layer 1
Layer 2:  h_1 -> [W_Q^2, W_K^2, W_V^2] -> Attention(Q_2, K_2, V_2) -> KV Cache Layer 2
Layer 3:  h_2 -> [W_Q^3, W_K^3, W_V^3] -> Attention(Q_3, K_3, V_3) -> KV Cache Layer 3
Layer 4:  h_3 -> [W_Q^4, W_K^4, W_V^4] -> Attention(Q_4, K_4, V_4) -> KV Cache Layer 4
Total Cached Layers: L = 4

Cross-Layer Attention (CLA-2 with Sharing Factor s = 2):
Layer 1:  h_0 -> [W_Q^1, W_K^1, W_V^1] -> Attention(Q_1, K_1, V_1) -> Shared KV Cache Group 1
Layer 2:  h_1 -> [W_Q^2]               -> Attention(Q_2, K_1, V_1) -> (Reuses Group 1)
Layer 3:  h_2 -> [W_Q^3, W_K^3, W_V^3] -> Attention(Q_3, K_3, V_3) -> Shared KV Cache Group 2
Layer 4:  h_3 -> [W_Q^4]               -> Attention(Q_4, K_3, V_3) -> (Reuses Group 2)
Total Cached Layers: L / 2 = 2

Mathematical Mechanics of Cross-Layer Attention

Cross-Layer Attention partitions the $L$ layers of a transformer network into contiguous sharing groups parameterized by a sharing factor $s \in \mathbb{Z}^+$ (where $s \ge 2$).

Let $g \in \{0, 1, \dots, \lceil L/s \rceil - 1\}$ denote the group index. The layers assigned to group $g$ are defined by the index set:

G_g = { g * s + 1, g * s + 2, ..., min((g + 1) * s, L) }

Within each group $G_g$, the first layer $l_{\text{anchor}} = g \cdot s + 1$ acts as the KV Producer (or anchor layer), while subsequent layers $l \in G_g \setminus \{l_{\text{anchor}}\}$ act as KV Consumers.

1. KV Producer Layer (Anchor Step)

The anchor layer computes queries, keys, and values from its input hidden representation $h_{l_{\text{anchor}}-1}$:

Q_{l_anchor} = h_{l_anchor - 1} W_Q^{l_anchor}
K_{l_anchor} = h_{l_anchor - 1} W_K^g
V_{l_anchor} = h_{l_anchor - 1} W_V^g

The resulting key and value tensors $K_{l_{\text{anchor}}}$ and $V_{l_{\text{anchor}}}$ are stored in the group's shared KV cache buffer. Scaled dot-product attention proceeds normally:

Attention_{l_anchor} = Softmax((Q_{l_anchor} K_{l_anchor}^T) / sqrt(d_k)) V_{l_anchor}

2. KV Consumer Layers (Sharing Step)

For any subsequent layer $l$ within group $g$ ($l > l_{\text{anchor}}$), the network does not instantiate or compute key and value projection matrices. Instead, layer $l$ projects only its unique query tensor from its current hidden state $h_{l-1}$:

Q_l = h_{l-1} W_Q^l

The attention computation at layer $l$ directly retrieves the cached keys and values generated by the preceding anchor layer:

Attention_l = Softmax((Q_l K_{l_anchor}^T) / sqrt(d_k)) V_{l_anchor}

Because only the anchor layers write to the cache, the total number of cached layers shrinks from $L$ to $\lceil L/s \rceil$. The total KV cache memory footprint becomes:

Memory_{CLA} = 2 * B * S * ceil(L / s) * H_kv * d_k * bytes_per_element

For CLA-2 ($s = 2$), the KV cache size is reduced by exactly 50%. For CLA-3 ($s = 3$), it is reduced by 66.7%.

Why Cross-Layer KV Sharing Works: Representational Dynamics

The core intuition behind Cross-Layer Attention lies in the functional asymmetry between queries versus keys and values in deep autoregressive architectures.

Query Specialization vs. Key-Value Inertia

• Queries ($Q$): Queries represent dynamic, step-specific retrieval instructions. As a representation traverses through feedforward sub-layers and residual connections, the query vectors must adapt dynamically to extract increasingly abstract contextual relationships. Restricting queries across layers cripples model expressivity.
• Keys and Values ($K, V$): Keys and values act as associative memory indices and contextual payload containers. In deep transformer networks, the representational geometry of token states changes gradually across consecutive layers. Probing studies show that adjacent layers exhibit high cosine similarity and low subspace drift in their key and value representations.

By decoupling queries from key-value generation, CLA preserves full query routing flexibility at every layer while eliminating redundant key-value state allocations across adjacent depth steps. The feedforward networks (MLPs) and residual streams between layers $l_{\text{anchor}}$ and $l_{\text{consumer}}$ continue to transform token hidden states, allowing the model to construct new query trajectories against the established key-value coordinate space.

Orthogonality Across KV Compression Paradigms

A critical property of Cross-Layer Attention is that it operates along an architectural dimension completely orthogonal to existing KV compression techniques.

| Compression Axis | Mechanism | Target Dimension | Representative Methods | | :--- | :--- | :--- | :--- | | Attention Heads | Share KV projections across query heads | $H_{kv}$ | MQA, GQA | | Channel / Subspace | Low-rank projection / Latent compression | $d_k, d_v$ | MLA (DeepSeek-V2/V3) | | Numerical Precision | Post-training quantization / Outlier rotation | $\text{bytes}$ | FP8, INT4 KV, QuaRot | | Sequence Length | Context pruning / Attention sinks / Sparsity | $S$ | StreamingLLM, H2O, SnapKV | | Layer Depth | Inter-layer projection sharing | $L$ | Cross-Layer Attention (CLA) |

Because CLA targets the depth dimension $L$, it combines multiplicatively with head-level and channel-level compression:

• GQA + CLA-2: An 8-head GQA model with 32 layers using CLA-2 maintains only 16 layers of 8-head KV caches, achieving double the throughput of standard GQA.
• MQA + CLA-2: Pairing single-head Multi-Query Attention with 2-layer sharing achieves an aggregate $2 \times H_q$ memory reduction over standard Multi-Head Attention (MHA).
• Quantized KV + CLA-2: An FP8-quantized KV cache deployed on a CLA-2 architecture yields a $4\times$ aggregate reduction in memory traffic compared to standard FP16 execution.

Empirical Validation and the Pareto Frontier

In their NeurIPS 2024 paper, Brandon et al. conducted comprehensive pre-training sweeps on 1B- and 3B-parameter transformer models across varied architectural configurations, head dimensions ($d_{\text{head}} \in \{32, 46, 64, 90, 128\}$), and sharing factors ($s \in \{2, 3, 4\}$).

Perplexity and Memory Trade-Offs

The empirical findings establish three core conclusions:

1. Near-Zero Degradation at $s=2$: In 1B-parameter pre-training runs, a model configured with $d_{\text{head}} = 128$, MQA, and CLA-2 achieved a $2\times$ reduction in KV cache footprint compared to an unmodified MQA baseline ($d_{\text{head}} = 128$), suffering an almost imperceptible validation perplexity degradation of just 0.04 points.
2. Advancing the Pareto Frontier: When comparing models at equal KV cache memory budgets, CLA consistently outperforms plain MQA with reduced head dimensions. For example, an MQA-CLA2 model with $d_{\text{head}} = 128$ matched the exact memory footprint of a baseline MQA model with $d_{\text{head}} = 64$, while achieving 0.21 to 0.48 points lower perplexity across standard language modeling benchmarks.
3. Diminishing Returns Beyond $s=2$: While CLA-3 and CLA-4 achieved Pareto improvements over narrow-head baselines, they exhibited steeper perplexity penalties than CLA-2 at matched memory footprints. Sharing across pairs of consecutive layers ($s=2$) represents the optimal balance between memory compression and representational capacity.

Systems and Hardware Implications for Production Serving

The architectural changes introduced by Cross-Layer Attention produce direct efficiency gains across GPU memory capacity, memory bandwidth utilization, and computational overhead.

1. Memory-Bandwidth Bound Decoding Throughput

During autoregressive token generation, language models operate in a strictly memory-bandwidth-bound regime. Each newly generated token requires loading the entire model parameter set and the full historical KV cache from high-bandwidth memory (HBM) into GPU SRAM to perform a single vector-matrix multiply per head.

By cutting the KV cache volume loaded per token step by $2\times$, CLA-2 halves the attention memory bus traffic. This directly reduces inter-token latency (ITL) and accelerates decoding throughput on memory-constrained hardware clusters.

2. Doubling PagedAttention Concurrency

Modern inference runtimes such as vLLM (PagedAttention), SGLang, and TensorRT-LLM manage KV caches using virtual memory paging tables.

With CLA-2, the engine allocates physical memory pages for only $L/2$ layers. For a fixed pool of GPU VRAM allocated to the KV cache, the serving runtime can host $2\times$ as many active sequences or double the maximum serviceable context length per node without spilling to host RAM or triggering out-of-memory (OOM) faults.

3. Reduced Prefill Compute and Parameter Count

In addition to inference savings, CLA reduces model parameter count and training FLOPs:

• Non-anchor layers eliminate $W_K$ and $W_V$ projection weights entirely, removing $(s - 1)/s \times 2 \times d_{\text{model}} \times (H_{kv} \cdot d_k)$ parameters from the network.
• During prompt prefill, the engine skips key-value matrix multiplications on all consumer layers, yielding minor FLOP reductions during prompt processing.

4. Zero Kernel Overhead

Unlike low-rank compression methods such as MLA, which require specialized matrix absorption transformations or custom RoPE-decoupled attention kernels to avoid decompressing latent vectors into SRAM, Cross-Layer Attention requires no custom GPU kernels.

Inference runtimes execute standard FlashAttention or FlashDecoding kernels. The serving engine simply passes the anchor layer's KV cache pointers to the consumer layer's kernel invocation, making CLA drop-in compatible with standard serving infrastructure.

Practical Considerations and Future Directions

While Cross-Layer Attention provides substantial memory efficiency benefits, adopting it requires training models with shared projections from scratch or applying specialized continued pre-training recipes:

• Pre-Training Requirement: Because key-value projections are tied across layers during forward passes, CLA cannot be applied zero-shot as a post-training compression technique to existing foundation checkpoints without distillation or weight adaptation.
• Layer Grouping Topology: While uniform adjacent grouping ($1-2, 3-4, \dots$) is standard, emerging research explores heterogeneous sharing schemes, such as maintaining independent KV projections in early layers where representations evolve rapidly, while applying wider sharing factors ($s=3$ or $s=4$) in deep, representationally static layers.

As context windows expand toward multi-million-token horizons and agentic architectures demand massive parallel rollouts, Cross-Layer Attention provides a mathematically sound, hardware-efficient mechanism for overcoming the transformer memory wall.

Sources

• Brandon, W., Mishra, M., Nrusimha, A., Panda, R., & Ragan-Kelley, J. (2024). Reducing Transformer Key-Value Cache Size with Cross-Layer Attention. Advances in Neural Information Processing Systems (NeurIPS 2024). https://arxiv.org/abs/2405.12981
• Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv preprint arXiv:1911.02150. https://arxiv.org/abs/1911.02150
• Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP 2023). https://arxiv.org/abs/2305.13245
• DeepSeek-AI. (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv preprint arXiv:2405.04434. https://arxiv.org/abs/2405.04434
• Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., & Stoica, I. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. Proceedings of the 29th ACM Symposium on Operating Systems Principles (SOSP 2023). https://arxiv.org/abs/2309.06180
• Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems (NeurIPS 2022). https://arxiv.org/abs/2205.14135