chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Diffusion] Add SGL-D diffusion efficient kernel skills (#19473)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-02-27 11:42:13 +08:00
committed by GitHub
parent 52c8a3632a
commit fe971b620a
1 changed files with 111 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
.claude/skills/diffusion/use-efficient-diffusion-kernels.md Normal file
View File

@@ -0,0 +1,111 @@
---
name: use-efficient-diffusion-kernels
description: Guidance for using SGLang Diffusion fused kernels and fast CUDA paths. Use when mapping fusion patterns in diffusion inference, choosing fused ops or attention backends, handling RoPE/QK norm performance pitfalls, or integrating new diffusion models with kernel-aware constraints.
---
# Use Efficient Diffusion Kernels
**Overview**
This skill focuses on SGLang Diffusion (`sglang.multimodal_gen`) kernel fusion patterns and fast CUDA paths. Prefer existing fused ops (Triton, CuTe DSL, sgl-kernel). Make constraints and fallbacks explicit.
**Key Files**
- `python/sglang/multimodal_gen/runtime/layers/layernorm.py`
- `python/sglang/multimodal_gen/runtime/layers/elementwise.py`
- `python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py`
- `python/sglang/jit_kernel/diffusion/triton/scale_shift.py`
- `python/sglang/jit_kernel/diffusion/triton/norm.py`
- `python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py`
- `python/sglang/jit_kernel/diffusion/triton/rotary.py`
- `python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py`
- `python/sglang/jit_kernel/norm.py`
- `python/sglang/multimodal_gen/runtime/platforms/cuda.py`
- `python/sglang/multimodal_gen/runtime/layers/attention/selector.py`
- `docs/diffusion/performance/attention_backends.md`
**Core Fusion Patterns**
1. Scale/Shift elementwise fusion (AdaLN modulation)
- Kernels: `fuse_scale_shift_kernel`, `fuse_scale_shift_gate_select01_kernel`
- Locations: `elementwise.py`, `layernorm.py`, `qwen_image.py`, `triton/scale_shift.py`
- Use cases: `x * (1 + scale) + shift` and `a * (k + b) + c`
- Constraints: `x` must be CUDA and contiguous. `scale/shift` support 0D/1D/2D/3D/4D broadcast. 4D `[B, F, 1, C]` requires `L % F == 0`.
- NPU fallback: `scale_shift.py` swaps to `npu_fallback` native path.
2. Norm + Scale/Shift fusion (CuTe DSL)
- Kernels: `fused_norm_scale_shift`, `fused_scale_residual_norm_scale_shift`
- Locations: `layernorm.py`, `cutedsl/scale_residual_norm_scale_shift.py`
- Use cases:
- `y = norm(x) * (1 + scale) + shift`
- `y = norm(residual + gate * x) * (1 + scale) + shift`
- Constraints: `D % 256 == 0` and `D <= 8192`. `x/residual/gate/scale/shift` must pass shape and stride validation. Dtypes limited to fp16/bf16/fp32.
- Behavior: CuTe DSL compilation cached by `(dtype, ndim, D, norm_type)`. `None` tensors replaced by scalar placeholders. If constraints fail, `layernorm.py` warns and falls back to native PyTorch.
3. Triton LayerNorm/RMSNorm fusion
- Kernels: `rms_norm_fn`, `layer_norm_fn`, `norm_infer`
- Locations: `triton/norm.py`, `layernorm.py`
- Use cases: fp32 RMSNorm with residual/dropout/rowscale/x1 branches, and inference-friendly `norm_infer`.
- Constraints: last dim must be contiguous, and `N * element_size < 64KB`.
4. Triton one-pass RMSNorm (small hidden size fast path)
- Kernel: `triton_one_pass_rms_norm`
- Locations: `triton/rmsnorm_onepass.py`, `layernorm.py`
- Use case: `hidden_size <= 128` in `RMSNorm.forward_cuda`.
5. Triton RoPE fusion
- Kernel: `apply_rotary_embedding`
- Locations: `triton/rotary.py`, `rotary_embedding/utils.py`
- Use case: GPT-J style RoPE when not Neox.
- Constraints: `head_size` must be even.
- NPU fallback: `npu_fallback.apply_rotary_embedding_native`.
**Faster CUDA Kernel Usage Points**
1. sgl-kernel RMSNorm and fused add RMSNorm
- Location: `layernorm.py`
- Behavior: CUDA uses `sgl_kernel.fused_add_rmsnorm` and `sgl_kernel.rmsnorm`. `hidden_size <= 128` uses Triton one-pass. ROCm falls back to native.
2. Attention backend selection (FlashAttention, Sage, SDPA)
- Locations: `platforms/cuda.py`, `attention/selector.py`, `docs/diffusion/performance/attention_backends.md`
- Behavior: CUDA prefers FlashAttention (FA3/FA4) when supported, otherwise Torch SDPA. Force via `--attention-backend` or `global_force_attn_backend`.
3. FlashInfer RoPE (Q/K inplace)
- Location: `rotary_embedding/utils.py`
- Behavior: `flashinfer.rope.apply_rope_with_cos_sin_cache_inplace` when available, otherwise Triton RoPE fallback.
**QK Norm Optimization**
- Entry point: `apply_qk_norm` in `layernorm.py`.
- Fast path: JIT fused inplace QK norm from `python/sglang/jit_kernel/norm.py` via `fused_inplace_qknorm`.
- Preconditions for fused path:
- CUDA only.
- `allow_inplace=True` and `q_eps == k_eps`.
- `can_use_fused_inplace_qknorm(head_dim, dtype)` returns true.
- Supported head dims: `64, 128, 256, 512, 1024`.
- Behavior: Fused path operates on `q` and `k` in place after reshaping to `[B, -1, head_dim]`. If preconditions fail, fall back to per-tensor RMSNorm.
**Common Entry Points in Diffusion Models**
- AdaLN modulation: `LayerNormScaleShift`, `RMSNormScaleShift`, `ScaleResidual*` in `layernorm.py`.
- Qwen-Image gating: `fuse_scale_shift_gate_select01_kernel` in `qwen_image.py`.
- QK norm: `apply_qk_norm` used in `flux.py`, `flux_2.py`, `qwen_image.py`, `zimage.py`, `wanvideo.py`, `ltx_2.py`, `hunyuanvideo.py`.
- RoPE: `_apply_rotary_emb` prefers Triton; Q/K RoPE prefers FlashInfer when present.
**Constraints and Fallbacks**
- `scale_shift` Triton requires CUDA + contiguous `x`. NPU swaps to native.
- CuTe DSL fused norms require `D % 256 == 0` and `D <= 8192`.
- Triton norm kernels error on feature size >= 64KB.
- FlashAttention requires fp16/bf16 and SM80+; otherwise SDPA.
**Integration Checklist for New Models**
1. Reuse `LayerNormScaleShift` or `ScaleResidual*` modules instead of re-implementing fusion logic.
2. Keep tensors contiguous and satisfy D alignment (`% 256`) and size (`<= 8192`) for CuTe fused paths.
3. Use `fuse_scale_shift_kernel` for AdaLN modulation and keep a PyTorch fallback.
4. Use `apply_qk_norm` and ensure head_dim is in the supported list for fused QK norm.
5. If using FlashInfer RoPE, avoid `pack qkv` and ensure Q/K are contiguous.
6. For attention, follow `selector.py` priority; override with CLI only if needed.
**When Extending or Modifying Kernels**
- Add `torch.library.custom_op` and `register_fake` for compile and meta support.
- Keep CuTe compile cache keys aligned to `(dtype, ndim, D)`.
- Avoid implicit broadcasts that force hidden `contiguous()` copies.
- Preserve NPU and ROCm fallback paths.