[chenfeiz0326]

Description

Fixes https://nvbugs/6919428, a perf regression on GB300 DeepSeek-V4-Pro FP4 TEP8 gen_only. Mean gen-worker device step time rose from 14.68 to 17.39 ms (+18.5%) after #19679 upgraded the base image to PyTorch 26.09, which brings in NCCL 2.31.2.

Root cause. When TP spans two nodes and the allreduce strategy is AUTO, each DeepSeek-V4 decoder layer issued two NCCL all-reduces per decode step:

1. The pre-MoE fused all-reduce (AllReduceFusionOp.RMS_NORM). MNNVLAllReduce has no plain RMS_NORM pattern, so this call fell back to NCCL plus a separate RMSNorm kernel.
2. The MoE output all-reduce. It was built without a dtype, so AllReduce never constructed MNNVLAllReduce.

These decode messages are 57 KB (4 tokens × 7168 × bf16). That is below runNCCLAllReduceSymmetric's input-registration threshold of about 228 KiB at 8 ranks, so the input is never window-registered. NCCL 2.30.7 still ran these calls on its symmetric AllReduce_AGxLLMC_R kernel (about 14.8 µs). NCCL 2.31.2 picks 1-channel RING/LL with symKernelId = Unknown (about 32.9 µs). That adds roughly 2.2 ms per decode step, which accounts for the regression. Preloading NCCL 2.31.2 into the old 26.08 image, with the old wheel, reproduces the whole regression. Rolling NCCL back is not an option, because PyTorch 26.09 requires 2.31 symbols.

Fix. This PR changes only DeepSeek-V4 modeling:

• Drop the pre-MoE all-reduce. DeepSeek-V4 attention always completes its TP reduction in o_b_proj (project_sparse_attn_output ignores all_reduce_params). The pre-MoE all-reduce therefore summed a tensor that was already replicated, scaling it by tp_size just before an RMSNorm that cancels the scale. DeepSeek-V4.1 already disables it for the same reason (modeling_deepseekv41.py: "Attention completes TP reduction before mHC; do not reduce the FFN input again."). disable_attn_allreduce is unchanged: once PRE_MOE_FUSION is False, its existing expression produces the same value as before.
• Pass dtype=config.torch_dtype to the DeepseekV4MoE AllReduce. Under AUTO, a multi-node TP group now takes MNNVLAllReduce, as o_b_proj already does.

With both changes, the DeepSeek-V4 TP decode path issues no NCCL all-reduce, so it no longer depends on which kernel NCCL selects. Removing the redundant collective also makes decode faster than the pre-regression baseline. Numerics change slightly: RMSNorm's eps used to act on a tp_size-scaled input, which is no longer the case.

Not in scope, and worth separate follow-ups:

• MNNVLAllReduce support for plain RMS_NORM (the kernel already has kARRMSNorm).
• Revisiting the small-message registration threshold for NCCL ≥ 2.31, which affects other models whose decode all-reduce goes through NCCL_SYMMETRIC.

Test Coverage

• New CPU-only unit tests in tests/unittest/_torch/modeling/test_modeling_deepseekv4.py. This file is already in l0_b200.yml and l0_dgx_b200.yml.
  • test_deepseek_v4_forward_moe_does_not_reduce_replicated_input, parametrized over TP1, TP4 and TP4+ADP, with eager fusion on and off. It checks that PRE_MOE_FUSION stays off and that forward_MoE issues no all-reduce on the replicated input.
  • test_deepseek_v4_decoder_forward_requests_attention_reduction, over the same cases. It runs DeepseekV4DecoderLayer.forward with a recording attention stub. It checks that attention is asked to reduce exactly when TP needs it, that the MoE output reduction is still requested, and that no other all-reduce is issued.
  • test_deepseek_v4_forward_moe_normalizes_unscaled_input runs a TP4 layer with the real RMSNorm and a large eps. It checks that the MoE input is the RMSNorm of the unscaled replicated input, not of a tp_size-scaled copy.
  • test_deepseek_v4_moe_allreduce_is_mnnvl_eligible checks that the MoE AllReduce receives the model dtype.
  • On the fca0803 CI wheel, the file's 61 non-GPU tests pass with this PR's model file. With the base model file, the four TP4-with-eager-fusion and MoE-dtype cases fail.
