[diffusion] Replay small Triton wrappers' launches per operand signature

[mickqian]

Motivation

Many of sglang's small Triton kernels cost more host time per call than GPU time. Before each launch, the Python wrapper runs operand checks, stride math, the device guard and Triton's JIT dispatch. For rmsnorm_scale at a Z-Image 512² shape that is ~28 µs on H200, against a ~9 µs launch. None of that work depends on tensor values, only on the operands' signature.

A py-spy profile of the denoising loop (H200, 2 GPUs, eager, 256² so that the host bounds the step) shows how much of each step this is. The share below is the host time spent inside sglang's diffusion fused-kernel wrappers and their predicates:

model	share of host time	largest wrappers
FLUX.1 TP2	≈32%	fused_layernorm_modulate_raw 14.6%, residual_gate_add 8.5%
FLUX.2 TP2	≈38%	fused_layernorm_modulate_raw 12.3%, residual_gate_add 8.2%, flux2_strided_qknorm_rope 6.2%, fused_packed_silu_mul_bitexact 6.1%
Qwen-Image-2.1 Ulysses2	≈28%	fused_layernorm_modulate_raw 10.6%, residual_gate_add 7.4%, qknorm_complex_rope_cuda 4.5%
Ideogram-4 TP2	≈26%	fused_rope_rotate_half_bitexact 8.1%, residual_gate_add 7.5%, modulate_scale_shift 4.9%

Whether this reaches end-to-end latency depends on whether the step is host-bound. That is the case on Grace hosts (GB200/GB300), with several TP ranks, and at low resolution.

#39878 removes the overhead for three Z-Image kernels by replaying recorded launches from a Rust extension. This PR does the same in pure Python, and does it generically, so that covering another kernel takes one decorator line and no change to the kernel or its wrapper.

Modifications

sglang/kernels/ops/diffusion/common/fast_launch.py adds @replay_launches:

• First call with a given signature. The signature is each tensor's shape, stride, dtype, device and data_ptr() % 16, plus the non-tensor arguments. On this call the wrapper runs normally while its Triton launches are recorded: compiled kernel, grid and arguments, with each tensor argument mapped by pointer to an input or a returned output.
• Later calls with that signature. They allocate the outputs with empty_strided and call the recorded compiled launchers directly. The kernels, grids and arguments are the same, so outputs are bit-identical.
• What a signature must satisfy to be replayed. The wrapper's GPU work is Triton launches into outputs it allocates with torch.empty* and returns. Recording enforces what it can see of that and otherwise leaves the signature on the normal path:
  • aten ops other than allocation and views (checked with a dispatch mode; a view is any op that mutates nothing and returns only aliases of its inputs, so reshape on inference tensors counts);
  • reads of tensor values (.item());
  • launches on tensors that are neither inputs nor outputs, or on inputs that share a pointer.
• Graph capture, compile and Triton versions.
  • Recording does not sync, so it also works inside CUDA graph capture.
  • Under torch.compile the decorator calls the wrapper itself, so it stays traceable.
  • A Triton whose launcher does not take every parameter (before 3.3) stays on the normal path.
  • SGLANG_DIFFUSION_TRITON_FAST_LAUNCH=0 turns it off.

The decorator is applied, one line each, to the hot Triton wrappers above:

• rmsnorm_scale, rmsnorm_tanh_residual, zimage_qk_rmsnorm_native
• fused_layernorm_modulate_raw
• fused_silu_mul_bitexact, fused_packed_silu_mul_bitexact
• flux2_strided_qknorm_rope, fused_rope_rotate_half_bitexact

Kernels and wrapper bodies are unchanged: existing files change by +18/−1 lines in total (one import and one decorator per wrapper, plus the package README and common/__init__ docstrings). The mechanism is 281 lines in the new module, and the tests are 343 lines. JIT CUDA (tvm_ffi) wrappers such as residual_gate_add are not covered, because their calls have no common hook to record.

Accuracy Tests

Replay calls the same compiled kernels with the same arguments, so outputs are bit-identical. Output image md5 matched main in every end-to-end configuration below.

test/registered/kernels/ops/diffusion/test_fast_launch.py:

• The decorated wrappers. For each of the 8, a replayed call equals the wrapper's own result bit for bit, and Triton's JIT is never called. This runs under both no_grad and inference_mode.
• The recording rules, using a toy Triton kernel:
  • a 16-byte-misaligned operand records its own launch;
  • a function that does aten compute, launches on a temporary, reads a tensor value or takes a list argument stays on the normal path;
  • an in-place launch, with or without a returned result, replays;
  • nested decorated functions replay as one;
  • scalars are keyed by type;
  • views replay under inference mode.
