[roberteg16]

Gemma 4's vision tower becomes a native vLLM module that packs the images (or video frames) of a batch into one sequence, so it runs on any --mm-encoder-attn-backend and its encoder CUDA graphs follow token budgets like Qwen3-VL. Draft for review before proposing it upstream. Depends on #1417 (encoder CUDA graphs: eager fallback for packed TORCH_SDPA), which is this PR's base branch, so the diff shows only the Gemma 4 change.

Why

Today the tower is transformers' Gemma4VisionModel with vLLM linears (vllm-project#43798): every image is padded to max_soft_tokens * 9 patches, attention is SDPA with a padding mask (it ignores --mm-encoder-attn-backend and never reaches FlashAttention), the pooler builds a one_hot over the padded batch, and encoder CUDA graphs capture padded rows per row width (vllm-project#46837, vllm-project#61365). Earlier attempts (vllm-project#43440, vllm-project#45574, issue vllm-project#43178) stopped at the padding; the review asked for MMEncoderAttention and ApplyRotaryEmb, which need the packed layout.

Changes

• gemma4_vision.py (new): built from vLLM layers: MMEncoderAttention with cu_seqlens, ReplicatedLinear, RMSNorm, get_act_fn, the axial 2D RoPE from get_rope (head_dim / 2 rotary dims per axis) applied with ApplyRotaryEmb on each half of the head (tables gathered once per forward), and moonvit's _build_merge_gather_idx / patch_merger_packed for the 3x3 pooling (averaged in fp32 like transformers). It validates that each image's patches form a row-major grid with sides that are multiples of 3, and refuses FlashInfer with GQA (one offset stride for q, k and v; maybe_recompute_cu_seqlens gains an explicit separate_qkv flag for that layout, which FP8 already used). Gemma 4's clipped linears are ReplicatedLinear plus the checkpoint's clamp bounds, keeping the .linear nesting. The MLP keeps transformers' tanh GELU with vLLM's GELUTanh (get_act_fn swaps it for the exact one on ROCm, for torch.compile, which this tower does not use). Parameter names match transformers: checkpoints and quantized towers load without a mapper.
• Host metadata: pixel_position_ids stay on the CPU (keep_on_cpu=True). prepare_encoder_metadata builds packed positions, cu_seqlens, max_seqlen, FlashInfer's sequence_lengths and the pooling indices without reading the GPU; the eager path, capture and replay share it and the same tower forward.
• gemma4_mm.py: images and video frames are encoded in one packed pass and split per item. Gone: resolution buckets, the memory-based _encoder_chunk (SDPA now bounds its memory itself, [Multimodal] Encoder CUDA graphs: eager fallback for packed TORCH_SDPA #1417), the LoRA padding path and the row-width graph axis. Vision LoRA counts are 9 x soft tokens, the patches the tower actually encodes, so the MM LoRA padding (and its per-image patch limit) is gone.
• Encoder CUDA graphs: buffers of budget * 9 patches. The attention grid launches work for every sequence slot and up to max_seqlen, so both are sized from what a replay can hold: budget // 35 sequences (every image or frame has at least half of the smallest max_soft_tokens; the manager does not bound a batch's frames) and the configured max_soft_tokens. Items outside those bounds (a per-request max_soft_tokens override) are marked eager_only ([Multimodal] Encoder CUDA graphs: eager fallback for packed TORCH_SDPA #1417) and run eager. With TORCH_SDPA (gfx1151's default), capture falls back to eager through [Multimodal] Encoder CUDA graphs: eager fallback for packed TORCH_SDPA #1417. Capture splits its dummy input into images of at most that size. Packed pixel values (eager and graph buffers) are in the model dtype.
• Derived models: DiffusionGemma and EmbeddingGemma2 use the same tower; Gemma 4 Unified (own embedder) moves the now host-side positions to the device and opts out of encoder graphs (it has no vision tower; with cudagraph_mm_encoder it used to fail at startup).

Design notes

• Replicated linears, no tensor or data parallelism: like the transformers tower they replace ([Bugfix] Convert Gemma4-MM ViT linear layers to vllm native impl vllm-project/vllm#43798). Gemma 4's vision towers are small (E2B/E4B: 16 layers, hidden 768; 26B/31B: 27 layers, hidden 1152), and TP/DP for them interacts with tower quantization and tower LoRA, which this PR keeps as is; DP is Support data-parallel multimodal encoder for Gemma-4 vllm-project/vllm#60314.
• Separate gate/up and q/k/v projections: each clipped linear clamps its input with its own checkpoint bounds, so the projections cannot share one GEMM (MergedColumnParallelLinear + GeluAndMul, QKVParallelLinear).
• RMSNorm: vLLM's casts to the model dtype before applying the weight, transformers' multiplies in fp32. The fp32 parity test cannot see that; the bf16 MMStar comparison below covers it.

Results (gfx1151, Strix Halo)

TTFT, one image per request, median of 15 (ms, 336x336 / 1024x800). Main runs SDPA with a mask regardless of the backend flag; gfx1151's default ViT backend is TORCH_SDPA, so "this PR, SDPA" is the like-for-like default.

Model	main	main + enc. graphs	this PR, SDPA	this PR, TRITON_ATTN	this PR, TRITON_ATTN + enc. graphs
Gemma 4 E4B	204.2 / 212.7	198.6 / 207.7	201.2 / 207.2	191.5 / 195.8	191.7 / 199.7
Gemma 4 26B-A4B AWQ	469.8 / 468.6	458.8 / 462.9	432.6 / 431.3	423.4 / 428.9	441.3 / 440.7

On gfx1151 the tower is compute bound (~200 ms per image on 26B), so padding a batch up to its budget costs more than the launches a graph saves; sizing the sequence slots and max_seqlen as above took 26B's graph TTFT from 464.8 to 439.3 ms. Graphs are expected to pay on GPUs where the tower is fast and launch-bound (not measured here: no CUDA hardware).

MMStar (the evaluation for this numerics change; tests/evals only has text GSM8K for Gemma 4), first 300 questions, greedy, main vs this PR (default backend, SDPA): same option letter in 293 (E2B), 296 (E4B), 293 (26B AWQ), 300 (31B AWQ) of 300; correct answers main / PR: E2B 68 / 67, E4B 159 / 160, 26B AWQ 54 / 55, 31B AWQ 106 / 106 (8-token answers). An earlier revision (exact GELU) also gave 300/300 on 31B AWQ and, with 64-token answers, E2B 112 / 117, E4B 166 / 164, 26B AWQ 177 / 177, 31B AWQ 189 / 190.

Encoder graph replay vs eager on E2B (TRITON_ATTN; mixed-size images, two videos, an image above the largest budget that falls back to eager): same shapes and per-item token counts, max difference 0.34 on values up to 6.4, the same as eager's own variance when an image is encoded alone or next to another (0.16).

Tests

• pytest tests/models/multimodal/test_gemma4_vision.py (new; tower vs transformers' padded tower, fp32, SDPA and Triton, with and without clipped linears, rtol/atol 1e-3): 4 passed.
• pytest tests/models/multimodal/processing/test_gemma4.py tests/models/multimodal/processing/test_gemma4_unified.py: passed (graph tests rewritten for token budgets; _encoder_chunk tests removed with the function).
• pytest tests/models/multimodal/pooling/test_embedding_gemma2.py -k contract: passed (borrowed methods updated).
• pytest tests/models/multimodal/generation/test_common.py -k "gemma4 or qwen2_5_vl": 7 passed (Qwen2.5-VL runs SDPA with host cu_seqlens on gfx1151).
• pytest tests/models/multimodal/generation/test_vit_cudagraph.py -k "gemma4 or qwen2_5_vl": 4 passed with gfx1151's default SDPA (capture falls back to eager with a warning); Gemma 4 also passed with TRITON_ATTN (graphs) on an earlier revision. Gemma 4 E4B with SDPA and cudagraph_mm_encoder: 196.3 / 204.6 ms TTFT.
• Encoder graph replay vs eager with TRITON_ATTN (script above) on this revision: same shapes and token counts, an image above the configured max_soft_tokens runs eager, and a 12-frame video replays with encoder_cudagraph_max_frames_per_batch=8.
• pytest tests/lora/test_gemma4_tp.py::test_gemma4_lora: passed on main and on this PR (teardown error: ray not installed).

Not covered / open points

• FLASH_ATTN and FLASHINFER (CUDA only) were not run. For FlashInfer, q, k and v are separate contiguous tensors, so every offset section uses the same per-token stride; this is passed through maybe_recompute_cu_seqlens(fp8_padded_hidden_size=...), which today only distinguishes the fused-QKV and FP8 layouts.
• Gemma 4 Unified inherits the encoder graph methods without a vision tower, as before.
• [Perf][Model][Multimodal] Batch Gemma 4 vision pooler and fuse vision/audio RMSNorm and 2D RoPE vllm-project/vllm#60478 patches transformers' vision tower (fused RMSNorm and 2D RoPE, batched pooler), which this PR removes; its fused 2D RoPE kernel (6-7x faster than two ApplyRotaryEmb calls on gfx1151) could move to layers/rotary_embedding/ in a follow-up. This PR supersedes [Perf][Model] Gemma 4: remove trailing image padding to enable FlashAttention in vision SDPA vllm-project/vllm#60577 (strip trailing padding). Their authors and issue [Feature]: Port Gemma4 vision encoder to MMEncoderAttention with FlashAttention support vllm-project/vllm#43178 are to be contacted before going upstream.

No open PR moves Gemma 4's vision tower to MMEncoderAttention (checked vllm-project#43440, vllm-project#45574, vllm-project#60478, vllm-project#60577, vllm-project#61365). AI-assisted (Claude Code); the submitter reviews every line.

🤖 Generated with Claude Code