[diffusion] FLUX.1 fused adaLN modulate (bit-exact) + RoPE cache hoist, LN-affine folding behind quality=high (H200 e2e -3.5% lossless / -6.9% high)

[BBuf]

Branch: k2-flux-modulate-rope, three commits on main (d4be483), standalone — no dependencies. Third wave of the FLUX.1 elementwise cleanup after #33819 (residual-gate fast path + gated GELU epilogue): FLUX.1 runs at ~97.8% GPU busy, so the remaining eager elementwise chains (13.4% of step kernel time) convert directly into wall-clock.

commit	scope	tier
7fd67406a6	RoPE cos/sin cache hoisted to one build per step	lossless (same tensor reused)
95675ae120	new bit-exact fused adaLN modulate kernel, wired unconditionally	lossless (torch.equal)
2eda4a7725	modulate folded into the LN affine (one aten kernel per site)	quality="high" mount

Motivation

Per 50-step 1024x1024 FLUX.1-dev image, the diffusers-style adaLN plumbing runs:

• Modulate chains: norm(x) * (1 + scale) + shift at 114 sites/step (19 dual blocks x 4: norm1/norm1_context/norm2/norm2_context + 38 single blocks x 1) = three eager elementwise kernels each — a tiny 1 + scale row op plus two full HBM passes over the [1, L, 3072] activation (mul 3.09 ms/step + add 3.14 ms/step + row ops 0.15 ms/step in the S1 trace).
• RoPE cache rebuilds: every attention call re-ran torch.cat([cos.float(), sin.float()], -1) on identical inputs — 57 launches/step (19 dual + 38 single), 2,850 per image, all but one producing the same tensor.

Modifications

Commit 1 — RoPE hoist (lossless by construction). FluxTransformer2DModel.forward builds the concatenated fp32 cache once per step (separately for the single-block relayout when SP text sharding is active) and passes the tensor down; FluxAttention keeps accepting the raw (cos, sin) tuple for other callers. Every attention consumes the same tensor values as before.

Commit 2 — bit-exact fused modulate (lossless, torch.equal). New JIT CUDA kernel modulate_scale_shift.cuh (sibling of residual_gate_add.cuh, same vec-8 row-broadcast tile structure) computing round(round(x * round(1+scale)) + shift) — each intermediate rounded back to the storage dtype exactly like the eager chain's kernel boundaries, which also structurally blocks FMA contraction. Because every eager op in the chain is elementwise, matching the per-op rounding is sufficient for bit-exactness (no reduction-order risk). Wired unconditionally via local FluxAdaLayerNormZero{,Single} subclasses (same parameters/state-dict as the diffusers parents) and the dual-block norm2 sites; guarded inputs (non-half dtype, B>1 non-contiguous chunk views, misalignment) fall back to the eager expression, torch.compiler.is_compiling() escape hatch as in #33819. Nunchaku branches (scale without 1+) untouched.

Commit 3 — LN-affine folding behind quality="high". The remaining lossless-path pair per site is the affine-free LayerNorm kernel + the fused modulate (two HBM passes). For high requests the modulate folds into the LN's elementwise affine: layer_norm(x, weight=1+scale, bias=shift) — one aten kernel, one pass. 1 + scale keeps the eager rounding of the [1, D] row; the remaining difference is that scale/shift apply in fp32 to the unrounded normalized value (bf16 rounding-order level, max abs diff ~0.05 on FLUX shapes) — close, but not bit-exact, hence mounted per batch by the DenoisingStage exactly like the fused linear+GELU epilogue (#33819): sites default off, lossless requests keep the bit-exact reference path (mount log: Mounted fused LN+modulate (affine folding) for quality=high). Per-call guard requires B == 1 (the folded affine is a [D] row).

Why the LN reduction itself is not fused on the lossless path

#33819 already audited the shared norm+scale/shift fusions for FLUX and dropped them (norm_scale_shift_native is Blackwell-gated; the Triton select01 layers are fp32 single-pass, not bit-exact). A T10-style (#33734 ERNIE) bit-exact replication was assessed against torch 2.11's vectorized_layer_norm_kernel: unlike the CuTe-DSL RMSNorm T10 replicated (documented ordered fadd chain + butterfly fold), aten's LayerNorm uses per-element Welford updates (mean + delta * (1.f/new_count)), count-weighted cuWelfordCombine merges, and the accurate multi-instruction rsqrtf — with the effective FMA-contraction choices baked into shipped libtorch SASS. Reproducing that bit-for-bit in Triton is a full task with real failure risk, so the LN fusion ships as the gated affine-folding instead (the fallback the plan anticipated), and the lossless tier keeps aten's LN kernel untouched.

Speed Tests (H200)

