diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index ba671795770a..4e5f6e19ba18 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/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,6 +320,7 @@ 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 @@ -322,9 +328,6 @@ def __init__( 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,9 +1303,10 @@ 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) @@ -1310,7 +1314,9 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: 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 diff --git a/tensorrt_llm/_torch/modules/qk_norm_attention.py b/tensorrt_llm/_torch/modules/qk_norm_attention.py index 0cd819c63a57..23b650831d45 100644 --- a/tensorrt_llm/_torch/modules/qk_norm_attention.py +++ b/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, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 6acd34ca02e3..bcc053abc1a2 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -50,10 +50,14 @@ from ..utils import is_gdn_replay_enabled from .config_utils import (MambaKVCacheParams, extract_mamba_kv_cache_params, extract_qwen4_exp_ple_cache_params, - get_layer_attention_window, is_gemma4_hybrid, - is_hybrid_linear, is_kimi_linear, is_mla, - is_nemotron_hybrid, is_qwen3_hybrid, is_qwen4_exp, + get_layer_attention_window, is_hybrid_linear, + is_kimi_linear, is_mla, is_nemotron_hybrid, + is_qwen3_hybrid, is_qwen4_exp, uses_vswa_kv_cache_layout) + +# isort: split +from .config_utils import (get_gemma4_layer_head_dim, + get_gemma4_layer_num_kv_heads, is_gemma4_hybrid) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder @@ -2333,7 +2337,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: @@ -2381,12 +2385,6 @@ 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 @@ -2394,25 +2392,15 @@ def _create_kv_cache_manager( # 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). @@ -2432,6 +2420,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 diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 03ede19b9e5c..276018fbf10b 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/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, @@ -123,8 +183,17 @@ def get_layer_attention_window( return sliding_window -def is_gemma4_hybrid(config): - """True for Gemma4 models with hybrid attention (different head_dim per layer type).""" +def is_gemma4_hybrid(config: transformers.PretrainedConfig) -> bool: + """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) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index cb5bf1becd0f..bd75a7bb4828 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -83,7 +83,12 @@ from ...logger import logger from ...mapping import Mapping from ..utils import maybe_compile -from .config_utils import uses_vswa_kv_cache_layout +from .config_utils import ( + get_gemma4_layer_head_dim, + get_gemma4_layer_num_kv_heads, + is_gemma4_hybrid, + uses_vswa_kv_cache_layout, +) from .connectors.kv_cache_connector import KvCacheConnectorManager from .kv_cache_stats import ( KVCacheV2IterationStatsReport, @@ -408,33 +413,53 @@ def _get_static_cache_size_layer_components( ) -> tuple[List[int], List[Optional[int]]]: config = model_config.pretrained_config - num_key_value_heads = getattr(config, "num_key_value_heads", config.num_attention_heads) - if isinstance(num_key_value_heads, Iterable): - num_key_value_heads = sum(num_key_value_heads) / len(num_key_value_heads) - mla = hasattr(config, "kv_lora_rank") and config.kv_lora_rank is not None if mla: - head_dim = config.kv_lora_rank + config.qk_rope_head_dim - kv_factor = 1 + cache_sizes_per_token = [config.kv_lora_rank + config.qk_rope_head_dim] + layer_indices = None + elif is_gemma4_hybrid(config): + # Static sizing runs before a manager exists. Gemma4 exposes its + # variable head and KV-head geometry only through per-layer configs, + # so mirror the constructor inputs here instead of reading ambiguous + # global attributes. + if num_layers is None: + layer_indices = mapping.pp_layers(model_config.get_num_attention_layers()) + else: + layer_indices = list(range(max(num_layers, 1))) + if not layer_indices: + layer_indices = [0] + + tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size + cache_sizes_per_token = [] + for layer_idx in layer_indices: + config_layer_idx = layer_idx % config.num_hidden_layers + num_kv_heads = get_gemma4_layer_num_kv_heads(config, config_layer_idx) + num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size + head_dim = get_gemma4_layer_head_dim(config, config_layer_idx) + cache_sizes_per_token.append(2 * head_dim * num_kv_heads) else: + num_key_value_heads = getattr(config, "num_key_value_heads", config.num_attention_heads) + if isinstance(num_key_value_heads, Iterable): + num_key_value_heads = sum(num_key_value_heads) / len(num_key_value_heads) tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size head_dim = getattr(config, "head_dim", None) if not isinstance(head_dim, int): head_dim = config.hidden_size // config.num_attention_heads head_dim = head_dim * num_key_value_heads // tp_size - kv_factor = 2 + cache_sizes_per_token = [2 * head_dim] + layer_indices = None - cache_size_per_token = kv_factor * head_dim quant_config = model_config.quant_config - if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(): - layer_size = cache_size_per_token - elif quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache(): - layer_size = math.ceil(cache_size_per_token / 2) + math.ceil(cache_size_per_token / 16) - else: + + def get_layer_size(cache_size_per_token: int) -> int: + if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(): + return cache_size_per_token + if quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache(): + return math.ceil(cache_size_per_token / 2) + math.ceil(cache_size_per_token / 16) assert quant_config is None or (not quant_config.quant_mode.has_kv_cache_quant()), ( "Quantized kv cache is not expected" ) - layer_size = cache_size_per_token * 2 + return cache_size_per_token * 2 if num_layers is None: total_attention_layers = model_config.get_num_attention_layers() @@ -442,7 +467,11 @@ def _get_static_cache_size_layer_components( else: local_layer_ids = list(range(max(num_layers, 1))) num_attention_layers = len(local_layer_ids) - layer_sizes = [layer_size] * num_attention_layers + if len(cache_sizes_per_token) == 1: + layer_sizes = [get_layer_size(cache_sizes_per_token[0])] * num_attention_layers + else: + layer_sizes = [get_layer_size(size) for size in cache_sizes_per_token] + assert len(layer_sizes) == num_attention_layers window_pattern = kv_cache_config.max_attention_window if kv_cache_config is not None else None # Static estimation accepts an unknown max_seq_len and treats recurrent-state diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 27d8a6f6e698..0603271bb771 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -81,7 +81,8 @@ set_per_request_prefill_cuda_graph_flag, set_torch_compiling, with_model_extra_attrs) from .breakable_cuda_graph_runner import BreakableCUDAGraphRunner -from .config_utils import is_mla +from .config_utils import (get_gemma4_layer_num_kv_heads, is_gemma4_hybrid, + is_mla) from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, CUDAGraphRunner, CUDAGraphRunnerConfig, EncoderCUDAGraphRunner, @@ -118,6 +119,25 @@ def resolve_mamba_metadata_cls(model: torch.nn.Module) -> Type[Mamba2Metadata]: return getattr(model, 'mamba_metadata_cls', None) or Mamba2Metadata +def _get_num_heads_per_kv(config) -> int: + """Return the largest GQA ratio required by the model's attention layers.""" + num_attention_heads = getattr(config, "num_attention_heads", None) + if is_gemma4_hybrid(config): + num_key_value_heads = [ + get_gemma4_layer_num_kv_heads(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + else: + num_key_value_heads = getattr(config, "num_key_value_heads", None) + + if isinstance(num_key_value_heads, (list, tuple)): + num_key_value_heads = min( + (kv for kv in num_key_value_heads if kv and kv > 0), default=0) + if num_attention_heads and num_key_value_heads: + return num_attention_heads // num_key_value_heads + return 1 + + def _make_single_token_context_graph_batch( scheduled_requests: ScheduledRequests, is_multimodal_decode_compatible: Optional[Callable[[LlmRequest], @@ -3463,22 +3483,8 @@ def _set_up_attn_metadata( self.attn_runtime_features.cache_reuse or self.attn_runtime_features.chunked_prefill) cache_indirection = self.cache_indirection_attention if self.attn_backend.Metadata is TrtllmAttentionMetadata else None - num_attention_heads = getattr(self.model.model_config.pretrained_config, - 'num_attention_heads', None) config = self.model.model_config.pretrained_config - - num_attention_heads = getattr(config, 'num_attention_heads', None) - num_key_value_heads = getattr(config, 'num_key_value_heads', None) - - # Calculate the number of attention heads per KV head (GQA ratio) - if isinstance(num_key_value_heads, (list, tuple)): - # Filter out invalid KV heads, default to 0 if no valid KV heads are found - num_key_value_heads = min( - (kv for kv in num_key_value_heads if kv and kv > 0), default=0) - if num_attention_heads and num_key_value_heads: - num_heads_per_kv = num_attention_heads // num_key_value_heads - else: - num_heads_per_kv = 1 + num_heads_per_kv = _get_num_heads_per_kv(config) metadata_cls = self.attn_backend.Metadata sparse_metadata_params = ( diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 76fba30b31ce..08d184a8b330 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -21,8 +21,11 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator -from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window, is_gemma4_hybrid +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + KVCacheManagerV2, + _get_static_cache_size_layer_components, +) from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MultimodalConfig, TorchLlmArgs @@ -732,6 +735,98 @@ def get_num_attention_layers(self) -> int: assert cache_cost == CacheCost(slope=64, intercept=0) +def test_gemma4_12b_v2_static_sizing_uses_per_layer_geometry() -> None: + gemma_layer_types = [ + "sliding_attention" if (layer_idx + 1) % 6 else "full_attention" for layer_idx in range(12) + ] + + class StrictGemma4TextConfig: + model_type = "gemma4_text" + num_hidden_layers = 12 + num_attention_heads = 16 + hidden_size = 3840 + vocab_size = 262_144 + layer_types = gemma_layer_types + per_layer_attributes = {"head_dim", "num_key_value_heads"} + per_layer_config = [ + SimpleNamespace( + head_dim=256 if layer_type == "sliding_attention" else 512, + num_key_value_heads=8 if layer_type == "sliding_attention" else 1, + ) + for layer_type in gemma_layer_types + ] + + def __getattribute__(self, name: str) -> object: + if name in { + "global_head_dim", + "head_dim", + "num_global_key_value_heads", + "num_key_value_heads", + }: + raise RuntimeError(f"ambiguous global per-layer attribute: {name}") + return super().__getattribute__(name) + + class FakeModelConfig: + quant_config = None + pretrained_config = StrictGemma4TextConfig() + + def get_num_attention_layers(self) -> int: + return self.pretrained_config.num_hidden_layers + + model_config = FakeModelConfig() + mapping = Mapping(world_size=1, tp_size=1, rank=0) + max_seq_len = 8192 + kv_cache_config = KvCacheConfig( + max_attention_window=[ + 128 if layer_type == "sliding_attention" else max_seq_len + for layer_type in model_config.pretrained_config.layer_types + ] + ) + + layer_sizes, attention_windows = _get_static_cache_size_layer_components( + model_config, + mapping, + max_seq_len=max_seq_len, + kv_cache_config=kv_cache_config, + ) + expected_layer_sizes = [ + 8192 if layer_type == "sliding_attention" else 2048 + for layer_type in model_config.pretrained_config.layer_types + ] + expected_windows = [ + 128 if layer_type == "sliding_attention" else None + for layer_type in model_config.pretrained_config.layer_types + ] + assert layer_sizes == expected_layer_sizes + assert attention_windows == expected_windows + assert CacheCost.from_raw( + KVCacheManagerV2.get_cache_size_per_token( + model_config, + mapping, + tokens_per_block=32, + max_seq_len=max_seq_len, + max_batch_size=1, + kv_cache_config=kv_cache_config, + ) + ) == CacheCost( + slope=2 * 2048, + intercept=10 * 128 * 8192, + ) + + +@pytest.mark.parametrize( + "model_type", + ["gemma4_unified", "gemma4_unified_audio", "qwen3_text"], +) +def test_is_gemma4_hybrid_rejects_non_text_configs(model_type: str) -> None: + config = SimpleNamespace( + model_type=model_type, + per_layer_attributes={"head_dim"}, + per_layer_config=[SimpleNamespace(head_dim=256), SimpleNamespace(head_dim=512)], + ) + assert not is_gemma4_hybrid(config) + + def test_creator_uses_v2_affine_cache_cost(): class FakeV2Manager(KVCacheManagerV2): @staticmethod diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 361052809f54..0b1a95e24af7 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -29,7 +29,7 @@ from tensorrt_llm._torch.pyexecutor.model_engine import ( PyTorchModelEngine, _build_request_multimodal_input, _filter_cuda_graph_batch_sizes, _get_context_prompt_lookahead_token, - _make_single_token_context_graph_batch) + _get_num_heads_per_kv, _make_single_token_context_graph_batch) from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, EncodeCudaGraphConfig, PrefillCudaGraphBackend, @@ -77,6 +77,29 @@ def register_kv_caches(self, kv_cache_tensor: torch.Tensor): pass +def test_gemma4_metadata_gqa_ratio_uses_per_layer_kv_heads(): + + class StrictGemma4TextConfig: + model_type = "gemma4_text" + num_hidden_layers = 12 + num_attention_heads = 16 + per_layer_attributes = ["num_key_value_heads"] + per_layer_config = [ + SimpleNamespace( + head_dim=256, + num_key_value_heads=8 if (layer_idx + 1) % 6 else 1, + ) for layer_idx in range(12) + ] + + def __getattribute__(self, name): + if name == "num_key_value_heads": + raise RuntimeError( + f"ambiguous global per-layer attribute: {name}") + return super().__getattribute__(name) + + assert _get_num_heads_per_kv(StrictGemma4TextConfig()) == 16 + + class DummyModel(torch.nn.Module): def __init__(self, dtype: torch.dtype): diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index b9156edc34b9..1ab8bc9d64e3 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -45,6 +45,11 @@ Gemma4TextModel, Gemma4TextScaledWordEmbedding, ) +from tensorrt_llm._torch.pyexecutor.config_utils import ( + get_gemma4_layer_head_dim, + get_gemma4_layer_num_kv_heads, + is_gemma4_hybrid, +) from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.mapping import Mapping @@ -143,6 +148,59 @@ def _make_model_config(config_dict): return ModelConfig(pretrained_config=cfg, mapping=mapping) +class _HeterogeneousGemma4TextConfig(Gemma4TextConfig): + """Simulate Transformers 5.14+ rejecting ambiguous global geometry.""" + + _AMBIGUOUS_GLOBAL_ATTRS = { + "global_head_dim", + "head_dim", + "num_global_key_value_heads", + "num_key_value_heads", + } + + def __getattribute__(self, name: str) -> object: + config_dict = object.__getattribute__(self, "__dict__") + is_global_config = "_heterogeneity_spec" in config_dict or isinstance( + config_dict.get("per_layer_config"), list + ) + if ( + name in object.__getattribute__(self, "_AMBIGUOUS_GLOBAL_ATTRS") + and config_dict.get("_reject_global_geometry", False) + and is_global_config + ): + raise RuntimeError(f"ambiguous global per-layer attribute: {name}") + return super().__getattribute__(name) + + +def _make_heterogeneous_model_config( + config_dict: dict[str, object] = GEMMA4_SMALL_CONFIG, +) -> ModelConfig: + """Build a Gemma4 config with only concrete per-layer geometry readable.""" + cfg = _HeterogeneousGemma4TextConfig(**deepcopy(config_dict)) + if getattr(cfg, "per_layer_attributes", None) is None: + cfg.per_layer_config = [ + SimpleNamespace( + head_dim=( + config_dict["head_dim"] + if layer_type == "sliding_attention" + else config_dict["global_head_dim"] + ), + num_key_value_heads=( + config_dict["num_key_value_heads"] + if layer_type == "sliding_attention" + else config_dict.get("num_global_key_value_heads") + or config_dict["num_key_value_heads"] + ), + ) + for layer_type in cfg.layer_types + ] + cfg.per_layer_attributes = {"head_dim", "num_key_value_heads"} + mapping = Mapping(world_size=1, tp_size=1, rank=0) + model_config = ModelConfig(pretrained_config=cfg, mapping=mapping) + cfg._reject_global_geometry = True + return model_config + + def _make_assistant_model_config(config_dict=GEMMA4_ASSISTANT_CONFIG): """Build a ModelConfig for a standalone Gemma4 assistant.""" cfg = Gemma4AssistantConfig(**deepcopy(config_dict)) @@ -222,6 +280,27 @@ def test_model_instantiation_basic(self): self.assertIsInstance(layer, Gemma4DecoderLayer) self.assertIsInstance(layer.self_attn, Gemma4Attention) + def test_model_instantiation_with_heterogeneous_transformers_config(self): + """Model construction must use concrete per-layer attention geometry.""" + model_config = _make_heterogeneous_model_config() + + self.assertTrue(is_gemma4_hybrid(model_config.pretrained_config)) + model = Gemma4ForCausalLM(model_config) + + for layer_idx, layer in enumerate(model.model.layers): + expected_head_dim = 64 if layer.is_sliding else 128 + expected_num_kv_heads = 2 if layer.is_sliding else 1 + self.assertEqual(layer.self_attn.head_dim, expected_head_dim) + self.assertEqual(layer.self_attn.num_key_value_heads, expected_num_kv_heads) + self.assertEqual( + get_gemma4_layer_head_dim(model_config.pretrained_config, layer_idx), + expected_head_dim, + ) + self.assertEqual( + get_gemma4_layer_num_kv_heads(model_config.pretrained_config, layer_idx), + expected_num_kv_heads, + ) + def test_model_instantiation_moe(self): """Create with MoE enabled and verify MoE layers exist.""" model_config = _make_model_config(GEMMA4_MOE_CONFIG) @@ -921,6 +1000,58 @@ def test_assistant_uses_target_kv_sources(self): } +class TestGemma4HeterogeneousKVCacheLayout(unittest.TestCase): + """Regression tests for Transformers per-layer config KV geometry.""" + + def test_12b_kv_cache_manager_uses_per_layer_geometry(self): + from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + + class CapturingKVCacheManagerV2(KVCacheManagerV2): + def __init__(self, *args: object, **kwargs: object) -> None: + self.args = args + self.kwargs = kwargs + + model_config = _make_heterogeneous_model_config(GEMMA4_12B_REAL_DIMS_CONFIG) + max_seq_len = 8192 + manager = _create_kv_cache_manager( + model_engine=None, + kv_cache_manager_cls=CapturingKVCacheManagerV2, + mapping=model_config.mapping, + kv_cache_config=KvCacheConfig(max_tokens=max_seq_len), + tokens_per_block=32, + max_seq_len=max_seq_len, + max_batch_size=1, + spec_config=None, + sparse_attention_config=None, + max_num_tokens=max_seq_len, + max_beam_width=1, + kv_connector_manager=None, + model_config=model_config, + dtype=torch.bfloat16, + is_draft=False, + kv_cache_type=object(), + ) + + layer_types = model_config.pretrained_config.layer_types + self.assertEqual( + manager.kwargs["head_dim"], + [256 if layer_type == "sliding_attention" else 512 for layer_type in layer_types], + ) + self.assertEqual( + manager.kwargs["num_kv_heads"], + [8 if layer_type == "sliding_attention" else 1 for layer_type in layer_types], + ) + self.assertEqual( + manager.args[0].max_attention_window, + [ + 128 if layer_type == "sliding_attention" else max_seq_len + for layer_type in layer_types + ], + ) + + def _build_gemma4_kv_cache_manager( config, num_blocks=4, @@ -947,24 +1078,12 @@ def _build_gemma4_kv_cache_manager( kv_dtype = tensorrt_llm.bindings.DataType.BF16 layer_types = config.layer_types - attention_k_eq_v = getattr(config, "attention_k_eq_v", False) - head_dim_per_layer = [] - num_kv_heads_per_layer = [] - for lt in layer_types: - is_sliding = lt == "sliding_attention" - use_k_eq_v = attention_k_eq_v and not is_sliding - if is_sliding: - head_dim_per_layer.append(config.head_dim) - num_kv_heads_per_layer.append(config.num_key_value_heads) - else: - head_dim_per_layer.append(getattr(config, "global_head_dim", config.head_dim)) - if use_k_eq_v: - num_kv_heads_per_layer.append( - getattr(config, "num_global_key_value_heads", None) - or config.num_key_value_heads - ) - else: - num_kv_heads_per_layer.append(config.num_key_value_heads) + head_dim_per_layer = [ + get_gemma4_layer_head_dim(config, layer_idx) for layer_idx in range(len(layer_types)) + ] + num_kv_heads_per_layer = [ + get_gemma4_layer_num_kv_heads(config, layer_idx) for layer_idx in range(len(layer_types)) + ] # Use scalar if all layers have same value head_dim = head_dim_per_layer if len(set(head_dim_per_layer)) > 1 else head_dim_per_layer[0] @@ -2588,27 +2707,31 @@ def test_trtllm_fused_prep_preserves_other_sliding_paths(self, _mock_is_sm_100f) self.assertIsNotNone(attn.rotary_emb) self.assertEqual(attn.attn.flashinfer_backend, "trtllm-gen") - def test_causal_lm_requires_trtllm_backend(self): - """Gemma4ForCausalLM must default to the TRTLLM attention backend.""" - defaults = Gemma4ForCausalLM.get_model_defaults(None) - self.assertIn("attn_backend", defaults, "get_model_defaults must set attn_backend") - self.assertEqual( - defaults["attn_backend"], "TRTLLM", "Gemma4 requires TRTLLM (exact uppercase)" - ) + def test_model_defaults_follow_architecture(self): + """Preserve TRTLLM on SM100f and FlashInfer FA2 elsewhere.""" + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4ForConditionalGeneration + + for is_sm100f, expected in ((True, "TRTLLM"), (False, "FLASHINFER")): + with ( + self.subTest(is_sm100f=is_sm100f), + unittest.mock.patch( + "tensorrt_llm._torch.models.modeling_gemma4.is_sm_100f", + return_value=is_sm100f, + ), + ): + self.assertEqual( + Gemma4ForCausalLM.get_model_defaults(None)["attn_backend"], expected + ) + self.assertEqual( + Gemma4ForConditionalGeneration.get_model_defaults(None)["attn_backend"], + expected, + ) def test_causal_lm_does_not_disable_cuda_graphs(self): """Gemma4ForCausalLM must not disable CUDA graphs.""" defaults = Gemma4ForCausalLM.get_model_defaults(None) self.assertNotIn("cuda_graph_config", defaults) - def test_conditional_gen_requires_trtllm_backend(self): - """Gemma4ForConditionalGeneration must also default to TRTLLM.""" - from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4ForConditionalGeneration - - defaults = Gemma4ForConditionalGeneration.get_model_defaults(None) - self.assertIn("attn_backend", defaults) - self.assertEqual(defaults["attn_backend"], "TRTLLM") - def test_conditional_gen_does_not_disable_cuda_graphs(self): """Gemma4ForConditionalGeneration must not disable CUDA graphs.""" from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4ForConditionalGeneration @@ -2623,26 +2746,39 @@ def test_external_shared_kv_mtp_defaults_to_flashinfer(self): spec_dec_mode = SimpleNamespace(is_mtp_eagle_one_model=lambda: True) llm_args = SimpleNamespace(speculative_config=SimpleNamespace(spec_dec_mode=spec_dec_mode)) - self.assertEqual( - Gemma4ForCausalLM.get_model_defaults(llm_args)["attn_backend"], "FLASHINFER" - ) - self.assertEqual( - Gemma4ForConditionalGeneration.get_model_defaults(llm_args)["attn_backend"], - "FLASHINFER", - ) + for is_sm100f in (True, False): + with ( + self.subTest(is_sm100f=is_sm100f), + unittest.mock.patch( + "tensorrt_llm._torch.models.modeling_gemma4.is_sm_100f", + return_value=is_sm100f, + ), + ): + self.assertEqual( + Gemma4ForCausalLM.get_model_defaults(llm_args)["attn_backend"], + "FLASHINFER", + ) + self.assertEqual( + Gemma4ForConditionalGeneration.get_model_defaults(llm_args)["attn_backend"], + "FLASHINFER", + ) - def test_attn_backend_dispatches_to_trtllm(self): - """Verify the Gemma4 default dispatches to TrtllmAttention.""" + def test_attn_backend_dispatches_by_architecture(self): + """Verify architecture defaults dispatch to their intended classes.""" from tensorrt_llm._torch.attention_backend.utils import get_attention_backend - defaults = Gemma4ForCausalLM.get_model_defaults(None) - backend_cls = get_attention_backend(defaults["attn_backend"]) - - self.assertEqual( - backend_cls.__name__, - "TrtllmAttention", - "TRTLLM must dispatch to TrtllmAttention", - ) + expected_classes = ((True, "TrtllmAttention"), (False, "FlashInferAttention")) + for is_sm100f, expected_class in expected_classes: + with ( + self.subTest(is_sm100f=is_sm100f), + unittest.mock.patch( + "tensorrt_llm._torch.models.modeling_gemma4.is_sm_100f", + return_value=is_sm100f, + ), + ): + defaults = Gemma4ForCausalLM.get_model_defaults(None) + backend_cls = get_attention_backend(defaults["attn_backend"]) + self.assertEqual(backend_cls.__name__, expected_class) def test_all_layers_use_trtllm_backend(self): """All Gemma4 layers use the default TRTLLM attention backend."""