[AMD] Diffusion - Enabel rocm miopen tuning on vae (#22428)

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

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: