[Gemma4] Enable Fast Prefill Optimization

[LucasWilkinson]

Summary

Add --kv-sharing-fast-prefill support for Gemma 4 models, porting the YOCO (You Only Cache Once) fast prefill optimization from Gemma3n. When enabled, the cross-decoder layers (KV-shared) skip prefill tokens and only process decode tokens, significantly reducing prefill latency and improving throughput under concurrent load.

shout-out to @sarckk for the original optimzation (#22628)

Test Plan

GSM8K accuracy (Gemma4-E4B, 5-shot)

# FP=OFF (baseline)
lm_eval --model vllm --tasks gsm8k --num_fewshot 5 \
  --model_args pretrained=google/gemma-4-E4B-it,gpu_memory_utilization=0.9,max_model_len=4096,tensor_parallel_size=1,trust_remote_code=True,attention_backend=TRITON_ATTN,kv_sharing_fast_prefill=False \
  --batch_size auto --apply_chat_template --fewshot_as_multiturn

# FP=ON (this PR)
lm_eval --model vllm --tasks gsm8k --num_fewshot 5 \
  --model_args pretrained=google/gemma-4-E4B-it,gpu_memory_utilization=0.9,max_model_len=4096,tensor_parallel_size=1,trust_remote_code=True,attention_backend=TRITON_ATTN,kv_sharing_fast_prefill=True \
  --batch_size auto --apply_chat_template --fewshot_as_multiturn

Serving benchmark

# Start server (without fast prefill)
vllm serve google/gemma-4-E4B-it \
  --port 8434 \
  --disable-log-stats \
  --no-enable-prefix-caching \
  --max-num-seqs 128 \
  --max-model-len 32768 \
  --max-num-batched-tokens 8192 \
  --attention-backend TRITON_ATTN \
  --trust-remote-code

# Start server (with fast prefill)
vllm serve google/gemma-4-E4B-it \
  --port 8434 \
  --disable-log-stats \
  --no-enable-prefix-caching \
  --max-num-seqs 128 \
  --max-model-len 32768 \
  --max-num-batched-tokens 8192 \
  --attention-backend TRITON_ATTN \
  --trust-remote-code \
  --kv-sharing-fast-prefill

# Run benchmark (after server is ready)
# concurrency=8
vllm bench serve \
  --backend vllm \
  --ignore-eos \
  --port 8434 \
  --model google/gemma-4-E4B-it \
  --dataset-name random \
  --max-concurrency 8 \
  --request-rate inf \
  --num-prompts 256 \
  --random-input-len 8192 \
  --random-output-len 150

# concurrency=32
vllm bench serve \
  --backend vllm \
  --ignore-eos \
  --port 8434 \
  --model google/gemma-4-E4B-it \
  --dataset-name random \
  --max-concurrency 32 \
  --request-rate inf \
  --num-prompts 256 \
  --random-input-len 8192 \
  --random-output-len 150

Test Results

GSM8K accuracy (Gemma4-E4B, 5-shot)

No accuracy regression:

	strict-match	flexible-extract
FP=OFF (baseline)	0.1054	0.1751
FP=ON (this PR)	0.1031	0.1850

Serving performance (Gemma4-E4B, 1xB200, ISL=8192, OSL=150, n=256)

concurrency=8

Metric	NORMAL	FAST_PREFILL	Delta
Throughput	4.22 req/s	5.06 req/s	+19.9%
Mean TTFT	570 ms	363 ms	-36.3%
Mean TPOT	8.90 ms	8.16 ms	-8.3%

concurrency=32

Metric	NORMAL	FAST_PREFILL	Delta
Throughput	6.53 req/s	9.07 req/s	+38.9%
Mean TTFT	942 ms	622 ms	-34.0%
Mean TPOT	26.43 ms	19.37 ms	-26.7%

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request introduces comprehensive support for the Gemma 4 model family, encompassing both text-only and multimodal (image, audio, and video) capabilities. Key additions include specialized proportional RoPE, logic to handle heterogeneous head dimensions, and custom parsers for reasoning and tool calls. The review feedback identifies several critical robustness and performance improvements: replacing process-terminating sys.exit(1) calls with ValueError exceptions, optimizing memory by conditionalizing tensor clones in the fast prefill path, ensuring global context consistency using try...finally blocks, and implementing bounds checking for batch sizes during profiling to prevent potential runtime crashes.

[RyanMullins]

LGTM. Shared layers don't compute so you can early exit depending on the config.

[lk-chen]

verified on TPU with same set up as vllm-project/tpu-inference#2126 (comment), MMMU-pro score is identical before/after this current PR. Performance metrics untested.

[bbrowning]

Just an FYI this may have created correctness issues in certain situations. In debugging #39392, reverting this appears to have fixed that problem for me. It could be this only impacts certain hardware or setups - I'm using a DGX Spark and the user in 39392 was using 8x RTX 4090.