diff --git a/vllm_omni/platforms/rocm/patch/__init__.py b/vllm_omni/platforms/rocm/patch/__init__.py new file mode 100644 index 000000000..ecf416ee3 --- /dev/null +++ b/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 diff --git a/vllm_omni/platforms/rocm/patch/worker/__init__.py b/vllm_omni/platforms/rocm/patch/worker/__init__.py new file mode 100644 index 000000000..824d33c5b --- /dev/null +++ b/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 diff --git a/vllm_omni/platforms/rocm/patch/worker/patch_groupnorm.py b/vllm_omni/platforms/rocm/patch/worker/patch_groupnorm.py new file mode 100644 index 000000000..7de866e73 --- /dev/null +++ b/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 diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index d7ad2b0f7..c2e3fae67 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/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"