A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 19:00 UTC
[ROCm] Run AITER RMSNorm pad fusion before AR RMS fusion (#42411)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96 committed
May 13, 2026 at 13:35 UTC
74dffae6666486d5b19c19d618045469c7932d15
1 file changed
+5
-3
vllm/compilation/passes/pass_manager.py
+5
-3
@@ -143,6 +143,11 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
143
if self.pass_config.fuse_gemm_comms:
144
self.passes += [AsyncTPPass(config)]
145
146
+ if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
147
+ # Run the more specific RMSNorm+router-pad fusion before
148
+ # AR+RMS, since both consume fused_add_rms_norm.
149
+ self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
150
+
151
if self.pass_config.fuse_allreduce_rms:
152
if rocm_aiter_ops.is_enabled():
153
self.passes += [RocmAiterAllReduceFusionPass(config)]
@@ -164,9 +169,6 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
169
if rocm_aiter_ops.is_enabled():
170
self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
171
167
- if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
168
- self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
169
-
172
if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled():
173
self.passes += [MLADualRMSNormFusionPass(config)]
174