minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-08-08 16:05:25 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent d747bd052e
commit 891445676c
3 changed files with 212 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py
+54 -24
View File
@@ -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
+103 -7
View File
@@ -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)