[Diffusion] Optimize SANA-WM convolution post-processing and streaming GDN (#38529)
This commit is contained in:
Xiaoyu Zhang
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user