[ROCm][Bugfix] Make intermediate_pad TP-aware in rocm_aiter_fused_experts (#44679)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>

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,