[Diffusion] Port the Wan VAE decoder fast paths to the Qwen-Image VAE (#38020)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

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
+  decoder rewrites used by FLUX.1/FLUX.2/Z-Image/SD3, and Wan / Qwen-Image
-  RMSNorm+SiLU.
+  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`.

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()
+                    x = _run_cached_causal_conv(self.time_conv, x, feat_cache, idx)
-                    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
                     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()
+            x = _run_cached_causal_conv(self.conv1, x, feat_cache, idx)
-            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
             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()
+            x = _run_cached_causal_conv(self.conv2, x, feat_cache, idx)
-            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
             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()
+            x = _run_cached_causal_conv(self.conv_in, x, feat_cache, idx)
-            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
             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()
+            x = _run_cached_causal_conv(self.conv_out, x, feat_cache, idx)
-            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
             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()
+            x = _run_cached_causal_conv(self.conv_in, x, feat_cache, idx)
-            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
             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()
+            x = _run_cached_causal_conv(self.conv_out, x, feat_cache, idx)
-            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
             feat_idx[0] += 1
         else:
             x = self.conv_out(x)

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:
+def _install_norm_silu(
-    """Wrap every decoder ``WanRMS_norm -> SiLU`` chain; ``None`` = fail closed."""
+    decoder: nn.Module,
-    from sglang.multimodal_gen.runtime.models.vaes.wanvae import (
+    gate: VaeFastPathGate,
-        WanResidualBlock,
+    *,
-        WanRMS_norm,
+    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.norm1, rms_norm_cls)
-        and _norm_fusable(m.norm2, WanRMS_norm)
+        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

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
+        """Install the quality-gated FLUX.2 / AutoencoderKL / Wan / Qwen-Image
-        fast paths.
+        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.",

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
 # -------------------------------------------------------------------------