[TRTLLM-14575][fix] MoE: fp32 accumulation in deferred MoEAllReduce finalize#16778
Conversation
…inalize The deferred fused-finalize path (do_finalize=false) in moefinalize_allreduce_fusion_kernel_oneshot_lamport accumulated the top-k weighted expert sum and the shared-expert add directly in DType (bf16), rounding after each of the top_k terms. The non-deferred in-kernel finalize (do_finalize=true) accumulates the same sum in fp32. This precision asymmetry degrades the routed MoE output, and with attention-DP disabled + MTP speculative decoding it drifts the target hidden states enough to lower the MTP acceptance length. Accumulate the top-k sum and the shared-expert add in a local fp32 array, rounding to DType only once when packing into the accumulator for the 128-bit Lamport all-reduce store. Transport is unchanged; the added work is a per-lane fp32 add, so the decode step time is unchanged. Measured on GLM-5.2 (NVFP4) with MTP speculative decoding (GB300, TP4/EP4, attention-DP off, batch size 1, decode-only): - before (bf16 deferred combine): acceptance length 2.28, ~152.7 tok/s/user - after (fp32 deferred combine): acceptance length 2.99, ~200.6 tok/s/user i.e. +31% decode throughput at an unchanged ~14.9 ms/step, with the acceptance length recovering to match the fp32 in-kernel finalize. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
|
/bot run |
df5994b to
3fa37f3
Compare
WalkthroughThe fused MoE one-shot Lamport kernel accumulates top-k weighted expert outputs and optional shared-expert outputs in fp32, then converts each accumulated result to ChangesMoE accumulation precision
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu (1)
595-595: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse complete, capitalized comments.
These comment fragments violate the C++ comment convention.
Proposed cleanup
- // * acc += scale(data) (fp32 accumulation) + // Accumulate scaled expert data in fp32. ... - // * Add shared expert output (fp32 accumulation) + // Add the shared-expert output using fp32 accumulation.As per coding guidelines, “capitalize and punctuate full-sentence comments.”
Also applies to: 603-603
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu` at line 595, Update the comments near the fp32 accumulation expressions at the referenced locations to use complete, capitalized sentences with appropriate ending punctuation, while preserving their existing meaning.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu`:
- Line 595: Update the comments near the fp32 accumulation expressions at the
referenced locations to use complete, capitalized sentences with appropriate
ending punctuation, while preserving their existing meaning.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8e52e2f9-546f-451e-b839-c58a8686af95
📒 Files selected for processing (1)
cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu
|
PR_Github #61198 [ run ] triggered by Bot. Commit: |
|
PR_Github #61198 [ run ] completed with state
|
|
/bot run --reuse-test |
1 similar comment
|
/bot run --reuse-test |
|
PR_Github #61351 [ run ] triggered by Bot. Commit: |
|
PR_Github #61351 [ run ] completed with state
|
xwang233
changed the title
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #61442 [ run ] triggered by Bot. Commit: |
|
PR_Github #61442 [ run ] completed with state |
3 participants
Description
The deferred fused MoE finalize path (
do_finalize=false) inmoefinalize_allreduce_fusion_kernel_oneshot_lamportaccumulated the top-kweighted expert sum and the shared-expert add directly in
DType(bf16),rounding after each of the
top_kterms. The non-deferred in-kernel finalize(
do_finalize=true) accumulates the same sum in fp32, so the two combine pathswere not numerically equivalent.
The bf16 per-term rounding biases the routed MoE output. With attention-DP
disabled + MTP speculative decoding, the drift in the target hidden states is
large enough to lower the MTP acceptance length.
Fix
Accumulate the top-k sum and the shared-expert add in a local fp32 array and
round to
DTypeonce, when packing into the accumulator for the 128-bitLamport all-reduce store. The Lamport transport is unchanged; the only added
work is a per-lane fp32 add, so the decode step time is unchanged.
Measured impact
GLM-5.2 (NVFP4) with MTP speculative decoding — GB300, TP4/EP4, attention-DP
off, batch size 1, decode-only:
≈ +31% decode throughput at unchanged step time, with the acceptance length
recovering to match the fp32 in-kernel finalize. Only decode is affected; the
prefill/context path already uses the non-deferred fp32 finalize.
Test Coverage
The modified kernel
moefinalize_allreduce_fusion_kernel_oneshot_lamportisexercised by the existing MoE finalize / AllReduce-fusion paths. The
accuracy/perf impact above was measured on an end-to-end NVFP4 MoE + MTP decode
workload; output remains coherent.
PR Checklist
or transport change.
Dev Engineer Review
DTypeonly once before the existing 128-bit Lamport all-reduce store, removing per-term rounding and fixing precision asymmetry.QA Engineer Review
No test changes.