From 58f1a0166295617d306971a4917ebeb09fede90b Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:22:41 +0000 Subject: [PATCH 1/6] [NVBUG-6700677][fix] Fix Gemma4 startup with per-layer KV geometry Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../_torch/attention/qk_norm_attention.py | 2 + tensorrt_llm/_torch/model_config.py | 31 ++- tensorrt_llm/_torch/models/modeling_gemma4.py | 150 +++++++++--- .../_torch/models/modeling_gemma4_unified.py | 3 +- .../_torch/models/modeling_gemma4mm.py | 3 + tensorrt_llm/_torch/pyexecutor/_util.py | 127 +++++----- .../_torch/pyexecutor/config_utils.py | 19 +- .../_torch/pyexecutor/engine/metadata.py | 14 +- .../kv_cache/kv_cache_manager_v2.py | 51 ++-- .../_torch/executor/engine/test_metadata.py | 30 ++- .../kv_cache/test_kv_cache_estimation.py | 125 +++++++++- .../_torch/modeling/test_modeling_gemma4.py | 231 ++++++++++++++---- .../modeling/test_qsa_runtime_wiring.py | 1 - 13 files changed, 596 insertions(+), 191 deletions(-) diff --git a/tensorrt_llm/_torch/attention/qk_norm_attention.py b/tensorrt_llm/_torch/attention/qk_norm_attention.py index 5618e152768d..52aa7ebe1e15 100644 --- a/tensorrt_llm/_torch/attention/qk_norm_attention.py +++ b/tensorrt_llm/_torch/attention/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/model_config.py b/tensorrt_llm/_torch/model_config.py index 13534230bad9..ac337124cd49 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/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,25 @@ # 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 + 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. @@ -321,6 +340,16 @@ 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.extra_attrs.get("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.extra_attrs["kv_cache_layer_specs"] = tuple(layer_specs) + def __setattr__(self, key, value): """ Prevent modification of frozen instance attributes. diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 260da1d0dbe2..f8db65198354 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -17,7 +17,18 @@ import copy import dataclasses import math -from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Literal, + Optional, + Protocol, + Sequence, + Tuple, + Union, + cast, +) import torch import torch.nn.functional as F @@ -56,7 +67,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 +100,79 @@ from transformers import Gemma4TextConfig # noqa: E402 +class _NamedLayerType(Protocol): + name: str + + +_LayerType = Union[str, _NamedLayerType] + + +class _Gemma4LayerGeometry(Protocol): + head_dim: int + num_key_value_heads: int + + +def _is_gemma4_sliding_layer(layer_type: _LayerType) -> bool: + """Return whether a Gemma4 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: Gemma4TextConfig, + layer_idx: int, +) -> Optional[_Gemma4LayerGeometry]: + """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 cast(Sequence[_Gemma4LayerGeometry], 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 _is_gemma4_sliding_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: 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 = _is_gemma4_sliding_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_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), + attention_window=( + int(sliding_window) + if sliding_window is not None + and _is_gemma4_sliding_layer(config.layer_types[layer_idx]) + 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 +312,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 @@ -240,22 +343,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() @@ -298,11 +391,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, @@ -316,6 +404,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 @@ -323,9 +412,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) @@ -1074,6 +1160,9 @@ def forward( # --------------------------------------------------------------------------- class Gemma4TextModel(DecoderModel): def __init__(self, model_config: ModelConfig[Gemma4TextConfig]): + model_config.set_kv_cache_layer_specs( + _get_gemma4_kv_cache_layer_specs(model_config.pretrained_config) + ) super().__init__(model_config) config = self.model_config pretrained = config.pretrained_config @@ -1301,9 +1390,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) @@ -1311,7 +1401,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/models/modeling_gemma4_unified.py b/tensorrt_llm/_torch/models/modeling_gemma4_unified.py index 0acd09ac63a0..e248cb9d069b 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4_unified.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 05476526e185..f48c21382897 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/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): diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9ad518e3db5f..7e83c22ff875 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -48,7 +48,7 @@ AttentionTypeCpp, create_kv_cache_transceiver, maybe_enable_fabric_memory_for_python_transceiver) from ..hostfunc import set_low_latency_dispatch -from ..model_config import ModelConfig +from ..model_config import ModelConfig, _get_kv_cache_layer_specs from ..models.modeling_multimodal_mixin import MultimodalModelMixin from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens, get_num_spec_layers, get_spec_decoder, @@ -58,10 +58,10 @@ from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, 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, - resolve_vocab_size, uses_vswa_kv_cache_layout) + get_layer_attention_window, is_hybrid_linear, + is_kimi_linear, is_mla, is_nemotron_hybrid, + is_qwen3_hybrid, is_qwen4_exp, resolve_vocab_size, + uses_vswa_kv_cache_layout) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder @@ -111,11 +111,17 @@ def _get_initial_lora_data_type( return None -def _non_hybrid_kv_cache_manager_cls(config, kv_cache_config: KvCacheConfig): - # Models with per-layer head_dim (e.g., Gemma4 hybrid attention) - # require KVCacheManagerV2 for per-layer buffer sizes. +def _has_variable_kv_cache_geometry(model_config: ModelConfig) -> bool: + layer_specs = _get_kv_cache_layer_specs(model_config) + return bool(layer_specs) and len({(spec.head_dim, spec.num_kv_heads) + for spec in layer_specs}) > 1 + + +def _non_hybrid_kv_cache_manager_cls(model_config: ModelConfig, + kv_cache_config: KvCacheConfig): + # Variable per-layer geometry requires V2's per-layer buffer sizes. needs_v2 = (kv_cache_config.use_kv_cache_manager_v2 is True - or is_gemma4_hybrid(config)) + or _has_variable_kv_cache_geometry(model_config)) return KVCacheManagerV2 if needs_v2 else KVCacheManager @@ -236,7 +242,8 @@ def get_kv_cache_manager_cls( if sparse_attn_config is not None: return get_sparse_attn_kv_cache_manager( sparse_attn_config, use_kv_cache_manager_v2=use_v2) - return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) + return _non_hybrid_kv_cache_manager_cls(model_config, + kv_cache_config) if sparse_attn_algorithm == "qsa" and not use_v2: raise ValueError( @@ -429,7 +436,7 @@ def get_kv_cache_manager_cls( return get_sparse_attn_kv_cache_manager(sparse_attn_config, use_kv_cache_manager_v2=use_v2) else: - return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) + return _non_hybrid_kv_cache_manager_cls(model_config, kv_cache_config) # --- KV cache cost model ------------------------------------------------------ @@ -741,7 +748,17 @@ def _derive_v2_layer_type_attention_windows( and issubclass(kv_cache_manager_cls, KVCacheManagerV2)): return None config = model_config.pretrained_config - derived_windows = _derive_layer_type_attention_windows(config, max_seq_len) + layer_specs = _get_kv_cache_layer_specs(model_config) + if layer_specs and any(spec.attention_window is not None + for spec in layer_specs): + derived_windows = _normalize_attention_windows( + [spec.attention_window for spec in layer_specs], max_seq_len) + if derived_windows is not None and not uses_vswa_kv_cache_layout( + derived_windows): + derived_windows = None + else: + derived_windows = _derive_layer_type_attention_windows( + config, max_seq_len) if derived_windows is None: return None if is_hybrid_linear(config): @@ -762,7 +779,7 @@ def _derive_v2_layer_type_attention_windows( def _get_num_pool_groups_for_estimation( - model_config: object, + model_config: ModelConfig, max_seq_len: int, fallback_attention_windows: Optional[List[Optional[int]]], ) -> int: @@ -771,17 +788,19 @@ def _get_num_pool_groups_for_estimation( Model sliding attention alone does not imply separate storage pools. Preserve distinctions required by hybrid layer types and page sizes. """ + config = model_config.pretrained_config + variable_geometry = _has_variable_kv_cache_geometry(model_config) if (fallback_attention_windows is not None - and not is_hybrid_linear(model_config)): + and not is_hybrid_linear(config)): normalized_windows = _normalize_attention_windows( fallback_attention_windows, max_seq_len) return 1 if normalized_windows is None else len(set(normalized_windows)) - layer_types = getattr(model_config, "layer_types", None) + layer_types = getattr(config, "layer_types", None) attention_windows = None - if (fallback_attention_windows is None and is_gemma4_hybrid(model_config)): + if (fallback_attention_windows is None and variable_geometry): attention_windows = _derive_layer_type_attention_windows( - model_config, max_seq_len) + config, max_seq_len) # Check whether KV storage uses configured or inferred attention windows, # which can be independently configured from the model's use of sliding-window # attention for computation. @@ -802,8 +821,8 @@ def _get_num_pool_groups_for_estimation( # type unless their page sizes differ (Gemma4). pool_types = set() for layer_type in layer_types: - if (not kv_cache_uses_attention_windows - and not is_gemma4_hybrid(model_config) and + if (not kv_cache_uses_attention_windows and not variable_geometry + and (_is_sliding_attention_layer(layer_type) or getattr(layer_type, "name", str(layer_type)).lower() == "full_attention")): @@ -1037,13 +1056,11 @@ def _validate_or_fallback_kv_cache_manager_v2( f"supported with " f"{incompat_str}. Disable the incompatible features to " f"run sparse-attention models.") - # Gemma4 hybrid uses per-layer head_dim that V1 would coerce to - # ``max(head_dim)``, changing per-layer KV byte sizes. - if is_gemma4_hybrid(config): + if _has_variable_kv_cache_geometry(model_config): raise NotImplementedError( - f"Gemma4 hybrid attention requires KVCacheManagerV2, " + f"Variable per-layer KV geometry requires KVCacheManagerV2, " f"which is not yet supported with {incompat_str}. " - f"Disable these features to run Gemma4 hybrid models.") + f"Disable these features to run this model.") if fp4_mla: raise NotImplementedError( "FP4 MLA requires Fp4MlaKVCacheManagerV2, which is " @@ -1610,7 +1627,7 @@ def _get_token_num_for_estimation(self) -> int: # (distinct layer types without sliding windows). num_pool_groups = 1 if self._is_kv_cache_manager_v2: - model_cfg = self._model_engine.model.model_config.pretrained_config + model_cfg = self._model_engine.model.model_config num_pool_groups = _get_num_pool_groups_for_estimation( model_cfg, self._model_engine.max_seq_len, @@ -3249,7 +3266,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, disable_overlap_scheduler: bool = False, @@ -3298,37 +3315,18 @@ 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) - 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 + + layer_specs = _get_kv_cache_layer_specs(_model_config) + if layer_specs: + head_dim = [spec.head_dim for spec in layer_specs] + num_key_value_heads = [spec.num_kv_heads for spec in layer_specs] + 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 # Derive per-layer max_attention_window for any model that publishes a # mixed sliding/full `layer_types` schedule (Gemma4 hybrid included) so V2 @@ -3359,7 +3357,8 @@ def _create_kv_cache_manager( derived_windows = None if (not is_draft and kv_cache_type == tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF - and (not estimating_kv_cache or is_gemma4_hybrid(config))): + and (not estimating_kv_cache + or _has_variable_kv_cache_geometry(_model_config))): derived_windows = _derive_v2_layer_type_attention_windows( kv_cache_config, kv_cache_manager_cls, _model_config, max_seq_len) if derived_windows is not None: @@ -3381,10 +3380,6 @@ def _create_kv_cache_manager( kv_cache_config = copy.copy(kv_cache_config) kv_cache_config.max_attention_window = derived_windows - # 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 - # Gemma4Attention. No layer_mask exclusion needed here. - if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(): kv_cache_dtype = tensorrt_llm.bindings.DataType.FP8 elif quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache( @@ -3410,8 +3405,8 @@ def _create_kv_cache_manager( if layer_mask is None: draft_config_for_kv = (getattr(model_engine.model, 'draft_config', None) if model_engine is not None else None) - # If num_key_value_heads is already a per-layer list (e.g., Gemma4 hybrid), - # use it directly; otherwise build from the scalar value. + # Use model-provided per-layer KV heads directly; otherwise expand the + # scalar value for speculative decoding when needed. if isinstance(num_key_value_heads, list): per_layer_num_kv_heads = num_key_value_heads else: @@ -3873,10 +3868,6 @@ def _create_kv_cache_manager( is_disagg=is_disagg, **manager_extra_kwargs, ) - # Note: Gemma4 KV sharing cache remapping is handled in Gemma4Attention - # via cache_layer_idx — shared layers use target layer's index for - # get_buffers(). No layer_offsets remapping needed here. - # Propagate the finalized chunked-prefill flag so KVCacheManager.fit_token_budget # only shrinks context chunks when the attention backend can consume a # partial context chunk. The flag is read from attn_runtime_features, which diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 8878090ebe18..24b094deb4d4 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -3,7 +3,7 @@ import dataclasses from collections.abc import Mapping as AbcMapping -from typing import TYPE_CHECKING, List, Optional, Sequence +from typing import TYPE_CHECKING, List, Optional, Protocol, Sequence, Union import torch import transformers @@ -18,6 +18,13 @@ from tensorrt_llm._torch.model_config import ModelConfig +class _NamedLayerType(Protocol): + name: str + + +_LayerType = Union[str, _NamedLayerType] + + def resolve_cache_transceiver_config( cache_transceiver_config: Optional[CacheTransceiverConfig]) -> None: """Resolve defaults and validate runtime-independent configuration.""" @@ -80,7 +87,7 @@ 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 @@ -132,14 +139,6 @@ 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).""" - 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) - and global_head_dim != head_dim) - - def is_hybrid_linear(config): return is_nemotron_hybrid(config) or is_qwen3_hybrid(config) or \ is_kimi_linear(config) or is_qwen4_exp(config) or is_glm5_next(config) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py index c9ed43c82ef9..5e6d8b1e8a29 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py @@ -14,7 +14,7 @@ AttentionMetadata, AttentionRuntimeFeatures, ) -from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.model_config import ModelConfig, _get_kv_cache_layer_specs from tensorrt_llm.mapping import Mapping from ..config_utils import is_mla @@ -26,9 +26,15 @@ __all__ = ["build_attention_metadata"] -def _get_num_heads_per_kv(pretrained_config: object) -> int: +def _get_num_heads_per_kv(model_config: ModelConfig) -> int: + pretrained_config = model_config.pretrained_config num_attention_heads = getattr(pretrained_config, "num_attention_heads", None) - num_key_value_heads = getattr(pretrained_config, "num_key_value_heads", None) + layer_specs = _get_kv_cache_layer_specs(model_config) + num_key_value_heads = ( + [spec.num_kv_heads for spec in layer_specs] + if layer_specs + else getattr(pretrained_config, "num_key_value_heads", None) + ) if isinstance(num_key_value_heads, (list, tuple)): num_key_value_heads = min( (heads for heads in num_key_value_heads if heads and heads > 0), @@ -61,7 +67,7 @@ def build_attention_metadata( attention_runtime_features.cache_reuse or attention_runtime_features.chunked_prefill ) if num_heads_per_kv is None: - num_heads_per_kv = _get_num_heads_per_kv(pretrained_config) + num_heads_per_kv = _get_num_heads_per_kv(model_config) sparse_attention_config = model_config.sparse_attention_config sparse_metadata_params = ( diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index e8069659892f..20594d4abd07 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -37,6 +37,7 @@ from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp +from tensorrt_llm._torch.model_config import _get_kv_cache_layer_specs from tensorrt_llm._torch.utils import helix_local_len, maybe_compile from tensorrt_llm._utils import ( TensorWrapper, @@ -686,34 +687,50 @@ def _get_static_cache_size_layer_components( is_external_draft: bool = False, ) -> 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) + layer_specs = _get_kv_cache_layer_specs(model_config) 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 layer_specs: + 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: + layer_spec = layer_specs[layer_idx % len(layer_specs)] + num_kv_heads = layer_spec.num_kv_heads + num_kv_heads = (num_kv_heads + tp_size - 1) // tp_size + cache_sizes_per_token.append(2 * layer_spec.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() @@ -721,7 +738,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/tests/unittest/_torch/executor/engine/test_metadata.py b/tests/unittest/_torch/executor/engine/test_metadata.py index 801d01b70f2f..f9f52ba5fe38 100644 --- a/tests/unittest/_torch/executor/engine/test_metadata.py +++ b/tests/unittest/_torch/executor/engine/test_metadata.py @@ -8,7 +8,11 @@ import pytest from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures -from tensorrt_llm._torch.pyexecutor.engine.metadata import build_attention_metadata +from tensorrt_llm._torch.model_config import KVCacheLayerSpec, ModelConfig +from tensorrt_llm._torch.pyexecutor.engine.metadata import ( + _get_num_heads_per_kv, + build_attention_metadata, +) pytestmark = pytest.mark.cpu_only @@ -95,3 +99,27 @@ def test_build_attention_metadata_forwards_shared_and_cache_inputs() -> None: assert metadata.draft_kv_cache_manager is draft_kv_cache_manager assert metadata.enable_context_mla_with_cached_kv assert metadata.num_heads_per_kv == 4 + + +def test_metadata_gqa_ratio_uses_model_provided_per_layer_kv_heads(): + class StrictConfig: + num_hidden_layers = 12 + num_attention_heads = 16 + + def __getattribute__(self, name): + if name == "num_key_value_heads": + raise RuntimeError(f"global geometry must not be read: {name}") + return super().__getattribute__(name) + + model_config = ModelConfig(pretrained_config=StrictConfig()) + model_config.set_kv_cache_layer_specs( + [ + KVCacheLayerSpec( + head_dim=256, + num_kv_heads=8 if (layer_idx + 1) % 6 else 1, + ) + for layer_idx in range(12) + ] + ) + + assert _get_num_heads_per_kv(model_config) == 16 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index d706fb826687..02c2a2ef9d80 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -19,7 +19,11 @@ import tensorrt_llm.bindings.internal.batch_manager as batch_manager from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures -from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.model_config import ( + KVCacheLayerSpec, + ModelConfig, + _get_kv_cache_layer_specs, +) from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor._util import ( CacheCost, @@ -30,7 +34,10 @@ get_mla_context_workspace_kv_len_cap, ) from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheManagerV2, + _get_static_cache_size_layer_components, +) from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode @@ -51,6 +58,11 @@ # --------------------------------------------------------------------------- +def test_unconfigured_mock_has_no_kv_cache_layer_specs() -> None: + """A lightweight mock must not fabricate iterable per-layer geometry.""" + assert _get_kv_cache_layer_specs(Mock()) is None + + def _make_mock_request(num_input_tokens, beam_width=1): """Create a mock request with the fields _get_token_num_for_estimation reads.""" req = Mock() @@ -661,8 +673,14 @@ def test_gemma4_hybrid_scales_by_num_pool_groups(): layer_types=layer_types, sliding_window=sliding_window, ) - config = hybrid._model_engine.model.model_config.pretrained_config - config.head_dim, config.global_head_dim = 256, 128 + hybrid._model_engine.model.model_config.kv_cache_layer_specs = tuple( + KVCacheLayerSpec( + head_dim=256 if layer_type == "sliding_attention" else 128, + num_kv_heads=1, + attention_window=sliding_window if layer_type == "sliding_attention" else None, + ) + for layer_type in layer_types + ) uniform = _make_creator( tpb, [_make_mock_request(max_seq_len - 1), _make_mock_request(1)], @@ -754,7 +772,10 @@ def test_uniform_window_overrides_mixed_attention_metadata() -> None: sliding_window=128, ) - assert _get_num_pool_groups_for_estimation(config, 4096, [512]) == 1 + assert ( + _get_num_pool_groups_for_estimation(SimpleNamespace(pretrained_config=config), 4096, [512]) + == 1 + ) def test_vswa_max_attention_window_fallback_scales(): @@ -808,8 +829,14 @@ def test_pool_scaling_prevents_mmmu_pro_underestimation(): layer_types=layer_types, sliding_window=sliding_window, ) - config = c._model_engine.model.model_config.pretrained_config - config.head_dim, config.global_head_dim = 256, 128 + c._model_engine.model.model_config.kv_cache_layer_specs = tuple( + KVCacheLayerSpec( + head_dim=256 if layer_type == "sliding_attention" else 128, + num_kv_heads=1, + attention_window=sliding_window if layer_type == "sliding_attention" else None, + ) + for layer_type in layer_types + ) total_tokens = c._get_token_num_for_estimation() per_pool_tokens = total_tokens // 2 # 2 pool groups @@ -975,6 +1002,90 @@ class UnsupportedKVCacheManagerV2(KVCacheManagerV2): assert single_kvcm == target == draft == CacheCost(slope=0, intercept=3 * 192 * 64) +@pytest.mark.parametrize( + ("tp_size", "enable_attention_dp", "expected_sliding_layer_size"), + [ + (1, False, 8192), + (8, False, 1024), + (8, True, 8192), + ], +) +def test_v2_static_sizing_uses_model_provided_per_layer_geometry( + tp_size: int, + enable_attention_dp: bool, + expected_sliding_layer_size: int, +) -> None: + layer_specs = tuple( + KVCacheLayerSpec( + head_dim=256 if (layer_idx + 1) % 6 else 512, + num_kv_heads=8 if (layer_idx + 1) % 6 else 1, + attention_window=128 if (layer_idx + 1) % 6 else None, + ) + for layer_idx in range(12) + ) + + class StrictConfig: + num_hidden_layers = 12 + num_attention_heads = 16 + hidden_size = 3840 + vocab_size = 262_144 + + def __getattribute__(self, name: str) -> object: + if name in {"head_dim", "num_key_value_heads"}: + raise RuntimeError(f"global geometry must not be read: {name}") + return super().__getattribute__(name) + + class FakeModelConfig: + quant_config = None + pretrained_config = StrictConfig() + kv_cache_layer_specs = layer_specs + + def get_num_attention_layers(self) -> int: + return self.pretrained_config.num_hidden_layers + + model_config = FakeModelConfig() + mapping = Mapping( + world_size=tp_size, + tp_size=tp_size, + rank=0, + enable_attention_dp=enable_attention_dp, + ) + max_seq_len = 8192 + kv_cache_config = KvCacheConfig( + max_attention_window=[ + spec.attention_window if spec.attention_window is not None else max_seq_len + for spec in layer_specs + ] + ) + + 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 = [ + expected_sliding_layer_size if spec.attention_window is not None else 2048 + for spec in layer_specs + ] + expected_windows = [spec.attention_window for spec in layer_specs] + 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 * expected_sliding_layer_size, + ) + + def test_creator_uses_v2_affine_cache_cost(): class FakeV2Manager(KVCacheManagerV2): @staticmethod diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index afdccd638c3f..f65d1268dd1a 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -46,6 +46,7 @@ Gemma4MoeRoutingMethod, Gemma4TextModel, Gemma4TextScaledWordEmbedding, + _get_gemma4_kv_cache_layer_specs, ) from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig @@ -145,6 +146,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)) @@ -224,6 +278,22 @@ 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() + model = Gemma4ForCausalLM(model_config) + layer_specs = model.model_config.kv_cache_layer_specs + self.assertIsNotNone(layer_specs) + self.assertGreater(len({(spec.head_dim, spec.num_kv_heads) for spec in layer_specs}), 1) + + 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(layer_specs[layer_idx].head_dim, expected_head_dim) + self.assertEqual(layer_specs[layer_idx].num_kv_heads, 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) @@ -924,6 +994,60 @@ 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.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) + model = Gemma4ForCausalLM(model_config) + model_config = model.model_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, @@ -934,9 +1058,8 @@ def _build_gemma4_kv_cache_manager( ): """Create KVCacheManagerV2 supporting Gemma4 per-layer head_dim / kv_heads. - Mirrors ``Gemma4Attention``'s layout (global kv heads only for K=V layers) - and the ``_util.py::is_gemma4_hybrid`` VSWA pool grouping so the page - sizes line up with what the model actually requests at runtime. + Uses the model adapter's normalized KV layer specs so the page sizes line + up with what the model requests at runtime. """ import tensorrt_llm from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -950,24 +1073,9 @@ 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) + layer_specs = _get_gemma4_kv_cache_layer_specs(config) + head_dim_per_layer = [spec.head_dim for spec in layer_specs] + num_kv_heads_per_layer = [spec.num_kv_heads for spec in layer_specs] # 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 +2696,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 +2735,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.backends.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.""" diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index 2907a2a7e744..ee5b97703635 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -324,7 +324,6 @@ def test_qsa_cache_manager_uses_resolved_index_geometry( dtype=torch.bfloat16, mamba_ssm_cache_dtype=torch.bfloat16, ) - monkeypatch.setattr(_util, "is_gemma4_hybrid", lambda config: False) monkeypatch.setattr(_util, "is_kimi_linear", lambda config: False) monkeypatch.setattr(_util, "is_mla", lambda config: False) monkeypatch.setattr(_util, "is_nemotron_hybrid", lambda config: False) From 1831abceb484a2017c19ca3d4c2d02dcdb9ec255 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:26:47 +0000 Subject: [PATCH 2/6] [https://nvbugs/6700677][fix] Resolve Gemma4 per-layer query geometry Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../_torch/attention/backends/interface.py | 4 +- .../_torch/attention/backends/trtllm.py | 2 +- tensorrt_llm/_torch/model_config.py | 1 + tensorrt_llm/_torch/models/modeling_gemma4.py | 58 ++++++---------- tensorrt_llm/_torch/pyexecutor/_util.py | 2 +- .../_torch/pyexecutor/engine/metadata.py | 22 +++--- .../pyexecutor/engine/runners/encoder.py | 4 +- .../engine/runners/encoder_decoder.py | 2 +- .../attention/sparse/test_sparse_mha.py | 2 +- .../attention/sparse/test_sparse_mqa_gqa.py | 2 +- .../_torch/executor/engine/test_metadata.py | 57 +++++++++++---- .../_torch/executor/engine/test_runners.py | 2 +- .../kv_cache/test_kv_cache_estimation.py | 69 +++++++++++++++++-- .../_torch/modeling/test_modeling_gemma4.py | 48 ++++++++++++- .../_torch/modeling/test_modeling_llama.py | 6 +- 15 files changed, 198 insertions(+), 83 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 7ccf04220604..a124fb319011 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -193,8 +193,8 @@ 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 + # The largest query-to-KV head ratio across layers sharing this metadata. + 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. diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 909e208129b0..1d7f183e6e85 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/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, diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index ac337124cd49..b4acdf012d86 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -76,6 +76,7 @@ class KVCacheLayerSpec: head_dim: int num_kv_heads: int + num_q_heads: int attention_window: Optional[int] = None diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index f8db65198354..7200191c36ca 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -17,18 +17,7 @@ import copy import dataclasses import math -from typing import ( - TYPE_CHECKING, - Any, - Dict, - Literal, - Optional, - Protocol, - Sequence, - Tuple, - Union, - cast, -) +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -100,33 +89,15 @@ from transformers import Gemma4TextConfig # noqa: E402 -class _NamedLayerType(Protocol): - name: str - - -_LayerType = Union[str, _NamedLayerType] - - -class _Gemma4LayerGeometry(Protocol): - head_dim: int - num_key_value_heads: int - - -def _is_gemma4_sliding_layer(layer_type: _LayerType) -> bool: - """Return whether a Gemma4 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: Gemma4TextConfig, layer_idx: int, -) -> Optional[_Gemma4LayerGeometry]: +) -> 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 cast(Sequence[_Gemma4LayerGeometry], per_layer_config)[layer_idx] + return per_layer_config[layer_idx] def _get_gemma4_layer_head_dim(config: Gemma4TextConfig, layer_idx: int) -> int: @@ -136,12 +107,20 @@ def _get_gemma4_layer_head_dim(config: Gemma4TextConfig, layer_idx: int) -> int: return layer_config.head_dim head_dim = config.head_dim - if _is_gemma4_sliding_layer(config.layer_types[layer_idx]): + 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) @@ -149,7 +128,7 @@ def _get_gemma4_layer_num_kv_heads(config: Gemma4TextConfig, layer_idx: int) -> return layer_config.num_key_value_heads num_kv_heads = config.num_key_value_heads - is_sliding = _is_gemma4_sliding_layer(config.layer_types[layer_idx]) + 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 @@ -162,10 +141,11 @@ def _get_gemma4_kv_cache_layer_specs(config: Gemma4TextConfig) -> list[KVCacheLa 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 _is_gemma4_sliding_layer(config.layer_types[layer_idx]) + and config.layer_types[layer_idx] == "sliding_attention" else None ), ) @@ -314,7 +294,10 @@ def __init__( config = model_config.pretrained_config geometry_layer_idx = layer_idx if geometry_layer_idx is None: - if getattr(config, "per_layer_attributes", 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." ) @@ -349,6 +332,7 @@ def __init__( use_k_eq_v = getattr(config, "attention_k_eq_v", False) and not is_sliding 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() @@ -393,7 +377,7 @@ def __init__( 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, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7e83c22ff875..4dd0f8a6f057 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3314,13 +3314,13 @@ def _create_kv_cache_manager( kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF hidden_size = config.hidden_size - num_attention_heads = config.num_attention_heads layer_specs = _get_kv_cache_layer_specs(_model_config) if layer_specs: head_dim = [spec.head_dim for spec in layer_specs] num_key_value_heads = [spec.num_kv_heads for spec in layer_specs] else: + 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): diff --git a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py index 5e6d8b1e8a29..8b130ab9877c 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py @@ -26,15 +26,15 @@ __all__ = ["build_attention_metadata"] -def _get_num_heads_per_kv(model_config: ModelConfig) -> int: +def _get_max_num_heads_per_kv(model_config: ModelConfig) -> int: + """Return the largest per-layer GQA ratio for shared mask allocation.""" + layer_specs = _get_kv_cache_layer_specs(model_config) + if layer_specs: + return max(spec.num_q_heads // spec.num_kv_heads for spec in layer_specs) + pretrained_config = model_config.pretrained_config num_attention_heads = getattr(pretrained_config, "num_attention_heads", None) - layer_specs = _get_kv_cache_layer_specs(model_config) - num_key_value_heads = ( - [spec.num_kv_heads for spec in layer_specs] - if layer_specs - else getattr(pretrained_config, "num_key_value_heads", None) - ) + num_key_value_heads = getattr(pretrained_config, "num_key_value_heads", None) if isinstance(num_key_value_heads, (list, tuple)): num_key_value_heads = min( (heads for heads in num_key_value_heads if heads and heads > 0), @@ -58,7 +58,7 @@ def build_attention_metadata( kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None = None, draft_kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None = None, enable_context_mla_with_cached_kv: bool | None = None, - num_heads_per_kv: int | None = None, + max_num_heads_per_kv: int | None = None, ) -> AttentionMetadata: """Construct attention metadata and resolve model-derived inputs.""" pretrained_config = model_config.pretrained_config @@ -66,8 +66,8 @@ def build_attention_metadata( enable_context_mla_with_cached_kv = is_mla(pretrained_config) and ( attention_runtime_features.cache_reuse or attention_runtime_features.chunked_prefill ) - if num_heads_per_kv is None: - num_heads_per_kv = _get_num_heads_per_kv(model_config) + if max_num_heads_per_kv is None: + max_num_heads_per_kv = _get_max_num_heads_per_kv(model_config) sparse_attention_config = model_config.sparse_attention_config sparse_metadata_params = ( @@ -86,6 +86,6 @@ def build_attention_metadata( enable_flash_mla=model_config.enable_flash_mla, enable_context_mla_with_cached_kv=enable_context_mla_with_cached_kv, cache_indirection=cache_indirection, - num_heads_per_kv=num_heads_per_kv, + max_num_heads_per_kv=max_num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py index 42146f9a0ec3..11db8b8c1d1b 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py @@ -340,7 +340,7 @@ def _create_attention_metadata( *, cache_indirection: torch.Tensor | None = None, enable_context_mla_with_cached_kv: bool | None = None, - num_heads_per_kv: int | None = None, + max_num_heads_per_kv: int | None = None, ) -> AttentionMetadata: metadata = build_attention_metadata( self._model.model_config, @@ -353,7 +353,7 @@ def _create_attention_metadata( cache_indirection=cache_indirection, kv_cache_manager=None, enable_context_mla_with_cached_kv=enable_context_mla_with_cached_kv, - num_heads_per_kv=num_heads_per_kv, + max_num_heads_per_kv=max_num_heads_per_kv, ) metadata.block_ids_per_seq = None metadata.kv_block_ids_per_seq = None diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py index 4caa001e58de..62b2e55871f9 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py @@ -85,7 +85,7 @@ def _build_attention_metadata( raise ValueError("Encoder sequence lengths and request IDs must have the same length.") metadata = self._create_attention_metadata( enable_context_mla_with_cached_kv=False, - num_heads_per_kv=1, + max_num_heads_per_kv=1, ) if not isinstance(metadata, (VanillaAttentionMetadata, TrtllmAttentionMetadata)): raise TypeError( diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mha.py b/tests/unittest/_torch/attention/sparse/test_sparse_mha.py index 1f8b62a6d591..9fa5f71cea56 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mha.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mha.py @@ -271,7 +271,7 @@ def create_generation_inputs(scenario: MhaGenerationScenario) -> MhaGenerationIn kv_cache_manager=kv_cache_manager, request_ids=request_ids, prompt_lens=list(scenario.past_kv_lens), - num_heads_per_kv=1, + max_num_heads_per_kv=1, runtime_features=AttentionRuntimeFeatures( has_speculative_draft_tokens=scenario.has_draft_tokens ), diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py b/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py index 37dc86cdd5fd..ff913bb1f8dd 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py @@ -1597,7 +1597,7 @@ def _create_generation_inputs(s: GenerationScenario) -> _GenerationInputs: request_ids=request_ids, prompt_lens=list(s.past_kv_lens), num_sparse_topk=num_sparse_topk, - num_heads_per_kv=s.q_heads_per_kv_head, + max_num_heads_per_kv=s.q_heads_per_kv_head, runtime_features=AttentionRuntimeFeatures(has_speculative_draft_tokens=s.has_draft_tokens), is_spec_decoding_enabled=s.has_draft_tokens, use_spec_decoding=s.has_draft_tokens, diff --git a/tests/unittest/_torch/executor/engine/test_metadata.py b/tests/unittest/_torch/executor/engine/test_metadata.py index f9f52ba5fe38..b1e0389d1431 100644 --- a/tests/unittest/_torch/executor/engine/test_metadata.py +++ b/tests/unittest/_torch/executor/engine/test_metadata.py @@ -10,7 +10,7 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures from tensorrt_llm._torch.model_config import KVCacheLayerSpec, ModelConfig from tensorrt_llm._torch.pyexecutor.engine.metadata import ( - _get_num_heads_per_kv, + _get_max_num_heads_per_kv, build_attention_metadata, ) @@ -55,7 +55,7 @@ def test_build_attention_metadata_resolves_model_derived_values() -> None: cache_indirection=None, ) - assert metadata.num_heads_per_kv == 4 + assert metadata.max_num_heads_per_kv == 4 assert metadata.enable_flash_mla assert metadata.enable_context_mla_with_cached_kv assert metadata.sparse_metadata_params is sparse_metadata_params @@ -87,7 +87,7 @@ def test_build_attention_metadata_forwards_shared_and_cache_inputs() -> None: kv_cache_manager=kv_cache_manager, draft_kv_cache_manager=draft_kv_cache_manager, enable_context_mla_with_cached_kv=True, - num_heads_per_kv=4, + max_num_heads_per_kv=4, ) assert metadata.max_num_requests == 4 @@ -98,28 +98,55 @@ def test_build_attention_metadata_forwards_shared_and_cache_inputs() -> None: assert metadata.kv_cache_manager is kv_cache_manager assert metadata.draft_kv_cache_manager is draft_kv_cache_manager assert metadata.enable_context_mla_with_cached_kv - assert metadata.num_heads_per_kv == 4 + assert metadata.max_num_heads_per_kv == 4 -def test_metadata_gqa_ratio_uses_model_provided_per_layer_kv_heads(): +@pytest.mark.parametrize( + ("head_pairs", "expected_ratio"), + [([(16, 8), (16, 1)], 16), ([(16, 8), (8, 2)], 4)], +) +def test_metadata_gqa_ratio_uses_model_provided_per_layer_heads( + head_pairs: list[tuple[int, int]], expected_ratio: int +) -> None: class StrictConfig: - num_hidden_layers = 12 - num_attention_heads = 16 + num_hidden_layers = 2 - def __getattribute__(self, name): - if name == "num_key_value_heads": + def __getattribute__(self, name: str) -> object: + if name in {"num_attention_heads", "num_key_value_heads"}: raise RuntimeError(f"global geometry must not be read: {name}") return super().__getattribute__(name) model_config = ModelConfig(pretrained_config=StrictConfig()) model_config.set_kv_cache_layer_specs( [ - KVCacheLayerSpec( - head_dim=256, - num_kv_heads=8 if (layer_idx + 1) % 6 else 1, - ) - for layer_idx in range(12) + KVCacheLayerSpec(head_dim=256, num_q_heads=q_heads, num_kv_heads=kv_heads) + for q_heads, kv_heads in head_pairs ] ) + metadata = build_attention_metadata( + model_config, + max_batch_size=4, + max_num_tokens=16, + max_beam_width=1, + attention_backend=_AttentionBackend, + attention_runtime_features=AttentionRuntimeFeatures(), + mapping=object(), + cache_indirection=None, + ) - assert _get_num_heads_per_kv(model_config) == 16 + assert metadata.max_num_heads_per_kv == expected_ratio + + +@pytest.mark.parametrize( + ("config", "expected_ratio"), + [ + (SimpleNamespace(num_attention_heads=8, num_key_value_heads=2), 4), + (SimpleNamespace(num_attention_heads=8, num_key_value_heads=[4, 2]), 4), + (SimpleNamespace(num_attention_heads=8, num_key_value_heads=[None, 0]), 1), + (SimpleNamespace(), 1), + ], +) +def test_metadata_gqa_ratio_preserves_flat_config_fallback( + config: SimpleNamespace, expected_ratio: int +) -> None: + assert _get_max_num_heads_per_kv(SimpleNamespace(pretrained_config=config)) == expected_ratio diff --git a/tests/unittest/_torch/executor/engine/test_runners.py b/tests/unittest/_torch/executor/engine/test_runners.py index 2bf30c67f1cf..886363d2a958 100644 --- a/tests/unittest/_torch/executor/engine/test_runners.py +++ b/tests/unittest/_torch/executor/engine/test_runners.py @@ -380,7 +380,7 @@ def test_no_kv_cache_runner_owns_and_reuses_attention_metadata() -> None: assert first.kv_cache_manager is None assert first.max_num_requests == 4 assert first.max_num_sequences == 4 - assert first.num_heads_per_kv == 4 + assert first.max_num_heads_per_kv == 4 assert first.block_ids_per_seq is None assert first.kv_block_ids_per_seq is None diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 02c2a2ef9d80..71488d562b1b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -677,6 +677,7 @@ def test_gemma4_hybrid_scales_by_num_pool_groups(): KVCacheLayerSpec( head_dim=256 if layer_type == "sliding_attention" else 128, num_kv_heads=1, + num_q_heads=1, attention_window=sliding_window if layer_type == "sliding_attention" else None, ) for layer_type in layer_types @@ -833,6 +834,7 @@ def test_pool_scaling_prevents_mmmu_pro_underestimation(): KVCacheLayerSpec( head_dim=256 if layer_type == "sliding_attention" else 128, num_kv_heads=1, + num_q_heads=1, attention_window=sliding_window if layer_type == "sliding_attention" else None, ) for layer_type in layer_types @@ -1019,6 +1021,7 @@ def test_v2_static_sizing_uses_model_provided_per_layer_geometry( KVCacheLayerSpec( head_dim=256 if (layer_idx + 1) % 6 else 512, num_kv_heads=8 if (layer_idx + 1) % 6 else 1, + num_q_heads=16, attention_window=128 if (layer_idx + 1) % 6 else None, ) for layer_idx in range(12) @@ -1031,7 +1034,7 @@ class StrictConfig: vocab_size = 262_144 def __getattribute__(self, name: str) -> object: - if name in {"head_dim", "num_key_value_heads"}: + if name in {"head_dim", "num_key_value_heads", "num_attention_heads"}: raise RuntimeError(f"global geometry must not be read: {name}") return super().__getattribute__(name) @@ -1071,6 +1074,14 @@ def get_num_attention_layers(self) -> int: expected_windows = [spec.attention_window for spec in layer_specs] assert layer_sizes == expected_layer_sizes assert attention_windows == expected_windows + full_layer_sizes, full_windows = _get_static_cache_size_layer_components( + model_config, + mapping, + max_seq_len=max_seq_len, + kv_cache_config=KvCacheConfig(), + ) + assert full_layer_sizes == expected_layer_sizes + assert full_windows == [None] * len(layer_specs) assert CacheCost.from_raw( KVCacheManagerV2.get_cache_size_per_token( model_config, @@ -1082,7 +1093,7 @@ def get_num_attention_layers(self) -> int: ) ) == CacheCost( slope=2 * 2048, - intercept=10 * 128 * expected_sliding_layer_size, + intercept=10 * 160 * expected_sliding_layer_size, ) @@ -1607,6 +1618,7 @@ def _create_manager_and_capture_config( kv_cache_config: KvCacheConfig, max_seq_len: int, manager_cls: type[KVCacheManager] | type[KVCacheManagerV2] = KVCacheManagerV2, + publish_layer_specs: bool = False, **factory_overrides: object, ) -> KvCacheConfig: """Run `_create_kv_cache_manager` with a recording subclass of `manager_cls` @@ -1622,6 +1634,16 @@ def __init__(self, kv_cache_config: KvCacheConfig, *args: object, **kwargs: obje model_config = Mock() model_config.pretrained_config = pretrained model_config.quant_config = None + if publish_layer_specs: + model_config.kv_cache_layer_specs = tuple( + KVCacheLayerSpec( + head_dim=256 if layer_type == _SLIDING else 512, + num_kv_heads=8 if layer_type == _SLIDING else 1, + num_q_heads=16, + attention_window=256 if layer_type == _SLIDING else None, + ) + for layer_type in pretrained.layer_types + ) factory_kwargs: dict[str, object] = dict( model_engine=None, @@ -1672,7 +1694,10 @@ def _eagle3_one_model_spec_config(num_draft_hidden_layers: int | None = None) -> ) -def test_create_kv_cache_manager_pool_ratio_arity_mismatch_keeps_single_window() -> None: +@pytest.mark.parametrize("publish_layer_specs", [False, True], ids=["layer_types", "layer_specs"]) +def test_create_kv_cache_manager_pool_ratio_arity_mismatch_keeps_single_window( + publish_layer_specs: bool, +) -> None: """A `pool_ratio` written for the single pool (one entry) does not match the two layer groups the derived windows would create. Rather than failing the manager's arity check at startup, the derivation is skipped with a warning @@ -1684,6 +1709,7 @@ def test_create_kv_cache_manager_pool_ratio_arity_mismatch_keeps_single_window() _mixed_schedule_pretrained([_SLIDING, _SLIDING, _FULL, _SLIDING]), kv_cache_config, max_seq_len=2048, + publish_layer_specs=publish_layer_specs, ) assert manager_config is kv_cache_config @@ -1694,7 +1720,10 @@ def test_create_kv_cache_manager_pool_ratio_arity_mismatch_keeps_single_window() assert "2 layer groups" in message -def test_create_kv_cache_manager_pool_ratio_per_layer_group_keeps_derivation() -> None: +@pytest.mark.parametrize("publish_layer_specs", [False, True], ids=["layer_types", "layer_specs"]) +def test_create_kv_cache_manager_pool_ratio_per_layer_group_keeps_derivation( + publish_layer_specs: bool, +) -> None: """One `pool_ratio` entry per derived layer group is the intended pairing.""" kv_cache_config = KvCacheConfig(pool_ratio=[0.5, 0.5]) @@ -1702,9 +1731,16 @@ def test_create_kv_cache_manager_pool_ratio_per_layer_group_keeps_derivation() - _mixed_schedule_pretrained([_SLIDING, _SLIDING, _FULL, _SLIDING]), kv_cache_config, max_seq_len=2048, + publish_layer_specs=publish_layer_specs, ) - assert manager_config.max_attention_window == [512, 512, 2048, 512] + sliding_window = 256 if publish_layer_specs else 512 + assert manager_config.max_attention_window == [ + sliding_window, + sliding_window, + 2048, + sliding_window, + ] assert manager_config.pool_ratio == [0.5, 0.5] @@ -1760,7 +1796,10 @@ def test_create_kv_cache_manager_keeps_appended_spec_layers_full_context( assert manager_config.max_attention_window == expected_windows -def test_create_kv_cache_manager_cross_pool_keeps_single_window_default() -> None: +@pytest.mark.parametrize("publish_layer_specs", [False, True], ids=["layer_types", "layer_specs"]) +def test_create_kv_cache_manager_cross_pool_keeps_single_window_default( + publish_layer_specs: bool, +) -> None: """The cross-attention pool stores encoder-side KV, which the decoder's `layer_types` do not describe.""" kv_cache_config = KvCacheConfig() @@ -1769,6 +1808,7 @@ def test_create_kv_cache_manager_cross_pool_keeps_single_window_default() -> Non _mixed_schedule_pretrained([_SLIDING, _SLIDING, _FULL, _SLIDING]), kv_cache_config, max_seq_len=2048, + publish_layer_specs=publish_layer_specs, kv_cache_type=batch_manager.CacheType.CROSS, num_layers=4, num_kv_heads=8, @@ -1777,3 +1817,20 @@ def test_create_kv_cache_manager_cross_pool_keeps_single_window_default() -> Non assert manager_config is kv_cache_config assert manager_config.max_attention_window is None + + +@pytest.mark.parametrize("guard", ["draft", "explicit_window"]) +def test_create_kv_cache_manager_layer_specs_preserve_window_guards(guard: str) -> None: + kv_cache_config = KvCacheConfig( + max_attention_window=[1024] if guard == "explicit_window" else None + ) + manager_config = _create_manager_and_capture_config( + _mixed_schedule_pretrained([_SLIDING, _SLIDING, _FULL, _SLIDING]), + kv_cache_config, + max_seq_len=2048, + publish_layer_specs=True, + is_draft=guard == "draft", + ) + + assert manager_config is kv_cache_config + assert manager_config.max_attention_window == ([1024] if guard == "explicit_window" else None) diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index f65d1268dd1a..a531ef4eb988 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -154,6 +154,7 @@ class _HeterogeneousGemma4TextConfig(Gemma4TextConfig): "head_dim", "num_global_key_value_heads", "num_key_value_heads", + "num_attention_heads", } def __getattribute__(self, name: str) -> object: @@ -172,12 +173,14 @@ def __getattribute__(self, name: str) -> object: def _make_heterogeneous_model_config( config_dict: dict[str, object] = GEMMA4_SMALL_CONFIG, + num_q_heads: list[int] | None = None, ) -> 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( + num_attention_heads=config_dict["num_attention_heads"], head_dim=( config_dict["head_dim"] if layer_type == "sliding_attention" @@ -193,6 +196,21 @@ def _make_heterogeneous_model_config( for layer_type in cfg.layer_types ] cfg.per_layer_attributes = {"head_dim", "num_key_value_heads"} + if num_q_heads is not None: + if isinstance(cfg.per_layer_config, list): + for layer_config, heads in zip(cfg.per_layer_config, num_q_heads, strict=True): + layer_config.num_attention_heads = heads + else: + cfg.per_layer_config = { + layer_idx: { + "head_dim": layer_config.head_dim, + "num_key_value_heads": layer_config.num_key_value_heads, + "num_attention_heads": heads, + } + for layer_idx, (layer_config, heads) in enumerate( + zip(cfg.per_layer_config, num_q_heads, strict=True) + ) + } mapping = Mapping(world_size=1, tp_size=1, rank=0) model_config = ModelConfig(pretrained_config=cfg, mapping=mapping) cfg._reject_global_geometry = True @@ -294,6 +312,31 @@ def test_model_instantiation_with_heterogeneous_transformers_config(self): self.assertEqual(layer_specs[layer_idx].head_dim, expected_head_dim) self.assertEqual(layer_specs[layer_idx].num_kv_heads, expected_num_kv_heads) + def test_model_instantiation_with_heterogeneous_query_heads(self) -> None: + """Attention modules and cache specs must agree on per-layer query heads.""" + num_q_heads = [ + 8 if layer_idx % 2 == 0 else 4 + for layer_idx in range(GEMMA4_SMALL_CONFIG["num_hidden_layers"]) + ] + model_config = _make_heterogeneous_model_config(num_q_heads=num_q_heads) + model = Gemma4ForCausalLM(model_config) + layer_specs = model.model_config.kv_cache_layer_specs + + for layer_idx, layer in enumerate(model.model.layers): + self.assertEqual(layer.self_attn.num_heads, num_q_heads[layer_idx]) + self.assertEqual(layer_specs[layer_idx].num_q_heads, num_q_heads[layer_idx]) + + def test_attention_requires_layer_idx_with_per_layer_config(self) -> None: + """A per-layer config must not silently select the first matching layer.""" + for per_layer_attributes in (None, set()): + with self.subTest(per_layer_attributes=per_layer_attributes): + config = SimpleNamespace(per_layer_config=[GEMMA4_SMALL_CONFIG]) + if per_layer_attributes is not None: + config.per_layer_attributes = per_layer_attributes + model_config = SimpleNamespace(pretrained_config=config) + with self.assertRaisesRegex(ValueError, "requires layer_idx"): + Gemma4Attention(model_config, is_sliding=True) + def test_model_instantiation_moe(self): """Create with MoE enabled and verify MoE layers exist.""" model_config = _make_model_config(GEMMA4_MOE_CONFIG) @@ -1027,7 +1070,6 @@ def __init__(self, *args: object, **kwargs: object) -> None: model_config=model_config, dtype=torch.bfloat16, is_draft=False, - kv_cache_type=object(), ) layer_types = model_config.pretrained_config.layer_types @@ -2760,6 +2802,10 @@ def test_attn_backend_dispatches_by_architecture(self): for is_sm100f, expected_class in expected_classes: with ( self.subTest(is_sm100f=is_sm100f), + unittest.mock.patch( + "tensorrt_llm._torch.attention.backends.utils.IS_FLASHINFER_AVAILABLE", + True, + ), unittest.mock.patch( "tensorrt_llm._torch.models.modeling_gemma4.is_sm_100f", return_value=is_sm100f, diff --git a/tests/unittest/_torch/modeling/test_modeling_llama.py b/tests/unittest/_torch/modeling/test_modeling_llama.py index 94a149dd252d..247f356f08d8 100644 --- a/tests/unittest/_torch/modeling/test_modeling_llama.py +++ b/tests/unittest/_torch/modeling/test_modeling_llama.py @@ -475,7 +475,7 @@ def test_llama_verification_with_kv_cache_relocation(self) -> None: kv_cache_manager=kv_cache_manager, request_ids=request_ids, prompt_lens=prompt_lens, - num_heads_per_kv=num_heads_per_kv, + max_num_heads_per_kv=num_heads_per_kv, ) position_ids = [torch.arange(0, input_ids.size(-1))] @@ -549,7 +549,7 @@ def run_forward(input_ids, position_ids, attn_metadata): use_spec_decoding=use_spec_decoding, is_spec_dec_tree=is_spec_dec_tree, is_spec_dec_dynamic_tree=is_spec_dec_dynamic_tree, - num_heads_per_kv=num_heads_per_kv, + max_num_heads_per_kv=num_heads_per_kv, ) cache = kv_cache_manager.kv_cache_map[request.py_request_id] assert cache.resize(input_ids.size(-1) + gen_input_ids_0.size(-1)) @@ -675,7 +675,7 @@ def run_forward(input_ids, position_ids, attn_metadata): use_spec_decoding=use_spec_decoding, is_spec_dec_tree=is_spec_dec_tree, is_spec_dec_dynamic_tree=False, - num_heads_per_kv=num_heads_per_kv, + max_num_heads_per_kv=num_heads_per_kv, ) attn_metadata_ref.spec_decoding_position_offsets = None From e64273eca93d4463ae0ad73914c8c221ca6ef059 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Sat, 10 Oct 2026 04:10:33 +0000 Subject: [PATCH 3/6] [https://nvbugs/6700677][test] Refresh Gemma4 geometry regression fixtures Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../test_layer_type_attention_windows.py | 1 - .../_torch/modeling/test_modeling_gemma4.py | 21 +++++++------------ 2 files changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/unittest/_torch/executor/test_layer_type_attention_windows.py b/tests/unittest/_torch/executor/test_layer_type_attention_windows.py index 0633af14b102..67f94ca6d15f 100644 --- a/tests/unittest/_torch/executor/test_layer_type_attention_windows.py +++ b/tests/unittest/_torch/executor/test_layer_type_attention_windows.py @@ -159,7 +159,6 @@ def test_invalid_sliding_window_falls_back_to_default(): # --------------------------------------------------------------------------- _FAMILY_PREDICATES = ( - "is_gemma4_hybrid", "is_kimi_linear", "is_mla", "is_nemotron_hybrid", diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index a531ef4eb988..f068a4809638 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -383,19 +383,14 @@ def test_per_layer_head_dim(self): config = model_config.pretrained_config for i, layer in enumerate(model.model.layers): - attn = layer.self_attn - if config.layer_types[i] == "sliding_attention": - self.assertEqual( - attn.head_dim, - config.head_dim, - f"Layer {i} (sliding) should have head_dim={config.head_dim}", - ) - else: - self.assertEqual( - attn.head_dim, - config.global_head_dim, - f"Layer {i} (full) should have head_dim={config.global_head_dim}", - ) + expected_head_dim = GEMMA4_SMALL_CONFIG[ + "head_dim" if config.layer_types[i] == "sliding_attention" else "global_head_dim" + ] + self.assertEqual( + layer.self_attn.head_dim, + expected_head_dim, + f"Layer {i} should have head_dim={expected_head_dim}", + ) def test_k_eq_v_attention(self): """Full attention layers with attention_k_eq_v=True should have v_norm.""" From 3165f6cd735b2ecec55cb1ec5756cd74f8407888 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Wed, 14 Oct 2026 13:36:11 +0000 Subject: [PATCH 4/6] [NVBUG-6700677][fix] Preserve target KV specs and align pool estimation Publish per-layer KV geometry from the target model so shared-KV assistant construction preserves the shared target specs. Resolve estimation windows through the manager helper and apply temporary pool sizing before budgeting. Validated 34 targeted tests with CI dependencies and 3 model-construction tests with Transformers 5.16.1. Both new regression scenarios fail on the prior production code. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_gemma4.py | 6 +- tensorrt_llm/_torch/pyexecutor/_util.py | 66 ++++------ .../kv_cache/test_kv_cache_estimation.py | 124 +++++++++++++++++- .../_torch/modeling/test_modeling_gemma4.py | 25 ++++ 4 files changed, 174 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 7200191c36ca..01eef6d0b49e 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -1144,9 +1144,6 @@ def forward( # --------------------------------------------------------------------------- class Gemma4TextModel(DecoderModel): def __init__(self, model_config: ModelConfig[Gemma4TextConfig]): - model_config.set_kv_cache_layer_specs( - _get_gemma4_kv_cache_layer_specs(model_config.pretrained_config) - ) super().__init__(model_config) config = self.model_config pretrained = config.pretrained_config @@ -1368,6 +1365,9 @@ 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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 4dd0f8a6f057..cbf65da37c43 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -781,62 +781,44 @@ def _derive_v2_layer_type_attention_windows( def _get_num_pool_groups_for_estimation( model_config: ModelConfig, max_seq_len: int, - fallback_attention_windows: Optional[List[Optional[int]]], + kv_cache_config: KvCacheConfig, ) -> int: - """Estimate V2 pool groups from KV storage windows and layer metadata. + """Estimate pool groups for the temporary V2 profiling manager. - Model sliding attention alone does not imply separate storage pools. - Preserve distinctions required by hybrid layer types and page sizes. + Hybrid recurrent/attention models also retain their distinct layer types + when automatic window derivation does not apply. """ config = model_config.pretrained_config - variable_geometry = _has_variable_kv_cache_geometry(model_config) - if (fallback_attention_windows is not None - and not is_hybrid_linear(config)): + attention_windows = kv_cache_config.max_attention_window + if attention_windows is None and _has_variable_kv_cache_geometry( + model_config): + attention_windows = _derive_v2_layer_type_attention_windows( + kv_cache_config, KVCacheManagerV2, model_config, max_seq_len) + if attention_windows is not None and not is_hybrid_linear(config): normalized_windows = _normalize_attention_windows( - fallback_attention_windows, max_seq_len) + attention_windows, max_seq_len) return 1 if normalized_windows is None else len(set(normalized_windows)) layer_types = getattr(config, "layer_types", None) - attention_windows = None - if (fallback_attention_windows is None and variable_geometry): - attention_windows = _derive_layer_type_attention_windows( - config, max_seq_len) - # Check whether KV storage uses configured or inferred attention windows, - # which can be independently configured from the model's use of sliding-window - # attention for computation. - kv_cache_uses_attention_windows = (fallback_attention_windows is not None - or attention_windows is not None) - - if attention_windows is not None: - normalized_windows = _normalize_attention_windows( - attention_windows, max_seq_len) - if normalized_windows is None: - return 1 - return len(set(normalized_windows)) - if isinstance(layer_types, (list, tuple)): - # These tags come from HF pretrained_config.layer_types. - # As above, sliding-window attention in the model does not imply windowed - # KV storage. Without windowed storage, sliding/full attention share a pool - # type unless their page sizes differ (Gemma4). + # Without distinct storage windows, sliding/full computation and + # different page geometries do not create separate window groups. pool_types = set() for layer_type in layer_types: - if (not kv_cache_uses_attention_windows and not variable_geometry - and - (_is_sliding_attention_layer(layer_type) - or getattr(layer_type, "name", - str(layer_type)).lower() == "full_attention")): + if attention_windows is None and ( + _is_sliding_attention_layer(layer_type) + or getattr(layer_type, "name", + str(layer_type)).lower() == "full_attention"): layer_type = "full_attention" pool_types.add(layer_type) if len(pool_types) > 1: return len(pool_types) - if fallback_attention_windows is not None: + if attention_windows is not None: normalized_windows = _normalize_attention_windows( - fallback_attention_windows, max_seq_len) + attention_windows, max_seq_len) if normalized_windows is not None: return len(set(normalized_windows)) - return 1 @@ -1631,7 +1613,7 @@ def _get_token_num_for_estimation(self) -> int: num_pool_groups = _get_num_pool_groups_for_estimation( model_cfg, self._model_engine.max_seq_len, - self._kv_cache_config.max_attention_window, + self._kv_cache_config, ) num_cache_blocks *= num_pool_groups @@ -1703,14 +1685,14 @@ def try_prepare_estimation(self) -> bool: "models, disable it.") if estimating_kv_cache: + # User-provided pool sizing can underprovision the temporary + # estimation cache and cause warmup to hang or fail. Override it + # before budgeting, then restore it in configure_kv_cache_capacity(). + self._kv_cache_config.pool_ratio = None estimate_max_tokens = self._get_token_num_for_estimation() max_tokens = min( estimate_max_tokens, self._kv_cache_config.max_tokens ) if self._kv_cache_config.max_tokens is not None else estimate_max_tokens - # User-provided pool sizing can underprovision the temporary - # estimation cache and cause warmup to hang or fail. Override it - # for estimation, then restore it in configure_kv_cache_capacity(). - self._kv_cache_config.pool_ratio = None self._kv_cache_config.avg_seq_len = self._max_seq_len if self._is_kv_cache_manager_v2: free_mem, _ = torch.cuda.mem_get_info() diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 71488d562b1b..c3d107f518fb 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -409,7 +409,7 @@ def _make_creator( c._model_engine.model = Mock(model_config=model_config) c._kv_cache_config = Mock( - free_gpu_memory_fraction=0.9, max_attention_window=max_attention_window + free_gpu_memory_fraction=0.9, max_attention_window=max_attention_window, pool_ratio=None ) # _get_token_num_for_estimation gates pool-group scaling on V2 (the @@ -774,11 +774,131 @@ def test_uniform_window_overrides_mixed_attention_metadata() -> None: ) assert ( - _get_num_pool_groups_for_estimation(SimpleNamespace(pretrained_config=config), 4096, [512]) + _get_num_pool_groups_for_estimation( + SimpleNamespace(pretrained_config=config), + 4096, + KvCacheConfig(max_attention_window=[512]), + ) == 1 ) +@pytest.mark.parametrize( + ("layer_types", "max_seq_len", "config_kwargs", "expected_windows", "expected_groups"), + [ + (None, 4096, {}, [128, 256, 4096], 3), + (["sliding_attention", "full_attention", "full_attention"], 4096, {}, [128, 256, 4096], 3), + (None, 256, {}, [128, 256, 256], 2), + (None, 4096, {"max_attention_window": [64]}, None, 1), + ( + ["sliding_attention", "full_attention", "full_attention"], + 4096, + {"pool_ratio": [1.0]}, + None, + 1, + ), + (None, 4096, {"pool_ratio": [0.25, 0.25, 0.5]}, [128, 256, 4096], 3), + ], + ids=[ + "missing_layer_types", + "conflicting_layer_types", + "clamped", + "explicit_window", + "pool_ratio_mismatch", + "pool_ratio_matches", + ], +) +def test_pool_group_estimation_uses_resolved_spec_windows( + layer_types: list[str] | None, + max_seq_len: int, + config_kwargs: dict[str, list[int] | list[float]], + expected_windows: list[int] | None, + expected_groups: int, +) -> None: + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace( + num_hidden_layers=3, + layer_types=layer_types, + sliding_window=512, + ), + kv_cache_layer_specs=tuple( + KVCacheLayerSpec( + head_dim=head_dim, num_q_heads=8, num_kv_heads=2, attention_window=window + ) + for head_dim, window in [(128, 128), (256, 256), (128, None)] + ), + ) + kv_cache_config = KvCacheConfig(**config_kwargs) + assert ( + _derive_v2_layer_type_attention_windows( + kv_cache_config, KVCacheManagerV2, model_config, max_seq_len + ) + == expected_windows + ) + assert ( + _get_num_pool_groups_for_estimation(model_config, max_seq_len, kv_cache_config) + == expected_groups + ) + + +def test_pool_group_estimation_keeps_uniform_geometry_profiling_full_context() -> None: + pretrained = _mixed_schedule_pretrained( + ["sliding_attention", "full_attention"], + sliding_window=128, + ) + model_config = SimpleNamespace(pretrained_config=pretrained) + kv_cache_config = KvCacheConfig() + manager_config = _create_manager_and_capture_config( + pretrained, + kv_cache_config, + max_seq_len=4096, + estimating_kv_cache=True, + ) + assert manager_config.max_attention_window is None + assert _get_num_pool_groups_for_estimation(model_config, 4096, kv_cache_config) == 1 + assert _derive_v2_layer_type_attention_windows( + kv_cache_config, KVCacheManagerV2, model_config, 4096 + ) == [128, 4096] + + +def test_estimation_budget_uses_temporary_pool_ratio() -> None: + max_seq_len = 4096 + creator = _make_creator( + 32, + [_make_mock_request(max_seq_len - 1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + ) + model_config = creator._model_engine.model.model_config + model_config.attn_backend = "FLASHINFER" + model_config.is_encoder_decoder = False + model_config.kv_cache_layer_specs = tuple( + KVCacheLayerSpec( + head_dim=head_dim, + num_q_heads=8, + num_kv_heads=2, + attention_window=window, + ) + for head_dim, window in [(128, 128), (256, 256), (128, None)] + ) + creator._kv_cache_config = KvCacheConfig(pool_ratio=[1.0]) + creator._skip_est = False + creator._max_seq_len = max_seq_len + creator._max_gpu_total_bytes_in = None + single_pool_tokens = creator._get_token_num_for_estimation() + + assert creator.try_prepare_estimation() + assert creator._kv_cache_config.pool_ratio is None + assert creator._kv_cache_config.max_tokens == 3 * single_pool_tokens + assert _derive_v2_layer_type_attention_windows( + creator._kv_cache_config, + KVCacheManagerV2, + model_config, + max_seq_len, + ) == [128, 256, max_seq_len] + + def test_vswa_max_attention_window_fallback_scales(): """When layer_types is absent but kv_cache_config.max_attention_window is a heterogeneous list (VSWA), we still scale by the number of distinct diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index f068a4809638..6b3a4022afa1 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -751,6 +751,31 @@ def test_ordered_embedding_combines_vocab_parallel_shards(self): torch.testing.assert_close(actual, expected) + def test_assistant_preserves_target_kv_cache_layer_specs(self) -> None: + target_config = _make_heterogeneous_model_config( + num_q_heads=[8] * GEMMA4_SMALL_CONFIG["num_hidden_layers"] + ) + target = Gemma4ForCausalLM(target_config) + target_specs = target_config.kv_cache_layer_specs + self.assertIsNotNone(target_specs) + registry = {} + target_config.extra_attrs["shared_registry"] = registry + + assistant_config = _make_assistant_model_config() + assistant_config.extra_attrs = target_config.extra_attrs + assistant_specs = _get_gemma4_kv_cache_layer_specs( + assistant_config.pretrained_config.text_config + ) + self.assertNotEqual(len(target_specs), len(assistant_specs)) + self.assertNotEqual(target_specs[0].num_q_heads, assistant_specs[0].num_q_heads) + + assistant = Gemma4AssistantForCausalLM(assistant_config) + + self.assertIs(assistant.model.model_config.extra_attrs, target_config.extra_attrs) + self.assertIs(target.model_config.kv_cache_layer_specs, target_specs) + self.assertIs(assistant.model.model_config.kv_cache_layer_specs, target_specs) + self.assertIs(assistant.model.model_config.extra_attrs["shared_registry"], registry) + def test_assistant_uses_target_kv_sources(self): model_config = _make_assistant_model_config() model_config.extra_attrs["_speculative_position_headroom"] = 2 * 4 From 41d5dacf7dc1ab8e2721f101b3fb69788fad726f Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:52:26 -0700 Subject: [PATCH 5/6] [https://nvbugs/6700677][fix] Keep Gemma4 KV layer specs per model config Store the published per-layer KV geometry on each ModelConfig instead of the extra_attrs registry that one-model drafters share with the target, so a drafter neither sees nor overwrites the target's specs. Also address the remaining review nits: clarify the Gemma4 backend-default and attention comments, let the backend decide sliding-layer RoPE fusion, describe max_num_heads_per_kv precisely, add a message to the KV-head count assert, and make the metadata ratio test discriminate against last-layer logic. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- .../_torch/attention/backends/interface.py | 3 ++- tensorrt_llm/_torch/model_config.py | 15 ++++++++--- tensorrt_llm/_torch/models/modeling_gemma4.py | 20 +++++++------- .../kv_cache/kv_cache_manager_v2.py | 5 +++- .../_torch/executor/engine/test_metadata.py | 2 +- .../_torch/modeling/test_modeling_gemma4.py | 27 ++++++++++++++++++- 6 files changed, 56 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index a124fb319011..1f50a31c329c 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -193,7 +193,8 @@ 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 largest query-to-KV head ratio across layers sharing this metadata. + # 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 diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index b4acdf012d86..21135d233db4 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -309,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 @@ -344,12 +352,12 @@ class ModelConfig(Generic[TConfig]): @property def kv_cache_layer_specs(self) -> Optional[Tuple[KVCacheLayerSpec, ...]]: """Return model-provided per-layer KV-cache geometry, if any.""" - return self.extra_attrs.get("kv_cache_layer_specs") + 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.extra_attrs["kv_cache_layer_specs"] = tuple(layer_specs) + self._kv_cache_layer_specs = tuple(layer_specs) def __setattr__(self, key, value): """ @@ -357,13 +365,14 @@ def __setattr__(self, key, value): 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) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index 01eef6d0b49e..27322ffacf61 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -392,8 +392,9 @@ def __init__( # 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, ) # Fix proportional RoPE for full-attention layers. @@ -439,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" @@ -1374,10 +1375,11 @@ def __init__( def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """Gemma4-specific defaults. - 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. + 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) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 20594d4abd07..0c5bbe950e52 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1552,7 +1552,10 @@ def __init__( (num_kv_heads + tp_size - 1) // tp_size for _ in range(self.num_layers) ] else: - assert len(num_kv_heads) == self.num_layers + assert len(num_kv_heads) == self.num_layers, ( + f"num_kv_heads list length ({len(num_kv_heads)}) must match " + f"num_layers ({self.num_layers})" + ) def append_to_kv_heads_per_layer( num_kv_heads_per_layer: List[int], kv_head: Optional[int] diff --git a/tests/unittest/_torch/executor/engine/test_metadata.py b/tests/unittest/_torch/executor/engine/test_metadata.py index b1e0389d1431..475cdf6b5d21 100644 --- a/tests/unittest/_torch/executor/engine/test_metadata.py +++ b/tests/unittest/_torch/executor/engine/test_metadata.py @@ -103,7 +103,7 @@ def test_build_attention_metadata_forwards_shared_and_cache_inputs() -> None: @pytest.mark.parametrize( ("head_pairs", "expected_ratio"), - [([(16, 8), (16, 1)], 16), ([(16, 8), (8, 2)], 4)], + [([(16, 8), (16, 1)], 16), ([(16, 8), (8, 2)], 4), ([(16, 1), (16, 8)], 16)], ) def test_metadata_gqa_ratio_uses_model_provided_per_layer_heads( head_pairs: list[tuple[int, int]], expected_ratio: int diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 6b3a4022afa1..e45f67f4e3e0 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -773,9 +773,34 @@ def test_assistant_preserves_target_kv_cache_layer_specs(self) -> None: self.assertIs(assistant.model.model_config.extra_attrs, target_config.extra_attrs) self.assertIs(target.model_config.kv_cache_layer_specs, target_specs) - self.assertIs(assistant.model.model_config.kv_cache_layer_specs, target_specs) + self.assertIsNone(assistant.model.model_config.kv_cache_layer_specs) self.assertIs(assistant.model.model_config.extra_attrs["shared_registry"], registry) + def test_draft_sharing_extra_attrs_keeps_its_own_kv_cache_layer_specs(self) -> None: + target_config = _make_heterogeneous_model_config( + num_q_heads=[8] * GEMMA4_SMALL_CONFIG["num_hidden_layers"] + ) + Gemma4ForCausalLM(target_config) + target_specs = target_config.kv_cache_layer_specs + self.assertIsNotNone(target_specs) + + draft_layers = 3 + draft_config = _make_heterogeneous_model_config( + {**GEMMA4_SMALL_CONFIG, "num_hidden_layers": draft_layers}, + num_q_heads=[4] * draft_layers, + ) + draft_config.extra_attrs = target_config.extra_attrs + # A drafter of another architecture publishes no specs and must not + # inherit the target's through the shared registry. + self.assertIsNone(draft_config.kv_cache_layer_specs) + + Gemma4ForCausalLM(draft_config) + + self.assertIs(target_config.kv_cache_layer_specs, target_specs) + draft_specs = draft_config.kv_cache_layer_specs + self.assertEqual(len(draft_specs), draft_layers) + self.assertTrue(all(spec.num_q_heads == 4 for spec in draft_specs)) + def test_assistant_uses_target_kv_sources(self): model_config = _make_assistant_model_config() model_config.extra_attrs["_speculative_position_headroom"] = 2 * 4 From 47355a04925d588bc879c4ffa78fb8865fe0b3e8 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:04:11 -0700 Subject: [PATCH 6/6] [https://nvbugs/6700677][fix] Let QKNormRoPEAttention accept rope_fusion=None Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- tensorrt_llm/_torch/attention/qk_norm_attention.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/attention/qk_norm_attention.py b/tensorrt_llm/_torch/attention/qk_norm_attention.py index 52aa7ebe1e15..0f370b292640 100644 --- a/tensorrt_llm/_torch/attention/qk_norm_attention.py +++ b/tensorrt_llm/_torch/attention/qk_norm_attention.py @@ -164,7 +164,7 @@ 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, ): @@ -178,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"