diff --git a/python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py b/python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py new file mode 100644 index 000000000..7af90e106 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Channels-last two-pass GroupNorm(+SiLU) Triton kernels. + +Relationship to ``group_norm_silu.py`` (``triton_group_norm_silu``): that +kernel serves the general NCHW-contiguous case (any channels-per-group, any +ndim, always applies SiLU) and keeps backing ``apply_group_norm_silu`` for +the HunyuanVAE / latent-upsampler paths. This module is a complementary +kernel for the channels_last VAE decoder fast path with a different contract: + +- activations whose channel dim is innermost ((N, H, W, C) channels_last + views or (N, L, C) rows) are normalized without any layout round-trip, + which is what makes a channels_last decoder run end-to-end without + nchwToNhwc transposes; +- fp32 statistics, with the affine transform folded into per-(batch, channel) + ``scale``/``shift`` in a separate finalize kernel, so the apply pass is a + pure elementwise kernel; +- optional SiLU epilogue (``apply_silu=False`` gives plain GroupNorm); +- restricted static shapes: power-of-two ``C <= 2048`` that ``num_groups`` + divides. Callers must treat a ``None`` return as "unsupported" and fall + back to their reference path. +""" + +import torch +import triton # type: ignore +import triton.language as tl # type: ignore + +_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32} +_MAX_CHANNELS = 2048 + + +@triton.jit +def _gn_partial_rows_kernel( + x_ptr, + psum_ptr, + psq_ptr, + rows, + rows_per_prog, + C: tl.constexpr, + BLOCK_R: tl.constexpr, +): + # x is (N, rows, C) with C innermost-contiguous (channels_last view). + chunk = tl.program_id(0).to(tl.int64) + n = tl.program_id(1).to(tl.int64) + nchunks = tl.num_programs(0) + row0 = chunk * rows_per_prog + cols = tl.arange(0, C) + acc_s = tl.zeros((C,), tl.float32) + acc_q = tl.zeros((C,), tl.float32) + x_base = x_ptr + n * rows * C + for r_off in range(0, rows_per_prog, BLOCK_R): + rs = row0 + r_off + tl.arange(0, BLOCK_R) + m = rs < rows + x = tl.load( + x_base + rs[:, None] * C + cols[None, :], + mask=m[:, None], + other=0.0, + ).to(tl.float32) + acc_s += tl.sum(x, 0) + acc_q += tl.sum(x * x, 0) + out_off = (n * nchunks + chunk) * C + cols + tl.store(psum_ptr + out_off, acc_s) + tl.store(psq_ptr + out_off, acc_q) + + +@triton.jit +def _gn_finalize_kernel( + psum_ptr, + psq_ptr, + w_ptr, + b_ptr, + ss_ptr, + nchunks, + group_numel, + eps, + C, + CPG: tl.constexpr, + BLOCK_K: tl.constexpr, +): + g = tl.program_id(0).to(tl.int64) + n = tl.program_id(1).to(tl.int64) + cols = g * CPG + tl.arange(0, CPG) + s = tl.zeros((), tl.float32) + q = tl.zeros((), tl.float32) + for k0 in range(0, nchunks, BLOCK_K): + ks = k0 + tl.arange(0, BLOCK_K) + m = ks < nchunks + offs = (n * nchunks + ks)[:, None] * C + cols[None, :] + s += tl.sum(tl.load(psum_ptr + offs, mask=m[:, None], other=0.0)) + q += tl.sum(tl.load(psq_ptr + offs, mask=m[:, None], other=0.0)) + mean = s / group_numel + var = q / group_numel - mean * mean + var = tl.maximum(var, 0.0) + rstd = tl.rsqrt(var + eps) + w = tl.load(w_ptr + cols).to(tl.float32) + b = tl.load(b_ptr + cols).to(tl.float32) + scale = w * rstd + shift = b - mean * scale + tl.store(ss_ptr + n * 2 * C + cols, scale) + tl.store(ss_ptr + (n * 2 + 1) * C + cols, shift) + + +@triton.jit +def _gn_apply_rows_kernel( + x_ptr, + ss_ptr, + y_ptr, + rows, + C: tl.constexpr, + BLOCK_R: tl.constexpr, + SILU: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + n = tl.program_id(1).to(tl.int64) + cols = tl.arange(0, C) + scale = tl.load(ss_ptr + n * 2 * C + cols) + shift = tl.load(ss_ptr + (n * 2 + 1) * C + cols) + rs = pid * BLOCK_R + tl.arange(0, BLOCK_R) + m = rs < rows + offs = n * rows * C + rs[:, None] * C + cols[None, :] + x = tl.load(x_ptr + offs, mask=m[:, None], other=0.0).to(tl.float32) + y = x * scale[None, :] + shift[None, :] + if SILU: + y = y * tl.sigmoid(y) + tl.store(y_ptr + offs, y.to(y_ptr.dtype.element_ty), mask=m[:, None]) + + +def _gn_silu_rows(x3, weight, bias, num_groups, eps, apply_silu): + """x3: (N, R, C) contiguous with C innermost. Returns same-shape tensor.""" + n_batch, rows, c = x3.shape + cpg = c // num_groups + block_r = max(1, 8192 // c) + rows_per_prog = block_r * 32 + nchunks = triton.cdiv(rows, rows_per_prog) + psum = torch.empty((n_batch, nchunks, c), device=x3.device, dtype=torch.float32) + psq = torch.empty_like(psum) + _gn_partial_rows_kernel[(nchunks, n_batch)]( + x3, psum, psq, rows, rows_per_prog, C=c, BLOCK_R=block_r, num_warps=4 + ) + ss = torch.empty((n_batch, 2, c), device=x3.device, dtype=torch.float32) + block_k = max(1, min(4096 // cpg, triton.next_power_of_2(nchunks))) + _gn_finalize_kernel[(num_groups, n_batch)]( + psum, + psq, + weight, + bias, + ss, + nchunks, + rows * cpg, + eps, + c, + CPG=cpg, + BLOCK_K=block_k, + num_warps=4, + ) + y3 = torch.empty_like(x3) + _gn_apply_rows_kernel[(triton.cdiv(rows, block_r), n_batch)]( + x3, ss, y3, rows, C=c, BLOCK_R=block_r, SILU=apply_silu, num_warps=4 + ) + return y3 + + +def _twopass_supported(x, weight, bias, num_groups) -> bool: + """Tensor-level support check shared by the 4D and rows entry points.""" + if not (x.is_cuda and not torch.is_grad_enabled()): + return False + if x.requires_grad or x.dtype not in _SUPPORTED_DTYPES: + return False + if weight is None or bias is None: + return False + c = x.shape[1] if x.dim() == 4 else x.shape[-1] + if weight.shape != (c,) or bias.shape != (c,): + return False + if num_groups < 1 or c % num_groups != 0: + return False + # tl.arange needs a power-of-two C; num_groups divides it, so the + # channels-per-group finalize block is a power of two as well. + return triton.next_power_of_2(c) == c and c <= _MAX_CHANNELS + + +def group_norm_silu_4d( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int, + eps: float, + apply_silu: bool = True, +) -> torch.Tensor | None: + """Fused GroupNorm(+SiLU) for a channels_last 4D (N, C, H, W) activation. + + Runs the rows kernel on the free (N, H*W, C) view (no layout copy) and + preserves the channels_last output layout. Returns ``None`` when the + input is unsupported; callers must fall back to their reference path. + """ + if x.dim() != 4 or not _twopass_supported(x, weight, bias, num_groups): + return None + n_batch, c, h, w = x.shape + # c > 1 and a non-trivial spatial extent make the channels_last check + # unambiguous (degenerate shapes are contiguous in both formats). + if not ( + c > 1 + and (h > 1 or w > 1) + and x.is_contiguous(memory_format=torch.channels_last) + ): + return None + x3 = x.permute(0, 2, 3, 1).reshape(n_batch, h * w, c) + y3 = _gn_silu_rows(x3, weight, bias, num_groups, eps, apply_silu) + return y3.reshape(n_batch, h, w, c).permute(0, 3, 1, 2) + + +def group_norm_silu_rows( + x3: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int, + eps: float, + apply_silu: bool = True, +) -> torch.Tensor | None: + """Fused GroupNorm(+SiLU) over (N, L, C) rows (C = channels, innermost). + + Returns ``None`` when the input is unsupported; callers must fall back. + """ + if x3.dim() != 3 or not x3.is_contiguous(): + return None + if not _twopass_supported(x3, weight, bias, num_groups): + return None + return _gn_silu_rows(x3, weight, bias, num_groups, eps, apply_silu) + + +__all__ = [ + "group_norm_silu_4d", + "group_norm_silu_rows", +] 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 new file mode 100644 index 000000000..ee961ae78 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA fast paths for the FLUX.2 VAE decoder (AutoencoderKLFlux2, diffusers +``Decoder``) on image workloads. + +All rewrites are mathematically exact re-associations of the original +operators (weight folding is done lazily in fp32 on first fast-path use and +written back to the model compute dtype). The wrappers are installed once at +VAE load and stay in place; each forward dispatches on a shared +request-scoped :class:`VaeFastPathGate`: requests with ``quality == "high"`` +run the fast paths, the ``"lossless"`` default runs the original module path +bit-for-bit. + +- channels_last: run the decoder in channels_last so cuDNN convolutions run + natively in NHWC (removes the nchwToNhwc/nhwcToNchw transpose kernels + around every conv). The parameter layout is swapped at decode entry to + match the gate, so lossless decodes always run the NCHW baseline kernels + bit-for-bit. The mid-block attention needs a layout-safe forward because + diffusers' ``AttnProcessor2_0`` calls ``.view`` on the 4D activation, + which is illegal for channels_last tensors. +- norm+SiLU: two-pass channels_last GroupNorm(+SiLU) Triton fusion (fp32 + statistics) for the ResnetBlock2D norm1/norm2 + SiLU chains and the + decoder ``conv_norm_out``/``conv_act`` tail, which upcast to fp32 under + autocast and dominate the decode profile. +- fused upsample: nearest-2x upsample + Conv2d(3x3, p1) == + ConvTranspose2d(k4, s2, p1) with a lazily-summed kernel. Removes the 4x + upsampled intermediate materialization. +- attention V/proj fold: fold the attention output projection into the V + projection of the single-head mid-block attention (softmax rows sum to 1, + so ``A @ (V W_v^T + b_v) W_o^T + b_o == A @ (V W_v'^T + b')``). + +Install is all-or-nothing and fail-closed: without Triton, or with any +attention block lacking the layout-safe rewrite, no wrapper is installed and +every request runs the unmodified decoder. +""" + +from types import MethodType + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +try: + from sglang.kernels.ops.diffusion.triton.group_norm_silu_twopass import ( + group_norm_silu_4d, + group_norm_silu_rows, + ) + + _HAS_TRITON = True +except ImportError: # pragma: no cover + _HAS_TRITON = False + + +class VaeFastPathGate: + """Mutable fast-path flag shared by every wrapper of one VAE. + + Published on the VAE as ``_sgl_vae_fast_path_gate``; ``DecodingStage`` + enables it for the duration of a decode when the request's ``quality`` + sampling param is ``"high"``. + """ + + __slots__ = ("enabled",) + + def __init__(self) -> None: + self.enabled = False + + +GATE_ATTR = "_sgl_vae_fast_path_gate" + + +# --------------------------------------------------------------------------- +# Fuse A: two-pass GroupNorm(+SiLU) fusion (channels_last Triton kernel) +# --------------------------------------------------------------------------- + + +class FusedGroupNormSiLU(nn.Module): + """GroupNorm + SiLU fused with the two-pass channels_last Triton kernel. + + fp32 statistics and affine/SiLU application, output in the input dtype. + Falls back to the original module chain (norm + F.silu, bit-identical to + the original norm + nn.SiLU pair) for unsupported inputs and whenever + the fast-path gate is disabled. + """ + + def __init__(self, norm: nn.GroupNorm, gate: VaeFastPathGate) -> None: + super().__init__() + # Keep the original GroupNorm state-dict layout (``weight``/``bias``) + # so checkpoint loading and component-accuracy weight transfer do not + # see wrapper-specific ``norm.*`` parameter names. + self.num_groups = norm.num_groups + self.num_channels = norm.num_channels + self.eps = norm.eps + self.affine = norm.affine + self.weight = norm.weight + self.bias = norm.bias + self._sgl_gate = gate + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self._sgl_gate.enabled and x.dim() == 4: + y = group_norm_silu_4d( + x, + self.weight, + self.bias, + self.num_groups, + self.eps, + apply_silu=True, + ) + if y is not None: + return y + return F.silu( + F.group_norm(x, self.num_groups, self.weight, self.bias, self.eps) + ) + + +def _install_norm_silu(decoder, resnet_cls, gate: VaeFastPathGate) -> int: + def _fusable(norm) -> bool: + return ( + type(norm) is nn.GroupNorm + and norm.affine + and norm.weight is not None + and norm.bias is not None + ) + + count = 0 + for m in decoder.modules(): + if ( + type(m) is resnet_cls + and m.time_emb_proj is None + # "default"/"group" apply plain norm2 (+ SiLU); "scale_shift" and + # "spatial" modify the activation between norm2 and SiLU. + and m.time_embedding_norm in ("default", "group") + and m.upsample is None + and m.downsample is None + and isinstance(m.nonlinearity, nn.SiLU) + and _fusable(m.norm1) + and _fusable(m.norm2) + ): + m.norm1 = FusedGroupNormSiLU(m.norm1, gate) + m.norm2 = FusedGroupNormSiLU(m.norm2, gate) + m.nonlinearity = nn.Identity() + count += 2 + if ( + type(getattr(decoder, "conv_norm_out", None)) is nn.GroupNorm + and isinstance(getattr(decoder, "conv_act", None), nn.SiLU) + and _fusable(decoder.conv_norm_out) + ): + decoder.conv_norm_out = FusedGroupNormSiLU(decoder.conv_norm_out, gate) + decoder.conv_act = nn.Identity() + count += 1 + return count + + +# --------------------------------------------------------------------------- +# Fuse B: nearest-2x upsample + Conv2d(3x3, p1) == ConvTranspose2d(k4, s2, p1) +# --------------------------------------------------------------------------- + +# Which 3x3 conv taps sum into each 4x4 transposed-conv tap (per spatial axis). +_UPSAMPLE_TAP_MAP = {0: (2,), 1: (1, 2), 2: (0, 1), 3: (0,)} + + +def _fold_upsample2x_conv2d_weight(conv: nn.Conv2d) -> torch.Tensor: + """Sum the 3x3 conv taps into the equivalent ConvTranspose2d(k4) kernel.""" + w = conv.weight.detach().float() # [Cout, Cin, 3, 3] + cout, cin = w.shape[:2] + wt = w.new_zeros(cin, cout, 4, 4) # ConvTranspose2d layout + for a in range(4): + for b in range(4): + acc = w.new_zeros(cout, cin) + for i in _UPSAMPLE_TAP_MAP[a]: + for j in _UPSAMPLE_TAP_MAP[b]: + acc += w[:, :, i, j] + wt[:, :, a, b] = acc.t() + wt = wt.to(conv.weight.dtype) + if conv.weight.is_contiguous(memory_format=torch.channels_last): + wt = wt.contiguous(memory_format=torch.channels_last) + return wt.to(conv.weight.device) + + +class FusedUpsample2xConv2d(nn.Module): + """ConvTranspose2d(k4, s2, p1) equivalent of diffusers Upsample2D + (nearest-2x interpolate + Conv2d(3x3, p1)). + + nearest 2x upsampling is pure pixel replication (no arithmetic), so the + fusion only re-associates the conv taps; the kernel is summed lazily in + fp32 on first fast-path use and written back to the conv dtype. With the + fast-path gate disabled the original Upsample2D runs bit-for-bit. + """ + + def __init__(self, upsample: nn.Module, gate: VaeFastPathGate) -> None: + super().__init__() + # Keep ``conv`` registered directly on the wrapper so parameter names + # remain ``...upsamplers.N.conv.*``. The unregistered original module + # is retained only to run its exact lossless forward implementation. + object.__setattr__(self, "_orig", upsample) + self.conv = upsample.conv + self.channels = upsample.channels + self._sgl_gate = gate + self._fused_weight = None + + def forward(self, hidden_states, output_size=None, *args, **kwargs): + if ( + not self._sgl_gate.enabled + or output_size is not None + or hidden_states.shape[1] != self.channels + ): + return self._orig(hidden_states, output_size=output_size) + conv = self.conv + w = self._fused_weight + if w is None: + w = _fold_upsample2x_conv2d_weight(conv) + self._fused_weight = w + return F.conv_transpose2d(hidden_states, w, conv.bias, stride=2, padding=1) + + +def _install_fused_upsample(decoder, upsample_cls, gate: VaeFastPathGate) -> int: + count = 0 + for blk in decoder.up_blocks: + upsamplers = getattr(blk, "upsamplers", None) + if not upsamplers: + continue + for i, up in enumerate(upsamplers): + if type(up) is not upsample_cls: + continue + conv = getattr(up, "conv", None) + if ( + up.use_conv + and not up.use_conv_transpose + and up.interpolate + and up.norm is None + and up.name == "conv" + and type(conv) is nn.Conv2d + and conv.kernel_size == (3, 3) + and conv.stride == (1, 1) + and conv.padding == (1, 1) + ): + upsamplers[i] = FusedUpsample2xConv2d(up, gate) + count += 1 + return count + + +# --------------------------------------------------------------------------- +# Fuse C: layout-safe single-head attention forward with the V/proj fold +# --------------------------------------------------------------------------- + + +def _fold_attn_vproj(m) -> tuple[torch.Tensor, torch.Tensor]: + w_v = m.to_v.weight.detach().float() + b_v = m.to_v.bias.detach().float() + w_o = m.to_out[0].weight.detach().float() + b_o = m.to_out[0].bias.detach().float() + dtype = m.to_v.weight.dtype + # Softmax rows sum to 1, so the folded bias broadcasts exactly like the + # original output-projection bias. + return (w_o @ w_v).to(dtype).contiguous(), (w_o @ b_v + b_o).to(dtype) + + +def _attn_fast_forward( + self, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + temb=None, + **cross_attention_kwargs, +): + if ( + not self._sgl_gate.enabled + or encoder_hidden_states is not None + or attention_mask is not None + or temb is not None + or hidden_states.ndim != 4 + ): + # Lossless requests (and anything unexpected) run the stock + # diffusers path; the decoder layout is NCHW in that case. + return type(self).forward( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + temb=temb, + **cross_attention_kwargs, + ) + + residual = hidden_states + batch_size, channels, height, width = hidden_states.shape + # Free view: with the gate enabled the decoder runs in channels_last. + hs = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels) + + if self.group_norm is not None: + gn = self.group_norm + y = group_norm_silu_rows( + hs, gn.weight, gn.bias, gn.num_groups, gn.eps, apply_silu=False + ) + hs = y if y is not None else gn(hs.transpose(1, 2)).transpose(1, 2) + + query = self.to_q(hs) + key = self.to_k(hs) + folded = self._sgl_folded_v + if folded is None: + folded = _fold_attn_vproj(self) + self._sgl_folded_v = folded + value = F.linear(hs, folded[0], folded[1]) + + out = F.scaled_dot_product_attention( + query.unsqueeze(1), key.unsqueeze(1), value.unsqueeze(1) + ) + out = out.squeeze(1).to(query.dtype) + out = self.to_out[1](out) # dropout (identity in eval) + + # The permuted view has channels_last strides, matching the residual. + out = out.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2) + if self.residual_connection: + out = out + residual + return out / self.rescale_output_factor + + +def _attn_fast_compatible(m, attn_cls, processor_cls) -> bool: + return ( + type(m) is attn_cls + and isinstance(m.processor, processor_cls) + and m.heads == 1 + and m.scale_qk + and m.spatial_norm is None + and not m.norm_cross + and getattr(m, "norm_q", None) is None + and getattr(m, "norm_k", None) is None + and getattr(m, "add_k_proj", None) is None + and type(m.to_q) is nn.Linear + and type(m.to_k) is nn.Linear + and type(m.to_v) is nn.Linear + and type(m.to_out[0]) is nn.Linear + and isinstance(m.to_out[1], nn.Dropout) + # The V/proj fold needs both biases present. + and m.to_v.bias is not None + and m.to_out[0].bias is not None + ) + + +# --------------------------------------------------------------------------- +# channels_last layout dispatch (swapped at decode entry to match the gate) +# --------------------------------------------------------------------------- + + +def _decoder_layout_forward(self, *args, **kwargs): + want_cl = self._sgl_gate.enabled + if want_cl != self._sgl_channels_last: + # Layout swaps are pure permutations of the parameter memory (values + # are bit-identical), so flipping back to contiguous restores the + # baseline NCHW cuDNN kernel selection exactly. + self.to( + memory_format=(torch.channels_last if want_cl else torch.contiguous_format) + ) + self._sgl_channels_last = want_cl + logger.info( + "FLUX.2 VAE: decoder switched to %s layout.", + "channels_last (NHWC)" if want_cl else "contiguous (NCHW)", + ) + return type(self).forward(self, *args, **kwargs) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def maybe_optimize_flux2_vae(vae: nn.Module) -> nn.Module: + """Install the quality-gated CUDA FLUX.2 VAE decoder fast paths.""" + 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." + ) + return vae + if not _HAS_TRITON: + # aten GroupNorm is ~2x slower on NHWC tensors than on NCHW, so the + # channels_last fast path is only a net win with the Triton + # 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." + ) + return vae + + decoder = vae.decoder + attn_modules = [ + m + for m in decoder.modules() + if _attn_fast_compatible(m, Attention, AttnProcessor2_0) + ] + n_attn_total = sum(1 for m in decoder.modules() if isinstance(m, Attention)) + if len(attn_modules) != n_attn_total: + # AttnProcessor2_0 `.view`s the 4D activation, which is illegal on + # 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; " + "skipping CUDA decoder fast paths.", + n_attn_total - len(attn_modules), + n_attn_total, + ) + return vae + + gate = VaeFastPathGate() + decoder._sgl_gate = gate + decoder._sgl_channels_last = False + decoder.forward = MethodType(_decoder_layout_forward, decoder) + n_up = _install_fused_upsample(decoder, Upsample2D, gate) + for m in attn_modules: + m._sgl_gate = gate + m._sgl_folded_v = None + m.forward = MethodType(_attn_fast_forward, m) + n_norm = _install_norm_silu(decoder, ResnetBlock2D, gate) + setattr(vae, GATE_ATTR, gate) + logger.info( + "FLUX.2 VAE: installed quality-gated decoder fast paths " + "(channels_last dispatch, %d fused upsamplers, %d fast attention " + "blocks, %d GroupNorm+SiLU fusions).", + n_up, + len(attn_modules), + n_norm, + ) + return vae diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py index 793381272..782fa4494 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py @@ -297,34 +297,45 @@ class DecodingStage(PipelineStage): assert vae is not None self.vae = vae - frames = self.decode(batch.latents, server_args, vae_dtype=vae_dtype) + # Request-scoped VAE fast-path gate (see flux2_vae_cuda_opt): + # quality == "high" opts this decode into the near-lossless fast + # paths; the "lossless" default keeps the bit-exact original + # module path. VAEs without installed wrappers have no gate. + gate = getattr(vae, "_sgl_vae_fast_path_gate", None) + if gate is not None: + gate.enabled = getattr(batch.sampling_params, "quality", None) == "high" + try: + frames = self.decode(batch.latents, server_args, vae_dtype=vae_dtype) - # decode trajectory latents if needed - if batch.return_trajectory_decoded: - assert ( - batch.trajectory_latents is not None - ), "batch should have trajectory latents" + # decode trajectory latents if needed + if batch.return_trajectory_decoded: + assert ( + batch.trajectory_latents is not None + ), "batch should have trajectory latents" - # 1. Batch trajectory decoding to improve GPU utilization - # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] - B, T, C, F, H, W = batch.trajectory_latents.shape - flat_latents = batch.trajectory_latents.view(B * T, C, F, H, W) + # 1. Batch trajectory decoding to improve GPU utilization + # batch.trajectory_latents is [batch_size, timesteps, channels, frames, height, width] + B, T, C, F, H, W = batch.trajectory_latents.shape + flat_latents = batch.trajectory_latents.view(B * T, C, F, H, W) - logger.info("decoding %s trajectory latents in batch", B * T) - # Use the optimized batch decode - all_decoded = self.decode( - flat_latents, server_args, vae_dtype=vae_dtype - ) + logger.info("decoding %s trajectory latents in batch", B * T) + # Use the optimized batch decode + all_decoded = self.decode( + flat_latents, server_args, vae_dtype=vae_dtype + ) - # 2. Reshape back - # Keep on GPU to allow faster vectorized post-processing - decoded_tensor = all_decoded.view(B, T, *all_decoded.shape[1:]) + # 2. Reshape back + # Keep on GPU to allow faster vectorized post-processing + decoded_tensor = all_decoded.view(B, T, *all_decoded.shape[1:]) - # Convert to list of tensors (per timestep) as expected by OutputBatch - # Each element in list is [B, channels, frames, H_out, W_out] - trajectory_decoded = [decoded_tensor[:, i] for i in range(T)] - else: - trajectory_decoded = None + # Convert to list of tensors (per timestep) as expected by OutputBatch + # Each element in list is [B, channels, frames, H_out, W_out] + trajectory_decoded = [decoded_tensor[:, i] for i in range(T)] + else: + trajectory_decoded = None + finally: + if gate is not None: + gate.enabled = False frames = server_args.pipeline_config.post_decoding(frames, server_args) diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index b32dfe30a..6fb73bb25 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -537,6 +537,28 @@ class CudaPlatformBase(Platform): def get_device_communicator_cls(cls) -> str: return "sglang.multimodal_gen.runtime.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa + @classmethod + def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module: + """Install the quality-gated FLUX.2 VAE decoder fast paths. + + Requests with quality == "high" run the fast paths; the "lossless" + default runs the original module path bit-for-bit. See + flux2_vae_cuda_opt for details. + """ + try: + from sglang.multimodal_gen.runtime.models.vaes.flux2_vae_cuda_opt import ( + maybe_optimize_flux2_vae, + ) + + vae = maybe_optimize_flux2_vae(vae) + except Exception: + logger.warning( + "Failed to apply CUDA FLUX.2 VAE optimizations; " + "using the unmodified VAE.", + exc_info=True, + ) + return vae + # NVML utils # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`, diff --git a/test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py b/test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py new file mode 100644 index 000000000..ba796ce0e --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py @@ -0,0 +1,59 @@ +"""Focused correctness checks for the FLUX.2 VAE CUDA fast path.""" + +import sys + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.upsampling import Upsample2D + +from sglang.kernels.ops.diffusion.triton import group_norm_silu_twopass as gn_kernel +from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt +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") + + +@torch.no_grad() +def test_flux2_vae_fastpath(): + torch.manual_seed(0) + gate = vae_opt.VaeFastPathGate() + gn = nn.GroupNorm(32, 128, eps=1e-6).to("cuda", torch.bfloat16) + x = torch.randn(1, 128, 64, 64, device="cuda", dtype=torch.bfloat16).to( + memory_format=torch.channels_last + ) + ref = F.silu(gn(x)) + fused_gn = vae_opt.FusedGroupNormSiLU(gn, gate) + assert set(fused_gn.state_dict()) == {"weight", "bias"} + assert torch.equal(fused_gn(x), ref) + assert ( + gn_kernel.group_norm_silu_4d(x.contiguous(), gn.weight, gn.bias, 32, 1e-6) + is None + ) + + gate.enabled = True + fast = fused_gn(x) + assert fast.is_contiguous(memory_format=torch.channels_last) + torch.testing.assert_close(fast.float(), ref.float(), atol=0.06, rtol=0) + + gate.enabled = False + up = Upsample2D(channels=32, use_conv=True).to("cuda", torch.bfloat16) + fused_up = vae_opt.FusedUpsample2xConv2d(up, gate) + assert set(fused_up.state_dict()) == {"conv.weight", "conv.bias"} + x = torch.randn(2, 32, 33, 29, device="cuda", dtype=torch.bfloat16) + ref = up(x) + assert torch.equal(fused_up(x), ref) + assert fused_up._fused_weight is None + + gate.enabled = True + fast = fused_up(x) + assert fused_up._fused_weight is not None + ref_range = ref.float().max() - ref.float().min() + relative_mse = F.mse_loss(fast.float(), ref.float()) / ref_range.square() + assert relative_mse < 3.2e-5 + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))