[TRTLLM-14806][feat] Prefer Python V2 transceiver for LlamaForCausalLM and Gemma3ForCausalLM

tensorrt_llm/_torch/models/modeling_gemma3.py

@@ -1,5 +1,5 @@
 import math
-from typing import Dict, Optional, Tuple
+from typing import Any, Dict, Literal, Optional, Tuple
 
 import torch
 from torch import nn
@@ -276,6 +276,13 @@ def forward(
 class Gemma3ForCausalLM(DecoderModelForCausalLM[Gemma3TextModel,
                                                 Gemma3TextConfig]):
 
+    @classmethod
+    def get_preferred_transceiver_runtime(
+        cls,
+        pretrained_config: Any = None,
+    ) -> Optional[Literal["CPP", "PYTHON"]]:
+        return "PYTHON"
+
     def __init__(
         self,
         model_config: ModelConfig[Gemma3TextConfig],

tensorrt_llm/_torch/models/modeling_llama.py

@@ -1,6 +1,6 @@
 import copy
 import os
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 import torch
 from PIL.Image import Image
@@ -1134,6 +1134,13 @@ def forward(
 @register_auto_model("LlamaForCausalLM")
 class LlamaForCausalLM(SpecDecOneEngineForCausalLM[LlamaModel, LlamaConfig]):
 
+    @classmethod
+    def get_preferred_transceiver_runtime(
+        cls,
+        pretrained_config: Any = None,
+    ) -> Optional[Literal["CPP", "PYTHON"]]:
+        return "PYTHON"
+
     def __init__(
         self,
         model_config: ModelConfig[LlamaConfig],

tests/integration/defs/accuracy/test_disaggregated_serving.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import concurrent
 import contextlib
 import functools
@@ -544,7 +559,8 @@ def run_parallel_test(model_name: str,
                       gen_instances: int,
                       test_sets: List[LlmapiAccuracyTestHarness],
                       ctx_model: str = None,
-                      gen_model: str = None):
+                      gen_model: str = None,
+                      cache_transceiver_backend: str = "DEFAULT"):
     total_ctx_gpus = ctx_tp * ctx_pp * ctx_instances
     total_gen_gpus = gen_tp * gen_pp * gen_instances
     if total_ctx_gpus + total_gen_gpus > get_device_count():
@@ -562,7 +578,7 @@ def run_parallel_test(model_name: str,
         "disable_overlap_scheduler": True,
         "kv_cache_config": kv_cache_config,
         "cache_transceiver_config": {
-            "backend": "DEFAULT",
+            "backend": cache_transceiver_backend,
             "max_tokens_in_buffer": 4096
         }
     }
@@ -572,7 +588,7 @@ def run_parallel_test(model_name: str,
         "disable_overlap_scheduler": True,
         "kv_cache_config": kv_cache_config,
         "cache_transceiver_config": {
-            "backend": "DEFAULT",
+            "backend": cache_transceiver_backend,
             "max_tokens_in_buffer": 4096
         }
     }
@@ -602,6 +618,9 @@ class TestLlama3_1_8BInstruct(LlmapiAccuracyTestHarness):
     MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"
     MODEL_PATH = f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct"
 
+    # Literal NIXL bypasses the harness's legacy UCX fallback. Omitting the
+    # runtime then exercises Llama's automatic preference for Python V2.
+
     @skip_pre_hopper
     @pytest.mark.skip_less_device(2)
     @pytest.mark.parametrize("ctx_disable_overlap_scheduler", [False, True])
@@ -618,7 +637,7 @@ def test_auto_dtype(self, ctx_disable_overlap_scheduler,
             }
         }
         ctx_server_config["cache_transceiver_config"] = {
-            "backend": "DEFAULT",
+            "backend": "NIXL",
             "max_tokens_in_buffer": 4096
         }
         gen_server_config = {
@@ -628,7 +647,7 @@ def test_auto_dtype(self, ctx_disable_overlap_scheduler,
             }
         }
         gen_server_config["cache_transceiver_config"] = {
-            "backend": "DEFAULT",
+            "backend": "NIXL",
             "max_tokens_in_buffer": 4096
         }
         disaggregated_server_config = {
@@ -756,7 +775,7 @@ def test_ngram(self):
             "disable_overlap_scheduler": True,
             "kv_cache_config": kv_cache_config,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             }
         }
@@ -765,7 +784,7 @@ def test_ngram(self):
             "speculative_config": speculative_decoding_config,
             "kv_cache_config": kv_cache_config,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             }
         }
