From 5b6b536fdc41b0f56bbf7cbb19c00eb544d6b16e Mon Sep 17 00:00:00 2001 From: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Date: Wed, 10 Jun 2026 15:10:50 -0500 Subject: [PATCH] [ROCm][Bugfix] Make intermediate_pad TP-aware in rocm_aiter_fused_experts (#44679) Signed-off-by: Rohan138 Co-authored-by: Andreas Karatzas --- .../layers/fused_moe/experts/rocm_aiter_moe.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index da02bec61d..5c2aa45560 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -340,6 +340,16 @@ def rocm_aiter_fused_experts( moe_config.intermediate_size_per_partition - moe_config.intermediate_size_per_partition_unpadded ) + # Round hidden_pad/intermediate_pad to match AITER's CK/FlyDSL MoE + # dispatch (currently pinned to v0.1.13.post1): + # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1073 + # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1099 + # TODO: Revisit this once we bump AITER to 0.1.15 with padding fixes + # for CK/FlyDSL MoE GEMM e.g. https://github.com/ROCm/aiter/pull/3401 + hidden_pad = hidden_pad // 128 * 128 + intermediate_pad = ( + intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1) + ) return rocm_aiter_ops.fused_moe( hidden_states, @@ -357,8 +367,8 @@ def rocm_aiter_fused_experts( doweight_stage1=apply_router_weight_on_input, num_local_tokens=num_local_tokens, output_dtype=output_dtype, - hidden_pad=hidden_pad // 128 * 128, - intermediate_pad=intermediate_pad // 64 * 64 * 2, + hidden_pad=hidden_pad, + intermediate_pad=intermediate_pad, bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None, bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None, moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,