You Only Cache Once, in vLLM

Aug 12, 2025

I recently landed a PR in vLLM (#22628) that allows Google’s Gemma 3n model to skip prefill entirely. This inference optimization is possible because the Gemma 3n architecture uses KV cache sharing in the style of the YOCO architecture. This blog post walks through the YOCO optimization and how I implemented it in vLLM.

YOCO overview

In You Only Cache Once: Decoder-Decoder Architectures for Language Models, researchers from Microsoft propose a decoder-decoder architecture for LLMs called YOCO, consisting of self-decoder and cross-decoder components. The attention layers in the self-decoder compute Key and Value vectors normally and cache them, while the cross-decoder layers re-use the KV cache from the self-decoder and perform cross-attention with their own computed Query vectors.

Because the cross-decoder layers do not emit any KV caches, YOCO has two advantages at inference time: first, KV cache memory is reduced and second, the prefill phase of inference can be skipped for the cross-decoder layers.

Gemma 3n

I had previously implemented cross-layer KV cache sharing in vLLM (#18212), so when I saw that Google’s newly released Gemma 3n model uses KV cache sharing, I knew I had to play around with it in vLLM.

As it turns out, Gemma 3n implements the KV cache sharing setup in the aforementioned YOCO paper. For example, the Gemma3n-E2B model has 5 sets of {SWA, SWA, SWA, SWA, Full Attention} layers for a total of 30 layers, of which the last 10 layers re-use the KV caches from the 19th (SWA) and 20th (Full Attention) layers (the SWA layers use KV from the 19th layer which is also SWA; the FA layers from the 20th layer which is FA). This means that during inference, only the first 20 (KV-cache-generating) layers need to be involved in prefill (where the prompt is processed through the layers), thus saving 10 layers worth of compute!

Implementation in vLLM

By the time I looked at the Gemma 3n model implementation in vLLM, KV cache sharing had already been enabled using the kv_sharing_target_layer_name arg I added in #18212:

This is good because we’re saving KV cache memory, but we weren’t getting the full benefit of YOCO: skipping prefill for the last 10 layers! So I decided to spend some time implementing this in vLLM. It turns out it was more challenging than I thought due to the way that the vLLM V1 engine is designed.

In the vLLM V1 design, the scheduler does not distinguish between prefill and decode. Instead, tokens for requests doing prefill and decode are batched together, as illustrated below:

Because vLLM batches prefill and decode requests together, it’s not trivial to skip computation for some layers during a single model forward pass. In comparison, the implementation would be much easier if prefill and decode were separate phases in the engine execution. For example, if you look at the code repo for the YOCO paper, you’ll see that they implement the prefill computation skip for cross-decoder by having separate loops for prefill and decode (generation), and passing skip_cross_decoder=True to the model during prefill which skips computation for the CrossDecoder module (see below).

To achieve the same effect in vLLM, I realized that we would have to essentially reduce the batch token size for the cross-decoder layers within the model’s forward pass. For example, let’s say we have two requests from a user: R1 and R2. For R1, the prompt length is 6 tokens and for R2, it is 2 tokens. Now let’s say we are serving a 24-layer model where the last 8 layers re-use the KV cache from the 9th-to-last layer. Then we might have a sequence of scheduling rounds in vLLM like this:

In Step 0, vLLM schedules 6 tokens from R1, which are all the prompt tokens. This would be the prefill step for R1, where we process the prompt tokens, populate the KV caches and generate the first output token (for simplicity, I’m ignoring the chunked prefill case).

In Step 1, the first output token generated from the previous step for R1 (token idx 6 in the diagram) is scheduled, which corresponds to a decode step for R1. But because vLLM doesn’t schedule prefill and decode separately, it has also scheduled 2 prompt tokens from R2. This is prefill step for R2. The total number of tokens scheduled in this scheduling round is 3. I’ll refer to this size as the batch token size, since this number becomes the batch dimension of the input to the model (e.g. [3, embedding_size]) during its forward pass.

Lastly, in Step 2, R1 and R2 are both in the decode phase after having completed prefilling. Token 7 is scheduled for R1 and Token 2 is scheduled for R2.

As illustrated in Fig 1, the batch token size remains the same within a single model forward pass, e.g. in step 0 we have 6 tokens so the input shape might be something like [6, embedding_size] and this will be the same for all 24 layers.

As I mentioned earlier in this article, we can actually do better since the last 8 layers don’t need to participate in the prefill computation for R1 and R2.

Let’s see which computation we can skip with the KV cache sharing setup we have. Looking at Fig 2:

• At Step 0, we need all 6 tokens to be processed by the self-decoder layers so that KV can be emitted and cached. However, when we get to the cross-decoder layers, we only need to process the last token of the prompt (token idx 5 in this case). Why is that?
• After Step 0, we should be producing the logits for the first output token (token idx 6). We should also be populating the KV caches for all 6 prompt tokens to attend to for subsequent token generation.
• To obtain these logits, we need to process the last prompt token (token idx 5), which involves processing its embedding through all the cross-decoder layers
• In the cross-decoder’s attention layers, remember we are doing cross-attention with the Keys and Values from the self-decoder layer (for token idx 0-4). These already exist from the self-decoder forward pass.
• Similarly, at Step 1 we only need to process the last prompt token for R2 (token idx 1) when we get to the cross-decoder.

This means that at every scheduling step, if we have more than 1 token from any single request in the (token) batch, we can collapse it to just a single token per request when we get to the cross-decoder layer (this trivially covers chunked prefill case as well, although we can technically skip any intermediate chunk that isn’t the last). In other words, the batch token size of the cross-decoder will be equal to the number of unique requests scheduled at that step. This gets us the inference-time benefit of the YOCO architecture: skipping computation during prefill.

Great! Are we done here? Well, unfortunately not quite. It turns out this works well in vLLM’s eager mode (--enforce-eager), but as the flag suggests vLLM V1 disables eager mode by default. The default setting that vLLM runs with (at the time of this writing) is piecewise torch.compile + CUDA graph, where the model is split by the attention layers and the submodules in-between are compiled and cudagraph captured separately.

This is problematic for us because the batch size reduction within the model forward pass breaks the torch.compile and CUDA graph:

1. torch.compile in vLLM assumes batch size remains the same within a single model forward. The traced graph will be specialized on the batch size, which leads to silent incorrectness if batch size changes within model forward pass.
2. CUDA graphs are shape specialized, meaning we will get incorrect results.

Adding torch.compile and cudagraph support

The solution I came up with (after extensive design discussions with the awesome vLLM committers) is to split the layers into self- and cross-decoder layers, and compile + graph capture them separately. For the 30-layer Gemma3n-E2B model, this means that the first 20 layers and other 10 layers will be grouped separately into independently compiled and CUDA graph captured modules.

This required several changes in the vLLM codebase to support:

• I had to propagate the logit positions from the model runner (logits_indices) to know which tokens to index into to create the batch for the cross-decoder layers. This was landed in PR#21590. I also needed to handle some edge cases around CUDA graph padding that gets added when the batch token size doesn’t match one of the captured graph sizes.
• Based on these logit indices, create a different attention metadata for the Attention layers in the cross-decoder as the token positions to attend to are now different (source)

def make_kv_sharing_fast_prefill_common_attn_metadata(
    common_attn_metadata: CommonAttentionMetadata,
) -> CommonAttentionMetadata:
    if common_attn_metadata.max_query_len == 1:
        # All requests are decode (assume 1 token for now)
        # Skip computing fast prefill path
        return common_attn_metadata

    assert common_attn_metadata.logits_indices_padded is not None
    assert common_attn_metadata.num_logits_indices is not None

    logits_indices_padded = common_attn_metadata.logits_indices_padded
    num_logits_indices = common_attn_metadata.num_logits_indices
    # Get rid of CUDAGraph padding, if any
    logits_indices = logits_indices_padded[:num_logits_indices]
    num_reqs = common_attn_metadata.num_reqs
    query_start_loc = common_attn_metadata.query_start_loc
    seq_lens = common_attn_metadata.seq_lens
    # Example inputs
    # num_reqs: 3
    # generation_indices:  [14, 18, 19, 27]
    # query_start_loc: [0, 15, 20, 28]
    # seq_lens:        [41, 31, 40]

    # Find how many decode indices belong to each request
    # request_ids: [0, 1, 1, 2]
    request_ids = torch.bucketize(logits_indices,
                                  query_start_loc[1:],
                                  right=True)

    # Figure out how many tokens are in each request
    # num_decode_tokens: [1, 2, 1]
    num_decode_tokens = torch.bincount(request_ids, minlength=num_reqs)

    # Calculate new query_start_loc with tokens in generation_indices
    # decode_query_start_loc: [0, 1, 3, 4]
    decode_query_start_loc = torch.empty(num_reqs + 1,
                                         device=query_start_loc.device,
                                         dtype=query_start_loc.dtype)

    decode_query_start_loc[0] = 0
    decode_query_start_loc[1:] = torch.cumsum(num_decode_tokens, dim=0)
    decode_max_query_len = int(num_decode_tokens.max().item())
    total_num_decode_tokens = int(num_decode_tokens.sum().item())

    common_attn_metadata = CommonAttentionMetadata(
        query_start_loc=decode_query_start_loc,
        query_start_loc_cpu=decode_query_start_loc.to("cpu",
                                                      non_blocking=True),
        seq_lens=seq_lens,
        seq_lens_cpu=seq_lens.to("cpu", non_blocking=True),
        num_computed_tokens_cpu=common_attn_metadata.num_computed_tokens_cpu,
        num_reqs=num_reqs,
        num_actual_tokens=total_num_decode_tokens,
        max_query_len=decode_max_query_len,
        max_seq_len=common_attn_metadata.max_seq_len,
        block_table_tensor=common_attn_metadata.block_table_tensor,
        slot_mapping=common_attn_metadata.slot_mapping,
        causal=True,
    )
    return common_attn_metadata

• Refactor the Gemma3nTextModel significantly to split the model into self-decoder and cross-decoder submodules if the prefill optimization is enabled:
• If --kv-sharing-fast-prefill flag is passed, we take a different self.fast_prefill_forward() path which uses the logits_indices_padded metadata passed to index into the subset of tokens for cross-decoder layers (i.e. batch size is reduced). We then merge it back to the output of self-decoder to get the final output.
• If --kv-sharing-fast-prefill flag is passed, we will compile self-decoder and cross-decoder submodules separately (this required more changes in PR#22269), and we also need to pre-allocate static buffers for CUDA graph replay. If it is not passed (default), we will still compile the top-level Gemma3TextModel
• A more subtle change: I also needed to change hidden_states shape from [altup_num_inputs, num_tokens,hidden_size] to [num_tokens,hidden_size, altup_num_inputs] to ensure num_tokens (batch size) comes at dim 0. We cannot have num_tokens be on dim=1 because creating a slice along dim=1 would a) cause torch.compile tensor stride assertions to fail, and b) resolving this by calling contiguous() on the slice would cause memory copy and therefore violate CUDA graph static address constraint.

After these changes, I was able to get self-decoder and cross-decoder piecewise compiled and CUDA graph captured separately to support the batch size reduction for YOCO!

Note

One caveat here is that when prompt_logprobs is enabled for a request, we can no longer use this optimization. This is because by skipping all but last prefill tokens, the logits for the prompt tokens will no longer be valid. For example, multiple choice question (MCQ) evals use prompt_logprobs to get the logprobs of continuation tokens (e.g. lm-evaluation-harness), so using --kv-sharing-fast-prefill will yield inaccurate results. To prevent this, completion requests with prompt logprobs will be rejected (HTTP 400 Bad Request) for online serving. For offline serving (i.e. with LLM class), it will raise an assertion error.

Results

I verified that the evals were on par after enabling this optimization:

Now for what I’ve been waiting for: the performance improvements

On my H100 setup, I performed a sweep over max-concurrency and random-input-len for benchmarking and found from 17% to 55% reduction in the average Time to First Token (TTFT). Interestingly, because vLLM batches prefill and decode together, the average Time Per Output Token (TPOT) also sees improvement across the board (prefill of one request takes shorter, which makes the whole batch including other decode requests shorter, leading to faster TPOT). Here’s a chart showing the same results more visually:

Take a look at the pull request #22628 for more details.

Conclusion

Overall, this was a fun exercise that took more effort and time than I initially expected going into this. One thing I’ve learned is that translating a research idea (i.e. YOCO prefill savings) from a paper demo to a production inference engine might be more tricky than you think because of different use cases and edge cases you may need to support. Getting this implementation finally landed in vLLM took multiple rounds of discussion with the wider vLLM community (special thanks to vLLM committers Chen Zhang and Lucas Wilkinson for the discussions) to support this feature in a way that fit well with vLLM’s design. It also took several PRs to get it merged: #21590, #22672, #22269 and finally #22628 to enable it for Gemma 3n. Update (April 2026): this prefill optimization was adopted by vLLM for the Gemma4 release as well with impressive speedups, see #38879).