[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>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Claude Fable 5
Claude Fable 5
parent
148f15b0af
commit
5dffa06fe1
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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__]))
|
||||
Reference in New Issue
Block a user