[diffusion] Generalize the FLUX.2 VAE decoder fast path to AutoencoderKL (Z-Image / FLUX.1) behind quality=high (#33818)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py

@@ -1,5 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-"""CUDA fast paths for the FLUX.2 VAE decoder (AutoencoderKLFlux2).
+"""CUDA fast paths for KL VAE decoders built on the diffusers ``Decoder``.
+
+Covers the FLUX.2 VAE (``AutoencoderKLFlux2``) and the generic
+``AutoencoderKL`` (FLUX.1 / Z-Image / SD3); both share the exact same
+decoder module family (``ResnetBlock2D`` GroupNorm+SiLU chains,
+``Upsample2D``, single-head mid-block ``Attention``).
 
 All rewrites are mathematically exact re-associations of the original
 operators. Wrappers are installed once at VAE load and dispatch on a
@@ -326,7 +331,8 @@ def _decoder_layout_forward(self, *args, **kwargs):
         )
         self._sgl_channels_last = want_cl
         logger.info(
-            "FLUX.2 VAE: decoder switched to %s layout.",
+            "%s: decoder switched to %s layout.",
+            self._sgl_label,
             "channels_last (NHWC)" if want_cl else "contiguous (NCHW)",
         )
     return type(self).forward(self, *args, **kwargs)
@@ -337,23 +343,16 @@ def _decoder_layout_forward(self, *args, **kwargs):
 # ---------------------------------------------------------------------------
 
 
-def maybe_optimize_flux2_vae(vae: nn.Module) -> nn.Module:
-    """Install the quality-gated CUDA FLUX.2 VAE decoder fast paths."""
+def _install_decoder_fast_paths(vae: nn.Module, label: str) -> nn.Module:
+    """Install the quality-gated fast paths on a diffusers ``Decoder`` VAE."""
     from diffusers.models.attention_processor import Attention, AttnProcessor2_0
-    from diffusers.models.autoencoders.vae import Decoder
     from diffusers.models.resnet import ResnetBlock2D
     from diffusers.models.upsampling import Upsample2D
 
-    from sglang.multimodal_gen.runtime.models.vaes.autoencoder_kl_flux2 import (
-        AutoencoderKLFlux2,
-    )
-
-    if not isinstance(vae, AutoencoderKLFlux2) or type(vae.decoder) is not Decoder:
-        return vae
     if getattr(vae, "_spatial_parallel_decode_enabled", False):
         logger.info(
-            "FLUX.2 VAE: spatial-parallel decode enabled; "
-            "skipping CUDA decoder fast paths."
+            "%s: spatial-parallel decode enabled; skipping CUDA decoder fast paths.",
+            label,
         )
         return vae
     if not _HAS_TRITON:
@@ -362,7 +361,7 @@ def maybe_optimize_flux2_vae(vae: nn.Module) -> nn.Module:
         # GroupNorm+SiLU fuse (measured: 97 -> 141 ms at 1024^2 with
         # channels_last alone vs 97 -> 29 ms with both).
         logger.warning(
-            "FLUX.2 VAE: Triton unavailable; skipping CUDA decoder fast paths."
+            "%s: Triton unavailable; skipping CUDA decoder fast paths.", label
         )
         return vae
 
@@ -378,8 +377,9 @@ def maybe_optimize_flux2_vae(vae: nn.Module) -> nn.Module:
         # channels_last tensors; without a layout-safe rewrite for every
         # attention block the layout switch cannot be applied (fail closed).
         logger.warning(
-            "FLUX.2 VAE: %d/%d attention blocks lack a layout-safe rewrite; "
+            "%s: %d/%d attention blocks lack a layout-safe rewrite; "
             "skipping CUDA decoder fast paths.",
+            label,
             n_attn_total - len(attn_modules),
             n_attn_total,
         )
@@ -387,6 +387,7 @@ def maybe_optimize_flux2_vae(vae: nn.Module) -> nn.Module:
 
     gate = VaeFastPathGate()
     decoder._sgl_gate = gate
+    decoder._sgl_label = label
     decoder._sgl_channels_last = False
     decoder.forward = MethodType(_decoder_layout_forward, decoder)
     n_up = _install_fused_upsample(decoder, Upsample2D, gate)
@@ -397,11 +398,37 @@ def maybe_optimize_flux2_vae(vae: nn.Module) -> nn.Module:
     n_norm = _install_norm_silu(decoder, ResnetBlock2D, gate)
     setattr(vae, GATE_ATTR, gate)
     logger.info(
-        "FLUX.2 VAE: installed quality-gated decoder fast paths "
+        "%s: installed quality-gated decoder fast paths "
         "(channels_last dispatch, %d fused upsamplers, %d fast attention "
         "blocks, %d GroupNorm+SiLU fusions).",
+        label,
         n_up,
         len(attn_modules),
         n_norm,
     )
     return vae
