[Diffusion] Fuse lossless Wan VAE post-ops for LongLive 2 I2V

[BBuf]

Motivation

LongLive 2's lossless Wan VAE decoder executes five separate FP32 pointwise passes after normalization. In the native I2V profile, the 464 norm/SiLU scopes take 690.82 ms across 3248 kernels. Keeping Torch's reduction and fusing only the exact post-ops reduces that to 275.26 ms / 1834 kernels, including first-sight reference checks and fallback calls.

On one H200, final-source I2V A–B–B–A groups reduce worker E2E by 5.53–5.88%. The representative final group reduces client E2E including video save from 5.980 to 5.665 s (5.27%). All lossless before/after videos are byte-identical. T2V worker latency improves, but retained client outliers prevent a T2V client E2E speed claim.

Modifications

• Keep native autocast-aware L2 normalization and epsilon clamp, then fuse divide, scale, affine multiplication, bias addition and SiLU into one JIT CUDA kernel. Inputs/gamma may be BF16, while the normalized intermediates and output remain FP32.
• Use explicit FP32 round-to-nearest operations and expf, with FMA contraction disabled and no fast math. Support dense NCDHW and channels-last-3D layouts through the existing JIT utilities.
• Integrate through the shared Wan VAE wrapper with per-signature integer-bit verification and permanent fallback on mismatch. Unsupported dtype/layout, gradient-enabled, compiled and unverified-capture calls retain the native path. Unsupported post-op inputs reuse the already-computed denominator.
• Preserve the existing high-quality full RMSNorm+SiLU dispatch. Keep three focused tests for autocast/layout/affine behavior, numerical boundaries and graph replay, and update fusion-coverage documentation.

