[TRTLLM-9465][fix] Fix accuracy issues with TP+Helix CP

[brb-nv]

Description

This MR fixes poor accuracy on GSM8K accuracy benchmark when Helix CP is used along with TP on gen server.

Background:

• When using helix parallelism, CP ranks are repurposed to TP ranks. This happens not just for the FFN part, but even for the o_proj part of attention layer.

Root cause:

• When there's a mix of TP & HelixCP (TP=2/CP=2 for example), all ranks become TP ranks for o_proj but there's an issue in the order of these ranks because cp_groups are interleaved.
• In above example, cp_groups are {0, 2}, {1,3} but repurposed tp_group is {0,1,2,3}. Note the difference in order of ranks highlighted in bold.
  So, o_proj has the wrong weights on a subset of ranks.

Fix:

• Pass down an override_tp_rank for weight loading so that the right weights are loaded in o_proj layers. This is the least disruptive approach to my understanding.
• While I also explored passing down a modified rank to ctor of o_proj's mapping (so that correct weight-loading would happen automatically), this seems to be a bad idea given how rank has been used for creating comm groups beforehand. This ends up causing a hang.

Test Coverage

$ pytest tests/integration/defs/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:none-pp1tp2cp2] -s -v
$ pytest tests/integration/defs/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp2cp2] -s -v
$ pytest tests/integration/defs/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo-cudagraph:none-pp1tp2cp2] -s -v
$ pytest tests/integration/defs/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo-cudagraph:with_padding-pp1tp2cp2] -s -v

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[coderabbitai]

📝 Walkthrough

Walkthrough

This change introduces an override mechanism for tensor-parallel rank in helix CP mappings. It adds get_helix_overridden_tp_rank() to the mapping layer, threads an override_tp_rank parameter through Attention and MLA output projections, and updates the Linear module's weight-loading path to use the override when available.

Changes

Cohort / File(s)	Summary
Helix TP rank override mechanism tensorrt_llm/mapping.py, tensorrt_llm/_torch/modules/linear.py	Added get_helix_overridden_tp_rank() method to compute effective TP rank when helix CP is repurposed for TP. Updated Linear.__init__ to accept optional override_tp_rank parameter, and modified load_weights_vanilla_helper to prefer override_tp_rank over module.tp_rank when computing effective rank for weight/bias shard loading.
Attention module integration tensorrt_llm/_torch/modules/attention.py	Updated MLA and Attention classes to compute and pass override_tp_rank_for_o_proj to their output projection Linear layers when mapping_with_cp provides helix CP configuration.
Unit tests tests/unittest/others/test_mapping.py	Added test_helix_overridden_tp_rank test method with multiple test cases covering pp/tp/cp combinations to validate get_helix_overridden_tp_rank() computation.
Integration test parametrization tests/integration/defs/accuracy/test_disaggregated_serving.py	Extended test_auto_dtype_with_helix with parametrized generation settings (gen_pp, gen_tp, gen_cp) and updated disaggregated server configuration to use dynamic tensor, pipeline, and context parallel sizes with new moe expert parallel size field.
Test configuration updates tests/integration/test_lists/qa/llm_function_core.txt, tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml, tests/integration/test_lists/test-db/l0_gb200_multi_nodes.yml	Updated test list entries and database configurations to reflect parametrized helix test variants with new pp/tp/cp suffixes (e.g., -pp1tp2cp2).

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 27.27% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description check	✅ Passed	The PR description clearly explains the background, root cause, and fix for the helix parallelism accuracy issue with specific examples and test coverage.
Title check	✅ Passed	The PR title '[TRTLLM-9465][fix] Fix accuracy issues with TP+Helix CP' accurately reflects the main changes: introducing an override mechanism for tensor-parallel rank when using helix CP mappings across attention and linear modules.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tests/integration/defs/accuracy/test_disaggregated_serving.py (1)

875-919: Helix gen_pp/gen_tp/gen_cp IDs here are out of sync with several test lists.

The new parametrization:

@pytest.mark.parametrize("gen_pp,gen_tp,gen_cp", [(1, 1, 2), (2, 1, 2)],
                         ids=["pp1tp1cp2", "pp2tp1cp2"])

is internally consistent, and the use of gen_ep = gen_tp * gen_cp for moe_expert_parallel_size matches the intended helix TP×CP repurposing.

However, multiple integration lists (qa/llm_function_core.txt, test-db/l0_gb200_multi_nodes.yml) still reference variants like ...test_auto_dtype_with_helix[...-pp1tp2cp2]. Those IDs do not exist under this parametrization, so the listed tests will never match a real parametrized case.

Please either:

• Change the ids here to pp1tp2cp2 / pp2tp2cp2 and update all test lists to the same naming convention, or
• Update the lists to use the existing pp1tp1cp2 / pp2tp1cp2 IDs (which reflect gen_tp=1, gen_cp=2), and keep the parametrization as-is.

