[diffusion] FLUX.1 bit-exact residual-gate fast path + tanh-GELU epilogue behind quality=high (H200 e2e -1.1% lossless / -4.3% high)

[BBuf]

Motivation

FLUX.1's transformer updates its residual streams through eager two-kernel gate * update + residual + ... pairs, and runs every FFN tanh-GELU as a separate bandwidth-bound kernel after the up-projection GEMM:

• Residual gates: 19 dual-stream blocks x 4 sites (img attn/mlp + txt attn/mlp) + 38 single-stream blocks x 1 site = 114 sites/step x 50 steps = 5,700 eager pairs per 1024x1024 image.
• tanh-GELU: 19 dual-stream blocks x 2 FFNs + 38 single-stream blocks x 1 MLP = 76 sites/step x 50 steps = 3,800 standalone GELU kernels per image, each over a [tokens, 12288] intermediate.

This PR wires two existing, already-shipped mechanisms into FLUX.1 (no new kernels, no protocol changes):

1. residual_gate_add_cuda (in-tree since [KDA-Pilot] Add diffusion residual-gate CUDA fast path for LTX2 #29361, used by LTX-2 and the ERNIE-Image PR [diffusion] ERNIE-Image bit-exact residual-gate fast path (H200 1024^2 e2e 16.17 -> 15.75 s) #33734 on their default paths): collapses each gate pair into one launch and, for half dtypes, reproduces the eager pair's rounding bit-for-bit — so it runs on the lossless default path with no quality gate.
2. cublasLt GELU epilogue ([diffusion] Fuse DiT FFN tanh-GELU into up-proj GEMM (cublasLt epilogue) behind quality=high (Qwen-Image 1024^2 denoise 12.36 -> 12.05 s on H200) #33536's fused_linear_gelu sites + DenoisingStage mount/unmount protocol, shipped for Qwen-Image/GLM-Image): folds the tanh-GELU into the up-proj GEMM epilogue. Not bit-exact (bf16 rounding-order level), so it is mounted only for quality="high" requests and the default path stays byte-identical.

History: why the GELU fusion is gated

An ungated FLUX GELU-epilogue fusion was merged once before (#28166, "[Diffusion] FLUX: fuse FeedForward GELU into up-proj GEMM (cublasLt epilogue)") and reverted in #28708 because the unconditional numerical deviation broke reference-model comparisons. The quality tier system (#33453) is exactly the mechanism that resolves that objection: the deviation is now opt-in per request, the default path is provably byte-identical to main (md5 below), and the mount/unmount protocol is the same one already shipped and validated for Qwen-Image and GLM-Image in #33536.

Modifications

All changes are FLUX-local (flux.py + two extended test files); the kernels and the mount protocol are untouched. Two commits: lossless residual-gate first, gated GELU second.

• _flux_residual_gate_add (same helper shape as LTX-2/[diffusion] ERNIE-Image bit-exact residual-gate fast path (H200 1024^2 e2e 16.17 -> 15.75 s) #33734: one-time runtime fallback, torch.compiler.is_compiling() escape hatch so compiled graphs never bake in a silent fallback, fp16/bf16-only where the kernel is bit-exact) wired into all 5 gate sites. The kernel's row-broadcast gate only covers [1, ..., 1, D], so batched [B>1, 1, D] gates fail the guard and fall back to the eager pair automatically (FLUX.1-dev uses embedded guidance, so serving batches are B=1 per forward and the gates are already dense [1, 1, D] slices — no extra copies needed).
• mark_fused_gelu_site on the three FLUX GELU up-projection layouts:
  • FluxGELU (TP>1 dual-stream FluxParallelFeedForward),
  • new FluxFusedGELUProj wrapper re-homing the shared diffusers-style FeedForward.net[0] (TP=1 dual-stream); it keeps the net.0.proj parameter path, so no checkpoint-mapping changes,
  • FluxSingleTransformerBlock.proj_mlp (single-stream, non-nunchaku branch).
    Sites start unmounted (bit-exact reference); the shared DenoisingStage mounts them per batch for quality="high" only (log-confirmed: Mounted fused linear+GELU (cublasLt epilogue) for quality=high). Nunchaku/quantized paths are never marked, and the mount's all-or-nothing static guard also rejects them.

What was audited and dropped

The adaptation matrix also flagged FLUX's norm2/norm2_context + scale/shift and the AdaLayerNormZero internals for the shared norm+scale/shift fusion. Dropped from this PR: the only bit-exact-capable shared implementation (norm_scale_shift_native) is hard-gated to Blackwell (capability >= 10, tuned for that launch geometry) while this validation lane is H200/SM90, and the Triton select01 layers are neither bit-exact (fp32 single-pass LayerNorm reduction vs eager) nor semantics-matched (dual-branch index selection). Per the matrix's own standard this was a rider item: it does not justify introducing gate complexity it cannot clear, so it is cut rather than gated (same disposition as the ERNIE PR's norm-fusion re-validation).

Speed Tests (H200)

FLUX.1-dev, 1024x1024, 50 steps, bf16, fixed prompt and seed=42. Client wall-clock in one warmed process, 10 consecutive generations, first dropped, median of 9 (#33451/#33536/#33734 protocol), single H200, CUDA_VISIBLE_DEVICES pinned. Weights: sha256-verified byte-identical mirror of black-forest-labs/FLUX.1-dev (gated repo; every LFS object checked against the official manifest).

Arm	E2E wall (s)	DenoisingStage (s)
main @ 604d356, lossless	7.481	6.802
this PR, lossless (default)	7.399 (-1.1%)	6.674 (-1.9%)
this PR, quality=high	7.158 (-4.3%)	6.524 (-4.1%)

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

Op	shape	eager	fused	speedup
residual gate	(1, 4096, 3072), gate (1,1,3072)	53.3 us	20.7 us	2.58x
residual gate	(1, 512, 3072), gate (1,1,3072)	12.9 us	13.2 us	0.98x
residual gate	(1, 4608, 3072), gate (1,1,3072)	59.4 us	23.1 us	2.57x
GELU up-proj	(1, 4096, 3072) -> 12288	491.1 us	455.5 us	1.08x
GELU up-proj	(1, 512, 3072) -> 12288	68.2 us	58.6 us	1.16x
GELU up-proj	(1, 4608, 3072) -> 12288	568.6 us	524.2 us	1.08x

Summed over sites: gates 2.6 ms/step (130 ms/image) vs measured 128 ms lossless denoise delta; GELU 2.5 ms/step (127 ms/image) vs measured 150 ms high-vs-lossless denoise delta. Sanity anchor: ERNIE-Image (#33734) got -2.6% e2e from 7,200 gate sites at hidden 4096 with CFG; FLUX's -1.1% lossless from 5,700 sites at hidden 3072 without CFG is proportionally consistent.

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 sites step aside via torch.compiler.is_compiling() and Inductor keeps its own fusion of these chains.

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="unsloth/FLUX.1-dev"  # sha256-verified mirror of the gated official repo, 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,
        num_inference_steps=50,
        quality="lossless",  # or "high" for the gated GELU-epilogue arm
    ))

Accuracy Tests

Check	Result
Kernel vs eager pair, torch.equal, real FLUX shapes (4096/512/4608 x 3072), bf16/fp16	True (atol=0, rtol=0)
GELU sites unmounted vs reference, torch.equal	True
Full 1024x1024 image, same prompt/seed, this PR lossless vs main	md5 identical (cad50fb5beddfda2a7898a595e5e303c on both, all 10 iterations)
Determinism within each arm (incl. quality=high)	md5 identical across iterations (unique=1)

quality=high vs lossless, same seed, full-image (3 prompts, first/middle/last sample images exported):

Prompt	PSNR (dB)	SSIM
cyberpunk city (main bench prompt)	34.68	0.9765
mountain lake sunrise	40.08	0.9923
craftsman portrait	35.67	0.9870

All comfortably above the quality="high" bar (PSNR > 25).

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

cyberpunk city

mountain lake sunrise

craftsman portrait

Unit tests (extended in place, base-b-kernel-unit):

• test_residual_gate_add.py: + real FLUX.1 1024^2 shapes (dual img/txt, single joint), bit-exact.
• test_fused_linear_gelu.py: + FluxFusedGELUProj site (gate off torch.equal vs reference, mounted assert_close, unmount restores bit-exactness).

19 passed

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

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #31073539937
Latest PR Test (Extra): ❌ Run #31073552155

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/31073539937/job/92530316814?pr=33819