[Diffusion] Enable lossless BCG for FLUX.1-dev

[BBuf]

Summary

Enable native FLUX.1-dev breakable CUDA graph execution for lossless requests. --enable-breakable-cuda-graph previously disabled itself for this model; the existing runner now captures its DiT segments and replays them for warmed shapes.

On 2× H200 at 1024×1024 / 50 steps, two A/B/B/A rounds reduce warmed engine E2E from 4.315577 s to 4.177754 s (-3.194%), and denoise from 4.119540 s to 3.911356 s (-5.054%). All eight formal PNGs are byte-identical. Peak reserved memory rises from 19.473 GiB to 20.707 GiB in the rank-0 perf dump.

The implementation adds FLUX.1-dev model IDs and FluxPipelineConfig to the existing support checks. FLUX keeps its unmasked, fixed 512-token T5 conditioning, so the generic runner can copy fresh conditioning/timestep/guidance inputs into the captured buffers. Existing TP graph registration, signature fallback, and quality-fusion checks apply. The patch also adds warmup/model-reference and conditioning-shape tests, plus a cookbook example.

Benchmark

• Native black-forest-labs/FLUX.1-dev at checkpoint 3de623fc3c33e44ffbe2bad470d0f45bccf2eb21; BF16, TP2 on the same two H200 GPUs throughout.
• PyTorch 2.13.0+cu130, Triton 3.7.1, driver 595.71.05; loaded NCCL reports 2.28.9.
• 1024×1024, 50 steps, guidance 3.5, seed 42, quality=lossless, dit=resident, other component residency auto; torch.compile disabled.
• Prompt: A futuristic cyberpunk city at night, neon lights reflecting on wet streets.
• Baseline: 482e9f257bb9d9ac57c2245c93768a96ec70edbe, eager. Measured candidate: 5e91bc39e7a421be979448661b1b5d93db35cb48, BCG. Final head 3d51468a54b8fad047b1d679c0968ac2919613e7 adds only tests/docs; its runtime file is byte-identical to the measured candidate.
• The preset uses one-step request warmup for eager. BCG forces synthetic server warmup/capture at the declared 1024×1024 shape (50 steps). Both warmup paths are excluded from measured request latency; the results compare those native operating modes.
• E2E is the warmed engine duration from the unprofiled perf dump, including encoding, denoise, and decode. Checkpoint download, startup/warmup, and final PNG writing are excluded. Profiler export wall time is excluded from all performance claims.

Run order	Eager E2E (s)	BCG E2E (s)	Eager denoise (s)	BCG denoise (s)
Round 1 — A1	4.315031		4.119472
Round 1 — B1		4.104417		3.905989
Round 1 — B2		4.105608		3.907831
Round 1 — A2	4.306322		4.105668
Round 2 — A1	4.337554		4.144473
Round 2 — B1		4.385315		3.917348
Round 2 — B2		4.115677		3.914257
Round 2 — A2	4.303404		4.108547
Mean	4.315577	4.177754	4.119540	3.911356

Per-round E2E reductions are 4.771% and 1.620%. Round 2 B1 spent 371.106 ms in text encoding versus roughly 98–102 ms in neighboring requests; that sample remains in the aggregate. No samples were discarded.

All benchmark rows and conditions · raw perf dumps and scripts

Profile comparison

The paired full-stage traces cover two denoise steps on rank 0, using the same checkpoint, shape, prompt, seed, topology, and respective source commits. Overlapping GPU stage annotations are unioned; no selected kernel crosses a stage boundary.

Two-step denoise diagnostic	Eager	BCG
GPU kernels	2,847	2,847
Host per-kernel launches, CUDA runtime + driver	2,847	157
cudaGraphLaunch calls	0	116
Captured segments per step	0	58
Cumulative GPU kernel time (ms)	144.160	147.284
Profiled CPU stage duration (ms)	278.175	160.611

BCG reduces recurring host launches; the cumulative device time does not improve. Most GEMMs, attention, and existing fused LN/modulation kernels keep their counts. The existing graph-aware custom all-reduce path changes from 228 push calls to 152 pull + 76 push calls. Collective durations include rank waiting, so they are not interpreted as pure communication bandwidth. The profiled CPU durations are diagnostic and differ substantially from the unprofiled timings above.

Eager profile digest · BCG profile digest. Raw trace sizes and SHA-256 hashes are recorded in those digests.

Existing FLUX optimizations from #34126, #34004, #33819, and #33818 were reviewed before choosing the graph replay path. Their fused arithmetic and quality gates are reused.

Output comparison

All eight 50-step benchmark PNGs have SHA-256 a713777ad67b4e82915aef972612f243d6f4cbd8db9aebdfdff976b4c7d2c7d3.

Eager, 50 steps	BCG, 50 steps

All saved benchmark images · Replay correctness results

On final head, persistent native generators at TP1 and TP2 each served six requests across three warmed resolutions, changed prompts and seeds, repeated an earlier request, then requested one deliberately uncaptured resolution. All 12 eager/BCG image pairs (24 PNGs) are byte-identical within their respective topology. Each rank captured exactly three graphs at warmup, with zero serving recaptures; only the final 640×512 request used the expected eager fallback. These six-step images validate replay state and are separate from the 50-step performance rows.

quality=high + BCG was tested and rejected by the existing DiT fusion compatibility check. The graph captures the lossless branches, so this PR does not claim high/extra-high BCG support. The baseline high eager image has SSIM 0.975800 and PSNR 32.3649 dB against lossless on this prompt; that separate existing quality tradeoff is not used in the lossless speedup.

Reproduce

Use an HF_TOKEN with access to the checkpoint. Run on each pinned source checkout and keep the same GPUs and request parameters:

CUDA_VISIBLE_DEVICES=0,1 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
SGLANG_DIFFUSION_SYNC_STAGE_PROFILING=1 \
NCCL_NVLS_ENABLE=0 \
sglang generate \
  --backend sglang \
  --model-path black-forest-labs/FLUX.1-dev \
  --revision 3de623fc3c33e44ffbe2bad470d0f45bccf2eb21 \
  --num-gpus 2 --tp-size 2 \
  --component-residency dit=resident \
  --enable-torch-compile false \
  --warmup-mode request \
  --quality lossless \
  --width 1024 --height 1024 \
  --num-inference-steps 50 --guidance-scale 3.5 --seed 42 \
  --prompt "A futuristic cyberpunk city at night, neon lights reflecting on wet streets" \
  --save-output --output-file-name eager.png \
  --perf-dump-path eager.json

For BCG, add --enable-breakable-cuda-graph --warmup-resolutions 1024x1024 and use distinct output/perf filenames; the runtime automatically changes warmup mode to server. Run A/B/B/A twice. For diagnostic traces, use two steps and add --profile --profile-all-stages; compare unprofiled requests for latency. The cookbook documents warmed resolution matching and the additional startup/graph memory cost.

Validation

• 52 warmup/padding/TP graph-capture unit tests and 16 subtests pass.
• Two 50-step A/B/B/A rounds; all eight saved PNGs byte-identical.
• Paired full-stage profiles; both two-step PNGs byte-identical.
• Persistent TP1/TP2 replay checks on final head: changed prompt/seed, three warmed resolutions, repeated request, and deliberate uncaptured-resolution fallback.
• Four changed files pass pre-commit; Mint build validation and broken-links checks pass.
• Task-owned FLUX checkpoint cache removed: 57,890,215,216 logical bytes / 9 weight files to zero, with an independent directory-absence check. Cleanup ledger.

Other FLUX variants, quantized transformer overrides, and other hardware require their own validation.

---

CI States

Latest PR Test (Base): ✅ Run #34300355195
Latest PR Test (Extra): ✅ Run #34300355947
Latest PR Test (AMD ROCm 7.2): ❌ Run #34300355213

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/34300355195/job/102322703368?pr=38591