Fix Gemma 4 for upcoming Transformers version

tests/config/test_model_arch_config.py

@@ -3,15 +3,20 @@
 """Tests for ModelArchitectureConfig and its integration with ModelConfig."""
 
 import json
+from copy import copy
 from pathlib import Path
 from types import SimpleNamespace
 from typing import cast
 
 import pytest
 from transformers import PretrainedConfig
+from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig
 
 from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
+from vllm.config.model_arch import ModelArchitectureConfig
+from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
 from vllm.transformers_utils.model_arch_config_convertor import (
+    Gemma4ModelArchConfigConvertor,
     ModelArchConfigConvertorBase,
 )
 
@@ -182,6 +187,220 @@ def test_legacy_modelopt_config_without_producer_is_normalized():
     assert convertor.get_quantization_config()["quant_method"] == "modelopt_fp4"
 
 
+def _layer(**overrides) -> ModelArchitectureConfig:
+    fields = dict(
+        architectures=["X"],
+        model_type="x",
+        text_model_type=None,
+        hidden_size=64,
+        total_num_hidden_layers=3,
+        total_num_attention_heads=8,
+        head_size=16,
+        vocab_size=32,
+        total_num_kv_heads=4,
+        num_experts=0,
+        num_experts_per_token=0,
+        quantization_config=None,
+        is_deepseek_mla=False,
+        is_mm_prefix_lm=False,
+        rswa_window=None,
+        derived_max_model_len_and_key=(1.0, None),
+    )
+    return ModelArchitectureConfig(**(fields | overrides))
+
+
+def test_from_layers_collapses_to_the_largest_layer():
+    """The whole-model config must size buffers for every layer."""
+    arch = ModelArchitectureConfig.from_layers(
+        [_layer(), _layer(), _layer(head_size=32, total_num_kv_heads=8)]
+    )
+
+    assert (arch.head_size, arch.total_num_kv_heads) == (32, 8)
+    assert [arch[i].head_size for i in range(3)] == [16, 16, 32]
+    assert [arch[i].total_num_kv_heads for i in range(3)] == [4, 4, 8]
+    # A layer view is itself homogeneous, so reading from it cannot recurse.
+    assert arch[0].per_layer_overrides is None
+
+
+def test_layer_view_follows_later_edits_to_the_whole_model_config():
+    """`model_arch_config` is edited after construction in a few places."""
+    arch = ModelArchitectureConfig.from_layers(
+        [_layer(), _layer(), _layer(head_size=32)]
+    )
+
+    arch.is_mm_prefix_lm = True
+
+    assert all(arch[i].is_mm_prefix_lm for i in range(3))
+    assert [arch[i].head_size for i in range(3)] == [16, 16, 32]
+
+
+def test_uniform_layers_stay_homogeneous():
+    """A checkpoint can vary attributes vLLM never reads."""
+    arch = ModelArchitectureConfig.from_layers([_layer()] * 3)
+
+    assert arch.per_layer_overrides is None
+    assert arch[1] is arch
+
+
+def test_getitem_rejects_negative_indices():
+    """Wrapping would only misbehave on heterogeneous models, so reject both."""
+    homogeneous = _layer()
+    heterogeneous = ModelArchitectureConfig.from_layers(
+        [_layer(), _layer(), _layer(head_size=32)]
+    )
+
+    for arch in (homogeneous, heterogeneous):
+        with pytest.raises(IndexError):
+            arch[-1]
+
+
+@pytest.mark.parametrize(
+    "varying",
+    [
+        # `bool` is an `int`; collapsing this one with `max` would make `use_mla`
+        # true model wide and silently discard every per-layer KV head count.
+        {"is_deepseek_mla": True},
+        {"quantization_config": {"quant_method": "fp8"}},
+        {"rswa_window": 512},
+    ],
+)
+def test_from_layers_rejects_fields_with_no_whole_model_value(varying: dict):
+    with pytest.raises(ValueError, match="varies across layers"):
+        ModelArchitectureConfig.from_layers([_layer(), _layer(), _layer(**varying)])
+
+
+def test_from_layers_rejects_a_layer_count_mismatch():
+    """Draft configs can inherit a target's per-layer spec."""
+    with pytest.raises(ValueError, match="per-layer configs"):
+        ModelArchitectureConfig.from_layers([_layer(), _layer(head_size=32)])
+
+
+def _gemma4_text_config(**overrides) -> Gemma4TextConfig:
+    """A six layer Gemma4 whose last layer is wider than the rest."""
+    fields = dict(
+        num_hidden_layers=6,
+        hidden_size=64,
+        num_attention_heads=8,
+        num_key_value_heads=4,
+        head_dim=16,
+        global_head_dim=32,
+        # Transformers forces the last layer to full attention.
+        layer_types=["sliding_attention"] * 5 + ["full_attention"],
+    )
+    return Gemma4TextConfig(**(fields | overrides))
+
+
+def test_gemma4_head_dims_vary_by_layer_type():
+    """Gemma4's full attention layers are wider than its sliding ones.
+
+    Transformers >= 5.15.0 says so in the config; this exercises the convertor
+    building the same per-layer view from the flat attributes used before that.
+    """
+    text_config = _gemma4_text_config(
+        num_global_key_value_heads=8, attention_k_eq_v=True
+    )
+
+    arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert()
+
+    assert (arch.head_size, arch.total_num_kv_heads) == (32, 8)
+    assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32]
+    assert [arch[i].total_num_kv_heads for i in range(6)] == [4] * 5 + [8]
+    # The model files resolve each layer through the same helper, so the KV cache
+    # vLLM allocates and the projections the model builds cannot disagree.
+    assert [gemma4_layer_config(text_config, i).head_dim for i in range(6)] == [
+        arch[i].head_size for i in range(6)
+    ]
+
+
+def test_gemma4_without_global_kv_heads():
+    """`num_global_key_value_heads` defaults to `None`, not to a head count."""
+    text_config = _gemma4_text_config()
+
+    arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert()
+
+    assert arch.total_num_kv_heads == 4
+    assert [arch[i].head_size for i in range(6)] == [16] * 5 + [32]
+
+
+def test_gemma4_layer_count_comes_from_num_hidden_layers():
+    """`dummy_hf_overrides` shrinks the stack but leaves `layer_types` long."""
+    text_config = _gemma4_text_config()
+    text_config.num_hidden_layers = 3
+
+    arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert()
+
+    assert arch.total_num_hidden_layers == 3
+    # Only sliding layers survive the truncation, so nothing varies.
+    assert arch.per_layer_overrides is None
+
+
+def test_gemma4_uniform_head_dims_are_homogeneous():
+    text_config = _gemma4_text_config(global_head_dim=16)
+
+    arch = Gemma4ModelArchConfigConvertor(text_config, text_config).convert()
+
+    assert arch.per_layer_overrides is None
+    assert arch[3] is arch
+
+
+class _HeterogeneousConfig(PretrainedConfig):
+    """A stand-in for the Transformers >= 5.15.0 heterogeneous config API.
+
+    No released Transformers has it, so the default per-layer seam has no other
+    way to be exercised. Mirrors the parts vLLM uses: per-layer configs are
+    shallow copies with the varying attributes applied and heterogeneity
+    stripped, so they do not recurse.
+    """
+
+    is_heterogeneous = True
+
+    def __init__(self, per_layer: dict[str, list], **kwargs):
+        super().__init__(**kwargs)
+        self._per_layer = per_layer
+
+    @property
+    def per_layer_config(self) -> list[PretrainedConfig]:
+        layers = []
+        for i in range(self.num_hidden_layers):
+            layer = copy(self)
+            layer.is_heterogeneous = False
+            for name, values in self._per_layer.items():
+                setattr(layer, name, values[i])
+            layers.append(layer)
+        return layers
+
+
+def test_transformers_heterogeneous_config_is_resolved_per_layer():
+    hf_config = _HeterogeneousConfig(
+        per_layer={"head_dim": [16, 16, 32], "num_key_value_heads": [4, 4, 8]},
+        num_hidden_layers=3,
+        hidden_size=64,
+        num_attention_heads=8,
+    )
+
+    arch = ModelArchConfigConvertorBase(hf_config, hf_config).convert()
+
+    assert (arch.head_size, arch.total_num_kv_heads) == (32, 8)
+    assert [arch[i].head_size for i in range(3)] == [16, 16, 32]
+    assert [arch[i].total_num_kv_heads for i in range(3)] == [4, 4, 8]
+
+
+def test_heterogeneous_config_varying_nothing_vllm_reads():
+    """Transformers prunes nothing for us: the diff has to notice."""
+    hf_config = _HeterogeneousConfig(
+        per_layer={"bos_token_id": [1, 2, 3]},
+        num_hidden_layers=3,
+        hidden_size=64,
+        num_attention_heads=8,
+        head_dim=16,
+    )
+
+    arch = ModelArchConfigConvertorBase(hf_config, hf_config).convert()
+
+    assert arch.per_layer_overrides is None
+    assert arch[2] is arch
+
+
 @pytest.mark.parametrize("model", BASE_MODELS_TO_TEST)
 def test_base_model_arch_config(model: str):
     """Test model architecture config for base models."""

