A high-throughput and memory-efficient inference and serving engine for LLMs
[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>
Rohan Potdar committed
Jun 10, 2026 at 15:10 UTC
5b6b536fdc41b0f56bbf7cbb19c00eb544d6b16e
1 file changed
+12
-2
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+12
-2
@@ -340,6 +340,16 @@ def rocm_aiter_fused_experts(
340
moe_config.intermediate_size_per_partition
341
- moe_config.intermediate_size_per_partition_unpadded
342
)
343
+ # Round hidden_pad/intermediate_pad to match AITER's CK/FlyDSL MoE
344
+ # dispatch (currently pinned to v0.1.13.post1):
345
+ # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1073
346
+ # https://github.com/ROCm/aiter/blob/v0.1.13.post1/aiter/fused_moe.py#L1099
347
+ # TODO: Revisit this once we bump AITER to 0.1.15 with padding fixes
348
+ # for CK/FlyDSL MoE GEMM e.g. https://github.com/ROCm/aiter/pull/3401
349
+ hidden_pad = hidden_pad // 128 * 128
350
+ intermediate_pad = (
351
+ intermediate_pad // 64 * 64 * (2 if moe_config.tp_size == 1 else 1)
352
+ )
353
354
return rocm_aiter_ops.fused_moe(
355
hidden_states,
@@ -357,8 +367,8 @@ def rocm_aiter_fused_experts(
367
doweight_stage1=apply_router_weight_on_input,
368
num_local_tokens=num_local_tokens,
369
output_dtype=output_dtype,
360
- hidden_pad=hidden_pad // 128 * 128,
361
- intermediate_pad=intermediate_pad // 64 * 64 * 2,
370
+ hidden_pad=hidden_pad,
371
+ intermediate_pad=intermediate_pad,
372
bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
373
bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
374
moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,