[AMD] diffusion refactor: move ROCM VAE optimization to Platform abstraction (#20496)

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -380,6 +380,11 @@ class Platform:
         """Whether to enable DIT layerwise offload by default on the current platform."""
         return True
 
+    @classmethod
+    def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
+        """Apply platform-specific optimizations to VAE after loading."""
+        return vae
+
     def get_attn_backend(self, *args, **kwargs) -> AttentionImpl:
         attention_cls_str = self.get_attn_backend_cls_str(*args, **kwargs)
         return resolve_obj_by_qualname(attention_cls_str)

python/sglang/multimodal_gen/runtime/platforms/rocm.py

@@ -182,6 +182,48 @@ class RocmPlatform(Platform):
     def get_device_communicator_cls(cls) -> str:
         return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # works for ROCm too
 
+    @classmethod
+    def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
+        """Replace nn.GroupNorm with AITer GroupNorm for improved ROCm VAE performance."""
+        if not envs.SGLANG_USE_ROCM_VAE:
+            return vae
+        try:
+            from aiter.ops.groupnorm import GroupNorm as AiterGroupNorm
+
+            count = cls._replace_groupnorm(vae, AiterGroupNorm)
+            if count > 0:
+                logger.info(
+                    "Replaced %d nn.GroupNorm modules with AITer GroupNorm in VAE",
+                    count,
+                )
+        except Exception:
+            logger.warning(
+                "Failed to apply AITer GroupNorm to VAE.",
+                exc_info=True,
+            )
+        return vae
+
+    @staticmethod
+    def _replace_groupnorm(module: torch.nn.Module, aiter_gn_cls: type) -> int:
+        count = 0
+        for name, child in module.named_children():
+            if isinstance(child, torch.nn.GroupNorm) and child.affine:
+                replacement = aiter_gn_cls(
+                    num_groups=child.num_groups,
+                    num_channels=child.num_channels,
+                    eps=child.eps,
+                    affine=True,
+                    device=child.weight.device,
+                    dtype=child.weight.dtype,
+                )
+                replacement.weight = child.weight
+                replacement.bias = child.bias
+                setattr(module, name, replacement)
+                count += 1
+            else:
+                count += RocmPlatform._replace_groupnorm(child, aiter_gn_cls)
+        return count
+
     @classmethod
     def enable_dit_layerwise_offload_for_wan_by_default(cls) -> bool:
         """ROCm performs better without DIT layerwise offload on Wan."""