diff --git a/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py b/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py index 4766ddd5f..432b2808d 100644 --- a/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py +++ b/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 ( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana.py b/python/sglang/multimodal_gen/runtime/models/dits/sana.py index 0893e0245..dcc1dd2b1 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana.py +++ b/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) diff --git a/test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py b/test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py new file mode 100644 index 000000000..eb0a01eac --- /dev/null +++ b/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__]))