[diffusion] GLM-Image bit-exact fused aten LayerNorm+modulate / qk-LN (H200 30-step denoise -8.1%) (#34008)

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

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<c10::BFloat16, float, false>``,
+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<bf16, 4> 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

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 * (
+        hidden_states = _glm_ln_modulate(
-            1 + scale_msa.unsqueeze(1)
+            self.norm, hidden_states, scale_msa, shift_msa, dtype
-        ) + shift_msa.unsqueeze(1)
+        )
-        encoder_hidden_states = norm_encoder_hidden_states * (
+        encoder_hidden_states = _glm_ln_modulate(
-            1 + c_scale_msa.unsqueeze(1)
+            self.norm_context, encoder_hidden_states, c_scale_msa, c_shift_msa, dtype
-        ) + c_shift_msa.unsqueeze(1)
+        )
 
         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:
+        if self.norm_q is not None and self.norm_k is not None:
-            query = self.norm_q(query).to(dtype=dtype)
+            query, key = _glm_qk_layernorm(self.norm_q, self.norm_k, query, key, dtype)
-        if self.norm_k is not None:
+        else:
-            key = self.norm_k(key).to(dtype=dtype)
+            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)
+        hidden_states = _glm_residual_gate_add(
-        encoder_hidden_states = (
+            hidden_states, ff_output, gate_mlp.unsqueeze(1)
-            encoder_hidden_states + ff_output_context * c_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
 

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__]))