tests/models/transformers/fusers/test_linear.py

@@ -15,6 +15,8 @@
     GLUFuser,
     PackedQKVFuser,
     QKVFuser,
+    packed_qkv,
+    qkv,
 )
 
 
@@ -640,3 +642,63 @@ def test_act_and_mul_derived_from_module(default_vllm_config):
     assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None
     with pytest.raises(ValueError, match="No AndMul equivalent"):
         GLUFuser._get_act_and_mul(nn.Dropout())
+
+
+def _wider_model_config(head_dim: int) -> SimpleNamespace:
+    """A model whose global head size is twice `head_dim`, as a wider layer
+    elsewhere in a heterogeneous checkpoint would make it."""
+    return SimpleNamespace(
+        model_config=SimpleNamespace(get_head_size=lambda: 2 * head_dim),
+        quant_config=None,
+    )
+
+
+@pytest.mark.parametrize(
+    "cls, fuser_module", [(FakeAttention, qkv), (PackedQKVAttention, packed_qkv)]
+)
+def test_head_counts_come_from_the_module_not_the_model(cls, fuser_module, monkeypatch):
+    """A layer narrower than the model-wide head size must not be miscounted.
+
+    On a heterogeneous checkpoint (Gemma 4) the model-wide head size is the
+    largest across layers, so deriving `total_num_heads = out_features //
+    head_size` from it undercounts heads on a narrower layer. The widths still
+    add up, so nothing raises below TP=4: the layer is just sharded wrong.
+    """
+    head_dim, heads, kv_heads = 8, 8, 4
+    vllm_config = _wider_model_config(head_dim)
+    with torch.device("meta"):
+        module = cls(hidden=32, head_dim=head_dim, heads=heads, kv_heads=kv_heads)
+
+    # Both replacements shard, so they need a TP group; only the head counts
+    # the fuser derives are under test here.
+    captured = {}
+    monkeypatch.setattr(
+        fuser_module,
+        "QKVParallelLinear",
+        lambda **kwargs: captured.update(kwargs) or nn.Identity(),
+    )
+    monkeypatch.setattr(
+        fuser_module, "replace_linear_class", lambda *a, **kw: nn.Identity()
+    )
+    fuser = get_fuser(module)
+    assert fuser is not None and fuser.validate(module, vllm_config)
+    fuser.update_attrs(module, "model.layers.0.self_attn", vllm_config)
+
+    assert captured["head_size"] == head_dim
+    assert captured["total_num_heads"] == heads
+    assert captured["total_num_kv_heads"] == kv_heads
+
+
+def test_validate_accepts_a_layer_the_model_wide_head_size_would_reject():
+    """`validate` gates fusion, so a wrong head size silently disables it."""
+    head_dim, heads, kv_heads = 8, 8, 3
+    vllm_config = _wider_model_config(head_dim)
+    with torch.device("meta"):
+        module = FakeAttention(
+            hidden=32, head_dim=head_dim, heads=heads, kv_heads=kv_heads
+        )
+
+    # kv width is 24, not a multiple of the model-wide 16, but is of this
+    # layer's 8.
+    fuser = get_fuser(module)
+    assert fuser is not None and fuser.validate(module, vllm_config)

