diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 3ff53deb4..6d0ad9f41 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -57,6 +57,7 @@ if TYPE_CHECKING: SGLANG_USE_RUNAI_MODEL_STREAMER: bool = True SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: bool = False SGLANG_USE_ROCM_VAE: bool = False + SGLANG_USE_ROCM_CUDNN_BENCHMARK: bool = False def get_default_cache_root() -> str: @@ -283,6 +284,8 @@ environment_variables: dict[str, Callable[[], Any]] = { ), # ROCm: use AITer GroupNorm in VAE for improved performance "SGLANG_USE_ROCM_VAE": _lazy_bool("SGLANG_USE_ROCM_VAE"), + # ROCm: enable cudnn.benchmark (MIOpen auto-tuning) for VAE conv layers + "SGLANG_USE_ROCM_CUDNN_BENCHMARK": _lazy_bool("SGLANG_USE_ROCM_CUDNN_BENCHMARK"), } # Add cache-dit Secondary Transformer Env Vars via programmatic generation to reduce duplication diff --git a/python/sglang/multimodal_gen/runtime/platforms/rocm.py b/python/sglang/multimodal_gen/runtime/platforms/rocm.py index 2ed9a86c3..b937d190f 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/rocm.py +++ b/python/sglang/multimodal_gen/runtime/platforms/rocm.py @@ -193,7 +193,19 @@ class RocmPlatform(Platform): @classmethod def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module: - """Replace nn.GroupNorm with AITer GroupNorm for improved ROCm VAE performance.""" + """Apply ROCm-specific optimizations to VAE. + + - Enable MIOpen benchmark mode so that the best convolution algorithm + is selected for each distinct input shape (benefits Conv3d-heavy VAE + decode). + - Replace nn.GroupNorm with AITer GroupNorm when available. + """ + if envs.SGLANG_USE_ROCM_CUDNN_BENCHMARK and not torch.backends.cudnn.benchmark: + torch.backends.cudnn.benchmark = True + logger.info( + "Enabled cudnn.benchmark (MIOpen auto-tuning) for VAE conv layers" + ) + if not envs.SGLANG_USE_ROCM_VAE: return vae try: