[fix] Fix mxfp4 triton MoE tp bug (#9473)

Signed-off-by: Hao Lu <14827759+hlu1@users.noreply.github.com>
2025-08-23 01:48:40 -07:00
parent c9dd70fbde
commit ccd3fb946e
3 changed files with 14 additions and 12 deletions
@@ -793,12 +793,11 @@ class GptOssForCausalLM(nn.Module):
            intermediate_size % mxfp4_block == 0
        ), f"{intermediate_size=} must be divisible by {mxfp4_block=}"
        intermediate_size_block = intermediate_size // mxfp4_block
-        if _is_sm100_supported:
-            per_rank_intermediate_size_block = math.ceil(
-                intermediate_size_block / moe_tp_size
-            )
-        else:
-            per_rank_intermediate_size_block = intermediate_size_block // moe_tp_size
+
+        per_rank_intermediate_size_block = math.ceil(
+            intermediate_size_block / moe_tp_size
+        )
+
        per_rank_intermediate_size = per_rank_intermediate_size_block * mxfp4_block

        # Calculate common slicing bounds for current rank