Refine diffusion skills and align JIT kernel docs with the new CI flow (#21283)
This commit is contained in:
Xiaoyu Zhang
@@ -98,7 +98,7 @@ python -m pytest python/sglang/multimodal_gen/test/
|
||||
|
||||
## Performance Tuning
|
||||
|
||||
For questions about optimal performance, fastest commands, VRAM reduction, or best flag combinations for a given model/GPU setup, **read the [diffusion-optimal-perf skill](skills/diffusion-optimal-perf/SKILL.md)**. It contains a complete table of all lossless and lossy optimization flags with trade-offs, quick recipes, and tips.
|
||||
For questions about optimal performance, fastest commands, VRAM reduction, or best flag combinations for a given model/GPU setup, **read the [sglang-diffusion-performance skill](skills/sglang-diffusion-performance/SKILL.md)**. It contains a complete table of lossless and lossy optimization flags with trade-offs, quick recipes, and tuning tips.
|
||||
|
||||
### Perf Measurement
|
||||
|
||||
|
||||
@@ -1,79 +0,0 @@
|
||||
---
|
||||
name: diffusion-kernel
|
||||
description: Index for SGLang Diffusion kernel development skills.
|
||||
---
|
||||
|
||||
# Diffusion Kernel Skills
|
||||
|
||||
## Rule: Follow User Kernel Language Preference
|
||||
|
||||
If the user explicitly states a preference for **Triton** or **CUDA**, follow that preference when implementing and optimizing kernels (even if the other option could work). Do not “pick for convenience”.
|
||||
|
||||
## Directory Layout
|
||||
|
||||
```
|
||||
python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/
|
||||
├── SKILL.md
|
||||
├── add-triton-kernel.md
|
||||
├── add-cuda-kernel.md
|
||||
├── diffusion-benchmark-and-profile.md
|
||||
├── nsight-profiler.md
|
||||
├── use-efficient-diffusion-kernels.md
|
||||
├── references/
|
||||
│ ├── kernel-templates.md # Copy-paste CUDA kernel templates (sglang JIT style)
|
||||
│ ├── troubleshooting.md # Build/perf/integration issues & fixes
|
||||
│ ├── h100-optimization-guide.md # H100 (sm_90) deep dive
|
||||
│ ├── a100-optimization-guide.md # A100 (sm_80) deep dive
|
||||
│ └── t4-optimization-guide.md # T4 (sm_75, FP16 only) deep dive
|
||||
└── scripts/
|
||||
├── bench_diffusion_rmsnorm.py # RMSNorm micro-benchmark vs PyTorch
|
||||
└── bench_diffusion_denoise.py # End-to-end denoise benchmark (sglang generate)
|
||||
```
|
||||
|
||||
## Index
|
||||
|
||||
Before running any benchmark, profiler, or kernel-validation command:
|
||||
- use `scripts/diffusion_skill_env.py` to derive the repo root from `sglang.__file__`
|
||||
- verify the repo is writable
|
||||
- export `FLASHINFER_DISABLE_VERSION_CHECK=1`
|
||||
- choose idle GPU(s) before starting perf work
|
||||
|
||||
- [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py)
|
||||
|
||||
Shared preflight helper for all diffusion skill commands. Use it to print the repo root, create benchmark/profile output directories, and choose idle GPUs before running `sglang generate`, torch profiler, nsys, or ncu.
|
||||
|
||||
- [add-triton-kernel.md](./add-triton-kernel.md)
|
||||
|
||||
Step-by-step guide for adding a new Triton kernel to SGLang Diffusion's `jit_kernel/diffusion/triton/` module, including authoring, autotune, `torch.compile` compatibility, integration, and tests. Use for fused elementwise ops, norm variants, RoPE variants, or when NPU/CPU fallback is needed.
|
||||
|
||||
- [add-cuda-kernel.md](./add-cuda-kernel.md)
|
||||
|
||||
Step-by-step guide for adding a JIT CUDA kernel. CUDA source goes in `jit_kernel/csrc/diffusion/<op>.cuh`; Python wrapper at `jit_kernel/diffusion/<op>.py`. Uses SGLang's JIT compilation system (`load_jit`, `cache_once`) and internal abstractions (`TensorMatcher`, `device::AlignedVector`, `host::LaunchKernel`, `device::warp::reduce_sum`). Use for bandwidth-bound reductions (RMSNorm, LayerNorm) or ops needing fine-grained vectorization and shared memory control. Adapted from [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels).
|
||||
|
||||
- [use-efficient-diffusion-kernels.md](./use-efficient-diffusion-kernels.md)
|
||||
|
||||
Practical guidance for using SGLang Diffusion fused kernels and fast CUDA paths, including constraints, fallbacks, and where the fused ops are wired into the runtime.
|
||||
|
||||
- [diffusion-benchmark-and-profile.md](./diffusion-benchmark-and-profile.md)
|
||||
|
||||
Denoise-stage benchmark and profiling guide for SGLang Diffusion models. Three profiling levels: Level 1 (torch.profiler — kernel time ranking), Level 2 (nsys — category breakdown), Level 3 (ncu — per-kernel bandwidth/occupancy/roofline analysis). **ncu is critical for kernel optimization** — always use it when writing or tuning custom kernels to verify hardware saturation.
|
||||
|
||||
- [nsight-profiler.md](./nsight-profiler.md)
|
||||
|
||||
Advanced profiling skill for NVIDIA Nsight Systems / Nsight Compute: collecting traces, reading reports, and interpreting kernel-level performance metrics.
|
||||
|
||||
## References (GPU optimization guides, templates, troubleshooting)
|
||||
|
||||
Loaded by `add-cuda-kernel.md`. Adapted from [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels).
|
||||
|
||||
- [references/kernel-templates.md](references/kernel-templates.md) — copy-paste ready sglang JIT CUDA templates: element-wise (SiLU), row-reduction (RMSNorm), fused AdaLN, Python wrapper, test, benchmark
|
||||
- [references/troubleshooting.md](references/troubleshooting.md) — build errors, performance issues, torch.compile compatibility, kernel injection pitfalls
|
||||
- [references/h100-optimization-guide.md](references/h100-optimization-guide.md) — H100 (sm_90): AlignedVector benchmarks, warp reductions, occupancy, TMA, PDL
|
||||
- [references/a100-optimization-guide.md](references/a100-optimization-guide.md) — A100 (sm_80): cp.async, TF32, 2:4 sparsity, H100→A100 migration checklist
|
||||
- [references/t4-optimization-guide.md](references/t4-optimization-guide.md) — T4 (sm_75): FP16 only, 320 GB/s bandwidth, 64 KB shared mem, 16 GB memory management
|
||||
|
||||
## Scripts (runnable benchmarks)
|
||||
|
||||
- [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py) — preflight helper: repo root discovery via `sglang.__file__`, write-access probe, benchmark/profile output directories, idle GPU selection
|
||||
- [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark: JIT CUDA vs PyTorch, correctness check, bandwidth efficiency analysis
|
||||
- [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark preset runner via `sglang generate`; save perf dumps by label and compare them with `compare_perf.py`
|
||||
+6
-6
@@ -1,11 +1,11 @@
|
||||
---
|
||||
name: add-new-diffusion-model
|
||||
description: Step-by-step guide for adding a new diffusion model to SGLang. Covers the recommended Hybrid Monolithic pipeline pattern (BeforeDenoisingStage), as well as when to use the Modular Composition Style. Includes pipeline config, model components, registration, and testing.
|
||||
name: sglang-diffusion-add-model
|
||||
description: Use when adding a new diffusion model or Diffusers pipeline to SGLang.
|
||||
---
|
||||
|
||||
# Tutorial: Adding a New Diffusion Model to SGLang
|
||||
# Add a Diffusion Model to SGLang
|
||||
|
||||
This tutorial walks through adding support for a new diffusion model. SGLang Diffusion supports two pipeline styles; choose the one that best fits your model.
|
||||
Use this skill when adding a new diffusion model or pipeline variant to `sglang.multimodal_gen`.
|
||||
|
||||
## Two Pipeline Styles
|
||||
|
||||
@@ -126,7 +126,7 @@ class MyModelTransformer2DModel(nn.Module):
|
||||
"""DiT model for MyModel.
|
||||
|
||||
Adapt from the Diffusers/reference implementation. Key points:
|
||||
- Use SGLang's fused LayerNorm/RMSNorm ops (see use-efficient-diffusion-kernels skill)
|
||||
- Use SGLang's fused LayerNorm/RMSNorm ops (see `existing-fast-paths.md` under the benchmark/profile skill)
|
||||
- Use SGLang's attention backend selector
|
||||
- Keep the same parameter naming as Diffusers for weight loading compatibility
|
||||
"""
|
||||
@@ -566,7 +566,7 @@ Before submitting, verify:
|
||||
- [ ] `_required_config_modules` lists all modules from `model_index.json`
|
||||
- [ ] `PipelineConfig` callbacks (`prepare_pos_cond_kwargs`, `get_freqs_cis`, etc.) match DiT's `forward()` signature
|
||||
- [ ] Latent scale/shift factors are correctly configured
|
||||
- [ ] Use fused kernels where possible (see `use-efficient-diffusion-kernels` skill)
|
||||
- [ ] Use fused kernels where possible (see `existing-fast-paths.md` under the benchmark/profile skill)
|
||||
- [ ] Weight names match Diffusers for automatic loading
|
||||
- [ ] **TP/SP support** considered for DiT model (recommended; reference `wanvideo.py` for TP+SP, `qwen_image.py` for USPAttention)
|
||||
- [ ] **Output quality verified** — generated images/videos are not noise; compared against Diffusers reference output
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
---
|
||||
name: sglang-diffusion-benchmark-profile
|
||||
description: Use when benchmarking denoise latency or profiling a diffusion bottleneck in SGLang.
|
||||
---
|
||||
|
||||
# SGLang Diffusion Benchmark and Profile
|
||||
|
||||
Use this skill when measuring denoise performance, finding the slow op, checking whether an existing fast path can solve it, or verifying a kernel change in `sglang.multimodal_gen`.
|
||||
|
||||
This skill covers diagnosis and fast-path reuse:
|
||||
- To write a new Triton kernel, use [../sglang-diffusion-triton-kernel/SKILL.md](../sglang-diffusion-triton-kernel/SKILL.md)
|
||||
- To write a new CUDA JIT kernel, use [../sglang-diffusion-cuda-kernel/SKILL.md](../sglang-diffusion-cuda-kernel/SKILL.md)
|
||||
|
||||
## Preflight
|
||||
|
||||
Before running any benchmark, profiler, or kernel-validation command:
|
||||
- use `scripts/diffusion_skill_env.py` to derive the repo root from `sglang.__file__`
|
||||
- verify the repo is writable
|
||||
- export `HF_TOKEN` before using gated Hugging Face models such as `black-forest-labs/FLUX.*`
|
||||
- export `FLASHINFER_DISABLE_VERSION_CHECK=1`
|
||||
- choose idle GPU(s) before starting perf work
|
||||
|
||||
## Main Reference
|
||||
|
||||
- [benchmark-and-profile.md](benchmark-and-profile.md) — canonical denoise benchmark and profiling workflow; includes `torch.profiler`, `nsys`, and `ncu`
|
||||
- [existing-fast-paths.md](existing-fast-paths.md) — map bottlenecks to existing fused kernels and runtime fast paths before writing new code
|
||||
- [nsight-profiler.md](nsight-profiler.md) — Nsight Systems / Nsight Compute metric interpretation
|
||||
- [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py) — preflight helper: repo root discovery via `sglang.__file__`, write-access probe, benchmark/profile output directories, idle GPU selection
|
||||
- [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark: JIT CUDA vs PyTorch, correctness check, bandwidth efficiency analysis
|
||||
- [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark preset runner via `sglang generate`; save perf dumps by label and compare them with `compare_perf.py`
|
||||
+12
-10
@@ -1,6 +1,6 @@
|
||||
---
|
||||
name: diffusion-benchmark-and-profile
|
||||
description: Denoise-stage benchmark and per-layer kernel profiling guide for SGLang Diffusion models. Use when measuring denoising latency, profiling DiT kernel breakdown with torch.profiler or nsys+gputrc2graph.py, investigating performance bottlenecks, or optimizing with custom Triton/CUDA kernels. Always verify output correctness before and after any optimization.
|
||||
name: benchmark-and-profile-reference
|
||||
description: Reference commands and workflow for denoise benchmarks and profiling in SGLang Diffusion.
|
||||
---
|
||||
|
||||
# SGLang Diffusion Benchmark and Profile Guide
|
||||
@@ -18,11 +18,12 @@ description: Denoise-stage benchmark and per-layer kernel profiling guide for SG
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
ENV_PY=python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py
|
||||
ENV_PY=python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py
|
||||
ROOT=$(python3 "$ENV_PY" print-root)
|
||||
cd "$ROOT"
|
||||
python3 "$ENV_PY" check-write-access >/dev/null
|
||||
|
||||
export HF_TOKEN=<your_hf_token> # required for gated repos such as black-forest-labs/FLUX.*
|
||||
export FLASHINFER_DISABLE_VERSION_CHECK=1
|
||||
export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 1)
|
||||
|
||||
@@ -53,6 +54,7 @@ Environment notes:
|
||||
- **Level 1 profiling**: `torch.profiler` (bundled with torch).
|
||||
- **Level 2 profiling**: `nsys`, `pandas`, `plotly`, `regex`, and `gputrc2graph.py` from the sglang repo.
|
||||
- All commands below assume you are inside the configured diffusion container shell and already `cd`'d to the repo root derived from `sglang.__file__`.
|
||||
- Export `HF_TOKEN` before running any command against a gated Hugging Face repo such as `black-forest-labs/FLUX.*`. Without it, the top-level `sglang generate` auto-detection can fail before model loading and report a misleading `Generate subcommand is not yet supported for model ...`.
|
||||
- Export `FLASHINFER_DISABLE_VERSION_CHECK=1` before any benchmark or profiler command.
|
||||
- Re-run `print-idle-gpus` before each perf command if GPU availability may have changed.
|
||||
- Keep benchmark commands within 4 GPUs or fewer.
|
||||
@@ -479,10 +481,10 @@ After pinpointing the slow op, choose the right tool:
|
||||
|
||||
| Scenario | Skill to use |
|
||||
|----------|-------------|
|
||||
| New fused elementwise, norm variant, RoPE variant | **`add-triton-kernel.md`** — Triton JIT, faster iteration, NPU fallback |
|
||||
| Bandwidth-bound reduction (RMSNorm) needing max vectorization | **`add-cuda-kernel.md`** — CUDA JIT with `AlignedVector`, warp reductions |
|
||||
| Attention or tile-based op needing shared memory tuning | **`add-cuda-kernel.md`** — full control over CUDA primitives |
|
||||
| Slow op already covered by existing fused kernel | **`use-efficient-diffusion-kernels.md`** — check constraints & enable |
|
||||
| New fused elementwise, norm variant, RoPE variant | **`sglang-diffusion-triton-kernel`** — Triton JIT, faster iteration, NPU fallback |
|
||||
| Bandwidth-bound reduction (RMSNorm) needing max vectorization | **`sglang-diffusion-cuda-kernel`** — CUDA JIT with `AlignedVector`, warp reductions |
|
||||
| Attention or tile-based op needing shared memory tuning | **`sglang-diffusion-cuda-kernel`** — full control over CUDA primitives |
|
||||
| Slow op already covered by an existing fused kernel | **`existing-fast-paths.md`** — check constraints and enable it |
|
||||
|
||||
**Quick decision rule**: start with Triton. Switch to CUDA JIT only when profiling shows Triton can't saturate hardware bandwidth.
|
||||
|
||||
@@ -531,9 +533,9 @@ TORCH_COMPILE_DEBUG=1 sglang generate ...
|
||||
→ for CUDA graph: use --graph-profiling node
|
||||
↓
|
||||
5. KERNEL OPTIMIZATION
|
||||
Existing fused kernel? → use-efficient-diffusion-kernels.md
|
||||
New Triton kernel? → add-triton-kernel.md
|
||||
New CUDA JIT kernel? → add-cuda-kernel.md
|
||||
Existing fused kernel? → existing-fast-paths.md
|
||||
New Triton kernel? → sglang-diffusion-triton-kernel
|
||||
New CUDA JIT kernel? → sglang-diffusion-cuda-kernel
|
||||
After writing kernel → ncu again to verify bandwidth/occupancy ★
|
||||
↓
|
||||
6. VERIFY CORRECTNESS
|
||||
+4
-9
@@ -1,12 +1,7 @@
|
||||
---
|
||||
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.
|
||||
---
|
||||
# SGLang Diffusion Fast Paths
|
||||
|
||||
# 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.
|
||||
Use this guide when mapping a diffusion bottleneck to an existing fused path in `sglang.multimodal_gen`.
|
||||
Prefer reuse before writing a new Triton or CUDA kernel.
|
||||
|
||||
**Key Files**
|
||||
- `python/sglang/multimodal_gen/runtime/layers/layernorm.py`
|
||||
@@ -114,4 +109,4 @@ This skill focuses on SGLang Diffusion (`sglang.multimodal_gen`) kernel fusion p
|
||||
- Keep CuTe compile cache keys aligned to `(dtype, ndim, D)`.
|
||||
- Avoid implicit broadcasts that force hidden `contiguous()` copies.
|
||||
- Preserve NPU and ROCm fallback paths.
|
||||
- **Always verify with ncu** (`ncu --set full`) and compare against both the unfused baseline and the hardware roofline. Do not rely on a single universal bandwidth/occupancy threshold; the right target depends on whether the kernel is memory-bound, compute-bound, or launch-limited. See `diffusion-benchmark-and-profile.md` Step 3.5 for the ncu workflow.
|
||||
- **Always verify with ncu** (`ncu --set full`) and compare against both the unfused baseline and the hardware roofline. Do not rely on a single universal bandwidth/occupancy threshold; the right target depends on whether the kernel is memory-bound, compute-bound, or launch-limited. See `benchmark-and-profile.md` in this directory for the canonical ncu workflow.
|
||||
+28
-9
@@ -2,21 +2,24 @@
|
||||
End-to-end denoise-stage benchmark presets for SGLang Diffusion.
|
||||
|
||||
Measures denoise latency (primary metric ★) and peak GPU memory.
|
||||
All model configs are kept in exact sync with diffusion-benchmark-and-profile.md.
|
||||
All model configs are kept in exact sync with benchmark-and-profile.md.
|
||||
|
||||
Usage:
|
||||
# Single model
|
||||
cd /path/to/sglang
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py --model flux
|
||||
|
||||
# Tag the run for later compare_perf.py usage
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux --label tuned
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py --model flux --label tuned
|
||||
|
||||
# All 10 preset models
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --all
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py --all
|
||||
|
||||
For gated Hugging Face repos such as FLUX, export HF_TOKEN first:
|
||||
export HF_TOKEN=<your_hf_token>
|
||||
|
||||
Input images required for image-guided models:
|
||||
ASSET_DIR=$(python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py print-assets-dir --mkdir)
|
||||
ASSET_DIR=$(python3 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py print-assets-dir --mkdir)
|
||||
wget -O "${ASSET_DIR}/cat.png" \
|
||||
https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png
|
||||
wget -O "${ASSET_DIR}/astronaut.jpg" \
|
||||
@@ -48,9 +51,10 @@ from diffusion_skill_env import (
|
||||
|
||||
REPO_ROOT = get_repo_root()
|
||||
ASSET_DIR = ensure_dir(get_assets_dir(REPO_ROOT))
|
||||
GATED_MODELS = {"flux", "flux2"}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model configs — kept in exact sync with diffusion-benchmark-and-profile.md
|
||||
# Model configs — kept in exact sync with benchmark-and-profile.md
|
||||
# Each entry produces the same `sglang generate` command as shown in that doc.
|
||||
# ---------------------------------------------------------------------------
|
||||
MODELS = {
|
||||
@@ -240,7 +244,7 @@ def build_sglang_cmd(
|
||||
) -> list[str]:
|
||||
"""
|
||||
Build the `sglang generate` command for the given model.
|
||||
Matches the commands in diffusion-benchmark-and-profile.md exactly.
|
||||
Matches the commands in benchmark-and-profile.md exactly.
|
||||
"""
|
||||
cfg = MODELS[model_key]
|
||||
|
||||
@@ -292,6 +296,21 @@ def run_benchmark_once(
|
||||
|
||||
env = os.environ.copy()
|
||||
env.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
|
||||
if env.get("HF_TOKEN") and not env.get("HUGGINGFACE_HUB_TOKEN"):
|
||||
env["HUGGINGFACE_HUB_TOKEN"] = env["HF_TOKEN"]
|
||||
|
||||
if model_key in GATED_MODELS and not (
|
||||
env.get("HF_TOKEN") or env.get("HUGGINGFACE_HUB_TOKEN")
|
||||
):
|
||||
print(f"\n{'=' * 64}")
|
||||
print(f"[{label.upper()}] {model_key}")
|
||||
print(" ERROR: this preset uses a gated Hugging Face repo.")
|
||||
print(" Export HF_TOKEN before running it, for example:")
|
||||
print(" export HF_TOKEN=<your_hf_token>")
|
||||
print(" Without a token, the top-level `sglang generate` model detection may")
|
||||
print(" fail early and report a misleading unsupported-model error.")
|
||||
return {"model": model_key, "label": label, "error": True, "elapsed_s": 0.0}
|
||||
|
||||
if not env.get("CUDA_VISIBLE_DEVICES"):
|
||||
env["CUDA_VISIBLE_DEVICES"] = ",".join(
|
||||
str(index) for index in pick_idle_gpus(required_gpus_for_model(model_key))
|
||||
@@ -370,7 +389,7 @@ def print_results_table(results: list[dict]):
|
||||
print()
|
||||
print("=" * 80)
|
||||
print("BENCHMARK RESULTS — Denoise Latency (primary metric ★)")
|
||||
print("(Models and params match diffusion-benchmark-and-profile.md)")
|
||||
print("(Models and params match benchmark-and-profile.md)")
|
||||
print("=" * 80)
|
||||
|
||||
print(
|
||||
@@ -438,7 +457,7 @@ def main():
|
||||
|
||||
print(f"Perf dump JSONs → {output_dir}")
|
||||
print(
|
||||
"Compare across runs: follow diffusion-benchmark-and-profile.md → Perf dump & before/after compare."
|
||||
"Compare across runs: follow benchmark-and-profile.md -> Perf dump & before/after compare."
|
||||
)
|
||||
|
||||
|
||||
+3
-3
@@ -9,7 +9,7 @@ Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kerne
|
||||
|
||||
Usage:
|
||||
cd /path/to/sglang
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py
|
||||
python3 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_rmsnorm.py
|
||||
|
||||
Requirements:
|
||||
# Run inside the configured SGLang diffusion container shell.
|
||||
@@ -34,7 +34,7 @@ import torch
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Import the JIT CUDA kernel.
|
||||
# When you implement add-cuda-kernel.md, the file will be at:
|
||||
# When you implement the sglang-diffusion-cuda-kernel workflow, the file will be at:
|
||||
# python/sglang/jit_kernel/diffusion/rmsnorm.py
|
||||
# ---------------------------------------------------------------------------
|
||||
try:
|
||||
@@ -45,7 +45,7 @@ except ImportError:
|
||||
JIT_AVAILABLE = False
|
||||
print(
|
||||
"WARNING: diffusion.rmsnorm JIT kernel not available. "
|
||||
"Run after implementing add-cuda-kernel.md."
|
||||
"Run after implementing the sglang-diffusion-cuda-kernel workflow."
|
||||
)
|
||||
|
||||
|
||||
+18
-14
@@ -1,28 +1,32 @@
|
||||
---
|
||||
name: add-cuda-kernel
|
||||
description: Step-by-step guide for adding a new JIT CUDA kernel to SGLang Diffusion. CUDA source files go in jit_kernel/csrc/diffusion/<op>.cuh; Python wrapper at jit_kernel/diffusion/<op>.py. Use when implementing optimized CUDA kernels for diffusion model operators (RMSNorm, RoPE, AdaLN, GEGLU, etc.) on NVIDIA GPUs (H100, A100). Covers kernel authoring with sglang abstractions, JIT compilation, Python wrapper, integration into the denoise stage, and benchmarking. Adapted from https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels.
|
||||
name: sglang-diffusion-cuda-kernel
|
||||
description: Use when writing or tuning a JIT CUDA diffusion kernel in SGLang.
|
||||
---
|
||||
|
||||
# Adding a CUDA Kernel to SGLang Diffusion (JIT Style)
|
||||
|
||||
Use this skill when Triton is not enough and you need vectorized loads, warp reductions, or tighter control over memory layout and occupancy.
|
||||
|
||||
> **Origin**: This skill is adapted from the [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels), rewritten to follow SGLang's JIT compilation system and internal abstractions.
|
||||
>
|
||||
> **Run environment first**: before compiling, benchmarking, or profiling any kernel from this guide, use `scripts/diffusion_skill_env.py` (or the setup block in `diffusion-benchmark-and-profile.md`) to `cd` to the repo root resolved from `sglang.__file__`, verify write access, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and pick an idle GPU.
|
||||
> **Run environment first**: before compiling, benchmarking, or profiling any kernel from this guide, use `../sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py` (or the setup block in `../sglang-diffusion-benchmark-profile/benchmark-and-profile.md`) to `cd` to the repo root resolved from `sglang.__file__`, verify write access, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and pick an idle GPU.
|
||||
>
|
||||
> **Extended references** (in this directory's `references/` and `scripts/`):
|
||||
> **Gated model note**: if you use any FLUX-based `sglang generate` examples from the references below, export `HF_TOKEN` first so the top-level CLI can recognize the gated Hugging Face repo as a diffusion model.
|
||||
>
|
||||
> **Extended references** (in this directory's `references/` and the sibling benchmark skill):
|
||||
> - [references/kernel-templates.md](references/kernel-templates.md) — copy-paste ready templates for element-wise, row-reduction (RMSNorm), fused AdaLN
|
||||
> - [references/troubleshooting.md](references/troubleshooting.md) — build errors, perf issues, integration pitfalls
|
||||
> - [references/h100-optimization-guide.md](references/h100-optimization-guide.md) — H100 (sm_90) deep dive
|
||||
> - [references/a100-optimization-guide.md](references/a100-optimization-guide.md) — A100 (sm_80) deep dive
|
||||
> - [references/t4-optimization-guide.md](references/t4-optimization-guide.md) — T4 (sm_75, FP16 only) deep dive
|
||||
> - [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark vs PyTorch
|
||||
> - [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark preset runner; compare perf dumps with `compare_perf.py`
|
||||
> - [../sglang-diffusion-benchmark-profile/scripts/bench_diffusion_rmsnorm.py](../sglang-diffusion-benchmark-profile/scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark vs PyTorch
|
||||
> - [../sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py](../sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark preset runner; compare perf dumps with `compare_perf.py`
|
||||
|
||||
## When to Use CUDA vs Triton
|
||||
|
||||
| Scenario | Use |
|
||||
|----------|-----|
|
||||
| Fused elementwise / norm variants / RoPE | **Triton** (`add-triton-kernel.md`) — faster iteration |
|
||||
| Fused elementwise / norm variants / RoPE | **Triton** (`sglang-diffusion-triton-kernel`) — faster iteration |
|
||||
| Bandwidth-bound reduction (RMSNorm, LayerNorm) requiring max vectorization | **CUDA** — full control over `__nv_bfloat162` / `float4` vectorization |
|
||||
| Attention pattern or tile-based ops needing shared memory tuning | **CUDA** — warp-level primitives, shared memory layout |
|
||||
| Prototype or NPU/CPU fallback needed | **Triton** — portable across backends |
|
||||
@@ -441,8 +445,8 @@ After correctness + benchmarking, you must collect **Nsight Compute (ncu)** data
|
||||
|
||||
Use the canonical docs in this directory (do not duplicate CLI details across multiple skills):
|
||||
|
||||
- `diffusion-benchmark-and-profile.md` → Step 3.5 (ncu workflow, including CUDA graph profiling)
|
||||
- `nsight-profiler.md` (metrics interpretation: bandwidth / occupancy / roofline / stall reasons)
|
||||
- `../sglang-diffusion-benchmark-profile/benchmark-and-profile.md` → Step 3.5 (ncu workflow, including CUDA graph profiling)
|
||||
- `../sglang-diffusion-benchmark-profile/nsight-profiler.md` (metrics interpretation: bandwidth / occupancy / roofline / stall reasons)
|
||||
|
||||
---
|
||||
|
||||
@@ -472,7 +476,7 @@ python/sglang/jit_kernel/diffusion/
|
||||
python/sglang/jit_kernel/tests/
|
||||
└── test_diffusion_rmsnorm.py # NEW: correctness tests
|
||||
|
||||
python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/
|
||||
python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/
|
||||
├── bench_diffusion_rmsnorm.py # Validated micro-benchmark used by this skill
|
||||
└── bench_diffusion_denoise.py # Preset runner for end-to-end perf dumps
|
||||
```
|
||||
@@ -502,10 +506,10 @@ python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/
|
||||
|
||||
### Other Diffusion Kernel Skills (this directory)
|
||||
|
||||
- **Triton alternative**: `add-triton-kernel.md` — prefer Triton unless bandwidth analysis shows CUDA needed
|
||||
- **Existing fused kernels**: `use-efficient-diffusion-kernels.md` — check here first before writing new kernels
|
||||
- **Profiling**: `diffusion-benchmark-and-profile.md` — workflow to identify bottleneck before implementing
|
||||
- **Nsight Compute deep dive**: `nsight-profiler.md` — full guide: occupancy analysis, roofline model, warp efficiency, kernel comparison
|
||||
- **Triton alternative**: `../sglang-diffusion-triton-kernel/SKILL.md` — prefer Triton unless bandwidth analysis shows CUDA needed
|
||||
- **Existing fused kernels**: `../sglang-diffusion-benchmark-profile/existing-fast-paths.md` — check here first before writing new kernels
|
||||
- **Profiling**: `../sglang-diffusion-benchmark-profile/benchmark-and-profile.md` — workflow to identify bottleneck before implementing
|
||||
- **Nsight Compute deep dive**: `../sglang-diffusion-benchmark-profile/nsight-profiler.md` — full guide: occupancy analysis, roofline model, warp efficiency, kernel comparison
|
||||
|
||||
### External
|
||||
|
||||
+1
-1
@@ -260,7 +260,7 @@ nsys profile -o profile_report python scripts/bench_diffusion_rmsnorm.py
|
||||
# - Stream utilization
|
||||
```
|
||||
|
||||
For end-to-end denoise profiling via `sglang generate`, see `diffusion-benchmark-and-profile.md` (Level 2: nsys + gputrc2graph.py).
|
||||
For end-to-end denoise profiling via `sglang generate`, see the sibling `sglang-diffusion-benchmark-profile` skill (Level 2: nsys + gputrc2graph.py).
|
||||
|
||||
### NVIDIA Nsight Compute (ncu)
|
||||
|
||||
+5
@@ -4,6 +4,8 @@ T4 is a Turing architecture GPU (GCP n1+T4, AWS g4dn) commonly used for cloud in
|
||||
Its key constraint for diffusion kernels: **no BF16 support** — FP16 only.
|
||||
|
||||
> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels)
|
||||
>
|
||||
> If you use the FLUX `sglang generate` example below, export `HF_TOKEN` first. `black-forest-labs/FLUX.*` is a gated Hugging Face repo, and without a token the top-level CLI can fail before model loading.
|
||||
|
||||
---
|
||||
|
||||
@@ -203,6 +205,9 @@ T4's 16 GB requires careful planning for large diffusion models.
|
||||
**sglang generate flags for T4:**
|
||||
|
||||
```bash
|
||||
# Required for gated FLUX repos:
|
||||
# export HF_TOKEN=<your_hf_token>
|
||||
|
||||
# Enable CPU offloading to fit within 16 GB
|
||||
sglang generate \
|
||||
--model-path=black-forest-labs/FLUX.1-dev \
|
||||
+2
@@ -306,6 +306,8 @@ dram__throughput.avg.pct_of_peak_sustained_elapsed \
|
||||
# - smsp__warp_issue_stalled_*: warp stall breakdown (memory_dependency / math_pipe)
|
||||
|
||||
# 6. System-level profiling (per-op breakdown inside sglang generate)
|
||||
# Required for gated FLUX repos:
|
||||
# export HF_TOKEN=<your_hf_token>
|
||||
nsys profile -o denoise_profile \
|
||||
sglang generate --model-path=black-forest-labs/FLUX.1-dev \
|
||||
--width=1024 --height=1024 --num-inference-steps=50 \
|
||||
+6
-5
@@ -1,14 +1,15 @@
|
||||
---
|
||||
name: diffusion-optimal-perf
|
||||
description: Guide for achieving optimal performance with SGLang-Diffusion. Covers all perf-related CLI flags, env vars, and best practices for lossless and lossy speedup.
|
||||
name: sglang-diffusion-performance
|
||||
description: Use when choosing the fastest SGLang Diffusion flags for a model, GPU, and VRAM budget.
|
||||
---
|
||||
|
||||
# SGLang-Diffusion: Optimal Performance Guide
|
||||
# SGLang Diffusion Performance Tuning
|
||||
|
||||
Use this guide when a user asks how to speed up diffusion inference, reduce latency, lower VRAM usage, or tune SGLang-Diffusion for production.
|
||||
Use this skill when the user wants the fastest command line, lower VRAM, or the right performance flags for a specific model and GPU setup.
|
||||
|
||||
Before running any `sglang generate` command below inside the diffusion container:
|
||||
- use `python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py` to derive the repo root, verify write access, and choose idle GPU(s)
|
||||
- use `python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py` to derive the repo root, verify write access, and choose idle GPU(s)
|
||||
- export `HF_TOKEN` first when the selected model lives in a gated Hugging Face repo such as `black-forest-labs/FLUX.*`
|
||||
- export `FLASHINFER_DISABLE_VERSION_CHECK=1`
|
||||
- `cd` to the repo root resolved from `sglang.__file__`
|
||||
|
||||
+10
-10
@@ -1,14 +1,14 @@
|
||||
---
|
||||
name: add-triton-kernel
|
||||
description: Step-by-step guide for adding a new Triton kernel to SGLang Diffusion's jit_kernel module. Use when implementing fused elementwise ops, norm variants, RoPE variants, or any other lightweight GPU kernel for diffusion models using Triton JIT. Covers kernel authoring, autotune, torch.compile compatibility, layer integration, and tests.
|
||||
name: sglang-diffusion-triton-kernel
|
||||
description: Use when writing or tuning a Triton diffusion kernel in SGLang.
|
||||
---
|
||||
|
||||
# Adding a Triton Kernel to SGLang Diffusion
|
||||
|
||||
This guide walks through adding a Triton kernel to `python/sglang/jit_kernel/diffusion/triton/`.
|
||||
Use this skill when authoring or integrating a Triton kernel in `python/sglang/jit_kernel/diffusion/triton/`.
|
||||
We use a fused elementwise operation as the running example: `y = x * (1 + scale) + shift` (AdaLN modulation).
|
||||
|
||||
Before compiling, benchmarking, or profiling any Triton kernel from this guide, use `scripts/diffusion_skill_env.py` (or the setup block in `diffusion-benchmark-and-profile.md`) to `cd` to the repo root resolved from `sglang.__file__`, verify write access, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and choose an idle GPU.
|
||||
Before compiling, benchmarking, or profiling any Triton kernel from this guide, use `../sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py` or the setup block in `../sglang-diffusion-benchmark-profile/benchmark-and-profile.md` to `cd` to the repo root resolved from `sglang.__file__`, verify write access, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and choose an idle GPU.
|
||||
|
||||
---
|
||||
|
||||
@@ -22,7 +22,7 @@ python/sglang/jit_kernel/diffusion/
|
||||
│ ├── rmsnorm_onepass.py # One-pass RMSNorm for small hidden size
|
||||
│ └── rotary.py # RoPE kernel
|
||||
└── cutedsl/
|
||||
└── ... # CuTe DSL kernels (see use-efficient-diffusion-kernels.md)
|
||||
└── ... # CuTe DSL kernels (see existing-fast-paths.md in the benchmark/profile skill)
|
||||
```
|
||||
|
||||
New Triton kernels go into `triton/<op_name>.py`.
|
||||
@@ -375,8 +375,8 @@ After correctness tests, you must use **ncu (Nsight Compute)** to validate hardw
|
||||
|
||||
To avoid duplicating ncu CLI details across multiple skills, this skill does not repeat command flags. Follow the canonical docs:
|
||||
|
||||
- `diffusion-benchmark-and-profile.md` → Step 3.5 (ncu workflow, including CUDA graph profiling)
|
||||
- `nsight-profiler.md` (metrics interpretation: bandwidth / occupancy / roofline / warp stalls)
|
||||
- `../sglang-diffusion-benchmark-profile/benchmark-and-profile.md` → Step 3.5 (ncu workflow, including CUDA graph profiling)
|
||||
- `../sglang-diffusion-benchmark-profile/nsight-profiler.md` (metrics interpretation: bandwidth / occupancy / roofline / warp stalls)
|
||||
|
||||
---
|
||||
|
||||
@@ -510,6 +510,6 @@ python/sglang/multimodal_gen/runtime/layers/layernorm.py # MODIFIED: integrat
|
||||
- `python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py` — `wrap_triton`, tiled one-pass reduction
|
||||
- `python/sglang/jit_kernel/diffusion/triton/norm.py` — complex autotune with many `constexpr` flags
|
||||
- `python/sglang/jit_kernel/diffusion/triton/rotary.py` — per-head grid, interleaved RoPE
|
||||
- `nsight-profiler.md` — full Nsight Compute guide: occupancy analysis, roofline model, warp efficiency, kernel comparison
|
||||
- `diffusion-benchmark-and-profile.md` — how to verify the kernel's impact on denoise latency
|
||||
- `use-efficient-diffusion-kernels.md` — overview of existing fused kernel entry points
|
||||
- `../sglang-diffusion-benchmark-profile/nsight-profiler.md` — full Nsight Compute guide: occupancy analysis, roofline model, warp efficiency, kernel comparison
|
||||
- `../sglang-diffusion-benchmark-profile/benchmark-and-profile.md` — how to verify the kernel's impact on denoise latency
|
||||
- `../sglang-diffusion-benchmark-profile/existing-fast-paths.md` — overview of existing fused kernel entry points
|
||||
Reference in New Issue
Block a user