Optionally, you might also add a @pytest.mark.skip_less_device on this test to protect local runs on machines with fewer GPUs than required by the (ctx_tp=4, gen_pp/gen_tp/gen_cp) configuration, but that's secondary to fixing the ID mismatch.

♻️ Duplicate comments (2)

tests/integration/test_lists/qa/llm_function_core.txt (1)

538-541: Helix test IDs here don’t match the parametrized IDs in test_auto_dtype_with_helix.

The parametrization in accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix defines IDs "pp1tp1cp2" and "pp2tp1cp2", but these list entries use "pp1tp2cp2". As written, these lines won’t match any generated test IDs, so the intended helix variants will not run from this list.

Please align the suffix here with the IDs in the test (or update the test IDs and all lists consistently).

tests/integration/test_lists/test-db/l0_gb200_multi_nodes.yml (1)

16-16: Multi-node helix test entries reference non-existent parametrized IDs.

Both new entries here use ...test_auto_dtype_with_helix[...-pp1tp2cp2], but the test in accuracy/test_disaggregated_serving.py only defines IDs "pp1tp1cp2" / "pp2tp1cp2" for (gen_pp, gen_tp, gen_cp). As a result, these lines won’t correspond to any concrete parametrized test.

Please update either:

• the IDs in the test to use pp1tp2cp2/pp2tp2cp2 and adjust all lists accordingly, or
• these YAML entries (and the QA list) to use the existing parametrized IDs.

Also applies to: 35-35

🧹 Nitpick comments (4)

tests/unittest/others/test_mapping.py (1)

16-16: Helix TP-rank override test coverage looks correct and thorough.

The new test_helix_overridden_tp_rank exercises multiple (pp_size, tp_size, cp_size) configurations and validates exactly the tp_rank * cp_size + cp_rank behavior implemented in MappingBase.get_helix_overridden_tp_rank, including multi-PP scenarios. This is solid coverage and aligns with the documented helix ordering.

If you later add more helix layouts, you might consider factoring HelixTestCase to module scope for reuse, but that’s not necessary now.

Also applies to: 18-18, 86-144

tensorrt_llm/mapping.py (1)

245-267: get_helix_overridden_tp_rank implementation matches the documented helix mapping.

The helper cleanly encodes the helix-to-repurposed-TP mapping as tp_rank * cp_size + cp_rank, consistent with the examples in the docstring and the new unit tests. It’s also harmless when cp_size == 1 (returns the original TP rank).

Optionally, you could assert or mention in the docstring that it’s intended for CpType.HELIX to guard accidental use with other CP types, but the current behavior is logically sound.

tensorrt_llm/_torch/modules/linear.py (1)

151-173: Override TP rank handling in Linear is locally correct; keep an eye on quantized paths.

Using effective_tp_rank = override_tp_rank or tp_rank in load_weights_vanilla_helper is a clean way to remap weight/bias shards for helix while keeping the rest of the module logic keyed off tp_rank. With override_tp_rank=None the behavior is unchanged, so this is backward compatible.

Two follow-ups to consider:

• Document in Linear’s docstring (or a brief comment) that override_tp_rank is a logical TP index override used only for weight/bias loading, currently for helix o_proj in MLA.
• If helix is ever combined with quantized o_proj variants, you’ll likely need the same effective TP rank propagated into the various load_weight_shard calls for scale tensors so that scales stay aligned with the remapped shards.

Also applies to: 2071-2113

tensorrt_llm/_torch/modules/attention.py (1)

797-798: Consider using logger.info_once instead of logger.warning.

Using logger.warning for expected, intentional behavior (helix parallelism configuration) creates unnecessary log noise. Since this message will appear for every MLA layer initialization when helix is enabled, consider using logger.info_once with a key to log this once per run.

🔎 Suggested change

         if mapping_with_cp is not None:
-            logger.warning(
-                "[MLA::__init__] Overriding mapping with CP detected.")
+            logger.info_once(
+                "[MLA::__init__] Overriding mapping with CP detected.",
+                key="mla_mapping_with_cp_override")
             self.mapping = mapping_with_cp

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 13ffe52 and ddc77db.

📒 Files selected for processing (8)

• tensorrt_llm/_torch/modules/attention.py
• tensorrt_llm/_torch/modules/linear.py
• tensorrt_llm/mapping.py
• tests/integration/defs/accuracy/test_disaggregated_serving.py
• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml
• tests/integration/test_lists/test-db/l0_gb200_multi_nodes.yml
• tests/unittest/others/test_mapping.py

🧰 Additional context used

