[https://nvbugs/6700677][fix] Fix Gemma4 startup on non-SM100 GPUs

tensorrt_llm/_torch/models/modeling_gemma4.py

@@ -63,6 +63,7 @@
 from ..modules.gemma4.fused_qkv import gemma4_fused_qkv_norm_rope_quant
 from ..modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
 from ..modules.rms_norm import RMSNorm
+from ..pyexecutor.config_utils import get_gemma4_layer_head_dim, get_gemma4_layer_num_kv_heads
 from ..speculative.interface import SpecMetadata
 from ..utils import ActivationType, Fp4QuantizedTensor, is_torch_compiling
 from .modeling_speculative import (
@@ -227,6 +228,25 @@ def __init__(
         self.is_sliding = is_sliding
         self.is_kv_shared = is_kv_shared
         config = model_config.pretrained_config
+        geometry_layer_idx = layer_idx
+        if geometry_layer_idx is None:
+            if getattr(config, "per_layer_attributes", None):
+                raise ValueError(
+                    "Gemma4Attention requires layer_idx with a heterogeneous Transformers config."
+                )
+            geometry_layer_idx = next(
+                (
+                    idx
+                    for idx, layer_type in enumerate(config.layer_types)
+                    if (layer_type == "sliding_attention") == is_sliding
+                ),
+                None,
+            )
+            if geometry_layer_idx is None:
+                raise ValueError(
+                    "Gemma4Attention could not infer layer_idx: no layer type "
+                    f"matches is_sliding={is_sliding}."
+                )
 
         # Native TRTLLM's SM100 FP8-KV path consumes BF16 Q/K/V and quantizes
         # while appending to the cache. Keep RoPE at the model layer so the
@@ -239,22 +259,12 @@ def __init__(
             and is_sm_100f()
         )
 
-        # Per-layer head_dim and kv heads
-        # Note: num_global_key_value_heads is only used when K=V (alternative
-        # attention). For non-K=V full layers, use regular num_key_value_heads.
+        # Transformers 5.14+ exposes heterogeneous geometry only through
+        # per_layer_config, while older versions keep flat global fields.
+        # Resolve both schemas without enabling ambiguous global access.
         use_k_eq_v = getattr(config, "attention_k_eq_v", False) and not is_sliding
-        if is_sliding:
-            layer_head_dim = config.head_dim
-            layer_num_kv_heads = config.num_key_value_heads
-        else:
-            layer_head_dim = getattr(config, "global_head_dim", config.head_dim)
-            if use_k_eq_v:
-                layer_num_kv_heads = (
-                    getattr(config, "num_global_key_value_heads", None)
-                    or config.num_key_value_heads
-                )
-            else:
-                layer_num_kv_heads = config.num_key_value_heads
+        layer_head_dim = get_gemma4_layer_head_dim(config, geometry_layer_idx)
+        layer_num_kv_heads = get_gemma4_layer_num_kv_heads(config, geometry_layer_idx)
 
         # Build RoPE params per layer type
         rope_params = RopeParams()
@@ -297,11 +307,6 @@ def __init__(
 
         self.use_k_eq_v = use_k_eq_v
 
-        # Temporarily override config.head_dim so the Attention base class
-        # picks up the correct per-layer head_dim.
-        original_head_dim = config.head_dim
-        config.head_dim = layer_head_dim
-
         super().__init__(
             hidden_size=config.hidden_size,
             num_attention_heads=config.num_attention_heads,
@@ -315,16 +320,14 @@ def __init__(
             dense_bias=False,
             config=model_config,
             q_scaling=q_scaling,
+            head_dim=layer_head_dim,
             # Full-attention layers use proportional RoPE whose active
             # frequencies are paired across the full head. Apply it at the
             # module layer because fused preprocessing cannot represent that
             # pairing with only the logical rotary dimension.
             rope_fusion=is_sliding and not self._use_trtllm_fused_qkv_prep,
         )
 
-        # Restore original config head_dim
-        config.head_dim = original_head_dim
-
         # Fix proportional RoPE for full-attention layers.
         #
         # HF proportional RoPE produces cos/sin of shape [seq, head_dim] (512)
@@ -1300,17 +1303,20 @@ def __init__(
     def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
         """Gemma4-specific defaults.
 
-        The TRTLLM attention backend uses trtllm-gen for the regular attention
-        phases and a Triton context phase for bidirectional multimodal masks.
-        External shared-KV MTP still requires FlashInfer attention metadata.
+        Preserve the existing TRTLLM default on datacenter Blackwell. Other
+        architectures use FlashInfer FA2 because the native MMHA backend does
+        not support Gemma4's 512-wide heads. External shared-KV MTP also
+        requires FlashInfer attention metadata.
         """
         speculative_config = getattr(llm_args, "speculative_config", None)
         spec_dec_mode = getattr(speculative_config, "spec_dec_mode", None)
         uses_external_shared_kv = (
             spec_dec_mode is not None and spec_dec_mode.is_mtp_eagle_one_model()
         )
         return {
-            "attn_backend": "FLASHINFER" if uses_external_shared_kv else "TRTLLM",
+            "attn_backend": (
+                "FLASHINFER" if uses_external_shared_kv or not is_sm_100f() else "TRTLLM"
+            ),
         }
 
     @classmethod

tensorrt_llm/_torch/modules/qk_norm_attention.py

@@ -166,6 +166,7 @@ def __init__(
         reduce_output: bool = True,
         rope_fusion: bool = True,
         mapping_with_cp: Optional[Mapping] = None,
+        head_dim: Optional[int] = None,
     ):
         self.pretrained_config = config.pretrained_config
 
@@ -201,6 +202,7 @@ def __init__(
             attn_output_gate=attn_output_gate,
             reduce_output=reduce_output,
             mapping_with_cp=mapping_with_cp,
+            head_dim=head_dim,
         )
 
         self.q_norm = RMSNorm(hidden_size=self.head_dim,

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -49,6 +49,8 @@
                            get_spec_decoder, should_use_separate_draft_kv_cache)
 from ..utils import is_gdn_replay_enabled
 from .config_utils import (MambaKVCacheParams, extract_mamba_kv_cache_params,
+                           get_gemma4_layer_head_dim,
+                           get_gemma4_layer_num_kv_heads,
                            get_layer_attention_window, is_gemma4_hybrid,
                            is_hybrid_linear, is_kimi_linear, is_mla,
                            is_nemotron_hybrid, is_qwen3_hybrid,
@@ -2272,7 +2274,7 @@ def _create_kv_cache_manager(
         layer_mask: Optional[List[bool]] = None,
         num_layers: Optional[int] = None,
         num_kv_heads: Optional[Union[int, List[int]]] = None,
-        head_dim: Optional[int] = None,
+        head_dim: Optional[Union[int, List[int]]] = None,
         kv_cache_type=None,
         is_disagg: bool = False,
         cold_page_codec_provider: Optional[object] = None) -> KVCacheManager:
@@ -2320,38 +2322,22 @@ def _create_kv_cache_manager(
 
     hidden_size = config.hidden_size
     num_attention_heads = config.num_attention_heads
-    num_key_value_heads = num_kv_heads if num_kv_heads is not None else getattr(
-        config, 'num_key_value_heads', num_attention_heads)
-    if not isinstance(head_dim, int):
-        head_dim = getattr(config, "head_dim", None)
-    if not isinstance(head_dim, int):
-        head_dim = hidden_size // num_attention_heads
 
     # Gemma4: build per-layer head_dim, num_kv_heads, and sliding window
     # for hybrid attention. Different layer types need different KV cache
     # pool groups (via max_attention_window) so FlashInfer page indices
     # are consistent within each group.
     if is_gemma4_hybrid(config):
         layer_types = config.layer_types
-        global_head_dim = config.global_head_dim
-        attention_k_eq_v = getattr(config, 'attention_k_eq_v', False)
-        num_global_kv_heads = (getattr(config, 'num_global_key_value_heads',
-                                       None) or num_key_value_heads)
         sliding_window = getattr(config, 'sliding_window', None)
-        head_dim_list = []
-        kv_heads_list = []
-        for lt in layer_types:
-            is_sliding = (lt == "sliding_attention")
-            if is_sliding:
-                head_dim_list.append(head_dim)
-                kv_heads_list.append(num_key_value_heads)
-            else:
-                head_dim_list.append(global_head_dim)
-                use_k_eq_v = attention_k_eq_v and not is_sliding
-                kv_heads_list.append(
-                    num_global_kv_heads if use_k_eq_v else num_key_value_heads)
-        head_dim = head_dim_list
-        num_key_value_heads = kv_heads_list
+        head_dim = [
+            get_gemma4_layer_head_dim(config, layer_idx)
+            for layer_idx in range(len(layer_types))
+        ]
+        num_key_value_heads = [
+            get_gemma4_layer_num_kv_heads(config, layer_idx)
+            for layer_idx in range(len(layer_types))
+        ]
 
         # Set per-layer max_attention_window so V2 creates separate pool
         # groups for sliding vs full attention layers (different page sizes).
@@ -2371,6 +2357,13 @@ def _create_kv_cache_manager(
                 if lt == "sliding_attention" else int(max_seq_len)
                 for lt in layer_types
             ]
+    else:
+        num_key_value_heads = num_kv_heads if num_kv_heads is not None else getattr(
+            config, 'num_key_value_heads', num_attention_heads)
+        if not isinstance(head_dim, int):
+            head_dim = getattr(config, "head_dim", None)
+        if not isinstance(head_dim, int):
+            head_dim = hidden_size // num_attention_heads
 
     # Note: Gemma4 KV sharing is handled at the model level — shared layers
     # use cache_layer_idx to read from the target layer's cache slot via

tensorrt_llm/_torch/pyexecutor/config_utils.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import dataclasses
-from typing import List, Optional, Sequence
+from typing import List, Optional, Protocol, Sequence, Union, cast
 
 import torch
 import transformers
@@ -11,6 +11,26 @@
 from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig
 from tensorrt_llm.logger import logger
 
+_GEMMA4_TEXT_MODEL_TYPES = {"gemma4_text", "gemma4_unified_text"}
+
+
+class _NamedLayerType(Protocol):
+    name: str
+
+
+_LayerType = Union[str, _NamedLayerType]
+
+
+class _Gemma4LayerGeometry(Protocol):
+    head_dim: int
+    num_key_value_heads: int
+
+
+class _Gemma4GeometryConfig(Protocol):
+    head_dim: int
+    num_key_value_heads: int
+    layer_types: Sequence[_LayerType]
+
 
 def resolve_cache_transceiver_config(
         cache_transceiver_config: Optional[CacheTransceiverConfig]) -> None:
@@ -71,12 +91,52 @@ def uses_vswa_kv_cache_layout(
                     for window in max_attention_windows))
 
 
-def _is_sliding_attention_layer(layer_type: object) -> bool:
+def _is_sliding_attention_layer(layer_type: _LayerType) -> bool:
     """Return whether a config layer type denotes sliding attention."""
     layer_type_name = getattr(layer_type, "name", str(layer_type)).lower()
     return "sliding" in layer_type_name
 
 
+def _get_gemma4_per_layer_config(
+    config: _Gemma4GeometryConfig,
+    layer_idx: int,
+) -> Optional[_Gemma4LayerGeometry]:
+    """Return a concrete Gemma4 layer config when Transformers provides one."""
+    per_layer_config = getattr(config, "per_layer_config", None)
+    if per_layer_config is None:
+        return None
+    return cast(Sequence[_Gemma4LayerGeometry], per_layer_config)[layer_idx]
+
+
+def get_gemma4_layer_head_dim(config: _Gemma4GeometryConfig,
+                              layer_idx: int) -> int:
+    """Return Gemma4's head dimension for one layer across HF config schemas."""
+    layer_config = _get_gemma4_per_layer_config(config, layer_idx)
+    if layer_config is not None:
+        return layer_config.head_dim
+
+    head_dim = config.head_dim
+    if _is_sliding_attention_layer(config.layer_types[layer_idx]):
+        return head_dim
+    global_head_dim = getattr(config, "global_head_dim", None)
+    return global_head_dim if global_head_dim is not None else head_dim
+
+
+def get_gemma4_layer_num_kv_heads(config: _Gemma4GeometryConfig,
+                                  layer_idx: int) -> int:
+    """Return Gemma4's KV-head count for one layer across HF config schemas."""
+    layer_config = _get_gemma4_per_layer_config(config, layer_idx)
+    if layer_config is not None:
+        return layer_config.num_key_value_heads
+
+    num_kv_heads = config.num_key_value_heads
+    is_sliding = _is_sliding_attention_layer(config.layer_types[layer_idx])
+    if not is_sliding and getattr(config, "attention_k_eq_v", False):
+        return getattr(config, "num_global_key_value_heads",
+                       None) or num_kv_heads
+    return num_kv_heads
+
+
 def get_layer_attention_window(
     config: object,
     layer_idx: int,
@@ -124,7 +184,16 @@ def get_layer_attention_window(
 
 
 def is_gemma4_hybrid(config):
-    """True for Gemma4 models with hybrid attention (different head_dim per layer type)."""
+    """True when Gemma4 requires per-layer attention geometry."""
+    model_type = str(getattr(config, "model_type", "")).lower()
+    if model_type not in _GEMMA4_TEXT_MODEL_TYPES:
+        return False
+
+    per_layer_attributes = getattr(config, "per_layer_attributes", None)
+    if per_layer_attributes is not None:
+        return not {"head_dim", "num_key_value_heads"
+                    }.isdisjoint(per_layer_attributes)
+
     global_head_dim = getattr(config, 'global_head_dim', None)
     head_dim = getattr(config, 'head_dim', None)
     return (global_head_dim is not None and isinstance(head_dim, int)