[ROCm] Add support for AITER GroupNorm (#3419)

Signed-off-by: Aleksi Vesanto <avesanto@amd.com>

vllm_omni/platforms/rocm/patch/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+def apply_patches():
+    """Apply all ROCm-specific patches."""
+    from vllm_omni.platforms.rocm.patch import worker  # noqa: F401

vllm_omni/platforms/rocm/patch/worker/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from vllm_omni.platforms.rocm.patch.worker import patch_groupnorm  # noqa: F401

vllm_omni/platforms/rocm/patch/worker/patch_groupnorm.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Patch ``initialize_model`` to replace VAE GroupNorm with AITER GroupNorm on ROCm."""
+
+import torch.nn as nn
+from vllm.logger import init_logger
+
+import vllm_omni.diffusion.registry as _registry_mod
+
+logger = init_logger(__name__)
+
+_original_initialize_model = _registry_mod.initialize_model
+
+
+def _replace_groupnorm_with_aiter(vae: nn.Module) -> bool:
+    from aiter.ops.groupnorm import GroupNorm as AiterGroupNorm
+
+    targets = [
+        (parent, name, child)
+        for parent in vae.modules()
+        for name, child in parent.named_children()
+        if isinstance(child, nn.GroupNorm) and child.affine
+    ]
+
+    for parent, name, child in targets:
+        new_group_norm = AiterGroupNorm(
+            num_groups=child.num_groups,
+            num_channels=child.num_channels,
+            eps=child.eps,
+            affine=True,
+            device=child.weight.device,
+            dtype=child.weight.dtype,
+        )
+        new_group_norm.weight = child.weight
+        new_group_norm.bias = child.bias
+        setattr(parent, name, new_group_norm)
+
+    return len(targets) > 0
+
+
+def _patched_initialize_model(od_config):
+    model = _original_initialize_model(od_config)
+
+    if hasattr(model, "vae"):
+        try:
+            from vllm._aiter_ops import is_aiter_found_and_supported
+
+            if is_aiter_found_and_supported() and _replace_groupnorm_with_aiter(model.vae):
+                logger.info("AITER GroupNorm is enabled for VAE.")
+        except Exception:
+            logger.warning("Failed to apply AITER GroupNorm to VAE.")
+
+    return model
+
+
+_registry_mod.initialize_model = _patched_initialize_model

vllm_omni/platforms/rocm/platform.py

@@ -10,6 +10,7 @@ from vllm.platforms.rocm import RocmPlatform
 
 from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
 from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
+from vllm_omni.platforms.rocm.patch import apply_patches
 
 logger = init_logger(__name__)
 
@@ -51,6 +52,10 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
 
     _omni_enum = OmniPlatformEnum.ROCM
 
+    def __init__(self):
+        super().__init__()
+        apply_patches()
+
     @classmethod
     def get_omni_ar_worker_cls(cls) -> str:
         return "vllm_omni.worker.gpu_ar_worker.GPUARWorker"