Repository navigation
[https://nvbugs/7015107][fix] Keep the Lamport all-reduce clear size in bytes - #21283
Conversation
…in bytes All one-shot Lamport kernels on a mapping share one workspace and clear the previous call's sub-buffer by a size that call left behind. Store it in bytes, not in elements of the writer's dtype, so a bf16 call followed by an fp32 call no longer over-clears into live data (hang) and the reverse no longer leaves stale data (wrong results). Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
@coderabbitai review |
✅ Action performed
|
|
No actionable comments were generated in the recent review.
|
| Layer / File(s) | Summary |
|---|---|
Byte-based Lamport clear sizes cpp/tensorrt_llm/kernels/communicationKernels/*, cpp/tensorrt_llm/kernels/moe/communication/moeAllReduceFusionKernels.cu, tensorrt_llm/_torch/distributed/allreduce_helper.py |
Lamport kernels now record communication-buffer sizes in bytes and round clear sizes up to float4 stores. Workspace comments specify that the clear size is measured in bytes. |
Multi-GPU shared-workspace validation tests/unittest/_torch/multi_gpu/test_allreduce.py, tests/unittest/_torch/multi_gpu/test_vocab_parallel_selection.py, tests/integration/test_lists/test-db/l0_dgx_b200.yml |
Tests cover mixed-dtype ONESHOT all-reduces, odd-token MiniMax RMS calls, and vocab-parallel exchanges around large bf16 all-reduces. The B200 test list adds the two-GPU all-reduce cases. |
Estimated code review effort: 3 (Moderate) | ~20 minutes
Change: Bug fix
Suggested reviewers: juney-nvidia
Merge Risk: 
Minimal · up to 731ca
-
This change keeps the Lamport all-reduce clear size in bytes, which prevents hangs or wrong results when calls with different dtypes share one workspace. It adds multi-GPU regression tests for that case. No concrete merge-blocking risk was found in the supplied context; CI reports the multi-GPU tests as pending the full pre-merge label.
Pre-merge checks: 
4 · 
1
-
Failed checks: 1 warning Passed checks: 4 passed-
Full details: Title check
-
Explanation
The title names the main change: the Lamport all-reduce clear size is now stored in bytes. It also carries an NVBugs ID and the [fix] type, which match the repository template. The changeset matches this summary: the one-shot fusion kernel, MiniMax reduce-RMS kernels, and MoE reduction/finalize kernels all switch the clear size to bytes, with 16-byte rounding on the read side.
Full details: Description check
-
Explanation
The description covers the problem, the fix, and the test coverage. It explains the mixed-dtype workspace clear-size bug and why the fix is in the kernels rather than the vocab-parallel exchange. It lists the new tests, the target kernels, and the reported validation results. It follows the template's Description and Test Coverage sections. The PR Checklist is reduced to a single checked box and does not show the individual template items. The description says the changes were tested on 4x GB300, but those results are author-reported and not verified in this review.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 15.79% 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 19 functions across 7 files.
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
… bf16 all-reduces test_vocab_parallel_exchange_after_large_allreduce runs the VocabParallelHead pair exchange on one-shot right after and before a large bf16 one-shot all-reduce. Both mixed-dtype tests start one rank late per round and fail on a thread timeout instead of hanging; the Lamport test runs in the DGX B200 2-GPU pre-merge stage. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
…e 16-byte stores The scalar MiniMax reduce-RMS kernel writes 4-byte variances, so at TP2 an odd token count left the tail of its sub-buffer uncleared. Readers now round the byte count up, and writers compute it in int64. test_minimax_allreduce_rms_odd_tokens covers it. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
@coderabbitai review |
✅ Action performed
|
|
/bot run |
|
PR_Github #81584 [ run ] triggered by Bot. Commit: |
|
PR_Github #81584 [ run ] completed with state
|
…dencies at module level The workers already import tensorrt_llm at module level. The Triton sample kernel stays a local import: cloudpickle ships the check by value and cannot pickle it. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
…e-bytes test_allreduce.py keeps both appended test groups: the Lamport mixed-dtype and odd-token tests and main's NCCL_SYMMETRIC unregistered-input test. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
|
/bot run |
|
PR_Github #82075 [ run ] triggered by Bot. Commit: |
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. The mHC all-reduce test now collects mismatches and fails after its last launch: a mismatch on one rank used to stop that rank while the other waited for it in the next all-reduce, so the test could hang instead of failing, as it did on GB200. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary[2-ipc] fails without the fix on every run (3/3, rank 0, MoE finalize) and both cases pass with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
Main NVIDIA#21283 changed the shared Lamport workspace dirty range from elements to bytes. Match that contract when clearing and rotating the Rubin reduce/finalize kernels, and update the existing interleave matrix expectations. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
Main NVIDIA#21283 changed the shared Lamport workspace dirty range from elements to bytes. Match that contract when clearing and rotating the Rubin reduce/finalize kernels, and update the existing interleave matrix expectations. Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
…r size in bytes Since NVIDIA#21283 the one-shot Lamport kernels keep the workspace's clear size in bytes. The fused mHC boundary all-reduce still wrote and read it in BF16 elements, so after a fused launch the next library all-reduce cleared only half of the fused launch's slots. When that buffer came around again, a rank could read a peer's stale slice before the fresh one arrived. A mismatch on one rank stopped that rank while the other waited for it in the next all-reduce, which is how test_mhc_allreduce_boundary_matches_allreduce_then_boundary hung on GB200. The fused kernel now records the bytes it pushed and clears the previous launch's bytes in whole 16-byte stores, as the library kernels do. Test (GB300): a fused launch followed by a library all-reduce leaves 2560 of the fused launch's 5120 slot words dirty without the fix, none with it. test_mhc_allreduce_boundary_matches_allreduce_then_boundary fails without the fix on every run (3/3) and passes with it (5/5). Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com>
Dev Engineer Review
The change stores the shared Lamport clear size in bytes. Readers round the byte count up to 16-byte stores. Writers calculate the size in 64-bit arithmetic. This aligns the one-shot all-reduce, MiniMax, and MoE kernels when calls with different dtypes share a workspace. The two-shot kernel and MNNVL workspace are outside this change because they use separate size conventions. No public API change is reported.
Review finding counts are unavailable. Reported CI results are mixed:
PR_Github #81584succeeded, while the/LLM/main/L0_MergeRequest_PR pipeline #67251was unstable. Multi-GPU CI requires theci: full pre-merge approvedlabel. The reported serve-pattern and performance checks are author-reported, not independently verified.QA Engineer Review
The change adds
test_allreduce_lamport_mixed_dtypeandtest_minimax_allreduce_rms_odd_tokensintest_allreduce.py, andtest_vocab_parallel_exchange_after_large_allreduceintest_vocab_parallel_selection.py. These tests cover shared-workspace reuse across dtype orders, odd-token MiniMax clearing, and vocab-parallel exchanges after large BF16 one-shot all-reduces. The B200 CI list adds the two all-reduce tests and retains the vocab-parallel module selector. Coverage verdict: sufficient for the reported byte-size, rounding, and reuse cases. CI results are mixed as noted above; the listed serve-pattern and performance results are reported by the author.Per-File QA Perspective
cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu— Verify byte-based clear sizes in the scalar and float4 Q/K paths, including odd-token inputs.cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu— Verify that the one-shot Lamport kernel reads and writes the shared clear size in bytes across mixed-dtype calls.cpp/tensorrt_llm/kernels/communicationKernels/allReduceWorkspace.cu— The workspace comment defineslayout_buffer[0]as a byte count. Verify that Lamport workspace users follow this convention.cpp/tensorrt_llm/kernels/moe/communication/moeAllReduceFusionKernels.cu— Verify byte-based size updates and rounded clearing in both MoE reduction paths.tensorrt_llm/_torch/distributed/allreduce_helper.py— The workspace comment now specifies bytes. No executable behavior change is reported.tests/unittest/_torch/multi_gpu/test_allreduce.py— The new tests cover mixed-dtype workspace reuse and odd-token MiniMax clearing. Both are listed in the B200 CI test list.tests/unittest/_torch/multi_gpu/test_vocab_parallel_selection.py— The new test covers vocab-parallel exchange after large all-reduces in both dtype orders. The module is listed in the B200 CI test list.tests/integration/test_lists/test-db/l0_dgx_b200.yml— Adds the two all-reduce test selectors to B200 CI and updates the vocab-parallel selection comment. No manual-QA list change is reported.Description
The one-shot Lamport all-reduce kernels share one workspace per mapping (
get_allreduce_workspace), and each call clears the sub-buffer that the previous call wrote, using a size the previous call left in the workspace.allreduce_fusion_kernel_oneshot_lamportstored that size in elements of its own dtype (params.size * NRanks), and the next call divided it by its own elements per 16-byte access. Calls of one dtype agree with each other, but mixed dtypes on the same mapping do not. A bf16 call followed by an fp32 call clears twice the bytes the bf16 call wrote. For a large bf16 message that spills into the live data buffer, and the ranks wait forever on-0.0. An fp32 call followed by a bf16 call clears half of what was written, and stale fp32 data is later read as already arrived, so the result is silently wrong.Main now hits this: the vocab-parallel LM-head token exchange builds
AllReduce(mapping, strategy=AUTO, dtype=torch.float32)(tensorrt_llm/_torch/speculative/vocab_parallel.py), which shares the workspace with the model's bf16 all-reduces. On GB300 every TP4 MiniMax-M3.1 + DSpark serve with default configs hangs in the autotuner warmup, which profiles bf16 ONESHOT up to the workspace limit and then the fp32 exchange. With the target all-reduces on NCCL the drafter's bf16 all-reduces still alternate with the fp32 exchange, and acceptance length varies from run to run.This PR stores the clear size in bytes and converts it to 16-byte accesses on the read side, in every kernel that uses the slot: the all-reduce fusion one-shot kernel, the MiniMax reduce-RMS kernels, and the MoE reduction/finalize all-reduce kernels. Readers round the byte count up to whole 16-byte stores, because the scalar MiniMax reduce-RMS kernel writes 4-byte variances: at TP2 an odd token count leaves a size that is not a multiple of 16, whose tail was never cleared before and returned stale variances three calls later. Writers compute the size in int64. Arithmetic within one dtype is otherwise unchanged. The two-shot kernel uses the separate non-Lamport flag and does not touch the slot, and the MNNVL kernels have their own workspace, which already counts in bytes.
This is the root fix, and I'd prefer it over #21279, which moves the vocab-parallel exchange to NCCL as a stopgap. With the kernel fixed, the exchange can stay on AUTO, and #21279 can be closed, or reverted if it has already merged.
Test Coverage
The new tests are listed in the DGX B200 pre-merge stages. The two
test_allreduce.pytests run in the 2-GPU block, and the vocab-parallel test runs through the existingtest_vocab_parallel_selection.pyentry in the 4-GPU block. The H100 multi-GPU stage collects all three as well. In every test that mixes dtypes, one rank starts each call late in each round, so the other rank reads that rank's part of the buffer before it arrives and takes any stale value. Apytest-timeoutthread timeout turns a hang into a failure.tests/unittest/_torch/multi_gpu/test_allreduce.py::test_allreduce_lamport_mixed_dtype(2 GPUs, both orders). It alternates a bf16 and an fp32 ONESHOT all-reduce of about 0.9x the workspace with small calls of the other dtype on one mapping, over three rounds of the triple buffer, and checks exact sums of small integers. On GB300 without the fix,bf16_then_fp32hangs andfp32_then_bf16returns 2048 wrong elements on the large bf16 call.tests/unittest/_torch/multi_gpu/test_vocab_parallel_selection.py::test_vocab_parallel_exchange_after_large_allreduce(2 GPUs, both orders). It runs theVocabParallelHeadpair exchange on the one-shot kernel right after a bf16 all-reduce filling 0.9x the workspace, and in the other order a large exchange followed by bf16 calls that reuse its sub-buffer. Untuned AUTO sends the exchange to NCCL, so the test pins the one-shot kernel that the autotuner profiles at warm-up. Without the fix, the first order hangs and the second returns 62250 wrong elements.tests/unittest/_torch/multi_gpu/test_allreduce.py::test_minimax_allreduce_rms_odd_tokens(2 GPUs). It runs twelve one-token MiniMax reduce-RMS calls whose per-rank inputs change every step and checks the normalized output. Without the rounding, every output element of the call that reuses the first call's sub-buffer is wrong (3072/3072).test_allreduce.py(fusion patterns, MoE reduction/finalize, MiniMax RMS/QK),test_ar_residual_norm.py,test_mnnvl_allreduce.pyandtest_vocab_parallel_selection.pypass on 4x GB300.PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.