+
+
+def maybe_optimize_flux2_vae(vae: nn.Module) -> nn.Module:
+    """Install the quality-gated CUDA FLUX.2 VAE decoder fast paths."""
+    from diffusers.models.autoencoders.vae import Decoder
+
+    from sglang.multimodal_gen.runtime.models.vaes.autoencoder_kl_flux2 import (
+        AutoencoderKLFlux2,
+    )
+
+    if not isinstance(vae, AutoencoderKLFlux2) or type(vae.decoder) is not Decoder:
+        return vae
+    return _install_decoder_fast_paths(vae, "FLUX.2 VAE")
+
+
+def maybe_optimize_autoencoder_kl(vae: nn.Module) -> nn.Module:
+    """Install the quality-gated CUDA fast paths on the generic
+    ``AutoencoderKL`` decoder (FLUX.1 / Z-Image / SD3)."""
+    from diffusers.models.autoencoders.vae import Decoder
+
+    from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL
+
+    if not isinstance(vae, AutoencoderKL) or type(vae.decoder) is not Decoder:
+        return vae
+    return _install_decoder_fast_paths(vae, "AutoencoderKL VAE")

python/sglang/multimodal_gen/runtime/platforms/cuda.py

@@ -550,7 +550,8 @@ class CudaPlatformBase(Platform):
 
     @classmethod
     def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
-        """Install the quality-gated FLUX.2 / Wan VAE decoder fast paths.
+        """Install the quality-gated FLUX.2 / AutoencoderKL / Wan VAE decoder
+        fast paths.
 
         Requests with quality == "high" run the fast paths; the "lossless"
         default runs the original module path bit-for-bit. See
@@ -558,6 +559,7 @@ class CudaPlatformBase(Platform):
         """
         try:
             from sglang.multimodal_gen.runtime.models.vaes.flux2_vae_cuda_opt import (
+                maybe_optimize_autoencoder_kl,
                 maybe_optimize_flux2_vae,
             )
             from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import (
@@ -565,6 +567,7 @@ class CudaPlatformBase(Platform):
             )
 
             vae = maybe_optimize_flux2_vae(vae)
+            vae = maybe_optimize_autoencoder_kl(vae)
             vae = maybe_optimize_wan_vae(vae)
         except Exception:
             logger.warning(

test/registered/kernels/ops/diffusion/test_autoencoder_kl_fastpath.py

@@ -0,0 +1,56 @@
+"""Install-path checks for the generic AutoencoderKL CUDA fast path."""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
+    StableDiffusion3VAEConfig,
+)
+from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt
+from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def _small_config():
+    config = StableDiffusion3VAEConfig()
+    config.arch_config.latent_channels = 2
+    config.arch_config.block_out_channels = (4, 4)
+    config.arch_config.down_block_types = ("DownEncoderBlock2D",) * 2
+    config.arch_config.up_block_types = ("UpDecoderBlock2D",) * 2
+    config.arch_config.layers_per_block = 1
+    config.arch_config.norm_num_groups = 1
+    config.arch_config.sample_size = 8
+    return config
+
+
+@torch.no_grad()
+def test_autoencoder_kl_fastpath_install():
+    torch.manual_seed(0)
+    vae = AutoencoderKL(_small_config()).to("cuda", torch.bfloat16).eval()
+    ref_names = {n for n, _ in vae.named_parameters()}
+    ref_sd = {k: v.clone() for k, v in vae.state_dict().items()}
+    z = torch.randn(1, 2, 8, 8, device="cuda", dtype=torch.bfloat16)
+    ref = vae.decode(z)
+
+    opt = vae_opt.maybe_optimize_autoencoder_kl(vae)
+    gate = getattr(opt, vae_opt.GATE_ATTR, None)
+    assert gate is not None and not gate.enabled
+    # Wrappers must not change parameter FQNs; strict load must round-trip.
+    assert {n for n, _ in opt.named_parameters()} == ref_names
+    opt.load_state_dict(ref_sd, strict=True)
+    # Gate off: bit-for-bit the original path.
+    assert torch.equal(opt.decode(z), ref)
+    # Gate on: fast path runs and stays close; gate off again restores exact.
+    gate.enabled = True
+    torch.testing.assert_close(opt.decode(z).float(), ref.float(), atol=0.1, rtol=0)
+    gate.enabled = False
+    assert torch.equal(opt.decode(z), ref)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))