@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70b84f0bcb

Fix Gemma 4 for upcoming Transformers version (#49797)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed Aug 10, 2026 at 08:48 UTC 70b84f0bcbb6d0a35b74b1035673a1c934089dbb
15 files changed +570 -106
tests/config/test_model_arch_config.py
+219
@@ -3,15 +3,20 @@
3 """Tests for ModelArchitectureConfig and its integration with ModelConfig."""
4
5 import json
6 +from copy import copy
7 from pathlib import Path
8 from types import SimpleNamespace
9 from typing import cast
10
11 import pytest
12 from transformers import PretrainedConfig
13 +from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
14
15 from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
16 +from vllm.config.model_arch import ModelArchitectureConfig
17 +from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
18 from vllm.transformers_utils.model_arch_config_convertor import (
19 + Gemma4ModelArchConfigConvertor,
20 ModelArchConfigConvertorBase,
21 )
22
@@ -182,6 +187,220 @@ def test_legacy_modelopt_config_without_producer_is_normalized():
187 assert convertor.get_quantization_config()["quant_method"] == "modelopt_fp4"
188
189
190 +def _layer(**overrides) -> ModelArchitectureConfig:
191 + fields = dict(
192 + architectures=["X"],
193 + model_type="x",
194 + text_model_type=None,
195 + hidden_size=64,
196 + total_num_hidden_layers=3,
197 + total_num_attention_heads=8,
198 + head_size=16,
199 + vocab_size=32,
200 + total_num_kv_heads=4,
201 + num_experts=0,
202 + num_experts_per_token=0,
203 + quantization_config=None,
204 + is_deepseek_mla=False,
205 + is_mm_prefix_lm=False,
206 + rswa_window=None,
207 + derived_max_model_len_and_key=(1.0, None),
208 + )
209 + return ModelArchitectureConfig(**(fields | overrides))
210 +
211 +
212 +def test_from_layers_collapses_to_the_largest_layer():
213 + """The whole-model config must size buffers for every layer."""
214 + arch = ModelArchitectureConfig.from_layers(
215 + [_layer(), _layer(), _layer(head_size=32, total_num_kv_heads=8)]
216 + )
217 +
218 + assert (arch.head_size, arch.total_num_kv_heads) == (32, 8)
219 + assert [arch[i].head_size for i in range(3)] == [16, 16, 32]
220 + assert [arch[i].total_num_kv_heads for i in range(3)] == [4, 4, 8]
221 + # A layer view is itself homogeneous, so reading from it cannot recurse.
222 + assert arch[0].per_layer_overrides is None
223 +
224 +
225 +def test_layer_view_follows_later_edits_to_the_whole_model_config():
226 + """`model_arch_config` is edited after construction in a few places."""
227 + arch = ModelArchitectureConfig.from_layers(
228 + [_layer(), _layer(), _layer(head_size=32)]
229 + )
230 +
231 + arch.is_mm_prefix_lm = True
232 +
233 + assert all(arch[i].is_mm_prefix_lm for i in range(3))
234 + assert [arch[i].head_size for i in range(3)] == [16, 16, 32]
235 +
236 +
237 +def test_uniform_layers_stay_homogeneous():
238 + """A checkpoint can vary attributes vLLM never reads."""
239 + arch = ModelArchitectureConfig.from_layers([_layer()] * 3)
240 +
241 + assert arch.per_layer_overrides is None
242 + assert arch[1] is arch
243 +
244 +
245 +def test_getitem_rejects_negative_indices():
246 + """Wrapping would only misbehave on heterogeneous models, so reject both."""
247 + homogeneous = _layer()
248 + heterogeneous = ModelArchitectureConfig.from_layers(
249 + [_layer(), _layer(), _layer(head_size=32)]
250 + )
251 +
252 + for arch in (homogeneous, heterogeneous):
253 + with pytest.raises(IndexError):
254 + arch[-1]
255 +
256 +
257 +@pytest.mark.parametrize(
258 + "varying",
259 + [
260 + # `bool` is an `int`; collapsing this one with `max` would make `use_mla`
261 + # true model wide and silently discard every per-layer KV head count.
262 + {"is_deepseek_mla": True},
263 + {"quantization_config": {"quant_method": "fp8"}},
264 + {"rswa_window": 512},
265 + ],
266 +)
267 +def test_from_layers_rejects_fields_with_no_whole_model_value(varying: dict):
268 + with pytest.raises(ValueError, match="varies across layers"):
269 + ModelArchitectureConfig.from_layers([_layer(), _layer(), _layer(**varying)])
270 +
271 +
272 +def test_from_layers_rejects_a_layer_count_mismatch():
273 + """Draft configs can inherit a target's per-layer spec."""
274 + with pytest.raises(ValueError, match="per-layer configs"):
275 + ModelArchitectureConfig.from_layers([_layer(), _layer(head_size=32)])
276 +
277 +
278 +def _gemma4_text_config(**overrides) -> Gemma4TextConfig:
279 + """A six layer Gemma4 whose last layer is wider than the rest."""
280 + fields = dict(
281 + num_hidden_layers=6,
282 + hidden_size=64,
283 + num_attention_heads=8,
284 + num_key_value_heads=4,
285 + head_dim=16,
286 + global_head_dim=32,
287 + # Transformers forces the last layer to full attention.
288 + layer_types=["sliding_attention"] * 5 + ["full_attention"],
289 + )
290 + return Gemma4TextConfig(**(fields | overrides))
291 +
292 +
293 +def test_gemma4_head_dims_vary_by_layer_type():
294 + """Gemma4's full attention layers are wider than its sliding ones.
295 +
296 + Transformers >= 5.15.0 says so in the config; this exercises the convertor
297 + building the same per-layer view from the flat attributes used before that.
298 + """
299 + text_config = _gemma4_text_config(
300 + num_global_key_value_heads=8, attention_k_eq_v=True
301 + )
302 +
303 + arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert()
304 +
305 + assert (arch.head_size, arch.total_num_kv_heads) == (32, 8)
306 + assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32]
307 + assert [arch[i].total_num_kv_heads for i in range(6)] == [4] * 5 + [8]
308 + # The model files resolve each layer through the same helper, so the KV cache
309 + # vLLM allocates and the projections the model builds cannot disagree.
310 + assert [gemma4_layer_config(text_config, i).head_dim for i in range(6)] == [
311 + arch[i].head_size for i in range(6)
312 + ]
313 +
314 +
315 +def test_gemma4_without_global_kv_heads():
316 + """`num_global_key_value_heads` defaults to `None`, not to a head count."""
317 + text_config = _gemma4_text_config()
318 +
319 + arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert()
320 +
321 + assert arch.total_num_kv_heads == 4
322 + assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32]
323 +
324 +
325 +def test_gemma4_layer_count_comes_from_num_hidden_layers():
326 + """`dummy_hf_overrides` shrinks the stack but leaves `layer_types` long."""
327 + text_config = _gemma4_text_config()
328 + text_config.num_hidden_layers = 3
329 +
330 + arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert()
331 +
332 + assert arch.total_num_hidden_layers == 3
333 + # Only sliding layers survive the truncation, so nothing varies.
334 + assert arch.per_layer_overrides is None
335 +
336 +
337 +def test_gemma4_uniform_head_dims_are_homogeneous():
338 + text_config = _gemma4_text_config(global_head_dim=16)
339 +
340 + arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert()
341 +
342 + assert arch.per_layer_overrides is None
343 + assert arch[3] is arch
344 +
345 +
346 +class _HeterogeneousConfig(PretrainedConfig):
347 + """A stand-in for the Transformers >= 5.15.0 heterogeneous config API.
348 +
349 + No released Transformers has it, so the default per-layer seam has no other
350 + way to be exercised. Mirrors the parts vLLM uses: per-layer configs are
351 + shallow copies with the varying attributes applied and heterogeneity
352 + stripped, so they do not recurse.
353 + """
354 +
355 + is_heterogeneous = True
356 +
357 + def __init__(self, per_layer: dict[str, list], **kwargs):
358 + super().__init__(**kwargs)
359 + self._per_layer = per_layer
360 +
361 + @property
362 + def per_layer_config(self) -> list[PretrainedConfig]:
363 + layers = []
364 + for i in range(self.num_hidden_layers):
365 + layer = copy(self)
366 + layer.is_heterogeneous = False
367 + for name, values in self._per_layer.items():
368 + setattr(layer, name, values[i])
369 + layers.append(layer)
370 + return layers
371 +
372 +
373 +def test_transformers_heterogeneous_config_is_resolved_per_layer():
374 + hf_config = _HeterogeneousConfig(
375 + per_layer={"head_dim": [16, 16, 32], "num_key_value_heads": [4, 4, 8]},
376 + num_hidden_layers=3,
377 + hidden_size=64,
378 + num_attention_heads=8,
379 + )
380 +
381 + arch = ModelArchConfigConvertorBase(hf_config, hf_config).convert()
382 +
383 + assert (arch.head_size, arch.total_num_kv_heads) == (32, 8)
384 + assert [arch[i].head_size for i in range(3)] == [16, 16, 32]
385 + assert [arch[i].total_num_kv_heads for i in range(3)] == [4, 4, 8]
386 +
387 +
388 +def test_heterogeneous_config_varying_nothing_vllm_reads():
389 + """Transformers prunes nothing for us: the diff has to notice."""
390 + hf_config = _HeterogeneousConfig(
391 + per_layer={"bos_token_id": [1, 2, 3]},
392 + num_hidden_layers=3,
393 + hidden_size=64,
394 + num_attention_heads=8,
395 + head_dim=16,
396 + )
397 +
398 + arch = ModelArchConfigConvertorBase(hf_config, hf_config).convert()
399 +
400 + assert arch.per_layer_overrides is None
401 + assert arch[2] is arch
402 +
403 +
404 @pytest.mark.parametrize("model", BASE_MODELS_TO_TEST)
405 def test_base_model_arch_config(model: str):
406 """Test model architecture config for base models."""
tests/models/transformers/fusers/test_linear.py
+62
@@ -15,6 +15,8 @@ from vllm.model_executor.models.transformers.fusers import (
15 GLUFuser,
16 PackedQKVFuser,
17 QKVFuser,
18 + packed_qkv,
19 + qkv,
20 )
21
22
@@ -640,3 +642,63 @@ def test_act_and_mul_derived_from_module(default_vllm_config):
642 assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None
643 with pytest.raises(ValueError, match="No AndMul equivalent"):
644 GLUFuser._get_act_and_mul(nn.Dropout())
645 +
646 +
647 +def _wider_model_config(head_dim: int) -> SimpleNamespace:
648 + """A model whose global head size is twice `head_dim`, as a wider layer
649 + elsewhere in a heterogeneous checkpoint would make it."""
650 + return SimpleNamespace(
651 + model_config=SimpleNamespace(get_head_size=lambda: 2 * head_dim),
652 + quant_config=None,
653 + )
654 +
655 +
656 +@pytest.mark.parametrize(
657 + "cls, fuser_module", [(FakeAttention, qkv), (PackedQKVAttention, packed_qkv)]
658 +)
659 +def test_head_counts_come_from_the_module_not_the_model(cls, fuser_module, monkeypatch):
660 + """A layer narrower than the model-wide head size must not be miscounted.
661 +
662 + On a heterogeneous checkpoint (Gemma 4) the model-wide head size is the
663 + largest across layers, so deriving `total_num_heads = out_features //
664 + head_size` from it undercounts heads on a narrower layer. The widths still
665 + add up, so nothing raises below TP=4: the layer is just sharded wrong.
666 + """
667 + head_dim, heads, kv_heads = 8, 8, 4
668 + vllm_config = _wider_model_config(head_dim)
669 + with torch.device("meta"):
670 + module = cls(hidden=32, head_dim=head_dim, heads=heads, kv_heads=kv_heads)
671 +
672 + # Both replacements shard, so they need a TP group; only the head counts
673 + # the fuser derives are under test here.
674 + captured = {}
675 + monkeypatch.setattr(
676 + fuser_module,
677 + "QKVParallelLinear",
678 + lambda **kwargs: captured.update(kwargs) or nn.Identity(),
679 + )
680 + monkeypatch.setattr(
681 + fuser_module, "replace_linear_class", lambda *a, **kw: nn.Identity()
682 + )
683 + fuser = get_fuser(module)
684 + assert fuser is not None and fuser.validate(module, vllm_config)
685 + fuser.update_attrs(module, "model.layers.0.self_attn", vllm_config)
686 +
687 + assert captured["head_size"] == head_dim
688 + assert captured["total_num_heads"] == heads
689 + assert captured["total_num_kv_heads"] == kv_heads
690 +
691 +
692 +def test_validate_accepts_a_layer_the_model_wide_head_size_would_reject():
693 + """`validate` gates fusion, so a wrong head size silently disables it."""
694 + head_dim, heads, kv_heads = 8, 8, 3
695 + vllm_config = _wider_model_config(head_dim)
696 + with torch.device("meta"):
697 + module = FakeAttention(
698 + hidden=32, head_dim=head_dim, heads=heads, kv_heads=kv_heads
699 + )
700 +
701 + # kv width is 24, not a multiple of the model-wide 16, but is of this
702 + # layer's 8.
703 + fuser = get_fuser(module)
704 + assert fuser is not None and fuser.validate(module, vllm_config)
vllm/config/model.py
+27 -10
@@ -886,9 +886,7 @@ class ModelConfig:
886 )
887 return supports_mm
888
889 - def get_model_arch_config(
890 - self,
891 - ) -> ModelArchitectureConfig:
889 + def get_model_arch_config(self) -> ModelArchitectureConfig:
890 convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get(
891 self.hf_config.model_type, ModelArchConfigConvertorBase
892 )
@@ -982,8 +980,8 @@ class ModelConfig:
980 """Determine which Transformers modeling backend class will be used if
981 `model_impl` is set to `transformers` or `auto`."""
982 cls = "Transformers"
985 - # If 'hf_config != hf_text_config' it's a nested config, i.e. multimodal
986 - cls += "MultiModal" if self.hf_config != self.hf_text_config else ""
983 + # If 'hf_config is not hf_text_config' it's a nested config, i.e. multimodal
984 + cls += "MultiModal" if self.hf_config is not self.hf_text_config else ""
985 cls += "MoE" if self.is_moe else ""
986 # Check if the architecture we're wrapping has defaults
987 runner = None
@@ -1473,21 +1471,40 @@ class ModelConfig:
1471 """Returns the total number of KV heads."""
1472 return self.model_arch_config.total_num_kv_heads
1473
1476 - def get_num_kv_heads(self, parallel_config: ParallelConfig) -> int:
1477 - """Returns the number of KV heads per GPU."""
1474 + def get_num_kv_heads(
1475 + self,
1476 + parallel_config: ParallelConfig,
1477 + arch_config: ModelArchitectureConfig | None = None,
1478 + ) -> int:
1479 + """Returns the number of KV heads per GPU.
1480 +
1481 + Pass ``arch_config`` (from ``model_arch_config[layer_idx]``) to size a
1482 + single layer of a heterogeneous model rather than the model as a whole.
1483 + """
1484 if self.use_mla:
1485 # When using MLA during decode it becomes MQA
1486 return 1
1487
1482 - total_num_kv_heads = self.get_total_num_kv_heads()
1488 + arch_config = arch_config or self.model_arch_config
1489 + total_num_kv_heads = arch_config.total_num_kv_heads
1490 # If tensor parallelism is used, we divide the number of KV heads by
1491 # the tensor parallel size. We will replicate the KV heads in the
1492 # case where the number of KV heads is smaller than the tensor
1493 # parallel size so each GPU has at least one KV head.
1494 return max(1, total_num_kv_heads // parallel_config.tensor_parallel_size)
1495
1489 - def get_num_attention_heads(self, parallel_config: ParallelConfig) -> int:
1490 - num_heads = self.model_arch_config.total_num_attention_heads
1496 + def get_num_attention_heads(
1497 + self,
1498 + parallel_config: ParallelConfig,
1499 + arch_config: ModelArchitectureConfig | None = None,
1500 + ) -> int:
1501 + """Returns the number of attention heads per GPU.
1502 +
1503 + Pass ``arch_config`` (from ``model_arch_config[layer_idx]``) to size a
1504 + single layer of a heterogeneous model rather than the model as a whole.
1505 + """
1506 + arch_config = arch_config or self.model_arch_config
1507 + num_heads = arch_config.total_num_attention_heads
1508 return num_heads // parallel_config.tensor_parallel_size
1509
1510 def get_num_experts(self) -> int:
vllm/config/model_arch.py
+77
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from copy import copy
4 +from dataclasses import fields as dataclass_fields
5 from typing import Any
6
7 from pydantic import ConfigDict
@@ -64,3 +66,78 @@ class ModelArchitectureConfig:
66
67 derived_max_model_len_and_key: tuple[float, str | None]
68 """Derived maximum model length and key from the hf config."""
69 +
70 + per_layer_overrides: "list[dict[str, Any]] | None" = None
71 + """Per-layer values for the fields that vary, `None` unless some field does.
72 +
73 + One dict per layer, holding only the fields whose value differs from the
74 + whole-model value above. Everything else is read from the whole-model config,
75 + so later edits to it are visible through `self[layer_idx]`."""
76 +
77 + def __getitem__(self, layer_idx: int) -> "ModelArchitectureConfig":
78 + """ModelArchitectureConfig for a specific layer.
79 +
80 + Returns `self` when no field varies by layer, so callers never need to
81 + branch on heterogeneity. Mirrors `PreTrainedConfig.per_layer_config[i]`.
82 + """
83 + if layer_idx < 0:
84 + # Homogeneous configs would ignore a negative index and heterogeneous
85 + # ones would wrap, so an off-by-one would only show up on the models
86 + # this exists for. Reject it for both.
87 + raise IndexError(f"layer index must not be negative, got {layer_idx}")
88 + if self.per_layer_overrides is None:
89 + return self
90 + layer = copy(self)
91 + object.__setattr__(layer, "per_layer_overrides", None)
92 + for name, value in self.per_layer_overrides[layer_idx].items():
93 + object.__setattr__(layer, name, value)
94 + return layer
95 +
96 + @classmethod
97 + def from_layers(
98 + cls, layers: "list[ModelArchitectureConfig]"
99 + ) -> "ModelArchitectureConfig":
100 + """Whole-model config for a checkpoint whose layers differ.
101 +
102 + Fields that agree across layers are taken as they are. Fields that differ
103 + are collapsed with `max`, so buffers are sized for the largest layer, and
104 + the differing values are kept per layer. No field is named here: which
105 + ones vary is whatever the checkpoint says.
106 + """
107 + if not layers:
108 + raise ValueError("a model must have at least one layer")
109 +
110 + merged: dict[str, Any] = {}
111 + overrides: list[dict[str, Any]] = [{} for _ in layers]
112 + for f in dataclass_fields(cls):
113 + if f.name == "per_layer_overrides":
114 + continue
115 + values = [getattr(layer, f.name) for layer in layers]
116 + if all(value == values[0] for value in values):
117 + merged[f.name] = values[0]
118 + continue
119 + # `bool` is an `int`, so an exact type check is what keeps a varying
120 + # flag from collapsing to `any`. `is_deepseek_mla` doing that would
121 + # make `use_mla` true model wide, and `get_num_kv_heads` then returns
122 + # 1 for every layer, discarding the overrides built here.
123 + if not all(type(value) in (int, float) for value in values):
124 + raise ValueError(
125 + f"{f.name!r} varies across layers and has no whole-model "
126 + f"value: {sorted(set(map(repr, values)))}. Only numeric "
127 + f"fields collapse (with `max`, to size buffers for the "
128 + f"largest layer); give this one an explicit rule in "
129 + f"ModelArchitectureConfig.from_layers."
130 + )
131 + merged[f.name] = max(values)
132 + for override, value in zip(overrides, values):
133 + if value != merged[f.name]:
134 + override[f.name] = value
135 +
136 + if len(layers) != merged["total_num_hidden_layers"]:
137 + raise ValueError(
138 + f"got {len(layers)} per-layer configs for a model with "
139 + f"{merged['total_num_hidden_layers']} layers"
140 + )
141 + # A checkpoint can be heterogeneous over attributes vLLM never reads, in
142 + # which case there is nothing to keep the layers apart for.
143 + return cls(**merged, per_layer_overrides=overrides if any(overrides) else None)
vllm/model_executor/model_loader/weight_utils.py
+2 -6
@@ -260,12 +260,8 @@ def get_quant_config(
260 and hf_quant_config.get("quant_method") == "compressed-tensors"
261 and "config_groups" in hf_quant_config
262 ):
263 - if hf_text_config is not None:
264 - n_heads = getattr(hf_text_config, "num_attention_heads", None)
265 - n_kv_heads = getattr(hf_text_config, "num_key_value_heads", None)
266 - else:
267 - n_heads = getattr(model_config.hf_config, "num_attention_heads", None)
268 - n_kv_heads = getattr(model_config.hf_config, "num_key_value_heads", None)
263 + n_heads = model_config.model_arch_config.total_num_attention_heads
264 + n_kv_heads = model_config.model_arch_config.total_num_kv_heads
265
266 hf_quant_config["total_num_heads"] = n_heads
267 hf_quant_config["total_num_kv_heads"] = (
vllm/model_executor/models/config.py
+20 -21
@@ -200,26 +200,29 @@ class Gemma4Config(VerifyAndUpdateConfig):
200 def verify_and_update_config(vllm_config: "VllmConfig") -> None:
201 """Configure attention for heterogeneous head dimensions.
202
203 - Gemma4 uses different head dimensions for sliding window
204 - (head_dim) vs full attention (global_head_dim) layers. The
205 - default FA3 on Hopper cannot handle head_dim > 256, which
206 - causes mixed backend selection and numerical divergence.
207 -
208 - When FA4 is available we force it for ALL layers, giving a
209 - uniform kernel path and avoiding the mixed FA3+FA4 penalty.
210 - When FA4 is not available we fall back to Triton.
203 + Gemma4 uses different head dimensions for sliding window vs full attention
204 + layers. The default FA3 on Hopper cannot handle head_dim > 256, which causes
205 + mixed backend selection and numerical divergence.
206 +
207 + When FA4 is available we force it for ALL layers, giving a uniform kernel path
208 + and avoiding the mixed FA3+FA4 penalty. When FA4 is not available we fall back
209 + to Triton.
210 """
212 - hf_text_config = vllm_config.model_config.hf_text_config
213 - head_dim = getattr(hf_text_config, "head_dim", None)
214 - global_head_dim = getattr(hf_text_config, "global_head_dim", None)
211 + model_config = vllm_config.model_config
212 + arch_config = model_config.model_arch_config
213 + layer_types = getattr(model_config.hf_text_config, "layer_types", None) or []
214 + head_dims = {
215 + layer_types[i]: arch_config[i].head_size
216 + for i in range(min(arch_config.total_num_hidden_layers, len(layer_types)))
217 + }
218
216 - if head_dim is None or global_head_dim is None or head_dim == global_head_dim:
219 + if len(set(head_dims.values())) <= 1:
220 return
221
222 from vllm.v1.attention.backends.fa_utils import is_fa_version_supported
223 from vllm.v1.attention.backends.registry import AttentionBackendEnum
224
222 - max_head_dim = max(head_dim, global_head_dim)
225 + max_head_dim = max(head_dims.values())
226
227 if is_fa_version_supported(4) and max_head_dim <= 512:
228 if (
@@ -229,20 +232,16 @@ class Gemma4Config(VerifyAndUpdateConfig):
232 ):
233 vllm_config.attention_config.flash_attn_version = 4
234 logger.info(
232 - "Gemma4 model has heterogeneous head dimensions "
233 - "(head_dim=%d, global_head_dim=%d). Using FA4 for "
235 + "Gemma4 model has heterogeneous head dimensions %s. Using FA4 for "
236 "all layers to avoid mixed FA3/FA4 penalty.",
235 - head_dim,
236 - global_head_dim,
237 + head_dims,
238 )
239 elif vllm_config.attention_config.backend is None:
240 vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
241 logger.info(
242 "Gemma4 model has heterogeneous head dimensions "
242 - "(head_dim=%d, global_head_dim=%d). FA4 not available, "
243 - "forcing TRITON_ATTN backend.",
244 - head_dim,
245 - global_head_dim,
243 + "%s. FA4 not available, forcing TRITON_ATTN backend.",
244 + head_dims,
245 )
246
247
vllm/model_executor/models/gemma4.py
+4 -24
@@ -63,6 +63,7 @@ from vllm.model_executor.model_loader.weight_utils import (
63 )
64 from vllm.platforms import current_platform
65 from vllm.sequence import IntermediateTensors
66 +from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
67 from vllm.triton_utils import tl, triton
68 from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata
69
@@ -380,7 +381,6 @@ class Gemma4Attention(nn.Module):
381 num_kv_heads: int,
382 head_dim: int,
383 max_position_embeddings: int,
383 - use_k_eq_v: bool = False,
384 cache_config: CacheConfig | None = None,
385 quant_config: QuantizationConfig | None = None,
386 attn_logits_soft_cap: float | None = None,
@@ -389,7 +389,6 @@ class Gemma4Attention(nn.Module):
389 super().__init__()
390 self.config = config
391 self.hidden_size = hidden_size
392 - self.use_k_eq_v = use_k_eq_v
392
393 tp_size = get_tensor_model_parallel_world_size()
394 self.tp_rank = get_tensor_model_parallel_rank()
@@ -570,27 +569,9 @@ class Gemma4DecoderLayer(nn.Module):
569 self.layer_idx = layer_idx
570
571 # Gemma4 uses different head dimensions for sliding vs full attention
573 - layer_type = config.layer_types[layer_idx]
574 - self.is_full_attention = layer_type == "full_attention"
575 - if self.is_full_attention:
576 - head_dim = getattr(config, "global_head_dim", config.head_dim)
577 - else:
578 - head_dim = config.head_dim
579 -
580 - # Determine if this full-attention layer uses k_eq_v
581 - # (laptop variant: no v_proj, K reused as V on full attention layers)
582 - use_k_eq_v = self.is_full_attention and getattr(
583 - config, "attention_k_eq_v", False
584 - )
585 -
586 - # For k_eq_v full-attention layers, use num_global_key_value_heads
587 - # as the KV head count when k_eq_v is enabled.
588 - if use_k_eq_v:
589 - num_kv_heads = getattr(
590 - config, "num_global_key_value_heads", config.num_key_value_heads
591 - )
592 - else:
593 - num_kv_heads = config.num_key_value_heads
572 + layer_config = gemma4_layer_config(config, layer_idx)
573 + head_dim = layer_config.head_dim
574 + num_kv_heads = layer_config.num_key_value_heads
575
576 self.self_attn = Gemma4Attention(
577 config=config,
@@ -599,7 +580,6 @@ class Gemma4DecoderLayer(nn.Module):
580 num_kv_heads=num_kv_heads,
581 head_dim=head_dim,
582 max_position_embeddings=config.max_position_embeddings,
602 - use_k_eq_v=use_k_eq_v,
583 cache_config=cache_config,
584 quant_config=quant_config,
585 attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None),
vllm/model_executor/models/gemma4_dspark.py
+9 -13
@@ -20,6 +20,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
20 VocabParallelEmbedding,
21 )
22 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23 +from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
24
25 from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer
26 from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal
@@ -38,18 +39,14 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
39 prefix: str,
40 ) -> None:
41 layer_idx = extract_layer_index(prefix)
41 - is_full = config.layer_types[layer_idx] == "full_attention"
42 - head_dim = (
43 - getattr(config, "global_head_dim", config.head_dim)
44 - if is_full
45 - else config.head_dim
46 - )
47 - use_k_eq_v = is_full and getattr(config, "attention_k_eq_v", False)
48 - num_kv_heads = (
49 - getattr(config, "num_global_key_value_heads", config.num_key_value_heads)
50 - if use_k_eq_v
51 - else config.num_key_value_heads
42 + layer_type = config.layer_types[layer_idx]
43 + self.use_k_eq_v = layer_type == "full_attention" and getattr(
44 + config, "attention_k_eq_v", False
45 )
46 +
47 + layer_config = gemma4_layer_config(config, layer_idx)
48 + head_dim = layer_config.head_dim
49 + num_kv_heads = layer_config.num_key_value_heads
50 super().__init__(
51 config=config,
52 hidden_size=config.hidden_size,
@@ -64,7 +61,6 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
61 )
62 self.is_kv_shared_layer = False
63 self.causal = _dflash_layer_causal(config, layer_idx)
67 - self.use_k_eq_v = use_k_eq_v
64 self.kv_size = self.num_kv_heads * self.head_dim
65 attn_bias = getattr(config, "attention_bias", False)
66 self.k_proj = ColumnParallelLinear(
@@ -76,7 +72,7 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
72 )
73 self.v_proj = (
74 None
79 - if use_k_eq_v
75 + if self.use_k_eq_v
76 else ColumnParallelLinear(
77 config.hidden_size,
78 self.total_num_kv_heads * self.head_dim,
vllm/model_executor/models/gemma4_mtp.py
+4 -15
@@ -45,6 +45,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
45 VocabParallelEmbedding,
46 )
47 from vllm.sequence import IntermediateTensors
48 +from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
49
50 from .gemma4 import Gemma4MLP, _get_text_config
51 from .utils import (
@@ -271,21 +272,9 @@ class Gemma4MTPDecoderLayer(nn.Module):
272 self.hidden_size = config.hidden_size
273
274 layer_idx = extract_layer_index(prefix)
274 - layer_type = config.layer_types[layer_idx]
275 - is_full_attention = layer_type == "full_attention"
276 - head_dim = (
277 - getattr(config, "global_head_dim", config.head_dim)
278 - if is_full_attention
279 - else config.head_dim
280 - )
281 -
282 - use_k_eq_v = is_full_attention and getattr(config, "attention_k_eq_v", False)
283 - if use_k_eq_v:
284 - num_kv_heads = getattr(
285 - config, "num_global_key_value_heads", config.num_key_value_heads
286 - )
287 - else:
288 - num_kv_heads = config.num_key_value_heads
275 + layer_config = gemma4_layer_config(config, layer_idx)
276 + head_dim = layer_config.head_dim
277 + num_kv_heads = layer_config.num_key_value_heads
278
279 self.self_attn = Gemma4MTPAttention(
280 config=config,
vllm/model_executor/models/transformers/base.py
+14 -6
@@ -522,9 +522,6 @@ class Base(
522 if qk_head_dim := qk_nope_head_dim + qk_rope_head_dim:
523 self.model_config.model_arch_config.head_size = qk_head_dim
524
525 - num_heads = self.model_config.get_num_attention_heads(self.parallel_config)
526 - head_size = self.model_config.get_head_size()
527 - num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
525 logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None)
526
527 pp_rank = self.pp_group.rank_in_group
@@ -532,11 +529,22 @@ class Base(
529 start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size)
530
531 for i in range(start, end):
532 + # `[i]` is the whole-model config unless the checkpoint is
533 + # heterogeneous, in which case it is this layer's own geometry.
534 + arch_config = self.model_config.model_arch_config[i]
535 + num_heads = self.model_config.get_num_attention_heads(
536 + self.parallel_config, arch_config
537 + )
538 + head_size = arch_config.head_size
539 + # Default to Llama scale, maybe updated in vllm_attention_forward
540 + scale = head_size**-0.5
541 + num_kv_heads = self.model_config.get_num_kv_heads(
542 + self.parallel_config, arch_config
543 + )
544 +
545 kwargs = dict(
546 num_heads=num_heads,
537 - # NOTE: We use Llama scale as default, if it's set by
538 - # Transformers, it's updated in vllm_attention_forward
539 - scale=head_size**-0.5,
547 + scale=scale,
548 cache_config=self.cache_config,
549 quant_config=self.quant_config,
550 prefix=f"{i}.attn",
vllm/model_executor/models/transformers/fusers/base.py
+11
@@ -57,6 +57,17 @@ class BaseFuser(ABC):
57 return {}
58
59
60 +def fused_head_size(module: nn.Module, vllm_config: "VllmConfig") -> int:
61 + """The head size of `module`, which head counts are derived from.
62 +
63 + Prefer the module's own `head_dim`, which Transformers sets per instance:
64 + the model-wide value is the largest head size across layers, so on a
65 + heterogeneous checkpoint it would divide out the wrong number of heads, and
66 + it is the text head size, so it does not describe a vision tower at all.
67 + """
68 + return getattr(module, "head_dim", None) or vllm_config.model_config.get_head_size()
69 +
70 +
71 def local_output_sizes(merged_name: str) -> str:
72 """Source for the per-rank widths of the merged linear `self.<merged_name>`."""
73 merged = f"self.{merged_name}"
vllm/model_executor/models/transformers/fusers/packed_qkv.py
+3 -2
@@ -12,6 +12,7 @@ from vllm.logger import init_logger
12 from vllm.model_executor.layers.linear import QKVParallelLinear
13 from vllm.model_executor.models.transformers.fusers.base import (
14 RewriteFuser,
15 + fused_head_size,
16 local_output_sizes,
17 )
18 from vllm.model_executor.models.transformers.fx_utils import (
@@ -122,7 +123,7 @@ class PackedQKVFuser(RewriteFuser):
123
124 def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
125 """Shapes must be compatible with a head-sharded packed GEMM."""
125 - head_size = vllm_config.model_config.get_head_size()
126 + head_size = fused_head_size(module, vllm_config)
127 qkv = module.get_submodule(self.qkv_name)
128 compatible = (
129 self.q_size % head_size == 0
@@ -137,7 +138,7 @@ class PackedQKVFuser(RewriteFuser):
138 self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
139 ) -> None:
140 quant_config = vllm_config.quant_config
140 - head_size = vllm_config.model_config.get_head_size()
141 + head_size = fused_head_size(module, vllm_config)
142 qkv_prefix = maybe_prefix(prefix, self.qkv_name)
143 qkv = module.get_submodule(self.qkv_name)
144 merged = QKVParallelLinear(
vllm/model_executor/models/transformers/fusers/qkv.py
+3 -2
@@ -12,6 +12,7 @@ from vllm.logger import init_logger
12 from vllm.model_executor.layers.linear import QKVParallelLinear
13 from vllm.model_executor.models.transformers.fusers.base import (
14 StackedFuser,
15 + fused_head_size,
16 local_output_sizes,
17 )
18 from vllm.model_executor.models.transformers.fx_utils import (
@@ -158,7 +159,7 @@ class QKVFuser(StackedFuser):
159 q = module.get_submodule(self.q_name)
160 k = module.get_submodule(self.k_name)
161 v = module.get_submodule(self.v_name)
161 - head_size = vllm_config.model_config.get_head_size()
162 + head_size = fused_head_size(module, vllm_config)
163 compatible = (
164 q.in_features == k.in_features == v.in_features
165 and len({proj.bias is None for proj in (q, k, v)}) == 1
@@ -174,7 +175,7 @@ class QKVFuser(StackedFuser):
175 self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
176 ) -> None:
177 quant_config = vllm_config.quant_config
177 - head_size = vllm_config.model_config.get_head_size()
178 + head_size = fused_head_size(module, vllm_config)
179 q = module.get_submodule(self.q_name)
180 k = module.get_submodule(self.k_name)
181 merged = QKVParallelLinear(
vllm/transformers_utils/configs/gemma4.py new
+34
@@ -0,0 +1,34 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Per-layer config resolution for Gemma4 and its variants."""
4 +
5 +from copy import copy
6 +
7 +from transformers import PretrainedConfig
8 +
9 +
10 +def gemma4_layer_config(
11 + text_config: PretrainedConfig, layer_idx: int
12 +) -> PretrainedConfig:
13 + """The Gemma4 text config as it applies to one layer.
14 +
15 + Gemma4 uses a larger head dimension on its full attention layers than on its
16 + sliding ones, and with `attention_k_eq_v` it uses more KV heads there too.
17 + Transformers >= 5.15.0 says so in the config itself; before that the values
18 + are flat attributes that `layer_types` picks between, so resolve them here.
19 + The result is homogeneous either way, so callers read `head_dim` and
20 + `num_key_value_heads` off it without caring which version they are on.
21 + """
22 + if getattr(text_config, "is_heterogeneous", False):
23 + return text_config.per_layer_config[layer_idx]
24 +
25 + layer = copy(text_config)
26 + if text_config.layer_types[layer_idx] == "full_attention":
27 + global_head_dim = getattr(text_config, "global_head_dim", None)
28 + layer.head_dim = global_head_dim or text_config.head_dim
29 + if getattr(text_config, "attention_k_eq_v", False):
30 + global_kv_heads = getattr(text_config, "num_global_key_value_heads", None)
31 + layer.num_key_value_heads = (
32 + global_kv_heads or text_config.num_key_value_heads
33 + )
34 + return layer
vllm/transformers_utils/model_arch_config_convertor.py
+81 -7
@@ -17,6 +17,7 @@ from vllm.transformers_utils.config import (
17 ConfigFormat,
18 get_safetensors_params_metadata,
19 )
20 +from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
21 from vllm.utils.torch_utils import common_broadcastable_dtype
22
23 logger = init_logger(__name__)
@@ -27,6 +28,37 @@ class ModelArchConfigConvertorBase:
28 self.hf_config = hf_config
29 self.hf_text_config = hf_text_config
30
31 + def get_per_layer_hf_configs(
32 + self,
33 + ) -> list[tuple[PretrainedConfig, PretrainedConfig]] | None:
34 + """`(hf_config, hf_text_config)` per layer, or `None` if homogeneous.
35 +
36 + This is the only place that decides whether a checkpoint is heterogeneous
37 + and how many layers it has, so a convertor can declare variation that
38 + Transformers does not express (see `Gemma4ModelArchConfigConvertor`).
39 +
40 + Each pair is converted independently and the results are diffed, so a
41 + convertor that overrides this must not also collapse a varying field in
42 + its getters: the collapse would then be applied per layer and every layer
43 + would report the same value.
44 + """
45 + config_het = getattr(self.hf_config, "is_heterogeneous", False)
46 + text_het = getattr(self.hf_text_config, "is_heterogeneous", False)
47 + if not (config_het or text_het):
48 + return None
49 + # The text config decides the layer count when it is the heterogeneous one:
50 + # its stack is what vLLM builds attention layers for.
51 + source = self.hf_text_config if text_het else self.hf_config
52 + return [
53 + (
54 + self.hf_config.per_layer_config[i] if config_het else self.hf_config,
55 + self.hf_text_config.per_layer_config[i]
56 + if text_het
57 + else self.hf_text_config,
58 + )
59 + for i in range(len(source.per_layer_config))
60 + ]
61 +
62 def get_architectures(self) -> list[str]:
63 # Sometimes we get here from `vllm_config.with_hf_config(text_config)` where
64 # `text_config` is a sub-config from a multi-modal model. If this is the case,
@@ -382,6 +414,35 @@ class ModelArchConfigConvertorBase:
414 return derived_max_model_len, max_len_key
415
416 def convert(self, supports_multimodal: bool = True) -> ModelArchitectureConfig:
417 + if (per_layer := self.get_per_layer_hf_configs()) is None:
418 + return self.convert_layer(supports_multimodal)
419 +
420 + if self.is_deepseek_mla():
421 + raise NotImplementedError(
422 + "Heterogeneous MLA models are not supported: `get_head_size` "
423 + "patches the config in place for them, which a per-layer "
424 + "conversion would apply to a throwaway copy."
425 + )
426 + # Convert each layer and let the result work out which fields differ.
427 + # Reading a varying attribute off the global config would raise, so the
428 + # whole-model config has to be built up from the layers.
429 + return ModelArchitectureConfig.from_layers(
430 + [
431 + type(self)(*configs).convert_layer(supports_multimodal)
432 + for configs in per_layer
433 + ]
434 + )
435 +
436 + def convert_layer(
437 + self, supports_multimodal: bool = True
438 + ) -> ModelArchitectureConfig:
439 + """Convert one homogeneous config, without resolving per-layer values.
440 +
441 + `convert` calls this once per layer, so it must not recurse back into
442 + per-layer resolution: the configs it receives are already layer-specific,
443 + and they can still carry the attributes that made the model look
444 + heterogeneous in the first place.
445 + """
446 model_arch_config = ModelArchitectureConfig(
447 architectures=self.get_architectures(),
448 model_type=self.hf_config.model_type,
@@ -622,13 +683,26 @@ class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
683 == "vision"
684 )
685
625 - def get_head_size(self) -> int:
626 - # Gemma4 uses dual head dimensions: head_dim (sliding attention)
627 - # and global_head_dim (full attention). Return the largest so
628 - # that attention backends allocate buffers large enough for both.
629 - head_dim = getattr(self.hf_text_config, "head_dim", 0)
630 - global_head_dim = getattr(self.hf_text_config, "global_head_dim", 0)
631 - return max(head_dim, global_head_dim) or super().get_head_size()
686 + def get_per_layer_hf_configs(
687 + self,
688 + ) -> list[tuple[PretrainedConfig, PretrainedConfig]] | None:
689 + # Gemma4 uses a larger head dimension, and sometimes more KV heads, on its
690 + # full attention layers than on its sliding ones. Transformers >= 5.15.0
691 + # says so in the config; before that the values are flat attributes picked
692 + # apart by `layer_types`, so build the per-layer configs here instead. Both
693 + # then reach the base convertor as ordinary homogeneous configs.
694 + if getattr(self.hf_text_config, "is_heterogeneous", False):
695 + return super().get_per_layer_hf_configs()
696 + text_config = self.hf_text_config
697 + if not getattr(text_config, "layer_types", None):
698 + return None
699 + # `num_hidden_layers` is what the stack is built from, and it is what
700 + # Transformers sizes `per_layer_config` by. `layer_types` can be longer,
701 + # e.g. when a test truncates the layer count to build a small model.
702 + return [
703 + (self.hf_config, gemma4_layer_config(text_config, layer_idx))
704 + for layer_idx in range(self.get_num_hidden_layers())
705 + ]
706
707
708 class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase):