[TRTLLM-11057][feat] Add Helix CP support for DSV3.2 - #11507
Conversation
809cbfa to
cd68c05
Compare
📝 WalkthroughWalkthroughThis PR introduces Context Parallel (CP) support for Helix-based distributed serving, including CP-major cache indexing for MLA KV caches, CP-aware reduction logic in DeepSeekV3 attention modules, Helix-specific sequence length scaling, and corresponding test coverage expansion with new Helix-based evaluation paths. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 1 | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)
1-3:⚠️ Potential issue | 🟠 MajorAdd required NVIDIA copyright header for this modified Python file.
This file is modified in this PR but does not include the mandated NVIDIA Apache‑2.0 header with the latest modification year (2026). Please add/update the header at the top.
As per coding guidelines: "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format. Include NVIDIA copyright header on ALL new files and update year on modified files."
tests/integration/defs/accuracy/test_disaggregated_serving.py (1)
1-3:⚠️ Potential issue | 🟠 MajorAdd required NVIDIA copyright header for this modified Python file.
This file is modified in this PR but does not include the mandated NVIDIA Apache‑2.0 header with the latest modification year (2026). Please add/update the header at the top.
As per coding guidelines: "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format. Include NVIDIA copyright header on ALL new files and update year on modified files."
cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp (1)
1-3:⚠️ Potential issue | 🟠 MajorUpdate NVIDIA copyright year to 2026.
This file is modified in this PR; update the header year to reflect the latest meaningful modification.
As per coding guidelines: "All source files must contain an NVIDIA copyright header with the year of latest meaningful modification. Use the Apache License 2.0 format. Include NVIDIA copyright header on ALL new files and update year on modified files."
cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp (1)
2-2:⚠️ Potential issue | 🟡 MinorCopyright year needs updating to 2026.
The file has been meaningfully modified in 2026, but the copyright header still reads
2023-2025. As per coding guidelines, all source files must contain an NVIDIA copyright header with the year of latest meaningful modification.- * SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
🤖 Fix all issues with AI agents
In `@cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp`:
- Around line 256-270: The local index variables in the sendBufferFun lambda
(cpDomainIdx, ppDomainIdx, cacheIdx) are computed once and never modified;
update their declarations to be const to follow the coding guideline that
non-mutated variables be const—modify the lambda inside sendBufferFun to declare
cpDomainIdx, ppDomainIdx and cacheIdx as const (using the existing expressions
involving processIdx, connectionsPerCPDomain and pPDomainSize) so behavior is
unchanged but the variables are immutable.
- Around line 252-270: The code assumes connections.size() is divisible by
cPDomainSize when computing connectionsPerCPDomain; add explicit validation in
the scope where connectionsPerCPDomain is computed (before the lambda that
defines sendBufferFun) to check divisibility and non-zero cPDomainSize, e.g.,
ensure cPDomainSize > 0 and connections.size() % cPDomainSize == 0, and also
validate that the computed cacheIdx (using cpDomainIdx = processIdx /
connectionsPerCPDomain and ppDomainIdx = (processIdx % connectionsPerCPDomain) %
pPDomainSize) will be < outputSplitCaches.size() for all processIdx values; if
any check fails, return or throw an error (or log via TLLM_CHECK_WITH_INFO) to
prevent out‑of‑range access in sendBufferFun.
🧹 Nitpick comments (1)
cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp (1)
1911-1966: Missing inline parameter comments for the last three test parameters.The new
UnevenLayertest instantiations omit the/*isIndexerKCache*/,/*indexerDimPerHead*/, and/*indexerKCacheQuantBlockSize*/inline comments on their final arguments (e.g., line 1928, 1947, 1966), unlike the nearby CP test instantiations (e.g., lines 1841–1843, 1863–1865) which include these comments. Adding them improves readability and consistency.
2ez4bz
left a comment
2ez4bz
left a comment
There was a problem hiding this comment.
Approving changes to tensorrt_llm/_torch/models, with a small clarification question.
cd68c05 to
2fa7c3c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #36106 [ run ] triggered by Bot. Commit: |
|
PR_Github #36106 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #36179 [ run ] triggered by Bot. Commit: |
|
PR_Github #36179 [ run ] completed with state
|
brb-nv
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #36197 [ run ] triggered by Bot. Commit: |
|
PR_Github #36197 [ run ] completed with state
|
2fa7c3c to
7a42fa1
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #36201 [ run ] triggered by Bot. Commit: |
|
PR_Github #36201 [ run ] completed with state
|
7a42fa1 to
1df6c52
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #36321 [ run ] triggered by Bot. Commit: |
|
PR_Github #36321 [ run ] completed with state
|
|
/bot run --only-multi-gpu-test --disable-fail-fast |
|
PR_Github #37578 [ run ] triggered by Bot. Commit: |
|
PR_Github #37567 [ run ] completed with state
|
|
PR_Github #37578 [ run ] completed with state
|
5b032e2 to
6c371fa
Compare
|
/bot run --only-multi-gpu-test --disable-fail-fast |
6c371fa to
ec9c715
Compare
|
/bot run --only-multi-gpu-test --disable-fail-fast |
|
PR_Github #37624 [ run ] triggered by Bot. Commit: |
|
PR_Github #37625 [ run ] triggered by Bot. Commit: |
|
PR_Github #37625 [ run ] completed with state
|
ec9c715 to
c2bec85
Compare
|
/bot run --only-multi-gpu-test --disable-fail-fast |
|
PR_Github #37911 [ run ] triggered by Bot. Commit: |
307379c to
ba5c970
Compare
|
/bot run --only-multi-gpu-test --disable-fail-fast |
|
PR_Github #37917 [ run ] triggered by Bot. Commit: |
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
This reverts commit c2bec85. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
d7b1fb0 to
1c114fd
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #37934 [ run ] triggered by Bot. Commit: |
|
PR_Github #37934 [ run ] completed with state |
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
8 participants
Description
This MR adds Helix support for DSV3.2. Also cherry-picks below MR from release branch to unblock this MR:
#11252
Test Coverage
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 thestage-listparameter 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.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip 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-pipelineReuse 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.
Summary by CodeRabbit
Release Notes
Bug Fixes
New Features
Tests