📓 Path-based instructions (2)

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces. Do not use tabs
Always maintain the namespace when importing in Python, even if only one class or function from a module is used
Python files should use snake_case naming: some_file.py
Python classes should use PascalCase naming: class SomeClass
Python functions and methods should use snake_case naming: def my_awesome_function():
Python local variables should use snake_case naming: my_variable = ...
Python variable names that start with a number should be prefixed with 'k': k_99th_percentile = ...
Python global variables should use upper snake_case with prefix 'G': G_MY_GLOBAL = ...
Python constants should use upper snake_case naming: MY_CONSTANT = ...
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings in Python for classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except to the smallest set of errors possible
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible, using the else block for logic

Files:

• tensorrt_llm/_torch/modules/linear.py
• tests/integration/defs/accuracy/test_disaggregated_serving.py
• tests/unittest/others/test_mapping.py
• tensorrt_llm/mapping.py
• tensorrt_llm/_torch/modules/attention.py

**/*.{cpp,h,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the year of its latest meaningful modification

Files:

• tensorrt_llm/_torch/modules/linear.py
• tests/integration/defs/accuracy/test_disaggregated_serving.py
• tests/unittest/others/test_mapping.py
• tensorrt_llm/mapping.py
• tensorrt_llm/_torch/modules/attention.py

🧠 Learnings (13)

📓 Common learnings

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 8063
File: tensorrt_llm/lora_manager.py:1080-1112
Timestamp: 2025-09-29T15:14:28.503Z
Learning: In tensorrt_llm/lora_manager.py, when calculating part_sizes for attn_qkv fused LoRA modules, the sizes are correctly multiplied by tp_size because model_config.num_heads and model_config.num_kv_heads are already divided by tp_size (per-TP-rank values), so multiplication is needed to get the original full concatenated dimension size. The interleave_fused_lora_weights_for_tp function provides proper validation.

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 8063
File: tensorrt_llm/lora_manager.py:1080-1112
Timestamp: 2025-09-29T15:14:28.503Z
Learning: In tensorrt_llm/lora_manager.py, when calculating part_sizes for attn_qkv fused LoRA modules, the sizes are correctly multiplied by tp_size because model_config.num_heads and model_config.num_kv_heads are already divided by tp_size (per-TP-rank values), so multiplication is needed to get the original full concatenated dimension size. The interleave_fused_lora_weights_for_tp function provides proper validation.

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 8063
File: tensorrt_llm/lora_manager.py:1080-1112
Timestamp: 2025-09-29T15:14:28.503Z
Learning: In tensorrt_llm/lora_manager.py, when calculating part_sizes for attn_qkv fused LoRA modules, the sizes are correctly multiplied by tp_size because model_config.num_heads and model_config.num_kv_heads are already divided by tp_size (per-TP-rank values), so multiplication is needed to get the original full concatenated dimension size. The interleave_fused_lora_weights_for_tp function provides proper validation with asserts for total size and TP divisibility.

📚 Learning: 2025-09-09T09:40:45.658Z

Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.

Applied to files:

• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml
• tests/integration/test_lists/test-db/l0_gb200_multi_nodes.yml

📚 Learning: 2025-08-26T09:49:04.956Z

Learnt from: pengbowang-nv
Repo: NVIDIA/TensorRT-LLM PR: 7192
File: tests/integration/test_lists/test-db/l0_dgx_b200.yml:56-72
Timestamp: 2025-08-26T09:49:04.956Z
Learning: In TensorRT-LLM test configuration files, the test scheduling system handles wildcard matching with special rules that prevent duplicate test execution even when the same tests appear in multiple yaml files with overlapping GPU wildcards (e.g., "*b200*" and "*gb200*").

Applied to files:

• tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml

📚 Learning: 2025-09-17T02:48:52.732Z

Learnt from: tongyuantongyu
Repo: NVIDIA/TensorRT-LLM PR: 7781
File: tests/integration/test_lists/waives.txt:313-313
Timestamp: 2025-09-17T02:48:52.732Z
Learning: In TensorRT-LLM, `tests/integration/test_lists/waives.txt` is specifically for waiving/skipping tests, while other test list files like those in `test-db/` and `qa/` directories are for different test execution contexts (pre-merge, post-merge, QA tests). The same test appearing in both waives.txt and execution list files is intentional - the test is part of test suites but will be skipped due to the waiver.

Applied to files:

• tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml

📚 Learning: 2025-07-28T17:06:08.621Z

Learnt from: moraxu
Repo: NVIDIA/TensorRT-LLM PR: 6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.

Applied to files:

• tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml
• tests/integration/test_lists/test-db/l0_gb200_multi_nodes.yml

📚 Learning: 2025-08-13T11:07:11.772Z

Learnt from: Funatiq
Repo: NVIDIA/TensorRT-LLM PR: 6754
File: tests/integration/test_lists/test-db/l0_a30.yml:41-47
Timestamp: 2025-08-13T11:07:11.772Z
Learning: In TensorRT-LLM test configuration files like tests/integration/test_lists/test-db/l0_a30.yml, TIMEOUT values are specified in minutes, not seconds.

Applied to files:

• tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml

📚 Learning: 2025-09-17T06:01:01.836Z

Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7785
File: tests/integration/defs/perf/utils.py:321-333
Timestamp: 2025-09-17T06:01:01.836Z
Learning: In test infrastructure code for disaggregated serving tests, prefer logging errors and continuing execution rather than raising exceptions on timeout, to avoid disrupting test cleanup and causing cascading failures.

Applied to files:

• tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml

📚 Learning: 2025-09-29T15:14:28.503Z

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 8063
File: tensorrt_llm/lora_manager.py:1080-1112
Timestamp: 2025-09-29T15:14:28.503Z
Learning: In tensorrt_llm/lora_manager.py, when calculating part_sizes for attn_qkv fused LoRA modules, the sizes are correctly multiplied by tp_size because model_config.num_heads and model_config.num_kv_heads are already divided by tp_size (per-TP-rank values), so multiplication is needed to get the original full concatenated dimension size. The interleave_fused_lora_weights_for_tp function provides proper validation with asserts for total size and TP divisibility.

Applied to files:

• tensorrt_llm/_torch/modules/attention.py

📚 Learning: 2025-09-29T15:14:28.503Z

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 8063
File: tensorrt_llm/lora_manager.py:1080-1112
Timestamp: 2025-09-29T15:14:28.503Z
Learning: In tensorrt_llm/lora_manager.py, when calculating part_sizes for attn_qkv fused LoRA modules, the sizes are correctly multiplied by tp_size because model_config.num_heads and model_config.num_kv_heads are already divided by tp_size (per-TP-rank values), so multiplication is needed to get the original full concatenated dimension size. The interleave_fused_lora_weights_for_tp function provides proper validation.

Applied to files:

• tensorrt_llm/_torch/modules/attention.py

📚 Learning: 2025-08-14T15:43:23.107Z

Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: tensorrt_llm/_torch/attention_backend/trtllm.py:259-262
Timestamp: 2025-08-14T15:43:23.107Z
Learning: In TensorRT-LLM's attention backend, tensor parameters in the plan() method are assigned directly without validation (dtype, device, contiguity checks). This maintains consistency across all tensor inputs and follows the pattern of trusting callers to provide correctly formatted tensors.

Applied to files:

• tensorrt_llm/_torch/modules/attention.py

📚 Learning: 2025-08-14T21:04:50.248Z

Learnt from: thorjohnsen
Repo: NVIDIA/TensorRT-LLM PR: 6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.

Applied to files:

• tensorrt_llm/_torch/modules/attention.py

📚 Learning: 2025-08-19T12:45:11.997Z

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.

Applied to files:

• tensorrt_llm/_torch/modules/attention.py

📚 Learning: 2025-08-14T06:36:40.701Z

Learnt from: timlee0212
Repo: NVIDIA/TensorRT-LLM PR: 6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.

Applied to files:

• tensorrt_llm/_torch/modules/attention.py

🧬 Code graph analysis (3)

tensorrt_llm/_torch/modules/linear.py (2)

tensorrt_llm/mapping.py (1)

• tp_rank (573-574)

tensorrt_llm/_torch/distributed/communicator.py (1)

• tp_size (64-65)

tests/unittest/others/test_mapping.py (1)

tensorrt_llm/mapping.py (4)

• CpType (25-33)
• rank (199-200)
• rank (203-210)
• get_helix_overridden_tp_rank (249-267)

tensorrt_llm/_torch/modules/attention.py (1)

tensorrt_llm/mapping.py (1)

• get_helix_overridden_tp_rank (249-267)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Pre-commit Check

🔇 Additional comments (1)

tensorrt_llm/_torch/modules/attention.py (1)

925-939: LGTM!

The override_tp_rank parameter is correctly wired through to the o_proj Linear layer, enabling proper weight loading for helix parallelism scenarios. The inline comment clearly documents the purpose.

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #29986 [ run ] triggered by Bot. Commit: b18df55

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #29987 [ run ] triggered by Bot. Commit: abd390e

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #29989 [ run ] triggered by Bot. Commit: 4e23631

[HuiGao-NV]

LGTM

[tensorrt-cicd]

PR_Github #29989 [ run ] completed with state SUCCESS. Commit: 4e23631
/LLM/main/L0_MergeRequest_PR pipeline #23070 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #30035 [ run ] triggered by Bot. Commit: 5973377

[tensorrt-cicd]

PR_Github #30035 [ run ] completed with state SUCCESS. Commit: 5973377
/LLM/main/L0_MergeRequest_PR pipeline #23112 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #30050 [ run ] triggered by Bot. Commit: 6a36120