diff --git a/docs/docs.json b/docs/docs.json index 8ece0ad7e..fa7fe0a0b 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -1608,6 +1608,7 @@ "docs/sglang-diffusion/performance-optimization", "docs/sglang-diffusion/deployment_cookbook", "docs/sglang-diffusion/attention_backends", + "docs/sglang-diffusion/fused_kernels", "docs/sglang-diffusion/parallelism", "docs/sglang-diffusion/ring_sp_performance", "docs/sglang-diffusion/encoder_parallel", diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx new file mode 100644 index 000000000..592e5a125 --- /dev/null +++ b/docs/docs/sglang-diffusion/fused_kernels.mdx @@ -0,0 +1,193 @@ +--- +title: "Fused Kernels" +description: "The fused CUDA/Triton kernels SGLang Diffusion ships, what each one replaces, and which are on by default." +tag: "preserve" +--- + +Diffusion transformers and VAEs spend a large share of their non-GEMM time on short elementwise chains — adaLN modulate, residual gating, QK-norm, RoPE, norm epilogues — each of which is a separate kernel launch and a separate HBM round trip in eager PyTorch. SGLang Diffusion replaces these chains with fused kernels under [`sglang/kernels/ops/diffusion`](https://github.com/sgl-project/sglang/tree/main/python/sglang/kernels/ops/diffusion). + +This page is an inventory: what each kernel fuses, what its numerical contract is, and which models use it. It is not a lever you tune — most of these kernels are on by default and require no flag. The one switch is `--quality`, described below. + +## Two numerical contracts + +Multi-step denoising amplifies a per-step rounding difference into visible quality loss, so "close enough" and "bit-exact" are different products here. Every kernel in the package falls into one of two classes. + +**Bit-exact — mounted unconditionally.** The kernel reproduces every rounding boundary of the eager chain, so `torch.equal` holds against the reference. Some go quite far to get there: the fused LayerNorm+modulate kernel replicates PyTorch's `vectorized_layer_norm_kernel` down to its Welford update order, guarded reciprocal, and warp-fold tree; the fused RMSNorm+scale/shift kernel replicates FlashInfer's CuTe-DSL `RMSNormKernel` fragment order and `shfl.bfly` fold. Because the dispatch they replicate can change underneath them, each one still verifies itself against the live eager chain on first sight and falls back permanently on any mismatch. + +**Not bit-exact — request-gated.** These differ from eager only at half-precision rounding-order level, but that is enough to matter, so they are mounted only for `quality="high"` requests, at batch boundaries, all-or-nothing per transformer. The default `quality="lossless"` runs the unmodified reference chain. + + +A plain fp32 single-pass norm fusion looks harmless and is not. On ERNIE-Image it moved the 50-step trajectory to 18.83 dB PSNR at `quality=high`, which is what motivated the bit-exact rewrite of that path. + + +## Enabling the request-gated set + +```bash +sglang generate --model-path MODEL_PATH --prompt "..." --quality high +``` + +The server default stays `lossless`; the OpenAI-compatible endpoints carry it per request. Images: + +```bash +curl -X POST http://${HOST}:${PORT}/v1/images/generations \ + -H 'Content-Type: application/json' \ + -d '{"model": "MODEL_PATH", "prompt": "...", "quality": "high"}' +``` + +Video, same field: + +```bash +curl -X POST http://${HOST}:${PORT}/v1/videos \ + -H 'Content-Type: application/json' \ + -d '{"model": "MODEL_PATH", "prompt": "...", "quality": "high"}' +``` + + +The `quality` field in a **video response** body is unrelated. It is Sora-compatible response metadata and is always reported as `"standard"`; it does not reflect the sampling quality the request ran with. + + +`quality` participates in the dynamic-batch signature, so mixed-quality traffic is batched separately and the transition happens safely at a batch boundary. Mounting is all-or-nothing: if any marked site on a transformer fails its static guards, no site on that transformer is fused. + +These fusion families mount under `quality="high"`: + +| Fusion | What it folds | +| --- | --- | +| Linear + tanh-GELU | Bias-add and GELU into the GEMM epilogue (cublasLt), removing the `[tokens, 4*dim]` intermediate round trip | +| LayerNorm + modulate | `layer_norm(x, weight=(1 + scale), bias=shift)` in place of affine-free LN plus a separate modulate | +| LTX-2 RMSNorm + modulate | `rms_norm(x) * (1 + scale) + shift` in one launch | +| Gate RMSNorm (BF16-native) | `RMSNorm + tanh + mul + add` in one pass | +| HunyuanVideo strided QK RMSNorm | Per-head QK RMSNorm over the packed QKV layout | + +## Kernel inventory + +34 operators are registered in the kernel registry across 38 implementations (some operators carry several backends). Backends are named by provenance, not device: `JIT` compiles under nvcc *and* hipcc, `TRITON` runs on CUDA and ROCm, `CUTE_DSL` needs CUTLASS, `FLYDSL` is ROCm gfx950 only, `AOT` comes from the `sgl_kernel` wheel. + +### Normalization + +| Operator | Backend | Contract | Replaces | +| --- | --- | --- | --- | +| `rmsnorm_scale_shift` | Triton | bit-exact | RMSNorm + `* (1 + scale) + shift` (4 kernels) | +| `scale_residual_norm_scale_shift` | Triton / CuTe-DSL / FlyDSL | bit-exact (Triton) | the above plus the preceding `residual + gate * update` | +| `layernorm_modulate` | Triton | bit-exact | affine-free LayerNorm + adaLN modulate | +| `qk_head_layernorm` | Triton | bit-exact | per-head LayerNorm on q/k | +| `qk_rmsnorm_native` | Triton | bit-exact | Z-Image per-head QK RMSNorm | +| `norm_scale_shift` | CuTe-DSL / FlyDSL | fp32 statistics | LN-or-RMS + scale/shift, many broadcast modes | +| `rmsnorm_scale`, `rmsnorm_tanh_residual` | Triton | bf16-native statistics | `RMSNorm(x) * scale`, `x + tanh(gate) * RMSNorm(y)` | +| `apply_group_norm_silu` | Triton | close | `GroupNorm + SiLU`, NCHW-contiguous | +| `group_norm_silu_4d`, `group_norm_silu_rows` | Triton | close | channels-last GroupNorm(+SiLU); what lets a VAE decoder run channels_last end to end with no `nchwToNhwc` transposes | +| `wan_rmsnorm_silu` | Triton | close | Wan VAE `channels_last_3d` RMSNorm + SiLU | + +### adaLN modulation and gating + +| Operator | Backend | Contract | Replaces | +| --- | --- | --- | --- | +| `modulate_scale_shift` | JIT CUDA | bit-exact | `x * (1 + scale) + shift` | +| `residual_gate_add` | JIT CUDA | bit-exact | `residual + gate * update` | +| `timestep_embedding` | JIT CUDA | close | sinusoidal timestep embedding | +| `temb_table_slices` | Triton | bit-exact | see note below | +| `ltx2_ada_values` | Triton | bit-exact | LTX-2 nine-way adaLN value split, slices come out contiguous | + + +`temb_table_slices` is worth knowing about. The eager `(scale_shift_table + temb.float()).chunk(6, dim=2)` materializes roughly 8 GB of fp32 at 704p/121f *and* hands six **strided** slices downstream, whose `.contiguous()` calls then copy each one again. The fused kernel produces the six slices in one pass, each naturally contiguous, so the downstream copies become no-ops. + + +### RoPE and QK-norm + +| Operator | Backend | Contract | Replaces | +| --- | --- | --- | --- | +| `fused_inplace_qknorm_rope` | JIT CUDA | one bf16 rounding step vs the split baseline; exact with `round_norm_before_rope=True` | separate QK-norm kernel + RoPE | +| `rope_rotate_half` | Triton | bit-exact | `chunk` → `cat(-x2, x1)` → two muls + add → `cat(tail)`, about 7 kernels per projection | +| `ltx2_qknorm_split_rope` | JIT CUDA | close (validated on B200) | LTX-2 QK-norm + split RoPE | +| `hunyuan_qkv_rope_pack` | Triton | bit-exact | QKV pack and RoPE in one pass | + +### Activation + +| Operator | Backend | Contract | Replaces | +| --- | --- | --- | --- | +| `silu_mul` | Triton | bit-exact | `F.silu(a) * b` for split-projection SwiGLU, where the concatenated `silu_and_mul` kernels do not apply without an extra full-width `cat` | +| `bias_silu`, `bias_glu` | Triton | bit-exact | Sana GLUMB conv bias + SiLU / GLU post-processing | +| `linear_gelu_tanh` | AOT (cublasLt) | not bit-exact, request-gated | bias-add and tanh-GELU folded into the GEMM epilogue | + +### Attention + +| Operator | Backend | Notes | +| --- | --- | --- | +| `sparse_linear_attn_fwd` | Triton | block-map, compression, and forward for sparse linear attention | +| `bigdn` | Triton | Sana-WM bidirectional gated delta-net; the chunkwise form splits phase A along the KV and Z streams so two blocks stay resident per SM, and stores `(I - P)` so phase B's MMA folds the identity-add in | + +### Data movement + +Every kernel here only moves values (plus zero fill, plus at most one same-order add), so each is bitwise identical to the aten chain it replaces. + +| Operator | Backend | Replaces | +| --- | --- | --- | +| `usp_merge_heads` | JIT CUDA | USP all-to-all output head merge (`permute` + `contiguous`) | +| `pack_qkv_destination_major` | Triton | Ulysses destination-major QKV pack | +| `varlen_pack_qkv`, `varlen_scatter_to_padded` | Triton | varlen gather/scatter around the masked attention path | +| `causal_conv3d_cat_pad` | JIT CUDA / Triton | causal Conv3d `cat` + `pad` | +| `cat_pad_channels_last_3d` | Triton | Wan causal VAE `cat + F.pad + contiguous` (three passes plus cache bookkeeping) in one pass | +| `dup_up3d_add` | Triton | `repeat_interleave + permute().contiguous() + add` | + +## Coverage by model + +Kernels are written against a specific eager chain in a specific model, so coverage is per-model rather than universal. + +| Model | Fused paths | +| --- | --- | +| FLUX.1 | LN+modulate, modulate, residual-gate add, linear+GELU | +| FLUX.2 | LN+modulate, packed SwiGLU, residual-gate add | +| Qwen-Image | linear+GELU, select-0/1 LN modulation | +| GLM-Image | LN+modulate, per-head qk LN, residual-gate add, linear+GELU | +| ERNIE-Image | RMSNorm+scale/shift, residual-gated variant, rotate-half RoPE, residual-gate add | +| Z-Image | BF16-native RMSNorm scale / tanh-residual, per-head QK RMSNorm | +| Ideogram 4 | gate RMSNorm, SwiGLU, rotate-half RoPE, modulate, residual-gate add | +| LTX-2 | QK-norm + split RoPE, ada-values split, RMSNorm+modulate, modulate, residual-gate add, linear+GELU | +| HunyuanVideo | QKV+RoPE pack, strided QK RMSNorm, linear+GELU | +| Sana | LN+modulate, GLUMB bias+SiLU / bias+GLU, residual-gate add | +| Sana-WM | bidirectional gated delta-net, fused QK inverse-RMS | +| Wan | temb table slices; VAE cat+pad and DupUp3D add, `channels_last_3d` RMSNorm+SiLU | +| Cosmos3 / Krea2 / MiniMax-H3 | QK-norm + RoPE (Krea2 also CuTe-DSL norm+scale/shift; MiniMax-H3 also indexed modulation) | +| FLUX.2 VAE / HunyuanVAE / latent upsampler | GroupNorm + SiLU (channels-last two-pass for FLUX.2) | + +## Inspecting what is registered + +Every kernel is described by a `KernelSpec` in the process-wide registry, so the inventory is queryable without importing any backend: + +```python +from sglang.kernels.registry import registry + +diffusion_ops = [op for op in registry.ops() if op.startswith("diffusion.")] +for spec in registry.get("diffusion.scale_residual_norm_scale_shift"): + print(spec.backend, spec.target, spec.capabilities) +``` + +Registration is metadata only — it imports neither torch nor a backend and triggers no JIT build. To pick a specific implementation of an operator that has several: + +```python +from sglang.kernels import select_kernel, KernelBackend + +fn = select_kernel( + "diffusion.scale_residual_norm_scale_shift", backend=KernelBackend.CUTE_DSL +).load() +``` + +## Importing the kernels + +Runtime code imports from the package, never from a submodule: + +```python +from sglang.kernels.ops.diffusion import fused_rmsnorm_scale_shift_bitexact +``` + +Resolution is lazy: the backends have disjoint, heavy dependencies (Triton, CUTLASS/CuTe-DSL, FlyDSL on ROCm, MLX on Apple), so an eager re-export would make every one of them an import-time requirement on every platform. Each public kernel is a predicate-plus-kernel pair — call `can_use_(...)` first and fall back to the reference chain when it returns `False`; the kernel raises on an unsupported input rather than silently returning `None`. + +The package `README.md` carries a selection matrix for the cases where several kernels look interchangeable and are not. The normalization domain alone holds more than a dozen implementations that differ by numerical contract, activation layout, and backend rather than by speed. + +## References + +- [Performance Optimization](./performance-optimization) +- [Attention Backends](./attention_backends) +- [Quantization](./quantization) +- [Profiling](./profiling) +- [`sglang/kernels/ops/diffusion`](https://github.com/sgl-project/sglang/tree/main/python/sglang/kernels/ops/diffusion) — source and selection matrix +- [RFC #29630](https://github.com/sgl-project/sglang/issues/29630) — the unified `sglang.kernels` namespace diff --git a/docs/docs/sglang-diffusion/performance-optimization.mdx b/docs/docs/sglang-diffusion/performance-optimization.mdx index 84234d0b3..e0206083d 100644 --- a/docs/docs/sglang-diffusion/performance-optimization.mdx +++ b/docs/docs/sglang-diffusion/performance-optimization.mdx @@ -61,6 +61,11 @@ These settings should preserve model behavior while changing residency, parallel Kernel choice dominates DiT latency or memory. Attention Backends + + Fused kernels + You want to know which elementwise chains are already fused, or to opt into the request-gated set. + Fused Kernels + Dynamic batching Serving many compatible requests concurrently. @@ -126,6 +131,7 @@ These techniques can change the denoising path, numerical representation, or gen - [Deployment and Performance Modes](./deployment_cookbook) - [Attention Backends](./attention_backends) +- [Fused Kernels](./fused_kernels) - [Sequence Parallelism](./ring_sp_performance) - [Caching Strategies](./caching-acceleration) - [Profiling](./profiling)