Skip to content

[Diffusion][FLUX.2] Fuse eager AdaLN and packed SwiGLU - #34616

Merged
BBuf merged 1 commit into
sgl-project:mainfrom
BBuf:bbuf/b300-flux2-eager-fusions
Aug 13, 2026
Merged

BBuf merged 1 commit into
sgl-project:mainfrom
BBuf:bbuf/b300-flux2-eager-fusions

Conversation

@BBuf

@BBuf BBuf commented Aug 12, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Speed up FLUX.2 eager inference with two bit-exact fusions:

  • LayerNorm + modulation at the single- and double-stream AdaLN sites
  • packed SwiGLU without materializing two sliced operands

Both paths retain the original expressions for torch.compile. The eager fast paths use per-signature first-use verification and fall back safely when layout, dtype or capture state is unsupported. The cleanup also makes two correctness constraints explicit:

  • a previously unseen SwiGLU signature never performs verification or a host sync during CUDA graph capture
  • the packed kernel rejects non-dense outer batch strides instead of flattening an unsupported view

B300 performance

FLUX.2 Klein Base, 1024x1024, seed 42, 50 steps, native backend, quality=lossless:

configuration denoise e2e vs latest-main eager
latest main eager 3.7310 s 3.9472 s baseline
#34616 only 3.2532 s 3.4687 s +12.81% / +12.12%
#34615 + #34616 2.9738 s 3.1957 s +20.30% / +19.04%
torch.compile with both PRs 2.9985 s 3.1777 s +19.63% / +19.49%

With both FLUX.2 PRs, eager is 0.83% faster than compile in denoise; E2E is within 0.57% in this single-request measurement.

Measured against main dbebc1deb42b00befa3d0de67265d7003994c1ad. The B300 validation tree is identical to this PR's final Git tree.

Output comparison

The before/after images use the same prompt, seed, resolution and step count. The generated PNGs are pixel-exact and have the same SHA-256 (6a103716dd891a301409e51bc6bd6dee2d6fb997fee21bfe0785fd504d6a5a39).

latest main eager (before) #34616 eager (after)
FLUX.2 latest-main output FLUX.2 output with PR 34616

Validation

  • B300 fusion exactness and fallback tests: 5 passed
  • coverage includes production sequence lengths, strided packed input, FP16 fallback, non-dense outer strides and a new signature encountered during capture
  • B300 native model run completed without backend fallback
  • final changed-file pre-commit suite passed
  • weight cache cleanup verified after the model runs

CI States

Latest PR Test (Base): ❌ Run #31687601642
Latest PR Test (Extra): ✅ Run #31692846540

@BBuf
BBuf force-pushed the bbuf/b300-flux2-eager-fusions branch from 565b96d to 8fec087 Compare August 13, 2026 09:39
@BBuf BBuf added the mergeable label Aug 13, 2026
@BBuf
BBuf merged commit 74c0322 into sgl-project:main Aug 13, 2026
141 of 196 checks passed
saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 16, 2026
Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

diffusion SGLang Diffusion jit-kernel mergeable run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant