From b688572999ae7dd29c8ceefb67e71ad9f25b422d Mon Sep 17 00:00:00 2001 From: Eric Tsai Date: Thu, 6 Aug 2026 20:27:19 -0700 Subject: [PATCH 1/8] [None][feat] Enable KVCacheManagerV2 by default for Gemma3 and Gemma4 Signed-off-by: Eric Tsai --- tensorrt_llm/_torch/models/modeling_gemma3.py | 15 +++++++++++++++ tensorrt_llm/_torch/models/modeling_gemma4.py | 10 ++++++++++ .../_torch/modeling/test_modeling_gemma3.py | 9 +++++++++ .../_torch/modeling/test_modeling_gemma4.py | 10 ++++++++++ 4 files changed, 44 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py index 0e41346e2f3e..866cef451dea 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3.py @@ -292,6 +292,21 @@ def __init__( hidden_size=model_config.pretrained_config.hidden_size, vocab_size=model_config.pretrained_config.vocab_size) + @classmethod + def get_model_defaults(cls, llm_args) -> dict: + return { + "kv_cache_config": { + "use_kv_cache_manager_v2": True, + }, + } + + @classmethod + def get_preferred_transceiver_runtime( + cls, + pretrained_config=None, + ): + return "PYTHON" + 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_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index a9aa54d6029c..46cdaabbb236 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -1256,8 +1256,18 @@ def get_model_defaults(cls, llm_args) -> dict: """ return { "attn_backend": "FLASHINFER", + "kv_cache_config": { + "use_kv_cache_manager_v2": True, + }, } + @classmethod + def get_preferred_transceiver_runtime( + cls, + pretrained_config=None, + ): + 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/tests/unittest/_torch/modeling/test_modeling_gemma3.py b/tests/unittest/_torch/modeling/test_modeling_gemma3.py index f252f4cbd2d6..1317b14a6d30 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma3.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma3.py @@ -684,3 +684,12 @@ def test_gemma3_local_context_mask_multi_image(self) -> None: ], device=device).bool() torch.testing.assert_close(attention_mask, expected_attention_mask) + + +def test_gemma3_model_defaults_select_v2(): + defaults = Gemma3ForCausalLM.get_model_defaults(object()) + assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True + + +def test_gemma3_prefers_python_transceiver(): + assert Gemma3ForCausalLM.get_preferred_transceiver_runtime() == "PYTHON" diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 5eff2c644555..752d12309d3e 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -4063,5 +4063,15 @@ def test_vision_attn_metadata_max_num_requests_is_not_one(self): ) +def test_gemma4_model_defaults_select_v2(): + defaults = Gemma4ForCausalLM.get_model_defaults(object()) + assert defaults["attn_backend"] == "FLASHINFER" + assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True + + +def test_gemma4_prefers_python_transceiver(): + assert Gemma4ForCausalLM.get_preferred_transceiver_runtime() == "PYTHON" + + if __name__ == "__main__": unittest.main() From 4b83678d64aa130ec69de29470df229e2487ef5a Mon Sep 17 00:00:00 2001 From: Eric Tsai Date: Thu, 6 Aug 2026 20:57:29 -0700 Subject: [PATCH 2/8] Add use_kv_cache_manager_v2=True to Gemma4ForConditionalGeneration.get_model_defaults Signed-off-by: Eric Tsai --- tensorrt_llm/_torch/models/modeling_gemma4mm.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index ab211c1a81a9..55ca23f29801 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -561,6 +561,9 @@ def get_model_defaults(cls, llm_args) -> dict: """Gemma4-specific defaults — see Gemma4ForCausalLM.get_model_defaults.""" return { "attn_backend": "FLASHINFER", + "kv_cache_config": { + "use_kv_cache_manager_v2": True, + }, } def _check_and_adjust_experts_implementation(self, *args, **kwargs): From aae7372b9560f2a08cb58f945ad59a8e712677af Mon Sep 17 00:00:00 2001 From: Eric Tsai Date: Mon, 10 Aug 2026 00:15:14 -0700 Subject: [PATCH 3/8] [None][fix] Address review: cover Gemma3 VLM, MM transceiver hook, resolution-path tests Signed-off-by: Eric Tsai --- tensorrt_llm/_torch/models/modeling_gemma3.py | 17 ++++-- .../_torch/models/modeling_gemma3vl.py | 27 ++++++++- tensorrt_llm/_torch/models/modeling_gemma4.py | 11 ++-- .../_torch/models/modeling_gemma4mm.py | 20 ++++++- .../_torch/modeling/test_modeling_gemma3.py | 59 ++++++++++++++++--- .../_torch/modeling/test_modeling_gemma4.py | 54 +++++++++++++++-- 6 files changed, 164 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py index 866cef451dea..ff2e8fe7db7b 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3.py @@ -1,5 +1,5 @@ import math -from typing import Any, Dict, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple import torch from torch import nn @@ -26,6 +26,9 @@ from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, register_auto_model) +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + class Gemma3TextScaledWordEmbedding(Embedding): @@ -293,7 +296,12 @@ def __init__( vocab_size=model_config.pretrained_config.vocab_size) @classmethod - def get_model_defaults(cls, llm_args) -> dict: + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: + """Gemma3-specific defaults. + + Enables the V2 KV-cache manager, which sizes the sliding-window and + full-attention pools independently for Gemma3's VSWA layout. + """ return { "kv_cache_config": { "use_kv_cache_manager_v2": True, @@ -303,8 +311,9 @@ def get_model_defaults(cls, llm_args) -> dict: @classmethod def get_preferred_transceiver_runtime( cls, - pretrained_config=None, - ): + 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, image_token_mask: torch.BoolTensor): diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py index 7600c6167fac..2b46e2bd090d 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 TYPE_CHECKING, Any, List, Literal, Optional, Tuple import torch from transformers import (AutoProcessor, AutoTokenizer, Gemma3Config, @@ -27,6 +27,9 @@ from .modeling_siglip import SiglipVisionModel from .modeling_utils import ModelConfig, filter_weights, register_auto_model +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + class Gemma3InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): @@ -177,6 +180,28 @@ def forward(self, vision_outputs: torch.Tensor): )) class Gemma3VLM(PreTrainedModel): + @classmethod + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: + """Gemma3 VLM defaults — same V2 KV-cache manager selection as + Gemma3ForCausalLM (the wrapped text model). + + The FLASHINFER attention backend is not set here: the text sub-model + config pins it internally in get_sub_model_config. + """ + return { + "kv_cache_config": { + "use_kv_cache_manager_v2": True, + }, + } + + @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 46cdaabbb236..c6050de95071 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 @@ -1264,8 +1266,9 @@ def get_model_defaults(cls, llm_args) -> dict: @classmethod def get_preferred_transceiver_runtime( cls, - pretrained_config=None, - ): + 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): diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 55ca23f29801..fd3184e6d6b1 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,7 +560,7 @@ 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", @@ -566,6 +569,19 @@ def get_model_defaults(cls, llm_args) -> dict: }, } + @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/tests/unittest/_torch/modeling/test_modeling_gemma3.py b/tests/unittest/_torch/modeling/test_modeling_gemma3.py index 1317b14a6d30..29c25ca256b8 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma3.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma3.py @@ -2,6 +2,7 @@ from copy import deepcopy from dataclasses import dataclass +import pytest import torch from _torch.helpers import make_hf_hybrid_cache_for_tests from parameterized import parameterized @@ -18,8 +19,16 @@ from tensorrt_llm._torch.models.checkpoints.hf.gemma3_weight_mapper import \ Gemma3HfWeightMapper from tensorrt_llm._torch.models.modeling_gemma3 import Gemma3ForCausalLM +from tensorrt_llm._torch.models.modeling_gemma3vl import Gemma3VLM +from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmapiKvCacheConfig +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, + _resolve_transceiver_runtime_auto, + apply_model_defaults_to_llm_args) from tensorrt_llm.mapping import Mapping GEMMA3_1B_CONFIG = { @@ -686,10 +695,46 @@ def test_gemma3_local_context_mask_multi_image(self) -> None: torch.testing.assert_close(attention_mask, expected_attention_mask) -def test_gemma3_model_defaults_select_v2(): - defaults = Gemma3ForCausalLM.get_model_defaults(object()) - assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True - - -def test_gemma3_prefers_python_transceiver(): - assert Gemma3ForCausalLM.get_preferred_transceiver_runtime() == "PYTHON" +@pytest.mark.parametrize( + ("architecture", "model_cls"), + [("Gemma3ForCausalLM", Gemma3ForCausalLM), + ("Gemma3ForConditionalGeneration", Gemma3VLM)], +) +def test_gemma3_defaults_resolve_to_v2(architecture: str, + model_cls: type) -> None: + """The checkpoint's architectures[0] must map to the class carrying the + defaults, and those defaults must survive the production resolution + ordering (transceiver first, then KV-cache manager) on the NIXL + disaggregated route — the route where a missing transceiver preference + silently downgrades the V2 default back to V1.""" + assert MODEL_CLASS_MAPPING[architecture] is model_cls + + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="auto"), + ) + original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 + defaults = model_cls.get_model_defaults(llm_args) + apply_model_defaults_to_llm_args(llm_args, defaults) + _resolve_transceiver_runtime_auto(llm_args, model_cls) + _resolve_kv_cache_manager_v2_auto(llm_args, + defaults, + original_setting=original_setting) + + assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + + +@pytest.mark.parametrize("model_cls", [Gemma3ForCausalLM, Gemma3VLM]) +def test_gemma3_explicit_user_setting_wins(model_cls: type) -> None: + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_config=LlmapiKvCacheConfig(use_kv_cache_manager_v2=False)) + original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 + defaults = model_cls.get_model_defaults(llm_args) + apply_model_defaults_to_llm_args(llm_args, defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, + defaults, + original_setting=original_setting) + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 752d12309d3e..93e83a067284 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -26,6 +26,7 @@ from types import SimpleNamespace from typing import TYPE_CHECKING +import pytest import torch from transformers import AutoConfig, Gemma4Config, Gemma4TextConfig @@ -45,7 +46,16 @@ Gemma4TextModel, Gemma4TextScaledWordEmbedding, ) +from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4ForConditionalGeneration +from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING from tensorrt_llm._utils import is_sm_100f +from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, TorchLlmArgs +from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmapiKvCacheConfig +from tensorrt_llm.llmapi.llm_utils import ( + _resolve_kv_cache_manager_v2_auto, + _resolve_transceiver_runtime_auto, + apply_model_defaults_to_llm_args, +) from tensorrt_llm.mapping import Mapping if TYPE_CHECKING: @@ -4063,14 +4073,46 @@ def test_vision_attn_metadata_max_num_requests_is_not_one(self): ) -def test_gemma4_model_defaults_select_v2(): - defaults = Gemma4ForCausalLM.get_model_defaults(object()) - assert defaults["attn_backend"] == "FLASHINFER" - assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True +@pytest.mark.parametrize( + ("architecture", "model_cls"), + [ + ("Gemma4ForCausalLM", Gemma4ForCausalLM), + ("Gemma4ForConditionalGeneration", Gemma4ForConditionalGeneration), + ], +) +def test_gemma4_defaults_resolve_to_v2(architecture: str, model_cls: type) -> None: + """The checkpoint's architectures[0] must map to the class carrying the + defaults, and those defaults must survive the production resolution + ordering (transceiver first, then KV-cache manager) on the NIXL + disaggregated route — the route where a missing transceiver preference + silently downgrades the V2 default back to V1.""" + assert MODEL_CLASS_MAPPING[architecture] is model_cls + + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig(backend="NIXL", transceiver_runtime="auto"), + ) + original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 + defaults = model_cls.get_model_defaults(llm_args) + apply_model_defaults_to_llm_args(llm_args, defaults) + _resolve_transceiver_runtime_auto(llm_args, model_cls) + _resolve_kv_cache_manager_v2_auto(llm_args, defaults, original_setting=original_setting) + assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + assert llm_args.attn_backend == "FLASHINFER" -def test_gemma4_prefers_python_transceiver(): - assert Gemma4ForCausalLM.get_preferred_transceiver_runtime() == "PYTHON" + +@pytest.mark.parametrize("model_cls", [Gemma4ForCausalLM, Gemma4ForConditionalGeneration]) +def test_gemma4_explicit_user_setting_wins(model_cls: type) -> None: + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", kv_cache_config=LlmapiKvCacheConfig(use_kv_cache_manager_v2=False) + ) + original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 + defaults = model_cls.get_model_defaults(llm_args) + apply_model_defaults_to_llm_args(llm_args, defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, defaults, original_setting=original_setting) + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False if __name__ == "__main__": From 17988fa7878d85be31bd6078d8fe15b19b367ae1 Mon Sep 17 00:00:00 2001 From: Eric Tsai Date: Mon, 10 Aug 2026 02:01:06 -0700 Subject: [PATCH 4/8] [None][doc] Add Gemma row to the KV cache manager selection table Signed-off-by: Eric Tsai --- docs/source/features/kvcache.md | 1 + 1 file changed, 1 insertion(+) 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 From 8c221e11872bf5644b96165aee233a0db89251ef Mon Sep 17 00:00:00 2001 From: Eric Tsai Date: Tue, 11 Aug 2026 02:43:42 -0700 Subject: [PATCH 5/8] [None][fix] Migrate Gemma V2 preference to get_preferred_kv_cache_manager_version Signed-off-by: Eric Tsai --- tensorrt_llm/_torch/models/modeling_gemma3.py | 21 ++++------ .../_torch/models/modeling_gemma3vl.py | 23 ++++------- tensorrt_llm/_torch/models/modeling_gemma4.py | 12 ++++-- .../_torch/models/modeling_gemma4mm.py | 8 ++-- .../_torch/modeling/test_modeling_gemma3.py | 41 ++++++++----------- .../_torch/modeling/test_modeling_gemma4.py | 27 ++++++------ tests/unittest/llmapi/test_llm_args.py | 5 +++ 7 files changed, 62 insertions(+), 75 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py index ff2e8fe7db7b..6881ff639255 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3.py @@ -1,5 +1,5 @@ import math -from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple +from typing import Any, Dict, Literal, Optional, Tuple import torch from torch import nn @@ -26,9 +26,6 @@ from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, register_auto_model) -if TYPE_CHECKING: - from tensorrt_llm.llmapi.llm_args import TorchLlmArgs - class Gemma3TextScaledWordEmbedding(Embedding): @@ -296,17 +293,15 @@ def __init__( vocab_size=model_config.pretrained_config.vocab_size) @classmethod - def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: - """Gemma3-specific defaults. + def get_preferred_kv_cache_manager_version(cls, + pretrained_config: Any = None + ) -> Literal["V2"]: + """Prefer KV cache manager V2 for Gemma3's VSWA layout. - Enables the V2 KV-cache manager, which sizes the sliding-window and - full-attention pools independently for Gemma3's VSWA layout. + V2 sizes the sliding-window and full-attention pools independently + instead of statically dividing memory between them. """ - return { - "kv_cache_config": { - "use_kv_cache_manager_v2": True, - }, - } + return "V2" @classmethod def get_preferred_transceiver_runtime( diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py index 2b46e2bd090d..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 TYPE_CHECKING, Any, List, Literal, Optional, Tuple +from typing import Any, List, Literal, Optional, Tuple import torch from transformers import (AutoProcessor, AutoTokenizer, Gemma3Config, @@ -27,9 +27,6 @@ from .modeling_siglip import SiglipVisionModel from .modeling_utils import ModelConfig, filter_weights, register_auto_model -if TYPE_CHECKING: - from tensorrt_llm.llmapi.llm_args import TorchLlmArgs - class Gemma3InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): @@ -181,18 +178,12 @@ def forward(self, vision_outputs: torch.Tensor): class Gemma3VLM(PreTrainedModel): @classmethod - def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: - """Gemma3 VLM defaults — same V2 KV-cache manager selection as - Gemma3ForCausalLM (the wrapped text model). - - The FLASHINFER attention backend is not set here: the text sub-model - config pins it internally in get_sub_model_config. - """ - return { - "kv_cache_config": { - "use_kv_cache_manager_v2": True, - }, - } + 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( diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index c6050de95071..5a3540f4a4ea 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -1258,11 +1258,17 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """ return { "attn_backend": "FLASHINFER", - "kv_cache_config": { - "use_kv_cache_manager_v2": True, - }, } + @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, diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index fd3184e6d6b1..5b7940c0af9b 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -564,11 +564,13 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """Gemma4-specific defaults — see Gemma4ForCausalLM.get_model_defaults.""" return { "attn_backend": "FLASHINFER", - "kv_cache_config": { - "use_kv_cache_manager_v2": True, - }, } + @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, diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma3.py b/tests/unittest/_torch/modeling/test_modeling_gemma3.py index 29c25ca256b8..883b216c7cd6 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma3.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma3.py @@ -20,15 +20,14 @@ Gemma3HfWeightMapper from tensorrt_llm._torch.models.modeling_gemma3 import Gemma3ForCausalLM from tensorrt_llm._torch.models.modeling_gemma3vl import Gemma3VLM -from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING +from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings.executor import KvCacheConfig from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmapiKvCacheConfig from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, - _resolve_transceiver_runtime_auto, - apply_model_defaults_to_llm_args) + _resolve_transceiver_runtime_auto) from tensorrt_llm.mapping import Mapping GEMMA3_1B_CONFIG = { @@ -700,41 +699,33 @@ def test_gemma3_local_context_mask_multi_image(self) -> None: [("Gemma3ForCausalLM", Gemma3ForCausalLM), ("Gemma3ForConditionalGeneration", Gemma3VLM)], ) -def test_gemma3_defaults_resolve_to_v2(architecture: str, - model_cls: type) -> None: +def test_gemma3_preference_resolves_to_v2(architecture: str, + model_cls: type) -> None: """The checkpoint's architectures[0] must map to the class carrying the - defaults, and those defaults must survive the production resolution + V2 preference, and the preference must survive the production resolution ordering (transceiver first, then KV-cache manager) on the NIXL disaggregated route — the route where a missing transceiver preference - silently downgrades the V2 default back to V1.""" - assert MODEL_CLASS_MAPPING[architecture] is model_cls + silently downgrades V2 back to V1.""" + assert get_registered_model_class(architecture) is model_cls llm_args = TorchLlmArgs( model="/tmp/dummy_model", cache_transceiver_config=CacheTransceiverConfig( backend="NIXL", transceiver_runtime="auto"), ) - original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 - defaults = model_cls.get_model_defaults(llm_args) - apply_model_defaults_to_llm_args(llm_args, defaults) _resolve_transceiver_runtime_auto(llm_args, model_cls) - _resolve_kv_cache_manager_v2_auto(llm_args, - defaults, - original_setting=original_setting) + assert _resolve_kv_cache_manager_v2_auto(llm_args, model_cls) is True assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True +@pytest.mark.parametrize("user_setting", [False, True]) @pytest.mark.parametrize("model_cls", [Gemma3ForCausalLM, Gemma3VLM]) -def test_gemma3_explicit_user_setting_wins(model_cls: type) -> None: - llm_args = TorchLlmArgs( - model="/tmp/dummy_model", - kv_cache_config=LlmapiKvCacheConfig(use_kv_cache_manager_v2=False)) - original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 - defaults = model_cls.get_model_defaults(llm_args) - apply_model_defaults_to_llm_args(llm_args, defaults) - _resolve_kv_cache_manager_v2_auto(llm_args, - defaults, - original_setting=original_setting) - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False +def test_gemma3_explicit_setting_wins(model_cls: type, + user_setting: bool) -> None: + llm_args = TorchLlmArgs(model="/tmp/dummy_model", + kv_cache_config=LlmapiKvCacheConfig( + use_kv_cache_manager_v2=user_setting)) + assert _resolve_kv_cache_manager_v2_auto(llm_args, + model_cls) is user_setting diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 93e83a067284..50aac989a964 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -47,7 +47,7 @@ Gemma4TextScaledWordEmbedding, ) from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4ForConditionalGeneration -from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING +from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class from tensorrt_llm._utils import is_sm_100f from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, TorchLlmArgs from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmapiKvCacheConfig @@ -4080,39 +4080,36 @@ def test_vision_attn_metadata_max_num_requests_is_not_one(self): ("Gemma4ForConditionalGeneration", Gemma4ForConditionalGeneration), ], ) -def test_gemma4_defaults_resolve_to_v2(architecture: str, model_cls: type) -> None: +def test_gemma4_preference_resolves_to_v2(architecture: str, model_cls: type) -> None: """The checkpoint's architectures[0] must map to the class carrying the - defaults, and those defaults must survive the production resolution - ordering (transceiver first, then KV-cache manager) on the NIXL + V2 preference, and the preference must survive the production resolution + ordering (defaults merge, transceiver, then KV-cache manager) on the NIXL disaggregated route — the route where a missing transceiver preference - silently downgrades the V2 default back to V1.""" - assert MODEL_CLASS_MAPPING[architecture] is model_cls + silently downgrades V2 back to V1.""" + assert get_registered_model_class(architecture) is model_cls llm_args = TorchLlmArgs( model="/tmp/dummy_model", cache_transceiver_config=CacheTransceiverConfig(backend="NIXL", transceiver_runtime="auto"), ) - original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 defaults = model_cls.get_model_defaults(llm_args) apply_model_defaults_to_llm_args(llm_args, defaults) _resolve_transceiver_runtime_auto(llm_args, model_cls) - _resolve_kv_cache_manager_v2_auto(llm_args, defaults, original_setting=original_setting) + assert _resolve_kv_cache_manager_v2_auto(llm_args, model_cls) is True assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True assert llm_args.attn_backend == "FLASHINFER" +@pytest.mark.parametrize("user_setting", [False, True]) @pytest.mark.parametrize("model_cls", [Gemma4ForCausalLM, Gemma4ForConditionalGeneration]) -def test_gemma4_explicit_user_setting_wins(model_cls: type) -> None: +def test_gemma4_explicit_setting_wins(model_cls: type, user_setting: bool) -> None: llm_args = TorchLlmArgs( - model="/tmp/dummy_model", kv_cache_config=LlmapiKvCacheConfig(use_kv_cache_manager_v2=False) + model="/tmp/dummy_model", + kv_cache_config=LlmapiKvCacheConfig(use_kv_cache_manager_v2=user_setting), ) - original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 - defaults = model_cls.get_model_defaults(llm_args) - apply_model_defaults_to_llm_args(llm_args, defaults) - _resolve_kv_cache_manager_v2_auto(llm_args, defaults, original_setting=original_setting) - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False + assert _resolve_kv_cache_manager_v2_auto(llm_args, model_cls) is user_setting if __name__ == "__main__": diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2d3d96f0aabb..2ebc658e0b5e 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -718,6 +718,11 @@ def test_registered_models_prefer_v2(self): "Qwen3_5ForCausalLM", "Qwen3_5MoeForConditionalGeneration", "Qwen3_5ForConditionalGeneration", + "Gemma3ForCausalLM", + "Gemma3ForConditionalGeneration", + "Gemma4ForCausalLM", + "Gemma4ForConditionalGeneration", + "Gemma4UnifiedForConditionalGeneration", ) for architecture in architectures: model_cls = get_registered_model_class(architecture) From e11ade7fa57e38456800d39a22bb8aa815768c95 Mon Sep 17 00:00:00 2001 From: Eric Tsai Date: Wed, 12 Aug 2026 01:49:23 -0700 Subject: [PATCH 6/8] [None][fix] Prefer the Python transceiver for DeepSeek-V4 and MiniMax-M3; add NIXL V2-retention test Signed-off-by: Eric Tsai --- .../_torch/models/modeling_deepseekv4.py | 13 +++++ .../_torch/models/modeling_minimaxm3.py | 23 ++++++++- .../_torch/modeling/test_modeling_gemma3.py | 45 ---------------- .../_torch/modeling/test_modeling_gemma4.py | 49 ------------------ tests/unittest/llmapi/test_llm_args.py | 51 +++++++++++++++++++ 5 files changed, 86 insertions(+), 95 deletions(-) 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_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/_torch/modeling/test_modeling_gemma3.py b/tests/unittest/_torch/modeling/test_modeling_gemma3.py index 883b216c7cd6..f252f4cbd2d6 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma3.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma3.py @@ -2,7 +2,6 @@ from copy import deepcopy from dataclasses import dataclass -import pytest import torch from _torch.helpers import make_hf_hybrid_cache_for_tests from parameterized import parameterized @@ -19,15 +18,8 @@ from tensorrt_llm._torch.models.checkpoints.hf.gemma3_weight_mapper import \ Gemma3HfWeightMapper from tensorrt_llm._torch.models.modeling_gemma3 import Gemma3ForCausalLM -from tensorrt_llm._torch.models.modeling_gemma3vl import Gemma3VLM -from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig -from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmapiKvCacheConfig -from tensorrt_llm.llmapi.llm_args import TorchLlmArgs -from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, - _resolve_transceiver_runtime_auto) from tensorrt_llm.mapping import Mapping GEMMA3_1B_CONFIG = { @@ -692,40 +684,3 @@ def test_gemma3_local_context_mask_multi_image(self) -> None: ], device=device).bool() torch.testing.assert_close(attention_mask, expected_attention_mask) - - -@pytest.mark.parametrize( - ("architecture", "model_cls"), - [("Gemma3ForCausalLM", Gemma3ForCausalLM), - ("Gemma3ForConditionalGeneration", Gemma3VLM)], -) -def test_gemma3_preference_resolves_to_v2(architecture: str, - model_cls: type) -> None: - """The checkpoint's architectures[0] must map to the class carrying the - V2 preference, and the preference must survive the production resolution - ordering (transceiver first, then KV-cache manager) on the NIXL - disaggregated route — the route where a missing transceiver preference - silently downgrades V2 back to V1.""" - assert get_registered_model_class(architecture) is model_cls - - 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 - - assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True - - -@pytest.mark.parametrize("user_setting", [False, True]) -@pytest.mark.parametrize("model_cls", [Gemma3ForCausalLM, Gemma3VLM]) -def test_gemma3_explicit_setting_wins(model_cls: type, - user_setting: bool) -> None: - llm_args = TorchLlmArgs(model="/tmp/dummy_model", - kv_cache_config=LlmapiKvCacheConfig( - use_kv_cache_manager_v2=user_setting)) - assert _resolve_kv_cache_manager_v2_auto(llm_args, - model_cls) is user_setting diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 50aac989a964..5eff2c644555 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -26,7 +26,6 @@ from types import SimpleNamespace from typing import TYPE_CHECKING -import pytest import torch from transformers import AutoConfig, Gemma4Config, Gemma4TextConfig @@ -46,16 +45,7 @@ Gemma4TextModel, Gemma4TextScaledWordEmbedding, ) -from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4ForConditionalGeneration -from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class from tensorrt_llm._utils import is_sm_100f -from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, TorchLlmArgs -from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmapiKvCacheConfig -from tensorrt_llm.llmapi.llm_utils import ( - _resolve_kv_cache_manager_v2_auto, - _resolve_transceiver_runtime_auto, - apply_model_defaults_to_llm_args, -) from tensorrt_llm.mapping import Mapping if TYPE_CHECKING: @@ -4073,44 +4063,5 @@ def test_vision_attn_metadata_max_num_requests_is_not_one(self): ) -@pytest.mark.parametrize( - ("architecture", "model_cls"), - [ - ("Gemma4ForCausalLM", Gemma4ForCausalLM), - ("Gemma4ForConditionalGeneration", Gemma4ForConditionalGeneration), - ], -) -def test_gemma4_preference_resolves_to_v2(architecture: str, model_cls: type) -> None: - """The checkpoint's architectures[0] must map to the class carrying the - V2 preference, and the preference must survive the production resolution - ordering (defaults merge, transceiver, then KV-cache manager) on the NIXL - disaggregated route — the route where a missing transceiver preference - silently downgrades V2 back to V1.""" - assert get_registered_model_class(architecture) is model_cls - - llm_args = TorchLlmArgs( - model="/tmp/dummy_model", - cache_transceiver_config=CacheTransceiverConfig(backend="NIXL", transceiver_runtime="auto"), - ) - defaults = model_cls.get_model_defaults(llm_args) - apply_model_defaults_to_llm_args(llm_args, defaults) - _resolve_transceiver_runtime_auto(llm_args, model_cls) - assert _resolve_kv_cache_manager_v2_auto(llm_args, model_cls) is True - - assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" - assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True - assert llm_args.attn_backend == "FLASHINFER" - - -@pytest.mark.parametrize("user_setting", [False, True]) -@pytest.mark.parametrize("model_cls", [Gemma4ForCausalLM, Gemma4ForConditionalGeneration]) -def test_gemma4_explicit_setting_wins(model_cls: type, user_setting: bool) -> None: - llm_args = TorchLlmArgs( - model="/tmp/dummy_model", - kv_cache_config=LlmapiKvCacheConfig(use_kv_cache_manager_v2=user_setting), - ) - assert _resolve_kv_cache_manager_v2_auto(llm_args, model_cls) is user_setting - - if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2ebc658e0b5e..21ccb682256b 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -718,6 +718,8 @@ def test_registered_models_prefer_v2(self): "Qwen3_5ForCausalLM", "Qwen3_5MoeForConditionalGeneration", "Qwen3_5ForConditionalGeneration", + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", "Gemma3ForCausalLM", "Gemma3ForConditionalGeneration", "Gemma4ForCausalLM", @@ -729,6 +731,55 @@ def test_registered_models_prefer_v2(self): 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. Models absent from this + list: MiniMax-M2 silently resolves to V1 on this route (its + disaggregated serving is unvalidated -- the missing preference is + deliberate); GLM 5.2 prefers the C++ transceiver for now. + """ + from tensorrt_llm._torch.models.modeling_utils import \ + get_registered_model_class + + architectures = ( + "DeepseekV3ForCausalLM", + "DeepseekV32ForCausalLM", + "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(): From f9b4a48bc849efb9ca213e6657f1f8025bc3db27 Mon Sep 17 00:00:00 2001 From: Eric Tsai Date: Wed, 12 Aug 2026 02:06:52 -0700 Subject: [PATCH 7/8] [None][fix] Drop duplicate Gemma3 transceiver hook (added upstream in #16787) Signed-off-by: Eric Tsai --- tensorrt_llm/_torch/models/modeling_gemma3.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py index 6881ff639255..628d1939797a 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3.py @@ -303,14 +303,6 @@ def get_preferred_kv_cache_manager_version(cls, """ 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, image_token_mask: torch.BoolTensor): device = image_token_mask.device sequence_length = len(image_token_mask) From a5c6e943bb335331863513a0e86596f083526f8a Mon Sep 17 00:00:00 2001 From: Eric Tsai Date: Wed, 12 Aug 2026 18:52:02 -0700 Subject: [PATCH 8/8] [None][test] Cover GLM 5.2 in the NIXL V2-retention test (#17283 made it prefer the Python transceiver) Signed-off-by: Eric Tsai --- tests/unittest/llmapi/test_llm_args.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 21ccb682256b..71963e36e551 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -735,10 +735,10 @@ 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. Models absent from this - list: MiniMax-M2 silently resolves to V1 on this route (its + 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); GLM 5.2 prefers the C++ transceiver for now. + deliberate). """ from tensorrt_llm._torch.models.modeling_utils import \ get_registered_model_class @@ -746,6 +746,7 @@ def test_registered_models_keep_v2_on_nixl(self): architectures = ( "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", + "GlmMoeDsaForCausalLM", "MistralLarge3ForCausalLM", "GptOssForCausalLM", "KimiK25ForConditionalGeneration",