diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 576927c59475..05864fee38a4 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -88,6 +88,7 @@ Models that select the V2 manager by default: | Hybrid Mamba (NemotronH, Qwen3-Next) | Attention KV and Mamba state pools must be sized together | | DeepSeek-V4 | Sparse attention attaches auxiliary per-layer buffers | | GPT-OSS | Sliding window on every other layer (VSWA), so the sliding-window and full-attention pools are sized independently | +| Gemma3 / Gemma4 (text and multimodal) | Alternating sliding-window and full-attention layers (VSWA); same independent pool sizing | Separately, Gemma4 hybrid attention and sparse-attention models are routed to V2 unconditionally: their per-layer buffer layouts cannot be represented by V1's diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 7bbcd0ce9186..6727f9e4a63d 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -2548,6 +2548,19 @@ def get_preferred_kv_cache_manager_version( """Prefer KV cache manager V2 for DeepSeek-V4.""" return "V2" + @classmethod + def get_preferred_transceiver_runtime( + cls, pretrained_config: object | None = None + ) -> Literal["PYTHON"]: + """Prefer the Python transceiver in disaggregated serving. + + DeepSeek-V4 runs DeepseekV4CacheManager, a KVCacheManagerV2 + subclass that the C++ transceiver cannot drive; the disaggregated + tests pin NIXL + PYTHON for the same reason. This routes the + fully-'auto' path to that combination. + """ + return "PYTHON" + def __init__(self, model_config: ModelConfig[PretrainedConfig]): model_config = _normalize_deepseek_v4_nvfp4_mixed_precision_config(model_config) self.mapping_with_cp = None diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py index 0e41346e2f3e..628d1939797a 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3.py @@ -292,6 +292,17 @@ def __init__( hidden_size=model_config.pretrained_config.hidden_size, vocab_size=model_config.pretrained_config.vocab_size) + @classmethod + def get_preferred_kv_cache_manager_version(cls, + pretrained_config: Any = None + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for Gemma3's VSWA layout. + + V2 sizes the sliding-window and full-attention pools independently + instead of statically dividing memory between them. + """ + return "V2" + def _get_token_type_mask(self, image_token_mask: torch.BoolTensor): device = image_token_mask.device sequence_length = len(image_token_mask) diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py index 7600c6167fac..422c943837ee 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3vl.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3vl.py @@ -1,6 +1,6 @@ import copy import dataclasses -from typing import List, Optional, Tuple +from typing import Any, List, Literal, Optional, Tuple import torch from transformers import (AutoProcessor, AutoTokenizer, Gemma3Config, @@ -177,6 +177,22 @@ def forward(self, vision_outputs: torch.Tensor): )) class Gemma3VLM(PreTrainedModel): + @classmethod + def get_preferred_kv_cache_manager_version(cls, + pretrained_config: Any = None + ) -> Literal["V2"]: + """Prefer KV cache manager V2 — same VSWA rationale as + Gemma3ForCausalLM (the wrapped text model).""" + return "V2" + + @classmethod + def get_preferred_transceiver_runtime( + cls, + pretrained_config: Any = None, + ) -> Optional[Literal["CPP", "PYTHON"]]: + """Prefer the Python transceiver so disaggregated serving over NIXL keeps V2.""" + return "PYTHON" + def __init__(self, model_config: ModelConfig[Gemma3Config]): if _is_mm_disagg(): raise NotImplementedError( diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index a9aa54d6029c..5a3540f4a4ea 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -16,7 +16,7 @@ import dataclasses import math -from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -63,6 +63,8 @@ from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + from .modeling_gemma4mm import Gemma4ForConditionalGeneration _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0" @@ -1247,7 +1249,7 @@ def __init__( super().__init__(Gemma4TextModel(model_config), model_config) @classmethod - def get_model_defaults(cls, llm_args) -> dict: + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """Gemma4-specific defaults. FlashInfer backend is required for hybrid attention (per-layer @@ -1258,6 +1260,23 @@ def get_model_defaults(cls, llm_args) -> dict: "attn_backend": "FLASHINFER", } + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """Prefer KV cache manager V2 for Gemma4's VSWA layout. + + Hybrid-attention checkpoints (per-layer head_dim) are routed to V2 + unconditionally regardless of this preference. + """ + return "V2" + + @classmethod + def get_preferred_transceiver_runtime( + cls, + pretrained_config: Any = None, + ) -> Optional[Literal["CPP", "PYTHON"]]: + """Prefer the Python transceiver so disaggregated serving over NIXL keeps V2.""" + return "PYTHON" + def _get_token_type_mask(self, mm_token_type_ids: torch.Tensor): """Build bidirectional attention mask from mm_token_type_ids. diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index ab211c1a81a9..5b7940c0af9b 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -25,7 +25,7 @@ import math from collections.abc import Sequence from itertools import groupby -from typing import Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple import torch import transformers @@ -62,6 +62,9 @@ from .modeling_multimodal_utils import _MULTIMODAL_ENV_NAME, _is_mm_disagg from .modeling_utils import ModelConfig, filter_weights, register_auto_model +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0" if Version(transformers.__version__) < Version(_MIN_TRANSFORMERS_FOR_GEMMA4): raise ImportError( @@ -557,12 +560,30 @@ class Gemma4MultimodalModelBase(MultimodalModelMixin, PreTrainedModel): supports_encoder_cache = True @classmethod - def get_model_defaults(cls, llm_args) -> dict: + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """Gemma4-specific defaults — see Gemma4ForCausalLM.get_model_defaults.""" return { "attn_backend": "FLASHINFER", } + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """Prefer KV cache manager V2 — see Gemma4ForCausalLM.""" + return "V2" + + @classmethod + def get_preferred_transceiver_runtime( + cls, + pretrained_config: Any = None, + ) -> Optional[Literal["CPP", "PYTHON"]]: + """Prefer the Python transceiver so disaggregated serving over NIXL keeps V2. + + Multimodal disaggregated serving is currently rejected in __init__, + but if it lands, the NIXL route must resolve to the Python + transceiver for _resolve_kv_cache_manager_v2_auto to keep V2. + """ + return "PYTHON" + def _check_and_adjust_experts_implementation(self, *args, **kwargs): # transformers 5.x ``PreTrainedModel.__init__`` calls this with an # ``experts_implementation`` argument and fails for VL wrapper models diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index e1978a577a3b..ec0fd8b36363 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -22,7 +22,7 @@ import copy import dataclasses import os -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Literal, Optional, Tuple from typing import Mapping as TMapping import torch @@ -2004,6 +2004,27 @@ def _fold_gemma_boundary_norm_weights(weights): class MiniMaxM3ForCausalLM(DecoderModelForCausalLM[MiniMaxM3Model, PretrainedConfig]): """Text-only M3 model.""" + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """Prefer KV cache manager V2 for MiniMax-M3. + + Sparse attention already routes M3 to a V2-core manager + unconditionally; declaring the preference keeps + ``kv_cache_config.use_kv_cache_manager_v2`` consistent with the + manager actually in use. + """ + return "V2" + + @classmethod + def get_preferred_transceiver_runtime(cls, pretrained_config: Any = None) -> Literal["PYTHON"]: + """Prefer the Python transceiver in disaggregated serving. + + M3 runs a V2-core cache manager, which the C++ transceiver cannot + drive; the KV-transfer unit test exercises the Python transceiver + directly. This routes the fully-'auto' path to that combination. + """ + return "PYTHON" + def __init__(self, model_config: "ModelConfig[PretrainedConfig]"): raw_pretrained = model_config.pretrained_config if is_minimax_m3_vl_config(raw_pretrained): diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2d3d96f0aabb..71963e36e551 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -718,12 +718,69 @@ def test_registered_models_prefer_v2(self): "Qwen3_5ForCausalLM", "Qwen3_5MoeForConditionalGeneration", "Qwen3_5ForConditionalGeneration", + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", + "Gemma3ForCausalLM", + "Gemma3ForConditionalGeneration", + "Gemma4ForCausalLM", + "Gemma4ForConditionalGeneration", + "Gemma4UnifiedForConditionalGeneration", ) for architecture in architectures: model_cls = get_registered_model_class(architecture) assert model_cls is not None assert model_cls.get_preferred_kv_cache_manager_version() == "V2" + def test_registered_models_keep_v2_on_nixl(self): + """Models preferring V2 and the Python transceiver keep V2 on NIXL. + + Both sentinels start at 'auto'; production resolves the transceiver + runtime first, then the KV cache manager. MiniMax-M2 is absent from + this list: it silently resolves to V1 on this route (its + disaggregated serving is unvalidated -- the missing preference is + deliberate). + """ + from tensorrt_llm._torch.models.modeling_utils import \ + get_registered_model_class + + architectures = ( + "DeepseekV3ForCausalLM", + "DeepseekV32ForCausalLM", + "GlmMoeDsaForCausalLM", + "MistralLarge3ForCausalLM", + "GptOssForCausalLM", + "KimiK25ForConditionalGeneration", + "NemotronHForCausalLM", + "NemotronHPuzzleForCausalLM", + "Qwen3NextForCausalLM", + "Qwen3_5MoeForCausalLM", + "Qwen3_5ForCausalLM", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5ForConditionalGeneration", + "DeepseekV4ForCausalLM", + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", + "Gemma3ForCausalLM", + "Gemma3ForConditionalGeneration", + "Gemma4ForCausalLM", + "Gemma4ForConditionalGeneration", + "Gemma4UnifiedForConditionalGeneration", + ) + for architecture in architectures: + model_cls = get_registered_model_class(architecture) + assert model_cls is not None, architecture + + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="auto"), + ) + _resolve_transceiver_runtime_auto(llm_args, model_cls) + assert _resolve_kv_cache_manager_v2_auto( + llm_args, model_cls) is True, architecture + assert (llm_args.cache_transceiver_config.transceiver_runtime == + "PYTHON"), architecture + @pytest.mark.cpu_only def test_KvCacheConfig_declaration():