@@ -807,7 +826,7 @@ def test_eagle3(self, overlap_scheduler, eagle3_one_model):
             "max_num_tokens": 13393 * 2,
             "max_batch_size": 1,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             },
             "cuda_graph_config": None,
@@ -822,7 +841,7 @@ def test_eagle3(self, overlap_scheduler, eagle3_one_model):
             "max_num_tokens": 13393 * 2,
             "max_batch_size": 16,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             },
             "cuda_graph_config": None,
@@ -908,14 +927,14 @@ def test_guided_decoding(self, backend: str, mocker):
             "disable_overlap_scheduler": True,
             "guided_decoding_backend": backend,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             }
         }
         gen_server_config = {
             "guided_decoding_backend": backend,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             }
         }
@@ -957,7 +976,7 @@ def test_guided_decoding_with_eagle3(self, backend: str,
             },
             "guided_decoding_backend": backend,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             }
         }
@@ -970,7 +989,7 @@ def test_guided_decoding_with_eagle3(self, backend: str,
             },
             "guided_decoding_backend": backend,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             }
         }
@@ -1003,7 +1022,8 @@ def test_tp_pp_symmetric(self, tp, pp, testset):
                                  gen_tp=tp,
                                  ctx_instances=1,
                                  gen_instances=1,
-                                 test_sets=[get_accuracy_task(testset)])
+                                 test_sets=[get_accuracy_task(testset)],
+                                 cache_transceiver_backend="NIXL")
 
     @parametrize_with_ids("ctx_pp", [2, 4])
     @parametrize_with_ids("gen_tp", [1, 2])
@@ -1020,7 +1040,8 @@ def test_ctx_pp_gen_tp_asymmetric(self, ctx_pp, gen_tp, testset):
                                  gen_tp=gen_tp,
                                  ctx_instances=1,
                                  gen_instances=1,
-                                 test_sets=[get_accuracy_task(testset)])
+                                 test_sets=[get_accuracy_task(testset)],
+                                 cache_transceiver_backend="NIXL")
 
     @pytest.mark.parametrize("testset", ["GSM8K", "MMLU"])
     def test_multi_instance(self, testset):
@@ -1032,7 +1053,8 @@ def test_multi_instance(self, testset):
                                  gen_tp=1,
                                  ctx_instances=2,
                                  gen_instances=2,
-                                 test_sets=[get_accuracy_task(testset)])
+                                 test_sets=[get_accuracy_task(testset)],
+                                 cache_transceiver_backend="NIXL")
 
 
 @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT)
@@ -1450,6 +1472,9 @@ class TestGemma3_1BInstruct(LlmapiAccuracyTestHarness):
     MODEL_NAME = "google/gemma-3-1b-it"
     MODEL_PATH = f"{llm_models_root()}/gemma/gemma-3-1b-it/"
 
+    # Literal NIXL bypasses the harness's legacy UCX fallback. Omitting the
+    # runtime then exercises Gemma's automatic preference for Python V2.
+
     @pytest.mark.skip_less_device(2)
     @pytest.mark.parametrize("block_reuse", [False, True])
     @skip_pre_hopper
@@ -1459,15 +1484,15 @@ def test_auto_dtype(self, block_reuse):
             "disable_overlap_scheduler": True,
             "cuda_graph_config": None,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             }
         }
         gen_server_config = {
             "disable_overlap_scheduler": False,
             "cuda_graph_config": None,
             "cache_transceiver_config": {
-                "backend": "DEFAULT",
+                "backend": "NIXL",
                 "max_tokens_in_buffer": 4096
             }
         }

tests/integration/defs/disaggregated/test_disaggregated.py

@@ -1523,9 +1523,16 @@ def extra_endpoints_test(_server_url: str):
         # Use helper function to validate all timing metrics comprehensively
         validate_timing_metrics(item, "perf_metrics test")
 
+    # This test validates the C++ transceiver's timing-metric semantics. Force
+    # DEFAULT to UCX so Llama's Python preference falls back to C++.
+    env = llm_venv._new_env | {
+        "TRTLLM_USE_NIXL_KVCACHE": "0",
+        "TRTLLM_USE_UCX_KVCACHE": "1",
+        "UCX_TLS": get_ucx_tls(),
+    }
     run_disaggregated_test(disaggregated_example_root,
                            "perf_metrics",
-                           env=llm_venv._new_env,
+                           env=env,
                            extra_endpoints_test=extra_endpoints_test,
                            model_path=llama_model_root,
                            cwd=llm_venv.get_working_directory(),
@@ -1559,10 +1566,16 @@ def test_disaggregated_kv_cache_time_output(disaggregated_test_root, llm_venv,
                         "TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 
     output_path = os.path.join(llm_venv.get_working_directory(), "cache_time")
+    env = llm_venv._new_env.copy()
+    # This test validates the C++ transceiver's CSV format. Selecting UCX for
+    # the DEFAULT backend also resolves the automatic runtime to C++.
+    env["TRTLLM_USE_NIXL_KVCACHE"] = "0"
+    env["TRTLLM_USE_UCX_KVCACHE"] = "1"
+    env["UCX_TLS"] = get_ucx_tls()
+    env["TRTLLM_KVCACHE_TIME_OUTPUT_PATH"] = output_path
     run_disaggregated_test(disaggregated_example_root,
                            "perf_metrics",
-                           env=llm_venv._new_env
-                           | {"TRTLLM_KVCACHE_TIME_OUTPUT_PATH": output_path},
+                           env=env,
                            model_path=llama_model_root,
                            cwd=llm_venv.get_working_directory())
     assert os.path.isdir(output_path)

tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py

@@ -998,8 +998,10 @@ def test_arbitrary_kv_cache_transfer(model, generation_overlap):
         KvCacheConfig(max_tokens=2048 * 8, enable_block_reuse=True)
         for _ in range(2)
     ]
+    # Arbitrary transfer uses the C++ serialized DataTransceiverState protocol.
     cache_transceiver_configs = [
-        CacheTransceiverConfig(backend="DEFAULT") for _ in range(2)
+        CacheTransceiverConfig(backend="DEFAULT", transceiver_runtime="CPP")
+        for _ in range(2)
     ]
     model_names = [model_path(model) for _ in range(2)]
     ranks = [0, 1]
@@ -1156,8 +1158,10 @@ def test_arbitrary_kv_cache_transfer_missing_blocks(model, generation_overlap):
         KvCacheConfig(max_tokens=2048 * 8, enable_block_reuse=True)
         for _ in range(2)
     ]
+    # Arbitrary transfer uses the C++ serialized DataTransceiverState protocol.
     cache_transceiver_configs = [
-        CacheTransceiverConfig(backend="DEFAULT") for _ in range(2)
+        CacheTransceiverConfig(backend="DEFAULT", transceiver_runtime="CPP")
+        for _ in range(2)
     ]
     model_names = [model_path(model) for _ in range(2)]
     ranks = [0, 1]

tests/unittest/llmapi/test_llm_args.py

@@ -24,6 +24,7 @@
 from tensorrt_llm._torch.auto_deploy.llm_args import \
     LlmArgs as AutoDeployLlmArgs
 from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_gemma3 import Gemma3ForCausalLM
 from tensorrt_llm._torch.models.modeling_llama import LlamaForCausalLM
 from tensorrt_llm._torch.virtual_memory import RestoreMode
 from tensorrt_llm.commands.serve import get_llm_args, is_non_default_or_required
@@ -3747,6 +3748,14 @@ def test_model_preference_adopted(self, explicit_auto):
         _resolve_transceiver_runtime_auto(args, _PreferPythonTransceiverModel)
         assert args.cache_transceiver_config.transceiver_runtime == "PYTHON"
 
+    @pytest.mark.parametrize("model_cls", [Gemma3ForCausalLM, LlamaForCausalLM])
+    def test_llama_and_gemma_model_preferences_adopted(self, model_cls):
+        """Llama and Gemma adopt Python when the runtime is left at auto."""
+        args = self._disagg_args()
+        assert args.cache_transceiver_config.transceiver_runtime == "auto"
+        _resolve_transceiver_runtime_auto(args, model_cls)
+        assert args.cache_transceiver_config.transceiver_runtime == "PYTHON"
+
     @pytest.mark.parametrize("explicit_runtime", ["CPP", "PYTHON", None])
     def test_explicit_value_not_overridden_by_model_preference(
             self, explicit_runtime):