diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 80ee489ad..781721c2d 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 8a86f7989..52c998ac8 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/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 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 30630bc39..39c7c9d41 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 6feb3b71f..b3f3ab6f1 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/platforms/rocm.py b/python/sglang/multimodal_gen/runtime/platforms/rocm.py index 34ddf020e..4f8952dc7 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/rocm.py +++ b/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."""