From 095ec6c997bfdd25d3864cb0ce77a6562a934b96 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sun, 16 Aug 2026 20:05:57 +0800 Subject: [PATCH] [diffusion][kernel] Accelerate Sana BCG with bit-exact conv post-processing (#34928) --- .../ops/diffusion/triton/sana_conv_post.py | 123 ++++++++++++++++++ .../runtime/models/dits/sana.py | 121 +++++++++++++++-- .../ops/diffusion/test_sana_conv_post.py | 60 +++++++++ 3 files changed, 296 insertions(+), 8 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/triton/sana_conv_post.py create mode 100644 test/registered/kernels/ops/diffusion/test_sana_conv_post.py diff --git a/python/sglang/kernels/ops/diffusion/triton/sana_conv_post.py b/python/sglang/kernels/ops/diffusion/triton/sana_conv_post.py new file mode 100644 index 000000000..2ee2b8a7b --- /dev/null +++ b/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", +] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana.py b/python/sglang/multimodal_gen/runtime/models/dits/sana.py index 303e561ef..cc12993d7 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana.py +++ b/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 diff --git a/test/registered/kernels/ops/diffusion/test_sana_conv_post.py b/test/registered/kernels/ops/diffusion/test_sana_conv_post.py new file mode 100644 index 000000000..89503bf21 --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_sana_conv_post.py @@ -0,0 +1,60 @@ +"""Sana GLUMB post-processing fusions must match the eager bf16 chain.""" + +import pytest +import torch +import torch.nn.functional as F + +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.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=3, stage="base-b-kernel-unit", runner_config="1-gpu-large") +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("channels", [2240, 11200]) +def test_sana_bias_silu_is_bit_exact(channels): + torch.manual_seed(0) + x = torch.randn( + (1, channels, 7, 5), + device="cuda", + dtype=torch.bfloat16, + ).to(memory_format=torch.channels_last) + bias = torch.randn(channels, device="cuda", dtype=torch.bfloat16) + + assert can_use_fused_bias_silu(x, bias) + actual = fused_bias_silu(x, bias) + expected = F.silu(x + bias[None, :, None, None]) + + assert actual.is_contiguous(memory_format=torch.channels_last) + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize("channels", [2240, 5600]) +def test_sana_bias_glu_is_bit_exact(channels): + torch.manual_seed(1) + x = torch.randn( + (1, 2 * channels, 7, 5), + device="cuda", + dtype=torch.bfloat16, + ).to(memory_format=torch.channels_last) + bias = torch.randn(2 * channels, device="cuda", dtype=torch.bfloat16) + + assert can_use_fused_bias_glu(x, bias) + actual = fused_bias_glu(x, bias) + biased = x + bias[None, :, None, None] + hidden, gate = torch.chunk(biased, 2, dim=1) + expected = hidden * F.silu(gate) + + assert actual.is_contiguous(memory_format=torch.channels_last) + assert torch.equal(actual, expected) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__]))