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

tensorrt_llm/_torch/attention/backends/interface.py

@@ -193,8 +193,9 @@ def set_adp_token_counts(self, counts: List[List[int]]) -> List[int]:
 
     _saved_tensors: Dict[str, torch.Tensor] = field(init=False,
                                                     default_factory=dict)
-    # The number of heads per kv head.
-    num_heads_per_kv: Optional[int] = 1
+    # Upper bound on query heads per KV head over the target model's layers
+    # (unsharded); sizes the spec-dec tree mask.
+    max_num_heads_per_kv: Optional[int] = 1
 
     multi_item_part_lens: Optional[list[list[int]]] = None
     """Additional token layout information for multi-item scoring.

tensorrt_llm/_torch/attention/backends/trtllm.py

@@ -1470,7 +1470,7 @@ def spec_decoding_param_prepare_for_blackwell(self) -> None:
             seqLenQ = self.spec_decoding_packed_mask.shape[
                 1] if self.spec_decoding_packed_mask is not None else 1
             mask_size = torch.ops.trtllm.spec_dec_tree_custom_mask_capacity(
-                self.max_num_requests, seqLenQ, self.num_heads_per_kv)
+                self.max_num_requests, seqLenQ, self.max_num_heads_per_kv)
             self.spec_decoding_bl_tree_mask = torch.zeros(
                 mask_size,
                 dtype=torch.uint32,

tensorrt_llm/_torch/attention/qk_norm_attention.py

@@ -164,8 +164,9 @@ def __init__(
         attn_output_gate: Optional[bool] = None,
         is_qk_norm: bool = True,
         reduce_output: bool = True,
-        rope_fusion: bool = True,
+        rope_fusion: Optional[bool] = True,
         mapping_with_cp: Optional[Mapping] = None,
+        head_dim: Optional[int] = None,
     ):
         self.pretrained_config = config.pretrained_config
 
@@ -177,9 +178,11 @@ def __init__(
         self._fuse_qk_norm_rope_gate = False
 
         # If fuse_qk_norm_rope is true, do not apply fused RoPE in attention OP, and self.rotary_emb
-        # will be skipped in the overridden apply_rope.
-        rope_fusion &= (not self.fuse_qk_norm_rope and not skip_rope
-                        and not attn_output_gate and not use_gemma_rms_norm)
+        # will be skipped in the overridden apply_rope. rope_fusion=None lets
+        # the attention backend decide.
+        if (self.fuse_qk_norm_rope or skip_rope or attn_output_gate
+                or use_gemma_rms_norm):
+            rope_fusion = False
         self.is_qk_norm = is_qk_norm
         assert not (fuse_qk_norm_rope and skip_rope
                     ), "Fusing qk norm and skipping rope is not supported"
@@ -201,6 +204,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/model_config.py

@@ -22,7 +22,7 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import (TYPE_CHECKING, Any, Dict, Generic, Iterator, List, Optional,
-                    TypeVar)
+                    Sequence, Tuple, TypeVar)
 
 import filelock
 import torch
@@ -69,6 +69,26 @@
 # Python's errno module omits Linux's EBADHANDLE value.
 _LINUX_EBADHANDLE = 521
 
+
+@dataclass(frozen=True)
+class KVCacheLayerSpec:
+    """Model-provided, unsharded KV-cache geometry for one attention layer."""
+
+    head_dim: int
+    num_kv_heads: int
+    num_q_heads: int
+    attention_window: Optional[int] = None
+
+
+def _get_kv_cache_layer_specs(
+    model_config: object, ) -> Optional[Sequence[KVCacheLayerSpec]]:
+    """Return concrete per-layer geometry from full or lightweight configs."""
+    layer_specs = getattr(model_config, "kv_cache_layer_specs", None)
+    if not isinstance(layer_specs, Sequence):
+        return None
+    return layer_specs
+
+
 _DEEPSEEK_V4_ARCHITECTURES = {"DeepseekV4ForCausalLM"}
 # V4.1 shares V4's routed-MoE formats and backend choice, but needs its own
 # dense FP8 policy for both original and ModelOpt-converted checkpoints.
@@ -289,6 +309,14 @@ class ModelConfig(Generic[TConfig]):
 
     extra_attrs: Dict = field(default_factory=dict, repr=False, init=False)
 
+    # Model-provided per-layer KV geometry. It is stored per config rather than
+    # in extra_attrs because one-model drafters share extra_attrs with the
+    # target while sizing their own KV caches.
+    _kv_cache_layer_specs: Optional[Tuple[KVCacheLayerSpec,
+                                          ...]] = field(default=None,
+                                                        repr=False,
+                                                        init=False)
+
     # cute dsl op configs
     use_cute_dsl_blockscaling_mm: bool = False
     use_cute_dsl_blockscaling_bmm: bool = False
@@ -321,19 +349,30 @@ class ModelConfig(Generic[TConfig]):
     # implementation. Set by the model loader, read by `_resolve_class`.
     modeling_v2_target: Optional[type] = None
 
+    @property
+    def kv_cache_layer_specs(self) -> Optional[Tuple[KVCacheLayerSpec, ...]]:
+        """Return model-provided per-layer KV-cache geometry, if any."""
+        return self._kv_cache_layer_specs
+
+    def set_kv_cache_layer_specs(
+            self, layer_specs: Sequence[KVCacheLayerSpec]) -> None:
+        """Publish unsharded per-layer KV geometry to generic runtime code."""
+        self._kv_cache_layer_specs = tuple(layer_specs)
+
     def __setattr__(self, key, value):
         """
         Prevent modification of frozen instance attributes.
         However, we allow modification of 'extra_attrs' attributes for torch.compile
         and 'pretrained_config' attributes for mutimodal models.
         'quant_config' is allowed to be modified to set different quantization for VLM.
+        '_kv_cache_layer_specs' is set by models that publish per-layer KV geometry.
         All the other attributes are frozen.
         This can be bypassed by manually setting '_frozen' to False. The design is
         to discourage modifying the attributes unintentionally.
         """
         if self._frozen:
             if key not in ('_frozen', 'extra_attrs', 'pretrained_config',
-                           'quant_config'):
+                           'quant_config', '_kv_cache_layer_specs'):
                 raise AttributeError(
                     f"Cannot modify ModelConfig.'{key}' - instance is frozen")
         super().__setattr__(key, value)

tensorrt_llm/_torch/models/modeling_gemma4.py

@@ -56,7 +56,7 @@
     rmsnorm_residual_add,
     rmsnorm_residual_add_scale,
 )
-from ..model_config import ModelConfig
+from ..model_config import KVCacheLayerSpec, ModelConfig
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.gated_mlp import GatedMLP
@@ -89,6 +89,70 @@
 from transformers import Gemma4TextConfig  # noqa: E402
 
 
+def _get_gemma4_per_layer_config(
+    config: Gemma4TextConfig,
+    layer_idx: int,
+) -> Optional[Gemma4TextConfig]:
+    """Return concrete layer geometry when Transformers provides it."""
+    per_layer_config = getattr(config, "per_layer_config", None)
+    if per_layer_config is None:
+        return None
+    return per_layer_config[layer_idx]
+
+
+def _get_gemma4_layer_head_dim(config: Gemma4TextConfig, layer_idx: int) -> int:
+    """Resolve one layer's head dimension across Gemma4 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 config.layer_types[layer_idx] == "sliding_attention":
+        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_q_heads(config: Gemma4TextConfig, layer_idx: int) -> int:
+    """Resolve the query-head count from concrete per-layer geometry."""
+    layer_config = _get_gemma4_per_layer_config(config, layer_idx)
+    if layer_config is not None:
+        return layer_config.num_attention_heads
+    return config.num_attention_heads
+
+
+def _get_gemma4_layer_num_kv_heads(config: Gemma4TextConfig, layer_idx: int) -> int:
+    """Resolve one layer's KV-head count across Gemma4 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 = config.layer_types[layer_idx] == "sliding_attention"
+    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_gemma4_kv_cache_layer_specs(config: Gemma4TextConfig) -> list[KVCacheLayerSpec]:
+    """Translate Gemma4 config schemas into model-agnostic KV geometry."""
+    sliding_window = getattr(config, "sliding_window", None)
+    return [
+        KVCacheLayerSpec(
+            head_dim=_get_gemma4_layer_head_dim(config, layer_idx),
+            num_kv_heads=_get_gemma4_layer_num_kv_heads(config, layer_idx),
+            num_q_heads=_get_gemma4_layer_num_q_heads(config, layer_idx),
+            attention_window=(
+                int(sliding_window)
+                if sliding_window is not None
+                and config.layer_types[layer_idx] == "sliding_attention"
+                else None
+            ),
+        )
+        for layer_idx in range(config.num_hidden_layers)
+    ]
+
+
 def _gemma4_rope_max_positions(model_config: ModelConfig) -> int:
     """Size RoPE for transient speculative positions beyond max_seq_len."""
     max_positions = model_config.pretrained_config.max_position_embeddings
@@ -228,6 +292,28 @@ 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)
+                or getattr(config, "per_layer_config", None) is not 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
@@ -240,22 +326,13 @@ 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)
+        layer_num_q_heads = _get_gemma4_layer_num_q_heads(config, geometry_layer_idx)
 
         # Build RoPE params per layer type
         rope_params = RopeParams()
@@ -298,14 +375,9 @@ 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,
+            num_attention_heads=layer_num_q_heads,
             num_key_value_heads=layer_num_kv_heads,
             max_position_embeddings=config.max_position_embeddings,
             bias=False,
@@ -316,16 +388,15 @@ 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,
+            # pairing with only the logical rotary dimension. Sliding layers
+            # let the backend decide whether it can fuse RoPE.
+            rope_fusion=None if is_sliding and not self._use_trtllm_fused_qkv_prep else False,
         )
 
-        # 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)
@@ -369,9 +440,9 @@ def __init__(
             # the rotate split, matching HF's rotate_half(head_dim//2) pairing.
             self.rotary_emb.head_dim = layer_head_dim
 
-        # When FlashInfer is explicitly selected, use trtllm-gen only on
-        # datacenter Blackwell. The default TRTLLM backend selects phased FMHA
-        # directly and does not expose this FlashInfer option.
+        # With FlashInfer, use trtllm-gen only on datacenter Blackwell and FA2
+        # elsewhere. The TRTLLM backend selects phased FMHA directly and does
+        # not expose this FlashInfer option.
         if hasattr(self.attn, "flashinfer_backend"):
             self.attn.flashinfer_backend = "trtllm-gen" if is_sm_100f() else "fa2"
 
@@ -1295,23 +1366,30 @@ def __init__(
                 "moe_ep_size>1 requires a Gemma4 MoE variant (only 26B-A4B-it today)."
             )
 
+        model_config.set_kv_cache_layer_specs(
+            _get_gemma4_kv_cache_layer_specs(model_config.pretrained_config)
+        )
         super().__init__(Gemma4TextModel(model_config), model_config)
 
     @classmethod
     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.
+        SM100f defaults to TRTLLM; other architectures default to FlashInfer
+        FA2 because the TRTLLM backend's non-trtllm-gen fallback (MMHA)
+        supports head sizes only up to 256, while Gemma4 full-attention layers
+        use 512. 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/models/modeling_gemma4_unified.py

@@ -315,8 +315,7 @@ def __init__(self, model_config: ModelConfig):
                 .to(self._device)
             )
 
-        # Surface the text (LLM) config so KV-cache sizing + is_gemma4_hybrid
-        # (global head dim differs from the per-layer head dim) read the text geometry.
+        # Surface the text (LLM) config and its normalized KV-cache geometry.
         self.post_config()
         self.is_loaded = True
 

tensorrt_llm/_torch/models/modeling_gemma4mm.py

@@ -917,6 +917,9 @@ def get_sub_model_config(
     def post_config(self):
         self.config = self.llm.config
         self.model_config.pretrained_config = self.llm.config
+        layer_specs = self.llm.model_config.kv_cache_layer_specs
+        if layer_specs is not None:
+            self.model_config.set_kv_cache_layer_specs(layer_specs)
 
     @property
     def draft_config(self):