[https://nvbugs/6858547][fix] Ignore cp_type with one CP rank and require HELIX with more

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(

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().

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)

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()

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,

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(

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

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

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"])