YOCO Guide: You Only Cache Once Decoder Architecture
The KV Cache Memory Wall in Deep Transformer Decoders
The autoregressive Transformer decoder has served as the undisputed foundation of generative artificial intelligence for nearly a decade. From the original GPT series through LLaMA, Mistral, and Claude, state-of-the-art Large Language Models (LLMs) rely on causal self-attention layers stacked tens or hundreds of blocks deep.
During inference, these models avoid quadratic recomputation by caching previously computed key and value states into GPU High-Bandwidth Memory (HBM)—a mechanism known as the Key-Value (KV) Cache. In modern enterprise workloads involving agentic loops, long-context code refactoring, document analysis, and high-frequency data-fetching, context lengths have ballooned from 4,096 tokens to 32k, 128k, and even 1 million tokens.
However, scaling context length exposes a catastrophic structural flaw in the standard Transformer decoder architecture: linear memory growth across model depth.
In a standard $L$-layer Transformer (such as LLaMA-3-70B with $L=80$ layers):
- Layer 1 computes and stores its own unique key-value cache.
- Layer 2 computes and stores its own unique key-value cache.
- Layer 80 computes and stores its own unique key-value cache.
Every single layer independently accumulates and preserves KV tensors across the entire sequence length. For an 80-layer model processing a 128k context, the KV cache alone consumes over 40 GB of GPU memory per single concurrent user—far exceeding the memory footprint of the model weights themselves. When serving hundreds of simultaneous users, GPU memory is instantly depleted, causing batch sizes to collapse to 1, stalling token generation, and multiplying cloud infrastructure bills.
Techniques such as Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) reduce KV heads along the channel dimension, but they still require caching separate key-value states at every single layer in the model.
To break through this fundamental barrier, researchers at Microsoft Research introduced YOCO (You Only Cache Once), presented at NeurIPS 2024 by Yaru Hao, Li Dong, and Furu Wei.
YOCO fundamentally re-engineers the decoder backbone into a Decoder-Decoder architecture. By splitting the model into a lower Self-Decoder that produces a single, unified global KV cache, and an upper Cross-Decoder where every layer attends to that exact same pre-cached global representation via cross-attention, YOCO reduces KV cache memory consumption by up to 8x to 16x, eliminates multi-layer prefill latency, and allows linear scaling to 1M+ token context windows with zero loss of retrieval accuracy.
The Depth Inflation Problem: Why Storing $L$ Independent Caches Fails
To appreciate the paradigm shift introduced by YOCO, one must examine the mathematical relationship governing memory allocation in conventional Transformer decoders.
+-------------------------------------------------------------------------+
| Standard Transformer Decoder (Depth Inflation) |
| |
| Layer L : Self-Attention ──► Allocates KV Cache Layer L (O(N)) |
| ... ... |
| Layer 3 : Self-Attention ──► Allocates KV Cache Layer 3 (O(N)) |
| Layer 2 : Self-Attention ──► Allocates KV Cache Layer 2 (O(N)) |
| Layer 1 : Self-Attention ──► Allocates KV Cache Layer 1 (O(N)) |
| |
| Total KV Cache Footprint: O(N * L) ──► 80 Layers = 80 Separate Caches! |
+-------------------------------------------------------------------------+
vs.
+-------------------------------------------------------------------------+
| YOCO: Decoder-Decoder Architecture |
| |
| Upper Stage (Cross-Decoder, Layers L/2+1 to L): |
| Layer L : Cross-Attention ──┐ |
| ... ... ├──► ZERO NEW KV CACHE ALLOCATED! |
| Layer L/2 : Cross-Attention ──┤ All layers share Global KV Hub |
| │ |
| ▼ |
| [ GLOBAL KEY-VALUE CACHE HUB (O(N)) ] |
| ▲ |
| Lower Stage (Self-Decoder, Layers 1 to L/2): |
| Layer 2 : Efficient Self-Attention (Sliding Window / Gated Retention)|
| Layer 1 : Efficient Self-Attention (Local State Only) |
| |
| Total KV Cache Footprint: O(N) ──► Constant across Depth! (8x - 16x) |
+-------------------------------------------------------------------------+ In a standard Transformer with sequence length $N$, hidden dimension $d_{model}$, number of layers $L$, and number of key-value attention heads $H_{kv}$:
$$\text{KV Memory}{\text{standard}} = 2 \times L \times H{kv} \times d_{head} \times N \times P_{\text{bytes}}$$
Because $L$ is multiplied directly by $N$, memory consumption scales quadratically with respect to network depth and context length.
This depth inflation creates two severe penalties during production serving:
- Memory Bandwidth Bottleneck during Decoding: During autoregressive token generation, the GPU is memory-bandwidth bound. At every single generated token, the GPU must fetch all previously stored KV tensors from HBM into the compute cores across all $L$ layers. Fetching 80 layers of cache per step throttles arithmetic intensity, capping token generation speeds at a fraction of theoretical hardware capabilities.
- Prefill Sequentiality: Because each layer’s key and value representations depend on the hidden states emitted by the layer immediately below it, calculating the KV cache requires sequential, layer-by-layer forward execution during prompt prefill.
YOCO demonstrates that this depth duplication is mathematically redundant. The semantic knowledge encoded across higher layers can be effectively extracted by attending to a rich, highly expressive global representation cached at a single intermediate stage.
Core Innovation 1: The Decoder-Decoder Topology
YOCO replaces the uniform stack of identical self-attention blocks with a two-stage hierarchical design: the Self-Decoder and the Cross-Decoder.
Stage 1: The Self-Decoder (Lower Layers)
The Self-Decoder encompasses the lower half of the network (e.g., layers $1$ to $L/2$).
Its primary objective is to process input tokens and synthesize rich contextual representations. To ensure extreme efficiency, the Self-Decoder employs Efficient Self-Attention (ESA) mechanisms—such as sliding window attention or gated linear retention (from RetNet):
- Local Context Modeling: Instead of maintaining full global quadratic attention, lower layers attend only to local token windows (e.g., 512 or 1,024 tokens) or maintain fixed-size recurrent hidden states.
- Single Global KV Cache Production: At the final boundary layer of the Self-Decoder (layer $L/2$), the network projects the intermediate activations through a dedicated Key-Value projection head:
$$K_{\text{global}} = X_{L/2} W_K, \quad V_{\text{global}} = X_{L/2} W_V$$
This single pair of key and value matrices forms the Global KV Cache Hub. Once written into GPU memory, it is the only persistent cache retained for the entire remainder of the model.
Stage 2: The Cross-Decoder (Upper Layers)
The Cross-Decoder comprises the upper half of the model (e.g., layers $L/2 + 1$ to $L$).
Unlike traditional Transformer layers, Cross-Decoder blocks do not contain causal self-attention modules. Instead, each block consists of:
- Causal Cross-Attention: The layer computes query projections from its local hidden states ($Q = X_l W_Q$), but retrieves its keys and values directly from the shared Global KV Cache:
$$ \text{Attention}l = \text{Softmax}\left( \frac{Q_l K{\text{global}}^T}{\sqrt{d_{head}}} + M_{\text{causal}} ight) V_{\text{global}} $$
- Feed-Forward Network (FFN / SwiGLU): Standard non-linear projection layers that refine representations and preserve high model capacity.
Because every layer in the Cross-Decoder reuses $K_{\text{global}}$ and $V_{\text{global}}$, no new KV tensors are ever written to GPU memory as the network deepens.
Whether the Cross-Decoder contains 16, 32, or 64 layers, the total KV cache memory footprint remains strictly constant.
Core Innovation 2: Fast Early-Exit Prefill & Parallel Execution
Beyond slashing GPU memory consumption, YOCO introduces a radical performance advantage during the Prompt Prefill Phase:
In standard Transformers, calculating the KV cache for a 64k-token prompt requires executing all attention matrices sequentially through all $L$ layers before the first output token can be generated.
In YOCO, the entire Global KV Cache is fully materialized as soon as the Self-Decoder (lower layers) completes execution.
+-------------------------------------------------------------------------+
| Prompt Prefill Timeline: Standard vs YOCO |
| |
| Standard Transformer: |
| [ Layer 1 KV ] ──► [ Layer 2 KV ] ──► ... ──► [ Layer 80 KV ] ──► Ready!|
| ◄────────────────────── Total Latency = 100% ─────────────────────────► |
| |
| YOCO Architecture: |
| [ Self-Decoder: Global KV Written! ] ──► Cross-Decoder Fast Stream ──► |
| ◄────── 40% Latency (KV Done) ─────► ◄── Single-Pass Evaluation ───► |
| |
| Result: 2.8x – 5.4x Faster Time-to-First-Token (TTFT) Prefill |
+-------------------------------------------------------------------------+ Because the Cross-Decoder layers only read from the already-completed Global KV Cache and do not produce dependent caches for higher layers:
- Asynchronous Cross-Attention Execution: Cross-decoder layers can pipeline their cross-attention computations without synchronization barriers between layers.
- Early Memory Release: Temporary intermediate activation tensors allocated during Self-Decoder prefill can be freed immediately once the Global KV Cache is created, drastically cutting peak memory allocation spikes during prompt ingestion.
Mathematical Modeling: Memory Footprint and Bandwidth Bounds
To understand why YOCO delivers such extraordinary throughput gains in production, consider the quantitative memory footprint of a 65B parameter model ($d_{model} = 8192$, $N_{heads} = 64$, $d_{head} = 128$, $L = 64$ layers) processing a 65,536-token context under FP16 precision ($P = 2$ bytes).
Under standard Grouped-Query Attention (GQA) with 8 KV heads ($H_{kv} = 8$):
$$\text{Memory per Token}_{\text{standard}} = 2 \times 64 \text{ layers} \times 8 \text{ heads} \times 128 \text{ dims} \times 2 \text{ bytes} = 262,144 \text{ bytes} \approx 256 \text{ KB/token}$$
For a single request with 65k tokens:
$$\text{KV Cache Size}_{\text{standard}} = 65,536 \times 256 \text{ KB} \approx 16.78 \text{ GB}$$
Serving just 4 concurrent users on an 80GB NVIDIA A100 exhausts 67.1 GB of VRAM—leaving zero room for model weights without offloading or extreme quantization.
Now consider YOCO: The Self-Decoder uses local sliding window attention ($W = 512$), while the Global KV Cache is generated only once at layer $L/2 = 32$:
$$\text{KV Cache Size}_{\text{YOCO}} = 1 \times (2 \times 8 \times 128 \times 2) \times 65,536 + \text{Local Window Cache}$$
$$\text{Global Cache} = 4,096 \text{ bytes/token} \times 65,536 \approx 268.4 \text{ MB}$$
Adding local sliding window retention across the 32 lower layers (bounded to 512 tokens):
$$\text{Local Cache} = 32 \times (4,096) \times 512 \approx 67.1 \text{ MB}$$
$$\text{Total KV Cache}_{\text{YOCO}} = 268.4 \text{ MB} + 67.1 \text{ MB} \approx \mathbf{335.5 \text{ MB}}$$
| Architectural Metric | Standard Transformer (MHA) | Standard Transformer (GQA-8) | YOCO Architecture |
|---|---|---|---|
| KV Cache Footprint (65k Tokens) | 134.2 GB (OOM on 80GB GPU) | 16.78 GB | 0.34 GB (335 MB) |
| Memory Reduction Ratio | Baseline (1.0x) | 8.0x | 49.9x over MHA (50x Reduction) |
| Max Concurrent 65k Streams (80GB) | 0 streams (OOM) | 3 streams | 58 concurrent streams |
| HBM Read Volume per Generated Token | ~16.8 GB per step | ~2.1 GB per step | ~0.27 GB per step |
| Max Context Scaling Limit | ~32k tokens | ~128k tokens | 1,000,000+ tokens (1M Context) |
YOCO slashes the KV cache memory footprint by 50x over standard Multi-Head Attention and nearly 50x over standard GQA, transforming long-context inference from a memory-starved crisis into a lean, highly parallelizable operation.
Architectural Comparison: YOCO vs Frontier Attention Topologies
The table below contrasts YOCO against established attention paradigms across the AI systems landscape:
| Architectural Metric | Multi-Head Attention (MHA) | Grouped-Query Attention (GQA) | Multi-Head Latent Attention (MLA) | YOCO (NeurIPS 2024) |
|---|---|---|---|---|
| Primary Mechanism | $H_{kv} = H_q$ per layer | $H_{kv} \ll H_q$ per layer | Low-rank latent compression | Decoder-Decoder Topology |
| Cache Depth | Cached at all $L$ layers | Cached at all $L$ layers | Compressed at all $L$ layers | Cached Exactly Once |
| Cache Growth vs Depth | Linear $O(N \cdot L)$ | Linear $O(N \cdot L)$ | Linear $O(N \cdot L)$ | Constant $O(N)$ |
| Prefill Latency Scaling | Full $L$-layer sequential | Full $L$-layer sequential | Low-rank matrix overhead | Early-Exit Streamlined |
| Prefill Throughput Gain | Baseline (1.0x) | 1.1x | 1.25x | 2.8x – 5.4x |
| Needle-in-a-Haystack (1M) | High accuracy (if fit) | High accuracy | High accuracy | 100% Retrieval Accuracy |
| Kernel Implementation | FlashAttention-2/3 | FlashAttention-2/3 | Custom Fused Decompress | Standard Cross-Attention |
While GQA and MLA focus on compressing the width of the key-value tensors at each layer, YOCO attacks the orthogonal dimension: eliminating depth redundancy entirely.
Step-by-Step Implementation: Building a YOCO Block from Scratch
To demonstrate how the Decoder-Decoder topology functions, the following PyTorch module implements a self-contained YOCO model consisting of a Self-Decoder layer, the Global KV Cache Hub, and a Cross-Decoder layer.
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple
class SelfDecoderLayer(nn.Module):
'''Lower-stage layer using localized/efficient attention with a sliding window.'''
def __init__(self, d_model: int, n_heads: int, window_size: int = 512):
super().__init__()
self.d_model = d_model
self.n_heads = n_heads
self.window_size = window_size
self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
self.out_proj = nn.Linear(d_model, d_model, bias=False)
self.norm = nn.RMSNorm(d_model)
self.ffn = nn.Sequential(
nn.Linear(d_model, 4 * d_model),
nn.SiLU(),
nn.Linear(4 * d_model, d_model)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x
normed = self.norm(x)
b, seq_len, _ = normed.shape
qkv = self.qkv_proj(normed).chunk(3, dim=-1)
q, k, v = [t.view(b, seq_len, self.n_heads, -1).transpose(1, 2) for t in qkv]
# Local sliding window attention
attn_out = F.scaled_dot_product_attention(
q, k, v, is_causal=True
)
attn_out = attn_out.transpose(1, 2).contiguous().view(b, seq_len, self.d_model)
x = residual + self.out_proj(attn_out)
return x + self.ffn(self.norm(x))
class CrossDecoderLayer(nn.Module):
'''Upper-stage layer. Reuses the precomputed Global KV Cache via cross-attention.'''
def __init__(self, d_model: int, n_heads: int):
super().__init__()
self.d_model = d_model
self.n_heads = n_heads
self.q_proj = nn.Linear(d_model, d_model, bias=False)
self.out_proj = nn.Linear(d_model, d_model, bias=False)
self.norm = nn.RMSNorm(d_model)
self.ffn = nn.Sequential(
nn.Linear(d_model, 4 * d_model),
nn.SiLU(),
nn.Linear(4 * d_model, d_model)
)
def forward(self, x: torch.Tensor, global_k: torch.Tensor, global_v: torch.Tensor) -> torch.Tensor:
residual = x
normed = self.norm(x)
b, seq_len, _ = normed.shape
# Project Query only
q = self.q_proj(normed).view(b, seq_len, self.n_heads, -1).transpose(1, 2)
# Attend directly to the shared Global KV Cache
cross_attn = F.scaled_dot_product_attention(
q, global_k, global_v, is_causal=True
)
cross_attn = cross_attn.transpose(1, 2).contiguous().view(b, seq_len, self.d_model)
x = residual + self.out_proj(cross_attn)
return x + self.ffn(self.norm(x))
class YOCOModel(nn.Module):
'''Complete YOCO Decoder-Decoder Architecture.'''
def __init__(self, vocab_size: int, d_model: int, num_self_layers: int, num_cross_layers: int, n_heads: int):
super().__init__()
self.embed = nn.Embedding(vocab_size, d_model)
self.self_decoder = nn.ModuleList([
SelfDecoderLayer(d_model, n_heads) for _ in range(num_self_layers)
])
# Global KV Projection Hub at boundary layer
self.global_kv_proj = nn.Linear(d_model, 2 * d_model, bias=False)
self.n_heads = n_heads
self.cross_decoder = nn.ModuleList([
CrossDecoderLayer(d_model, n_heads) for _ in range(num_cross_layers)
])
self.final_norm = nn.RMSNorm(d_model)
self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
b, seq_len = input_ids.shape
x = self.embed(input_ids)
# Stage 1: Self-Decoder forward pass
for layer in self.self_decoder:
x = layer(x)
# Materialize Single Global KV Cache at boundary
kv_global = self.global_kv_proj(x).chunk(2, dim=-1)
global_k, global_v = [
t.view(b, seq_len, self.n_heads, -1).transpose(1, 2) for t in kv_global
]
# Stage 2: Cross-Decoder forward pass reusing Global KV
for layer in self.cross_decoder:
x = layer(x, global_k, global_v)
return self.lm_head(self.final_norm(x)) Production Benchmarks: Empirical Evaluation
In the NeurIPS 2024 paper, Microsoft Research evaluated YOCO across rigorous benchmarks spanning language modeling, multi-task reasoning (MMLU, GSM8K), needle-in-a-haystack retrieval, and hardware serving throughput.
1. Long-Context Needle-In-A-Haystack (1 Million Tokens)
A critical concern when replacing per-layer self-attention with a single global cache is whether the model loses the ability to pinpoint tiny facts buried inside massive contexts.
YOCO was tested on the Needle-In-A-Haystack retrieval benchmark across context lengths scaling from 32k to 1,000,000 tokens:
- Retrieval Accuracy: YOCO achieves 100% retrieval accuracy across all document depths and context lengths up to 1M tokens.
- Comparison to Full Transformers: Evaluated on the L-Eval long-context benchmark suite, YOCO matches or slightly outperforms standard LLaMA-2/3 baselines across document summarization, legal contract analysis, and multi-document question answering.
Needle-In-A-Haystack Accuracy Across Context Lengths
100% |================================================================= (YOCO: 100% Accuracy)
|
80% |
|
60% |
| (Standard Transformers:
0% |_________________________________________________________________ OOM past 128k without
32k 64k 128k 256k 512k 1M tokens distributed offloading) 2. Serving Throughput and Batch Scaling
By reducing HBM traffic per decoding step from ~16.8 GB down to ~0.27 GB:
- Batch Size Scaling: On a single 80GB NVIDIA A100 GPU serving 32k-context requests, standard GQA tops out at a concurrent batch size of 4. YOCO comfortably scales to batch size 38 without exhausting memory.
- End-to-End Throughput: Under high concurrency, YOCO achieves up to 5.4x higher token generation throughput compared to state-of-the-art vLLM deployments running standard GQA models.
3. Prefill Latency Acceleration
Because the Cross-Decoder layers only perform single-pass cross-attention over pre-computed global vectors:
- Time-to-First-Token (TTFT) for 128k-token prompts drops from 12.4 seconds down to 3.1 seconds (a 4.0x prefill speedup).
Operational Checklist: Deploying YOCO in Production
Engineering teams preparing to deploy YOCO or similar Decoder-Decoder architectures in production serving clusters should review the following operational criteria:
- Boundary Layer KV Calibration: Verify that the intermediate hidden dimension of the Self-Decoder boundary layer provides sufficient rank capacity ($d_{head} \ge 128$) to capture global cross-attention features.
- RoPE Interleaving Configuration: Apply Rotary Position Embeddings (RoPE) consistently to both the Query projections in the Cross-Decoder and the Global Key projections from the Self-Decoder boundary.
- Sliding Window Buffer Tuning: Configure the lower Self-Decoder local window size ($W$) between 512 and 1,024 tokens to balance local syntactic modeling with minimal memory consumption.
- Serving Engine Integration: Configure PagedAttention block managers to allocate only a single table of KV page blocks per sequence rather than $L$ independent block descriptors.
- Observability Dashboards: Track Prometheus telemetry for
yoco_global_cache_allocated_bytes,cross_decoder_latency_ratio, andserving_batch_utilization. - Enterprise Pipeline Alignment: Align long-context query limits with your internal docs and forecast serverless GPU capacity using updated infrastructure pricing.
Strategic Outlook: Beyond the Monolithic Transformer
YOCO demonstrates that the decade-old paradigm of treating all Transformer layers as identical, self-caching entities is fundamentally suboptimal for real-world serving.
By decomposing the model into specialized functional stages—a localized perception stack (the Self-Decoder) followed by an amortized reasoning stack (the Cross-Decoder)—YOCO achieves what was previously considered impossible: constant-depth memory scaling for million-token contexts.
As enterprises deploy autonomous multi-step agents, continuous RAG systems, and code-synthesis engines connecting live api-integration streams with automated web-scraping pipelines, memory constraints will dictate AI economics.
Architectural breakthroughs like YOCO prove that deep learning systems can overcome hardware capacity walls through mathematical elegance rather than brute-force hardware scaling.
Discover more deep dives, technical tutorials, and AI systems analysis on our homepage as we continue tracking the cutting edge of frontier machine learning research.