[Diffusion] Optimize SANA-WM convolution post-processing and streaming GDN (#38529)

python/sglang/kernels/ops/diffusion/activation/sana_conv_post_triton.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Bit-exact post-processing kernels for Sana's channels-last GLUMB convs."""
+"""Bit-exact post-processing kernels for Sana's GLUMB convs."""
 
 from __future__ import annotations
 
@@ -11,10 +11,21 @@ from sglang.kernels.ops.diffusion.common.numerics import round_bf16_to_fp32
 
 
 @triton.jit
-def _bias_silu_kernel(out_ptr, x_ptr, bias_ptr, numel, channels: tl.constexpr):
+def _bias_silu_kernel(
+    out_ptr,
+    x_ptr,
+    bias_ptr,
+    numel,
+    channels: tl.constexpr,
+    spatial: tl.constexpr,
+    channels_last: tl.constexpr,
+):
     offsets = tl.program_id(0).to(tl.int64) * 1024 + tl.arange(0, 1024)
     mask = offsets < numel
-    channel = offsets % channels
+    if channels_last:
+        channel = offsets % channels
+    else:
+        channel = (offsets // spatial) % channels
     x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
     bias = tl.load(bias_ptr + channel, mask=mask, other=0.0).to(tl.float32)
     # nn.Conv2d applies its bf16 bias before nn.SiLU, so preserve the
@@ -30,40 +41,50 @@ def _bias_glu_kernel(
     bias_ptr,
     out_numel,
     channels: tl.constexpr,
+    spatial: tl.constexpr,
+    channels_last: tl.constexpr,
+    has_bias: tl.constexpr,
 ):
     offsets = tl.program_id(0).to(tl.int64) * 1024 + tl.arange(0, 1024)
     mask = offsets < out_numel
-    channel = offsets % channels
-    pixel = offsets // channels
-    in_base = pixel * (2 * channels) + channel
+    if channels_last:
+        channel = offsets % channels
+        pixel = offsets // channels
+        in_base = pixel * (2 * channels) + channel
+        gate_offset = channels
+    else:
+        channel = (offsets // spatial) % channels
+        batch = offsets // (channels * spatial)
+        in_base = batch * (2 * channels * spatial) + offsets % (channels * spatial)
+        gate_offset = channels * spatial
 
     hidden = tl.load(x_ptr + in_base, mask=mask, other=0.0).to(tl.float32)
-    gate = tl.load(x_ptr + in_base + channels, mask=mask, other=0.0).to(tl.float32)
-    hidden_bias = tl.load(bias_ptr + channel, mask=mask, other=0.0).to(tl.float32)
-    gate_bias = tl.load(bias_ptr + channels + channel, mask=mask, other=0.0).to(
-        tl.float32
-    )
-
-    hidden = round_bf16_to_fp32(hidden + hidden_bias)
-    gate = round_bf16_to_fp32(gate + gate_bias)
+    gate = tl.load(x_ptr + in_base + gate_offset, mask=mask, other=0.0).to(tl.float32)
+    if has_bias:
+        hidden_bias = tl.load(bias_ptr + channel, mask=mask, other=0.0).to(tl.float32)
+        gate_bias = tl.load(bias_ptr + channels + channel, mask=mask, other=0.0).to(
+            tl.float32
+        )
+        hidden = round_bf16_to_fp32(hidden + hidden_bias)
+        gate = round_bf16_to_fp32(gate + gate_bias)
     # SiLU materializes a bf16 tensor before the following multiply in eager.
     gate = round_bf16_to_fp32(gate * tl.sigmoid(gate))
     tl.store(out_ptr + offsets, hidden * gate, mask=mask)
 
 
-def _is_channels_last_bf16(x: torch.Tensor) -> bool:
+def _is_dense_bf16(x: torch.Tensor) -> bool:
     return (
         x.is_cuda
         and x.dtype is torch.bfloat16
         and x.dim() == 4
         and x.numel() > 0
-        and x.is_contiguous(memory_format=torch.channels_last)
+        and (x.is_contiguous() or x.is_contiguous(memory_format=torch.channels_last))
     )
 
 
 def can_use_fused_bias_silu(x: torch.Tensor, bias: torch.Tensor) -> bool:
     return (
-        _is_channels_last_bf16(x)
+        _is_dense_bf16(x)
         and bias.is_cuda
         and bias.dtype is x.dtype
         and bias.device == x.device
@@ -79,38 +100,62 @@ def fused_bias_silu(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
     out = torch.empty_like(x, memory_format=torch.preserve_format)
     with torch.cuda.device(x.device):
         _bias_silu_kernel[(triton.cdiv(x.numel(), 1024),)](
-            out, x, bias, x.numel(), channels=x.shape[1]
+            out,
+            x,
+            bias,
+            x.numel(),
+            channels=x.shape[1],
+            spatial=x.shape[2] * x.shape[3],
+            channels_last=x.is_contiguous(memory_format=torch.channels_last),
         )
     return out
 
 
-def can_use_fused_bias_glu(x: torch.Tensor, bias: torch.Tensor) -> bool:
+def can_use_fused_bias_glu(x: torch.Tensor, bias: torch.Tensor | None) -> bool:
     return (
-        _is_channels_last_bf16(x)
+        _is_dense_bf16(x)
         and x.shape[1] % 2 == 0
-        and bias.is_cuda
-        and bias.dtype is x.dtype
-        and bias.device == x.device
-        and bias.dim() == 1
-        and bias.shape[0] == x.shape[1]
-        and bias.is_contiguous()
+        and (
+            bias is None
+            or (
+                bias.is_cuda
+                and bias.dtype is x.dtype
+                and bias.device == x.device
+                and bias.dim() == 1
+                and bias.shape[0] == x.shape[1]
+                and bias.is_contiguous()
+            )
+        )
     )
 
 
-def fused_bias_glu(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+def fused_bias_glu(x: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor:
+    """Apply optional bias then ``hidden * silu(gate)`` along the channel axis.
+
+    Pass no bias for an already biased native depthwise-convolution output:
+    splitting that convolution's bias can change its accumulation rounding.
+    """
     if not can_use_fused_bias_glu(x, bias):
         raise RuntimeError("unsupported input for Sana fused bias-GLU")
     batch, double_channels, height, width = x.shape
     channels = double_channels // 2
+    channels_last = x.is_contiguous(memory_format=torch.channels_last)
     out = torch.empty(
         (batch, channels, height, width),
         dtype=x.dtype,
         device=x.device,
-        memory_format=torch.channels_last,
+        memory_format=torch.channels_last if channels_last else torch.contiguous_format,
     )
     with torch.cuda.device(x.device):
         _bias_glu_kernel[(triton.cdiv(out.numel(), 1024),)](
-            out, x, bias, out.numel(), channels=channels
+            out,
+            x,
+            bias,
+            out.numel(),
+            channels=channels,
+            spatial=height * width,
+            channels_last=channels_last,
+            has_bias=bias is not None,
         )
     return out
 

python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py

@@ -13,11 +13,21 @@ import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.embeddings import get_1d_rotary_pos_embed
 
+from sglang.kernels.ops import diffusion as diffusion_ops
+from sglang.kernels.ops.diffusion import BitExactFusionGate, tensors_equal
 from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
 
+_SANA_WM_CONV_POST = BitExactFusionGate(
+    "SANA-WM conv post-processing", per_signature=True
+)
+
+_SANA_WM_GDN_REVERSE = BitExactFusionGate(
+    "SANA-WM reverse GDN scan", per_signature=True
+)
+
 _SANA_WM_TRITON_GDN_DISABLED_REASON: Optional[str] = None
 _SANA_WM_TRITON_GDN_FALLBACK_LOGGED = False
 _SANA_WM_TRITON_CAM_GDN_DISABLED_REASON: Optional[str] = None
@@ -1032,11 +1042,67 @@ class GLUMBConvTemp(nn.Module):
         )
         nn.init.zeros_(self.t_conv.weight)
 
-    def _apply_spatial(self, x: torch.Tensor) -> torch.Tensor:
+    def _spatial_glu_reference(self, x: torch.Tensor) -> torch.Tensor:
         x = self.inverted_conv(x)
         x = self.depth_conv(x)
         a, g = x.chunk(2, dim=1)
-        return self.point_conv(a * self.glu_act(g))
+        return a * self.glu_act(g)
+
+    def _spatial_glu(self, x: torch.Tensor) -> torch.Tensor:
+        if (
+            not _SANA_WM_CONV_POST.disabled
+            and x.is_cuda
+            and x.dtype is torch.bfloat16
+            and x.is_contiguous()
+            and x.numel() > 0
+            and not torch.is_grad_enabled()
+            and not torch.compiler.is_compiling()
+        ):
+            conv = self.inverted_conv.conv
+            sig = (
+                x.shape,
+                x.stride(),
+                x.device,
+                x.dtype,
+                conv.weight.shape,
+                conv.weight.stride(),
+                self.depth_conv.conv.weight.stride(),
+                torch.backends.cudnn.enabled,
+                torch.backends.cudnn.benchmark,
+                torch.backends.cudnn.deterministic,
+                torch.backends.cudnn.allow_tf32,
+            )
+            verified = _SANA_WM_CONV_POST.is_verified(sig)
+            if verified or not torch.cuda.is_current_stream_capturing():
+                try:
+                    raw = F.conv2d(
+                        x,
+                        conv.weight,
+                        None,
+                        conv.stride,
+                        conv.padding,
+                        conv.dilation,
+                        conv.groups,
+                    )
+                    hidden = diffusion_ops.fused_bias_silu(raw, conv.bias)
+                    # Native depthwise conv accumulates bias before rounding;
+                    # preserve it and fuse only its following SiLU/multiply.
+                    out = diffusion_ops.fused_bias_glu(self.depth_conv(hidden), None)
+                except Exception as exc:
+                    _SANA_WM_CONV_POST.on_exception(exc, logger=logger)
+                else:
+                    if verified:
+                        return out
+                    return _SANA_WM_CONV_POST.accept_or_fallback(
+                        out,
+                        self._spatial_glu_reference(x),
+                        sig=sig,
+                        logger=logger,
+                    )
+        return self._spatial_glu_reference(x)
+
+    def _apply_spatial(self, x: torch.Tensor) -> torch.Tensor:
+        return self.point_conv(self._spatial_glu(x))
 
     def _apply_spatial_autochunked(self, x: torch.Tensor) -> torch.Tensor:
         """Avoid oversized Conv2d calls on long videos while keeping short path fused."""
@@ -1698,40 +1764,7 @@ def _single_path_delta_scan_bidirectional(
 # ---------------------------------------------------------------------------
 
 
-def _gdn_scan_cached(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    q_rot: torch.Tensor,
-    k_rot: torch.Tensor,
-    beta: torch.Tensor,
-    decay: torch.Tensor,
-    *,
-    init_state_kv: Optional[torch.Tensor] = None,
-    init_state_z: Optional[torch.Tensor] = None,
-    eps: float = 1e-6,
-) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
-    """Chunk-causal main-branch GDN scan for streaming `forward_long`.
-
-    The forward pass seeds/returns ``(state_kv, state_z)`` so chunks stay
-    continuous; the backward pass is intra-chunk and stateless. Returns
-    ``(out, (state_kv, state_z))``.
-    """
-    (num_fwd, den_fwd), (state_kv, state_z) = _gdn_scan_forward_stateful(
-        q,
-        k,
-        v,
-        q_rot,
-        k_rot,
-        beta,
-        decay,
-        init_state_kv=init_state_kv,
-        init_state_z=init_state_z,
-        eps=eps,
-        return_components=True,
-        return_state=True,
-    )
-
+def _gdn_scan_backward_reference(q, k, v, q_rot, k_rot, beta, decay, eps):
     B, H, D, N = q.shape
     T = beta.shape[2]
     S = N // T
@@ -1767,6 +1800,167 @@ def _gdn_scan_cached(
 
     num_bwd = flip_back(num_bwd_flipped, D)
     den_bwd = flip_back(den_bwd_flipped, 1)
+    return num_bwd, den_bwd
+
+
+def _single_path_delta_scan_backward_reference(q_rot, k_rot, v, beta, decay):
+    B, H, D, N = q_rot.shape
+    T = beta.shape[2]
+    S = N // T
+
+    def to_time(x):
+        return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4)
+
+    def from_time(x):
+        return x.permute(0, 1, 3, 2, 4).reshape(B, H, D, N)
+
+    q_rot_bwd = from_time(torch.flip(to_time(q_rot), dims=[2]))
+    k_rot_bwd = from_time(_flip_and_shift(to_time(k_rot), dim=2, shift_val=0.0))
+    v_bwd = from_time(_flip_and_shift(to_time(v), dim=2, shift_val=0.0))
+    beta_bwd = _flip_and_shift(beta, dim=2, shift_val=0.0)
+    decay_bwd = _flip_and_shift(decay, dim=2, shift_val=1.0)
+
+    out_bwd_flipped = _single_path_delta_scan_forward(
+        q_rot_bwd,
+        k_rot_bwd,
+        v_bwd,
+        beta_bwd,
+        decay_bwd,
+    )
+    out_bwd = torch.flip(out_bwd_flipped.view(B, H, D, T, S), dims=[3]).reshape(
+        B, H, D, N
+    )
+    return out_bwd
+
+
+def _sana_wm_reverse_scan_impl(q_rot, k_rot, v, beta, decay, q=None, k=None):
+    """Traverse the exclusive backward recurrence without flipped full videos.
+
+    The old flip/shift cat followed by reshape materializes contiguous K/V.
+    Keep that layout so cuBLAS sees the same leading dimensions. Queries can
+    be selected directly from their original layout. The synthetic first
+    update has zero K/V/beta and unit decay, leaving the zero state unchanged.
+    Its query matmuls are retained, including their nonfinite-input behavior.
+    """
+    B, H, D, N = q_rot.shape
+    T = beta.shape[2]
+    S = N // T
+    k_rot = k_rot.contiguous().view(B, H, D, T, S)
+    v = v.contiguous().view(B, H, D, T, S)
+    q_rot = q_rot.view(B, H, D, T, S)
+    main_branch = q is not None
+    if main_branch:
+        q = q.view(B, H, D, T, S)
+        k = k.contiguous().view(B, H, D, T, S)
+    beta = beta.unsqueeze(3) if beta.ndim == 4 else beta.view(B, H, T, 1, 1)
+    decay = decay.view(B, H, T, 1, 1)
+    state_kv = torch.zeros(B, H, D, D, device=q_rot.device, dtype=q_rot.dtype)
+    state_z = (
+        torch.zeros(B, H, D, 1, device=q_rot.device, dtype=q_rot.dtype)
+        if main_branch
+        else None
+    )
+    nums, dens = [None] * T, [None] * T
+    for i in range(T - 1, -1, -1):
+        if i + 1 < T:
+            j = i + 1
+            kt, vt = k_rot[:, :, :, j], v[:, :, :, j]
+            bt, gt = beta[:, :, j], decay[:, :, j]
+            state_kv = state_kv * gt
+            if main_branch:
+                state_z = state_z * gt
+            delta_v = (vt - torch.matmul(state_kv, kt)) * bt
+            state_kv = state_kv + torch.matmul(delta_v, kt.transpose(-1, -2))
+            if main_branch:
+                key = k[:, :, :, j]
+                delta_z = (1.0 - torch.matmul(state_z.transpose(-1, -2), key)) * bt
+                state_z = state_z + torch.matmul(key, delta_z.transpose(-1, -2))
+        nums[i] = torch.matmul(state_kv, q_rot[:, :, :, i])
+        if main_branch:
+            dens[i] = torch.matmul(state_z.transpose(-1, -2), q[:, :, :, i])
+    num = torch.stack(nums, dim=2).permute(0, 1, 3, 2, 4).reshape(B, H, D, N)
+    if not main_branch:
+        return num
+    den = torch.stack(dens, dim=2).permute(0, 1, 3, 2, 4).reshape(B, H, 1, N)
+    return num, den
+
+
+def _sana_wm_reverse_scan(q_rot, k_rot, v, beta, decay, *, reference, q=None, k=None):
+    inputs = (q_rot, k_rot, v, beta, decay) + (() if q is None else (q, k))
+    if (
+        not _SANA_WM_GDN_REVERSE.disabled
+        and not torch.is_grad_enabled()
+        and not torch.compiler.is_compiling()
+        and all(
+            x.is_cuda and x.dtype is torch.float32 and x.numel() > 0 for x in inputs
+        )
+    ):
+        sig = (
+            q is not None,
+            tuple((x.shape, x.stride(), x.device, x.dtype) for x in inputs),
+            torch.get_float32_matmul_precision(),
+        )
+        verified = _SANA_WM_GDN_REVERSE.is_verified(sig)
+        if verified or not torch.cuda.is_current_stream_capturing():
+            try:
+                out = _sana_wm_reverse_scan_impl(q_rot, k_rot, v, beta, decay, q, k)
+            except Exception as exc:
+                _SANA_WM_GDN_REVERSE.on_exception(exc, logger=logger)
+            else:
+                if verified:
+                    return out
+                return _SANA_WM_GDN_REVERSE.accept_or_fallback(
+                    out, reference(), sig=sig, equal=tensors_equal, logger=logger
+                )
+    return reference()
+
+
+def _gdn_scan_cached(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_rot: torch.Tensor,
+    k_rot: torch.Tensor,
+    beta: torch.Tensor,
+    decay: torch.Tensor,
+    *,
+    init_state_kv: Optional[torch.Tensor] = None,
+    init_state_z: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    """Chunk-causal main-branch GDN scan for streaming `forward_long`.
+
+    The forward pass seeds/returns ``(state_kv, state_z)`` so chunks stay
+    continuous; the backward pass is intra-chunk and stateless. Returns
+    ``(out, (state_kv, state_z))``.
+    """
+    (num_fwd, den_fwd), (state_kv, state_z) = _gdn_scan_forward_stateful(
+        q,
+        k,
+        v,
+        q_rot,
+        k_rot,
+        beta,
+        decay,
+        init_state_kv=init_state_kv,
+        init_state_z=init_state_z,
+        eps=eps,
+        return_components=True,
+        return_state=True,
+    )
+
+    num_bwd, den_bwd = _sana_wm_reverse_scan(
+        q_rot,
+        k_rot,
+        v,
+        beta,
+        decay,
+        q=q,
+        k=k,
+        reference=lambda: _gdn_scan_backward_reference(
+            q, k, v, q_rot, k_rot, beta, decay, eps
+        ),
+    )
     out = (num_fwd + num_bwd) / (den_fwd + den_bwd + eps)
     return out, (state_kv, state_z)
 
@@ -1795,31 +1989,15 @@ def _single_path_delta_scan_cached(
         return_state=True,
     )
 
-    B, H, D, N = q_rot.shape
-    T = beta.shape[2]
-    S = N // T
-
-    def to_time(x):
-        return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4)
-
-    def from_time(x):
-        return x.permute(0, 1, 3, 2, 4).reshape(B, H, D, N)
-
-    q_rot_bwd = from_time(torch.flip(to_time(q_rot), dims=[2]))
-    k_rot_bwd = from_time(_flip_and_shift(to_time(k_rot), dim=2, shift_val=0.0))
-    v_bwd = from_time(_flip_and_shift(to_time(v), dim=2, shift_val=0.0))
-    beta_bwd = _flip_and_shift(beta, dim=2, shift_val=0.0)
-    decay_bwd = _flip_and_shift(decay, dim=2, shift_val=1.0)
-
-    out_bwd_flipped = _single_path_delta_scan_forward(
-        q_rot_bwd,
-        k_rot_bwd,
-        v_bwd,
-        beta_bwd,
-        decay_bwd,
-    )
-    out_bwd = torch.flip(out_bwd_flipped.view(B, H, D, T, S), dims=[3]).reshape(
-        B, H, D, N
+    out_bwd = _sana_wm_reverse_scan(
+        q_rot,
+        k_rot,
+        v,
+        beta,
+        decay,
+        reference=lambda: _single_path_delta_scan_backward_reference(
+            q_rot, k_rot, v, beta, decay
+        ),
     )
     return out_fwd + out_bwd, state_kv