diff --git a/python/sglang/kernels/ops/diffusion/activation/sana_conv_post_triton.py b/python/sglang/kernels/ops/diffusion/activation/sana_conv_post_triton.py index 69d4f4f5a..f4b1252c2 100644 --- a/python/sglang/kernels/ops/diffusion/activation/sana_conv_post_triton.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py index cad7a50c4..ad097b7c0 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_components.py +++ b/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 diff --git a/test/registered/kernel/diffusion/test_sana_wm_conv_post.py b/test/registered/kernel/diffusion/test_sana_wm_conv_post.py new file mode 100644 index 000000000..1f33b8625 --- /dev/null +++ b/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() diff --git a/test/registered/kernel/diffusion/test_sana_wm_reverse_scan.py b/test/registered/kernel/diffusion/test_sana_wm_reverse_scan.py new file mode 100644 index 000000000..476e2949b --- /dev/null +++ b/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()