forked from Karylab-cklius/vllm
[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:
@@ -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)]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user