Repository navigation
[https://nvbugs/6932241][fix] Scope NCCL window leases in helix attention tests - #20588
Conversation
…tion tests The helix MHA/MLA unit tests call the attention module directly, so the NCCL-window tensors produced around the o_proj all-reduce have no enclosing tensor-lease scope. Their leases are released by tensor destructors, which aborts both ranks under TLLM_NCCL_WINDOW_LIFECYCLE_ASSERT and leaves the MPIPoolExecutor parent waiting until the stage timeout. Run the distributed generation step inside nccl_window_tensor_scope and clone its output for the accuracy check, and remove the waivers for nvbugs 6932241 and 6932265. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
WalkthroughThe distributed MHA and MLA tests now run generation inside ChangesDistributed attention tests
Priority: ➖ Normal Estimated code review effort: 2 (Simple) | ~10 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 2 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Title check The title clearly identifies the fix, the NCCL window lease scope, and the affected helix attention tests. It follows the required NVBugs and type format. Description check The description explains the symptoms, root cause, fix, affected release, removed waivers, test coverage, and CI validation. It includes the required sections and completed review checkbox. 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.
- Fix all pre-merge checks with AI
✨ Finishing Touches
-
🧪 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.
|
@yunruis, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~ Assignment explanation
Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once. |
|
/bot run --disable-fail-fast |
|
PR_Github #78871 [ run ] triggered by Bot. Commit: |
|
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements. |
|
PR_Github #78871 [ run ] completed with state |
4 participants
Dev Engineer Review
The MHA and MLA tests now run generation inside
nccl_window_tensor_scopeand clone the output before leaving the scope. This keeps the output available for accuracy checks while the NCCL window lease is scoped. The change also removes two waivers. No review findings were provided, so finding counts are unavailable.QA Engineer Review
The change modifies two test files and
tests/integration/test_lists/waives.txt. It removes the broad_torch/attentionwaiver and thetest_mha_helix_distributedwaiver. No matching helix test entries were found in thetest-db/orqa/lists. No test results were supplied. Coverage verdict: needs follow-up.Per-File QA Perspective
tests/unittest/_torch/attention/multi_gpu/test_mha_helix.py: Covers distributed MHA attention behavior and accuracy after scoped generation. No matching registration was found in the CI or manual-QA lists.tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py: Covers distributed MLA attention behavior and accuracy after scoped generation. No matching registration was found in the CI or manual-QA lists.tests/integration/test_lists/waives.txt: Removes two skip waivers, allowing the associated tests to run instead of being waived.Description
Fixes https://nvbugs/6932241 and its duplicate https://nvbugs/6932265 (release/1.3 only).
Symptom. On
DGX_B300-4_GPUs-PyTorch-Post-Merge-1,test_mha_helix_distributed[nccl-...batch=1, ctx_len=64...]never finishes. It runs into the 2400 s stage timeout and is reported as "Test terminated unexpectedly". The aggregatedunittest/_torch/attention ...entry (6932265) fails because of the same hang: in every failing build its culprit isIN-FLIGHT .../test_mha_helix.py::test_mha_helix_distributed[nccl-...].Root cause. #17705 (NCCL symmetric-window lease scoping, release/1.3 only) gives an NCCL-window lease to the
o_projGEMM output and to the NCCL-symmetric all-reduce output. These leases are meant to be released by an enclosingnccl_window_tensor_scope; the model engine installs one around each decoder layer. A lease released by its tensor destructor instead logs:and calls
std::abort()(TLLM_NCCL_WINDOW_LIFECYCLE_ASSERTdefaults to 1). The helix MHA/MLA tests call the attention module directly, outside any scope. Both MPI ranks compute the correct result (0 mismatches) and then abort when_run_attention_distributedreturns.MPIPoolExecutordoesn't notice that its spawned ranks died, so the parent waits until the stage timeout.The hang is deterministic, not flaky. It reproduces on every run, for every comms medium (nccl, fifo_v1, fifo_v2), and in both
test_mha_helixandtest_mla_helix. CI only ever showed the first parametrization because it blocks everything after it. On release/1.3 post-merge, builds #15–#28 passed all 48 MHA/MLA helix cases; #29 (the first build containing #17705) through #31 hung, after which the waivers landed.Fix. This follows the pattern #17705 applied to other tests that call modules directly. The distributed generation step runs inside
nccl_window_tensor_scope(input_gen), and its output is cloned so the accuracy check doesn't depend on the lease lifetime. The waivers for both bugs are removed. Test-only change.main isn't affected: #17705 targeted release/1.3 and isn't on main. Main post-merge #3015–#3027 pass all 48 MHA/MLA helix cases on B300. Locally on B300, main L0_PostMerge #3010 passes the 6932241 command 10/10, and the 6932265 aggregate command completes with all helix cases passing.
Test Coverage
On B300 ×4 (
umb-b300-dp-127), using the exact failing CI environment: the release-1.3 L0_PostMerge #30 build tarball and its CI image (pytorch-26.08-...-202610082043-19909). The helix test files and the NCCL-window code are identical between #30 and the release/1.3 tip.test_mha_helix.py::test_mha_helix_distributed[nccl-scenario: ...batch=1, ctx_len=64...]Fatal Python error: Aborted)test_mha_helix.py -k fifo_v2test_mla_helix.py -k nccltest_mha_helix.py test_mla_helix.py(all comms media and scenarios)unittest/_torch/attention --ignore=...test_mha_helixrocketkvmodel tests needLLM_MODELS_ROOT.CI:
/bot run --extra-stage "DGX_B300-4_GPUs-PyTorch-Post-Merge-1"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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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
To see a list of available CI bot commands, please comment
/bot help.🤖 Generated with Claude Code