diff --git a/tensorrt_llm/_torch/models/modeling_gemma3.py b/tensorrt_llm/_torch/models/modeling_gemma3.py index 6777ab606469..0e41346e2f3e 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 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], diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 78851f9de78f..a44a42f05a38 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/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], diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 749ed07f50a0..014df89124b0 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/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,7 +1484,7 @@ 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 } } @@ -1467,7 +1492,7 @@ def test_auto_dtype(self, block_reuse): "disable_overlap_scheduler": False, "cuda_graph_config": None, "cache_transceiver_config": { - "backend": "DEFAULT", + "backend": "NIXL", "max_tokens_in_buffer": 4096 } } diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 60bac5aa1700..a3c443b89617 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/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) diff --git a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py index 97634b686ea5..beec99cff550 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py +++ b/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] diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 2fa8772fd073..5e21caf3ee2c 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/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):