From 50c1bf0db0ac6b962c6879f0b3138a0f4a50a0c2 Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <158081477+Dayuxiaoshui@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:02:33 +0800 Subject: [PATCH] [Diffusion] Port the Wan VAE decoder fast paths to the Qwen-Image VAE (#38020) Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../existing-fast-paths.md | 14 +- .../models/vaes/autoencoder_kl_qwenimage.py | 213 +++++++++--------- .../runtime/models/vaes/wan_vae_cuda_opt.py | 185 +++++++++++++-- .../multimodal_gen/runtime/platforms/cuda.py | 6 +- .../ops/diffusion/test_model_fast_paths.py | 155 +++++++++++++ 5 files changed, 446 insertions(+), 127 deletions(-) diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md index d997f1f3f..c252b3a36 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md @@ -47,6 +47,7 @@ framework-specific optimization workflow. - `python/sglang/multimodal_gen/runtime/models/vaes/fast_path_gate.py` - `python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py` - `python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py` +- `python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py` - `python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py` - `test/registered/kernels/ops/diffusion/test_qwen_image_modulation.py` - `test/registered/kernels/ops/diffusion/test_group_norm_silu.py` @@ -128,8 +129,10 @@ framework-specific optimization workflow. Wan cublasLt/NVFP4 GELU, Qwen added-QKV, GLM/Qwen/Hunyuan/LTX fused GELU, LTX RMSNorm+modulate, Hunyuan QK RMSNorm, Ideogram gated RMSNorm, LingBot RMSNorm, SANA-Video linear attention, generic KL VAE - decoder rewrites used by FLUX.1/FLUX.2/Z-Image/SD3, and Wan VAE - RMSNorm+SiLU. + decoder rewrites used by FLUX.1/FLUX.2/Z-Image/SD3, and Wan / Qwen-Image + VAE RMSNorm+SiLU (the Qwen-Image VAE is the Wan 2.1 VAE; its gate also + re-expresses the `Resample` upsample input with canonical NHWC strides so + the 2D conv runs channels_last end-to-end). - Do not confuse request `--quality` with `--output-quality`, which controls output-file compression rather than model math. - Validation: `test_quality_gate.py`, `test_fused_ln_modulate.py`, @@ -223,7 +226,12 @@ framework-specific optimization workflow. - Numerical contract: these are bit-exact data-movement / same-order-add replacements and run independently of the request-gated Wan RMSNorm+SiLU path. Unsupported layouts or padding fall back to the aten chain. -- Validation: `test/registered/kernels/ops/diffusion/test_wan_causal_cache.py`. +- The Qwen-Image VAE (`autoencoder_kl_qwenimage.py`) uses the same + `cat_pad_channels_last_3d` + compact-cache helper for every causal conv + slot; single-frame image decodes keep the compact cache at the reference + size (one frame) so peak memory does not grow. +- Validation: `test/registered/kernels/ops/diffusion/test_wan_causal_cache.py` + and the `test_qwen_vae_*` cases in `test_model_fast_paths.py`. 15. Helios paired transposed RoPE - Kernel: `fused_inplace_helios_qk_rope`. diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py index d42ee7759..c106319f2 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py @@ -37,13 +37,90 @@ from sglang.multimodal_gen.runtime.models.vaes.common import ( has_decode_parallel_world, should_run_spatial_shard_parallel_decode, ) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +if current_platform.is_cuda(): + try: + from sglang.kernels.ops.diffusion import cat_pad_channels_last_3d + except ImportError: # pragma: no cover + cat_pad_channels_last_3d = None +else: + cat_pad_channels_last_3d = None + logger = init_logger(__name__) # pylint: disable=invalid-name CACHE_T = 2 +def _fused_conv_cache_supported(conv: nn.Module, x: torch.Tensor) -> bool: + """The fused cat+pad writes channels_last_3d, so it only pays off (and is + only exercised) when the conv weights are channels_last_3d too; the kernel + symbol being present already implies a CUDA platform.""" + return ( + cat_pad_channels_last_3d is not None + and type(conv) is QwenImageCausalConv3d + and x.dim() == 5 + and x.is_cuda + and conv.weight.is_contiguous(memory_format=torch.channels_last_3d) + and not torch.compiler.is_compiling() + ) + + +def _run_cached_causal_conv( + conv: nn.Module, + x: torch.Tensor, + cache_list: list, + idx: int, +) -> torch.Tensor: + """Run one causal conv, consuming and refreshing its feature-cache slot. + + Same contract as the Wan VAE helper (this VAE is the Wan 2.1 VAE): the + fast path builds the conv input (cache frames + hidden state + padding) + directly in channels_last_3d with one kernel and takes the next cache + entry as one compact copy of that input's unpadded tail, instead of the + per-chunk ``clone``/``cat``/``F.pad``/relayout bookkeeping. The compact + cache holds exactly the reference cache values, so fused and fallback + chunks can interleave. Pure data movement plus zero fill, so the conv + input is bit-identical to the aten chain. Falls back to the original op + chain whenever the fused kernel does not support the request. + """ + cache = cache_list[idx] + is_rep = isinstance(cache, str) # "Rep" marker from QwenImageResample + payload = None if is_rep else cache + if _fused_conv_cache_supported(conv, x) and ( + payload is None or (payload.device == x.device and payload.dtype == x.dtype) + ): + # Emit exactly the frames the reference bookkeeping would keep: the + # first chunk of a clip stores only its own tail (``x[:, :, -CACHE_T:]``), + # later chunks and the "Rep" slot always hold ``CACHE_T`` frames. This + # keeps single-frame (image) decodes from pinning a second full-size + # frame per conv site. + keep_t = ( + CACHE_T if (payload is not None or is_rep) else min(CACHE_T, x.shape[2]) + ) + pair = cat_pad_channels_last_3d(x, payload, conv._padding, keep_cache_t=keep_t) + if pair is not None: + inp, cache_list[idx] = pair + return nn.Conv3d.forward(conv, inp) + # Original aten path (bit-identical bookkeeping). + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and payload is not None: + # cache last frame of last two chunk + cache_x = torch.cat( + [payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + elif cache_x.shape[2] < 2 and is_rep: + cache_x = torch.cat( + [torch.zeros_like(cache_x).to(cache_x.device), cache_x], + dim=2, + ) + out = conv(x) if payload is None else conv(x, payload) + cache_list[idx] = cache_x + return out + + class QwenImageCausalConv3d(nn.Conv3d): r""" A custom 3D causal convolution layer with feature caching support. @@ -88,6 +165,18 @@ class QwenImageCausalConv3d(nn.Conv3d): def forward(self, x, cache_x=None): padding = list(self._padding) + if ( + any(padding) + and _fused_conv_cache_supported(self, x) + and ( + cache_x is None + or (cache_x.device == x.device and cache_x.dtype == x.dtype) + ) + ): + # Bit-exact: cat + pad + channels_last_3d relayout in one pass. + inp = cat_pad_channels_last_3d(x, cache_x, padding) + if inp is not None: + return super().forward(inp) x = causal_conv3d_cat_pad(x, cache_x, padding) return super().forward(x) @@ -138,6 +227,10 @@ class QwenImageUpsample(nn.Upsample): """ def forward(self, x): + if current_platform.is_amp_supported(): + # nearest-exact is a pure gather, so bf16/fp16 is bit-identical to + # the fp32 round trip; skip its three full-size copies. + return super().forward(x) return super().forward(x.float()).type_as(x) @@ -220,36 +313,7 @@ class QwenImageResample(nn.Module): feat_cache[idx] = "Rep" feat_idx[0] += 1 else: - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if ( - cache_x.shape[2] < 2 - and feat_cache[idx] is not None - and feat_cache[idx] != "Rep" - ): - # cache last frame of last two chunk - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :] - .unsqueeze(2) - .to(cache_x.device), - cache_x, - ], - dim=2, - ) - if ( - cache_x.shape[2] < 2 - and feat_cache[idx] is not None - and feat_cache[idx] == "Rep" - ): - cache_x = torch.cat( - [torch.zeros_like(cache_x).to(cache_x.device), cache_x], - dim=2, - ) - if feat_cache[idx] == "Rep": - x = self.time_conv(x) - else: - x = self.time_conv(x, feat_cache[idx]) - feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.time_conv, x, feat_cache, idx) feat_idx[0] += 1 x = x.reshape(b, 2, c, t, h, w) @@ -327,18 +391,7 @@ class QwenImageResidualBlock(nn.Module): if feat_cache is not None: idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - - x = self.conv1(x, feat_cache[idx]) - feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv1, x, feat_cache, idx) feat_idx[0] += 1 else: x = self.conv1(x) @@ -352,18 +405,7 @@ class QwenImageResidualBlock(nn.Module): if feat_cache is not None: idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - - x = self.conv2(x, feat_cache[idx]) - feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv2, x, feat_cache, idx) feat_idx[0] += 1 else: x = self.conv2(x) @@ -421,7 +463,16 @@ class QwenImageAttentionBlock(nn.Module): x = x.view(batch_size, time, channels, height, width) x = x.permute(0, 2, 1, 3, 4) - x = x + identity + gate = getattr(self, "_sgl_gate", None) + if gate is not None and gate.enabled and not torch.compiler.is_compiling(): + # ``identity`` carries the decoder's channels_last_3d layout; + # putting it first makes the sum inherit that layout so the next + # block's fused RMSNorm+SiLU applies. The add is commutative, but + # the layout changes the reduction order of an eager norm that + # follows under fp32 autocast, so this stays quality-gated. + x = identity + x + else: + x = x + identity if self.spatial_parallel: x = chunk_height_for_parallel_decode(x) return x @@ -565,18 +616,7 @@ class QwenImageEncoder3d(nn.Module): def forward(self, x, feat_cache=None, feat_idx=[0]): if feat_cache is not None: idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv_in(x, feat_cache[idx]) - feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv_in, x, feat_cache, idx) feat_idx[0] += 1 else: x = self.conv_in(x) @@ -596,18 +636,7 @@ class QwenImageEncoder3d(nn.Module): x = self.nonlinearity(x) if feat_cache is not None: idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv_out(x, feat_cache[idx]) - feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv_out, x, feat_cache, idx) feat_idx[0] += 1 else: x = self.conv_out(x) @@ -828,18 +857,7 @@ class QwenImageDecoder3d(nn.Module): ## conv1 if feat_cache is not None: idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv_in(x, feat_cache[idx]) - feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv_in, x, feat_cache, idx) feat_idx[0] += 1 else: x = self.conv_in(x) @@ -856,18 +874,7 @@ class QwenImageDecoder3d(nn.Module): x = self.nonlinearity(x) if feat_cache is not None: idx = feat_idx[0] - cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and feat_cache[idx] is not None: - # cache last frame of last two chunk - cache_x = torch.cat( - [ - feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), - cache_x, - ], - dim=2, - ) - x = self.conv_out(x, feat_cache[idx]) - feat_cache[idx] = cache_x + x = _run_cached_causal_conv(self.conv_out, x, feat_cache, idx) feat_idx[0] += 1 else: x = self.conv_out(x) diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py b/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py index eb53a88f7..fa9ecea36 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -"""CUDA fast path for the Wan VAE decoder (AutoencoderKLWan). +"""CUDA fast path for the Wan-family VAE decoders (AutoencoderKLWan and the +Qwen-Image VAE, which is the Wan 2.1 VAE under other class names). Fuses every decoder ``WanRMS_norm -> SiLU`` chain into one Triton kernel on the channels_last_3d layout. Wrappers are installed once at VAE load and @@ -58,6 +59,62 @@ class FusedWanRMSNormSiLU(nn.Module): return F.silu(F.normalize(x, dim=1) * self.scale * self.gamma + self.bias) +class GatedChannelsLastUpsample(nn.Module): + """Wan-style ``Resample`` upsample that keeps the decoder channels_last. + + ``Resample.forward`` reaches its 2D ``Upsample -> Conv2d`` through + ``x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)``. On a channels_last_3d + tensor with ``b * t == 1`` that view keeps a degenerate batch stride which + ``is_contiguous(channels_last)`` accepts but aten's + ``suggest_memory_format`` does not, so the upsample and the conv2d run + their NCHW kernels (cuDNN converts back and forth internally) and every + up block starts with an NCDHW tensor that the fused RMSNorm+SiLU cannot + take. With the gate on, re-express the same memory with canonical NHWC + strides (a pure view) so the conv2d runs channels_last end-to-end. The + NHWC conv2d is not guaranteed to pick the same cuDNN algorithm as the NCHW + one, so the stride canonicalisation stays behind the quality gate; with + the gate off the layout is left exactly as the eager module sees it. + """ + + def __init__(self, upsample: nn.Upsample, gate: VaeFastPathGate) -> None: + super().__init__() + self._sgl_upsample = upsample + self._sgl_gate = gate + + def forward(self, x: torch.Tensor) -> torch.Tensor: + up = self._sgl_upsample + if torch.compiler.is_compiling() or x.dim() != 4: + return up(x) + _, c, h, w = x.shape + canonical = (h * w * c, 1, w * c, c) + # ``is_contiguous(channels_last)`` guarantees every dim of size > 1 + # already has its canonical stride; only size-1 dims (the merged + # batch dim here) can carry a stray one, and their stride is never + # used to address memory. So ``as_strided`` to the canonical strides + # is a pure re-labelling of the same elements, for any batch size. + if ( + self._sgl_gate.enabled + and x.stride() != canonical + and x.is_contiguous(memory_format=torch.channels_last) + ): + x = x.as_strided(x.shape, canonical) + return up(x) + + +def _install_module_gates( + decoder: nn.Module, gate: VaeFastPathGate, module_classes: tuple[type, ...] +) -> int: + """Hand the gate to layout-sensitive modules that consult ``_sgl_gate`` in + their own forward (the attention block's residual operand order, so the + sum keeps the channels_last_3d layout).""" + count = 0 + for m in decoder.modules(): + if type(m) in module_classes: + m._sgl_gate = gate + count += 1 + return count + + def _is_plain_silu(act: object) -> bool: return isinstance(act, nn.SiLU) and not act.inplace @@ -70,34 +127,42 @@ def _norm_fusable(norm: object, wan_rms_norm_cls: type) -> bool: ) -def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None: - """Wrap every decoder ``WanRMS_norm -> SiLU`` chain; ``None`` = fail closed.""" - from sglang.multimodal_gen.runtime.models.vaes.wanvae import ( - WanResidualBlock, - WanRMS_norm, - ) +def _install_norm_silu( + decoder: nn.Module, + gate: VaeFastPathGate, + *, + residual_block_cls: type, + rms_norm_cls: type, + label: str, +) -> int | None: + """Wrap every decoder ``RMS_norm -> SiLU`` chain; ``None`` = fail closed. - res_blocks = [m for m in decoder.modules() if isinstance(m, WanResidualBlock)] + ``residual_block_cls`` / ``rms_norm_cls`` name the model's own Wan-style + residual block and channel-first RMSNorm (``WanResidualBlock`` / + ``WanRMS_norm`` or their Qwen-Image twins). + """ + res_blocks = [m for m in decoder.modules() if isinstance(m, residual_block_cls)] eligible = [ m for m in res_blocks - if type(m) is WanResidualBlock + if type(m) is residual_block_cls and _is_plain_silu(m.nonlinearity) - and _norm_fusable(m.norm1, WanRMS_norm) - and _norm_fusable(m.norm2, WanRMS_norm) + and _norm_fusable(m.norm1, rms_norm_cls) + and _norm_fusable(m.norm2, rms_norm_cls) ] if len(eligible) != len(res_blocks): logger.warning( - "Wan VAE: %d/%d residual blocks non-standard; skipping fast path.", + "%s: %d/%d residual blocks non-standard; skipping fast path.", + label, len(res_blocks) - len(eligible), len(res_blocks), ) return None if not ( - _norm_fusable(getattr(decoder, "norm_out", None), WanRMS_norm) + _norm_fusable(getattr(decoder, "norm_out", None), rms_norm_cls) and _is_plain_silu(getattr(decoder, "nonlinearity", None)) ): - logger.warning("Wan VAE: non-standard output head; skipping fast path.") + logger.warning("%s: non-standard output head; skipping fast path.", label) return None count = 0 @@ -111,32 +176,114 @@ def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None: return count + 1 +def _install_channels_last_upsample( + decoder: nn.Module, gate: VaeFastPathGate, upsample_cls: type +) -> int: + """Wrap the 2D nearest upsample of every ``Resample`` (index 0 of its + ``resample`` Sequential, no parameters, so state_dict names are kept).""" + count = 0 + for m in decoder.modules(): + seq = getattr(m, "resample", None) + if ( + isinstance(seq, nn.Sequential) + and len(seq) >= 1 + and type(seq[0]) is upsample_cls + ): + seq[0] = GatedChannelsLastUpsample(seq[0], gate) + count += 1 + return count + + def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module: """Install the quality-gated CUDA Wan VAE decoder fast path.""" from sglang.multimodal_gen.runtime.models.vaes.wanvae import ( AutoencoderKLWan, WanDecoder3d, + WanResidualBlock, + WanRMS_norm, ) if not isinstance(vae, AutoencoderKLWan): return vae + return _maybe_optimize_wan_family_vae( + vae, + decoder_cls=WanDecoder3d, + residual_block_cls=WanResidualBlock, + rms_norm_cls=WanRMS_norm, + label="Wan VAE", + ) + + +def maybe_optimize_qwen_image_vae(vae: nn.Module) -> nn.Module: + """Install the quality-gated CUDA Qwen-Image VAE decoder fast path. + + The Qwen-Image VAE is the Wan 2.1 VAE (same ``F.normalize``-based + channel-first RMSNorm, same residual block layout), so it takes the same + ``RMSNorm -> SiLU`` fusion under the same gate. + """ + from sglang.multimodal_gen.runtime.models.vaes.autoencoder_kl_qwenimage import ( + AutoencoderKLQwenImage, + QwenImageAttentionBlock, + QwenImageDecoder3d, + QwenImageResidualBlock, + QwenImageRMS_norm, + QwenImageUpsample, + ) + + if not isinstance(vae, AutoencoderKLQwenImage): + return vae + return _maybe_optimize_wan_family_vae( + vae, + decoder_cls=QwenImageDecoder3d, + residual_block_cls=QwenImageResidualBlock, + rms_norm_cls=QwenImageRMS_norm, + upsample_cls=QwenImageUpsample, + gated_module_classes=(QwenImageAttentionBlock,), + label="Qwen-Image VAE", + ) + + +def _maybe_optimize_wan_family_vae( + vae: nn.Module, + *, + decoder_cls: type, + residual_block_cls: type, + rms_norm_cls: type, + label: str, + upsample_cls: type | None = None, + gated_module_classes: tuple[type, ...] = (), +) -> nn.Module: decoder = getattr(vae, "decoder", None) - if type(decoder) is not WanDecoder3d: + if type(decoder) is not decoder_cls: return vae if decoder.use_parallel_decode and decoder.world_size > 1: - logger.info("Wan VAE: spatial-parallel decode; skipping fast path.") + logger.info("%s: spatial-parallel decode; skipping fast path.", label) return vae if not _HAS_TRITON: - logger.warning("Wan VAE: Triton unavailable; skipping fast path.") + logger.warning("%s: Triton unavailable; skipping fast path.", label) return vae gate = VaeFastPathGate() - n_norm = _install_norm_silu(decoder, gate) + n_norm = _install_norm_silu( + decoder, + gate, + residual_block_cls=residual_block_cls, + rms_norm_cls=rms_norm_cls, + label=label, + ) if n_norm is None: return vae + n_up = 0 + if upsample_cls is not None: + n_up = _install_channels_last_upsample(decoder, gate, upsample_cls) + if gated_module_classes: + _install_module_gates(decoder, gate, gated_module_classes) register_vae_fast_path_gate(vae, gate) logger.info( - "Wan VAE: installed quality-gated fast path (%d RMSNorm+SiLU fusions).", + "%s: installed quality-gated fast path (%d RMSNorm+SiLU fusions, " + "%d channels_last upsamples).", + label, n_norm, + n_up, ) return vae diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index ad6a4df46..844ae4006 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -745,8 +745,8 @@ class CudaPlatformBase(Platform): @classmethod def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module: - """Install the quality-gated FLUX.2 / AutoencoderKL / Wan VAE decoder - fast paths. + """Install the quality-gated FLUX.2 / AutoencoderKL / Wan / Qwen-Image + VAE decoder fast paths. Requests with quality="extra-high" or "high" run the fast paths; the "lossless" default runs the original module path bit-for-bit. See @@ -758,12 +758,14 @@ class CudaPlatformBase(Platform): maybe_optimize_flux2_vae, ) from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import ( + maybe_optimize_qwen_image_vae, maybe_optimize_wan_vae, ) vae = maybe_optimize_flux2_vae(vae) vae = maybe_optimize_autoencoder_kl(vae) vae = maybe_optimize_wan_vae(vae) + vae = maybe_optimize_qwen_image_vae(vae) except Exception: logger.warning( "Failed to apply CUDA VAE optimizations; using the unmodified VAE.", diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index 85d6b984e..238fde2e5 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py @@ -114,11 +114,18 @@ from sglang.multimodal_gen.runtime.models.dits.sana import ( from sglang.multimodal_gen.runtime.models.dits.sana import ( sana_ln_modulate, ) +from sglang.multimodal_gen.runtime.models.vaes import ( + autoencoder_kl_qwenimage as qwen_vae, +) from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt +from sglang.multimodal_gen.runtime.models.vaes import ( + wan_vae_cuda_opt, +) from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import use_vae_fast_path from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import ( FusedWanRMSNormSiLU, + GatedChannelsLastUpsample, VaeFastPathGate, ) from sglang.multimodal_gen.runtime.models.vaes.wanvae import WanRMS_norm @@ -1053,6 +1060,154 @@ def test_wan_vae_rejects_empty_input() -> None: assert not can_use_wan_rmsnorm_silu(x, gamma, None) +# ------------------------------------------------------------------------- +# Qwen-Image VAE (Wan 2.1 VAE) -- lossless causal-conv data movement and the +# quality-gated RMSNorm+SiLU / channels_last upsample fast path +# ------------------------------------------------------------------------- + + +def _qwen_causal_conv(cin, cout, channels_last=True): + conv = qwen_vae.QwenImageCausalConv3d(cin, cout, 3, padding=1).to( + "cuda", torch.bfloat16 + ) + if channels_last: + conv.weight.data = conv.weight.data.to(memory_format=torch.channels_last_3d) + return conv + + +@torch.no_grad() +def test_qwen_vae_causal_conv_cat_pad_is_bit_exact() -> None: + # Fused cat + pad + relayout must reproduce the aten chain for no cache, + # a one-frame cache (partial zero fill) and a full two-frame cache. + torch.manual_seed(0) + conv = _qwen_causal_conv(8, 16) + ref = _qwen_causal_conv(8, 16, channels_last=False) + ref.load_state_dict(conv.state_dict()) + x = _wan_cl3d((1, 8, 1, 12, 10), torch.bfloat16) + for cache_t in (0, 1, 2): + cache = _wan_cl3d((1, 8, cache_t, 12, 10), torch.bfloat16) if cache_t else None + # The reference is the eager chain on the same (channels_last) weights. + expected = nn.Conv3d.forward( + conv, + qwen_vae.causal_conv3d_cat_pad(x, cache, list(conv._padding)).contiguous( + memory_format=torch.channels_last_3d + ), + ) + assert qwen_vae._fused_conv_cache_supported(conv, x) + assert torch.equal(conv(x, cache), expected) + # Contiguous weights take the original path unchanged. + assert not qwen_vae._fused_conv_cache_supported(ref, x) + + +@torch.no_grad() +def test_qwen_vae_run_cached_causal_conv_matches_eager_bookkeeping() -> None: + # Stream three single-frame chunks (the image decode shape, and the first + # chunks of a clip) through one conv slot on the fused and the eager path; + # outputs and cache contents must agree, including the "Rep" slot. + torch.manual_seed(0) + fused = _qwen_causal_conv(8, 8) + eager = _qwen_causal_conv(8, 8, channels_last=False) + eager.load_state_dict(fused.state_dict()) + chunks = [_wan_cl3d((1, 8, 1, 6, 6), torch.bfloat16) for _ in range(3)] + for first in (None, "Rep"): + fused_cache, eager_cache = [first], [first] + for x in chunks: + out_f = qwen_vae._run_cached_causal_conv(fused, x, fused_cache, 0) + out_e = qwen_vae._run_cached_causal_conv(eager, x, eager_cache, 0) + assert torch.equal(out_f, out_e) + # The fused path only ever stores tensors that the eager path would + # also consume: same values, possibly wider (never narrower). + fc, ec = fused_cache[0], eager_cache[0] + assert fc.shape[2] >= ec.shape[2] + assert torch.equal(fc[:, :, -ec.shape[2] :], ec) + + +@torch.no_grad() +def test_qwen_vae_upsample_skips_fp32_round_trip_bit_exactly() -> None: + up = qwen_vae.QwenImageUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact") + x = torch.randn(1, 8, 12, 10, device="cuda", dtype=torch.bfloat16).contiguous( + memory_format=torch.channels_last + ) + assert torch.equal(up(x), nn.Upsample.forward(up, x.float()).type_as(x)) + + +@torch.no_grad() +def test_qwen_vae_gate_dispatch() -> None: + torch.cuda.manual_seed(0) + norm = qwen_vae.QwenImageRMS_norm(96, images=False).to( + device="cuda", dtype=torch.bfloat16 + ) + norm.gamma.add_(torch.randn_like(norm.gamma)) + gate = VaeFastPathGate() + fused = FusedWanRMSNormSiLU(norm, gate) + assert [n for n, _ in fused.named_parameters()] == ["gamma"] + x = _wan_cl3d((1, 96, 1, 10, 14), torch.bfloat16) + assert torch.equal(fused(x), nn.SiLU()(norm(x))) + gate.enabled = True + expected = wan_rmsnorm_silu(x, norm.gamma, rms_scale=float(norm.scale)) + assert torch.equal(fused(x), expected) + + +@torch.no_grad() +def test_qwen_vae_channels_last_upsample_keeps_layout() -> None: + gate = VaeFastPathGate() + up = GatedChannelsLastUpsample( + qwen_vae.QwenImageUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + gate, + ) + x5 = _wan_cl3d((1, 8, 1, 6, 6), torch.bfloat16) + x = x5.permute(0, 2, 1, 3, 4).reshape(1, 8, 6, 6) # degenerate batch stride + ref = nn.Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact")(x) + off = up(x) + assert torch.equal(off, ref) and off.is_contiguous() # NCHW, as today + gate.enabled = True + on = up(x) + assert torch.equal(on, ref) + assert on.is_contiguous(memory_format=torch.channels_last) + assert on.stride() == (12 * 12 * 8, 1, 12 * 8, 8) + + +@torch.no_grad() +def test_qwen_vae_decoder_install_is_gated_and_close() -> None: + torch.manual_seed(0) + dec = qwen_vae.QwenImageDecoder3d( + dim=16, z_dim=4, dim_mult=[1, 1], num_res_blocks=1, temperal_upsample=[False] + ).to("cuda", torch.bfloat16) + for m in dec.modules(): + if isinstance(m, nn.Conv3d): + m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) + for p in dec.parameters(): + p.add_(0.1 * torch.randn_like(p)) + keys = set(dec.state_dict()) + z = _wan_cl3d((1, 4, 1, 6, 6), torch.bfloat16) + ref = dec(z, feat_cache=[None] * 64, feat_idx=[0]) + + gate = VaeFastPathGate() + n_norm = wan_vae_cuda_opt._install_norm_silu( + dec, + gate, + residual_block_cls=qwen_vae.QwenImageResidualBlock, + rms_norm_cls=qwen_vae.QwenImageRMS_norm, + label="test", + ) + n_up = wan_vae_cuda_opt._install_channels_last_upsample( + dec, gate, qwen_vae.QwenImageUpsample + ) + # mid block (2 resnets) + 2 up blocks x (num_res_blocks + 1) resnets + # -> 6 resnets x 2 norms + norm_out; only the first up block upsamples. + assert n_norm == 13 and n_up == 1 + assert set(dec.state_dict()) == keys + assert torch.equal(dec(z, feat_cache=[None] * 64, feat_idx=[0]), ref) + gate.enabled = True + out = dec(z, feat_cache=[None] * 64, feat_idx=[0]) + assert out.shape == ref.shape + # Not bit-exact (fp32 statistics in the fused norm), but a bf16-rounding + # level perturbation: bound the relative error energy over the whole + # output rather than per element (random weights, unnormalised scale). + rel = ((out.float() - ref.float()).norm() / ref.float().norm()).item() + assert rel < 2e-2, rel + + # ------------------------------------------------------------------------- # FLUX.2 VAE -- fused GroupNorm+SiLU and folded 2x upsample conv # -------------------------------------------------------------------------