Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

This PR adds an option to enable prefill optimization for Gemma3n model with --kv-sharing-fast-prefill.

Background

In You Only Cache Once (https://arxiv.org/abs/2405.05254), self-decoder layers generate KV caches while cross-decoder layers use cross-attention and reuse the shared KV cache. As only self-decoder layers generate KV caches, cross-decoder layers don't need to do prefill. Below is a figure from the YOCO paper on the prefill optimization:

Design

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

When we skip tokens corresponding to prefill in the cross-decoder layers, we therefore will have the batch size reduced during model forward for the cross-decoder layers:

Without optimization enabled (baseline)

With optimization enabled (--kv-sharing-fast-prefill)

With this change, we can no longer compile the top-level model for 2 reasons:

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, so we will get incorrect results.

Solution: we split the layers into self- and cross-decoder layers, and compile + graph capture them separately. For Gemma3n-E2B which has 30 layers, the first 20 layers and other 10 layers will be grouped separately into independently compiled and CUDA graph captured modules.

Other changes required in this PR:

• Build attention metadata builder subclass for eligible layers so it can call make_kv_sharing_fast_prefill_common_attn_metadata to create an attention metadata excluding all prefill tokens. This requires passing logits_indices to CommonAttentionMetadata
• Create a subclass of attention metadata for eligible layers which isinstance of KVSharingFastPrefillAttentionMetadata. This has two additional metadata (logits_indices_padded and num_logits_indices) which are required for indexing into hidden states in the model implementation to match the shapes that the new attention metadata expects
• Changes to Gemma3n model implementation.
  • Need 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.
  • 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, 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
  • Attention group changes. After this PR, the attn groups looks for Gemma3n-2B looks like this:

- attn_groups[0] (non-sliding window layers)
  - attn_groups[0][0]: 4, 9, 14, 19 
  - attn_groups[0][1]: 24, 29

- attn_groups[1]
  - attn_groups[1][0] layers: 0, 1, 2, 3

- attn_groups[2]
  - attn_groups[2][0] layers 5, 6, 7, 8

- attn_groups[3]
  - attn_groups[3][0] layers: 10, 11, 12, 13

- attn_groups[4] (sliding window layers)
  - attn_groups[4][0] layers: 15, 16, 17, 18
  - attn_groups[4][1] layers: 20, 21, 22, 23, 25, 26, 27, 28

Compared to trunk, the only difference is there are extra groups for attn_groups[0] and attn_groups[4] for layers which need a separate attention metadata builder for the fast prefill path. Previously it looks like this:

- attn_groups[0] (non-sliding window layers)
  - attn_groups[0][0]: 4, 9, 14, 19, 24, 29 

# attn_groups[1], attn_groups[2] and attn_groups[3] same

- attn_groups[4] (sliding window layers)
  - attn_groups[4][0] layers: 15, 16, 17, 18, 20, 21, 22, 23, 25, 26, 27, 28

Follow ups

• Address IMA issue when using fast prefill with hybrid KV cache manager (currently hybrid KV manager is disabled if --kv-sharing-fast-prefill is set) (EDIT: seems like [BugFix] Fix for IMA in FA3 varlen combine #22967 fixed it)
• Land fix for incorrect weakref with multiple compile units (Add option to disable weakref conversion for last piecewise cudagraph in a module #22282) to get rid of self_decoder_hidden_state.clone() in Gemma3n

Test Plan

Evals

PORT=8000
vllm serve google/gemma-3n-E2B-it --disable-log-requests

lm_eval --model local-completions --tasks $TASKNAME \
    --model_args model=google/gemma-3n-E2B-it,base_url=http://127.0.0.1:$PORT/v1/completions,num_concurrent=200,tokenized_requests=False --batch_size auto --apply_chat_template --fewshot_as_multiturn

ran gsm8k, mmlu, mmlu pro

Unit tests

pytest tests/v1/worker/test_gpu_model_runner.py -k "test_init_kv_cache"

pytest tests/v1/e2e/test_kv_sharing_fast_prefill.py::test_kv_sharing_fast_prefill

Performance

VLLM_DISABLE_COMPILE_CACHE=1 python -m vllm.entrypoints.openai.api_server --model google/gemma-3n-E2B-it --disable-log-requests -tp 1 --port 8000 --no-enable-prefix-caching --max-num-seqs 128 --max-model-len=32768 --max_num_batched_token=8192 --kv-sharing-fast-prefill

Perform sweep over max-concurrency and random-input-len,
$num_reqs = 256
max-num-batched-tokens = 8192
max-num-seqs = 128

python benchmarks/benchmark_serving.py     --backend vllm     --ignore-eos     --port 8000     --model google/gemma-3n-E2B-it     --dataset-name random --max-concurrency 8 --request-rate inf --num-prompts $num_reqs         --random-input-len 8192 --random-output-len 150

Test Result

Evals

Evals on par

gsm8k: 0.5451 and mmlu_pro:0.3424 with hybrid memory allocator enabled as well

Unit tests

Unit tests all pass

Performance

Mean TTFT and TPOT (ms):