[trtllm-agent]

Description

With context_parallel_size 1, the PyTorch executor still checked cp_config["cp_type"] in its request path:

• merge_requests (request_utils.py) and PyExecutor._update_request_states accepted only HELIX and raised NotImplementedError (Unsupported cp type ...) for any other cp_type.
• DecoderInputPreparer.prepare_inputs raised NotImplementedError for RING.

CpConfig defaults cp_type to ULYSSES. So with one CP rank, a cp_config with the default cp_type (for example cp_config: {}) or ULYSSES made the first request raise NotImplementedError inside the executor loop, and RING failed during warmup. Mapping.has_cp_helix() and has_cp_ulysses(), the decoder warmup and the attention modules already apply a cp_type only when cp_size > 1.

With context_parallel_size > 1, TorchLlmArgs accepted a cp_type other than HELIX, for which merge_requests and _update_request_states raise NotImplementedError. Where the Attention or MLA constructor (or some other module) checks it, such a configuration failed later, at executor initialization. By static reading, a model whose layers have no such check (for example Glm5NextForCausalLM) could start and reach that raise on the first request.

With this change:

• The three request-path checks above apply only when cp_size > 1. With one CP rank, a cp_type other than HELIX no longer changes how requests are merged, updated and prepared: they are handled as without a cp_config.
• TorchLlmArgs raises ValueError for a cp_config whose cp_type is not HELIX (including CpConfig(), whose default is ULYSSES) when context_parallel_size > 1. That configuration now fails at argument validation, before any model is built, instead of in the executor.

Not changed:

• HELIX handling, for any context_parallel_size.
• context_parallel_size > 1 without a cp_config.
• KV cache sizing: with a ULYSSES cp_config and one CP rank, startup without kv_cache_config.max_tokens still fails (Quota not set), as before.

Test Coverage

• tests/unittest/_torch/executor/test_request_utils.py (CPU): with one CP rank, merge_requests merges the default, ULYSSES and RING cp_config like no cp_config, aggregated and disaggregated. With cp_size 2, ULYSSES and RING still raise.
• tests/unittest/_torch/executor/test_py_executor.py (CPU): with one CP rank, _update_request_states takes the TP path for ULYSSES, RING and HELIX. ULYSSES with cp_size 2 still raises.
• tests/unittest/_torch/executor/test_pytorch_model_engine.py: with one CP rank, prepare_inputs accepts ULYSSES and RING. RING with cp_size 2 still raises.
• tests/unittest/llmapi/test_llm_args.py (CPU): with context_parallel_size 2, TorchLlmArgs rejects CpConfig(), ULYSSES and RING, and accepts HELIX and no cp_config. With context_parallel_size 1, it accepts CpConfig() and RING.
• New tests/unittest/_torch/executor/test_single_cp_rank_cp_config.py (GPU, Qwen3-0.6B): LLM with CpConfig() and with CpConfig(cp_type="RING") at context_parallel_size 1 generates the same greedy tokens as without a cp_config.
• Without the code change, 12 of the new unit cases fail: the one-CP-rank cases (except ULYSSES in prepare_inputs and HELIX in _update_request_states) and the three TorchLlmArgs rejections. The cp_size 2 guards and the accepted configurations pass. Both cases of the new LLM API test fail. With the change, all pass.
• Nearby tests on H200:
  • CPU-only tests of test_request_utils.py, test_py_executor.py, test_executor_request_queue.py, tests/unittest/executor/test_forward_error_survival.py, tests/unittest/others/test_mapping.py, tests/unittest/llmapi/test_llm_args.py, tests/unittest/api_stability and tests/unittest/llmapi/test_trtllm_serve_config_schemas.py: 1083 passed without the change, 1102 with it (the 19 new CPU-only cases).
  • test_pytorch_model_engine.py: 116 passed without the change, 118 with it. The same 2 DecoderForwardTestCase::test_mrope_* tests fail in both runs, with TypeError: create_model_engine_and_kvcache() missing 1 required positional argument: 'test_case'.
  • scripts/generate_llm_args_golden_manifest.py --check and scripts/generate_trtllm_serve_schemas.py --check pass with the change.
• trtllm-serve on H200 with Qwen2.5-0.5B-Instruct and kv_cache_config.max_tokens: 8192. Each run sends six requests: chat, greedy chat, completion, streaming chat, chat with n=2, and chat with a long prompt.
  • One GPU, cp_type: ULYSSES, cp_config: {}, and ULYSSES with disable_overlap_scheduler: true: before, the first request failed with Unsupported cp type ULYSSES raised in the executor loop. After, all six return 200.
  • One GPU, cp_type: RING: before, startup failed in warmup. After, all six return 200.
  • Two GPUs, pipeline_parallel_size 2 with cp_type: ULYSSES: before, the first request failed with Unsupported cp type ULYSSES raised in the executor loop. After, all six return 200.
  • After the change, the greedy outputs of these runs equal those of the run without cp_config.
  • One GPU, without cp_config and with cp_type: HELIX: all six return 200 before and after, with the same outputs.
  • Two GPUs, context_parallel_size 2 with ULYSSES or RING: before, argument validation passed and executor initialization failed in Attention construction (CP type must be HELIX for Attention). After, trtllm-serve stops at argument validation with the new ValueError, before the executor starts.
  • Two GPUs, context_parallel_size 2 with HELIX: argument validation passes with the change. Executor initialization then stops with the existing Helix context parallelism is supported only on disaggregated generation servers error.

The new tests are collected by existing CI entries: unittest/_torch/executor in l0_cpu.yml (CPU-only tests) and l0_h100.yml (GPU tests), and unittest/llmapi/test_llm_args.py in l0_cpu.yml.

PR Checklist

