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

[BBuf]

Motivation

GLM-Image's transformer block (hidden 4096, 30 blocks, 32 heads x head_dim
128, 30 steps, true CFG = 2 forwards/step) still runs its adaLN chain in
eager pieces — 242 bare aten layer_norm launches per step plus a
hand-written modulate/gate chain around them:

• GlmImageAdaLayerNormZero runs bare nn.LayerNorm (no affine) plus a
  hand-written norm * (1 + scale) + shift for both the image and text
  streams — 4 aten kernels per stream per block;
• the qk-norm is nn.LayerNorm(128) per head on (1, S, 32, 128): aten's
  vectorized_layer_norm_kernel spends a full 128-thread block per
  128-element row (warp 0 does 4 elements per lane, warps 1-3 idle), and
  the two launches cost 301 us per block at the real shape — 10.1% of
  denoise GPU time in the profiler trace;
• the final GlmImageAdaLayerNormContinuous repeats the LN+modulate
  pattern once per forward;
• the ff gate hidden + ff_out * gate is another eager pair per stream.

This PR fuses each site into one Triton kernel that is numerically
re-engineered to be rounding-faithful to the eager chain, step for step,
so it is bit-exact (torch.equal, full-image md5-identical) and wired on
the default lossless path with no quality gate — the standard #33734 /
#33854 / #33886 used.

What "rounding-faithful" means here

nn.LayerNorm on bf16 dispatches to aten's
vectorized_layer_norm_kernel<BFloat16, float>. The Triton kernels
replicate that kernel's exact arithmetic DAG, verified instruction by
instruction against the sm_90 SASS of torch 2.11's libtorch_cuda.so:

• the (32, 4) block's per-thread serial Welford accumulation over
  aligned_vector<bf16, 4> loads (thread t owns vectors t, t+128, ...),
  with both FMA contractions nvcc makes (mean' = FFMA(delta, recip, mean),
  sigma2' = FFMA(delta, val - mean', sigma2)) and its guarded
  4-instruction reciprocal fast path (MUFU.RCP; FFMA(x,r,-1); FADD.FTZ(-e,-0); FFMA(r,e,r)) reproduced via inline PTX;
• the SHFL.DOWN 16/8/4/2/1 fold tree through cuWelfordCombine
  (count-weighted merge, zero-count branch included) expressed with
  order-exact tl.reshape/tl.permute/tl.split folds, then the
  (0,2)/(1,3)/(0,1) inter-warp combines;
• sigma2 / float(N) as an IEEE correctly-rounded div.rn.f32, and
  rstd = rsqrtf(...) as rsqrt.approx.f32 with the 2^24/2^12 subnormal
  rescale branch;
• the output cvt.rn.bf16(FMUL(rstd, FADD(x, -mean))) (no FMA), then the
  aten modulate chain's per-op bf16 boundaries (round(1+scale),
  round(y*(1+scale)), round(+shift)) via the bitcast RNE helper from
  [diffusion] ERNIE-Image bit-exact fused RMSNorm+scale/shift (H200 1024^2 e2e 15.63 -> 15.00 s, denoise -3.3%) #33854;
• the qk kernel replicates the same DAG for head_dim <= 128 rows, where
  only vectors 0..head_dim/4-1 of the 128 aten threads carry data (the
  remaining lanes and warps 1..3 enter the fold with count 0), and handles
  q and k in one launch, 64 rows per program.

Because bit-exactness is a property of which LayerNorm implementation aten
dispatches to, the model-side wrappers verify torch.equal against the
eager chain once at runtime on the first real input and permanently fall
back to eager on any mismatch. The fast path can never silently change
outputs.

