From 3aaa92d2a2a5303adec471b29efba1f2c224eccf Mon Sep 17 00:00:00 2001 From: BBuf <1182563586@qq.com> Date: Wed, 19 Aug 2026 11:08:10 +0800 Subject: [PATCH 1/2] [docs] Add a fused-kernels page for SGLang Diffusion The diffusion kernels that replace eager elementwise chains -- adaLN modulate, residual gating, QK-norm, RoPE, norm epilogues -- had no user-facing page. What exists is a package README aimed at kernel authors, so there was nowhere to point someone who wants to know what is already fused, whether it changes their output, or how to turn on the parts that do. Adds `sglang-diffusion/fused_kernels` under Performance Optimization: - The two numerical contracts the package is organized around: bit-exact fusions mounted unconditionally (each self-verifying against the live eager chain on first sight), and the ones that differ at half-precision rounding level and therefore mount only for `quality="high"`. - How to opt into the request-gated set, with the CLI and the OpenAI-compatible request field. - An inventory by operator domain -- what each of the 34 registered operators replaces, on which backend, under which contract. - Which models use which fused paths, since these kernels are written against a specific eager chain in a specific model rather than being general-purpose operators. - How to query the registry and pin a backend, and the import contract. Every table row is derived from the merged tree rather than written from memory: the operator list and backends come from the registry `_SPECS`, the model coverage from the actual import sites under `multimodal_gen/runtime`, and both Python snippets were executed against the registry before landing. Cross-linked from the Performance Optimization overview. Co-Authored-By: Claude Opus 5 (1M context) --- docs/docs.json | 1 + docs/docs/sglang-diffusion/fused_kernels.mdx | 181 ++++++++++++++++++ .../performance-optimization.mdx | 6 + 3 files changed, 188 insertions(+) create mode 100644 docs/docs/sglang-diffusion/fused_kernels.mdx diff --git a/docs/docs.json b/docs/docs.json index 8ece0ad7efb6..fa7fe0a0b309 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 000000000000..bfb43e8166e3 --- /dev/null +++ b/docs/docs/sglang-diffusion/fused_kernels.mdx @@ -0,0 +1,181 @@ +--- +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 --prompt "..." --quality high +``` + +The server default stays `lossless`; the OpenAI-compatible endpoints carry it per request: + +```bash +curl -X POST http://${HOST}:${PORT}/v1/images/generations \ + -H 'Content-Type: application/json' \ + -d '{"model": "", "prompt": "...", "quality": "high"}' +``` + +`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 84234d0b36ee..e0206083d937 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) From ab914bf65e936befac1b85be9abd75d82acf6c6f Mon Sep 17 00:00:00 2001 From: BBuf <1182563586@qq.com> Date: Wed, 19 Aug 2026 12:01:48 +0800 Subject: [PATCH 2/2] [docs] Document quality=high on the video endpoint `quality` reaches SamplingParams on /v1/videos as well, just by a different route than on images: VideoGenerationsRequest sets extra="allow" and does not declare the field, so it arrives in model_extra and video_api's _extra_value() forwards it into build_sampling_params. The image request declares it and maps OpenAI's "auto" default to None first. Also flags the confusable one: VideoResponse.quality is Sora-compatible response metadata, hardcoded to "standard", and says nothing about the sampling quality the request ran with. --- docs/docs/sglang-diffusion/fused_kernels.mdx | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx index bfb43e8166e3..592e5a125e21 100644 --- a/docs/docs/sglang-diffusion/fused_kernels.mdx +++ b/docs/docs/sglang-diffusion/fused_kernels.mdx @@ -23,17 +23,29 @@ A plain fp32 single-pass norm fusion looks harmless and is not. On ERNIE-Image i ## Enabling the request-gated set ```bash -sglang generate --model-path --prompt "..." --quality high +sglang generate --model-path MODEL_PATH --prompt "..." --quality high ``` -The server default stays `lossless`; the OpenAI-compatible endpoints carry it per request: +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": "", "prompt": "...", "quality": "high"}' + -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"`: