[Diffusion] Fuse LongCat GELU+cat and support Edit-Turbo BCG

[BBuf]

Motivation

LongCat's single-stream DiT blocks materialize BF16 tanh GELU and then concatenate it with attention output. The native Edit-Turbo profile shows 20 such sites per denoising step. A fused kernel removes the intermediate store/reload and one launch per site while preserving the eager BF16 result.

Edit-Turbo also lacks BCG registration. Simply registering it does not work for a non-square edit: the synthetic square input changes the warmup output grid and VL prefix, so the real request misses the captured signature. This patch makes the existing model-owned synthetic warmup hook preserve the explicitly requested aspect ratio and enables BCG for Edit-Turbo.

On one H200, the combined patch reduces native saved-image request E2E in two A–B–B–A groups: worker 2.46% / 2.65%, client including PNG save 2.89% / 2.62%. Baseline is original eager; candidate uses the fused kernel, corrected warmup and actual BCG. All valid lossless outputs are byte-identical.

Modifications

• Add a vectorized JIT CUDA cat(attention, gelu(mlp, approximate="tanh")) kernel. It copies attention unchanged, matches aten's FP32 operation order and BF16 store, and queries device SM count/occupancy at runtime.
• Use it in the LongCat single block behind BF16, layout, alignment, device and first-sight exactness checks. Preserve the existing quality-mounted GEMM+GELU path. Grad-enabled, compiled and unsupported calls retain the eager path; unverified graph capture does not attempt a host equality check.
• Build the synthetic edit image at the explicit warmup resolution. Preserve the complete VL prefix plus 512-token prompt body and image position IDs. Register Edit-Turbo for BCG; ordinary Edit is not added to the model-ID allowlist.
• Keep focused CUDA rounding/replay tests and the edit warmup geometry/BCG padding regressions.

This follows the lossless rounding approach in #38530 and #38396. Native edit conditioning comes from #35829 and is unchanged for user requests. There is no checkpoint, sampling, attention or quantization change.

Benchmark

meituan-longcat/LongCat-Image-Edit-Turbo@6a7262de5549f0bf0ec54c08ef7d283ef41f3214, 1× H200, BF16/lossless, resident components, no torch.compile. Prompt: “Make the cat wear a red hat.”, seed 42, 8 steps, CFG 1, prompt rewrite disabled. The fixed input produces 1264×848; the explicit warmup resolution is 1264×848 in both arms. DiT has 10 joint blocks and 20 single blocks; the single-stream sequence has 9233 tokens (8374 image + 859 conditioning), width 3072, MLP 12288.

Baseline 993d1fccbaafe3e79d91567d2fc1d665cc94fa50; candidate 10410a20b6d2716c1489a3132dc79c775dc297bb. Torch 2.13.0+cu130, Triton 3.7.1, CUDA 13.0, driver 595.71.05. Environment.

Fresh native CLI processes on the same idle GPU, two complete A–B–B–A groups. A=original eager; B=combined BCG. Seconds, excluding load/warmup:

Group	Run	Worker E2E	Denoise	Client incl. PNG save
1	A1	1.5705	1.2700	1.73
1	B1	1.5331	1.2493	1.68
1	B2	1.5342	1.2524	1.68
1	A2	1.5743	1.2718	1.73
2	A1	1.5663	1.2668	1.72
2	B1	1.5283	1.2490	1.68
2	B2	1.5260	1.2462	1.67
2	A2	1.5710	1.2693	1.72

Group	Mean metric	Baseline	Candidate	Reduction
1	Worker E2E	1.5724	1.5336	2.46%
1	Client incl. PNG save	1.7300	1.6800	2.89%
1	Denoise	1.2709	1.2508	1.58%
2	Worker E2E	1.5687	1.5272	2.65%
2	Client incl. PNG save	1.7200	1.6750	2.62%
2	Denoise	1.2681	1.2476	1.61%

Worker timing is native perf-dump total_duration_ms, measured before output saving. Client timing wraps the native request and saved PNG, logged to 0.01 s precision. Profile/export requests are excluded. Peak reserved memory is 32.805 GiB baseline, 31.328 GiB candidate eager, 35.248 GiB candidate BCG; graphs need additional memory. All raw perf dumps, source SHAs, configs and logs · Audited evidence.

