diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index 39c7c9d41..30630bc39 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -22,11 +22,42 @@ 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): """ @@ -60,6 +91,7 @@ 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: @@ -153,6 +185,25 @@ 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 @@ -164,6 +215,8 @@ 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