vllm/config/model.py

@@ -886,9 +886,7 @@ def _supports_multimodal_for_mm_prefix(self) -> bool:
             )
         return supports_mm
 
-    def get_model_arch_config(
-        self,
-    ) -> ModelArchitectureConfig:
+    def get_model_arch_config(self) -> ModelArchitectureConfig:
         convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get(
             self.hf_config.model_type, ModelArchConfigConvertorBase
         )
@@ -982,8 +980,8 @@ def _get_transformers_backend_cls(self) -> str:
         """Determine which Transformers modeling backend class will be used if
         `model_impl` is set to `transformers` or `auto`."""
         cls = "Transformers"
-        # If 'hf_config != hf_text_config' it's a nested config, i.e. multimodal
-        cls += "MultiModal" if self.hf_config != self.hf_text_config else ""
+        # If 'hf_config is not hf_text_config' it's a nested config, i.e. multimodal
+        cls += "MultiModal" if self.hf_config is not self.hf_text_config else ""
         cls += "MoE" if self.is_moe else ""
         # Check if the architecture we're wrapping has defaults
         runner = None
@@ -1473,21 +1471,40 @@ def get_total_num_kv_heads(self) -> int:
         """Returns the total number of KV heads."""
         return self.model_arch_config.total_num_kv_heads
 
