[docs] Add a fused-kernels page for SGLang Diffusion (#35436)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Claude Opus 5
Claude Opus 5
parent
8a1e6e4e46
commit
9113fc6d93
@@ -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",
|
||||
|
||||
@@ -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.
|
||||
|
||||
<Note>
|
||||
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.
|
||||
</Note>
|
||||
|
||||
## 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"}'
|
||||
```
|
||||
|
||||
<Warning>
|
||||
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.
|
||||
</Warning>
|
||||
|
||||
`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 |
|
||||
|
||||
<Tip>
|
||||
`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.
|
||||
</Tip>
|
||||
|
||||
### 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_<op>(...)` 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
|
||||
@@ -61,6 +61,11 @@ These settings should preserve model behavior while changing residency, parallel
|
||||
<td style={{padding: "9px 12px"}}>Kernel choice dominates DiT latency or memory.</td>
|
||||
<td style={{padding: "9px 12px"}}><a href="./attention_backends">Attention Backends</a></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500}}>Fused kernels</td>
|
||||
<td style={{padding: "9px 12px"}}>You want to know which elementwise chains are already fused, or to opt into the request-gated set.</td>
|
||||
<td style={{padding: "9px 12px"}}><a href="./fused_kernels">Fused Kernels</a></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500}}>Dynamic batching</td>
|
||||
<td style={{padding: "9px 12px"}}>Serving many compatible requests concurrently.</td>
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user