Scope of the gain: kernel-only eager ABBA did not establish saved-request E2E improvement (worker 0.19%). Combined eager denoise improved 1.62% / 1.57%, but saved-client outliers of 2.87 s and 5.79 s prevent an eager E2E claim. These rows are retained in the evidence. The headline is the repeated combined BCG configuration above; it does not attribute the whole gain to graph replay alone.

Profile and operator benchmark

The complete synchronized baseline denoising step has 699 GPU kernels. Joining aten operators to kernels through trace External id identifies 20 GELU +20 channel concat kernels, 5.7457 ms. The candidate eager step has 679 kernels, including 20 fused calls, 3.2138 ms. Other joint-block GELUs remain unchanged.

The candidate BCG full trace records 248 cudaGraphLaunch calls = 31 segments × 8 steps, with no serving-signature miss or eager fallback. Its complete-step slice has 31 graph launches, 36 individual launches and 679 GPU kernels. All three step slices have zero boundary-crossing kernels. Baseline trace · Candidate eager · Candidate BCG · Replay audit. Full native traces, source attribution and three-table triage accompany each slice.

Archived marker benchmark, rotating inputs, default median, exact benchmark-candidate wrapper including output allocation:

Tokens	Attention /MLP width	Eager GELU+cat (us)	Fused (us)
859	3072 /12288	26.4200	14.8378
4608	3072 /12288	134.5206	73.8934
9233	3072 /12288	266.8966	147.5473

Raw marker results. NCU full/source reports confirm DRAM traffic 981.626→548.146 MB for the production pair. NCU timing is 302.336→183.360 us under profiling clocks and is kept separate from marker/native timing. Fused source stalls point to the MLP load and GELU arithmetic; the kernel has 32 registers and no spills. A tanh.approx experiment changed 60 finite BF16 outputs and was rejected. Baseline NCU · Fused NCU, including report files.

Output comparison

All 30 valid lossless saved images, including both ABBA groups, eager experiments and native profiles, are pixel- and byte-identical: max error 0, SSIM 1, PSNR ∞. Input and before/after images were visually inspected.

Input · Original baseline PNG · Original optimized PNG · Every image's SHA256 and metrics.

PNG SHA256: de4186303dcbcd776b58663209938a88bd27201e79a08df3433b313683e7e165.

The separate same-candidate lossless/high eager/BCG A–B–B–A matrix checks applicability. Rejected graph rows are not performance evidence. High eager produces SSIM 0.99224 / PSNR 32.84 dB against lossless in both rows. High+BCG is explicitly rejected because its request-scoped projection GELU differs from the lossless warmup graph; the headline stays lossless.

Matrix cell	Requested mode	Applicability
quality-lossless-a1	Eager	Valid
quality-lossless-b1	BCG	Valid
quality-lossless-b2	BCG	Valid
quality-lossless-a2	Eager	Valid
quality-high-a1	Eager	Valid
quality-high-b1	BCG	Rejected: cannot be used with breakable cuda graphs
quality-high-b2	BCG	Rejected: cannot be used with breakable cuda graphs
quality-high-a2	Eager	Valid

Validation and reproduction

Validation scope: The test counts, logs and GPU measurements below are historical results from the recorded benchmark revisions, before test cleanup. The current diff retains two GELU rounding/replay tests and the edit warmup geometry/BCG padding regressions. This cleanup does not change runtime code or claim a new GPU measurement.

• 59 tests passed on H200: 6 kernel cases, 7 LongCat dispatch/model cases, 46 config/padding cases plus 7 subtests. Kernel tests cover all 65,280 finite BF16 encodings including signed zero, production shape, preserved inputs, changed-input graph replay and invalid inputs.
• All changed files passed pre-commit, including registered test collection and executable test entries.
• Campaign-owned checkpoint cache was removed after validation: zero residual files, weight files or bytes. Cleanup audit.

Run each revision with the same pinned checkpoint and fixed input:

