[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
+    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,20 +41,30 @@ 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
+    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)
+    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.
@@ -51,19 +72,19 @@ def _bias_glu_kernel(
     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 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(
+def _gdn_scan_backward_reference(q, k, v, q_rot, k_rot, beta, decay, eps):
-    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,
-    )
-
     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
+    out_bwd = _sana_wm_reverse_scan(
-    T = beta.shape[2]
+        q_rot,
-    S = N // T
+        k_rot,
-
+        v,
-    def to_time(x):
+        beta,
-        return x.view(B, H, D, T, S).permute(0, 1, 3, 2, 4)
+        decay,
-
+        reference=lambda: _single_path_delta_scan_backward_reference(
-    def from_time(x):
+            q_rot, k_rot, v, beta, decay
-        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_fwd + out_bwd, state_kv
 

test/registered/kernel/diffusion/test_sana_wm_conv_post.py

@@ -0,0 +1,141 @@
+"""SANA-WM conv post-processing parity, rounding, and graph replay."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+import torch.nn.functional as F
+
+import sglang.multimodal_gen.runtime.models.dits.sana_wm_components as wm
+from sglang.kernels.ops.diffusion import BitExactFusionGate
+from sglang.kernels.ops.diffusion.activation.sana_conv_post_triton import (
+    can_use_fused_bias_glu,
+    can_use_fused_bias_silu,
+    fused_bias_glu,
+    fused_bias_silu,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestSanaWMConvPost(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA required")
+        self.original_gate = wm._SANA_WM_CONV_POST
+        wm._SANA_WM_CONV_POST = BitExactFusionGate("test", per_signature=True)
+        self.addCleanup(setattr, wm, "_SANA_WM_CONV_POST", self.original_gate)
+        torch.manual_seed(42)
+
+    @torch.inference_mode()
+    def test_post_kernels_both_layouts_and_bias_modes(self):
+        for shape in [(3, 34, 17, 23), (2, 13440, 22, 40), (2, 34, 1, 1)]:
+            for layout in [torch.contiguous_format, torch.channels_last]:
+                x = torch.randn(shape, device="cuda", dtype=torch.bfloat16).to(
+                    memory_format=layout
+                )
+                bias = torch.randn(shape[1], device="cuda", dtype=x.dtype)
+                biased = x + bias[None, :, None, None]
+                self.assertTrue(torch.equal(fused_bias_silu(x, bias), F.silu(biased)))
+                for b, z in [(bias, biased), (None, x)]:
+                    a, g = z.chunk(2, dim=1)
+                    actual = fused_bias_glu(x, b)
+                    self.assertTrue(torch.equal(actual, a * F.silu(g)))
+                    self.assertTrue(actual.is_contiguous(memory_format=layout))
+        # Use a real spatial slice; width-one tensors remain contiguous.
+        sliced = torch.empty(2, 34, 7, 10, device="cuda", dtype=x.dtype)[:, :, :, ::2]
+        self.assertFalse(can_use_fused_bias_silu(sliced, bias))
+        self.assertFalse(can_use_fused_bias_glu(sliced, None))
+        self.assertFalse(can_use_fused_bias_glu(x.float(), None))
+
+    @staticmethod
+    def reference(module, x):
+        z = module.depth_conv(F.silu(module.inverted_conv.conv(x)))
+        a, g = z.chunk(2, dim=1)
+        return a * F.silu(g)
+
+    @torch.inference_mode()
+    def test_native_convolution_and_depthwise_bias_rounding(self):
+        for channels, hidden, shape in [
+            (32, 96, (3, 32, 17, 23)),
+            (2240, 6720, (14, 2240, 22, 40)),
+        ]:
+            module = wm.GLUMBConvTemp(channels, hidden).cuda().bfloat16()
+            for seed in (0, 1):
+                torch.manual_seed(seed)
+                x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+                self.assertTrue(
+                    torch.equal(module._spatial_glu(x), self.reference(module, x))
+                )
+                self.assertTrue(wm._SANA_WM_CONV_POST.verified)
+                self.assertFalse(wm._SANA_WM_CONV_POST.disabled)
+            # Guard the numerical distinction which forbids bias extraction.
+            z = module.inverted_conv(x)
+            conv = module.depth_conv.conv
+            split = F.conv2d(z, conv.weight, None, padding=1, groups=conv.groups)
+            split = split + conv.bias[None, :, None, None]
+            self.assertFalse(torch.equal(conv(z), split))
+
+    @torch.inference_mode()
+    def test_graph_replay_updates_inputs_and_streaming_tail(self):
+        module = wm.GLUMBConvTemp(32, 96).cuda().bfloat16()
+        # Exercise a nonzero temporal convolution instead of its zero init.
+        module.t_conv.weight.normal_(std=0.01)
+        x = torch.randn(2, 3 * 7 * 11, 32, device="cuda", dtype=torch.bfloat16)
+        tail = torch.randn(2, 32, 1, 77, device="cuda", dtype=x.dtype)
+        module(x, (3, 7, 11), ffn_tail=tail, save_ffn_tail=True)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual, actual_tail = module(
+                x, (3, 7, 11), ffn_tail=tail, save_ffn_tail=True
+            )
+        x.add_(0.25)
+        tail.mul_(0.5)
+        graph.replay()
+        with patch.object(wm._SANA_WM_CONV_POST, "disabled", True):
+            expected, expected_tail = module(
+                x, (3, 7, 11), ffn_tail=tail, save_ffn_tail=True
+            )
+        self.assertTrue(torch.equal(actual, expected))
+        self.assertTrue(torch.equal(actual_tail, expected_tail))
+
+    @torch.inference_mode()
+    def test_unverified_capture_and_mismatch_use_reference(self):
+        module = wm.GLUMBConvTemp(32, 96).cuda().bfloat16()
+        x = torch.randn(3, 32, 17, 23, device="cuda", dtype=torch.bfloat16)
+        expected = self.reference(module, x)
+        with patch.object(wm.diffusion_ops, "fused_bias_silu") as fused:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                actual = module._spatial_glu(x)
+            graph.replay()
+            fused.assert_not_called()
+        self.assertTrue(torch.equal(actual, expected))
+        with patch.object(
+            wm.diffusion_ops, "fused_bias_glu", return_value=torch.zeros_like(expected)
+        ):
+            actual = module._spatial_glu(x)
+        self.assertTrue(torch.equal(actual, expected))
+        self.assertTrue(wm._SANA_WM_CONV_POST.disabled)
+        with patch.object(wm.diffusion_ops, "fused_bias_silu") as fused:
+            actual = module._spatial_glu(x)
+            fused.assert_not_called()
+        self.assertTrue(torch.equal(actual, expected))
+
+    def test_grad_enabled_uses_differentiable_reference(self):
+        module = wm.GLUMBConvTemp(32, 96).cuda().bfloat16()
+        x = torch.randn(
+            2, 32, 7, 11, device="cuda", dtype=torch.bfloat16, requires_grad=True
+        )
+        with patch.object(wm.diffusion_ops, "fused_bias_silu") as fused:
+            actual = module._spatial_glu(x)
+            actual.float().sum().backward()
+            fused.assert_not_called()
+        self.assertIsNotNone(x.grad)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernel/diffusion/test_sana_wm_reverse_scan.py

@@ -0,0 +1,160 @@
+"""Exact streaming GDN output/state parity without materialized reversed video."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+import sglang.multimodal_gen.runtime.models.dits.sana_wm_components as wm
+from sglang.kernels.ops.diffusion import BitExactFusionGate
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestSanaWMReverseScan(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA required")
+        original = wm._SANA_WM_GDN_REVERSE
+        wm._SANA_WM_GDN_REVERSE = BitExactFusionGate("test", per_signature=True)
+        self.addCleanup(setattr, wm, "_SANA_WM_GDN_REVERSE", original)
+        torch.manual_seed(42)
+
+    def inputs(self, frames, transposed=False, frame_beta=False):
+        b, h, d, s = 1, 20, 112, 880
+        shape = (b, h, frames * s, d) if transposed else (b, h, d, frames * s)
+        xs = [torch.randn(shape, device="cuda") * 0.03 for _ in range(5)]
+        if transposed:
+            xs = [x.transpose(-1, -2) for x in xs]
+        beta_shape = (b, h, frames) if frame_beta else (b, h, frames, s)
+        return (
+            *xs,
+            torch.rand(beta_shape, device="cuda") * 0.02,
+            torch.rand(b, h, frames, device="cuda") * 0.5 + 0.4,
+        )
+
+    def assert_nested_equal(self, a, b):
+        if isinstance(a, tuple):
+            for x, y in zip(a, b, strict=True):
+                self.assert_nested_equal(x, y)
+        else:
+            self.assertTrue(torch.equal(a, b))
+
+    @torch.inference_mode()
+    def test_multiple_chunks_outputs_and_carried_states(self):
+        for transposed in (False, True):
+            for frame_beta in (False, True):
+                reference_state = candidate_state = (None, None)
+                reference_cam = candidate_cam = None
+                for frames in (1, 4, 3, 2):
+                    q, k, v, qr, kr, beta, decay = self.inputs(
+                        frames, transposed, frame_beta
+                    )
+                    with patch.object(wm._SANA_WM_GDN_REVERSE, "disabled", True):
+                        a, reference_state = wm._gdn_scan_cached(
+                            q,
+                            k,
+                            v,
+                            qr,
+                            kr,
+                            beta,
+                            decay,
+                            init_state_kv=reference_state[0],
+                            init_state_z=reference_state[1],
+                        )
+                        c, reference_cam = wm._single_path_delta_scan_cached(
+                            qr, kr, v, beta, decay, init_state_kv=reference_cam
+                        )
+                    b, candidate_state = wm._gdn_scan_cached(
+                        q,
+                        k,
+                        v,
+                        qr,
+                        kr,
+                        beta,
+                        decay,
+                        init_state_kv=candidate_state[0],
+                        init_state_z=candidate_state[1],
+                    )
+                    d, candidate_cam = wm._single_path_delta_scan_cached(
+                        qr, kr, v, beta, decay, init_state_kv=candidate_cam
+                    )
+                    self.assert_nested_equal(
+                        (a, reference_state, c, reference_cam),
+                        (b, candidate_state, d, candidate_cam),
+                    )
+                    self.assertFalse(wm._SANA_WM_GDN_REVERSE.disabled)
+        self.assertTrue(wm._SANA_WM_GDN_REVERSE.verified)
+
+    @torch.inference_mode()
+    def test_changed_graph_inputs_and_initial_states(self):
+        q, k, v, qr, kr, beta, decay = self.inputs(3, True)
+        state = torch.randn(1, 20, 112, 112, device="cuda") * 0.001
+        z = torch.randn(1, 20, 112, 1, device="cuda") * 0.001
+
+        def run():
+            return (
+                wm._gdn_scan_cached(
+                    q, k, v, qr, kr, beta, decay, init_state_kv=state, init_state_z=z
+                ),
+                wm._single_path_delta_scan_cached(
+                    qr, kr, v, beta, decay, init_state_kv=state
+                ),
+            )
+
+        run()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            actual = run()
+        q.add_(0.01)
+        qr.mul_(0.7)
+        beta.mul_(0.8)
+        state.add_(0.001)
+        z.neg_()
+        graph.replay()
+        with patch.object(wm._SANA_WM_GDN_REVERSE, "disabled", True):
+            expected = run()
+        self.assert_nested_equal(actual, expected)
+
+    @torch.inference_mode()
+    def test_unverified_capture_and_mismatch_fallback(self):
+        _, _, v, qr, kr, beta, decay = self.inputs(2)
+
+        def reference():
+            return wm._single_path_delta_scan_backward_reference(qr, kr, v, beta, decay)
+
+        expected = reference()
+        with patch.object(wm, "_sana_wm_reverse_scan_impl") as fast:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                actual = wm._sana_wm_reverse_scan(
+                    qr, kr, v, beta, decay, reference=reference
+                )
+            graph.replay()
+            fast.assert_not_called()
+        self.assertTrue(torch.equal(actual, expected))
+        with patch.object(
+            wm, "_sana_wm_reverse_scan_impl", return_value=torch.ones_like(expected)
+        ):
+            actual = wm._sana_wm_reverse_scan(
+                qr, kr, v, beta, decay, reference=reference
+            )
+        self.assertTrue(wm._SANA_WM_GDN_REVERSE.disabled)
+        self.assertTrue(torch.equal(actual, expected))
+
+    @torch.inference_mode()
+    def test_synthetic_zero_update_keeps_nonfinite_query_behavior(self):
+        q, k, v, qr, kr, beta, decay = self.inputs(1)
+        qr[..., 0] = float("nan")
+        q[..., 1] = float("inf")
+        expected = wm._gdn_scan_backward_reference(q, k, v, qr, kr, beta, decay, 1e-6)
+        actual = wm._sana_wm_reverse_scan_impl(qr, kr, v, beta, decay, q, k)
+        for a, b in zip(actual, expected, strict=True):
+            torch.testing.assert_close(a, b, rtol=0, atol=0, equal_nan=True)
+
+
+if __name__ == "__main__":
+    unittest.main()