The existing Wan full-norm fusion (#33546, #35981) and the open conv-bias/full-norm tiling work in #38650 are distinct from this lossless post-only path. The shared wrapper also serves other Wan/Qwen VAEs; this PR's native model evidence covers LongLive 2 T2V and I2V. Other model families are protected by the first-sight exactness gate but are not claimed as benchmarked here.

Source review and execution flow. Native SiLU arithmetic was checked against the installed Torch source.

Benchmark

Model Rabinovich/LongLive-2.0-5B-Diffusers@37664a64cfc38cdc5fabe88576961ef10646e4f1. Native SGLang eager, 1×H200 / TP1, unquantized BF16 DiT, native VAE/autocast precision, all components resident. Automatic performance mode resolves CPU/layerwise offload to false in both arms. No torch.compile or approximate cache. Four steps, CFG1, flow shift5, seed42, 61 frames at24fps (2.5417s), no audio.

• I2V: native helper requests960×928; the input-image aspect handling produces 1152×768 in all24 accepted I2V videos. The helper's supported-resolution advisory is retained. Prompt: The cat starts walking slowly towards the camera.
• T2V: requested and actual832×480. Prompt: A curious raccoon.
• Both modes use the same native helper warmup in every arm: 960×928×29 frames, two of four steps. Warmup is excluded from request timing; the requested warmup flags do not override that helper's synthetic request geometry.

Baseline 993d1fccbaafe3e79d91567d2fc1d665cc94fa50; final measured/PR candidate be05e7e6c24252db11458621f2d06f3969be08d1. Earlier r1/r2 use a98b0875c4; final r3/r4 include the reviewed fallback fix and unique custom-op name. The CUDA device implementation and marker/NCU harness are unchanged. Torch2.13.0+cu130, Triton3.7.1, CUDA13.0, driver595.71.05.

Environment · fixed input image.

Each group uses fresh native CLI processes on the same idle GPU in A–B–B–A order. Seconds; reductions are calculated from arithmetic means, with no sample removal.

Workload / group	Worker E2E	Denoise	Client E2E including save
i2v-r1	5.6111 → 5.2703 (6.07%)	1.5028 → 1.5025 (0.01%)	6.0000 → 5.6600 (5.67%)
i2v-r2	5.6117 → 5.2834 (5.85%)	1.5018 → 1.5000 (0.12%)	6.0100 → 5.8400 (2.83%)
i2v-r3	5.5972 → 5.2878 (5.53%)	1.4977 → 1.5007 (-0.20%)	5.9800 → 5.6650 (5.27%)
i2v-r4	5.6159 → 5.2855 (5.88%)	1.5038 → 1.5011 (0.18%)	6.5400 → 5.7200 (12.54%)
t2v-r1	2.3925 → 2.2465 (6.10%)	0.5832 → 0.5813 (0.32%)	2.6900 → 4.2350 (-57.43%)
t2v-r2	2.3906 → 2.2615 (5.40%)	0.5827 → 0.5843 (-0.28%)	2.6950 → 2.5650 (4.82%)
t2v-r3	2.3984 → 2.2782 (5.01%)	0.5858 → 0.5856 (0.04%)	2.7350 → 2.5800 (5.67%)
t2v-r4	2.3905 → 2.2704 (5.02%)	0.5848 → 0.5821 (0.47%)	2.6850 → 3.8750 (-44.32%)

All four I2V groups clear1.5% in both saved-request worker and client means. I2V r4A1 has a7.08s baseline client outlier, inflating r4's12.54% client reduction; 5.27% from final r3 is the representative client figure. Denoise is essentially unchanged because the optimization is in VAE decoding.

T2V r1B1 and r4B2 clients take5.93s and5.18s despite worker times2.2457s and2.2708s. These delays occur after the output-saved log. A separate fixed diagnostic ABBA also reproduces a post-save delay on the unmodified baseline (6.09s client /2.3921s worker). Active-thread py-spy and GC callbacks did not determine its cause. No T2V client E2E speedup is claimed, and none of these samples is removed.

Diagnostic evidence. The original all-workloads audit retains qualified:false; workload qualification explicitly accepts I2V alone.

All32 unprofiled native A/B measurements

Group	Arm	Worker E2E	Denoise	Client E2E incl. save	Peak reserved GiB
i2v-r1	A1	5.6134	1.5025	6.00	59.6426
i2v-r1	B1	5.2737	1.5001	5.66	61.0215
i2v-r1	B2	5.2668	1.5049	5.66	61.0215
i2v-r1	A2	5.6089	1.5030	6.00	59.6426
i2v-r2	A1	5.6152	1.5060	6.03	59.6426
i2v-r2	B1	5.2896	1.4984	6.01	61.0215
i2v-r2	B2	5.2771	1.5016	5.67	61.0215
i2v-r2	A2	5.6082	1.4976	5.99	59.6426
i2v-r3	A1	5.5936	1.4966	5.98	59.6426
i2v-r3	B1	5.2842	1.5018	5.66	61.0195
i2v-r3	B2	5.2915	1.4995	5.67	61.0195
i2v-r3	A2	5.6008	1.4988	5.98	59.6426
i2v-r4	A1	5.6242	1.5043	7.08	59.6426
i2v-r4	B1	5.2861	1.4992	5.77	61.0195
i2v-r4	B2	5.2848	1.5029	5.67	61.0195
i2v-r4	A2	5.6077	1.5032	6.00	59.6426
t2v-r1	A1	2.3958	0.5853	2.69	59.6426
t2v-r1	B1	2.2457	0.5804	5.93	61.0195
t2v-r1	B2	2.2472	0.5823	2.54	61.0195
t2v-r1	A2	2.3893	0.5811	2.69	59.6426
t2v-r2	A1	2.3927	0.5824	2.71	59.6426
t2v-r2	B1	2.2613	0.5838	2.57	61.0195
t2v-r2	B2	2.2618	0.5848	2.56	61.0195
t2v-r2	A2	2.3885	0.5830	2.68	59.6426
t2v-r3	A1	2.3965	0.5867	2.70	59.6426
t2v-r3	B1	2.2829	0.5873	2.57	61.0195
t2v-r3	B2	2.2736	0.5838	2.59	61.0195
t2v-r3	A2	2.4002	0.5849	2.77	59.6426
t2v-r4	A1	2.3875	0.5836	2.68	59.6426
t2v-r4	B1	2.2701	0.5796	2.57	61.0195
t2v-r4	B2	2.2708	0.5845	5.18	61.0195
t2v-r4	A2	2.3934	0.5860	2.69	59.6426

Worker E2E is the native perf-dump stage duration before saving, excluding loading/warmup. Client E2E is the native scheduler-request/output-save timer logged to0.01s. Profile, marker, sampler and initial admission runs are excluded from performance claims. Final peak reserved memory increases 59.6426→61.0195GiB (+1.377GiB).

All raw perf dumps, commands, source SHAs and native logs.

Profile and kernel benchmark

Full native all-stage profiles are used because this causal loop does not advance the denoise-step profiler. CPU norm-wrapper scopes are joined to kernels through CUDA launch correlations; nested stage totals are not double-counted. Each full request has464 decoder norm/SiLU scopes.

Workload / source	Attributed kernels	Attributed GPU ms	Fused post-op calls	Fused GPU ms
t2v-baseline-profile-retry1	3248	312.4739	0	0.0000
t2v-norm-final-profile	1834	124.1225	434	47.5257
i2v-baseline-profile	3248	690.8231	0	0.0000
i2v-norm-final-profile	1834	275.2550	434	107.9946

Each final trace has434 fused calls (415 channels-last and19 NCDHW), with30 scopes retaining fallback. The first-sight reference work remains included. Cumulative GPU times explain the change and are not substituted for E2E latency.

Original full traces, source attribution and three-table triage.

Archived whole-chain marker benchmark, rotating inputs, median microseconds. Both arms include the native norm/clamp and normal output allocation:

Input N,C,T,H,W	Layout	Torch µs	Fused µs
1,1024,1,30,52	NCDHW	51.008	24.288
1,1024,1,30,52	channels-last-3D	45.696	19.360
1,1024,2,60,104	NCDHW	216.192	63.520
1,1024,2,60,104	channels-last-3D	208.448	52.800
1,256,4,240,416	NCDHW	1484.960	357.696
1,256,4,240,416	channels-last-3D	1519.456	364.128

NCU isolates the five post-ops at the last channels-last shape, with the same precomputed native denominator in both arms:

NCU measurement	Five Torch kernels	One fused kernel
Sum of kernel duration	1599.136µs	315.104µs
DRAM reads	1,842,186,240B	206,129,664B
DRAM writes	1,930,885,632B	385,019,392B

The isolated chain is 5.07× faster, with 84.33% less measured DRAM traffic. The fused launch uses30 registers/thread,1056 CTAs×256threads, zero local load/store sectors and70.62% active warps. Combined DRAM throughput is38.99% of peak; source/PM counters identify denominator-load long-scoreboard and SiLU dependency/MIO stalls. There is no claim of saturated HBM or tensor-core execution for this pointwise kernel.

Marker log · NCU analysis, source/PM samples and next experiment · Full/source .ncu-rep reports.

Output comparison

The native audit covers56 request records:48 accepted videos and eight invalid BCG requests. All42 accepted lossless videos are byte-identical within each mode; independent decoding verifies61/61 identical frames (SSIM1, PSNR∞),24fps and no audio. Both lossless contact sheets were visually inspected.

Original I2V baseline MP4 · Original optimized MP4 · High-quality comparator MP4.

Original T2V baseline MP4 · Original optimized MP4 · High-quality comparator MP4.

Lossless SHA256:

• I2V: 07027ed1e3cc3bd9fb97ff6316226010252f5509366437d89d9aeeb28423ede0
• T2V: fb15120992e8e5fc5bf50fb15ffe31e653f42d8e896bd620ef773015bb44f5b4

The separate high/lossless and eager/BCG applicability matrix has the same output in baseline, initial candidate and final candidate at each high mode. Compared with each mode's lossless reference:

Quality/mode	Mean SSIM	Worst-frame SSIM	Mean PSNR dB	Worst-frame PSNR dB
high T2V eager	0.955288	0.945272	36.5040	35.3677
high I2V eager	0.980385	0.977651	43.4208	42.7313

Both clear the video thresholds0.92/24dB. Actual LongLive2 BCG is explicitly disabled at lossless and high in both arms; all eight probes are excluded from performance. Standalone CUDA graph replay tests do not imply native BCG support.

T2V quality matrix · I2V quality matrix · Per-frame hashes.

Test

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 three autocast/layout, numerical-boundary and graph-replay tests. This cleanup does not change runtime code or claim a new GPU measurement.

H200: 9 pytest cases plus16 subtests passed. Coverage includes all65,280 finite BF16 patterns, FP32 SiLU/subnormal boundaries, BF16/FP32 input and affine tensors, both layouts, nonintegral sqrt(C), zero norms, changed-input/affine CUDA replay, first-sight/mismatch/capture fallback, unsupported dtype/layout, denominator reuse, compile guard, gradients, VAE-gate behavior and import isolation.

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python FLASHINFER_DISABLE_VERSION_CHECK=1 \
python -m pytest -q \
  test/registered/kernels/ops/diffusion/test_wan_norm_silu_post.py \
  python/sglang/multimodal_gen/test/unit/test_vae_fast_path_gate.py \
  python/sglang/multimodal_gen/test/unit/test_diffusion_import_isolation.py

Pre-cleanup validation covered all seven files and included Mint documentation and broken-link checks with Node22. The retained tests now use the single-GPU CUDA CI lane; the archived local GPU results were collected on H200.

Test log · Pre-commit · Documentation build · Link check.

Reproduce

Use the pinned model revision and fixed input image, with the same native helper command in baseline and candidate checkouts. The task-owned snapshot's main ref is pinned to the revision above; offline mode prevents a moving download.

CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=8 \
FLASHINFER_DISABLE_VERSION_CHECK=1 SGLANG_DIFFUSION_SYNC_STAGE_PROFILING=1 \
sglang generate --backend=sglang \
  --model-path=Rabinovich/LongLive-2.0-5B-Diffusers \
  --prompt='The cat starts walking slowly towards the camera.' \
  --image-path=/path/to/input-cat.png --seed=42 \
  --width=960 --height=928 --num-frames=61 \
  --num-inference-steps=4 --guidance-scale=1.0 \
  --num-gpus=1 --tp-size=1 --enable-torch-compile=false \
  --warmup-resolutions=960x928 --warmup-num-frames=61 \
  --quality=lossless --save-output --warmup-mode=request \
  --perf-dump-path=/tmp/longlive2-i2v.json \
  --output-path=/tmp/longlive2-i2v --output-file-name=longlive2-i2v

For T2V, use --prompt='A curious raccoon', remove --image-path, and set width832/height480 and warmup-resolutions832x480. For diagnostic native profiles add --profile --profile-all-stages and set SGLANG_DIFFUSION_TORCH_PROFILER_DIR; exclude those runs from latency claims.

Exact runners, fixed ABBA schedule and audit scripts. After the model cycle, 689 owned files, including five weight files (24,211,983,385 bytes total), were removed and zero residual files/weights/bytes were verified. The shared cache was untouched. Cleanup ledger · Artifact SHA256 manifest.

Upstream CI follow-up

Base CI, attempt 2: both single-GPU JIT partitions, multi-GPU JIT, JIT benchmark, multimodal unit tests, BCG, RTX 5090 diffusion, and four 1-GPU diffusion partitions passed. The workflow is red because unchanged CPU fixtures lack split_index / _prev_prefill_end_ts; later B200, component-accuracy, remaining 1-GPU and 2-GPU jobs stopped at PR-health fast-fail before their tests, followed by incomplete diffusion coverage. The local H200 numerical, profile and native saved-output evidence above passed.

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: 5fa8ec4ae7efb10ee6e7dc936370023cdca0b444. 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.

---

CI States

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

[mickqian]

Exact-head CI follow-up: CPU(6) job106020065558 fails test_pdmux_split_prefill_schedules_auxiliary_output_copy because its batch fixture lacks split_index; CPU(9) job106020065542 fails test_disagg_idle_step_counters.py because Scheduler lacks _prev_prefill_end_ts. The B200 kernel job106020065430 and H100 partition3 job106020087281 explicitly fast-fail on the CPU root before model/kernel testing. These failures are not evidence of a Wan VAE regression. Update: main now contains the upstream scheduler fixture repair in #40411, commit22f02cc3399dcd59380fc546de8d3c89fba6fa3a. Please merge updated main instead of rerunning the stale fixtures or cherry-picking the earlier #39206 variants. The upstream repair includes completion-clock/burst accounting and additional split-prefill/overlap tests. This diagnoses these root failures only, not a claim that all checks have passed.

[BBuf]

https://github.com/sgl-project/sglang/actions/runs/35871872451/job/107527379198?pr=40405