[diffusion][kernel] Accelerate Sana BCG with bit-exact conv post-processing (#34928)

python/sglang/kernels/ops/diffusion/triton/sana_conv_post.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Bit-exact post-processing kernels for Sana's channels-last GLUMB convs."""
+
+from __future__ import annotations
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.kernels.ops.diffusion.triton.numerics import round_bf16_to_fp32
+
+
+@triton.jit
+def _bias_silu_kernel(out_ptr, x_ptr, bias_ptr, numel, channels: tl.constexpr):
+    offsets = tl.program_id(0).to(tl.int64) * 1024 + tl.arange(0, 1024)
+    mask = offsets < numel
+    channel = offsets % 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
+    # intermediate bf16 rounding boundary rather than contracting the chain.
+    biased = round_bf16_to_fp32(x + bias)
+    tl.store(out_ptr + offsets, biased * tl.sigmoid(biased), mask=mask)
+
+
+@triton.jit
+def _bias_glu_kernel(
+    out_ptr,
+    x_ptr,
+    bias_ptr,
+    out_numel,
+    channels: 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
+
+    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)
+    # 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:
+    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)
+    )
+
+
+def can_use_fused_bias_silu(x: torch.Tensor, bias: torch.Tensor) -> bool:
+    return (
+        _is_channels_last_bf16(x)
+        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()
+    )
+
+
+def fused_bias_silu(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    if not can_use_fused_bias_silu(x, bias):
+        raise RuntimeError("unsupported input for Sana fused bias-SiLU")
+    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]
+        )
+    return out
+
+
+def can_use_fused_bias_glu(x: torch.Tensor, bias: torch.Tensor) -> bool:
+    return (
+        _is_channels_last_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()
+    )
+
+
+def fused_bias_glu(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    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
+    out = torch.empty(
+        (batch, channels, height, width),
+        dtype=x.dtype,
+        device=x.device,
+        memory_format=torch.channels_last,
+    )
+    with torch.cuda.device(x.device):
+        _bias_glu_kernel[(triton.cdiv(out.numel(), 1024),)](
+            out, x, bias, out.numel(), channels=channels
+        )
+    return out
+
+
+__all__ = [
+    "can_use_fused_bias_glu",
+    "can_use_fused_bias_silu",
+    "fused_bias_glu",
+    "fused_bias_silu",
+]

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

@@ -6,11 +6,18 @@ import torch.nn.functional as F
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding
 
 from sglang.kernels.ops.diffusion.bitexact_gate import BitExactFusionGate
+from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add
 from sglang.kernels.ops.diffusion.triton.layernorm_modulate import (
     can_use_fused_layernorm_modulate,
     fused_layernorm_modulate_raw,
     is_plain_layer_norm,
 )
+from sglang.kernels.ops.diffusion.triton.sana_conv_post import (
+    can_use_fused_bias_glu,
+    can_use_fused_bias_silu,
+    fused_bias_glu,
+    fused_bias_silu,
+)
 from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig
 from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
 from sglang.multimodal_gen.runtime.layers.linear import MergedColumnParallelLinear
@@ -29,6 +36,8 @@ _SANA_LN_MOD = BitExactFusionGate("Sana fused LN+modulate", per_signature=True)
 _SANA_LN_MOD_SIGS = _SANA_LN_MOD.verified_sigs
 assert _SANA_LN_MOD_SIGS is not None
 _SANA_LN_MOD_DISABLED = False
+_SANA_CONV_SILU = BitExactFusionGate("Sana fused conv bias-SiLU")
+_SANA_CONV_GLU = BitExactFusionGate("Sana fused conv bias-GLU")
 
 
 def _eager_ln_modulate(
@@ -149,6 +158,103 @@ def _mps_safe_conv2d(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor:
     ).to(x.dtype)
 
 
+def _use_sana_bcg_fast_path(x: torch.Tensor) -> bool:
+    if torch.compiler.is_compiling() or not x.is_cuda:
+        return False
+    return torch.cuda.is_current_stream_capturing() or (
+        torch.cuda.current_stream() != torch.cuda.default_stream()
+    )
+
+
+def _conv2d_without_bias(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor:
+    return F.conv2d(
+        x,
+        conv.weight,
+        None,
+        conv.stride,
+        conv.padding,
+        conv.dilation,
+        conv.groups,
+    )
+
+
+def _sana_conv_bias_silu(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor:
+    if conv.bias is None or not _use_sana_bcg_fast_path(x):
+        return F.silu(_mps_safe_conv2d(conv, x))
+
+    raw = _conv2d_without_bias(conv, x)
+    if not can_use_fused_bias_silu(raw, conv.bias):
+        return F.silu(raw + conv.bias[None, :, None, None])
+    verified = _SANA_CONV_SILU.verified
+    if not verified and not _SANA_CONV_SILU.can_attempt_once():
+        return F.silu(raw + conv.bias[None, :, None, None])
+    try:
+        out = fused_bias_silu(raw, conv.bias)
+    except Exception as exc:
+        _SANA_CONV_SILU.on_exception(exc, logger=logger)
+        return F.silu(raw + conv.bias[None, :, None, None])
+    if verified:
+        return out
+    return _SANA_CONV_SILU.accept_or_fallback(
+        out,
+        F.silu(raw + conv.bias[None, :, None, None]),
+        logger=logger,
+        mismatch_msg=(
+            "Sana fused conv bias-SiLU path is not bit-exact on this "
+            "platform; falling back to eager"
+        ),
+    )
+
+
+def _sana_conv_bias_glu(conv: nn.Conv2d, x: torch.Tensor) -> torch.Tensor:
+    if conv.bias is None or not _use_sana_bcg_fast_path(x):
+        hidden_states = _mps_safe_conv2d(conv, x)
+        hidden_states, gate = torch.chunk(hidden_states, 2, dim=1)
+        return hidden_states * F.silu(gate)
+
+    raw = _conv2d_without_bias(conv, x)
+    if not can_use_fused_bias_glu(raw, conv.bias):
+        hidden_states, gate = torch.chunk(
+            raw + conv.bias[None, :, None, None], 2, dim=1
+        )
+        return hidden_states * F.silu(gate)
+    verified = _SANA_CONV_GLU.verified
+    if not verified and not _SANA_CONV_GLU.can_attempt_once():
+        hidden_states, gate = torch.chunk(
+            raw + conv.bias[None, :, None, None], 2, dim=1
+        )
+        return hidden_states * F.silu(gate)
+    try:
+        out = fused_bias_glu(raw, conv.bias)
+    except Exception as exc:
+        _SANA_CONV_GLU.on_exception(exc, logger=logger)
+        hidden_states, gate = torch.chunk(
+            raw + conv.bias[None, :, None, None], 2, dim=1
+        )
+        return hidden_states * F.silu(gate)
+    if verified:
+        return out
+    biased = raw + conv.bias[None, :, None, None]
+    hidden_states, gate = torch.chunk(biased, 2, dim=1)
+    return _SANA_CONV_GLU.accept_or_fallback(
+        out,
+        hidden_states * F.silu(gate),
+        logger=logger,
+        mismatch_msg=(
+            "Sana fused conv bias-GLU path is not bit-exact on this "
+            "platform; falling back to eager"
+        ),
+    )
+
+
+def _sana_residual_gate_add(
+    residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
+) -> torch.Tensor:
+    if torch.compiler.is_compiling():
+        return residual + gate * update
+    return residual_gate_add(residual, update, gate)
+
+
 def _mps_match_dtype(tensor: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
     if ref.device.type == "mps" and tensor.dtype != ref.dtype:
         return tensor.to(dtype=ref.dtype)
@@ -225,11 +331,8 @@ class GLUMBConv(nn.Module):
         self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False)
 
     def forward(self, hidden_states):
-        hidden_states = _mps_safe_conv2d(self.conv_inverted, hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
-        hidden_states = _mps_safe_conv2d(self.conv_depth, hidden_states)
-        hidden_states, gate = torch.chunk(hidden_states, 2, dim=1)
-        hidden_states = hidden_states * self.nonlinearity(gate)
+        hidden_states = _sana_conv_bias_silu(self.conv_inverted, hidden_states)
+        hidden_states = _sana_conv_bias_glu(self.conv_depth, hidden_states)
         hidden_states = _mps_safe_conv2d(self.conv_point, hidden_states)
         return hidden_states
 
@@ -376,7 +479,7 @@ class SanaTransformerBlock(nn.Module):
 
         norm_hidden = _sana_ln_modulate(self.norm1, hidden_states, scale_msa, shift_msa)
         attn_output = self.attn1(norm_hidden)
-        hidden_states = hidden_states + gate_msa * attn_output
+        hidden_states = _sana_residual_gate_add(hidden_states, attn_output, gate_msa)
 
         attn_output = self.attn2(
             hidden_states, encoder_hidden_states, encoder_attention_mask
@@ -387,7 +490,7 @@ class SanaTransformerBlock(nn.Module):
         norm_hidden = norm_hidden.unflatten(1, (height, width)).permute(0, 3, 1, 2)
         ff_output = self.ff(norm_hidden)
         ff_output = ff_output.flatten(2, 3).permute(0, 2, 1)
-        hidden_states = hidden_states + gate_mlp * ff_output
+        hidden_states = _sana_residual_gate_add(hidden_states, ff_output, gate_mlp)
 
         return hidden_states
 
@@ -484,7 +587,9 @@ class SanaTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         post_patch_width = width // p
 
         hidden_states = _mps_safe_conv2d(self.patch_embed["proj"], hidden_states)
-        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        # One layout conversion here prevents every downstream LayerNorm from
+        # copying the transposed patch view independently.
+        hidden_states = hidden_states.flatten(2).transpose(1, 2).contiguous()
 
         timestep_emb, embedded_timestep = self.time_embed(
             timestep, hidden_dtype=hidden_states.dtype