[diffusion] FLUX.2 VAE decoder fast path behind quality=high (H200: 1024^2 97.6->29.2 ms, 2048^2 437.2->168.5 ms)

[BBuf]

Branch: pr-flux2-vae-fastpath, single commit eada083612 on main (7ba393d, post-#33453), standalone — no dependencies. 5 files, +920/−28.

Motivation

FLUX.2 image generation spends 97.6 ms (1024^2) / 437.2 ms (2048^2) per image in VAE decode on H200. Profiling the decoder (diffusers Decoder) shows two structural problems:

• the GroupNorm -> SiLU chains upcast to fp32 under autocast and are ~46% of decode time — and they get worse if you naively switch to channels_last, because aten GroupNorm is ~2x slower on NHWC and the existing triton_group_norm_silu kernel forces a .contiguous() (NCHW) round-trip on channels_last inputs;
• the decoder runs NCHW, so every cuDNN conv pays nchwToNhwc/nhwcToNchw transpose kernels, and the D512 mid-block runs a single-head attention whose diffusers AttnProcessor2_0 forward is layout-unsafe (calls .view on the 4D activation).

This PR adds (a) a layout-aware two-pass GroupNorm(+SiLU) Triton kernel and (b) a request-scoped, quality-gated FLUX.2 VAE decoder fast path built on it. All decoder rewrites are mathematically exact re-associations of the original operators (weight folding is done lazily in fp32 on first fast-path use).

Modifications

1. Kernel — new python/sglang/kernels/ops/diffusion/triton/group_norm_silu_twopass.py (+232):

• Channels-last only. Pass 1: per-(batch, channel) partial sum / sum-of-squares reduction over (N, L, C) rows with C innermost (serves channels_last 4D via a free permute view). Finalize: per-(batch, group) fp32 mean/rstd with the affine transform folded into per-(batch, channel) scale/shift. Pass 2: fused elementwise apply + optional SiLU.
• Public API group_norm_silu_4d / group_norm_silu_rows returns None for unsupported inputs (non-CUDA, grad enabled, non-pow2 C, C > 2048, not channels_last); callers must fall back.
• Coexists with the existing general group_norm_silu.py (any channels-per-group, always-SiLU, NCHW via .contiguous()), which keeps serving the HunyuanVAE / latent-upsampler paths unchanged; on NCHW the two are a wash (63–219 us vs 81–210 us), so this kernel deliberately has no NCHW variant — gate-off decodes run the original aten path, never this kernel. Documented in the module docstring.

2. FLUX.2 fast paths + gate — runtime/models/vaes/flux2_vae_cuda_opt.py (+421):

• CudaPlatform.optimize_vae (existing platform hook called by the VAE loader, +22 in platforms/cuda.py) installs always-present fast-path wrappers on the FLUX.2 decoder; each wrapper's forward dispatches on a shared per-VAE VaeFastPathGate (a module-local bool holder published on the VAE as _sgl_vae_fast_path_gate).
• Gate semantics: quality == "high" ([diffusion] Restrict request-level quality to two validated tiers: lossless (default) and high #33453 tier) runs the fast paths; the "lossless" default runs the original module path bit-for-bit (verified byte-identical below), so golden-output CI is unaffected.
• Install is all-or-nothing and fail-closed: without Triton, or with any attention block lacking the layout-safe rewrite, nothing is installed and every request runs the unmodified decoder.
• channels_last decoding: parameter layout swapped at decode entry to match the gate (layout swaps are pure permutations), so lossless decodes always run the baseline NCHW kernels bit-for-bit.
• GroupNorm+SiLU fusion via the new two-pass kernel (ResnetBlock2D norm1/norm2 + SiLU chains and the conv_norm_out/conv_act tail).
• Fused nearest-2x upsample + Conv2d(3x3, p1) == ConvTranspose2d(k4, s2, p1) with a lazily-summed kernel (nearest upsampling is pure pixel replication, so the fusion only re-associates conv taps).
• Layout-safe single-head attention forward with the output projection folded into V (softmax rows sum to 1, so A @ (V W_v^T + b_v) W_o^T + b_o == A @ (V W'^T) + b'). Attention stays on the stock SDPA backend here; the exact chunked D512 composition for >= 16k tokens is a follow-up PR (recovers 168.1 -> ~147 ms at 2048^2).

3. DecodingStage wiring — pipelines_core/stages/decoding.py (+39/−27, mostly the try/finally re-indent):

• DecodingStage.forward sets the gate for the duration of a decode iff the request's quality sampling param is "high" and always restores it in a finally; VAEs without installed wrappers have no gate attribute (no-op).

Follow-up PRs: (1) chunked D512 attention composition; (2) Qwen-Image VAE fast paths.

Speedup

All numbers H200, bf16 decode. Three levels:

Kernel microbenchmark (FLUX.2 VAE decoder norm shapes, G=32, us per call, mean of 50; before = aten fp32 GroupNorm+SiLU on channels_last, after = new kernel on channels_last):

shape	aten CL	existing kernel fed CL	new CL	speedup vs aten / vs existing
(1, 512, 128, 128)	365.9	93.0	76.0	4.8x / 1.2x
(1, 512, 256, 256)	1394.6	284.0	121.0	11.5x / 2.3x
(1, 512, 512, 512)	5863.9	1430.5	311.1	18.9x / 4.6x
(1, 256, 1024, 1024)	16310.2	7544.4	913.8	17.8x / 8.3x
(1, 128, 1024, 1024)	7965.2	3698.7	307.9	25.9x / 12.0x

(HunyuanVAE-style 5D shapes are untouched — no call site of the existing kernel changes.)

Component harness (frozen decode-input latents, steady-state mean of 5 decodes, this branch):

workload	lossless (= main)	quality=high (this PR)	speedup
flux2 1024^2	97.6 ms	29.3 ms	3.33x
flux2 2048^2	437.2 ms	168.1 ms	2.60x

Lossless-path timing on this branch is unregressed (97.1 / 437.0 ms vs main baseline 97.6 / 437.2 ms); the minimized branch reproduces the fast path at 29.2 / 168.5 ms (run-to-run noise of the numbers above).

End-to-end latency (sglang generate, FLUX.2-klein-4B, H200, prompt "A futuristic cyberpunk city at night, neon lights reflecting on wet streets", seed 42, --dit-layerwise-offload false --warmup --save-output). Metric = per-request wall-clock measured client-side around the full request round-trip ("Pixel data generated successfully in X seconds"), which includes the VAE decode GPU time, the final sync, and the PNG write. Per cell: one warmed process serving sequential identical requests, first request dropped (absorbs the post-warmup pickup gap), median of >= 5 clean steady-state requests (min–max); requests hit by neighbor-GPU interference on the shared box are excluded via a denoise-time stability filter (DenoisingStage is untouched by this PR, so its per-request time must be flat):

resolution	default (lossless)	--quality high	Δ per request	component-harness decode Δ
1024^2	0.740 s (0.69–0.80)	0.640 s (0.64–0.66)	−0.100 s	−68 ms (97.6 → 29.3)
2048^2	2.705 s (2.64–2.90)	2.430 s (2.41–2.50)	−0.275 s	−269 ms (437.2 → 168.1)

The e2e deltas match the component-harness decode deltas within run-to-run noise (a 4-step klein request is denoise-dominated, so decode's share of the request bounds the relative win). The component harness table above stays the authoritative decode-time measurement.

Cross-checks: (1) the pipeline stage-timer total (total_duration_ms, printed as "Warmed-up request processed in X") stops at decode kernel launch without a CUDA sync, so it is decode-blind by construction — and it is flat between the two modes (medians of 5 independent cold-process runs per cell: 0.41 / 0.41 s at 1024^2, 1.47 / 1.46 s at 2048^2; perf-dump totals 407 / 407 ms and 1454 / 1457 ms), which independently confirms the gate adds no overhead to TextEncoding/Denoising or host-side dispatch. (2) --quality high is accepted end-to-end by the generic sampling-param validation, and the log confirms the fast path engages only there — "FLUX.2 VAE: decoder switched to channels_last (NHWC) layout" appears once per process (lazy fold, not per request) and never in lossless runs.

End-to-end quality validation (maintainer bar for quality=high: PSNR > 25 dB vs the lossless output, same seed and prompt, saved PNGs):

prompt	resolution	PSNR (dB)	SSIM	bar	PASS
cyberpunk city (main)	1024^2	56.70	0.99949	PSNR > 25	yes
cyberpunk city (main)	2048^2	57.85	0.99930	PSNR > 25	yes
mountain lake at sunrise	1024^2	59.11	0.99900	PSNR > 25	yes
elderly fisherman close-up portrait	1024^2	58.06	0.99912	PSNR > 25	yes

All pairs sit 30+ dB above the bar (max per-pixel diff 4–6/255): the deviation is bf16 rounding-order level, not a visible quality change; the fp32 ground-truth check below confirms the fast path is not statistically further from an fp32 reference than the baseline bf16 decoder.

Accuracy

• Lossless default is bit-exact end-to-end: gate-off decodes of the frozen FLUX.2 latents on this branch produce byte-identical PNGs to the origin/main baseline (verified with cmp at 1024^2 and 2048^2), so golden-output CI is unaffected.
• fp32 ground-truth check (frozen latents, 1024^2, full decoder): fast-path bf16 error vs an fp32 VAE ground truth is statistically identical to the baseline bf16 error — mean abs err 0.000828 vs 0.000824, p99 0.003051 vs 0.003027.
• e2e quality: --quality high vs default output PSNR = 56.7–59.1 dB and SSIM >= 0.999 across 3 prompts / 2 resolutions (>= 55 dB internal bar; maintainer bar PSNR > 25 dB — see the e2e quality table above).
• Unit tests: test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py (+206, one file, 7 cases: channels_last kernel numerics fp32/bf16 vs fp32 reference, rows variant as used by attention, unsupported-input None contract incl. NCHW; gate-off torch.equal + no eager folding + fast-path PSNR for GroupNorm+SiLU and fused upsample; attention V/proj fold exactness). 7 passed on H200.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests.
• Provide accuracy and speed benchmark results.
• Follow the SGLang code style guidance.

Sample outputs (lossless vs --quality high)

FLUX.2-klein-4B, seed 42. Left: default (lossless, bit-exact reference path). Right: --quality high (fast path). Full-resolution originals; PSNR/SSIM per pair in the quality table above.

	lossless (default)	--quality high
cyberpunk city, 1024x1024
cyberpunk city, 2048x2048
mountain lake, 1024x1024
fisherman portrait, 1024x1024

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ⏳ Run #30906077610
Latest PR Test (Extra): ⏳ Run #30906077382

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/30906077610/job/91981729409?pr=33451