• Existing test_modeling_deepseekv41.py decoder-reduction tests (does_not_reduce_replicated_input, pre_fusion_requires_unnormalized_input) still pass. I ran them, plus the new tests, on the GenPostMergeBuilds 2911 wheel with this PR's modeling_deepseekv4.py overlaid: 9 + 8 passed.
• Accuracy: accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_tep_mtp_separate_draft_kv_cache (TEP4 + MTP, enable_attention_dp=False, which is the decoder path this PR changes). Passed on the 2911 wheel with this PR's file: GSM8K 95.30 vs 95.07 for unpatched main (reference 95.11); MTP acceptance length 1.9584 vs 1.9581 (min 1.860).
• Perf, on GB300 oci-jhb, for the case aggr_upload-gen_only_no_context-gb300_deepseek-v4-pro-fp4_8k1k_con8_ctx1_dep4_gen4_tep8_eplb0_mtp3_ccb-NIXL. Each arm uses a CI-built wheel with its own harness and pinned image. The good and bad anchors are the L0_PostMerge 3010 and 3012 artifacts that produced the bug's numbers; main is GenPostMergeBuilds 2911. Patched arms overlay the modified modeling_deepseekv4.py into the installed wheel.

Arm (CI wheel, own harness + pinned image)	Mean device step time per rep (ms)	Median per rep (ms)	Output tok/s
good cc63e66 (26.08, NCCL 2.30.7)	14.86, 14.87, 14.66	14.86, 14.86, 14.64	259.5, 259.3, 262.9
bad c490057 (26.09, NCCL 2.31.2)	17.80, 17.49, 17.26	17.29, 17.23, 17.24	217.8, 221.5, 224.2
bad c490057 + this fix	14.00, 13.75, 15.28	13.96, 13.73, 13.93	276.6, 281.2, 253.6
main fca0803 (GenPostMergeBuilds 2911)	19.09, 18.93	17.25, 17.05	—
main fca0803 + this PR's file	14.03, 14.65	13.88, 13.77	—

On main, the same wheel and image give 19.01 → 14.34 ms mean (−24.6%) and 17.15 → 13.82 ms median (−19.4%). The fixed build also beats the pre-regression good commit by about 6% on the median. Some reps on the 26.09 image include bursts of slow iterations (> 20 ms), which inflate the mean but not the median; the fix leaves them in place.

PR Checklist

• 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

Dev Engineer Review

• DeepseekV4MoE now passes config.torch_dtype to AllReduce, allowing AUTO to select MNNVL where supported.
• DeepseekV4DecoderLayer no longer issues a pre-MoE all-reduce. Attention’s output projection already completes the TP reduction. The change also removes the prior environment-variable and attention-DP controls for pre-MoE fusion.
• RMSNorm now receives the unscaled replicated input. The PR reports a slight numerical change. Validate output accuracy and backend selection in the target runtime. Plain RMS_NORM support in MNNVL and NCCL’s small-message registration threshold remain out of scope.
• Review finding counts are unavailable. The PR reports passing tests and performance improvements, but also reports remaining slow-iteration bursts. These results are author-reported, not independently verified here.

QA Engineer Review

• Modified tests/unittest/_torch/modeling/test_modeling_deepseekv4.py. New tests cover unscaled-input normalization, attention reduction requests across TP and attention-DP configurations, and MoE AllReduce dtype and backend eligibility.
• The module appears in tests/integration/test_lists/test-db/l0_b200.yml and l0_dgx_b200.yml. l0_rtx_pro_6000.yml lists the specific test_deepseek_v4_sanity[fp8-ds-mla] test. No matching manual-QA list entry was found.
• Coverage verdict: needs follow-up. Unit tests cover the changed reduction logic and configuration. They do not independently establish runtime MNNVL selection or validate numerical and performance behavior.

Per-File QA Perspective

• tensorrt_llm/_torch/models/modeling_deepseekv4.py: Verify that eligible multi-node TP groups select MNNVL and that removing the pre-MoE reduction preserves expected outputs. Check accuracy with RMSNorm receiving the unscaled input.
• tests/unittest/_torch/modeling/test_modeling_deepseekv4.py: Tests cover normalization input scaling, reduction requests, and MoE all-reduce configuration. The module is listed in B200 CI lists; one sanity test is listed for RTX Pro 6000, and no matching manual-QA entry was found.