• The PR title follows [JIRA ticket/NVBugs ID/GitHub issue/None][type] summary.
• New unit tests cover the change.
• The PR description states the behavior change: context_parallel_size > 1 with a cp_type other than HELIX is refused at argument validation.

Dev Engineer Review

• With context_parallel_size == 1, request merging, request-state updates, and decoder input preparation no longer reject or act on non-HELIX cp_type values. This makes these paths behave as if cp_config were absent.
• With context_parallel_size > 1, TorchLlmArgs now rejects non-HELIX cp_config values during argument validation. The request paths retain their existing unsupported-type errors as defense in depth.
• The change leaves HELIX handling and KV cache sizing unchanged. A one-rank ULYSSES setup without kv_cache_config.max_tokens still fails with Quota not set.
• The added coverage targets one-rank behavior and multi-rank validation. No review-finding severity counts or independent test-run results were supplied.

QA Engineer Review

• Modified unit tests cover request-state updates, decoder input preparation, request merging, argument validation, and generated-token equivalence with no config versus one-rank default or RING config. The token-equivalence test uses Qwen3-0.6B.
• The reported H200 trtllm-serve checks cover one-rank ULYSSES/RING and two-GPU pipeline-parallel ULYSSES. The reported multi-rank outcomes distinguish early validation errors for ULYSSES/RING from the existing disaggregated-generation-server restriction for HELIX.
• test_llm_args.py is listed in CI test lists. test_pytorch_model_engine.py is listed with a selector for a different test, so the new cases are not shown as selected. The other changed test files were not found in the CI or manual-QA lists.
• Coverage verdict: sufficient for the described behavior, with follow-up needed if the new model-engine cases are expected to run through the listed CI selector.

Per-File QA Perspective

• tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py: Verify that single-rank non-HELIX configs use the usual decoder preparation path and that multi-rank unsupported types still fail.
• tensorrt_llm/_torch/pyexecutor/py_executor.py: Verify that single-rank request-state updates use the TP path for configured CP types and that multi-rank unsupported types remain rejected.
• tensorrt_llm/_torch/pyexecutor/request_utils.py: Verify that single-rank request merging matches the no-config result for aggregated and disaggregated requests. Verify that multi-rank non-HELIX types still raise.
• tensorrt_llm/llmapi/llm_args.py: Verify early validation for multi-rank non-HELIX cp_config, and acceptance of HELIX or no config at multiple ranks. test_llm_args.py appears in CI lists.
• tests/unittest/_torch/executor/test_py_executor.py: Adds request-state tests for all CP types at one rank and an unsupported ULYSSES case at two ranks. Not found in the integration test lists.
• tests/unittest/_torch/executor/test_pytorch_model_engine.py: Adds decoder preparation tests for one-rank ULYSSES/RING and two-rank RING. The file is listed in CI with a selector for a different test; the new cases are not included by that selector.
• tests/unittest/_torch/executor/test_request_utils.py: Adds aggregated and disaggregated merge checks at one rank and error checks at two ranks. Not found in the integration test lists.
• tests/unittest/_torch/executor/test_single_cp_rank_cp_config.py: Adds a generated-token comparison between no config and one-rank default/RING config. Not found in the integration test lists.
• tests/unittest/llmapi/test_llm_args.py: Adds argument-validation cases for default, ULYSSES, RING, HELIX, and absent config at one or multiple ranks. Listed in CI.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/TensorRT-LLM
• Branch: main
• File: .coderabbit.yaml
• Commit: 4f613f2
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 858f582d-edb4-4488-9f89-22ffd7b92a23

Commits

Reviewing files that changed from the base of the PR and between 4f613f2 and 0bc05cd.

Files selected for processing: 9

• tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/_torch/pyexecutor/request_utils.py
• tensorrt_llm/llmapi/llm_args.py
• tests/unittest/_torch/executor/test_py_executor.py
• tests/unittest/_torch/executor/test_pytorch_model_engine.py
• tests/unittest/_torch/executor/test_request_utils.py
• tests/unittest/_torch/executor/test_single_cp_rank_cp_config.py
• tests/unittest/llmapi/test_llm_args.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The changes validate non-HELIX context-parallel configurations when the parallel size exceeds one. Executor paths now apply cp_type handling only when context parallelism uses multiple ranks. Tests cover configuration validation, request processing, input preparation, and generated token IDs.

Changes

Context-parallel configuration handling

Layer / File(s)	Summary
Validate multi-rank configurations tensorrt_llm/llmapi/llm_args.py, tests/unittest/llmapi/test_llm_args.py	TorchLlmArgs rejects configured cp_type values other than HELIX when context_parallel_size > 1. Tests cover accepted and rejected configurations.
Gate executor behavior by CP size tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py, tensorrt_llm/_torch/pyexecutor/py_executor.py, tensorrt_llm/_torch/pyexecutor/request_utils.py, tests/unittest/_torch/executor/*	Input preparation, request-state updates, and request merging apply cp_type-specific handling only when context parallelism uses multiple ranks. Tests cover single-rank and multi-rank behavior and compare generated token IDs.

Priority: Normal

Estimated code review effort: 2 (Simple) | ~12 minutes

Change: Bug fix

Suggested reviewers: lori-ren

Merge Risk: Minimal · up to 0bc05

No actionable merge-blocking risk was identified.

Pre-merge checks: 4 · 1

Failed checks: 1 inconclusive

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the main change: ignore cp_type with one CP rank and require HELIX with multiple CP ranks. It uses the required NVBugs and fix format.
Description check		The description includes the required Description, Test Coverage, and PR Checklist sections. It clearly explains the issue, solution, behavior boundaries, test coverage, and known baseline failures.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 5 functions across 8 files. (1 skipped: 1 too large.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}