[https://nvbugs/6402048][fix] Mirror MXFP4's pattern in MXFP8CutlassFusedMoEMethod — round…

[trtllm-agent]

Summary

• Root cause: MXFP8CutlassFusedMoEMethod.create_weights hard-asserted intermediate_size_per_partition % 128 == 0, which breaks whenever moe_tp_size > 1 shards a non-128-aligned intermediate axis (e.g. Qwen1.5-MoE-A2.7B intermediate_size=1408 with TTP moe_tp_size=4 gives 352).
• Fix: Mirror MXFP4's pattern in MXFP8CutlassFusedMoEMethod — round intermediate_size_per_partition up to weight_alignment for storage, and override load_expert_w{3_w1,2}_weight plus pad-before-shard in load_quant_scales so per-shard e4m3 weights and UE8M0 SF tensors land aligned; padded entries are zeros so reduce-scatter output matches the unpadded reference.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6402048

Summary by CodeRabbit

• Bug Fixes

  • Improved handling of expert weight and scale loading so tensor shapes are aligned consistently, reducing the chance of shape-related issues during model loading and sharding.
  • Updated scale loading behavior to better match padded weight layouts for more reliable mixed-precision MoE execution.

• Tests

  • Removed a test skip entry, so one previously waived MoE test selection will now run again.

[coderabbitai]

Caution

Review failed

An error occurred during the review process. Please try again later.

📝 Walkthrough

Walkthrough

MXFP8 Cutlass fused-MoE weight and scale-factor loading now pads the intermediate-size axis to a computed alignment before buffer allocation and before sharding checkpoint tensors, affecting create_weights, load_expert_w3_w1_weight, load_expert_w2_weight, and load_quant_scales. A test waiver entry was also removed.

Changes

MXFP8 padding alignment

Layer / File(s)	Summary
Padded weight/scale allocation and loading overrides tensorrt_llm/_torch/modules/fused_moe/quantization.py	create_weights rounds intermediate_size_per_partition up to weight_alignment for both w3_w1 and w2 weight/scale buffers; load_expert_w3_w1_weight/load_expert_w2_weight pad checkpoint tensors to per-shard alignment before delegating to base sharding logic.
Padded scale loading in load_quant_scales tensorrt_llm/_torch/modules/fused_moe/quantization.py	load_quant_scales pads w1_sf/w3_sf and w2_sf tensors to computed alignment before column/row sharding, replacing direct unpadded shard loads.

Test waiver cleanup

Layer / File(s)	Summary
Removed obsolete waiver entry tests/integration/test_lists/waives.txt	Removes the skip waiver for test_configurable_moe_multi_gpu with the CUTLASS -k filter.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#15825: Both PRs edit tests/integration/test_lists/waives.txt for the same test_configurable_moe_multi_gpu waiver entry, one adding/updating it and this one removing it.

Suggested reviewers: taylor-yb-lee, nvchenghaoz, yingguo-trt

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[xxi-nv]

Superseded by #15865, which takes a test-skip route (adding QuantAlgo.MXFP8 to should_skip_cutlass's tp_alignment_quants) consistent with how the sibling non-128-aligned quant algos (W8A16 / NVFP4 / W4A8_AWQ) are already handled, instead of an in-kernel padding change. Closing this one.