diff --git a/vllm/config/model.py b/vllm/config/model.py index e5ac17909ba6..86e46ef4000f 100644 --- a/vllm/config/model.py +++ b/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: diff --git a/vllm/config/model_arch.py b/vllm/config/model_arch.py index 0b99df22b880..4a0b13ef53f7 100644 --- a/vllm/config/model_arch.py +++ b/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) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 4ffd6b92d6e4..826bd1fe91d2 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/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"] = ( diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 64d606c2890b..ccbb0465a693 100644 --- a/vllm/model_executor/models/config.py +++ b/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, ) diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index a22ed4c328c1..ada414850360 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/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, diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index da10a0427136..8fad74dc7ace 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/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, diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 234ae9570b21..3be85463082a 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/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, diff --git a/vllm/transformers_utils/configs/gemma4.py b/vllm/transformers_utils/configs/gemma4.py index c57ba5b7ecf2..d7b29778d29f 100644 --- a/vllm/transformers_utils/configs/gemma4.py +++ b/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 diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 250aee503786..9e6935b675a3 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -17,6 +17,7 @@ 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 @@ def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig 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, @@ -161,6 +193,18 @@ def get_num_experts(self) -> int: num_experts = self.get_num_experts_from_block_configs() return num_experts + def get_num_experts_per_token(self) -> int: + """Returns the number of routed experts selected per token.""" + return ( + getattr(self.hf_text_config, "num_experts_per_tok", 0) + or getattr(self.hf_text_config, "num_selected_experts", 0) + or getattr(self.hf_text_config, "num_experts_per_token", 0) + ) + + def rswa_window(self) -> int | None: + """Returns the Reference Sliding Window Attention window size.""" + return getattr(self.hf_text_config, "rswa_window", None) + @final @classmethod def get_torch_dtype( @@ -290,8 +334,10 @@ def is_deepseek_mla(self) -> bool: ) return False - def is_mm_prefix_lm(self) -> bool: + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: """Whether to use bidirectional attention for mm positions.""" + if not supports_multimodal: + return False if hasattr(self.hf_config, "is_mm_prefix_lm"): return bool(self.hf_config.is_mm_prefix_lm) # fallback to list of known models @@ -342,7 +388,36 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: derived_max_model_len = tmp_max_len return derived_max_model_len, max_len_key - def convert(self) -> ModelArchitectureConfig: + 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, @@ -354,9 +429,11 @@ def convert(self) -> ModelArchitectureConfig: vocab_size=self.get_vocab_size(), total_num_kv_heads=self.get_total_num_kv_heads(), num_experts=self.get_num_experts(), + num_experts_per_token=self.get_num_experts_per_token(), quantization_config=self.get_quantization_config(), is_deepseek_mla=self.is_deepseek_mla(), - is_mm_prefix_lm=self.is_mm_prefix_lm(), + is_mm_prefix_lm=self.is_mm_prefix_lm(supports_multimodal), + rswa_window=self.rswa_window(), derived_max_model_len_and_key=self.derive_max_model_len_and_key(), ) @@ -384,7 +461,7 @@ def get_total_num_kv_heads(self) -> int: ) return enc_num_kv_heads - def is_mm_prefix_lm(self) -> bool: + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: return False @@ -562,19 +639,75 @@ def get_num_hidden_layers(self) -> int: class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase): - def is_mm_prefix_lm(self) -> bool: + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: + if not supports_multimodal: + return False return ( getattr(self.hf_text_config, "use_bidirectional_attention", None) == "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_total_num_kv_heads(self) -> int: + """Handle heterogeneous Gemma4 configs from Transformers v5. + + Gemma4 uses more KV heads on full-attention layers when attention_k_eq_v + is enabled. Transformers v5 heterogeneous configs forbid direct access to + config.num_key_value_heads; must read from per_layer_config instead. + """ + text_config = self.hf_text_config + + # Transformers >= 5.15.0: heterogeneous config with per_layer_config + if getattr(text_config, "is_heterogeneous", False): + per_layer_configs = getattr(text_config, "per_layer_config", []) + if per_layer_configs: + # Return MAX kv heads across all layers (needed for KV cache sizing) + return max( + getattr(layer_cfg, "num_key_value_heads", 0) + for layer_cfg in per_layer_configs + ) + + # Transformers < 5.15.0: old-style with layer_types + flat attributes + if getattr(text_config, "layer_types", None): + num_kv_heads = getattr(text_config, "num_key_value_heads", None) + # Check if full_attention layers use more KV heads + if getattr(text_config, "attention_k_eq_v", False): + global_kv_heads = getattr( + text_config, "num_global_key_value_heads", None + ) + if global_kv_heads and num_kv_heads: + return max(num_kv_heads, global_kv_heads) + elif global_kv_heads: + return global_kv_heads + if num_kv_heads: + return num_kv_heads + + # Fallback: homogeneous config or no KV head variation + return getattr( + text_config, + "num_key_value_heads", + self.get_total_num_attention_heads(), + ) + + 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()) + ] # hf_config.model_type -> convertor class