You Only Cache Once: Decoder-Decoder Architectures for Language

You Only Cache Once:

Decoder-Decoder Architectures for Language Models

Yutao Sun

∗ †‡

Li Dong

∗ †

Yi Zhu

†

Shaohan Huang

†

Jianyong Wang

‡

Furu Wei

†⋄

†

Microsoft Research

‡

Tsinghua University

https://aka.ms/GeneralAI

Figure 1: We propose a decoder-decoder architecture, YOCO, for large language model, which only

caches key/value once. YOCO markedly reduces the KV cache memory and the prefilling time, while

being scalable in terms of training tokens, model size, and context length. The inference cost is

reported to be 512K as the context length, and Figures 7–10 present more results for different lengths.

∗

Equal contribution. ⋄ Corresponding author.

1 Introduction

use a bidirectional encoder to encode input and a unidirectional decoder to generate output. Both

of the above layouts struggle with autoregressive generation due to bidirectionality. Specifically,

encoders have to encode the whole input and output tokens again for the next generation step.

the parameters of encoder, especially for multi-turn conversation. Third, decoder-only language

models, such as GPT [

20

], generate tokens autoregressively. By caching the previously

This compelling feature establishes the decoder-only language model as the standard option.

capacity of one H100-80GB GPU. In addition, the prefilling latency of long-sequence input is

extremely high. For instance, using four H100 GPUs, the 7B language model (augmented with

Flash-Decoding [

DHMS23

] and kernel fusion) requires about 110 seconds to prefill 450K tokens, and

380 seconds for 1M length. The above bottlenecks make it difficult to deploy long-context language

models in practice.

In this work, we propose a decoder-decoder architecture, YOCO, for large language models, which

only caches KV pairs once. Specifically, we stack cross-decoder upon self-decoder. Given an input

sequence, the self-decoder utilizes efficient self-attention to obtain KV caches. Then the cross-decoder

layers employ cross-attention to reuse the shared KV caches. The decoder-decoder architecture is

conceptually similar to encoder-decoder, but the whole model behaves more like a decoder-only

allows for more efficient system design for distributed long-sequence training. In addition, we

gating mechanism.

We conduct extensive experiments to show that YOCO achieves favorable language modeling perfor-

mance and has many advantages in terms of inference efficiency. Experimental results demonstrate

that YOCO can be scaled up with more training tokens, larger model size, and longer context length.

Specifically, we scale up the 3B YOCO model to trillions of training tokens, attaining results on par

with prominent Transformer language models, such as StableLM [

TBMR

]. Moreover, the scaling

curves ranging from 160M to 13B show that YOCO are competitive compared to Transformer. We

1M tokens. The prefill stage is speeded up by

for the 1M context and

for the 32K input.

For example, for a 512K context, YOCO reduces the Transformer prefilling latency from 180 seconds

to less than six seconds. The results position YOCO as a strong candidate model architecture for

future large language models with native long-sequence support.

2

caches. As the self-decoder utilizes an efficient attention module, the cache size is bounded to a constant, which

can be ignored compared to global caches when the sequence length is large.

2

Figure 2: Overview of the decoder-decoder architecture. Self-decoder generates the global KV cache.

self-decoder and cross-decoder. Specifically, YOCO is stacked with

are self-decoder while the rest modules are cross-decoder. Given an input sequence

|x|

,

the input embeddings are packed into

|x|

, where

model

is hidden

/2

l

l−1

.

Both self- and cross-decoder follow a similar block layout (i.e., interleaved attention and feed-forward

network) as in Transformer [

VSP

17

]. We also include pre-RMSNorm [

Sha20

RMSNorm [ZS19] is used for LN(·). Causal masking is used for efficient self-attention.

3

the inference capacity bottleneck becomes KV caches (Figure 7b), our method enables us to serve

4