[diffusion] Sana: bit-exact fused aten LayerNorm+modulate under BCG (H200 denoise -4.8%) (#34015)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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

@@ -14,7 +14,9 @@ which ``nn.LayerNorm`` dispatches to for bf16 rows with ``N % 4 == 0`` and
 16-byte-aligned buffers; SASS-level derivation in PR #34008):
 
 - 128 aten threads per row; thread ``t`` serially Welford-pushes the
-  4-element vectors ``t, t+128, ...``, each scalar as
+  4-element vectors ``t, t+128, ... < N/4`` (for ``N % 512 != 0`` the
+  trailing threads stop one iteration early and enter the fold with a
+  smaller count), each scalar as
   ``mean' = fma(delta, rcp(count+1), mean)``,
   ``m2' = fma(delta, val - mean', m2)``, with ``rcp`` being nvcc's
   guarded-reciprocal fast path (``_rcp4``).
@@ -230,14 +232,30 @@ def _layernorm_modulate_kernel(
 
     # pass 1: per-"thread" serial Welford in aten's exact element order.
     # Out-of-range rows compute garbage that is never stored.
-    for i in tl.static_range(D // 512):
+    for i in tl.static_range((D + 511) // 512):
         cols = i * 512 + lanes[:, None] * 4 + tl.arange(0, 4)[None, :]
-        x4 = tl.load(
-            x_ptr + row_base[:, None, None] + cols[None, :, :],
-            mask=row_mask[:, None, None],
-            other=0.0,
-        ).to(tl.float32)
-        mean, m2, cnt = _push_vec4(x4, mean, m2, cnt, row_mask, ROWS, 128, MASKED=False)
+        if (i + 1) * 512 <= D:
+            x4 = tl.load(
+                x_ptr + row_base[:, None, None] + cols[None, :, :],
+                mask=row_mask[:, None, None],
+                other=0.0,
+            ).to(tl.float32)
+            mean, m2, cnt = _push_vec4(
+                x4, mean, m2, cnt, row_mask, ROWS, 128, MASKED=False
+            )
+        else:
+            # partial tail chunk (D % 512 != 0): aten threads whose vector
+            # index i*128+t reaches N/4 skip the iteration, leaving their
+            # Welford state untouched.
+            vec_valid = (i * 128 + lanes < D // 4)[None, :]
+            x4 = tl.load(
+                x_ptr + row_base[:, None, None] + cols[None, :, :],
+                mask=row_mask[:, None, None] & vec_valid[:, :, None],
+                other=0.0,
+            ).to(tl.float32)
+            mean, m2, cnt = _push_vec4(
+                x4, mean, m2, cnt, vec_valid, ROWS, 128, MASKED=True
+            )
 
     # warp fold trees, then the (0,2)/(1,3)/(0,1) inter-warp combines.
     mean = tl.reshape(mean, (ROWS * 4, 32), can_reorder=False)
@@ -256,27 +274,30 @@ def _layernorm_modulate_kernel(
     batch = row_offs // seq_len
 
     # pass 2: normalize + modulate, in aten's rounding order.
-    for i in tl.static_range(D // 512):
+    for i in tl.static_range((D + 511) // 512):
         cols = i * 512 + tl.arange(0, 512)
+        mask = row_mask[:, None]
+        if (i + 1) * 512 > D:
+            mask = mask & (cols < D)[None, :]
         x = tl.load(
             x_ptr + row_base[:, None] + cols[None, :],
-            mask=row_mask[:, None],
+            mask=mask,
             other=0.0,
         ).to(tl.float32)
         y = _round_bf16_to_fp32(rstd * (x - mean))
         sc = tl.load(
             scale_ptr + batch[:, None] * scale_row_stride + cols[None, :],
-            mask=row_mask[:, None],
+            mask=mask,
             other=0.0,
         ).to(tl.float32)
         sh = tl.load(
             shift_ptr + batch[:, None] * scale_row_stride + cols[None, :],
-            mask=row_mask[:, None],
+            mask=mask,
             other=0.0,
         ).to(tl.float32)
         one_plus = _round_bf16_to_fp32(1.0 + sc)
         y = _round_bf16_to_fp32(y * one_plus) + sh
-        tl.store(y_ptr + row_base[:, None] + cols[None, :], y, mask=row_mask[:, None])
+        tl.store(y_ptr + row_base[:, None] + cols[None, :], y, mask=mask)
 
 
 @triton.jit
@@ -376,7 +397,7 @@ def can_use_fused_layernorm_modulate(
         and x.dim() == 3
         and x.numel() > 0
         and x.is_contiguous()
-        and x.shape[-1] % 512 == 0
+        and x.shape[-1] % 4 == 0
         and x.shape[-1] <= 8192
         and _is_bf16_cuda(scale)
         and _is_bf16_cuda(shift)
@@ -396,16 +417,16 @@ def _fake_ln_modulate(
     return torch.empty_like(x)
 
 
-@register_custom_op(
-    op_name="triton_fused_layernorm_modulate",
-    mutates_args=[],
-    fake_impl=_fake_ln_modulate,
-)
-def fused_layernorm_modulate(
+def fused_layernorm_modulate_raw(
     x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, eps: float
 ) -> torch.Tensor:
     """``LN(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)``, bit-exact
-    vs the eager aten chain (LayerNorm without affine)."""
+    vs the eager aten chain (LayerNorm without affine).
+
+    Direct-call variant without the ``torch.ops`` dispatch (which costs tens
+    of microseconds per call); use it on CPU-launch-bound eager hot paths
+    (e.g. Sana), and the registered custom op under ``torch.compile``.
+    """
     batch, seq_len, hidden = x.shape
     n_rows = batch * seq_len
     rows = 2
@@ -424,13 +445,22 @@ def fused_layernorm_modulate(
             D=hidden,
             ROWS=rows,
             # H200-tuned: 38.5us at (1, 4096, 4096) vs the 121.8us eager
-            # chain.  ROWS=1 + 4 warps triggers pathological Triton layout
-            # conversions in the fold stage (47-58us).
-            num_warps=4 if hidden >= 4096 else 2,
+            # chain, 14.3us at Sana's (2, 1024, 2240) vs 43.1us.  ROWS=1 +
+            # 4 warps triggers pathological Triton layout conversions in
+            # the fold stage (47-58us).
+            num_warps=4 if hidden >= 2048 else 2,
         )
     return out
 
 
+fused_layernorm_modulate = register_custom_op(
+    fused_layernorm_modulate_raw,
+    op_name="triton_fused_layernorm_modulate",
+    mutates_args=[],
+    fake_impl=_fake_ln_modulate,
+)
+
+
 def can_use_fused_qk_head_layernorm(q: torch.Tensor, k: torch.Tensor) -> bool:
     head_dim = q.shape[-1] if q.dim() == 4 else 0
     return (

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

@@ -5,6 +5,11 @@ import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding
 
+from sglang.kernels.ops.diffusion.triton.layernorm_modulate import (
+    can_use_fused_layernorm_modulate,
+    fused_layernorm_modulate_raw,
+    is_plain_layer_norm,
+)
 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
@@ -17,6 +22,101 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
 
+_SANA_FUSED_LN_MOD_DISABLED = False
+# (shape/stride/dtype/eps) signatures whose fused output was torch.equal-
+# verified against the live eager chain.
+_SANA_FUSED_LN_MOD_OK_SIGS: set = set()
+
+
+def _eager_ln_modulate(
+    norm: nn.LayerNorm,
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+) -> torch.Tensor:
+    return norm(x) * (1 + scale) + shift
+
+
+def _sana_ln_modulate(
+    norm: nn.LayerNorm,
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+) -> torch.Tensor:
+    """Single-kernel ``LN(x) * (1 + scale) + shift``, bit-exact vs eager.
+
+    ``scale`` / ``shift`` are Sana's ``(batch, 1, dim)`` adaLN rows.  Each
+    new input signature is verified with ``torch.equal`` against the eager
+    chain (bit-exactness depends on which LayerNorm kernel aten dispatches
+    to); any mismatch disables the fast path permanently.
+
+    The fusion only engages in CUDA-graph contexts (stream capture, or the
+    non-default stream the breakable-CUDA-graph runner warms up and captures
+    on, #33989), where replay pays no Python launch cost.  Sana's eager mode
+    is CPU-launch bound and one Triton launch costs more Python than the
+    whole aten chain it replaces (+14% forward wall), so default-stream
+    serving keeps the eager chain.
+
+    ``x`` reaches every site transposed (patch-embed / GLUMBConv permute
+    layout); aten's LayerNorm contiguizes internally, and the fast path
+    issues the same copy explicitly.
+    """
+    global _SANA_FUSED_LN_MOD_DISABLED
+
+    if _SANA_FUSED_LN_MOD_DISABLED or torch.compiler.is_compiling() or not x.is_cuda:
+        return _eager_ln_modulate(norm, x, scale, shift)
+
+    capturing = torch.cuda.is_current_stream_capturing()
+    if not capturing and torch.cuda.current_stream() == torch.cuda.default_stream():
+        return _eager_ln_modulate(norm, x, scale, shift)
+
+    sig = (
+        x.shape,
+        x.stride(),
+        x.dtype,
+        scale.shape,
+        scale.stride(),
+        shift.stride(),
+        norm.eps,
+    )
+    if sig in _SANA_FUSED_LN_MOD_OK_SIGS:
+        return fused_layernorm_modulate_raw(
+            x.contiguous(), scale[:, 0], shift[:, 0], norm.eps
+        )
+    if capturing:
+        # unverified signature during capture: cannot sync-verify here
+        return _eager_ln_modulate(norm, x, scale, shift)
+
+    if (
+        x.dtype is torch.bfloat16
+        and x.dim() == 3
+        and scale.dim() == 3
+        and scale.shape[1] == 1
+        and shift.shape == scale.shape
+        and is_plain_layer_norm(norm, x.shape[-1])
+    ):
+        x_c = x.contiguous()
+        if not can_use_fused_layernorm_modulate(x_c, scale[:, 0], shift[:, 0]):
+            return _eager_ln_modulate(norm, x, scale, shift)
+        try:
+            out = fused_layernorm_modulate_raw(x_c, scale[:, 0], shift[:, 0], norm.eps)
+        except Exception as exc:
+            logger.warning_once(f"Disabling Sana fused LN+modulate fast path: {exc}")
+            _SANA_FUSED_LN_MOD_DISABLED = True
+        else:
+            ref = _eager_ln_modulate(norm, x, scale, shift)
+            if torch.equal(out, ref):
+                _SANA_FUSED_LN_MOD_OK_SIGS.add(sig)
+                return out
+            logger.warning_once(
+                "Sana fused LN+modulate fast path is not bit-exact against "
+                "this platform's LayerNorm dispatch; falling back to eager"
+            )
+            _SANA_FUSED_LN_MOD_DISABLED = True
+            return ref
+
+    return _eager_ln_modulate(norm, x, scale, shift)
+
 
 def _mps_safe_linear(linear: nn.Linear, x: torch.Tensor) -> torch.Tensor:
     if x.device.type != "mps":
@@ -96,11 +196,9 @@ class SanaModulatedNorm(nn.Module):
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
 
     def forward(self, x, temb, scale_shift_table):
-        x = self.norm(x)
         scale_shift_table = _mps_match_dtype(scale_shift_table, temb)
         shift, scale = (scale_shift_table[None] + temb[:, None]).chunk(2, dim=1)
-        x = x * (1 + scale) + shift
-        return x
+        return _sana_ln_modulate(self.norm, x, scale, shift)
 
 
 class GLUMBConv(nn.Module):
@@ -271,8 +369,7 @@ class SanaTransformerBlock(nn.Module):
             scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
         ).chunk(6, dim=1)
 
-        norm_hidden = self.norm1(hidden_states)
-        norm_hidden = norm_hidden * (1 + scale_msa) + shift_msa
+        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
 
@@ -281,8 +378,7 @@ class SanaTransformerBlock(nn.Module):
         )
         hidden_states = hidden_states + attn_output
 
-        norm_hidden = self.norm2(hidden_states)
-        norm_hidden = norm_hidden * (1 + scale_mlp) + shift_mlp
+        norm_hidden = _sana_ln_modulate(self.norm2, hidden_states, scale_mlp, shift_mlp)
         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)

test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py

@@ -0,0 +1,55 @@
+"""Sana fused LN+modulate fast path must stay bit-exact vs eager."""
+
+import pytest
+import torch
+
+import sglang.multimodal_gen.runtime.models.dits.sana as sana
+from sglang.multimodal_gen.runtime.models.dits.sana import (
+    _eager_ln_modulate,
+    _sana_ln_modulate,
+)
+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(
+    "shape,nmod,transposed",
+    [
+        ((2, 1024, 2240), 6, False),
+        ((2, 1024, 2240), 2, False),
+        ((1, 1024, 2240), 6, True),
+        ((1, 37, 2240), 6, False),
+    ],
+)
+def test_sana_fused_ln_modulate_is_bit_exact(shape, nmod, transposed):
+    # (., 1024, 2240) is the real Sana 1024px shape; hidden 2240 % 512 != 0
+    # exercises the kernel's partial tail chunk.  nmod mirrors the two adaLN
+    # chunk layouts, transposed the permuted layout the Sana DiT serves.
+    torch.manual_seed(0)
+    batch, seq, hidden = shape
+    norm = torch.nn.LayerNorm(hidden, eps=1e-6, elementwise_affine=False).cuda()
+    x = (torch.randn(batch, seq, hidden, device="cuda") * 4).bfloat16()
+    if transposed:
+        x = x.permute(0, 2, 1).contiguous().permute(0, 2, 1)
+    emb = torch.randn(batch, nmod, hidden, device="cuda").bfloat16()
+    shift, scale = emb.chunk(nmod, dim=1)[0], emb.chunk(nmod, dim=1)[-1]
+    # default-stream eager serving must stay on the untouched eager chain
+    n_sigs = len(sana._SANA_FUSED_LN_MOD_OK_SIGS)
+    _sana_ln_modulate(norm, x, scale, shift)
+    assert len(sana._SANA_FUSED_LN_MOD_OK_SIGS) == n_sigs
+    # the fusion engages on non-default streams (the BCG warmup/capture path)
+    with torch.cuda.stream(torch.cuda.Stream()):
+        out = _sana_ln_modulate(norm, x, scale, shift)
+        assert len(sana._SANA_FUSED_LN_MOD_OK_SIGS) == n_sigs + 1  # verified
+        out2 = _sana_ln_modulate(norm, x, scale, shift)  # verified-sig lane
+    torch.cuda.synchronize()
+    assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift))
+    assert torch.equal(out2, out) and not sana._SANA_FUSED_LN_MOD_DISABLED
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__]))