[Model] Complete port of Transformers v5 heterogeneous config fix

vllm/config/model.py

@@ -805,8 +805,11 @@ 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
+        # Use 'is not' instead of '!=' to avoid triggering __eq__ which accesses
+        # attributes that may raise AmbiguousGlobalPerLayerAttributeError on
+        # Transformers v5 heterogeneous configs (e.g. Gemma-4).
+        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
@@ -1289,21 +1292,31 @@ 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:
+    def get_num_kv_heads(
+        self,
+        parallel_config: ParallelConfig,
+        arch_config: ModelArchitectureConfig | None = None,
+    ) -> int:
         """Returns the number of KV heads per GPU."""
         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:
+        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
@@ -47,6 +49,9 @@ class ModelArchitectureConfig:
     num_experts: int
     """Number of experts in the model."""
 
+    num_experts_per_token: int
+    """Number of routed experts selected per token."""
+
     quantization_config: dict[str, Any] | None
     """Quantization configuration dictionary containing quantization parameters."""
 
@@ -56,5 +61,78 @@ class ModelArchitectureConfig:
     is_mm_prefix_lm: bool
     """Whether the model uses image bidirectional attention."""
 
+    rswa_window: int | None
+    """Reference Sliding Window Attention window size (None disables R-SWA)."""
+
     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:
+            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"
+            )
+        return cls(**merged, per_layer_overrides=overrides if any(overrides) else None)

vllm/model_executor/model_loader/weight_utils.py

@@ -270,12 +270,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

@@ -59,26 +59,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.
-
-        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.
+        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.
         """
-        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 (
@@ -88,20 +91,16 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
             ):
                 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

@@ -564,6 +564,7 @@ def __init__(
         self.is_full_attention = layer_type == "full_attention"
         head_dim = layer_config.head_dim
         num_kv_heads = layer_config.num_key_value_heads
+        num_heads = layer_config.num_attention_heads
 
         # Determine if this full-attention layer uses k_eq_v
         # (laptop variant: no v_proj, K reused as V on full attention layers)
@@ -574,7 +575,7 @@ def __init__(
         self.self_attn = Gemma4Attention(
             config=config,
             hidden_size=self.hidden_size,
-            num_heads=config.num_attention_heads,
+            num_heads=num_heads,
             num_kv_heads=num_kv_heads,
             head_dim=head_dim,
             max_position_embeddings=config.max_position_embeddings,

vllm/model_executor/models/gemma4_mtp.py

@@ -276,11 +276,12 @@ def __init__(
         layer_config = gemma4_layer_config(config, layer_idx)
         head_dim = layer_config.head_dim
         num_kv_heads = layer_config.num_key_value_heads
+        num_heads = layer_config.num_attention_heads
 
         self.self_attn = Gemma4MTPAttention(
             config=config,
             hidden_size=self.hidden_size,
-            num_heads=config.num_attention_heads,
+            num_heads=num_heads,
             num_kv_heads=num_kv_heads,
             head_dim=head_dim,
             max_position_embeddings=config.max_position_embeddings,

vllm/model_executor/models/transformers/base.py

@@ -505,9 +505,6 @@ def create_attention_instances(self) -> dict[int, Attention]:
         """
         text_config = self.text_config
 
-        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)
 
         # In encoder models, the attention layers will have `is_causal=False`
@@ -528,6 +525,19 @@ def create_attention_instances(self) -> dict[int, Attention]:
 
         attention_instances = {}
         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_flash_attention_forward
+            scale = head_size**-0.5
+            num_kv_heads = self.model_config.get_num_kv_heads(
+                self.parallel_config, arch_config
+            )
+
             # Handle interleaved sliding window attention
             per_layer_sliding_window = None
             if (
@@ -544,9 +554,7 @@ def create_attention_instances(self) -> dict[int, Attention]:
             attention_instances[i] = attn_cls(
                 num_heads=num_heads,
                 head_size=head_size,
-                # NOTE: We use Llama scale as default, if it's set by
-                # Transformers, it's updated in vllm_flash_attention_forward
-                scale=head_size**-0.5,
+                scale=scale,
                 num_kv_heads=num_kv_heads,
                 cache_config=self.cache_config,
                 quant_config=self.quant_config,

vllm/transformers_utils/configs/gemma4.py

@@ -25,10 +25,10 @@ def gemma4_layer_config(
     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
+        head_dim = getattr(text_config, "head_dim", None)
+        layer.head_dim = global_head_dim or 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
-            )
+            num_kv_heads = getattr(text_config, "num_key_value_heads", None)
+            layer.num_key_value_heads = global_kv_heads or num_kv_heads
     return layer