---
title: 'YOCO: Efficient KV Caching for LLMs'
url: https://www.emergentmind.com/topics/yoco-you-only-cache-once
type: topic
---

# YOCO: Efficient KV Caching for LLMs

YOCO (You Only Cache Once) is a decoder-decoder architecture for large language models that addresses the memory and computational bottlenecks associated with key-value (KV) caching in standard Transformer decoders. By restructuring the decoder into a two-stage pipeline—comprising a self-decoder for efficient context encoding, followed by a cross-decoder that reuses a single, global KV cache—YOCO achieves an order-of-magnitude reduction in memory footprint and a substantial improvement in prefill and generation throughput, without compromising language modeling performance or context length capabilities. Extensions such as YOCO++ and Universal YOCO (YOCO-U) further enhance this efficiency-quality tradeoff with residual KV composition and recursive depth scaling.

## 1. Architecture and Computational Flow

YOCO divides a standard $L$-layer decoder-only Transformer into two stacked modular components:
- **Self-Decoder**: The first $L/2$ layers, using an efficient self-attention scheme (e.g., sliding-window or gated retention), encode the full prompt context into a global key-value cache.
- **Cross-Decoder**: The remaining $L/2$ layers consume the tokens autoregressively, performing cross-attention into the global cache produced by the self-decoder.

The model operates as follows:
\[
X^{(0)} \in \R^{N\times d}, \quad
X^{(l)} = 
\begin{cases}
  \mathrm{SelfDecoder}(X^{(l-1)}) & l=1,\ldots,L/2\\
  \mathrm{CrossDecoder}(X^{(l-1)}, \hat K, \hat V) & l=L/2+1,\ldots,L
\end{cases}
\]
At the self-decoder–cross-decoder interface, a single set of keys and values is constructed:
\[
M = X^{(L/2)}, \qquad
\hat K = \mathrm{LN}(M) W_K,\qquad
\hat V = \mathrm{LN}(M) W_V
\]
For $l > L/2$, each cross-decoder layer computes
\[
Q^{(l)} = \mathrm{LN}(X^{(l-1)}) W_Q^{(l)}, \qquad
Y^{(l)} = \mathrm{Attention}(Q^{(l)}, \hat K, \hat V) + X^{(l-1)}
\]
\[
X^{(l)} = \mathrm{SwiGLU}(\mathrm{LN}(Y^{(l)})) + Y^{(l)}
\]
Causal masking is maintained throughout to preserve autoregressive dependencies.

The self-decoder employs either sliding-window or gated retention attention, both of which incur $\mathcal{O}(1)$ KV memory due to localized or recurrent state organization. The cross-decoder's cross-attention is always to the same precomputed $(\hat K, \hat V)$, and no new cache is materialized during generation.

## 2. Caching Mechanism and Memory Complexity

The conventional Transformer caches KV pairs for every layer and every input token: memory use is $\mathcal{O}(N \cdot L \cdot d)$ for prompt length $N$, depth $L$, and hidden dimension $d$. YOCO, by contrast, caches only:
- The single global cache at layer $L/2$: $\mathcal{O}(N \cdot d)$;
- Small local buffers from efficient self-attention: $\mathcal{O}(C \cdot L \cdot d)$ (window size $C \ll N$), often negligible.

Total KV cache requirement for YOCO is $
\mathcal{O}((N + C\,L) d) \approx \mathcal{O}(N d)
$, representing an $L$-fold memory reduction relative to standard practice [2405.05254, 2604.01220].

This constant-cache property is preserved during generation—new tokens append only to the global cache, obviating per-layer synchronization. During decoding, cross-attention retrieves from the fixed cache, with no need to recompute embeddings or store additional per-layer states.

## 3. Prefill Strategy and Inference Efficiency

YOCO’s bifurcated computation enables an early-exit “prefill” optimization:
- Classic Transformers perform $L$ layers of self-attention during prefill, with $\mathcal{O}(L N^2 d)$ time.
- YOCO limits prefill to just the $L/2$ self-decoder layers, applying only efficient attention per token: cost is $\mathcal{O}(L N d)$ — linear, rather than quadratic, in sequence length.

Empirical latency benchmarks (H100-80GB GPU, YOCO$_\mathrm{gRet}$) establish prefill speedups ranging from $\sim$2.8$\times$ at $N=32$K tokens (3.2s vs 9.1s) to $\sim$38$\times$ at $N=1$M (10s vs 380s). Post-prefill, generation switches to the cross-decoder, which behaves identically to a normal decoder, so there is no discrepancy in output or generation semantics [2405.05254].

## 4. Model Scaling, Long-Context Performance, and Benchmark Results

At increasing context lengths and model sizes, YOCO demonstrates:
- **Memory reduction**: At $N=1$M, a 3B-parameter model requires $\sim$12.4 GB (YOCO) versus $\sim$114 GB (Transformer), a factor of 9.4 reduction.
- **Throughput gains**: At $N=512$K, YOCO sustains 43.1 tokens/s, compared to 4.5 tokens/s for the baseline Transformer—a $9.6\times$ speedup.
- **Comparable or improved model quality**: Across parameter counts (160M–13B) and training budgets, YOCO matches or slightly exceeds Llama-style Transformer validation loss.

Needle-in-a-haystack (NIAH) retrieval tasks validate YOCO’s long-context capabilities:
- Single-needle accuracy $\geq 0.98$ across depths up to 1M tokens.
- Multi-needle retrieval at 128K input length: accuracy $(N=1,2,4,8) = (0.98, 0.98, 0.84, 0.56)$, matching or exceeding alternative LLMs.

For language modeling across 1M-token prompts, per-token negative log-likelihood steadily improves with context length, demonstrating actual exploitation of the extended history [2405.05254].

## 5. Extensions: YOCO++, Universal YOCO (YOCO-U), and Efficiency-Capacity Tradeoffs

### YOCO++

YOCO++ enhances the original YOCO by fusing each self-decoder layer's KV with that of the bottom layer using learnable weighted residuals:
\[
\tilde K^{(i)} = \lambda (\alpha_{i,1} K^{(1)} + \alpha_{i,2} K^{(i)}), \quad
\tilde V^{(i)} = \lambda (\alpha_{i,1} V^{(1)} + \alpha_{i,2} V^{(i)}), \quad i=2,\ldots,m
\]
with $\lambda > 1$, typically set to 35, and $\alpha_{i,*}$ learned end-to-end. This composition raises expressivity while preserving all the computational and memory benefits of YOCO.

Empirical evaluations (TinyLlama 1.1B, $L=22$, $m=11$):
- Inference throughput and prefill latency at various context lengths match those of YOCO ($\sim$50% faster than vanilla Transformer).
- YOCO++ attains the lowest training loss and highest average zero-shot accuracy (48.99%) compared to YOCO (47.98%) and FusedKV variants, and outperforms the standard Transformer (48.37%).
- Ablations show both the residual connection and the scaling factor $\lambda$ are necessary for best results [2604.13556].

### Universal YOCO (YOCO-U)

Universal YOCO [Editor’s term: "YOCO-U"] incorporates recursion in the self-decoder. Instead of stacking more layers, it iteratively applies the shallow ($L/2$-layer) self-decoder block $T$ times with shared parameters:
\[
\text{USD}(X) = \underbrace{\mathrm{SD} \circ \cdots \circ \mathrm{SD}}_{T\ \text{times}} (X)
\]
The global cache is constructed after recursion. This approach enhances effective depth and representational capacity without increasing cache size, enabling high efficiency even under test-time scaling.

Benchmarks indicate YOCO-U:
- Lowers validation loss over non-recursive YOCO at equal FLOPs and converges in fewer training tokens.
- Achieves higher downstream task performance (+4.45 to +5.30%\ absolute) and a 24.4% average improvement in math reasoning benchmarks.
- Matches RINS in generalization, at substantially lower cache memory (62MB for YOCO-U vs 1.28GB for RINS at 16K context).
- Maintains near-linear prefill cost and one-copy KV caching.

Ablation studies confirm greater returns from recursion in the shallow self-decoder than in deeper blocks or from simple width increases. Diminishing performance gains with increased recursion iterations suggest representational convergence after a small number of repeats [2604.01220].

## 6. Practical Considerations, Integration, and Trade-offs

To deploy YOCO or its variants:
- **Integration**: Implement or use an inference engine supporting the YOCO cache protocol. In the self-decoder, produce each layer’s K,V, optionally fusing with the bottom layer (YOCO++), and cache as prescribed.
- **Compression rate**: Most evaluations use $m=L/2$ (caching only 50% of layers); reduced cache rates require empirical performance validation.
- **Scaling factor**: For YOCO++, set $\lambda$ in $[20,50]$; $\lambda=35$ is default.
- **Quality-memory trade-off**: YOCO++ at 50% compression matches or improves full-Transformer accuracy and loss, with half the memory and compute during prefill. Aggressive compression (much less than 50%) can degrade performance.

YOCO’s memory and compute scaling—$\mathcal{O}(N\,d)$ for KV cache, $\mathcal{O}(L\,N\,d)$ for prefill, and no extra decode overhead—enable context lengths and speeds not feasible with standard architectures. YOCO-U further decouples depth and memory, facilitating depth scaling without cache inflation.

## 7. Future Directions and Comparative Context

YOCO’s core principle—decoupling representational depth from cache growth by splitting attention into "efficient self-decoder" and "cross-decoder"—constitutes a paradigm shift for scalable autoregressive modeling. Empirical evidence supports its architectural and efficiency claims across pretraining, downstream tasks, and extreme-context retrieval [2405.05254, 2604.13556, 2604.01220].

Future research avenues include:
- Exploring adaptive/conditional recursion within the self-decoder.
- Incorporating advanced subquadratic efficient-attention mechanisms.
- Extending YOCO’s design to multimodal, encoder-decoder, or retrieval-augmented configurations.

By offering a method to increase context length and model depth while maintaining tractable hardware requirements, YOCO and its extensions present a robust foundation for next-generation LLM inference and serve as the basis for new variants in cross-layer KV compression and parameter-efficient scaling.

Source: https://www.emergentmind.com/topics/yoco-you-only-cache-once