The ff gate reuses residual_gate_add_cuda (bit-exact by construction,
#33734) behind the same one-time-disable wrapper.

Modifications

• python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py
  (new): the two Triton kernels + guards, registered as custom ops
  (torch.compile-safe fake impls).
• glm_image.py: _glm_ln_modulate / _glm_qk_layernorm /
  _glm_residual_gate_add wrappers (first-call self-verification, one-time
  disable on exception/mismatch, [diffusion] ERNIE-Image bit-exact fused RMSNorm+scale/shift (H200 1024^2 e2e 15.63 -> 15.00 s, denoise -3.3%) #33854 style) wired at
  GlmImageAdaLayerNormZero (both streams),
  GlmImageAdaLayerNormContinuous, the attention qk-norm, and both ff-gate
  sites.
• test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py
  (new): bit-exactness at the real GLM-Image shapes plus an assertion that
  the fast paths were actually taken (not silently disabled).

Speed Tests (H200)

zai-org/GLM-Image, 1024x1024, default 30-step preset (true CFG, 2
forwards/step), bf16, eager, fixed prompt, seed=42. One warmed process,
10 sequential requests, first dropped, median of 9, GPU 0.

arm	DenoisingStage (s)	s/step	GlmImageAR (s)	e2e wall (s)	e2e minus AR (s)
main (d4be483), run 1	6.385	0.2127	32.09	38.785	6.70
this PR	5.870	0.1956	35.37	41.559	6.19
main, run 2 (A/B/A)	6.391	0.2129	33.24	39.952	6.71

DenoisingStage -8.1% (pooled main 6.388 -> 5.870 s), bit-exact
(md5-identical) output.

E2e caveat: ~83% of e2e is the GlmImageAR autoregressive prior-token
stage (a separate GlmImageForConditionalGeneration model this PR does not
touch), which drifted 32.1 -> 35.4 -> 33.2 s across the three arms with
host load (other tenants on the box) — the A/B/A bracket shows the drift is
environmental, not tree-dependent (main's own e2e moved +1.2 s between its
two runs while its denoise stayed within 6 ms). Net of the AR stage, e2e
drops 6.70/6.71 -> 6.19 s (-7.7%, i.e. -0.52 s/request, -1.3% of total
e2e).

Kernel microbenchmark at the real shapes (triton.testing.do_bench):

Site	eager	fused	speedup
LN + norm*(1+scale)+shift (1, 4096, 4096) bf16 (aten LN + 3 aten elementwise)	121.8 us	38.5 us	3.16x
qk nn.LayerNorm(128) (1, 4360, 32, 128) bf16, q and k (2 aten launches)	301.2 us	62 us	4.86x

Attribution (built-in --profile, same 5-step denoise window on both
trees, kernel self-time):

kernel family	main	this PR
aten vectorized_layer_norm_kernel (norm1 x2, qk x2, norm_out)	141.3 ms / 10.8% (n=1517)	0
aten modulate/gate elementwise (mul + adds)	64.8 ms / 4.9%	0
_qk_ln_head_kernel (new)	-	24.1 ms / 2.0%
_ln_modulate_kernel (new)	-	17.6 ms / 1.5%
residual_gate_add (reused #33734 kernel)	-	9.6 ms / 0.8%
whole window	1313.4 ms	1202.5 ms (-8.4%)

Remaining window is GEMMs (72.8%) and FlashAttention (13.6%), consistent
with the -8.1% DenoisingStage wall-clock delta.

Benchmark command

Both arms run the identical protocol (this PR is lossless-only; arms = main
tree @ d4be483 vs PR tree), default eager configuration
(performance_mode=auto, enable_torch_compile=false):

from sglang.multimodal_gen import DiffGenerator

gen = DiffGenerator.from_pretrained(
    model_path="zai-org/GLM-Image", num_gpus=1, local_mode=True
)
for i in range(10):
    gen.generate(sampling_params_kwargs=dict(
        prompt=PROMPT, seed=42, width=1024, height=1024,  # default 30-step preset
    ))

Accuracy Tests

Check	Result
Both kernels vs eager chain, torch.equal, real shapes (1,4096,4096) LN+mod, (1,4360,32,128) qk, plus text-stream/odd shapes	True (registered unit test, 6 passed)
Full 1024x1024 images, 2 prompts x 2 seeds, this PR vs main	md5 identical in all 4 pairs (f86659da..., 91882448..., 7a2f3b00..., ac01683d...)
Runtime self-verification	wired; first call per site verifies torch.equal vs eager and permanently disables the fast path on mismatch
Fallback warnings in the 14-request PR arm	none

Unit test (registered, base-b-kernel-unit):

6 passed

pre-commit run --files ...: all hooks pass.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ⏳ Run #31235033025
Latest PR Test (Extra): ✅ Run #31235968414