FLUX.1-dev (sha256-verified unsloth/FLUX.1-dev mirror), 1024x1024, 50 steps, bf16, embedded guidance (single forward, B=1), fixed prompt, seed=42. One warmed process, 10 sequential generations, first dropped, median of 9 (#33451/#33536/#33819 protocol), single pinned H200.

Arm	E2E wall (ms)	DenoisingStage (ms)
main @ d4be483, lossless	7442.9	6809.5
+ commit 1 (RoPE hoist)	7396.0 (-0.6%)	6789.8 (-0.3%)
this PR, lossless (default)	7181.9 (-3.5%)	6607.3 (-3.0%)
main, quality=high (GELU epilogue)	7218.7	6648.4
this PR, quality=high	6926.1	6405.5

• Gated tier vs its own baseline (both at quality=high): denoise -3.7%, wall -4.1%.
• Combined vs main lossless: denoise -5.9%, wall -6.9%.

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

Op	shape	eager	fused	speedup
modulate chain	(1, 4096, 3072), rows (1, 3072)	68.5 us	15.9 us	4.32x
modulate chain	(1, 512, 3072)	25.4 us	13.2 us	1.92x
modulate chain	(1, 4608, 3072)	76.8 us	17.8 us	4.33x
LN + fused modulate vs folded LN	(1, 4096, 3072)	38.3 us	27.3 us	1.40x
LN + fused modulate vs folded LN	(1, 512, 3072)	23.3 us	20.1 us	1.16x
RoPE cos/sin cat	(4608, 128) fp32 x2	7.5 us/call	1 call/step instead of 57	—

Summed over sites: RoPE 0.43 ms/step vs measured 0.39 ms/step commit-1 delta; modulate ~4.7 ms/step (micro) vs 3.65 ms/step measured; LN folding ~1.0 ms/step (micro) vs ~0.8 ms/step measured (high-arm delta net of the lossless wins). Microbenchmarks run standalone slightly above the in-graph deltas, as expected for launch-dense eager chains.

Accuracy Tests

Check	Result
kernel unit tests (test_modulate_scale_shift.py, 12 cases: FLUX shapes, bf16/fp16, batched rows, adaLN chunk views)	torch.equal vs eager chain, all pass
module A/B: FluxAdaLayerNormZero{,Single} vs diffusers parents, L in {4096, 512, 4608}	bit-exact (torch.equal, all outputs)
e2e md5, lossless: main x10 / commit-1 x10 / this PR x10 / final HEAD x2 iter images	all 32 images one identical md5 (cad50fb5...)
test_fused_ln_modulate.py (folded-LN closeness + mount protocol)	pass (max abs diff 0.047 at bf16, rounding-order level)
adjacent suites (test_transformer_quant, test_fused_linear_gelu, test_residual_gate_add)	45 passed
pre-commit (pinned) on all changed files, per commit	clean

quality=high image quality (vs the bit-exact lossless reference, same seed; #33819 table protocol):

Prompt	PSNR	SSIM
cyberpunk city (bench prompt)	35.48 dB	0.9785
mountain lake sunrise	36.40 dB	0.9873
elderly craftsman portrait	37.95 dB	0.9916

Context: the already-merged GELU-epilogue high tier alone measures 34.67 dB / 0.9765 on the bench prompt against the same reference — the LN folding does not degrade the high tier's quality band (this PR high vs main high: 35.64 dB / 0.9786). Sample images: /data/solmig/k2_runs/{main,mainhigh,k2high,mainref}/.

Benchmark configuration

• Machine: ion8-h200 (H200, container sglang-bbuf), CUDA_VISIBLE_DEVICES=2, torch 2.11.0+cu130.
• Model: unsloth/FLUX.1-dev from the shared HF cache (byte-identical mirror of the gated black-forest-labs/FLUX.1-dev).
• Driver: /data/solmig/prA_work/pra_e2e.py (the [diffusion] FLUX.1 bit-exact residual-gate fast path + tanh-GELU epilogue behind quality=high (H200 e2e -1.1% lossless / -4.3% high) #33819 harness) — DiffGenerator.from_pretrained(model_path, num_gpus=1, local_mode=True), prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets", seed=42, width=height=1024, num_inference_steps=50, quality={lossless|high}, 10 iterations, per-iter PNGs for md5/PSNR.
• Arms/worktrees: main /data/solmig/k2_main_tree (d4be483), commit-1 /data/solmig/k2_rope_tree (7fd67406a6), PR /data/solmig/k2_tree (HEAD). Raw logs and records JSON: /data/solmig/k2_runs/{main,rope,k2,mainhigh,k2high,mainref,k2final}*.
• DenoisingStage times parsed from the scheduler logs ([DenoisingStage] finished in X seconds), same drop-first/median-of-9 policy as wall.

Side-by-side sample (lossless left — md5-identical to main — vs quality=high right):

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ⏳ Run #31235036039
Latest PR Test (Extra): ✅ Run #31236056496

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/31235036039/job/93045930915?pr=34004