CUDA_VISIBLE_DEVICES=0 SGLANG_DIFFUSION_SYNC_STAGE_PROFILING=1 \
sglang generate --backend sglang \
  --model-path meituan-longcat/LongCat-Image-Edit-Turbo \
  --revision 6a7262de5549f0bf0ec54c08ef7d283ef41f3214 \
  --image-path input.jpg --prompt 'Make the cat wear a red hat.' \
  --num-gpus 1 --num-inference-steps 8 --guidance-scale 1 --seed 42 \
  --enable-prompt-rewrite=false --performance-mode manual --quality lossless \
  --enable-torch-compile=false --warmup-mode request \
  --warmup-resolutions 1264x848 --save-output \
  --output-path outputs --output-file-name sample --perf-dump-path perf.json

Add --enable-breakable-cuda-graph for candidate BCG. Use --profile --profile-all-stages only for separate diagnostic requests. Exact runner and preparation scripts · SHA256 manifest.

Additional LongCat checkpoint coverage

The same committed kernel was also tested on ordinary LongCat-Image-Edit@7b54ef423aa7854be7861600024be5c56ab7875a and LongCat-Image@d2ea50b79a930074c37b9b97ce45e3b2ea8cf4d8. Both use one H200, BF16/lossless, 50 steps, CFG 4.5, seed 42, resident components, no compile or prompt rewrite, and the same baseline/candidate SHAs above. Edit uses the same input and red-hat prompt at 1264×848. T2I uses 1024×1024 and “A red panda reading a book beside a sunlit window.” One complete ABBA per mode:

Family	Mode	Mean metric	Baseline s	Candidate s	Reduction
Edit	eager	Worker E2E	16.4149	16.1459	1.64%
Edit	eager	Client incl. PNG save	16.6050	17.7550	-6.93%
Edit	eager	Denoise	16.0308	15.8122	1.36%
T2I	eager	Worker E2E	7.0458	6.9293	1.65%
T2I	eager	Client incl. PNG save	7.2700	7.1150	2.13%
T2I	eager	Denoise	6.9239	6.8064	1.70%
T2I	BCG	Worker E2E	7.0778	6.9764	1.43%
T2I	BCG	Client incl. PNG save	7.2700	7.3150	-0.62%
T2I	BCG	Denoise	6.9546	6.8516	1.48%

Family	Mode	Run	Worker E2E s	Denoise s	Client incl. PNG save s
Edit	eager	A1	16.4306	16.0181	16.66
Edit	eager	B1	16.1430	15.8090	16.29
Edit	eager	B2	16.1489	15.8154	19.22
Edit	eager	A2	16.3992	16.0435	16.55
T2I	eager	A1	7.0449	6.9229	7.24
T2I	eager	B1	6.9341	6.8087	7.12
T2I	eager	B2	6.9246	6.8041	7.11
T2I	eager	A2	7.0467	6.9249	7.30
T2I	bcg	A1	7.0729	6.9491	7.26
T2I	bcg	B1	6.9736	6.8478	7.46
T2I	bcg	B2	6.9792	6.8555	7.17
T2I	bcg	A2	7.0827	6.9601	7.28

Ordinary Edit's candidate B2 has a 19.22 s saved-client outlier, so these rows do not establish an eager saved-request gain. Its first B2 admission was rejected while the preceding CUDA context drained; the retained B2 is a uniquely named retry after waiting for idle GPU outside the request timer. Ordinary Edit BCG remains explicitly disabled and is excluded. T2I BCG compares actual replay in both arms; the saved-client mean does not improve. T2I eager is a positive single-group observation, not an additional repeated headline claim.

The complete ordinary Edit CFG iteration contains 80 GELU/cat kernels taking 12.2241 ms, versus 40 fused calls taking 6.9521 ms. Native paired traces and three-table triage are included for both checkpoints. T2I's BCG profile records four CFG iterations: 8 model calls ×31 segments =248 actual graph launches, with no serving miss/fallback. All selected CFG slices have zero boundary-crossing kernels.

Output comparison: all seven ordinary Edit PNGs and all eleven T2I PNGs are byte- and pixel-identical within each checkpoint. Both outputs were visually inspected. SHA256: Edit 34c9c9ef854a64203d71188161f77929cb1564f63a87b9d73d522d59cf8eeb72; T2I 0443d562b6da18257dbc587617b6ec0bcfac5b078c8bce18243fdb3ce50f7475.

Checkpoint	Baseline	Candidate
Ordinary Edit
T2I

Ordinary Edit raw rows, native logs and traces · T2I raw rows, native logs and traces. Both owned weight caches were deleted after validation, with zero residual files/bytes: Edit cleanup, T2I cleanup.

