[diffusion] Ideogram 4: fuse RMSNorm modulate/gate chains via the Z-Image Triton suite behind quality=high (H200 e2e -2.9%/-3.4%)

[BBuf]

Reuses the Z-Image bf16-native norm kernels already in-tree (kernels/ops/diffusion/triton/zimage_native_norm.py) — no new kernels; adds the mount/unmount site protocol for them (same shape as #33536's fused_linear_gelu). 4 files, +283/−28.

Motivation

The adaptation-matrix audit (main @ 407a65d) flagged Ideogram 4's transformer block as form-identical to Z-Image's modulate/gate structure, for which a fused Triton suite already ships on Z-Image's default path:

• Gate sites: x + tanh(gate) * RMSNorm(out) — 2 sites/block x 34 blocks = 68 sites/step; eager = tanh + RMSNorm + mul + add over the [tokens, 4608] stream.
• Modulate sites (rider, covered by the same suite): RMSNorm(x) * (1 + scale) — 2 sites/block x 34 blocks; eager = RMSNorm + add + mul.

At V4_FAST_20 (guidance 1.0, single forward per step) that is 1,360 gate + 1,360 modulate eager chains per 1024x1024 image; V4_INSTANT_8 has 544 + 544. zimage_rmsnorm_tanh_residual / zimage_rmsnorm_scale fuse each chain into one kernel.

Why quality=high and not the lossless default

Z-Image mounts these kernels unconditionally because they reproduce its own native-bf16 reference RMSNorm by design. Ideogram's reference norm is F.rms_norm (fp32 internal statistics), so on Ideogram the fused path is close but not bit-exact — probed at real shapes ([1, 4096..4736, 4608], bf16): max abs diff 0.0625, i.e. bf16 rounding of the norm statistics. Following the #33536 precedent (and the same disposition as the FLUX.1 GELU gate in #33819), the sites are mounted only for quality="high" requests; the "lossless" default runs the unmodified reference chain byte-for-byte (md5 proof below).

Modifications

• kernels/ops/diffusion/fused_gate_rmsnorm.py (new) — mount/unmount site protocol for the existing Z-Image Triton kernels, mirroring fused_linear_gelu: mark_fused_gate_rmsnorm_site(module, norm_attrs), all-or-nothing mount_fused_gate_rmsnorm with static guards (bf16 1-D norm weight, hidden <= 8192 kernel limit, Triton available), unmount_... resets to the bit-exact reference. Thin fused_rmsnorm_scale / fused_rmsnorm_tanh_residual wrappers re-export the kernels, which keep their own runtime shape/stride/dtype guards and return None to fall back per call.
• runtime/models/dits/ideogram.py — the four per-block chains route through _norm_scale / _gate_residual helpers: fused when the site is mounted, torch.compiler.is_compiling() escape hatch (compiled graphs never bake in the Triton path), eager fallback otherwise. The reference expressions are unchanged from main (verified torch.equal at operator level and md5-identical e2e). The kernels consume the raw adaln .chunk() slices directly (strided-gate support is native), so no extra copies.
• runtime/pipelines_core/stages/denoising.py — _maybe_toggle_fused_gelu generalized to _maybe_toggle_quality_fusions: the same batch-boundary toggle now mounts/unmounts both quality=high fusion families (cublasLt linear+GELU and gate-RMSNorm). Models without marked sites are no-ops, exactly as before; log-confirmed on Ideogram: Mounted fused gate RMSNorm (Z-Image Triton suite) for quality=high.

Speed Tests (H200)

cocktailpeanut/ideogram-4-fp8 — byte-layout-identical ungated mirror of ideogram-ai/ideogram-4-fp8 (official repo and both fal/ideogram-v4-* distilled repos are HF-gated) — 1024x1024, fixed prompt, seed=42, guidance 1.0 presets (single DiT forward per step). Client wall-clock in one warmed process, 10 consecutive generations, first dropped, median of 9 (#33451/#33536/#33734 protocol), single H200 (CUDA_VISIBLE_DEVICES=7).

Preset	Arm	E2E wall (s)	DenoisingStage (ms)
V4_FAST_20	this PR, lossless (default)	8.482	8259
V4_FAST_20	this PR, quality=high	8.239 (-2.9%)	7977 (-3.4%)
V4_INSTANT_8	this PR, lossless (default)	3.488	3243
V4_INSTANT_8	this PR, quality=high	3.371 (-3.4%)	3118 (-3.9%)

main @ 604d356 single-shot sanity anchors agree with the lossless arm within noise (FAST_20 8.548 s / denoise 8219 ms; INSTANT_8 3.498 s / denoise 3224 ms), as they must — the lossless path is byte-identical (below).

Attribution closes against kernel microbenchmarks at the real shapes (bf16, 200 iters after warmup, H200):

Op	shape	eager	fused	speedup
gate+residual	(1, 4096, 4608), gate (1,1,4608)	107.5 us	34.5 us	3.12x
norm+modulate	(1, 4096, 4608), scale (1,1,4608)	77.7 us	39.2 us	1.98x
gate+residual	(1, 4736, 4608)	123.6 us	38.0 us	3.26x
norm+modulate	(1, 4736, 4608)	89.7 us	38.7 us	2.32x

Summed over sites: ~223 us/block x 34 blocks ≈ 7.6 ms/step ≈ 152 ms per FAST_20 image from the four chains alone; measured denoise delta is 282 ms (the fused calls also drop the standalone per-block tanh launches and their sync/launch overhead at 2.4 it/s). Per-step saving is preset-independent (~14 ms/step), so the few-step INSTANT preset shows the larger relative win.

Benchmark configuration

All arms above run the default eager configuration (performance_mode=auto; every archived server log shows enable_torch_compile: false, regional_compile: false) — i.e. the default production path for this model. Runs with --enable-torch-compile are unaffected: the fused helpers step aside via torch.compiler.is_compiling() and the reference chain runs under Inductor.

Benchmark driver (local mode, one warmed process; 10 sequential requests, first dropped, median of 9):

from sglang.multimodal_gen import DiffGenerator

gen = DiffGenerator.from_pretrained(
    model_path="cocktailpeanut/ideogram-4-fp8"  # byte-layout-identical ungated mirror, 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,
        preset="V4_FAST_20",  # or "V4_INSTANT_8"
        quality="lossless",  # or "high" to mount the fused gate-RMSNorm sites
    ))

Accuracy Tests

Check	Result
Reference (unmounted) helpers vs main expressions, torch.equal, real shapes	True
Full 1024x1024 image, same prompt/seed, this PR lossless vs main @ 604d356	md5 identical, both presets x both prompts (e.g. a742d0b0... FAST_20 p0 on both)
quality=high determinism	stable across repeated runs

quality=high vs lossless, same seed, full-image:

Preset / prompt	PSNR (dB)	SSIM
FAST_20 / fox photo	37.91	0.9981
FAST_20 / typography poster	58.07	0.9996
INSTANT_8 / fox photo	41.13	0.9989
INSTANT_8 / typography poster	58.00	0.9996

All comfortably above the quality="high" bar (PSNR > 25); the typography prompt — Ideogram's signature use case — is essentially indistinguishable.

Side-by-side samples (lossless left, quality=high right, FAST_20):

fox photo

typography poster

Unit tests (new test_fused_gate_rmsnorm.py, base-b-kernel-unit, mirrors test_fused_linear_gelu.py): fused vs F.rms_norm reference at the Ideogram hidden size incl. adaln-strided chunk gates; mount all-or-nothing guard (one fp32 site rejects the whole model); unmount restores the reference. Existing test_zimage_native_norm.py untouched and passing (8/8 across both files on H200).

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #31075920352
Latest PR Test (Extra): ❌ Run #31075922438

[BBuf]

/rerun-failed-ci

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/31075920352/job/92587075296?pr=33822