diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py b/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py index ae6951d10..326987689 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py +++ b/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") diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 397286073..04f015cdc 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/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( diff --git a/test/registered/kernels/ops/diffusion/test_autoencoder_kl_fastpath.py b/test/registered/kernels/ops/diffusion/test_autoencoder_kl_fastpath.py new file mode 100644 index 000000000..5ee99b0ad --- /dev/null +++ b/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"]))