CI follow-up (2026-09-21)

Merged main 80da4432d085ed4d6166ef643d9fd2b829dbb0c5, including the upstream scheduler fixture repair #40411 and offline DFlash fixture repair #40427. Validated merge head: 15d83b34e132bed6867f62ec36bc75af46211aab. Merged-head validation and all logs, source delta in this PR's files, pre-commit. The original E2E and output comparisons above remain measurements of their explicitly pinned benchmark heads; this follow-up does not claim new model-request timings. The common repaired CPU fixture suite passes 49 tests + 57 subtests.

Current-head H100 lane audit

Single-GPU partition 3, job 106118222114 completed with LTX 2.3 HQ performance failures after the initial execution plus six retries. Consistency checks passed, but failures moved between E2E, denoise steps and VAE decode; the last decode was 2953.94 ms versus 1681.54 ms. Runner: h100-novita-host3-gpu-4, H100 80GB, driver 580.126.09, Torch 2.13.0. Its log warns that CUDA_VISIBLE_DEVICES is unset; that warning alone does not establish contention or incorrect device assignment. The checked-in H100 reference metadata names source revision 2ee0d38a8553d186faad9c50893480e3c1d79270 and 2026-08-16, so it is not a contemporaneous main-versus-PR control. Component-accuracy/Base-B health-check skips are not successful model tests. Thresholds remain unchanged and no blanket rerun was requested. Source-versus-environment attribution remains unresolved without a controlled same-runner comparison and runtime telemetry.

---

CI States

Latest PR Test (Base): 🚫 Run #35954508425
Latest PR Test (Extra): ⚠️ Not enabled -- add run-ci-extra label to opt in.
Latest PR Test (AMD ROCm 10): ⏳ Run #35954508075

[mickqian]

Exact-head CI follow-up (10410a2): test_disagg_idle_gap_0_prefill_normal fails because the Scheduler fixture lacks _prev_prefill_end_ts. Log: https://github.com/sgl-project/sglang/actions/runs/35470402130/job/105970288701 . This CPU test-fixture issue is fixed on main by #40411. Please merge current main and run CI on the resulting head; rerunning this unchanged head will not repair the fixture. I cannot push to the head repository. This diagnosis concerns this CPU failure only and does not establish that every GPU failure is unrelated.

[mickqian]

Current-head CI follow-up for 15d83b3: the real completed failure is LTX 2.3 HQ in single-GPU partition 3 (https://github.com/sgl-project/sglang/actions/runs/35525996118/job/106118222114). It exhausted initial execution plus six retries, with changing E2E/denoise/decode performance failures; the final decode was 2953.94 ms versus the 1681.54 ms limit. Component accuracy job106118191509 did not run accuracy tests: its health check fast-failed on this LTX job. Inspected Base-B job106119984607 likewise failed its health check, not its serving tests. The run still has pending work, so this is not a complete run-success diagnosis. These logs do not establish a LongCat correctness regression or prove a runner root cause. Preserve thresholds and investigate the LTX execution environment before another expensive retry. I have no push permission to the head repository and have not changed this branch.

[mickqian]

Exact-head maintenance follow-up for 4b1789e: an isolated merge-tree against live main 81cb895 confirms one conflict in python/sglang/kernels/ops/diffusion/init.py. Preserve both new can_use_fused_gelu_tanh_cat and fused_gelu_tanh_cat exports, while taking the main-side load_extension_with_recovery mapping to sglang.srt.utils.cpp_extension_loader (not ext.loader). I have no push permission to this fork.

The new Base-C root failure is separate: test_dp_attention_local_control_broadcast.py fails during server initialization with TCPStore EADDRINUSE on port 41135 on h100-novita-temp-gpu-4567. Exact log: https://github.com/sgl-project/sglang/actions/runs/35871934132/job/107319279203 . Six other failed Base-C jobs explicitly fast-failed on this root job; they did not run their model tests. One targeted rerun request was refused by gh with job cannot be rerun while run attempt 1 still has queued B300 work; no rerun was accepted. Reassess a targeted retry after the run settles and the current head is refreshed. No code or threshold changes were made, and this port collision does not establish a LongCat regression.

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/35954508425/job/107570310394?pr=40384