Fix Gemma 4 for upcoming Transformers version (#49797)

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

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 @@ from vllm.model_executor.models.transformers.fusers import (
     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 @@ class ModelConfig:
             )
         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 @@ class ModelConfig:
         """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 @@ class ModelConfig:
         """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:

vllm/config/model_arch.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from copy import copy
+from dataclasses import fields as dataclass_fields
 from typing import Any
 
 from pydantic import ConfigDict
@@ -64,3 +66,78 @@ class ModelArchitectureConfig:
 
     derived_max_model_len_and_key: tuple[float, str | None]
     """Derived maximum model length and key from the hf config."""
+
+    per_layer_overrides: "list[dict[str, Any]] | None" = None
+    """Per-layer values for the fields that vary, `None` unless some field does.
+
+    One dict per layer, holding only the fields whose value differs from the
+    whole-model value above. Everything else is read from the whole-model config,
+    so later edits to it are visible through `self[layer_idx]`."""
+
+    def __getitem__(self, layer_idx: int) -> "ModelArchitectureConfig":
+        """ModelArchitectureConfig for a specific layer.
+
+        Returns `self` when no field varies by layer, so callers never need to
+        branch on heterogeneity. Mirrors `PreTrainedConfig.per_layer_config[i]`.
+        """
+        if layer_idx < 0:
+            # Homogeneous configs would ignore a negative index and heterogeneous
+            # ones would wrap, so an off-by-one would only show up on the models
+            # this exists for. Reject it for both.
+            raise IndexError(f"layer index must not be negative, got {layer_idx}")
+        if self.per_layer_overrides is None:
+            return self
+        layer = copy(self)
+        object.__setattr__(layer, "per_layer_overrides", None)
+        for name, value in self.per_layer_overrides[layer_idx].items():
+            object.__setattr__(layer, name, value)
+        return layer
+
+    @classmethod
+    def from_layers(
+        cls, layers: "list[ModelArchitectureConfig]"
+    ) -> "ModelArchitectureConfig":
+        """Whole-model config for a checkpoint whose layers differ.
+
+        Fields that agree across layers are taken as they are. Fields that differ
+        are collapsed with `max`, so buffers are sized for the largest layer, and
+        the differing values are kept per layer. No field is named here: which
+        ones vary is whatever the checkpoint says.
+        """
+        if not layers:
+            raise ValueError("a model must have at least one layer")
+
+        merged: dict[str, Any] = {}
+        overrides: list[dict[str, Any]] = [{} for _ in layers]
+        for f in dataclass_fields(cls):
+            if f.name == "per_layer_overrides":
+                continue
+            values = [getattr(layer, f.name) for layer in layers]
+            if all(value == values[0] for value in values):
+                merged[f.name] = values[0]
+                continue
+            # `bool` is an `int`, so an exact type check is what keeps a varying
+            # flag from collapsing to `any`. `is_deepseek_mla` doing that would
+            # make `use_mla` true model wide, and `get_num_kv_heads` then returns
+            # 1 for every layer, discarding the overrides built here.
+            if not all(type(value) in (int, float) for value in values):
+                raise ValueError(
+                    f"{f.name!r} varies across layers and has no whole-model "
+                    f"value: {sorted(set(map(repr, values)))}. Only numeric "
+                    f"fields collapse (with `max`, to size buffers for the "
+                    f"largest layer); give this one an explicit rule in "
+                    f"ModelArchitectureConfig.from_layers."
+                )
+            merged[f.name] = max(values)
+            for override, value in zip(overrides, values):
+                if value != merged[f.name]:
+                    override[f.name] = value
+
+        if len(layers) != merged["total_num_hidden_layers"]:
+            raise ValueError(
+                f"got {len(layers)} per-layer configs for a model with "
+                f"{merged['total_num_hidden_layers']} layers"
+            )
+        # A checkpoint can be heterogeneous over attributes vLLM never reads, in
+        # which case there is nothing to keep the layers apart for.
+        return cls(**merged, per_layer_overrides=overrides if any(overrides) else None)

vllm/model_executor/model_loader/weight_utils.py

@@ -260,12 +260,8 @@ def get_quant_config(
         and hf_quant_config.get("quant_method") == "compressed-tensors"
         and "config_groups" in hf_quant_config
     ):
-        if hf_text_config is not None:
-            n_heads = getattr(hf_text_config, "num_attention_heads", None)
-            n_kv_heads = getattr(hf_text_config, "num_key_value_heads", None)
-        else:
-            n_heads = getattr(model_config.hf_config, "num_attention_heads", None)
-            n_kv_heads = getattr(model_config.hf_config, "num_key_value_heads", None)
+        n_heads = model_config.model_arch_config.total_num_attention_heads
+        n_kv_heads = model_config.model_arch_config.total_num_kv_heads
 
         hf_quant_config["total_num_heads"] = n_heads
         hf_quant_config["total_num_kv_heads"] = (

vllm/model_executor/models/config.py

@@ -200,26 +200,29 @@ class Gemma4Config(VerifyAndUpdateConfig):
     def verify_and_update_config(vllm_config: "VllmConfig") -> None:
         """Configure attention for heterogeneous head dimensions.
 
-        Gemma4 uses different head dimensions for sliding window
-        (head_dim) vs full attention (global_head_dim) layers. The
-        default FA3 on Hopper cannot handle head_dim > 256, which
-        causes mixed backend selection and numerical divergence.
+        Gemma4 uses different head dimensions for sliding window vs full attention
+        layers. The default FA3 on Hopper cannot handle head_dim > 256, which causes
+        mixed backend selection and numerical divergence.
 
-        When FA4 is available we force it for ALL layers, giving a
-        uniform kernel path and avoiding the mixed FA3+FA4 penalty.
-        When FA4 is not available we fall back to Triton.
+        When FA4 is available we force it for ALL layers, giving a uniform kernel path
+        and avoiding the mixed FA3+FA4 penalty. When FA4 is not available we fall back
+        to Triton.
         """
-        hf_text_config = vllm_config.model_config.hf_text_config
-        head_dim = getattr(hf_text_config, "head_dim", None)
-        global_head_dim = getattr(hf_text_config, "global_head_dim", None)
+        model_config = vllm_config.model_config
+        arch_config = model_config.model_arch_config
+        layer_types = getattr(model_config.hf_text_config, "layer_types", None) or []
+        head_dims = {
+            layer_types[i]: arch_config[i].head_size
+            for i in range(min(arch_config.total_num_hidden_layers, len(layer_types)))
+        }
 
-        if head_dim is None or global_head_dim is None or head_dim == global_head_dim:
+        if len(set(head_dims.values())) <= 1:
             return
 
         from vllm.v1.attention.backends.fa_utils import is_fa_version_supported
         from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
-        max_head_dim = max(head_dim, global_head_dim)
+        max_head_dim = max(head_dims.values())
 
         if is_fa_version_supported(4) and max_head_dim <= 512:
             if (
@@ -229,20 +232,16 @@ class Gemma4Config(VerifyAndUpdateConfig):
             ):
                 vllm_config.attention_config.flash_attn_version = 4
                 logger.info(
-                    "Gemma4 model has heterogeneous head dimensions "
-                    "(head_dim=%d, global_head_dim=%d). Using FA4 for "
+                    "Gemma4 model has heterogeneous head dimensions %s. Using FA4 for "
                     "all layers to avoid mixed FA3/FA4 penalty.",
-                    head_dim,
-                    global_head_dim,
+                    head_dims,
                 )
         elif vllm_config.attention_config.backend is None:
             vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
             logger.info(
                 "Gemma4 model has heterogeneous head dimensions "
-                "(head_dim=%d, global_head_dim=%d). FA4 not available, "
-                "forcing TRITON_ATTN backend.",
-                head_dim,
-                global_head_dim,
+                "%s. FA4 not available, forcing TRITON_ATTN backend.",
+                head_dims,
             )
 
 

vllm/model_executor/models/gemma4.py

@@ -63,6 +63,7 @@ from vllm.model_executor.model_loader.weight_utils import (
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata
 
@@ -380,7 +381,6 @@ class Gemma4Attention(nn.Module):
         num_kv_heads: int,
         head_dim: int,
         max_position_embeddings: int,
-        use_k_eq_v: bool = False,
         cache_config: CacheConfig | None = None,
         quant_config: QuantizationConfig | None = None,
         attn_logits_soft_cap: float | None = None,
@@ -389,7 +389,6 @@ class Gemma4Attention(nn.Module):
         super().__init__()
         self.config = config
         self.hidden_size = hidden_size
-        self.use_k_eq_v = use_k_eq_v
 
         tp_size = get_tensor_model_parallel_world_size()
         self.tp_rank = get_tensor_model_parallel_rank()
@@ -570,27 +569,9 @@ class Gemma4DecoderLayer(nn.Module):
         self.layer_idx = layer_idx
 
         # Gemma4 uses different head dimensions for sliding vs full attention
-        layer_type = config.layer_types[layer_idx]
-        self.is_full_attention = layer_type == "full_attention"
-        if self.is_full_attention:
-            head_dim = getattr(config, "global_head_dim", config.head_dim)
-        else:
-            head_dim = config.head_dim
-
-        # Determine if this full-attention layer uses k_eq_v
-        # (laptop variant: no v_proj, K reused as V on full attention layers)
-        use_k_eq_v = self.is_full_attention and getattr(
-            config, "attention_k_eq_v", False
-        )
-
-        # For k_eq_v full-attention layers, use num_global_key_value_heads
-        # as the KV head count when k_eq_v is enabled.
-        if use_k_eq_v:
-            num_kv_heads = getattr(
-                config, "num_global_key_value_heads", config.num_key_value_heads
-            )
-        else:
-            num_kv_heads = config.num_key_value_heads
+        layer_config = gemma4_layer_config(config, layer_idx)
+        head_dim = layer_config.head_dim
+        num_kv_heads = layer_config.num_key_value_heads
 
         self.self_attn = Gemma4Attention(
             config=config,
@@ -599,7 +580,6 @@ class Gemma4DecoderLayer(nn.Module):
             num_kv_heads=num_kv_heads,
             head_dim=head_dim,
             max_position_embeddings=config.max_position_embeddings,
-            use_k_eq_v=use_k_eq_v,
             cache_config=cache_config,
             quant_config=quant_config,
             attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None),

vllm/model_executor/models/gemma4_dspark.py

@@ -20,6 +20,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
 
 from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer
 from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal
@@ -38,18 +39,14 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
         prefix: str,
     ) -> None:
         layer_idx = extract_layer_index(prefix)
-        is_full = config.layer_types[layer_idx] == "full_attention"
-        head_dim = (
-            getattr(config, "global_head_dim", config.head_dim)
-            if is_full
-            else config.head_dim
-        )
-        use_k_eq_v = is_full and getattr(config, "attention_k_eq_v", False)
-        num_kv_heads = (
-            getattr(config, "num_global_key_value_heads", config.num_key_value_heads)
-            if use_k_eq_v
-            else config.num_key_value_heads
+        layer_type = config.layer_types[layer_idx]
+        self.use_k_eq_v = layer_type == "full_attention" and getattr(
+            config, "attention_k_eq_v", False
         )
+
+        layer_config = gemma4_layer_config(config, layer_idx)
+        head_dim = layer_config.head_dim
+        num_kv_heads = layer_config.num_key_value_heads
         super().__init__(
             config=config,
             hidden_size=config.hidden_size,
@@ -64,7 +61,6 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
         )
         self.is_kv_shared_layer = False
         self.causal = _dflash_layer_causal(config, layer_idx)
-        self.use_k_eq_v = use_k_eq_v
         self.kv_size = self.num_kv_heads * self.head_dim
         attn_bias = getattr(config, "attention_bias", False)
         self.k_proj = ColumnParallelLinear(
@@ -76,7 +72,7 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
         )
         self.v_proj = (
             None
-            if use_k_eq_v
+            if self.use_k_eq_v
             else ColumnParallelLinear(
                 config.hidden_size,
                 self.total_num_kv_heads * self.head_dim,

vllm/model_executor/models/gemma4_mtp.py

@@ -45,6 +45,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding,
 )
 from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
 
 from .gemma4 import Gemma4MLP, _get_text_config
 from .utils import (
@@ -271,21 +272,9 @@ class Gemma4MTPDecoderLayer(nn.Module):
         self.hidden_size = config.hidden_size
 
         layer_idx = extract_layer_index(prefix)
-        layer_type = config.layer_types[layer_idx]
-        is_full_attention = layer_type == "full_attention"
-        head_dim = (
-            getattr(config, "global_head_dim", config.head_dim)
-            if is_full_attention
-            else config.head_dim
-        )
-
-        use_k_eq_v = is_full_attention and getattr(config, "attention_k_eq_v", False)
-        if use_k_eq_v:
-            num_kv_heads = getattr(
-                config, "num_global_key_value_heads", config.num_key_value_heads
-            )
-        else:
-            num_kv_heads = config.num_key_value_heads
+        layer_config = gemma4_layer_config(config, layer_idx)
+        head_dim = layer_config.head_dim
+        num_kv_heads = layer_config.num_key_value_heads
 
         self.self_attn = Gemma4MTPAttention(
             config=config,

vllm/model_executor/models/transformers/base.py

@@ -522,9 +522,6 @@ class Base(
                 if qk_head_dim := qk_nope_head_dim + qk_rope_head_dim:
                     self.model_config.model_arch_config.head_size = qk_head_dim
 
-        num_heads = self.model_config.get_num_attention_heads(self.parallel_config)
-        head_size = self.model_config.get_head_size()
-        num_kv_heads = self.model_config.get_num_kv_heads(self.parallel_config)
         logits_soft_cap = getattr(text_config, "attn_logit_softcapping", None)
 
         pp_rank = self.pp_group.rank_in_group
@@ -532,11 +529,22 @@ class Base(
         start, end = get_pp_indices(text_config.num_hidden_layers, pp_rank, pp_size)
 
         for i in range(start, end):
+            # `[i]` is the whole-model config unless the checkpoint is
+            # heterogeneous, in which case it is this layer's own geometry.
+            arch_config = self.model_config.model_arch_config[i]
+            num_heads = self.model_config.get_num_attention_heads(
+                self.parallel_config, arch_config
+            )
+            head_size = arch_config.head_size
+            # Default to Llama scale, maybe updated in vllm_attention_forward
+            scale = head_size**-0.5
+            num_kv_heads = self.model_config.get_num_kv_heads(
+                self.parallel_config, arch_config
+            )
+
             kwargs = dict(
                 num_heads=num_heads,
-                # NOTE: We use Llama scale as default, if it's set by
-                # Transformers, it's updated in vllm_attention_forward
-                scale=head_size**-0.5,
+                scale=scale,
                 cache_config=self.cache_config,
                 quant_config=self.quant_config,
                 prefix=f"{i}.attn",

vllm/model_executor/models/transformers/fusers/base.py

@@ -57,6 +57,17 @@ class BaseFuser(ABC):
         return {}
 
 
+def fused_head_size(module: nn.Module, vllm_config: "VllmConfig") -> int:
+    """The head size of `module`, which head counts are derived from.
+
+    Prefer the module's own `head_dim`, which Transformers sets per instance:
+    the model-wide value is the largest head size across layers, so on a
+    heterogeneous checkpoint it would divide out the wrong number of heads, and
+    it is the text head size, so it does not describe a vision tower at all.
+    """
+    return getattr(module, "head_dim", None) or vllm_config.model_config.get_head_size()
+
+
 def local_output_sizes(merged_name: str) -> str:
     """Source for the per-rank widths of the merged linear `self.<merged_name>`."""
     merged = f"self.{merged_name}"

vllm/model_executor/models/transformers/fusers/packed_qkv.py

@@ -12,6 +12,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import QKVParallelLinear
 from vllm.model_executor.models.transformers.fusers.base import (
     RewriteFuser,
+    fused_head_size,
     local_output_sizes,
 )
 from vllm.model_executor.models.transformers.fx_utils import (
@@ -122,7 +123,7 @@ class PackedQKVFuser(RewriteFuser):
 
     def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
         """Shapes must be compatible with a head-sharded packed GEMM."""
-        head_size = vllm_config.model_config.get_head_size()
+        head_size = fused_head_size(module, vllm_config)
         qkv = module.get_submodule(self.qkv_name)
         compatible = (
             self.q_size % head_size == 0
@@ -137,7 +138,7 @@ class PackedQKVFuser(RewriteFuser):
         self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
     ) -> None:
         quant_config = vllm_config.quant_config
-        head_size = vllm_config.model_config.get_head_size()
+        head_size = fused_head_size(module, vllm_config)
         qkv_prefix = maybe_prefix(prefix, self.qkv_name)
         qkv = module.get_submodule(self.qkv_name)
         merged = QKVParallelLinear(

vllm/model_executor/models/transformers/fusers/qkv.py

@@ -12,6 +12,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import QKVParallelLinear
 from vllm.model_executor.models.transformers.fusers.base import (
     StackedFuser,
+    fused_head_size,
     local_output_sizes,
 )
 from vllm.model_executor.models.transformers.fx_utils import (
@@ -158,7 +159,7 @@ class QKVFuser(StackedFuser):
         q = module.get_submodule(self.q_name)
         k = module.get_submodule(self.k_name)
         v = module.get_submodule(self.v_name)
-        head_size = vllm_config.model_config.get_head_size()
+        head_size = fused_head_size(module, vllm_config)
         compatible = (
             q.in_features == k.in_features == v.in_features
             and len({proj.bias is None for proj in (q, k, v)}) == 1
@@ -174,7 +175,7 @@ class QKVFuser(StackedFuser):
         self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
     ) -> None:
         quant_config = vllm_config.quant_config
-        head_size = vllm_config.model_config.get_head_size()
+        head_size = fused_head_size(module, vllm_config)
         q = module.get_submodule(self.q_name)
         k = module.get_submodule(self.k_name)
         merged = QKVParallelLinear(

vllm/transformers_utils/configs/gemma4.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-layer config resolution for Gemma4 and its variants."""
+
+from copy import copy
+
+from transformers import PretrainedConfig
+
+
+def gemma4_layer_config(
+    text_config: PretrainedConfig, layer_idx: int
+) -> PretrainedConfig:
+    """The Gemma4 text config as it applies to one layer.
+
+    Gemma4 uses a larger head dimension on its full attention layers than on its
+    sliding ones, and with `attention_k_eq_v` it uses more KV heads there too.
+    Transformers >= 5.15.0 says so in the config itself; before that the values
+    are flat attributes that `layer_types` picks between, so resolve them here.
+    The result is homogeneous either way, so callers read `head_dim` and
+    `num_key_value_heads` off it without caring which version they are on.
+    """
+    if getattr(text_config, "is_heterogeneous", False):
+        return text_config.per_layer_config[layer_idx]
+
+    layer = copy(text_config)
+    if text_config.layer_types[layer_idx] == "full_attention":
+        global_head_dim = getattr(text_config, "global_head_dim", None)
+        layer.head_dim = global_head_dim or text_config.head_dim
+        if getattr(text_config, "attention_k_eq_v", False):
+            global_kv_heads = getattr(text_config, "num_global_key_value_heads", None)
+            layer.num_key_value_heads = (
+                global_kv_heads or text_config.num_key_value_heads
+            )
+    return layer

vllm/transformers_utils/model_arch_config_convertor.py

@@ -17,6 +17,7 @@ from vllm.transformers_utils.config import (
     ConfigFormat,
     get_safetensors_params_metadata,
 )
+from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
 from vllm.utils.torch_utils import common_broadcastable_dtype
 
 logger = init_logger(__name__)
@@ -27,6 +28,37 @@ class ModelArchConfigConvertorBase:
         self.hf_config = hf_config
         self.hf_text_config = hf_text_config
 
+    def get_per_layer_hf_configs(
+        self,
+    ) -> list[tuple[PretrainedConfig, PretrainedConfig]] | None:
+        """`(hf_config, hf_text_config)` per layer, or `None` if homogeneous.
+
+        This is the only place that decides whether a checkpoint is heterogeneous
+        and how many layers it has, so a convertor can declare variation that
+        Transformers does not express (see `Gemma4ModelArchConfigConvertor`).
+
+        Each pair is converted independently and the results are diffed, so a
+        convertor that overrides this must not also collapse a varying field in
+        its getters: the collapse would then be applied per layer and every layer
+        would report the same value.
+        """
+        config_het = getattr(self.hf_config, "is_heterogeneous", False)
+        text_het = getattr(self.hf_text_config, "is_heterogeneous", False)
+        if not (config_het or text_het):
+            return None
+        # The text config decides the layer count when it is the heterogeneous one:
+        # its stack is what vLLM builds attention layers for.
+        source = self.hf_text_config if text_het else self.hf_config
+        return [
+            (
+                self.hf_config.per_layer_config[i] if config_het else self.hf_config,
+                self.hf_text_config.per_layer_config[i]
+                if text_het
+                else self.hf_text_config,
+            )
+            for i in range(len(source.per_layer_config))
+        ]
+
     def get_architectures(self) -> list[str]:
         # Sometimes we get here from `vllm_config.with_hf_config(text_config)` where
         # `text_config` is a sub-config from a multi-modal model. If this is the case,
@@ -382,6 +414,35 @@ class ModelArchConfigConvertorBase:
         return derived_max_model_len, max_len_key
 
     def convert(self, supports_multimodal: bool = True) -> ModelArchitectureConfig:
+        if (per_layer := self.get_per_layer_hf_configs()) is None:
+            return self.convert_layer(supports_multimodal)
+
+        if self.is_deepseek_mla():
+            raise NotImplementedError(
+                "Heterogeneous MLA models are not supported: `get_head_size` "
+                "patches the config in place for them, which a per-layer "
+                "conversion would apply to a throwaway copy."
+            )
+        # Convert each layer and let the result work out which fields differ.
+        # Reading a varying attribute off the global config would raise, so the
+        # whole-model config has to be built up from the layers.
+        return ModelArchitectureConfig.from_layers(
+            [
+                type(self)(*configs).convert_layer(supports_multimodal)
+                for configs in per_layer
+            ]
+        )
+
+    def convert_layer(
+        self, supports_multimodal: bool = True
+    ) -> ModelArchitectureConfig:
+        """Convert one homogeneous config, without resolving per-layer values.
+
+        `convert` calls this once per layer, so it must not recurse back into
+        per-layer resolution: the configs it receives are already layer-specific,
+        and they can still carry the attributes that made the model look
+        heterogeneous in the first place.
+        """
         model_arch_config = ModelArchitectureConfig(
             architectures=self.get_architectures(),
             model_type=self.hf_config.model_type,
@@ -622,13 +683,26 @@ class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
             == "vision"
         )
 
-    def get_head_size(self) -> int:
-        # Gemma4 uses dual head dimensions: head_dim (sliding attention)
-        # and global_head_dim (full attention).  Return the largest so
-        # that attention backends allocate buffers large enough for both.
-        head_dim = getattr(self.hf_text_config, "head_dim", 0)
-        global_head_dim = getattr(self.hf_text_config, "global_head_dim", 0)
-        return max(head_dim, global_head_dim) or super().get_head_size()
+    def get_per_layer_hf_configs(
+        self,
+    ) -> list[tuple[PretrainedConfig, PretrainedConfig]] | None:
+        # Gemma4 uses a larger head dimension, and sometimes more KV heads, on its
+        # full attention layers than on its sliding ones. Transformers >= 5.15.0
+        # says so in the config; before that the values are flat attributes picked
+        # apart by `layer_types`, so build the per-layer configs here instead. Both
+        # then reach the base convertor as ordinary homogeneous configs.
+        if getattr(self.hf_text_config, "is_heterogeneous", False):
+            return super().get_per_layer_hf_configs()
+        text_config = self.hf_text_config
+        if not getattr(text_config, "layer_types", None):
+            return None
+        # `num_hidden_layers` is what the stack is built from, and it is what
+        # Transformers sizes `per_layer_config` by. `layer_types` can be longer,
+        # e.g. when a test truncates the layer count to build a small model.
+        return [
+            (self.hf_config, gemma4_layer_config(text_config, layer_idx))
+            for layer_idx in range(self.get_num_hidden_layers())
+        ]
 
 
 class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase):