vllm - 💡(How to fix) Fix [Bug] Gemma 4 31B crashes on k_eq_v full-attention layers (QKV split shape mismatch) [1 participants]

Description

Loading google/gemma-4-31b-it on vLLM 0.20.0 crashes immediately with a tensor shape mismatch during the QKV split.

Error

RuntimeError: split_with_sizes expects split_sizes to sum exactly to the size of dimension -1 (18432), but got split_sizes=[16384, 2048, 2048] summing to 20480

Root Cause

Gemma 4 31B has two attention layer types:

• Sliding layers — standard GQA with separate Q, K, V projections
• Full-attention (k_eq_v) layers — K and V share the same projection; the checkpoint has q_proj and k_proj on disk but no v_proj

Gemma4Attention routes all layers through QKVParallelLinear, then splits the output as [q_size, kv_size, kv_size]. For k_eq_v full-attention layers the actual packed tensor is only [q_size + kv_size] (no V column), so the split sums to q_size + 2×kv_size but the tensor is q_size + kv_size — causing the crash.

The _weight_iterator attempted to paper over this by duplicating k_proj weights as v_proj, but QKVParallelLinear still expects both slots to be present in the packed layout.

Affected models

• google/gemma-4-31b-it (first Gemma 4 model with mixed sliding/full-attention layers)
• Gemma 4 E4B and other small variants are not affected

Fix

PR #41253 fixes this by giving k_eq_v full-attention layers separate q_proj / k_proj ColumnParallelLinear modules and setting v = k directly in the forward pass, matching the checkpoint layout exactly. The _weight_iterator duplication hack is also removed.

Environment

• vLLM 0.20.0
• Bug also present in current main