• Graph capture, compile and the off switch. A recorded launch replays inside CUDA graph capture, and a new signature records there without a sync. torch.compile(fullgraph=True) traces the wrapper. The off switch runs the wrapper.

Passing on H200 (sm90) and GB300 (sm103), together with every existing suite that calls a decorated wrapper:

• test_norm, test_model_fast_paths, test_activation, test_flux2_fp8_norm_quant, test_flux2_strided_qknorm_rope;
• test_layernorm_modulate, test_sites, test_shape_specialization, test_h3_vae_fused_norm, test_longcat_image_norm_modulate, test_fused_qk_rmsnorm_rope_gate;
• test_zimage_qknorm_fusion, test_qwen_image21_cuda, test_minimax_h3_vae_parallel_modes.

Speed Tests and Profiling

Base 8afe107, torch 2.13, triton 3.7.1.

Host time per call (µs, 256²-class shapes):

wrapper	H200 main	H200 this PR	GB300 main	GB300 this PR
rmsnorm_scale	27.95	13.55	30.79	14.71
rmsnorm_tanh_residual	32.02	14.97	34.03	15.60
zimage_qk_rmsnorm_native	24.11	12.46	25.67	13.68
fused_layernorm_modulate_raw	31.32	14.40	33.16	15.82
fused_silu_mul_bitexact	19.76	14.28	29.32	21.74
fused_packed_silu_mul_bitexact	20.99	11.24	23.22	12.99
flux2_strided_qknorm_rope	26.70	18.29	29.60	19.94
fused_rope_rotate_half_bitexact	22.05	15.73	31.29	22.46

With SGLANG_DIFFUSION_TRITON_FAST_LAUNCH=0 the numbers match main.

End to end (client wall clock, median, eager):

GB300, one GPU. Servers of main and this PR are resident on the same GPU, requests are rotated between them, and processes are pinned to the GPU's NUMA node.

	main	this PR		this PR faster in
Z-Image-Turbo 512²	435.6 / 416.2 ms	374.7 / 371.0 ms	−14.0% / −10.9%	15 and 20 rotations
Z-Image-Turbo 1024²	447.9 ms	404.3 ms	−9.7%
FLUX.1-dev 512²	3219 ms	3098 ms	−3.8%	10/10
FLUX.1-dev 1024²	3211 ms	3092 ms	−3.7%	10/10
Ideogram-4 512²	1712 ms	1671 ms	−2.4%	10/10
Ideogram-4 1024²	2661 ms	2663 ms	no change	5/10

The previous revision of this PR applied the cache by hand inside each of the three Z-Image wrappers. It was measured in the same rotation: Z-Image 512² 357.5 / 355.1 ms, 1024² 405.2 ms. The decorator keeps most of that gain without touching the wrappers. Its per-call cost measured the same as the hand-written path, and every call replayed (counted in a live server).

GB300 ×2, Qwen-Image-2.1 TP2. main, this PR and a later JIT-replay experiment were resident together, with requests rotated over 8 rounds. The JIT-replay experiment was dropped: replay costs more than the thin tvm-ffi wrappers spend.

	main	this PR
512²	1546.7 ms	1354.5 ms	−12.4%
1024²	1625.5 ms	1461.8 ms	−10.1%

B200 ×2, TP2 / Ulysses2 (x86 host, AMD EPYC 9575F). Servers were swapped each round, 2 rounds.

• FLUX.1-dev 512²: main 1335–1357 ms, this PR 1240–1297 ms (−5 to −8%). 1024² is bimodal on both, and this PR is at or below main.
• Ideogram-4, Qwen-Image-2.1 and FLUX.2-dev: within noise at both sizes. These cells are GPU-bound on this host.

H200, Z-Image one GPU: no change (512² 236.5 vs 234.9 ms; 1024² 862.7 vs 862.4 ms). The step is GPU-bound there.

The gain shows up where the step is host-bound: on Grace hosts, and at low resolution with several ranks. The JIT CUDA wrappers that remain (residual_gate_add, qknorm_complex_rope_cuda, modulate_scale_shift) and custom_all_reduce are the next host-time items for FLUX/Qwen/Ideogram.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. (Package README layout line.)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): 🚫 Run #38355101182
Latest PR Test (Extra): ❌ Run #38355100702
Latest PR Test (AMD ROCm 10): ❌ Run #38355101222

[mickqian]

/tag-and-rerun-ci