-    def get_num_kv_heads(self, parallel_config: ParallelConfig) -> int:
-        """Returns the number of KV heads per GPU."""
+    def get_num_kv_heads(
+        self,
+        parallel_config: ParallelConfig,
+        arch_config: ModelArchitectureConfig | None = None,
+    ) -> int:
+        """Returns the number of KV heads per GPU.
+
+        Pass ``arch_config`` (from ``model_arch_config[layer_idx]``) to size a
+        single layer of a heterogeneous model rather than the model as a whole.
+        """
         if self.use_mla:
             # When using MLA during decode it becomes MQA
             return 1
 
-        total_num_kv_heads = self.get_total_num_kv_heads()
+        arch_config = arch_config or self.model_arch_config
+        total_num_kv_heads = arch_config.total_num_kv_heads
         # If tensor parallelism is used, we divide the number of KV heads by
         # the tensor parallel size. We will replicate the KV heads in the
         # case where the number of KV heads is smaller than the tensor
         # parallel size so each GPU has at least one KV head.
         return max(1, total_num_kv_heads // parallel_config.tensor_parallel_size)
 
-    def get_num_attention_heads(self, parallel_config: ParallelConfig) -> int:
-        num_heads = self.model_arch_config.total_num_attention_heads
+    def get_num_attention_heads(
+        self,
+        parallel_config: ParallelConfig,
+        arch_config: ModelArchitectureConfig | None = None,
+    ) -> int:
+        """Returns the number of attention heads per GPU.
+
+        Pass ``arch_config`` (from ``model_arch_config[layer_idx]``) to size a
+        single layer of a heterogeneous model rather than the model as a whole.
+        """
+        arch_config = arch_config or self.model_arch_config
+        num_heads = arch_config.total_num_attention_heads
         return num_heads // parallel_config.tensor_parallel_size
 
     def get_num_experts(self) -> int: