diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index b891c9f01865..3ba2528e7403 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -725,7 +725,8 @@ def prepare_inputs( buffers, and the attention and speculative metadata are prepared. """ set_per_request_prefill_cuda_graph_flag(False) - if "cp_type" in self.mapping.cp_config: + # A cp_type takes effect only with context parallelism on. + if "cp_type" in self.mapping.cp_config and self.mapping.cp_size > 1: cp_type = self.mapping.cp_config["cp_type"] if cp_type not in (CpType.HELIX, CpType.ULYSSES): raise NotImplementedError( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e0502b4cd181..c6882bc160b7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -9027,7 +9027,8 @@ def _update_request_states_tp(self, scheduled_requests: ScheduledRequests): @nvtx_range("_update_request_states") def _update_request_states(self, scheduled_requests: ScheduledRequests): cp_config = self.dist.cp_config - if 'cp_type' in cp_config: + # A cp_type takes effect only with context parallelism on. + if 'cp_type' in cp_config and self.dist.cp_size > 1: cp_type = cp_config['cp_type'] if cp_type == CpType.HELIX: # Take the usual route with _update_request_states_tp(). diff --git a/tensorrt_llm/_torch/pyexecutor/request_utils.py b/tensorrt_llm/_torch/pyexecutor/request_utils.py index aa8f5636bd6c..1883d9c2fb9f 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/request_utils.py @@ -334,7 +334,9 @@ def merge_requests( """Merge RequestQueueItems to LlmRequests based on CP configuration. This is a router function that dispatches to the appropriate merge function - based on the CP (Context Parallelism) configuration. + based on the CP (Context Parallelism) configuration. A cp_type other than + HELIX takes effect only with context parallelism on (cp_size > 1); with a + single CP rank, requests are merged as without CP. Args: new_requests: List of RequestQueueItem objects. @@ -351,7 +353,7 @@ def merge_requests( List of LlmRequest objects. Raises: - NotImplementedError: If cp_type is not supported. + NotImplementedError: If cp_size > 1 and cp_type is not supported. """ if "cp_type" in cp_config: cp_type = cp_config["cp_type"] @@ -365,7 +367,7 @@ def merge_requests( exclude_last_generation_logits=exclude_last_generation_logits, ) return merge_requests_to_llm_requests(new_requests, exclude_last_generation_logits) - else: + elif cp_size > 1: raise NotImplementedError(f"Unsupported cp type {cp_type.name}.") return merge_requests_to_llm_requests(new_requests, exclude_last_generation_logits) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e02c42ae24f5..92f0fd91ff7e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -7973,6 +7973,23 @@ def validate_helix_tokens_per_block(self) -> 'TorchLlmArgs': return self + @model_validator(mode='after') + def validate_cp_type_with_context_parallelism(self) -> 'TorchLlmArgs': + """Reject a cp_type other than HELIX when context parallelism is on. + + The PyTorch executor implements context parallelism only with HELIX: + with context_parallel_size > 1, the request path raises for any other + cp_type. + """ + if (self.context_parallel_size > 1 and self.cp_config is not None + and self.cp_config.cp_type != CpType.HELIX): + raise ValueError( + f"cp_config.cp_type {self.cp_config.cp_type.name} is not " + f"supported with context_parallel_size " + f"{self.context_parallel_size}: the PyTorch backend supports " + "context parallelism only with cp_type HELIX.") + return self + def warn_on_unstable_feature_usage(self) -> 'TorchLlmArgs': """Warn on unstable feature usage.""" set_fields = self.model_dump(exclude_unset=True).keys() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 067c082067a6..1343eaaf5e32 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -65,6 +65,7 @@ from tensorrt_llm.bindings.internal.batch_manager import LlmRequestType from tensorrt_llm.llmapi import DisaggScheduleStyle from tensorrt_llm.llmapi.llm_args import EncodeCudaGraphConfig, MTPDecodingConfig +from tensorrt_llm.mapping import CpType from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError pytestmark = pytest.mark.cpu_only @@ -3455,6 +3456,25 @@ def test_update_request_states_uses_common_dummy_cleanup(): stub._terminate_attention_dp_dummy_requests.assert_called_once_with() +@pytest.mark.parametrize("cp_type", [CpType.ULYSSES, CpType.RING, CpType.HELIX]) +def test_update_request_states_single_cp_rank_takes_tp_path(cp_type): + stub = Mock(dist=Mock(cp_config={"cp_type": cp_type}, cp_size=1)) + scheduled_requests = Mock(context_requests=[]) + + PyExecutor._update_request_states(stub, scheduled_requests) + + stub._update_request_states_tp.assert_called_once_with(scheduled_requests) + + +def test_update_request_states_rejects_unsupported_cp_type_with_context_parallelism(): + stub = Mock(dist=Mock(cp_config={"cp_type": CpType.ULYSSES}, cp_size=2)) + + with pytest.raises(NotImplementedError, match="Unsupported cp type ULYSSES"): + PyExecutor._update_request_states(stub, Mock(context_requests=[])) + + stub._update_request_states_tp.assert_not_called() + + def test_overlap_adp_preserves_legacy_role_without_forward_intent_collective(): stub = _StubADPExecutor( max_num_tokens=4096, diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 42fe06ecb65e..bd2ec024191c 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -2491,6 +2491,48 @@ def test_prepare_inputs_returns_steady_gen_fast_inputs(self) -> None: inputs._can_use_steady_gen_fast_prepare.assert_called_once_with( batch, new_tokens, None, None, True) + def test_prepare_inputs_single_cp_rank_ignores_cp_type(self) -> None: + for cp_type in (CpType.ULYSSES, CpType.RING): + with self.subTest(cp_type=cp_type): + inputs = _input_preparer() + inputs.mapping = Mapping(world_size=1, + cp_size=1, + cp_config={"cp_type": cp_type}) + inputs._config = SimpleNamespace(spec_config=None) + inputs._can_use_steady_gen_fast_prepare = Mock( + return_value=True) + fast_inputs = PreparedInputs({"prepared": True}) + inputs._apply_steady_gen_fast_prepare = Mock( + return_value=fast_inputs) + + result = inputs.prepare_inputs( + ScheduledRequests(), + Mock(), + Mock(), + new_tensors_device=SimpleNamespace( + new_tokens=torch.zeros(1)), + enable_spec_decode=False, + runtime_draft_len=0, + is_dummy=True) + + self.assertIs(result.kwargs, fast_inputs.kwargs) + + def test_prepare_inputs_rejects_unsupported_cp_type_with_context_parallelism( + self) -> None: + inputs = _input_preparer() + inputs.mapping = Mapping(world_size=2, + cp_size=2, + cp_config={"cp_type": CpType.RING}) + + with self.assertRaisesRegex(NotImplementedError, + "Unsupported cp_type RING"): + inputs.prepare_inputs(ScheduledRequests(), + Mock(), + Mock(), + enable_spec_decode=False, + runtime_draft_len=0, + is_dummy=True) + def test_prepare_inputs_reports_tree_draft_length(self) -> None: inputs = _input_preparer() inputs._config = SimpleNamespace(spec_config=SimpleNamespace( diff --git a/tests/unittest/_torch/executor/test_request_utils.py b/tests/unittest/_torch/executor/test_request_utils.py index bfa853b8ca63..052a09a2d16e 100644 --- a/tests/unittest/_torch/executor/test_request_utils.py +++ b/tests/unittest/_torch/executor/test_request_utils.py @@ -44,6 +44,7 @@ MultimodalItemMetadata, MultimodalItemPartition, ) +from tensorrt_llm.llmapi.llm_args import CpConfig from tensorrt_llm.mapping import CpType from tensorrt_llm.sampling_params import LogprobMode @@ -1396,6 +1397,59 @@ def test_merge_requests_agg_helix_keeps_full_prompt(): assert llm_request.total_input_len_cp == len(input_tokens) +@pytest.mark.parametrize("is_disagg", [False, True], ids=["agg", "disagg"]) +@pytest.mark.parametrize( + "cp_config", + [ + CpConfig().model_dump(exclude_none=True), + {"cp_type": CpType.ULYSSES}, + {"cp_type": CpType.RING}, + ], + ids=["default", "ulysses", "ring"], +) +def test_merge_requests_single_cp_rank_merges_as_without_cp(cp_config, is_disagg): + """With one CP rank, a cp_type other than HELIX merges like no cp_config.""" + input_tokens = [1, 2, 3, 4, 5, 6, 7] + + def merged(config): + executor_request = trtllm.Request( + input_token_ids=input_tokens, + max_tokens=5, + streaming=False, + sampling_config=trtllm.SamplingConfig(), + output_config=trtllm.OutputConfig(), + ) + return merge_requests( + [RequestQueueItem(id=1, request=executor_request)], + cp_config=config, + cp_rank=0, + cp_size=1, + exclude_last_generation_logits=False, + is_disagg=is_disagg, + ) + + result = merged(cp_config) + expected = merged({}) + + assert len(result) == len(expected) == 1 + assert result[0].request_id == expected[0].request_id == 1 + assert result[0].get_tokens(0) == expected[0].get_tokens(0) == input_tokens + assert result[0].prompt_len == expected[0].prompt_len == len(input_tokens) + + +@pytest.mark.parametrize("cp_type", [CpType.ULYSSES, CpType.RING]) +def test_merge_requests_rejects_unsupported_cp_type_with_context_parallelism(cp_type): + with pytest.raises(NotImplementedError, match=f"Unsupported cp type {cp_type.name}"): + merge_requests( + [RequestQueueItem(1, Mock())], + cp_config={"cp_type": cp_type}, + cp_rank=0, + cp_size=2, + exclude_last_generation_logits=False, + is_disagg=False, + ) + + def test_get_from_waiting_queue(): """Test getting items from waiting queue.""" # Add items to waiting queue diff --git a/tests/unittest/_torch/executor/test_single_cp_rank_cp_config.py b/tests/unittest/_torch/executor/test_single_cp_rank_cp_config.py new file mode 100644 index 000000000000..027cff34c867 --- /dev/null +++ b/tests/unittest/_torch/executor/test_single_cp_rank_cp_config.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A cp_config with context_parallel_size 1 serves requests like no cp_config.""" + +import pytest +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import CpConfig + +PROMPTS = ["The capital of France is", "Hello, my name is"] + + +def _generate_token_ids(cp_config): + with LLM( + model=str(llm_models_root() / "Qwen3" / "Qwen3-0.6B"), + cp_config=cp_config, + kv_cache_config=KvCacheConfig(max_tokens=4096), + ) as llm: + outputs = llm.generate(PROMPTS, SamplingParams(max_tokens=8, temperature=0.0)) + return [output.outputs[0].token_ids for output in outputs] + + +@pytest.fixture(scope="module") +def reference_token_ids(): + return _generate_token_ids(cp_config=None) + + +@pytest.mark.parametrize( + "cp_config", + [CpConfig(), CpConfig(cp_type="RING")], + ids=["default", "ring"], +) +def test_single_cp_rank_cp_config_matches_no_cp_config(reference_token_ids, cp_config): + assert _generate_token_ids(cp_config) == reference_token_ids diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index ede5add24e59..86c21948e61d 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -41,7 +41,7 @@ # fmt: off from tensorrt_llm.llmapi.llm_args import (BaseLlmArgs, BlockReuseConfig, CacheTransceiverConfig, CalibConfig, - ContextChunkingPolicy, + ContextChunkingPolicy, CpConfig, CudaGraphConfig, DecodeCudaGraphConfig, DecodingBaseConfig, @@ -300,6 +300,38 @@ def test_rejection_sampling_still_gated_on_context_parallel(): speculative_config=spec_cfg) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "cp_config", + [CpConfig(), + CpConfig(cp_type="ULYSSES"), + CpConfig(cp_type="RING")], + ids=["default", "ulysses", "ring"]) +def test_context_parallelism_rejects_cp_type_other_than_helix(cp_config): + with pytest.raises(ValueError, + match="supports context parallelism only with cp_type " + "HELIX"): + TorchLlmArgs(model="/tmp/dummy_model", + context_parallel_size=2, + cp_config=cp_config) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "context_parallel_size,cp_config", [(1, CpConfig()), + (1, CpConfig(cp_type="RING")), + (2, CpConfig(cp_type="HELIX")), + (2, None)], + ids=["one_rank_default", "one_rank_ring", "helix", "no_cp_config"]) +def test_context_parallelism_accepts_supported_cp_config( + context_parallel_size, cp_config): + args = TorchLlmArgs(model="/tmp/dummy_model", + context_parallel_size=context_parallel_size, + cp_config=cp_config) + assert args.context_parallel_size == context_parallel_size + assert args.cp_config == cp_config + + @pytest.mark.cpu_only @pytest.mark.parametrize("backend", ["DEFAULT", "UCX", "NIXL", "MOONCAKE", "MPI"])