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

python/sglang/multimodal_gen/envs.py

@@ -56,6 +56,7 @@ if TYPE_CHECKING:
     # model loading
     SGLANG_USE_RUNAI_MODEL_STREAMER: bool = True
     SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: bool = False
+    SGLANG_USE_ROCM_VAE: bool = False
 
 
 def get_default_cache_root() -> str:
@@ -276,6 +277,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "SGLANG_USE_RUNAI_MODEL_STREAMER": _lazy_bool(
         "SGLANG_USE_RUNAI_MODEL_STREAMER", "true"
     ),
+    # ROCm: use AITer GroupNorm in VAE for improved performance
+    "SGLANG_USE_ROCM_VAE": _lazy_bool("SGLANG_USE_ROCM_VAE"),
 }
 
 # Add cache-dit Secondary Transformer Env Vars via programmatic generation to reduce duplication

python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py

@@ -16,6 +16,7 @@ from sglang.multimodal_gen.runtime.loader.utils import (
     skip_init_modules,
 )
 from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
     get_diffusers_component_config,
@@ -116,6 +117,7 @@ class VAELoader(ComponentLoader):
                     logger.info(
                         "VAE: converted %d Conv3d weights to channels_last_3d", n
                     )
+            vae = current_platform.optimize_vae(vae)
             return vae
 
         # Load from ModelRegistry (standard VAE classes)
@@ -153,4 +155,5 @@ class VAELoader(ComponentLoader):
             if n > 0:
                 logger.info("VAE: converted %d Conv3d weights to channels_last_3d", n)
 
+        vae = current_platform.optimize_vae(vae)
         return vae

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py

@@ -22,42 +22,11 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
 )
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
-from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 
 logger = init_logger(__name__)
 
-_is_hip = current_platform.is_hip()
-_use_aiter_vae = get_bool_env_var("SGLANG_USE_ROCM_VAE") and _is_hip
-
-
-def _replace_groupnorm_with_aiter(module: torch.nn.Module) -> int:
-    """Recursively replace nn.GroupNorm with AITer GroupNorm in a module tree.
-
-    Returns the number of replaced modules.
-    """
-    from aiter.ops.groupnorm import GroupNorm as AiterGroupNorm
-
-    count = 0
-    for name, child in module.named_children():
-        if isinstance(child, torch.nn.GroupNorm) and child.affine:
-            aiter_gn = AiterGroupNorm(
-                num_groups=child.num_groups,
-                num_channels=child.num_channels,
-                eps=child.eps,
-                affine=True,
-                device=child.weight.device,
-                dtype=child.weight.dtype,
-            )
-            aiter_gn.weight = child.weight
-            aiter_gn.bias = child.bias
-            setattr(module, name, aiter_gn)
-            count += 1
-        else:
-            count += _replace_groupnorm_with_aiter(child)
-    return count
-
 
 def _ensure_tensor_decode_output(decode_output):
     """
@@ -91,7 +60,6 @@ class DecodingStage(PipelineStage):
         super().__init__()
         self.vae: ParallelTiledVAE = vae
         self.pipeline = weakref.ref(pipeline) if pipeline else None
-        self._aiter_gn_applied = False
 
     @property
     def parallelism_type(self) -> StageParallelismType:
@@ -185,25 +153,6 @@ class DecodingStage(PipelineStage):
         image = (image / 2 + 0.5).clamp(0, 1)
         return image
 
-    def _apply_aiter_groupnorm(self):
-        """Replace nn.GroupNorm with AITer GroupNorm (called once after VAE load)."""
-        if self._aiter_gn_applied:
-            return
-        self._aiter_gn_applied = True
-        try:
-            count = _replace_groupnorm_with_aiter(self.vae)
-            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. "
-                "Model may be in an inconsistent state.",
-                exc_info=True,
-            )
-
     def load_model(self):
         # load vae if not already loaded (used for memory constrained devices)
         pipeline = self.pipeline() if self.pipeline else None
@@ -215,8 +164,6 @@ class DecodingStage(PipelineStage):
             if pipeline:
                 pipeline.add_module("vae", self.vae)
             self.server_args.model_loaded["vae"] = True
-        if _use_aiter_vae:
-            self._apply_aiter_groupnorm()
 
     def offload_model(self):
         # Offload models if needed

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."""