From 5dffa06fe1c3068d7680bad84715bd78ee2315fb Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 8 Aug 2026 13:26:11 +0800 Subject: [PATCH] [diffusion] GLM-Image bit-exact fused aten LayerNorm+modulate / qk-LN (H200 30-step denoise -8.1%) (#34008) Co-authored-by: Claude Fable 5 --- .../diffusion/triton/layernorm_modulate.py | 486 ++++++++++++++++++ .../runtime/models/dits/glm_image.py | 189 ++++++- .../diffusion/test_glm_image_ln_modulate.py | 55 ++ 3 files changed, 713 insertions(+), 17 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py create mode 100644 test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py diff --git a/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py b/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py new file mode 100644 index 000000000..4766ddd5f --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused LayerNorm + adaLN modulate Triton kernels for bf16 activations. + +Two fusions, each replacing an eager multi-kernel chain with a single +launch while reproducing the eager results bit for bit (``torch.equal``), +so callers need no quality gate: + +- ``fused_layernorm_modulate``: ``LN(x) * (1 + scale) + shift`` +- ``fused_qk_head_layernorm``: per-head ``LN(q)`` / ``LN(k)`` + +Numerics contract (replicates torch 2.11's +``at::native::vectorized_layer_norm_kernel``, +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 + ``mean' = fma(delta, rcp(count+1), mean)``, + ``m2' = fma(delta, val - mean', m2)``, with ``rcp`` being nvcc's + guarded-reciprocal fast path (``_rcp4``). +- Lane states fold via shfl.down offsets 16, 8, 4, 2, 1 through + ``cuWelfordCombine`` (argument order: self = lower lane, other = upper; + non-positive counts fold to zeros); the 4 warp results then combine + pairwise (0,2), (1,3), (0,1). +- ``rstd = rsqrtf(div.rn.f32(m2, N) + eps)``: correctly-rounded fp32 + division, then ``MUFU.RSQ`` with a 2^24/2^12 rescale for subnormals. +- The normalized output is ``cvt.rn.bf16(rstd * (x - mean))`` (mul then + add, no fma); the modulate chain then rounds to bf16 after every op: + ``round(1 + scale)``, ``round(y * that)``, ``round(prod + shift)``. + +For the qk kernel (per-head rows, ``dim_head % 4 == 0``, ``dim_head <= +128``) only the first ``dim_head / 4`` of the 128 aten threads carry data; +the remaining lanes and warps enter the fold with ``count == 0`` and are +reproduced faithfully. + +Bit-exactness holds only for the dispatch above, so callers must verify +``torch.equal`` against the live eager chain once at runtime and fall back +on mismatch (see ``glm_image.py``). +""" + +from __future__ import annotations + +import torch +import triton # type: ignore +import triton.language as tl # type: ignore + +from sglang.srt.utils.custom_op import register_custom_op + +_FLT_MIN = tl.constexpr(1.1754943508222875e-38) + + +@triton.jit +def _round_bf16_to_fp32(value): + # RNE round of an fp32 value to bf16 precision, staying in fp32 registers + # (also blocks any fmul+fadd contraction across the boundary). + bits = value.to(tl.int32, bitcast=True) + rounding_bias = 0x7FFF + ((bits >> 16) & 1) + rounded_bits = (bits + rounding_bias) & -65536 + return rounded_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def _rcp4(x): + # nvcc's reciprocal fast path (always taken for our integer counts). + return tl.inline_asm_elementwise( + asm="""{ + .reg .f32 r0, e, e2; + rcp.approx.f32 r0, $1; + fma.rn.f32 e, $1, r0, 0fBF800000; + sub.ftz.f32 e2, 0f80000000, e; + fma.rn.f32 $0, r0, e2, r0; + }""", + constraints="=f,f", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _div_rn(x, y): + # IEEE correctly-rounded fp32 division. + return tl.inline_asm_elementwise( + asm="div.rn.f32 $0, $1, $2;", + constraints="=f,f,f", + args=[x, y], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _rsqrt_approx(x): + return tl.inline_asm_elementwise( + asm="rsqrt.approx.f32 $0, $1;", + constraints="=f,f", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _rsqrtf(x): + # CUDA rsqrtf: MUFU.RSQ with a 2^24 / 2^12 rescale for subnormal inputs. + p = tl.abs(x) < _FLT_MIN + xs = tl.where(p, x * 16777216.0, x) + r = _rsqrt_approx(xs) + return tl.where(p, r * 4096.0, r) + + +@triton.jit +def _welford_push(val, mean, m2, cnt, valid, MASKED: tl.constexpr): + # ``valid`` masks lanes whose aten thread does not execute this + # iteration (their state must stay untouched). + delta = val - mean + new_cnt = cnt + 1.0 + recip = _rcp4(new_cnt) + new_mean = tl.fma(delta, recip, mean) + t = val - new_mean + new_m2 = tl.fma(delta, t, m2) + if MASKED: + new_mean = tl.where(valid, new_mean, mean) + new_m2 = tl.where(valid, new_m2, m2) + new_cnt = tl.where(valid, new_cnt, cnt) + return new_mean, new_m2, new_cnt + + +@triton.jit +def _welford_combine(mean_b, m2_b, cnt_b, mean_a, m2_a, cnt_a): + # b = self / lower lane, a = other / upper lane; the op order matters. + count = cnt_a + cnt_b + pos = count > 0.0 + coef = _rcp4(tl.where(pos, count, 1.0)) + delta = mean_b - mean_a + n_b = coef * cnt_b + d2 = delta * delta + n_a = cnt_a * coef + s = m2_a + m2_b + t1 = n_b * mean_b + mean = tl.fma(mean_a, n_a, t1) + t2 = cnt_a * d2 + m2 = tl.fma(n_b, t2, s) + mean = tl.where(pos, mean, 0.0) + m2 = tl.where(pos, m2, 0.0) + return mean, m2, count + + +@triton.jit +def _split_halves(x, rows: tl.constexpr, half: tl.constexpr): + # (rows, 2*half) -> two (rows, half) tensors pairing lane i with i+half, + # the tree a shfl.down fold with the largest offset first produces. + x = tl.reshape(x, (rows, 2, half), can_reorder=False) + x = tl.permute(x, (0, 2, 1)) + return tl.split(x) + + +@triton.jit +def _fold_halves(mean, m2, cnt, rows: tl.constexpr, half: tl.constexpr): + mb, ma = _split_halves(mean, rows, half) + sb, sa = _split_halves(m2, rows, half) + cb, ca = _split_halves(cnt, rows, half) + return _welford_combine(mb, sb, cb, ma, sa, ca) + + +@triton.jit +def _fold_tree_32(mean, m2, cnt, rows: tl.constexpr): + # shfl.down offsets 16, 8, 4, 2, 1 over the 32 lane states -> (rows, 1). + mean, m2, cnt = _fold_halves(mean, m2, cnt, rows, 16) + mean, m2, cnt = _fold_halves(mean, m2, cnt, rows, 8) + mean, m2, cnt = _fold_halves(mean, m2, cnt, rows, 4) + mean, m2, cnt = _fold_halves(mean, m2, cnt, rows, 2) + mean, m2, cnt = _fold_halves(mean, m2, cnt, rows, 1) + return mean, m2, cnt + + +@triton.jit +def _push_vec4( + x4, + mean, + m2, + cnt, + valid, + rows: tl.constexpr, + lanes: tl.constexpr, + MASKED: tl.constexpr, +): + # Push the 4 elements of one aligned_vector in exact serial + # order. x4 is (rows, lanes, 4) fp32. + g = tl.reshape(x4, (rows, lanes, 2, 2), can_reorder=False) + p02, p13 = tl.split(g) # elements (0, 2) / (1, 3) + e0, e2 = tl.split(tl.reshape(p02, (rows, lanes, 1, 2), can_reorder=False)) + e1, e3 = tl.split(tl.reshape(p13, (rows, lanes, 1, 2), can_reorder=False)) + e0 = tl.reshape(e0, (rows, lanes), can_reorder=False) + e1 = tl.reshape(e1, (rows, lanes), can_reorder=False) + e2 = tl.reshape(e2, (rows, lanes), can_reorder=False) + e3 = tl.reshape(e3, (rows, lanes), can_reorder=False) + mean, m2, cnt = _welford_push(e0, mean, m2, cnt, valid, MASKED) + mean, m2, cnt = _welford_push(e1, mean, m2, cnt, valid, MASKED) + mean, m2, cnt = _welford_push(e2, mean, m2, cnt, valid, MASKED) + mean, m2, cnt = _welford_push(e3, mean, m2, cnt, valid, MASKED) + return mean, m2, cnt + + +@triton.jit +def _layernorm_modulate_kernel( + y_ptr, + x_ptr, + scale_ptr, + shift_ptr, + seq_len, + n_rows, + scale_row_stride, + eps, + D: tl.constexpr, + ROWS: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + row_offs = pid * ROWS + tl.arange(0, ROWS) + row_mask = row_offs < n_rows + row_base = row_offs * D + + lanes = tl.arange(0, 128) + mean = tl.zeros((ROWS, 128), dtype=tl.float32) + m2 = tl.zeros((ROWS, 128), dtype=tl.float32) + cnt = tl.zeros((ROWS, 128), dtype=tl.float32) + + # 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): + 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) + + # warp fold trees, then the (0,2)/(1,3)/(0,1) inter-warp combines. + mean = tl.reshape(mean, (ROWS * 4, 32), can_reorder=False) + m2 = tl.reshape(m2, (ROWS * 4, 32), can_reorder=False) + cnt = tl.reshape(cnt, (ROWS * 4, 32), can_reorder=False) + mean, m2, cnt = _fold_tree_32(mean, m2, cnt, ROWS * 4) + mean = tl.reshape(mean, (ROWS, 4), can_reorder=False) + m2 = tl.reshape(m2, (ROWS, 4), can_reorder=False) + cnt = tl.reshape(cnt, (ROWS, 4), can_reorder=False) + mean, m2, cnt = _fold_halves(mean, m2, cnt, ROWS, 2) + mean, m2, cnt = _fold_halves(mean, m2, cnt, ROWS, 1) + + denom = tl.zeros((ROWS, 1), dtype=tl.float32) + D + rstd = _rsqrtf(_div_rn(m2, denom) + eps) # (ROWS, 1) + + batch = row_offs // seq_len + + # pass 2: normalize + modulate, in aten's rounding order. + for i in tl.static_range(D // 512): + cols = i * 512 + tl.arange(0, 512) + x = tl.load( + x_ptr + row_base[:, None] + cols[None, :], + mask=row_mask[:, None], + 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], + other=0.0, + ).to(tl.float32) + sh = tl.load( + shift_ptr + batch[:, None] * scale_row_stride + cols[None, :], + mask=row_mask[:, None], + 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]) + + +@triton.jit +def _qk_ln_head_one( + dst, + src, + pid, + n_rows, + eps, + D: tl.constexpr, + D_POW2: tl.constexpr, + ROWS: tl.constexpr, +): + row_offs = pid * ROWS + tl.arange(0, ROWS) + row_mask = row_offs < n_rows + row_base = row_offs * D + + lanes = tl.arange(0, 32) + lane_valid = (lanes < D // 4)[None, :] + mean = tl.zeros((ROWS, 32), dtype=tl.float32) + m2 = tl.zeros((ROWS, 32), dtype=tl.float32) + cnt = tl.zeros((ROWS, 32), dtype=tl.float32) + + cols = lanes[:, None] * 4 + tl.arange(0, 4)[None, :] + x4 = tl.load( + src + row_base[:, None, None] + cols[None, :, :], + mask=row_mask[:, None, None] & lane_valid[:, :, None], + other=0.0, + ).to(tl.float32) + mean, m2, cnt = _push_vec4(x4, mean, m2, cnt, lane_valid, ROWS, 32, MASKED=True) + + mean, m2, cnt = _fold_tree_32(mean, m2, cnt, ROWS) + # inter-warp combines with the all-zero warps 1..3 of the aten block: + # (0,2) with zero, then (0,1) where warp 1 combined two zero warps. + zero = tl.zeros((ROWS, 1), dtype=tl.float32) + mean, m2, cnt = _welford_combine(mean, m2, cnt, zero, zero, zero) + mean, m2, cnt = _welford_combine(mean, m2, cnt, zero, zero, zero) + + denom = tl.zeros((ROWS, 1), dtype=tl.float32) + D + rstd = _rsqrtf(_div_rn(m2, denom) + eps) + + cols2 = tl.arange(0, D_POW2) + out_mask = row_mask[:, None] & (cols2 < D)[None, :] + x = tl.load(src + row_base[:, None] + cols2[None, :], mask=out_mask, other=0.0).to( + tl.float32 + ) + y = rstd * (x - mean) + tl.store(dst + row_base[:, None] + cols2[None, :], y, mask=out_mask) + + +@triton.jit +def _qk_ln_head_kernel( + q_out_ptr, + k_out_ptr, + q_ptr, + k_ptr, + n_rows, + eps, + D: tl.constexpr, + D_POW2: tl.constexpr, + ROWS: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + if tl.program_id(1) == 0: + _qk_ln_head_one(q_out_ptr, q_ptr, pid, n_rows, eps, D, D_POW2, ROWS) + else: + _qk_ln_head_one(k_out_ptr, k_ptr, pid, n_rows, eps, D, D_POW2, ROWS) + + +def is_plain_layer_norm(norm: torch.nn.Module, hidden: int) -> bool: + """True for a bare ``nn.LayerNorm((hidden,))`` without affine params.""" + return ( + type(norm) is torch.nn.LayerNorm + and norm.weight is None + and norm.bias is None + and tuple(norm.normalized_shape) == (hidden,) + ) + + +def _is_bf16_cuda(t: torch.Tensor) -> bool: + return t.is_cuda and t.dtype is torch.bfloat16 + + +def _mod_row_stride(t: torch.Tensor, batch: int, hidden: int) -> int | None: + # (batch, hidden) modulation rows, possibly strided views of a chunked + # adaLN projection; the last dim must be packed. + if t.dim() != 2 or t.shape != (batch, hidden) or t.stride(1) != 1: + return None + return t.stride(0) if batch > 1 else hidden + + +def can_use_fused_layernorm_modulate( + x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor +) -> bool: + if not ( + _is_bf16_cuda(x) + and x.dim() == 3 + and x.numel() > 0 + and x.is_contiguous() + and x.shape[-1] % 512 == 0 + and x.shape[-1] <= 8192 + and _is_bf16_cuda(scale) + and _is_bf16_cuda(shift) + and scale.device == x.device + and shift.device == x.device + ): + return False + batch, _, hidden = x.shape + q = _mod_row_stride(scale, batch, hidden) + v = _mod_row_stride(shift, batch, hidden) + return q is not None and v is not None and q == v + + +def _fake_ln_modulate( + x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, eps: float +) -> torch.Tensor: + 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( + 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).""" + batch, seq_len, hidden = x.shape + n_rows = batch * seq_len + rows = 2 + out = torch.empty_like(x) + stride = _mod_row_stride(scale, batch, hidden) + with torch.cuda.device(x.device): + _layernorm_modulate_kernel[(triton.cdiv(n_rows, rows),)]( + out, + x, + scale, + shift, + seq_len, + n_rows, + stride, + eps, + 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, + ) + return out + + +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 ( + _is_bf16_cuda(q) + and _is_bf16_cuda(k) + and q.device == k.device + and q.dim() == 4 + and q.shape == k.shape + and head_dim % 4 == 0 + and 0 < head_dim <= 128 + and q.numel() > 0 + and q.is_contiguous() + and k.is_contiguous() + ) + + +def _fake_qk_ln( + q: torch.Tensor, k: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.empty_like(q), torch.empty_like(k) + + +@register_custom_op( + op_name="triton_fused_qk_head_layernorm", + mutates_args=[], + fake_impl=_fake_qk_ln, +) +def fused_qk_head_layernorm( + q: torch.Tensor, k: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-head ``nn.LayerNorm(dim_head)`` (no affine) over q and k in one + launch, bit-exact vs the eager aten kernel.""" + head_dim = q.shape[-1] + n_rows = q.numel() // head_dim + rows = 64 + q_out = torch.empty_like(q) + k_out = torch.empty_like(k) + with torch.cuda.device(q.device): + _qk_ln_head_kernel[(triton.cdiv(n_rows, rows), 2)]( + q_out, + k_out, + q, + k, + n_rows, + eps, + D=head_dim, + D_POW2=triton.next_power_of_2(head_dim), + ROWS=rows, + # H200-tuned: 62us at (1, 4360, 32, 128) vs the 301us of the two + # aten launches (one 128-thread block per head_dim-element row). + num_warps=2, + ) + return q_out, k_out diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index 9ce0fcddd..ebc066b4a 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py @@ -23,6 +23,17 @@ from sglang.kernels.ops.diffusion.fused_linear_gelu import ( fused_linear_gelu_tanh, mark_fused_gelu_site, ) +from sglang.kernels.ops.diffusion.residual_gate_add import ( + can_use_residual_gate_add_cuda, + residual_gate_add_cuda, +) +from sglang.kernels.ops.diffusion.triton.layernorm_modulate import ( + can_use_fused_layernorm_modulate, + can_use_fused_qk_head_layernorm, + fused_layernorm_modulate, + fused_qk_head_layernorm, + is_plain_layer_norm, +) from sglang.multimodal_gen.configs.models.dits.glmimage import GlmImageDitConfig from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_parallel_rank, @@ -64,6 +75,147 @@ logger = init_logger(__name__) _is_cuda = current_platform.is_cuda() +_GLM_FUSED_LN_MOD_DISABLED = False +_GLM_FUSED_LN_MOD_VERIFIED = False +_GLM_FUSED_QK_LN_DISABLED = False +_GLM_FUSED_QK_LN_VERIFIED = False +_GLM_RESIDUAL_GATE_CUDA_DISABLED = False + + +def _eager_ln_modulate( + norm: nn.LayerNorm, + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + return norm(x).to(dtype=dtype) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +def _glm_ln_modulate( + norm: nn.LayerNorm, + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + """Single-kernel ``LN(x) * (1 + scale) + shift``, bit-exact vs eager. + + Bit-exactness depends on which LayerNorm kernel aten dispatches to, so + the first call verifies ``torch.equal`` against the eager chain and + disables the fast path permanently on any mismatch. + """ + global _GLM_FUSED_LN_MOD_DISABLED, _GLM_FUSED_LN_MOD_VERIFIED + + if ( + not _GLM_FUSED_LN_MOD_DISABLED + and _is_cuda + and dtype is x.dtype + and is_plain_layer_norm(norm, x.shape[-1]) + and can_use_fused_layernorm_modulate(x, scale, shift) + and (_GLM_FUSED_LN_MOD_VERIFIED or not torch.compiler.is_compiling()) + ): + try: + out = fused_layernorm_modulate(x, scale, shift, norm.eps) + except Exception as exc: + if torch.compiler.is_compiling(): + raise + logger.warning_once(f"Disabling GLM fused LN+modulate fast path: {exc}") + _GLM_FUSED_LN_MOD_DISABLED = True + else: + if _GLM_FUSED_LN_MOD_VERIFIED: + return out + ref = _eager_ln_modulate(norm, x, scale, shift, dtype) + if torch.equal(out, ref): + _GLM_FUSED_LN_MOD_VERIFIED = True + return out + logger.warning_once( + "GLM fused LN+modulate fast path is not bit-exact against " + "this platform's LayerNorm dispatch; falling back to eager" + ) + _GLM_FUSED_LN_MOD_DISABLED = True + return ref + + return _eager_ln_modulate(norm, x, scale, shift, dtype) + + +def _glm_qk_layernorm( + norm_q: nn.LayerNorm, + norm_k: nn.LayerNorm, + query: torch.Tensor, + key: torch.Tensor, + dtype: torch.dtype, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-head LayerNorm over q and k in one launch, bit-exact vs eager. + + First call verifies ``torch.equal`` against the eager pair and falls + back permanently on any mismatch. + """ + global _GLM_FUSED_QK_LN_DISABLED, _GLM_FUSED_QK_LN_VERIFIED + + if ( + not _GLM_FUSED_QK_LN_DISABLED + and _is_cuda + and dtype is query.dtype + and dtype is key.dtype + and is_plain_layer_norm(norm_q, query.shape[-1]) + and is_plain_layer_norm(norm_k, key.shape[-1]) + and norm_q.eps == norm_k.eps + and can_use_fused_qk_head_layernorm(query, key) + and (_GLM_FUSED_QK_LN_VERIFIED or not torch.compiler.is_compiling()) + ): + try: + q_out, k_out = fused_qk_head_layernorm(query, key, norm_q.eps) + except Exception as exc: + if torch.compiler.is_compiling(): + raise + logger.warning_once(f"Disabling GLM fused qk-LayerNorm fast path: {exc}") + _GLM_FUSED_QK_LN_DISABLED = True + else: + if _GLM_FUSED_QK_LN_VERIFIED: + return q_out, k_out + q_ref = norm_q(query).to(dtype=dtype) + k_ref = norm_k(key).to(dtype=dtype) + if torch.equal(q_out, q_ref) and torch.equal(k_out, k_ref): + _GLM_FUSED_QK_LN_VERIFIED = True + return q_out, k_out + logger.warning_once( + "GLM fused qk-LayerNorm fast path is not bit-exact against " + "this platform's LayerNorm dispatch; falling back to eager" + ) + _GLM_FUSED_QK_LN_DISABLED = True + return q_ref, k_ref + + return norm_q(query).to(dtype=dtype), norm_k(key).to(dtype=dtype) + + +def _glm_residual_gate_add( + residual: torch.Tensor, + update: torch.Tensor, + gate: torch.Tensor, +) -> torch.Tensor: + """Single-kernel ``residual + gate * update``, bit-exact vs the eager pair. + + Half dtypes only: for fp32 the kernel would contract to an fma (one + rounding) and stop being bit-exact. + """ + global _GLM_RESIDUAL_GATE_CUDA_DISABLED + + if ( + not _GLM_RESIDUAL_GATE_CUDA_DISABLED + and residual.dtype in (torch.float16, torch.bfloat16) + and can_use_residual_gate_add_cuda(residual, update, gate) + ): + try: + return residual_gate_add_cuda(residual, update, gate) + except Exception as exc: + if torch.compiler.is_compiling(): + raise + logger.warning_once(f"Disabling GLM residual-gate CUDA fast path: {exc}") + _GLM_RESIDUAL_GATE_CUDA_DISABLED = True + + return residual + gate * update + class GlmImageLayerKVCache: """KV cache for GlmImage model.""" @@ -278,10 +430,6 @@ class GlmImageAdaLayerNormZero(nn.Module): temb: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: dtype = hidden_states.dtype - norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) - norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to( - dtype=dtype - ) emb, _ = self.linear(temb) ( @@ -299,12 +447,12 @@ class GlmImageAdaLayerNormZero(nn.Module): c_gate_mlp, ) = emb.chunk(12, dim=1) - hidden_states = norm_hidden_states * ( - 1 + scale_msa.unsqueeze(1) - ) + shift_msa.unsqueeze(1) - encoder_hidden_states = norm_encoder_hidden_states * ( - 1 + c_scale_msa.unsqueeze(1) - ) + c_shift_msa.unsqueeze(1) + hidden_states = _glm_ln_modulate( + self.norm, hidden_states, scale_msa, shift_msa, dtype + ) + encoder_hidden_states = _glm_ln_modulate( + self.norm_context, encoder_hidden_states, c_scale_msa, c_shift_msa, dtype + ) return ( hidden_states, @@ -517,10 +665,13 @@ class GlmImageAttention(torch.nn.Module): value = value.unflatten(2, (self.num_local_kv_heads, -1)) # 2. QK normalization - if self.norm_q is not None: - query = self.norm_q(query).to(dtype=dtype) - if self.norm_k is not None: - key = self.norm_k(key).to(dtype=dtype) + if self.norm_q is not None and self.norm_k is not None: + query, key = _glm_qk_layernorm(self.norm_q, self.norm_k, query, key, dtype) + else: + if self.norm_q is not None: + query = self.norm_q(query).to(dtype=dtype) + if self.norm_k is not None: + key = self.norm_k(key).to(dtype=dtype) # 3. Rotational positional embeddings applied to latent stream if image_rotary_emb is not None: @@ -686,9 +837,11 @@ class GlmImageTransformerBlock(nn.Module): ff_output = self.ff(norm_hidden_states) ff_output_context = self.ff(norm_encoder_hidden_states) - hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1) - encoder_hidden_states = ( - encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1) + hidden_states = _glm_residual_gate_add( + hidden_states, ff_output, gate_mlp.unsqueeze(1) + ) + encoder_hidden_states = _glm_residual_gate_add( + encoder_hidden_states, ff_output_context, c_gate_mlp.unsqueeze(1) ) return hidden_states, encoder_hidden_states @@ -779,6 +932,8 @@ class GlmImageAdaLayerNormContinuous(nn.Module): # *** NO SiLU here *** emb = self.linear(conditioning_embedding.to(x.dtype)) scale, shift = torch.chunk(emb, 2, dim=1) + if is_plain_layer_norm(self.norm, x.shape[-1]): + return _glm_ln_modulate(self.norm, x, scale, shift, x.dtype) x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] return x diff --git a/test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py b/test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py new file mode 100644 index 000000000..d2bb151d6 --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py @@ -0,0 +1,55 @@ +"""GLM-Image fused LN+modulate / qk-LN fast paths must stay bit-exact vs eager.""" + +import pytest +import torch + +import sglang.multimodal_gen.runtime.models.dits.glm_image as glm_image +from sglang.multimodal_gen.runtime.models.dits.glm_image import ( + _eager_ln_modulate, + _glm_ln_modulate, + _glm_qk_layernorm, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=4, 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", [(1, 4096, 4096), (2, 301, 4096), (1, 1, 2560)]) +def test_fused_ln_modulate_is_bit_exact(shape): + # (1, 4096, 4096) is the real GLM-Image image-stream shape (1024^2, + # hidden 4096); the others cover the text stream and another hidden. + torch.manual_seed(0) + batch, seq, hidden = shape + norm = torch.nn.LayerNorm(hidden, eps=1e-5, elementwise_affine=False).cuda() + x = (torch.randn(batch, seq, hidden, device="cuda") * 8).bfloat16() + emb = torch.randn(batch, 12 * hidden, device="cuda").bfloat16() + chunks = emb.chunk(12, dim=1) # strided adaLN projection views + shift, scale = chunks[0], chunks[2] + out = _glm_ln_modulate(norm, x, scale, shift, x.dtype) + assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift, x.dtype)) + assert glm_image._GLM_FUSED_LN_MOD_VERIFIED + assert not glm_image._GLM_FUSED_LN_MOD_DISABLED + + +@pytest.mark.parametrize("shape", [(1, 4360, 32, 128), (2, 37, 3, 40), (1, 129, 5, 64)]) +def test_fused_qk_head_layernorm_is_bit_exact(shape): + # (1, 4360, 32, 128) is the real GLM-Image q/k shape (text + image + # tokens, 32 heads of dim 128); the others cover partially-filled warps. + torch.manual_seed(1) + batch, seq, heads, head_dim = shape + norm_q = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda() + norm_k = torch.nn.LayerNorm(head_dim, eps=1e-5, elementwise_affine=False).cuda() + q = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16() + k = (torch.randn(batch, seq, heads, head_dim, device="cuda") * 5).bfloat16() + q_out, k_out = _glm_qk_layernorm(norm_q, norm_k, q, k, q.dtype) + assert torch.equal(q_out, norm_q(q).to(q.dtype)) + assert torch.equal(k_out, norm_k(k).to(k.dtype)) + assert glm_image._GLM_FUSED_QK_LN_VERIFIED + assert not glm_image._GLM_FUSED_QK_LN_DISABLED + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__]))