[ROCm] Run AITER RMSNorm pad fusion before AR RMS fusion (#42411)

Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
This commit is contained in:
akii96
2026-05-13 18:35:12 +08:00
committed by GitHub
parent 97c4317bf5
commit 74dffae666
+5 -3
View File
@@ -143,6 +143,11 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
if self.pass_config.fuse_gemm_comms:
self.passes += [AsyncTPPass(config)]
if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
# Run the more specific RMSNorm+router-pad fusion before
# AR+RMS, since both consume fused_add_rms_norm.
self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
if self.pass_config.fuse_allreduce_rms:
if rocm_aiter_ops.is_enabled():
self.passes += [RocmAiterAllReduceFusionPass(config)]
@@ -164,9 +169,6 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
if rocm_aiter_ops.is_enabled():
self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled():
self.passes += [MLADualRMSNormFusionPass(config)]