From da25b471e3c4194c6ed0c7891bc39736a388791f Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sat, 4 Apr 2026 21:43:52 +0800 Subject: [PATCH] Align diffusion nightly presets and broaden skill discovery (#22099) --- .../test_norm_tanh_mul_add_norm_scale.py | 0 .../sglang-diffusion-ako4all-kernel/SKILL.md | 17 +- .../SKILL.md | 35 +- .../benchmark-and-profile.md | 657 +++++------------- .../existing-fast-paths.md | 64 +- .../nsight-profiler.md | 286 -------- .../scripts/bench_diffusion_denoise.py | 59 +- .../scripts/bench_diffusion_rmsnorm.py | 210 ------ .../scripts/diffusion_skill_env.py | 1 - .../sglang-diffusion-cuda-kernel/SKILL.md | 516 -------------- .../references/a100-optimization-guide.md | 283 -------- .../references/h100-optimization-guide.md | 364 ---------- .../references/kernel-templates.md | 570 --------------- .../references/t4-optimization-guide.md | 340 --------- .../references/troubleshooting.md | 330 --------- .../sglang-diffusion-performance/SKILL.md | 29 +- .../sglang-diffusion-triton-kernel/SKILL.md | 515 -------------- .../multimodal_gen/runtime/utils/profiler.py | 17 +- scripts/ci/utils/diffusion/run_comparison.py | 26 +- 19 files changed, 381 insertions(+), 3938 deletions(-) rename python/sglang/jit_kernel/tests/{ => diffusion}/test_norm_tanh_mul_add_norm_scale.py (100%) delete mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/nsight-profiler.md delete mode 100755 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_rmsnorm.py delete mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/SKILL.md delete mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/a100-optimization-guide.md delete mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/h100-optimization-guide.md delete mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/kernel-templates.md delete mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/t4-optimization-guide.md delete mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/troubleshooting.md delete mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-triton-kernel/SKILL.md diff --git a/python/sglang/jit_kernel/tests/test_norm_tanh_mul_add_norm_scale.py b/python/sglang/jit_kernel/tests/diffusion/test_norm_tanh_mul_add_norm_scale.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_norm_tanh_mul_add_norm_scale.py rename to python/sglang/jit_kernel/tests/diffusion/test_norm_tanh_mul_add_norm_scale.py diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md index ca6ed3c61..94c3a884b 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md @@ -6,7 +6,7 @@ description: Use when optimizing an existing SGLang diffusion kernel with AKO4AL # SGLang Diffusion AKO4ALL Kernel Use this skill to run the full AKO4ALL-based optimization loop for an existing SGLang diffusion kernel. -It packages the workflow we used for diffusion Triton and JIT kernel tuning: bootstrap a custom AKO harness, benchmark and profile the kernel, iterate with `ncu`, port the best version back to `sglang`, then validate with targeted tests and model-level denoise runs. +It is the default implementation path once the benchmark/profile skill has already shown that a hotspot is real and not covered by an existing fast path. This workflow bootstraps a custom AKO harness, benchmarks and profiles the kernel, iterates with `ncu`, ports the best version back to `sglang`, then validates with targeted tests and model-level denoise runs. This skill assumes a sibling repo layout like: @@ -21,16 +21,19 @@ If `AKO4ALL/` is missing under the current base directory, clone it first. ## Use This Skill When - tuning an existing diffusion Triton, CUDA JIT, CuTeDSL, or runtime-integrated kernel in `sglang` +- `sglang-diffusion-benchmark-profile` has already ruled out an existing in-repo fast path or overlap family - creating a custom AKO4ALL harness for a real diffusion kernel instead of using the default benchmark tasks - validating that a kernel-level win transfers to Qwen, FLUX, Wan, Hunyuan, MOVA, or other diffusion denoise latency - preparing PR artifacts such as microbench tables, `ncu` before/after data, and proof image outputs -Do not use this skill when adding a brand-new kernel from scratch with no existing SGLang integration. -For that, start from the sibling kernel-authoring skills first: +Do not start here when the bottleneck has not been proven yet. +First use [../sglang-diffusion-benchmark-profile/SKILL.md](../sglang-diffusion-benchmark-profile/SKILL.md) to: +- measure the real denoise regression +- collect the perf dump baseline +- capture one representative `torch.profiler` trace +- rule out existing merged fast paths -- Triton: [../sglang-diffusion-triton-kernel/SKILL.md](../sglang-diffusion-triton-kernel/SKILL.md) -- CUDA JIT: [../sglang-diffusion-cuda-kernel/SKILL.md](../sglang-diffusion-cuda-kernel/SKILL.md) -- Denoise benchmark/profile: [../sglang-diffusion-benchmark-profile/SKILL.md](../sglang-diffusion-benchmark-profile/SKILL.md) +If a future specialized optimization skill matches the kernel family better than AKO4ALL, hand off there instead. The diagnosis contract stays the same. ## Mandatory AKO4ALL Preflight @@ -55,8 +58,6 @@ By default it uses the existing `origin` URL, or `AKO4ALL_URL` if you need to ov - Record the target shapes, dtypes, model families, and whether the kernel is on a hot path. - Reuse existing unit tests and benchmark entry points when they already exist. -If the implementation work is primarily Triton or CUDA authoring, read only the relevant sibling skill from the list above. - ### 2. Bootstrap the AKO Harness Inside the clean `AKO4ALL` repo: diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md index 9553b064c..d91f2c2a2 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md @@ -5,11 +5,16 @@ description: Use when benchmarking denoise latency or profiling a diffusion bott # 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`. +Use this skill when measuring denoise performance, finding the slow op, checking whether an existing fast path can solve it, or verifying that a hotspot is real before any kernel work 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) +This skill is diagnosis-first. It owns: +- checked-in denoise benchmark presets +- perf dump collection and before/after comparison +- `torch.profiler` trace capture and quick hotspot ranking +- mapping hot kernels back to known fast paths and fusion families +- handing confirmed kernel work to a specialized optimization skill such as [../sglang-diffusion-ako4all-kernel/SKILL.md](../sglang-diffusion-ako4all-kernel/SKILL.md) + +This skill does not own low-level kernel authoring or standalone Nsight workflows. ## Preflight @@ -22,9 +27,21 @@ Before running any benchmark, profiler, or kernel-validation command: ## 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 +- [benchmark-and-profile.md](benchmark-and-profile.md) — canonical denoise benchmark, perf dump, and `torch.profiler` workflow; uses the checked-in nightly-aligned presets, including `LTX-2` two-stage +- [existing-fast-paths.md](existing-fast-paths.md) — map bottlenecks to existing fused kernels, packed QKV paths, fused `QK norm + RoPE`, and distributed overlap patterns before proposing new code - [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` +- [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark preset runner via `sglang generate`; use `--list-models` to inspect preset order, then save perf dumps by label and compare them with `compare_perf.py` + +## Opportunity Discovery Rule + +Before calling a diffusion hotspot "new", first classify it with `existing-fast-paths.md`. + +Always rule out these existing families first: +- merged Z-Image residual-form modulation +- fused diffusion `QK norm + RoPE` +- NVFP4 / Nunchaku packed QKV +- Nunchaku fused GELU MLP +- Ulysses / USP attention overlap +- turbo-layer async all-to-all overlap +- `torch.compile` compute / communication reorder +- dual-stream diffusion execution diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md index 5d54535cc..be4aac215 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md @@ -1,19 +1,30 @@ --- name: benchmark-and-profile-reference -description: Reference commands and workflow for denoise benchmarks and profiling in SGLang Diffusion. +description: Reference commands and workflow for denoise benchmarks, perf dumps, and torch.profiler analysis in SGLang Diffusion. --- # SGLang Diffusion Benchmark and Profile Guide **Primary Metric: Denoise Latency** - Denoise latency is the total DiT forward-pass time across all inference steps. -- It is the dominant cost for diffusion inference, typically more than 80% of end-to-end time. -- It is the **sole optimization target** for kernel work. -- End-to-end latency is a secondary sanity check only. +- It is the dominant cost for diffusion inference and the main optimization target. +- End-to-end latency and peak memory are secondary sanity checks. -> **Correctness First**: Faster but incorrect output is not an improvement. Always compare generated images/videos against a reference baseline before and after any change. +> **Correctness First**: Faster but incorrect output is not an improvement. Always compare generated images or videos against a reference baseline before and after any change. ---- +## Scope + +This guide intentionally stops at: +- checked-in denoise benchmarks +- structured perf dumps +- `torch.profiler` trace capture +- hotspot ranking +- mapping hotspots to known fast paths + +If the hotspot survives this checklist, hand the work to +`sglang-diffusion-ako4all-kernel` or another specialized kernel-optimization +skill. Do not grow this skill back into a general Nsight or kernel-authoring +guide. ## Prerequisites @@ -30,7 +41,6 @@ export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 1) ASSET_DIR=$(python3 "$ENV_PY" print-assets-dir --mkdir) BENCH_DIR=$(python3 "$ENV_PY" print-output-dir --kind benchmarks --mkdir) PROFILE_DIR=$(python3 "$ENV_PY" print-output-dir --kind profiles --mkdir) -NCU_DIR=$(python3 "$ENV_PY" print-output-dir --kind ncu --mkdir) export PROFILE_DIR check() { @@ -42,24 +52,17 @@ check() { check "sglang" python3 -c "import sglang" check "torch+CUDA" python3 -c "import torch; assert torch.cuda.is_available()" check "torch.profiler" python3 -c "import torch.profiler" -check "nsys (Level 2)" which nsys -check "ncu (Level 3)" which ncu -check "pandas" python3 -c "import pandas" -check "plotly" python3 -c "import plotly" -check "regex" python3 -c "import regex" ``` Environment notes: -- **Minimum for benchmarking**: `sglang`, `torch` with CUDA. -- **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. +- all commands below assume you are inside the configured diffusion container shell +- export `HF_TOKEN` before any gated Hugging Face model run +- export `FLASHINFER_DISABLE_VERSION_CHECK=1` before any benchmark or profiler run +- re-run `print-idle-gpus` before each perf command if GPU availability may have changed +- keep benchmark commands within 4 GPUs or fewer + +Download input images required by some presets: -Download input images required by some models: ```bash wget -O "${ASSET_DIR}/cat.png" \ https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png @@ -67,129 +70,89 @@ wget -O "${ASSET_DIR}/mova_single_person.jpg" \ https://github.com/OpenMOSS/MOVA/raw/main/assets/single_person.jpg ``` ---- +## Benchmark Presets -## Benchmark Commands +Treat +`python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py` +as the source of truth for preset order. -All commands include `--warmup` and `--enable-torch-compile` for real production performance. Add `--perf-dump-path .json` for machine-readable output. +Nightly diffusion comparison is server/API based (`sglang serve` plus requests). +This skill stays on `sglang generate` for local benchmarking and profiling, but +the first 9 presets in `bench_diffusion_denoise.py` are aligned to nightly on +model, prompt, negative prompt, reference image, size, frames, fps, seed, GPU +count, and any explicitly overridden sampling or parallelism flags. -Nightly diffusion comparison is server/API based (`sglang serve` + OpenAI-compatible requests). The commands below stay on `sglang generate` for local profiling, but the first 8 presets are aligned to nightly on model, prompt, reference image, steps, guidance scale, GPU count, and parallelism flags. +List the current preset order: -If you want a checked-in preset runner instead of copying commands manually, use `scripts/bench_diffusion_denoise.py --model --label ` or `--list-models`. It writes the same perf dump JSONs used by `compare_perf.py`. +```bash +PYTHONPATH=python python3 \ + python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py \ + --list-models +``` -### Preset Catalog +Run one preset and save a perf dump: + +```bash +PYTHONPATH=python python3 \ + python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py \ + --model ltx2 \ + --label baseline \ + --output-dir "${BENCH_DIR}" +``` + +Run the full preset sweep: + +```bash +PYTHONPATH=python python3 \ + python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py \ + --all \ + --label prXXXX \ + --output-dir "${BENCH_DIR}" +``` Nightly-aligned presets come first; skill-only presets stay available after them. | Preset | Model | Nightly | Notes | | --- | --- | --- | --- | -| `flux` | `black-forest-labs/FLUX.1-dev` | Yes: `flux1_dev_t2i_1024` | Aligned to nightly prompt + `--dit-layerwise-offload false` | +| `flux` | `black-forest-labs/FLUX.1-dev` | Yes: `flux1_dev_t2i_1024` | Aligned to nightly prompt plus `--dit-layerwise-offload false` | | `flux2` | `black-forest-labs/FLUX.2-dev` | Yes: `flux2_dev_t2i_1024` | Aligned to nightly prompt, 50 steps, guidance 4.0 | -| `qwen` | `Qwen/Qwen-Image-2512` | Yes: `qwen_image_2512_t2i_1024` | Aligned to nightly prompt/steps; no extra offload overrides | -| `qwen-edit` | `Qwen/Qwen-Image-Edit-2511` | Yes: `qwen_image_edit_2511` | Uses nightly cat image + edit prompt | -| `zimage` | `Tongyi-MAI/Z-Image-Turbo` | Yes: `zimage_turbo_t2i_1024` | Aligned to nightly prompt + guidance 4.0 | +| `qwen` | `Qwen/Qwen-Image-2512` | Yes: `qwen_image_2512_t2i_1024` | Aligned to nightly prompt and steps | +| `qwen-edit` | `Qwen/Qwen-Image-Edit-2511` | Yes: `qwen_image_edit_2511` | Uses the nightly cat image and edit prompt | +| `zimage` | `Tongyi-MAI/Z-Image-Turbo` | Yes: `zimage_turbo_t2i_1024` | Aligned to nightly prompt and guidance 4.0 | | `wan-t2v` | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | Yes: `wan22_t2v_a14b_720p` | Aligned to nightly CFG-parallel 4-GPU launch | -| `wan-ti2v` | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | Yes: `wan22_ti2v_5b_720p` | Uses nightly cat image + motion prompt | -| `wan-i2v` | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Yes: `wan22_i2v_a14b_720p` | Added to match nightly; aligned to CFG-parallel 4-GPU launch | +| `wan-ti2v` | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | Yes: `wan22_ti2v_5b_720p` | Uses the nightly cat image and motion prompt | +| `ltx2` | `Lightricks/LTX-2` | Yes: `ltx2_twostage_t2v` | Uses `LTX2TwoStagePipeline`; nightly-aligned prompt, negative prompt, 1536x1024, 121 frames, fps 24, seed 1234 | +| `wan-i2v` | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Yes: `wan22_i2v_a14b_720p` | Aligned to nightly CFG-parallel 4-GPU launch | | `hunyuanvideo` | `hunyuanvideo-community/HunyuanVideo` | No | Skill-only extra preset | | `mova-720p` | `OpenMOSS-Team/MOVA-720p` | No | Skill-only extra preset | | `helios` | `BestWishYsh/Helios-Base` | No | Skill-only extra preset | -### Perf dump & before/after compare +For Wan2.2 video models, remember the difference between **nightly alignment** +and **best latency tuning**: +- the nightly-aligned 4-GPU commands intentionally keep `--enable-cfg-parallel --ulysses-degree=2` so CFG and ring behavior stay covered +- do not assume that is the fastest topology +- for pure latency tuning, benchmark pure Ulysses too, for example `--ulysses-degree=4 --ring-degree=1` on 4 GPUs, and on 8 GPUs compare pure `--ulysses-degree=8` against `--enable-cfg-parallel --ulysses-degree=4` -For every benchmark run, always write a perf dump JSON: +### Manual command example: LTX-2 Two-Stage -```bash -sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/.json" -``` - -Before/after comparison (outputs a Markdown table suitable for PR descriptions): - -```bash -# Baseline (on main branch or before changes) -sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/baseline.json" - -# New (after changes) -sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/new.json" - -python3 python/sglang/multimodal_gen/benchmarks/compare_perf.py \ - "${BENCH_DIR}/baseline.json" "${BENCH_DIR}/new.json" -``` - -### FLUX.1-dev (Nightly: `flux1_dev_t2i_1024`) ```bash sglang generate \ - --model-path=black-forest-labs/FLUX.1-dev \ - --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets" \ - --width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \ - --seed=42 --save-output --enable-torch-compile --warmup \ - --dit-layerwise-offload false + --model-path=Lightricks/LTX-2 \ + --pipeline-class-name=LTX2TwoStagePipeline \ + --prompt="A beautiful sunset over the ocean" \ + --negative-prompt="shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \ + --width=1536 --height=1024 \ + --num-frames=121 --fps=24 \ + --seed=1234 --num-gpus=1 \ + --save-output --enable-torch-compile --warmup ``` -### FLUX.2-dev (Nightly: `flux2_dev_t2i_1024`) -```bash -sglang generate \ - --model-path=black-forest-labs/FLUX.2-dev \ - --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets" \ - --width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \ - --seed=42 --save-output --enable-torch-compile --warmup \ - --dit-layerwise-offload false -``` +After [PR #20707](https://github.com/sgl-project/sglang/pull/20707), +`LTX2TwoStagePipeline` is a native path. The spatial upsampler and distilled +LoRA are auto-resolved from the same model snapshot unless you override them. -### Qwen-Image-2512 (Nightly: `qwen_image_2512_t2i_1024`) -```bash -sglang generate \ - --model-path=Qwen/Qwen-Image-2512 \ - --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets" \ - --width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \ - --seed=42 --save-output --enable-torch-compile --warmup -``` +### Manual command example: Wan2.2-I2V-A14B 720P -### Qwen-Image-Edit-2511 (Nightly: `qwen_image_edit_2511`) -```bash -sglang generate \ - --model-path=Qwen/Qwen-Image-Edit-2511 \ - --prompt="Make the cat wear a red hat" \ - --image-path="${ASSET_DIR}/cat.png" \ - --width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \ - --seed=42 --save-output --enable-torch-compile --warmup -``` - -### Z-Image-Turbo (Nightly: `zimage_turbo_t2i_1024`) -```bash -sglang generate \ - --model-path=Tongyi-MAI/Z-Image-Turbo \ - --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets" \ - --width=1024 --height=1024 --num-inference-steps=9 --guidance-scale=4.0 \ - --seed=42 --save-output --enable-torch-compile --warmup -``` - -### Wan2.2-T2V-A14B 720P (Nightly: `wan22_t2v_a14b_720p`) -```bash -# Select four idle GPUs first: -# export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 4) -sglang generate \ - --model-path=Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --prompt="A cat and a dog baking a cake together in a kitchen." \ - --720p --num-inference-steps=2 --num-frames=81 \ - --guidance-scale=5.0 --seed=42 --save-output \ - --num-gpus=4 --enable-cfg-parallel --ulysses-degree=2 \ - --text-encoder-cpu-offload --pin-cpu-memory \ - --warmup --enable-torch-compile -``` - -### Wan2.2-TI2V-5B 720P (Nightly: `wan22_ti2v_5b_720p`) -```bash -sglang generate \ - --model-path=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ - --prompt="The cat starts walking slowly towards the camera." \ - --image-path="${ASSET_DIR}/cat.png" \ - --num-frames=81 --720p --num-inference-steps=50 --guidance-scale=5.0 \ - --seed=42 --save-output \ - --enable-torch-compile --warmup -``` - -### Wan2.2-I2V-A14B 720P (Nightly: `wan22_i2v_a14b_720p`) ```bash # Select four idle GPUs first: # export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 4) @@ -204,63 +167,52 @@ sglang generate \ --warmup --enable-torch-compile ``` -### HunyuanVideo (Skill-only, not nightly) +After [PR #21390](https://github.com/sgl-project/sglang/pull/21390), +`Wan2.2-I2V-A14B` uses the 720p max-area config by default, and explicit +`--width/--height` overrides control the target area while preserving the +reference-image aspect ratio. + +## Perf Dump Workflow + +For every benchmark run, write a perf dump JSON: + ```bash -sglang generate \ - --model-path=hunyuanvideo-community/HunyuanVideo \ - --text-encoder-cpu-offload --pin-cpu-memory \ - --prompt="A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." \ - --save-output --num-frames=65 --width=848 --height=480 \ - --num-inference-steps=30 \ - --warmup --enable-torch-compile +sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/.json" ``` -### MOVA-720p (Skill-only, not nightly) +Before/after comparison: + ```bash -# Select four idle GPUs first: -# export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 4) -sglang generate \ - --model-path=OpenMOSS-Team/MOVA-720p \ - --prompt="A man in a blue blazer and glasses speaks in a formal indoor setting, framed by wooden furniture and a filled bookshelf. Quiet room acoustics underscore his measured tone as he delivers his remarks. At one point, he says, \"I would also believe that this advance in AI recently wasn’t unexpected.\"" \ - --image-path="${ASSET_DIR}/mova_single_person.jpg" \ - --adjust-frames=false \ - --num-gpus=4 --ring-degree=1 --ulysses-degree=4 \ - --num-frames=193 --fps=24 \ - --num-inference-steps=2 \ - --enable-torch-compile --save-output --warmup +python3 python/sglang/multimodal_gen/benchmarks/compare_perf.py \ + "${BENCH_DIR}/baseline.json" \ + "${BENCH_DIR}/new.json" ``` -### Helios-Base (Skill-only, not nightly) +Always keep: +- denoise latency +- end-to-end latency +- peak GPU memory +- exact command line, model shape, dtype, and GPU topology + +## `torch.profiler` Workflow + +### 1. Establish the baseline + ```bash -sglang generate \ - --model-path=BestWishYsh/Helios-Base \ - --prompt="A curious raccoon" \ - --width=640 --height=384 --num-frames=33 \ - --dit-layerwise-offload false --dit-cpu-offload false \ - --text-encoder-cpu-offload false --vae-cpu-offload false \ - --seed=42 --save-output --enable-torch-compile --warmup +PYTHONPATH=python python3 \ + python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py \ + --model flux \ + --label baseline \ + --output-dir "${BENCH_DIR}" ``` -**Key metrics** (all models): denoise latency ★, end-to-end latency, peak GPU memory. +Keep model shape, seed, and GPU topology fixed for every comparison. Save one +reference image or video before changing code. ---- +### 2. Capture a representative trace -## Performance Bottleneck Workflow - -### Step 1: Identify the Slow DiT Operation - -Add `--log-level=info` and observe: -- **Denoise loop latency** ★ — primary target -- Per-step DiT latency — denoise ÷ steps - -### Step 2: Profile with torch.profiler (Level 1) - -**Compile-safety rule for fused or rewritten kernels** -- Any new kernel must be checked for `torch.compile` graph breaks before trusting its benchmark result. -- If a direct Python/library call triggers tracing issues, wrap it as a custom op first. -- For external libraries, use `register_custom_op_from_extern(...)`. -- For SGLang JIT kernels, use `@register_custom_op(...)` and keep the JIT/module loading inside the custom op body. -- Re-run `torch._dynamo.explain` on representative shapes and verify the optimized path still gets `graph_count=1` and `graph_break_count=0`. +By default SGLang profiles the denoising stage. The default sampling window is +5 profiled timesteps after warmup. ```bash SGLANG_TORCH_PROFILER_DIR="${PROFILE_DIR}/torch" \ @@ -269,334 +221,89 @@ sglang generate \ --prompt="A futuristic cyberpunk city at night" \ --width=1024 --height=1024 --num-inference-steps=50 \ --seed=42 --enable-torch-compile --warmup \ - --profile --num-profiled-timesteps 3 + --profile ``` -Parse the trace without a browser: +Use `--profile-all-stages` only when you really need text encoder, VAE, or +other non-denoise stages too. + +The generated trace path is printed in the console and also lands under +`./logs/` or `SGLANG_TORCH_PROFILER_DIR`. Open it in Perfetto if you want a +timeline view: +- https://ui.perfetto.dev/ + +### 3. Rank the hot CUDA kernels + +Use this parser for a quick top-k table without opening a browser: + ```python -import gzip, json, collections, glob, os +import collections +import glob +import gzip +import json +import os log_dir = os.environ.get("SGLANG_TORCH_PROFILER_DIR", "./logs") -trace_path = sorted(glob.glob(f"{log_dir}/*.trace.json.gz"), key=os.path.getmtime, reverse=True)[0] +trace_path = sorted( + glob.glob(f"{log_dir}/*.trace.json.gz"), + key=os.path.getmtime, + reverse=True, +)[0] + with gzip.open(trace_path, "rb") as f: data = json.loads(f.read()) cuda_ops = collections.defaultdict(lambda: {"total_us": 0, "count": 0}) -for e in data.get("traceEvents", []): - if e.get("cat") in ("kernel", "gpu_memcpy") and "dur" in e: - cuda_ops[e.get("name","unknown")]["total_us"] += e["dur"] - cuda_ops[e.get("name","unknown")]["count"] += 1 +for event in data.get("traceEvents", []): + if event.get("cat") in ("kernel", "gpu_memcpy") and "dur" in event: + cuda_ops[event.get("name", "unknown")]["total_us"] += event["dur"] + cuda_ops[event.get("name", "unknown")]["count"] += 1 -print(f"{'Kernel':<80} {'Total(ms)':>10} {'Count':>6}") -for name, s in sorted(cuda_ops.items(), key=lambda x: -x[1]["total_us"])[:30]: - print(f"{name:<80} {s['total_us']/1000:>10.3f} {s['count']:>6}") +print(f"{'Kernel':<90} {'Total(ms)':>10} {'Count':>6}") +for name, stat in sorted(cuda_ops.items(), key=lambda item: -item[1]["total_us"])[:30]: + print(f"{name:<90} {stat['total_us'] / 1000:>10.3f} {stat['count']:>6}") ``` -Add `record_function` scopes in the DiT block for per-layer attribution: -```python -with torch.profiler.record_function(f"dit_block_{idx}.attn"): - x = self.attn(x) -with torch.profiler.record_function(f"dit_block_{idx}.norm"): - x = self.norm(x) -``` +If you need better attribution, add `record_function(...)` scopes around DiT +attention, norm, modulation, MLP, or communication boundaries and re-run. -**Expected dominant kernels per DiT sub-component:** +### 4. Classify the hotspot with `existing-fast-paths.md` -| Sub-component | Expected kernel | -|--------------|-----------------| -| QKV / output / MLP projections | `cutlass_gemm` / `ampere_*_gemm` | -| Attention | `flash_attn_fwd` / `fmha_*` (FA3/FA4) | -| AdaLN modulation | `fuse_scale_shift_kernel` | -| RMSNorm / LayerNorm | `sgl_kernel_rmsnorm` / Triton norm | -| SiLU gate | `vectorized_elementwise_kernel` | -| RoPE | `apply_rotary_embedding` (Triton) | -| QK Norm | `fused_inplace_qknorm` (JIT) | +Do not jump from a hot kernel straight into new code. First classify it against +the known merged families. -### Step 3: Deep CUDA Kernel Breakdown (Level 2 — nsys) +| What the trace shows | First interpretation | +| --- | --- | +| `fused_inplace_qknorm_rope` missing, but separate qk norm plus rope show up | Check whether the fused diffusion `QK norm + RoPE` path should have engaged | +| `to_q -> to_k -> to_v` on NVFP4 or Nunchaku FLUX-family checkpoints | Treat as a packed-QKV fast-path miss or checkpoint-format mismatch | +| `fused_norm_tanh_mul_add*` missing on Z-Image | Treat as a missing merged modulation path, not a new fusion request | +| `all_to_all`, ring attention, or async A2A dominate | Classify against Ulysses, USP, or turbo-layer overlap first | +| split `fc1 -> gelu -> quant -> fc2.lora_down` on Nunchaku FLUX | Treat as a missing fused GELU MLP path | +| attention kernels dominate | Confirm backend, topology, and shape guards before proposing a new kernel | -Workflow: -- **Pass A**: collect the `nsys` trace. -- **Pass B**: measure wall-clock runtime without profiling. -- Write the non-profiled wall-clock time into `ELAPSED_SEC`. +If the hot path is already covered by a merged optimization family, fix the +enablement, shape guard, backend choice, or checkpoint mapping first. -```bash -# Pass A — collect nsys trace (skip warmup with --delay) -nsys profile -t cuda -o "${PROFILE_DIR}/flux_dev" -f true \ - --trace-fork-before-exec=true --delay 120 --duration 60 \ - sglang generate \ - --model-path=black-forest-labs/FLUX.1-dev \ - --prompt="A futuristic cyberpunk city at night" \ - --width=1024 --height=1024 --num-inference-steps=50 \ - --seed=42 --enable-torch-compile --warmup +### 5. Hand off only real kernel work -# Pass B — measure wall-clock time without profiling -time sglang generate --model-path=black-forest-labs/FLUX.1-dev \ - --width=1024 --height=1024 --num-inference-steps=50 --seed=42 \ - --enable-torch-compile --warmup -# Record ELAPSED_SEC from Pass B -``` +Only after the hotspot survives the fast-path checklist: -Create classification JSON at `examples/profiler/nsys_profile_tools/sglang_diffusion_engine_model.json`: -```json -{ - "sglang": { - "diffusion": { - "gemm|nvjet|cutlass": "gemm", - "flash|fmha|fwd_flash": "attn", - "fuse_scale_shift|scale_shift_gate": "adaln_modulation", - "_norm_|Norm|rmsnorm|fused_add_rmsnorm": "norm", - "rotary|rope": "rope", - "act_and_mul|silu|gelu": "activation", - "ncclDevKernel|all_gather|all_reduce": "nccl_comm", - "triton": "triton_kernel", - "CUDA mem": "non-gpu-H_D_memops", - ".*": "misc" - } - } -} -``` +1. save a baseline perf dump +2. save a representative `torch.profiler` trace +3. note the exact model, shape, dtype, and GPU topology +4. hand the work to `sglang-diffusion-ako4all-kernel` or another future specialized optimization skill -Notes: -- `gputrc2graph.py` only recognizes `sglang,diffusion,...` after this JSON file exists in `examples/profiler/nsys_profile_tools/`. -- If you only want a quick structural check, set `ELAPSED_SEC=0`. The report will still generate, but `CPU(non-GPU)` time can be inflated. +This skill intentionally stops here. It tells you whether you are looking at: +- a missing existing optimization +- a configuration or backend problem +- or a real kernel opportunity worth handing off -Run analysis: -```bash -ELAPSED_SEC=12.34 -cd "$ROOT/examples/profiler/nsys_profile_tools" -python3 gputrc2graph.py \ - --in_file "${PROFILE_DIR}/flux_dev.nsys-rep,sglang,diffusion,${ELAPSED_SEC}" \ - --out_dir "${PROFILE_DIR}/analysis" \ - --title "FLUX.1-dev denoise kernel breakdown" +## Minimal Merge Checklist -# Read results -python3 - << 'EOF' -import os -import pandas as pd - -df = pd.read_csv(f"{os.environ['PROFILE_DIR']}/analysis/result.csv") -summary = df.groupby("Category")["Elapsed Time (sec)"].sum().sort_values(ascending=False) -total = summary.sum() -for cat, sec in summary.items(): - print(f"{cat:<30} {sec:>8.3f}s ({sec/total*100:>5.1f}%)") -EOF -``` - -**What the category breakdown tells you:** - -| Category high | Investigation | -|--------------|---------------| -| `gemm` dominant | Check tensor parallelism; QKV/MLP bottleneck | -| `attn` dominant | Verify FA3/FA4 is active | -| `adaln_modulation` high | Verify fused `fuse_scale_shift_kernel` is used | -| `norm` high | Verify `sgl_kernel_rmsnorm` / CuTe DSL path; check D alignment | -| `nccl_comm` high | Multi-GPU: tune Ulysses degree | -| `triton_kernel` high | Identify which Triton kernel; consider CUDA replacement | -| `non-gpu-H_D_memops` high | Accidental CPU offload or `.cpu()` calls mid-denoising | -| `CPU(non-GPU)` high | Python dispatch overhead / torch.compile graph breaks | - -### Step 3.5: Per-Kernel Deep Analysis (Level 3 — ncu) - -**CRITICAL**: `ncu` (Nsight Compute) is the essential tool for kernel-level optimization. While nsys and torch.profiler tell you **which** kernels are slow, only ncu tells you **why** — memory bandwidth utilization, compute throughput, occupancy limiters, warp stall reasons, and roofline position. **Always use ncu when optimizing or writing custom kernels.** - -#### When to use ncu - -- After writing a new Triton or CUDA kernel — verify it saturates hardware bandwidth -- When a kernel shows up as a top bottleneck in Level 1/2 profiling -- When comparing your fused kernel vs PyTorch baseline or torch.compile output -- When tuning Triton autotune configs (block sizes, num_warps) -- When profiling `sglang generate`, add `--target-processes all` so child worker processes are included - -#### Basic ncu workflow - -```bash -# 1. Profile a specific kernel by name (skip warmup launches, collect 3 invocations) -ncu --target-processes all \ - --kernel-name "_fused_gated_residual_add_kernel" \ - --launch-skip 10 --launch-count 3 \ - --set full \ - -o "${NCU_DIR}/gated_residual" \ - sglang generate \ - --model-path=black-forest-labs/FLUX.1-dev \ - --prompt="test" --width=1024 --height=1024 \ - --num-inference-steps=5 --seed=42 - -# 2. Profile all kernels in a short run (use few steps to limit time) -ncu --target-processes all \ - --launch-skip 50 --launch-count 200 \ - --set full \ - -o "${NCU_DIR}/all_kernels" \ - sglang generate \ - --model-path=black-forest-labs/FLUX.1-dev \ - --prompt="test" --width=1024 --height=1024 \ - --num-inference-steps=3 --seed=42 - -# 3. For CUDA graph mode, keep --graph-profiling=node on the ncu side. -# Note: `--enable-piecewise-cuda-graph` is a server flag, not a valid -# `sglang generate` flag, so do not append it here. -ncu --target-processes all \ - --graph-profiling node \ - --kernel-name "_fused_gated_residual_add_kernel" \ - --launch-skip 5 --launch-count 3 \ - --set full \ - -o "${NCU_DIR}/gated_residual_cudagraph" \ - sglang generate \ - --model-path=black-forest-labs/FLUX.1-dev \ - --prompt="test" --width=1024 --height=1024 \ - --num-inference-steps=5 --seed=42 -``` - -#### Reading ncu results (CLI, no GUI needed) - -```bash -# Summary of all profiled kernels -ncu --import "${NCU_DIR}/gated_residual.ncu-rep" --page raw --csv 2>/dev/null | head -50 - -# Key metrics to extract: -ncu --import "${NCU_DIR}/gated_residual.ncu-rep" \ - --page details --csv 2>/dev/null | python3 -c " -import csv, sys -reader = csv.DictReader(sys.stdin) -key_metrics = { - 'gpu__time_duration.avg': 'Duration', - 'sm__throughput.avg.pct_of_peak_sustained_elapsed': 'Compute (SM) Throughput', - 'dram__throughput.avg.pct_of_peak_sustained_elapsed': 'DRAM Throughput', - 'l1tex__throughput.avg.pct_of_peak_sustained_elapsed': 'L1/TEX Cache Throughput', - 'sm__warps_active.avg.pct_of_peak_sustained_active': 'Achieved Occupancy', - 'launch__occupancy_limit_registers': 'Block Limit Registers', - 'launch__occupancy_limit_shared_mem': 'Block Limit Shared Mem', -} -for row in reader: - name = row.get('Metric Name', '') - if any(alias in name or metric in name for metric, alias in key_metrics.items()): - print(f'{name:<60} {row.get(\"Metric Value\",\"\")}') -" -``` - -#### Interpreting ncu results for kernel optimization - -| Metric | Good | Action if bad | -|--------|------|--------------| -| DRAM throughput > 80% peak | Memory-bound, near optimal | Already saturating HBM — fuse with adjacent ops to reduce total memory traffic | -| DRAM throughput < 50% peak | Not saturating memory bandwidth | Check coalescing, increase vector width, tune BLOCK sizes | -| SM throughput > 60% peak | Compute-bound, near optimal | Reduce arithmetic, use faster instructions (e.g., FMA) | -| SM throughput < 30% peak | Underutilized compute | Increase occupancy, reduce warp stalls, check instruction mix | -| Achieved occupancy > 50% | Acceptable for most kernels | — | -| Achieved occupancy < 25% | Too few active warps | Reduce register pressure or shared memory; increase block size | - -#### Comparing before/after with ncu - -```bash -# Profile baseline kernel -ncu --target-processes all \ - --kernel-name "vectorized_elementwise_kernel" \ - --launch-skip 10 --launch-count 3 --set full \ - -o "${NCU_DIR}/baseline" ./program - -# Profile optimized kernel -ncu --target-processes all \ - --kernel-name "_fused_gated_residual_add_kernel" \ - --launch-skip 10 --launch-count 3 --set full \ - -o "${NCU_DIR}/optimized" ./program - -# Compare key metrics -for report in baseline optimized; do - echo "=== $report ===" - ncu --import "${NCU_DIR}/${report}.ncu-rep" \ - --page details --csv 2>/dev/null | grep -E "time_duration|throughput.*pct|occupancy" -done -``` - -**Decision rule after ncu analysis:** -- Kernel already at >80% DRAM bandwidth → fuse with neighbors to reduce total traffic -- Kernel at <50% DRAM bandwidth → tune block sizes, fix coalescing, increase vectorization -- Kernel compute-bound (SM util high, DRAM low) → reduce FLOPs or switch to a faster algorithm -- Low occupancy → reduce registers (simplify kernel) or increase block size in autotune configs - -### Step 4: Apply Kernel Optimization - -After pinpointing the slow op, choose the right tool: - -| Scenario | Skill to use | -|----------|-------------| -| 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. - -Both kernel types use SGLang's JIT compilation: -- **Triton**: `python/sglang/jit_kernel/diffusion/triton/.py` -- **CUDA JIT**: `python/sglang/jit_kernel/csrc/diffusion/.cuh` + wrapper `python/sglang/jit_kernel/diffusion/.py` - -### Step 5: torch.compile Coverage - -```bash -TORCH_COMPILE_DEBUG=1 sglang generate ... -``` -- Dynamic shape changes trigger recompilation → fix resolution and frame count when benchmarking -- `tensor.item()` in conditional branches causes graph breaks → rewrite as tensor ops - -### Step 6: Multi-GPU Efficiency (Wan2.2-T2V-A14B / MOVA) - -- Verify `--ulysses-degree` evenly divides `--num-gpus` -- Keep the command shape fixed when comparing kernels; for quick checks, reduce only `--num-inference-steps` -- If a run OOMs or jitters because of host contention, first confirm there are no leaked scheduler processes on the chosen GPU set - ---- - -## Optimization Workflow Summary - -``` -0. BASELINE - sglang generate --seed=42 --save-output → save reference images/videos - ↓ -1. BENCHMARK - Run benchmark commands above → record denoise latency baseline - ↓ -2. LEVEL 1 PROFILE (torch.profiler) - --profile --num-profiled-timesteps 3 - → parse .trace.json.gz → rank ops by CUDA time - → identify slow DiT layer (norm / attn / mlp / rope / adaln) - ↓ -3. LEVEL 2 PROFILE (nsys + gputrc2graph.py) - → result.csv category breakdown (gemm / attn / adaln / norm / triton / cpu) - → confirm where GPU time is concentrated - ↓ -4. LEVEL 3 PROFILE (ncu — per-kernel deep analysis) ★ CRITICAL - → ncu --set full on target kernel(s) - → extract DRAM bandwidth util, SM throughput, achieved occupancy - → determine if kernel is memory-bound, compute-bound, or latency-bound - → for CUDA graph: use --graph-profiling node - ↓ -5. KERNEL OPTIMIZATION - 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 - sglang generate --seed=42 --save-output → diff against reference - If output differs beyond tolerance → reject optimization - ↓ -7. RE-BENCHMARK - Verify denoise latency improvement; no regression on other models -``` - ---- - -## Checklist Before Merging - -### Correctness (must pass first) -- [ ] Reference outputs collected with `--seed=42 --save-output` **before** any change -- [ ] After change: regenerate with identical args and compare -- [ ] No visible quality degradation in generated images / videos -- [ ] Correctness verified on all benchmark models - -### Performance (only after correctness passes) -- [ ] All benchmark models executed; denoise latency ★, end-to-end, peak memory recorded -- [ ] No regression in denoise latency vs. previous baseline (±2% tolerance) -- [ ] New kernel shows measurable improvement on at least 2 models -- [ ] No new torch.compile graph breaks introduced -- [ ] Results reproducible with all offloads disabled and fixed `--seed=42` +- [ ] fixed-shape baseline perf dump saved +- [ ] fixed-shape new perf dump saved +- [ ] `compare_perf.py` table generated +- [ ] one representative `torch.profiler` trace saved +- [ ] hotspot classified against `existing-fast-paths.md` +- [ ] reference image or video checked for correctness +- [ ] any remaining kernel work handed to a specialized optimization skill diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md index eb98b24dd..a0e0c4c60 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md @@ -1,7 +1,8 @@ # SGLang Diffusion Fast Paths -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. +Use this guide when mapping a diffusion bottleneck to an existing fused path or +distributed overlap pattern in `sglang.multimodal_gen`. Prefer reuse and +configuration first before handing the problem to a specialized kernel-optimization skill. **Key Files** - `python/sglang/multimodal_gen/runtime/layers/layernorm.py` @@ -35,19 +36,29 @@ Prefer reuse before writing a new Triton or CUDA kernel. - 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 +3. Z-Image fused tanh/gate modulation +- Kernels: `fused_norm_tanh_mul_add`, `fused_norm_tanh_mul_add_norm_scale` +- Locations: `layernorm.py`, `cutedsl/norm_tanh_mul_add_norm_scale.py`, `zimage.py` +- Use cases: + - `y = tanh(gate) * norm(x) + shift` + - `y, y2 = tanh(gate) * norm(x) + shift`, then `y2 = norm(y) * (1 + scale)` +- Constraints: same CuTe DSL envelope as the norm+scale/shift family in practice: contiguous last dim, fp16/bf16/fp32, and `D % 256 == 0`, `D <= 8192`. +- Validation: `python/sglang/jit_kernel/tests/diffusion/test_norm_tanh_mul_add_norm_scale.py` +- Behavior: this is already a merged fast path, so if Z-Image traces show the unfused chain, treat it as a missing or regressed existing optimization before proposing a new kernel. + +4. 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) +5. 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`. - `torch.compile` note: keep this path behind the custom-op wrapper in `rmsnorm_onepass.py`; direct `wrap_triton` can recompile on dynamic row counts. -5. Triton RoPE fusion +6. Triton RoPE fusion - Kernel: `apply_rotary_embedding` - Locations: `triton/rotary.py`, `rotary_embedding/utils.py` - Use case: GPT-J style RoPE when not Neox. @@ -83,12 +94,53 @@ Prefer reuse before writing a new Triton or CUDA kernel. - 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. +**QK Norm + RoPE Optimization** + +- Entry point: `apply_qk_norm_rope` in `layernorm.py`. +- Fast path: JIT fused inplace QK norm + RoPE from `python/sglang/jit_kernel/diffusion/qknorm_rope.py` via `fused_inplace_qknorm_rope`. +- Toggle: `SGLANG_ENABLE_FUSED_QKNORM_ROPE=1` keeps the fused path enabled by default. +- Preconditions for fused path: + - CUDA only. + - `allow_inplace=True` and `q_eps == k_eps`. + - `q` / `k` are contiguous 4D tensors with the same shape. + - `q.dtype` is `fp16` or `bf16`, and norm weights match tensor dtype. + - `can_use_fused_inplace_qknorm_rope(head_dim, rope_dim, is_neox, dtype)` returns true. + - Supported head dims: `64, 128, 256`. +- Behavior: `apply_qk_norm_rope` prefers the fused JIT kernel when all guards pass; otherwise it falls back to `apply_qk_norm(...)` plus `apply_flashinfer_rope_qk_inplace(...)`. + +**Nunchaku Fused GELU MLP** + +- Entry point: `_fused_gelu_mlp` in `runtime/models/dits/flux.py`. +- Fast path: Nunchaku checkpoints can fuse `fc1 GEMM + GELU + shift + re-quant + fc2.lora_down` before the second GEMM instead of materializing a standalone GELU activation. +- Scope: this is a model-specific fast path for Nunchaku-quantized FLUX-family checkpoints. +- Workflow rule: if a Nunchaku trace shows split `fc1 -> gelu -> quant -> fc2.lora_down`, treat it as a missing existing fast path before proposing a new fusion. + +**NVFP4 / Nunchaku Packed QKV** + +- Entry points: `runtime/models/dits/flux.py`, `runtime/models/dits/flux_2.py`, and the FLUX config remapping in `configs/models/dits/flux.py`. +- Fast path: quantized FLUX-family checkpoints can store attention projections in packed QKV form, and SGLang intentionally switches to `MergedColumnParallelLinear` paths such as `to_qkv`, `to_added_qkv`, and `to_qkv_mlp_proj` instead of separate `to_q`, `to_k`, `to_v`. +- FLUX.2 NVFP4 note: `flux_2.py` explicitly enables fused packed QKV when `quant_config` is `ModelOptFp4Config`, because the NVFP4 checkpoint stores image-attention QKV packed on disk. +- Nunchaku note: raw and converted Nunchaku checkpoint names are remapped onto fused `to_qkv` / `to_added_qkv` names in `configs/models/dits/flux.py`; correctness on NVFP4-style checkpoints also depends on quant metadata such as `wtscale` and attention `wcscales`. +- Workflow rule: if an NVFP4 or Nunchaku trace shows split `to_q -> to_k -> to_v` where packed QKV is expected, treat it as a missing quantized fast path or checkpoint-format mismatch before proposing a new attention fusion. + **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`. +- Z-Image residual-form modulation: `fused_norm_tanh_mul_add` and `fused_norm_tanh_mul_add_norm_scale` in `zimage.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`. +- QK norm + RoPE: `apply_qk_norm_rope` in `layernorm.py`; use this path when the model wants fused attention prep instead of separate QK norm and RoPE calls. +- Nunchaku fused GELU MLP: `_fused_gelu_mlp` in `flux.py` for quantized FLUX-family checkpoints. +- NVFP4 / packed QKV attention: `to_qkv`, `to_added_qkv`, and `to_qkv_mlp_proj` in FLUX-family quantized paths. - RoPE: `_apply_rotary_emb` prefers Triton; Q/K RoPE prefers FlashInfer when present. +**Existing Overlap / Communication Families** + +- Ulysses / USP attention: treat `all_to_all`, `ring_attn`, and head / sequence reshards as an existing distributed attention family, not a new overlap idea. +- Turbo-layer async all-to-all: `all_to_all_single(..., async_op=True)` plus staged waits already form an existing overlap family in `turbo_layer.py`. +- TorchInductor compute / communication reorder: `torch._inductor.config.reorder_for_compute_comm_overlap = True` can already partially overlap compiled denoise traces. +- Dual-stream diffusion models: `use_dual_stream = True` in models such as `hunyuan3d.py` is an existing overlap family. +- Workflow rule: if a hotspot is communication-heavy, rule out these in-repo overlap families before proposing a brand new overlap design. + **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`. @@ -109,4 +161,4 @@ Prefer reuse before writing a new Triton or CUDA kernel. - 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 `benchmark-and-profile.md` in this directory for the canonical ncu workflow. +- If none of the families above match, package the evidence from the benchmark/profile skill and hand the kernel work to a specialized optimization skill such as `sglang-diffusion-ako4all-kernel`. diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/nsight-profiler.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/nsight-profiler.md deleted file mode 100644 index 5d15cb399..000000000 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/nsight-profiler.md +++ /dev/null @@ -1,286 +0,0 @@ ---- -name: nsight-profiler -description: Expert skill for NVIDIA Nsight Systems and Nsight Compute profiling tools. Configure profiling sessions, analyze kernel reports, interpret occupancy metrics, roofline model data, memory bandwidth bottlenecks, and warp execution efficiency. -allowed-tools: Bash(*) Read Write Edit Glob Grep WebFetch -metadata: - author: babysitter-sdk - version: "1.0.0" - category: performance-profiling - backlog-id: SK-002 - source: "Adapted from https://github.com/lobehub/lobehub (.agents/skills/nsight-profiler)" ---- - -> **Source**: This skill is adapted from the [lobehub/lobehub](https://github.com/lobehub/lobehub) open-source repository (`.agents/skills/nsight-profiler`). Original author: `babysitter-sdk`. - -# nsight-profiler - -You are **nsight-profiler** - a specialized skill for NVIDIA Nsight Systems and Nsight Compute profiling tools. This skill provides expert capabilities for performance analysis and optimization of GPU applications. - -## Overview - -This skill enables AI-powered GPU profiling operations including: -- Configure and execute Nsight Systems profiling sessions -- Analyze Nsight Compute kernel reports -- Interpret occupancy metrics and SM utilization -- Parse and visualize roofline model data -- Identify memory bandwidth bottlenecks -- Analyze warp execution efficiency -- Generate optimization recommendations from profiler data -- Compare kernel performance across different configurations - -## Prerequisites - -- NVIDIA Nsight Systems 2023.1+ -- NVIDIA Nsight Compute 2023.1+ -- CUDA Toolkit 11.0+ -- GPU with compute capability 7.0+ (for full profiling features) - -## Capabilities - -### 1. Nsight Systems Profiling - -System-wide performance analysis: - -```bash -# Basic system profile -nsys profile -o report ./cuda_program - -# Profile with CUDA API tracing -nsys profile -t cuda,nvtx,osrt -o report ./cuda_program - -# Capture GPU metrics -nsys profile --gpu-metrics-device=all -o report ./cuda_program - -# Profile specific duration -nsys profile -d 10 -o report ./cuda_program - -# Export to multiple formats (one type per command) -nsys export -t sqlite report.nsys-rep -nsys export -t json report.nsys-rep - -# Generate summary statistics -nsys stats report.nsys-rep -``` - -Child-process note: -- If the target launches worker processes, add `--trace-fork-before-exec=true`. -- For `sglang generate`, this is usually required to capture the real worker trace. - -### 2. Nsight Compute Profiling - -Detailed kernel analysis: - -```bash -# Profile all kernels -ncu --target-processes all -o profile ./cuda_program - -# Profile specific kernel -ncu --target-processes all --kernel-name myKernel -o profile ./cuda_program - -# Full metric collection -ncu --target-processes all --set full -o profile ./cuda_program - -# Roofline analysis -ncu --target-processes all --set roofline -o profile ./cuda_program - -# Memory analysis -ncu --target-processes all --section MemoryWorkloadAnalysis -o profile ./cuda_program - -# Import an existing report -ncu --import profile.ncu-rep --page details --csv -``` - -Child-process note: -- If the target forks or spawns subprocesses, use `--target-processes all`. -- For `sglang generate`, this is the safer default. - -### 3. Occupancy Analysis - -Analyze and optimize occupancy: - -```bash -# Collect occupancy metrics -ncu --section Occupancy -o occupancy ./cuda_program - -# Key metrics to analyze: -# - Achieved Occupancy -# - Theoretical Occupancy -# - Block Limit (registers, shared memory, warps) -# - Occupancy Limiter -``` - -```cuda -// Query occupancy in code -int numBlocks; -int blockSize = 256; -cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &numBlocks, myKernel, blockSize, sharedMemSize); - -float occupancy = (numBlocks * blockSize) / - (float)deviceProp.maxThreadsPerMultiProcessor; -printf("Theoretical Occupancy: %.2f%%\n", occupancy * 100); -``` - -### 4. Roofline Model Analysis - -Performance bound analysis: - -```bash -# Generate roofline data -ncu --set roofline -o roofline ./cuda_program - -# Key metrics: -# - Achieved FLOP/s -# - Achieved Memory Bandwidth -# - Arithmetic Intensity (FLOP/byte) -# - Ridge Point -``` - -Interpretation guide: -- Below memory roofline: Memory bound -- Below compute roofline: Compute bound -- At peak: Optimal utilization - -### 5. Memory Bandwidth Analysis - -Identify memory bottlenecks: - -```bash -# Memory analysis sections -ncu --section MemoryWorkloadAnalysis \ - --section MemoryWorkloadAnalysis_Chart \ - --section MemoryWorkloadAnalysis_Tables \ - -o memory ./cuda_program -``` - -Key metrics: -- Global Load/Store Throughput -- L1/L2 Cache Hit Rate -- Shared Memory Bandwidth -- Memory Transactions per Request - -### 6. Warp Execution Analysis - -Analyze warp efficiency: - -```bash -# Warp state analysis -ncu --section WarpStateStatistics -o warp ./cuda_program - -# Scheduler statistics -ncu --section SchedulerStatistics -o scheduler ./cuda_program -``` - -Key metrics: -- Warp Cycles Per Issued Instruction -- Eligible Warps Per Active Cycle -- Active Warps Per Scheduler -- Stall Reasons (memory, sync, execution) - -### 7. Kernel Comparison - -Compare kernel variants: - -```bash -# Step 1: Profile baseline -ncu --target-processes all --set full -o baseline ./program_v1 - -# Step 2: Profile optimized version -ncu --target-processes all --set full -o optimized ./program_v2 - -# Step 3: Export both profiles to CSV, then compare with Python (no GUI needed) -# Note: --import can only be specified once; --page diff is not a valid page value. -ncu --import baseline.ncu-rep --page details --csv > baseline_details.csv -ncu --import optimized.ncu-rep --page details --csv > optimized_details.csv - -python3 -c " -import csv -def load(p): - return {r.get('Metric Name',''): r.get('Metric Value','') - for r in csv.DictReader(open(p))} -b = load('baseline_details.csv') -o = load('optimized_details.csv') -for k in sorted(set(b) | set(o)): - bv, ov = b.get(k,''), o.get(k,'') - if bv != ov: - print(f'{k[:55]:<55} {bv} -> {ov}') -" -``` - -### 8. Performance Recommendations - -Automated analysis: - -```bash -# Get optimization recommendations -ncu --section SpeedOfLight \ - --section SpeedOfLight_RooflineChart \ - -o speedoflight ./cuda_program - -# Export with recommendations -ncu --import profile.ncu-rep --page details --csv > details.csv -``` - -## Common Profiling Workflows - -### Workflow 1: Initial Performance Assessment - -```bash -# Step 1: System overview -nsys profile -t cuda -o system_overview ./program -nsys stats system_overview.nsys-rep - -# Step 2: Identify hot kernels -ncu --launch-skip 10 --launch-count 5 -o hot_kernels ./program - -# Step 3: Deep dive on bottleneck kernel -ncu --kernel-name hotKernel --set full -o detailed ./program -``` - -### Workflow 2: Memory Optimization - -```bash -# Analyze memory access patterns -ncu --section SourceCounters \ - --section MemoryWorkloadAnalysis \ - --kernel-name targetKernel \ - -o memory_analysis ./program - -# Check for coalescing issues -ncu --metrics l1tex__t_sectors_pipe_lsu_mem_global_op_ld.sum,\ -l1tex__t_requests_pipe_lsu_mem_global_op_ld.sum \ - -o coalescing ./program -``` - -### Workflow 3: Occupancy Optimization - -```bash -# Profile with occupancy focus -ncu --section Occupancy \ - --section LaunchStatistics \ - -o occupancy ./program -``` - -**Interpreting occupancy limiters** (from the `Occupancy` section report): - -| Limiter shown | Fix | -|---------------|-----| -| `Registers` | Reduce register pressure: use fewer local variables, add `maxnreg` hint | -| `Shared Memory` | Decrease shared memory allocation or use 32-bit instead of 64-bit | -| `Block Size` | Increase threads per block; ensure block size is a multiple of warp size (32) | -| `Warp Limit` | Already at theoretical max for this SM; no action needed | - -> **For Triton kernels**: block sizes are controlled via `@triton.autotune` configs, not CLI flags. To test occupancy at different block sizes, add or modify the `triton.Config({"BLOCK_C": N}, num_warps=W)` entries in the autotune list and re-run. Do **not** pass `--block-size` as a CLI argument — the Triton benchmark script does not accept it. - -## Dependencies - -- Nsight Systems 2023.1+ -- Nsight Compute 2023.1+ -- CUDA Toolkit 11.0+ - -## Constraints - -- Full profiling requires root/admin privileges -- Some metrics only available on specific GPU architectures -- Profiling adds overhead; results may differ from production -- Nsight Compute profiles one kernel invocation at a time by default diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py index bc52f4a43..002c13429 100755 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py @@ -12,7 +12,7 @@ Usage: # Tag the run for later compare_perf.py usage python3 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py --model flux --label tuned - # All 11 preset models + # All 12 preset models python3 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py --all # Show preset order, model path, and nightly mapping @@ -157,7 +157,23 @@ MODELS = { "--guidance-scale=5.0", ], }, - # 8. Nightly: wan22_i2v_a14b_720p + # 8. Nightly: ltx2_twostage_t2v + "ltx2": { + "nightly_case_id": "ltx2_twostage_t2v", + "path": "Lightricks/LTX-2", + "prompt": "A beautiful sunset over the ocean", + "negative_prompt": "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static.", + "seed": 1234, + "extra_args": [ + "--pipeline-class-name=LTX2TwoStagePipeline", + "--width=1536", + "--height=1024", + "--num-frames=121", + "--fps=24", + "--num-gpus=1", + ], + }, + # 9. Nightly: wan22_i2v_a14b_720p # Requires: /inputs/diffusion_benchmark/figs/cat.png "wan-i2v": { "nightly_case_id": "wan22_i2v_a14b_720p", @@ -176,7 +192,7 @@ MODELS = { "--pin-cpu-memory", ], }, - # 9. Skill-only extra preset + # 10. Skill-only extra preset "hunyuanvideo": { "path": "hunyuanvideo-community/HunyuanVideo", "prompt": "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window.", @@ -189,7 +205,7 @@ MODELS = { "--num-inference-steps=30", ], }, - # 10. Skill-only extra preset + # 11. Skill-only extra preset # Requires: /inputs/diffusion_benchmark/figs/mova_single_person.jpg "mova-720p": { "path": "OpenMOSS-Team/MOVA-720p", @@ -205,7 +221,7 @@ MODELS = { "--num-inference-steps=2", ], }, - # 11. Skill-only extra preset + # 12. Skill-only extra preset "helios": { "path": "BestWishYsh/Helios-Base", "prompt": "A curious raccoon", @@ -278,8 +294,9 @@ def build_sglang_cmd( "--log-level=info", ] - if seed is not None: - cmd.append(f"--seed={seed}") + effective_seed = cfg.get("seed", seed) + if effective_seed is not None: + cmd.append(f"--seed={effective_seed}") if "negative_prompt" in cfg: cmd.append(f"--negative-prompt={cfg['negative_prompt']}") @@ -364,20 +381,28 @@ def run_benchmark_once( float(total_ms) / 1000.0 if total_ms is not None else None ) - # denoise latency: accept the canonical "DenoisingStage" plus - # model-specific variants such as "MOVADenoisingStage" and - # "HeliosChunkedDenoisingStage". - # steps = [{"name": "DenoisingStage", "duration_ms": 1234.5}, ...] + # denoise latency: sum all true denoise/refinement stages. + # This accepts variants such as "MOVADenoisingStage", + # "HeliosChunkedDenoisingStage", and the LTX-2 two-stage pair + # "LTX2AVDenoisingStage" + "LTX2RefinementStage", while excluding + # setup stages like "QwenImageLayeredBeforeDenoisingStage". denoise_latency_s = None + denoise_stage_total_ms = 0.0 for step in perf.get("steps", []): step_name = step.get("name") if ( isinstance(step_name, str) - and "DenoisingStage" in step_name and step.get("duration_ms") is not None + and ( + step_name.endswith("DenoisingStage") + or step_name.endswith("RefinementStage") + ) + and "BeforeDenoisingStage" not in step_name ): - denoise_latency_s = float(step["duration_ms"]) / 1000.0 - break + denoise_stage_total_ms += float(step["duration_ms"]) + + if denoise_stage_total_ms > 0.0: + denoise_latency_s = denoise_stage_total_ms / 1000.0 # fallback: sum all per-step durations from denoise_steps_ms # denoise_steps_ms = [{"step": 0, "duration_ms": 100.5}, ...] @@ -432,7 +457,9 @@ def print_results_table(results: list[dict]): print("-" * 92) print() - print("★ Denoise latency = total DiT forward pass time across all inference steps.") + print( + "★ Denoise latency = sum of stages ending with DenoisingStage plus any RefinementStage." + ) print( " Compare two runs with python/sglang/multimodal_gen/benchmarks/compare_perf.py." ) @@ -447,7 +474,7 @@ def main(): choices=list(MODELS.keys()), help="Model to benchmark (default: flux)", ) - parser.add_argument("--all", action="store_true", help="Benchmark all 11 models") + parser.add_argument("--all", action="store_true", help="Benchmark all 12 models") parser.add_argument( "--list-models", action="store_true", diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_rmsnorm.py b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_rmsnorm.py deleted file mode 100755 index 55ca7123f..000000000 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_rmsnorm.py +++ /dev/null @@ -1,210 +0,0 @@ -""" -Micro-benchmark for the SGLang Diffusion JIT CUDA RMSNorm kernel. - -Compares: - 1. SGLang JIT CUDA kernel (diffusion_rmsnorm) - 2. PyTorch baseline (torch.nn.functional.rms_norm) - -Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels - -Usage: - cd /path/to/sglang - 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. - # This script auto-selects an idle GPU when CUDA_VISIBLE_DEVICES is unset. -""" - -import os -import sys -import time -from pathlib import Path -from typing import Tuple - -SCRIPT_DIR = Path(__file__).resolve().parent -if str(SCRIPT_DIR) not in sys.path: - sys.path.insert(0, str(SCRIPT_DIR)) - -from diffusion_skill_env import configure_runtime_env - -configure_runtime_env(required_gpus=1) - -import torch - -# --------------------------------------------------------------------------- -# Import the JIT CUDA kernel. -# When you implement the sglang-diffusion-cuda-kernel workflow, the file will be at: -# python/sglang/jit_kernel/diffusion/rmsnorm.py -# --------------------------------------------------------------------------- -try: - from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm - - JIT_AVAILABLE = True -except ImportError: - JIT_AVAILABLE = False - print( - "WARNING: diffusion.rmsnorm JIT kernel not available. " - "Run after implementing the sglang-diffusion-cuda-kernel workflow." - ) - - -def pytorch_rmsnorm( - x: torch.Tensor, - weight: torch.Tensor | None = None, - eps: float = 1e-6, -) -> torch.Tensor: - """Reference PyTorch implementation of RMSNorm.""" - hidden = x.shape[-1] - return torch.nn.functional.rms_norm( - x.float(), (hidden,), weight.float() if weight is not None else None, eps=eps - ).to(x.dtype) - - -def benchmark_kernel( - func, - args, - warmup: int = 20, - iterations: int = 100, -) -> Tuple[float, float]: - """Benchmark a kernel function. Returns (avg_ms, min_ms).""" - for _ in range(warmup): - func(*args) - torch.cuda.synchronize() - - times = [] - for _ in range(iterations): - torch.cuda.synchronize() - t0 = time.perf_counter() - func(*args) - torch.cuda.synchronize() - times.append((time.perf_counter() - t0) * 1000) - - return sum(times) / len(times), min(times) - - -def run_benchmark(): - print("=" * 72) - print("SGLang Diffusion RMSNorm Micro-Benchmark: JIT CUDA vs PyTorch") - print("=" * 72) - print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', '')}") - print(f"Device: {torch.cuda.get_device_name(0)}") - cap = torch.cuda.get_device_capability() - print(f"Compute Capability: sm_{cap[0]}{cap[1]}") - print() - - if not JIT_AVAILABLE: - print("Skipping JIT kernel benchmark (kernel not available).") - return - - # Determine dtype: T4 (sm_75) has no BF16 - dtype = torch.bfloat16 if cap >= (8, 0) else torch.float16 - print(f"Dtype: {dtype}") - print() - - # Typical DiT hidden sizes for sglang diffusion models: - # FLUX.1-dev: hidden=3072 - # Qwen-Image: hidden=2048 - # Wan2.2: hidden=4096 - configs = [ - # (batch_tokens, hidden_size, has_weight) - (1024, 2048, True), # Qwen-Image: 1 sample × 1024 tokens - (4096, 2048, True), # Qwen-Image: larger batch - (1024, 3072, True), # FLUX: 1 sample × 1024 tokens - (4096, 3072, True), # FLUX: larger - (4096, 4096, True), # Wan2.2 - (4096, 2048, False), # no-weight (elementwise_affine=False) - (16384, 3072, True), # long sequence - ] - - print( - f"{'Config':<32} {'JIT(ms)':>10} {'PyTorch(ms)':>12} {'Speedup':>9} {'Weight'}" - ) - print("-" * 72) - - total_speedup = 0 - n = 0 - - for batch_tokens, hidden, has_weight in configs: - x = torch.randn(batch_tokens, hidden, dtype=dtype, device="cuda") - weight = torch.ones(hidden, dtype=dtype, device="cuda") if has_weight else None - - jit_avg, _ = benchmark_kernel( - diffusion_rmsnorm, (x, weight, 1e-6), warmup=20, iterations=100 - ) - pt_avg, _ = benchmark_kernel( - pytorch_rmsnorm, (x, weight, 1e-6), warmup=20, iterations=100 - ) - - speedup = pt_avg / jit_avg - total_speedup += speedup - n += 1 - - w_str = "yes" if has_weight else "no " - cfg = f"[{batch_tokens}×{hidden}]" - print(f"{cfg:<32} {jit_avg:>10.3f} {pt_avg:>12.3f} {speedup:>8.2f}x {w_str}") - - print("-" * 72) - print(f"{'Average Speedup':>56} {total_speedup / n:.2f}x") - print() - - # ----------------------------------------------------------------------- - # Correctness check - # ----------------------------------------------------------------------- - print("Correctness Check (BF16 tolerance 0.02):") - x = torch.randn(4096, 3072, dtype=dtype, device="cuda") - weight = torch.ones(3072, dtype=dtype, device="cuda") - - out_jit = diffusion_rmsnorm(x, weight=weight, eps=1e-6) - out_ref = pytorch_rmsnorm(x, weight=weight, eps=1e-6) - - max_diff = (out_jit - out_ref).abs().max().item() - rel_diff = ((out_jit - out_ref).abs() / (out_ref.abs() + 1e-8)).max().item() - passed = max_diff < 0.02 - - print(f" Max absolute diff: {max_diff:.2e}") - print(f" Max relative diff: {rel_diff:.2e}") - print(f" Correctness: {'PASS ✓' if passed else 'FAIL ✗'}") - print() - - # ----------------------------------------------------------------------- - # Memory bandwidth analysis - # ----------------------------------------------------------------------- - print("Memory Bandwidth Analysis:") - bt, hid = 4096, 3072 - x = torch.randn(bt, hid, dtype=dtype, device="cuda") - weight = torch.ones(hid, dtype=dtype, device="cuda") - - bytes_per_elem = dtype.itemsize - total_bytes = ( - bt * hid + hid + bt * hid - ) * bytes_per_elem # read x + read w + write out - jit_avg, _ = benchmark_kernel(diffusion_rmsnorm, (x, weight, 1e-6)) - - bandwidth_gbps = (total_bytes / 1e9) / (jit_avg / 1000) - device_name = torch.cuda.get_device_name(0) - if "H200" in device_name: - theoretical_bw = 4800 # H200 SXM: ~4.8 TB/s - else: - theoretical_bw = { - (9, 0): 3350, # H100: 3.35 TB/s - (8, 0): 2000, # A100 80GB - }.get( - cap, 320 - ) # T4: 320 GB/s - efficiency = bandwidth_gbps / theoretical_bw * 100 - - print(f" Shape: [{bt} × {hid}] dtype: {dtype}") - print(f" Total data: {total_bytes / 1e6:.1f} MB") - print(f" Achieved: {bandwidth_gbps:.1f} GB/s") - print(f" Theoretical ({device_name}): {theoretical_bw} GB/s") - print(f" Bandwidth efficiency: {efficiency:.1f}%") - print() - print("Target: ≥ 30% efficiency (H200/H100/A100), ≥ 40% (T4)") - - -if __name__ == "__main__": - if not torch.cuda.is_available(): - print("CUDA not available.") - else: - run_benchmark() diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py index 6148c0630..cffbdebf7 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/scripts/diffusion_skill_env.py @@ -9,7 +9,6 @@ from pathlib import Path OUTPUT_DIR_NAMES = { "benchmarks": Path("outputs/diffusion_benchmarks"), "profiles": Path("outputs/diffusion_profiles"), - "ncu": Path("outputs/ncu_reports"), } diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/SKILL.md deleted file mode 100644 index 34f728fb9..000000000 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/SKILL.md +++ /dev/null @@ -1,516 +0,0 @@ ---- -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 `../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. -> -> **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 -> - [../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** (`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 | - -For most diffusion-model elementwise ops, **start with Triton**. Switch to CUDA when profiling shows Triton can't reach hardware bandwidth limits. - -## Directory Layout - -``` -python/sglang/jit_kernel/ -├── csrc/ -│ ├── diffusion/ # JIT CUDA source files for diffusion kernels (this skill) -│ │ ├── timestep_embedding.cuh # existing example -│ │ ├── rmsnorm.cuh # NEW: add here -│ │ └── adaln.cuh # NEW: add here -│ └── elementwise/ # shared JIT CUDA csrc (non-diffusion) -├── diffusion/ -│ ├── triton/ # Triton kernels (scale_shift, norm, rope, ...) -│ ├── cutedsl/ # CuTe DSL kernels -│ └── rmsnorm.py # NEW: CUDA JIT Python wrapper (add here) -├── timestep_embedding.py # existing CUDA diffusion kernel Python wrapper (legacy) -``` - -New diffusion CUDA kernel source files go into `python/sglang/jit_kernel/csrc/diffusion/.cuh`. -The Python wrapper goes at `python/sglang/jit_kernel/diffusion/.py` -(inside `diffusion/`, alongside the `triton/` and `cutedsl/` subdirectories). - ---- - -## SGLang Kernel Abstractions (Required) - -Always use these — do **not** use raw CUDA primitives directly. - -```cpp -#include // TensorMatcher, SymbolicSize, SymbolicDevice -#include // fp16_t, bf16_t, fp32_t, dtype_trait, packed_t -#include // RuntimeCheck, div_ceil -#include // LaunchKernel, SGL_DEVICE, type aliases -#include // AlignedVector — 128-bit vector loads -#include // warp::reduce_sum, warp::reduce_max -#include // device::math::rsqrt, sqrt, ... -#include // tile::Memory (strided access pattern) -``` - -Key types: `fp16_t` = `__half`, `bf16_t` = `__nv_bfloat16`, `fp32_t` = `float`. -Packed variants: `fp16x2_t`, `bf16x2_t`. Use `packed_t` for the 2-element alias. - ---- - -## Step 1: Write the CUDA Kernel - -Create `python/sglang/jit_kernel/csrc/diffusion/rmsnorm.cuh` (RMSNorm as example). - -### 1a. Vectorized RMSNorm Kernel - -```cpp -#include -#include -#include -#include -#include -#include - -#include -#include - -namespace { - -// --------------------------------------------------------------- -// RMSNorm kernel: y = x / rms(x) * weight -// T = fp16_t | bf16_t | fp32_t -// kVecN = vectorized elements per load (8 for fp16/bf16, 4 for fp32) -// --------------------------------------------------------------- -template -__global__ void rmsnorm_kernel( - T* __restrict__ dst, - const T* __restrict__ src, - const T* __restrict__ weight, // may be nullptr if no affine weight - uint32_t hidden_size, - uint32_t n_vecs, // hidden_size / kVecN - float eps) -{ - using vec_t = device::AlignedVector; - - const uint32_t row = blockIdx.x; - const T* row_src = src + row * hidden_size; - T* row_dst = dst + row * hidden_size; - - // --- Pass 1: accumulate sum of squares (vectorized) --- - float sum_sq = 0.f; - for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { - vec_t v; - v.load(row_src, vi); - #pragma unroll - for (int i = 0; i < kVecN; ++i) { - float val = static_cast(v[i]); - sum_sq += val * val; - } - } - - // --- Warp reduction --- - sum_sq = device::warp::reduce_sum(sum_sq); - - // --- Block reduction via shared memory --- - __shared__ float smem[32]; - if (threadIdx.x % 32 == 0) { - smem[threadIdx.x / 32] = sum_sq; - } - __syncthreads(); - if (threadIdx.x < 32) { - sum_sq = (threadIdx.x < blockDim.x / 32) ? smem[threadIdx.x] : 0.f; - sum_sq = device::warp::reduce_sum(sum_sq); - } - __syncthreads(); - - const float rms_inv = device::math::rsqrt(sum_sq / static_cast(hidden_size) + eps); - - // --- Pass 2: normalize + apply weight (vectorized) --- - for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { - vec_t v_in, v_w, v_out; - v_in.load(row_src, vi); - if (weight != nullptr) { - v_w.load(weight, vi); - } - #pragma unroll - for (int i = 0; i < kVecN; ++i) { - float val = static_cast(v_in[i]) * rms_inv; - if (weight != nullptr) { - val *= static_cast(v_w[i]); - } - v_out[i] = static_cast(val); - } - v_out.store(row_dst, vi); - } -} - -// --------------------------------------------------------------- -// Launcher -// --------------------------------------------------------------- -template -void rmsnorm( - tvm::ffi::TensorView dst, - tvm::ffi::TensorView src, - tvm::ffi::TensorView weight, // pass empty / nullptr for no-weight case - float eps) -{ - using namespace host; - - // Validate - SymbolicSize B{"batch_tokens"}, H{"hidden_size"}; - SymbolicDevice device; - device.set_options(); - - TensorMatcher({B, H}) - .with_dtype() - .with_device(device) - .verify(dst) - .verify(src); - - const uint32_t num_rows = static_cast(B.unwrap()); - const uint32_t hidden = static_cast(H.unwrap()); - const DLDevice dev = device.unwrap(); - - RuntimeCheck(hidden % (16 / sizeof(T)) == 0, - "rmsnorm: hidden_size must be divisible by vector width, got ", hidden); - - constexpr int kVecN = 16 / sizeof(T); // 128-bit vector: 8×fp16/bf16, 4×fp32 - const uint32_t n_vecs = hidden / kVecN; - - // Thread count: enough warps to cover n_vecs, max 512 threads - uint32_t threads = std::min(n_vecs, 512u); - threads = (threads + 31) / 32 * 32; // round up to warp boundary - - const T* w_ptr = (weight.data_ptr() != nullptr) - ? static_cast(weight.data_ptr()) : nullptr; - - LaunchKernel(num_rows, threads, dev)( - rmsnorm_kernel, - static_cast(dst.data_ptr()), - static_cast(src.data_ptr()), - w_ptr, - hidden, - n_vecs, - eps); -} - -} // namespace -``` - ---- - -## Step 2: Python Wrapper - -Create `python/sglang/jit_kernel/diffusion/rmsnorm.py`: - -```python -from __future__ import annotations -from typing import TYPE_CHECKING - -import torch - -from sglang.jit_kernel.utils import ( - cache_once, - is_arch_support_pdl, - load_jit, - make_cpp_args, -) - -if TYPE_CHECKING: - from tvm_ffi.module import Module - - -@cache_once -def _jit_rmsnorm_module(hidden_size: int, dtype: torch.dtype) -> Module: - args = make_cpp_args(hidden_size, is_arch_support_pdl(), dtype) - return load_jit( - "diffusion_rmsnorm", - *args, - cuda_files=["diffusion/rmsnorm.cuh"], # relative to csrc/ - cuda_wrappers=[("rmsnorm", f"RMSNormKernel<{args}>::run")], - ) - - -def diffusion_rmsnorm( - src: torch.Tensor, - weight: torch.Tensor | None = None, - eps: float = 1e-6, - out: torch.Tensor | None = None, -) -> torch.Tensor: - """ - RMSNorm for diffusion DiT layers. - - y = x / rms(x) * weight (weight=None → no affine scaling) - - Supported fast path: float16 / bfloat16. - For unsupported combinations (for example some float32 configs), - fall back to torch.nn.functional.rms_norm. - """ - assert src.is_cuda, "src must be a CUDA tensor" - assert src.dtype in (torch.float16, torch.bfloat16, torch.float32) - hidden_size = src.shape[-1] - - if out is None: - out = torch.empty_like(src) - - w = weight if weight is not None else torch.ones(hidden_size, dtype=src.dtype, device=src.device) - - module = _jit_rmsnorm_module(hidden_size, src.dtype) - module.rmsnorm(src.reshape(-1, hidden_size), w, out.reshape(-1, hidden_size), eps) - return out -``` - -**Key rules for the wrapper:** -- Use `cache_once` — never `functools.lru_cache` (breaks `torch.compile`) -- Include every compile-time specialization parameter in the cache key (`hidden_size`, PDL support, dtype here) -- `cuda_files` are relative to `python/sglang/jit_kernel/csrc/` -- `cuda_wrappers`: `(python_name, cpp_template_instantiation)` - ---- - -## Step 3: Integrate into Runtime (Optional, After Standalone Validation) - -The kernel replaces a slow operator inside the DiT forward pass. Find the correct module in: - -``` -python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py -python/sglang/multimodal_gen/runtime/models/dits/.py -``` - -There is no built-in `SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS` hook in the runtime. After the standalone test/benchmark passes, wire the new kernel into the actual execution path explicitly. A minimal pattern is to monkey-patch the target RMSNorm modules before `torch.compile` or any CPU offload setup: - -```python -from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm - -def _patch_rmsnorm(model: torch.nn.Module) -> None: - for name, module in model.named_modules(): - cls_name = type(module).__name__ - if cls_name in ("RMSNorm", "LlamaRMSNorm") or "RMSNorm" in cls_name: - eps = getattr(module, "eps", getattr(module, "variance_epsilon", 1e-6)) - has_weight = hasattr(module, "weight") and module.weight is not None - - if has_weight: - def _make_fwd(mod, epsilon): - def forward(x): - return diffusion_rmsnorm(x, weight=mod.weight, eps=epsilon) - return forward - module.forward = _make_fwd(module, eps) - else: - def _make_fwd_noweight(epsilon): - def forward(x): - return diffusion_rmsnorm(x, weight=None, eps=epsilon) - return forward - module.forward = _make_fwd_noweight(eps) -``` - -**Critical:** inject kernels **before** `torch.compile` and before any CPU offload is enabled. - ---- - -## Step 4: Key Kernel Patterns Reference - -### Diffusion-Specific Operators - -| Operator | Kernel Pattern | Notes | -|----------|---------------|-------| -| **RMSNorm** | 2-pass row reduction + vectorized normalize | Weight may be `None` (`elementwise_affine=False`) | -| **AdaLN modulation** | `y = norm(x) * (1 + scale) + shift` | Fuse norm + scale + shift in one pass | -| **RoPE 3D** | Read `(t, h, w)` cos/sin tables, apply to `(q, k)` | Layout: `[batch, t*h*w, heads, head_dim]` | -| **GEGLU** | Split last dim → `gate * silu(linear)` | Input `[B, L, 2*H]` → output `[B, L, H]` | -| **SiLU gate** | `out = a * sigmoid(a)` fused | Avoid separate elementwise ops | - -### Vectorized Memory Access - -```cpp -// BF16: 8 elements × 2 bytes = 16 bytes per vector load (AlignedVector) -// FP16: 8 elements × 2 bytes = 16 bytes (AlignedVector) -// FP32: 4 elements × 4 bytes = 16 bytes (AlignedVector) -constexpr int kVecN = 16 / sizeof(T); -using vec_t = device::AlignedVector; -``` - -### Warp / Block Reductions - -```cpp -// Warp reduction (within 32 threads) -float result = device::warp::reduce_sum(partial); - -// Block reduction via shared memory (see rmsnorm example above) -__shared__ float smem[32]; -// ... write warp-leaders into smem, sync, reduce again -``` - -### Thread Configuration - -```cpp -// Element-wise (RoPE, GEGLU, SiLU): simple 1D grid -constexpr uint32_t kBlock = 256; -uint32_t grid = host::div_ceil(total_elements, kBlock); -LaunchKernel(grid, kBlock, dev)(kernel, ...); - -// Row reduction (RMSNorm, LayerNorm): one block per row -uint32_t threads = std::min(hidden_size / kVecN, 512u); -threads = (threads + 31) / 32 * 32; -LaunchKernel(num_rows, threads, dev)(kernel, ...); -``` - ---- - -## Step 5: GPU Architecture Targets - -| GPU | Compute Cap | Memory BW | BF16 | Key Note | -|-----|------------|-----------|------|----------| -| H100 | sm_90 | 3.35 TB/s | Yes | Primary target; 132 SMs, 192 KB shared mem/SM | -| A100 | sm_80 | 2.0 TB/s | Yes | 108 SMs, 164 KB shared mem/SM | -| T4 | sm_75 | 320 GB/s | **No** | FP16 only; no `__nv_bfloat16` | - -If kernel requires SM90+ features (e.g., TMA, wgmma), raise a clear error: - -```python -if torch.cuda.get_device_capability()[0] < 9: - raise RuntimeError("This kernel requires SM90 (H100/Hopper) or later") -``` - -**Grid sizing for H100** (132 SMs): aim for grid multiples of 132 for good occupancy. - ---- - -## Step 6: Tests - -For this tutorial kernel, the repo now includes a verified regression test at `python/sglang/jit_kernel/tests/test_diffusion_rmsnorm.py`. Model new kernel tests after it: - -```python -import pytest -import torch -from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm - - -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("shape", [(1, 2048), (4, 3072), (16, 4096)]) -@pytest.mark.parametrize("has_weight", [True, False]) -def test_rmsnorm_correctness(dtype, shape, has_weight): - batch, hidden = shape - src = torch.randn(batch, hidden, dtype=dtype, device="cuda") - weight = torch.randn(hidden, dtype=dtype, device="cuda") if has_weight else None - - out_jit = diffusion_rmsnorm(src, weight=weight, eps=1e-6) - - # Reference: torch.nn.functional - ref = torch.nn.functional.rms_norm( - src.float(), (hidden,), weight.float() if weight is not None else None, eps=1e-6 - ).to(dtype) - - tol = {"rtol": 1e-2, "atol": 1e-2} if dtype != torch.float32 else {"rtol": 1e-5, "atol": 1e-6} - torch.testing.assert_close(out_jit, ref, **tol) - - -if __name__ == "__main__": - import sys - sys.exit(pytest.main([__file__, "-v", "-s"])) -``` - ---- - -## Step 7: Benchmark - -For the RMSNorm example in this skill, use the checked-in micro-benchmark script `scripts/bench_diffusion_rmsnorm.py`. For new kernels, follow the same structure or model a `triton.testing` benchmark after `python/sglang/jit_kernel/benchmark/bench_rmsnorm.py`. - ---- - -## Step 8: Profile with Nsight Compute (required) - -After correctness + benchmarking, you must collect **Nsight Compute (ncu)** data to validate: - -- Whether the kernel reaches reasonable bandwidth/throughput (avoid false positives where it is “faster” but under-utilizes hardware) -- Whether there are clear occupancy / register / shared memory limiters - -Use the canonical docs in this directory (do not duplicate CLI details across multiple skills): - -- `../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) - ---- - -## Common Pitfalls - -| Issue | Fix | -|-------|-----| -| `RMSNorm weight is None` | Use `type(module).__name__` check; pass `None` weight explicitly | -| `isinstance(m, torch.nn.RMSNorm)` misses diffusers variants | Use `"RMSNorm" in type(m).__name__` | -| Kernel patched after `torch.compile` | Inject **before** any compile call | -| Kernel patched after `enable_model_cpu_offload()` | Inject **before** CPU offload | -| `hidden_size` not divisible by `kVecN` | Add `RuntimeCheck(hidden % kVecN == 0, ...)` in launcher | -| `torch.compile` fails with custom CUDA kernel | Register as `@torch.library.custom_op` or use Triton instead | -| T4 GPU with BF16 kernel | Gate on compute capability; T4 is `sm_75`, no native BF16 | - ---- - -## Summary of Files - -``` -python/sglang/jit_kernel/csrc/diffusion/ -└── rmsnorm.cuh # NEW: JIT CUDA kernel source - -python/sglang/jit_kernel/diffusion/ -└── rmsnorm.py # NEW: Python wrapper + load_jit - -python/sglang/jit_kernel/tests/ -└── test_diffusion_rmsnorm.py # NEW: correctness tests - -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 -``` - ---- - -## References - -### This Skill's Extended Docs (references/ and scripts/) - -| File | Contents | -|------|----------| -| [references/kernel-templates.md](references/kernel-templates.md) | Copy-paste templates: element-wise, RMSNorm, AdaLN, Python wrapper, test, benchmark | -| [references/troubleshooting.md](references/troubleshooting.md) | Build errors, perf issues, torch.compile compatibility, debugging checklist | -| [references/h100-optimization-guide.md](references/h100-optimization-guide.md) | H100 (sm_90): memory hierarchy, warp reductions, occupancy, vectorization benchmarks | -| [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, low bandwidth, tile size limits, memory constraints | -| [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) | Micro-benchmark: JIT CUDA RMSNorm vs PyTorch, correctness check, bandwidth analysis | -| [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) | End-to-end preset runner. Save perf dumps per label, then compare with `compare_perf.py` | - -### SGLang Internals - -- **JIT system**: `add-jit-kernel` skill (`sglang/.claude/skills/add-jit-kernel/SKILL.md`) -- **JIT utils**: `python/sglang/jit_kernel/utils.py` — `cache_once`, `load_jit`, `make_cpp_args` -- **Abstractions**: `python/sglang/jit_kernel/include/sgl_kernel/` — `tensor.h`, `utils.cuh`, `vec.cuh`, `warp.cuh`, `math.cuh`, `tile.cuh` -- **Real csrc examples**: `python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh`, `python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh` - -### Other Diffusion Kernel Skills (this directory) - -- **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 - -- [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) — original source adapted for this skill diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/a100-optimization-guide.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/a100-optimization-guide.md deleted file mode 100644 index af3a24c09..000000000 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/a100-optimization-guide.md +++ /dev/null @@ -1,283 +0,0 @@ -# A100 GPU Optimization Guide — SGLang Diffusion JIT Kernels - -Deep dive into A100-specific optimizations for diffusion model CUDA kernels in SGLang's JIT system. - -> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) - ---- - -## A100 Ampere Architecture Overview - -| Component | A100 40GB | A100 80GB | Notes | -|-----------|-----------|-----------|-------| -| Compute Capability | sm_80 | sm_80 | Use `"-arch=sm_80"` in `extra_cuda_cflags` | -| SMs | 108 | 108 | Grid: aim for multiples of 108 | -| Shared Memory | 164 KB/SM | 164 KB/SM | Configurable: 48/96/164 KB | -| L2 Cache | 40 MB | 40 MB | Less than H100 (50 MB) | -| Memory Bandwidth | 1.55 TB/s | 2.0 TB/s | HBM2e | -| Max Threads/SM | 2048 | 2048 | Same as H100 | -| Tensor Cores | 3rd gen | 3rd gen | FP16, BF16, TF32, INT8, INT4 | - -### A100 vs H100 Comparison - -| Feature | A100 | H100 | Impact on JIT Kernels | -|---------|------|------|-----------------------| -| Memory BW | 2.0 TB/s | 3.35 TB/s | H100 ~67% faster for memory-bound ops | -| SMs | 108 | 132 | Adjust persistent kernel grid sizing | -| Shared Mem/SM | 164 KB | 192 KB | Reduce max tile sizes on A100 | -| L2 Cache | 40 MB | 50 MB | Attention tile reuse still works well | -| TMA | No | Yes | Can't use `cp.async.bulk` on A100 | -| FP8 | No | Yes | Use FP16/BF16 only on A100 | - ---- - -## Memory Access Optimization - -Same coalescing and vectorization rules as H100; lower bandwidth makes them even more critical. - -### `AlignedVector` Vectorization (same pattern as H100) - -```cpp -#include - -constexpr int kVecN = 16 / sizeof(T); // 8 for bf16/fp16, 4 for fp32 -using vec_t = device::AlignedVector; - -vec_t v; -v.load(src, vi); -// ... process elements ... -v.store(dst, vi); -``` - -**Expected A100 performance (BF16 RMSNorm):** - -| Implementation | A100 (ms) | H100 (ms) | A100 Speedup | -|:---|:---:|:---:|:---:| -| Scalar loads | ~0.10 | 0.065 | 1.00x | -| `AlignedVector` | ~0.03 | 0.019 | ~3x | - -**Target bandwidth**: 30–40% of A100's 2.0 TB/s = 600–800 GB/s. - -### Shared Memory Configuration - -```cpp -// A100 max: 164 KB/SM -cudaFuncSetAttribute( - your_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - 164 * 1024 // 164 KB max on A100 -); -``` - -Attention tile sizes for A100: - -``` -BLOCK_SIZE_M = 128 (Q block) -BLOCK_SIZE_N = 64 (K,V block) -Tile = 128×64×2 = 16 KB (FP16) — fits in 164 KB shared mem -``` - ---- - -## Occupancy Tuning - -**Grid sizing for A100 (108 SMs):** - -```cpp -#include - -// Cap blocks to SM × occupancy (same pattern as H100) -static const uint32_t max_occ = host::runtime::get_blocks_per_sm(kernel, kBlockSize); -static const uint32_t num_sm = host::runtime::get_sm_count(device.unwrap().device_id); -const uint32_t num_blocks = std::min(num_sm * max_occ, host::div_ceil(n, kBlockSize)); -``` - -**Recommended block sizes (same as H100):** - -| Kernel Type | Threads/Block | Notes | -|-------------|---------------|-------| -| Element-wise | 256 | High occupancy | -| Row reduction | 512 | Full reduction per row | -| Tiled/attention | 256 | Balance shared mem | - ---- - -## A100-Specific Features - -### Async Memory Copy (sm_80) - -A100 introduced `cp.async` for overlapping compute and memory. Use this in custom kernels for prefetching: - -```cuda -#if __CUDA_ARCH__ >= 800 -// Async copy from global to shared (A100+) -__pipeline_memcpy_async(smem_ptr, global_ptr, bytes); -__pipeline_commit(); -__pipeline_wait_prior(0); -#endif -``` - -### TF32 Mode (A100 specific) - -Enables FP32-range with FP16-like throughput for GEMM. Enable in Python: - -```python -# Enable TF32 for matmuls (A100+) -torch.backends.cuda.matmul.allow_tf32 = True -torch.backends.cudnn.allow_tf32 = True -``` - -TF32 is automatic for FP32 GEMMs via cuBLAS — no kernel changes needed. - -### Structural Sparsity (2:4) - -A100 tensor cores support 50% structured sparsity: - -```python -from torch.sparse import to_sparse_semi_structured -sparse_weight = to_sparse_semi_structured(dense_weight) -# ~2x GEMM speedup for matmul with sparse weight -``` - ---- - -## JIT Compilation for A100 - -```python -return load_jit( - "my_kernel", - *args, - cuda_files=["diffusion/my_kernel.cuh"], - cuda_wrappers=[("my_kernel", f"my_kernel<{args}>")], - extra_cuda_cflags=[ - "-O3", - "--use_fast_math", - "-arch=sm_80", # A100 only; omit for multi-arch - ], -) -``` - -**Multi-arch (A100 + H100):** - -```python -extra_cuda_cflags=[ - "-O3", "--use_fast_math", - "-gencode=arch=compute_80,code=sm_80", # A100 - "-gencode=arch=compute_90,code=sm_90", # H100 -] -``` - -Runtime arch guard (in Python wrapper): - -```python -cap = torch.cuda.get_device_capability() -if cap < (8, 0): - raise RuntimeError(f"This kernel requires sm_80 (A100) or later, got sm_{cap[0]}{cap[1]}") -``` - ---- - -## H100 → A100 Migration Checklist - -When porting an H100-optimized kernel to A100: - -| Item | H100 | A100 | Change Required | -|------|------|------|-----------------| -| Shared memory | 192 KB | 164 KB | Reduce `cudaFuncSetAttribute` size | -| Grid sizing | ×132 SMs | ×108 SMs | `get_sm_count()` handles automatically | -| TMA bulk copy | Available | **Not available** | Remove `cp.async.bulk`; use standard `__pipeline_memcpy_async` | -| FP8 | Available | **Not available** | Fall back to FP16/BF16 | -| PDL | Supported | Supported | `.enable_pdl(true)` works on sm_80 | -| Warp shuffles | Same | Same | No changes | -| `AlignedVector` | Same | Same | No changes | - -**Conditional compilation:** - -```cuda -#if __CUDA_ARCH__ >= 900 - // H100-only: TMA, FP8, thread block clusters - #define USE_TMA 1 -#elif __CUDA_ARCH__ >= 800 - // A100: cp.async, TF32, 2:4 sparsity - #define USE_ASYNC_COPY 1 -#endif -``` - ---- - -## Precision Notes - -| Type | Available on A100 | Notes | -|------|-------------------|-------| -| FP16 | Yes | Good, watch overflow in attention | -| BF16 | Yes | Preferred for training and inference | -| TF32 | Yes (A100 specific) | Auto for FP32 GEMMs | -| FP8 | **No** | H100 only | - ---- - -## Performance Profiling - -### NVIDIA Nsight Systems (nsys) - -```bash -nsys profile -o a100_profile python scripts/bench_diffusion_rmsnorm.py - -# Key metrics to watch: -# - Kernel duration -# - Memory transfer time -# - GPU idle time -# - Stream utilization -``` - -### NVIDIA Nsight Compute (ncu) - -```bash -# Full metrics -ncu --set full -o a100_metrics.ncu-rep \ - python scripts/bench_diffusion_rmsnorm.py - -# Specific metrics for bandwidth / occupancy checks -ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\ -dram__throughput.avg.pct_of_peak_sustained_elapsed \ - python scripts/bench_diffusion_rmsnorm.py - -# Key metrics for A100 diffusion kernels: -# - Achieved occupancy (sm__warps_active.avg.pct_of_peak_sustained_active) -# - Memory throughput (dram__throughput.avg.pct_of_peak_sustained_elapsed) -# → Target: 30–40% of 2.0 TB/s (600–800 GB/s) for vectorized kernels -# - Compute throughput (sm__throughput.avg.pct_of_peak_sustained_elapsed) -# - Warp stall reasons (smsp__warp_issue_stalled_*.avg.pct_of_peak_sustained_active) -# - Kernel time (gpu__time_duration.avg) -``` - -### Common A100 Performance Issues - -1. **Memory bound below target**: `dram__throughput` < 30% - - Fix: Use `AlignedVector` (128-bit vector loads) - -2. **Low occupancy**: Grid too small for 108 SMs - - Fix: Use `runtime::get_sm_count()` persistent kernel pattern - -3. **No TF32 for FP32 GEMMs**: torch.backends.cuda.matmul.allow_tf32 not set - - Fix: `torch.backends.cuda.matmul.allow_tf32 = True` - ---- - -## Best Practices Summary (A100) - -1. **Bandwidth**: Even more critical than H100 — profile with `ncu` first -2. **Vectorization**: `AlignedVector` gives ~3x over scalar -3. **TF32**: Enable for any FP32 matmul workload -4. **Shared memory**: Cap at 164 KB; use `cudaFuncSetAttribute` -5. **Grid sizing**: Multiples of 108 SMs via `runtime::get_sm_count` -6. **cp.async**: Use for prefetching in tiled kernels -7. **Multi-arch**: Build for both `sm_80` and `sm_90` to support both GPUs -8. **Same abstractions**: `AlignedVector`, `TensorMatcher`, `LaunchKernel` work identically - -## Reference Benchmark Results (A100 80GB, BF16) - -| Kernel | Shape | A100 (ms) | H100 (ms) | H100 Speedup | -|--------|-------|-----------|-----------|--------------| -| RMSNorm | [2, 1024, 2048] | ~0.08 | 0.054 | 1.5x | -| GEGLU | [2, 1024, 4096] | ~0.05 | 0.030 | 1.7x | diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/h100-optimization-guide.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/h100-optimization-guide.md deleted file mode 100644 index a47d082e4..000000000 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/h100-optimization-guide.md +++ /dev/null @@ -1,364 +0,0 @@ -# H100 GPU Optimization Guide — SGLang Diffusion JIT Kernels - -Deep dive into H100-specific optimizations for diffusion model CUDA kernels, written for SGLang's JIT kernel system. - -> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) - ---- - -## H100 Hopper Architecture Overview - -| Component | Specification | Optimization Implication | -|-----------|---------------|--------------------------| -| Compute Capability | sm_90 | Use `extra_cuda_cflags=["-arch=sm_90"]` in `load_jit` | -| SMs | 132 | Grid: aim for multiples of 132 | -| Shared Memory | 192 KB/SM | Configurable: 96/144/192 KB | -| L2 Cache | 50 MB | Tile K,V of attention to fit in L2 | -| Memory Bandwidth | 3.35 TB/s | BF16 vectorized: achieves ~38% (~1.27 TB/s) | -| Max Threads/SM | 2048 | Max 16 blocks of 128 threads per SM | -| Warp Size | 32 | All reductions use `warp::reduce_sum` | -| Registers | 64K 32-bit/SM | 255 per thread max | - -### New Hopper Features (sm_90+) - -1. **Thread Block Clusters** — groups cooperating via Distributed Shared Memory -2. **TMA (Tensor Memory Accelerator)** — hardware-accelerated bulk copies -3. **FP8 support** — native 8-bit floating point in tensor cores -4. **PDL (Programmatic Dependent Launch)** — enable with `.enable_pdl(true)` in `LaunchKernel` - -Gate sm_90+ features with a runtime check before calling `load_jit`: - -```python -if torch.cuda.get_device_capability()[0] < 9: - raise RuntimeError("This kernel requires H100 (sm_90+)") -``` - ---- - -## Memory Hierarchy Optimization - -### Coalesced Global Memory Access - -```cpp -// GOOD: threads read consecutive addresses → 128-byte transaction per warp -uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; -fp16_t val = src[idx]; - -// BAD: strided access → multiple transactions, lower effective bandwidth -uint32_t idx = threadIdx.x * stride; // avoid stride > 1 -``` - -**Transaction sizes**: 32 bytes minimum, 128 bytes optimal (full warp, FP32). - -### Vectorized Memory Access with `AlignedVector` - -SGLang's `AlignedVector` provides 128-bit (16-byte) vector loads. Always use this instead of raw pointer reinterprets. - -```cpp -#include - -// 16 bytes per load: 8×bf16_t, 8×fp16_t, or 4×fp32_t -constexpr int kVecN = 16 / sizeof(T); -using vec_t = device::AlignedVector; - -// Load -vec_t v; -v.load(src, vi); // loads src[vi * kVecN .. vi * kVecN + kVecN - 1] - -// Process -#pragma unroll -for (int i = 0; i < kVecN; ++i) { - float val = static_cast(v[i]); - // ... compute ... - v[i] = static_cast(result); -} - -// Store -v.store(dst, vi); -``` - -**RMSNorm benchmark (H100 80GB, BF16):** - -| Implementation | Time (ms) | Speedup | -|:---|:---:|:---:| -| Scalar loads | 0.065 | 1.00x | -| `AlignedVector` | 0.019 | **3.37x** | - -Bandwidth achieved: **~38% of 3.35 TB/s** = 1.27 TB/s. - -### L2 Cache Utilization (50 MB) - -For attention, tile K and V so they stay in L2 while Q iterates: - -``` -BLOCK_SIZE_M = 128 (Q block) -BLOCK_SIZE_N = 64 (K,V block) -With head_dim=64: tile = 128×64×2 = 16 KB (FP16), multiple tiles fit in L2 -``` - -### Shared Memory Configuration - -Request max shared memory for attention kernels: - -```cpp -// In launcher (after selecting kernel function pointer): -cudaFuncSetAttribute( - your_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - 192 * 1024 // 192 KB max on H100 -); -``` - -Shared memory has 32 banks (4 bytes/bank). Avoid conflicts with padding: - -```cpp -__shared__ float data[32][33]; // 33 instead of 32 → no bank conflict -``` - ---- - -## Warp & CTA Reductions (SGLang Abstractions) - -Use `sgl_kernel/warp.cuh` and `sgl_kernel/cta.cuh` — never raw `__shfl_xor_sync`. - -```cpp -#include -#include - -// Warp-level sum (uses __shfl_xor_sync internally) -float result = device::warp::reduce_sum(partial); - -// Warp-level max -float mx = device::warp::reduce_max(val); - -// CTA-wide max via shared memory -__shared__ float smem[32]; -device::cta::reduce_max(val, smem, -1e38f); -// smem[0] holds the result after __syncthreads() -``` - -**Block reduction pattern for RMSNorm:** - -```cpp -// 1. Warp reduction -sum_sq = device::warp::reduce_sum(sum_sq); - -// 2. Write warp leaders to smem -__shared__ float smem_r[32]; -if (threadIdx.x % 32 == 0) smem_r[threadIdx.x / 32] = sum_sq; -__syncthreads(); - -// 3. Final warp reduction over warp leaders -if (threadIdx.x < 32) { - sum_sq = (threadIdx.x < blockDim.x / 32) ? smem_r[threadIdx.x] : 0.f; - sum_sq = device::warp::reduce_sum(sum_sq); -} -__syncthreads(); -``` - ---- - -## Occupancy Tuning - -``` -Occupancy = Active Warps per SM / Max Warps per SM (64) - -Limiting factors on H100: - 1. Registers: 65536 / (threads_per_block × regs_per_thread) - 2. Shared Memory: 192 KB / smem_per_block - 3. Threads: 2048 / threads_per_block -``` - -**Recommended block sizes:** - -| Kernel Type | Threads/Block | Warps | Reasoning | -|-------------|---------------|-------|-----------| -| Element-wise (RoPE, GEGLU) | 256 | 8 | High occupancy, simple | -| Row reduction (RMSNorm, LayerNorm) | 256–512 | 8–16 | Enough threads for full reduction | -| Tiled (attention) | 256 | 8 | Balance shared mem and registers | - -**Persistent kernel pattern** (cap grid to SM × occupancy): - -```cpp -#include - -static const uint32_t max_occ = host::runtime::get_blocks_per_sm(kernel, kBlockSize); -static const uint32_t num_sm = host::runtime::get_sm_count(device.unwrap().device_id); -const uint32_t num_blocks = std::min(num_sm * max_occ, host::div_ceil(n, kBlockSize)); -host::LaunchKernel(num_blocks, kBlockSize, device.unwrap())(kernel, params); -``` - ---- - -## Precision and Numerical Stability - -| Type | Exponent Bits | Mantissa Bits | Range | Use Case | -|------|--------------|---------------|-------|----------| -| FP16 | 5 | 10 | ±65504 | Inference; attention score overflow risk | -| BF16 | 8 | 7 | ±3.39×10³⁸ | Training/inference preferred; safer for attn | -| FP32 | 8 | 23 | ±3.39×10³⁸ | Accumulation only | - -**Mixed precision pattern** (always accumulate in FP32): - -```cpp -// Input via AlignedVector -vec_t v; -v.load(src, vi); -float acc = 0.f; -#pragma unroll -for (int i = 0; i < kVecN; ++i) { - float val = static_cast(v[i]); // promote to FP32 - acc += val * val; -} -// Output -v[i] = static_cast(fp32_result); // demote back -``` - ---- - -## Diffusion-Specific Patterns - -### DiT Block Operators - -| Operator | Pattern | Key Constraint | -|----------|---------|----------------| -| **RMSNorm** | 2-pass row reduction | weight may be `None` | -| **AdaLN** | `norm(x) * (1 + scale) + shift` | fuse norm+scale+shift | -| **RoPE 3D** | `[B, t*h*w, heads, head_dim]` | layout: `seq = t*h*w` | -| **GEGLU** | `gelu(gate) * value`, input `[B,L,2H]` | don't use for LTX-Video (uses GELU) | -| **SiLU gate** | `x * sigmoid(x)` | fuse with MLP linear | - -### Online Softmax (for custom attention) - -```cuda -// Numerically stable without materializing full [seq×seq] score matrix -float row_max = -INFINITY, row_sum = 0.f; -for each K block: - compute local_scores - new_max = max(row_max, max(local_scores)) - rescale = exp(row_max - new_max) - row_sum = row_sum * rescale + sum(exp(local_scores - new_max)) - out_acc = out_acc * rescale + softmax(local_scores) @ V_block - row_max = new_max -``` - ---- - -## Profiling and Debugging - -### NVIDIA Nsight Systems (nsys) - -System-wide profiling to see kernel durations, memory transfers, and GPU idle time: - -```bash -nsys profile -o profile_report python scripts/bench_diffusion_rmsnorm.py - -# Key metrics to watch: -# - Kernel duration -# - Memory transfer time -# - GPU idle time -# - Stream utilization -``` - -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) - -Detailed per-kernel analysis for tuning individual JIT CUDA kernels: - -```bash -# Full metrics — use when you need everything (slow) -ncu --set full -o metrics.ncu-rep \ - python scripts/bench_diffusion_rmsnorm.py - -# Specific metrics — use for targeted bandwidth / occupancy checks -ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\ -dram__throughput.avg.pct_of_peak_sustained_elapsed \ - python scripts/bench_diffusion_rmsnorm.py - -# Key metrics for diffusion JIT kernels: -# - Achieved occupancy (sm__warps_active.avg.pct_of_peak_sustained_active) -# - Memory throughput (dram__throughput.avg.pct_of_peak_sustained_elapsed) -# - Compute throughput (sm__throughput.avg.pct_of_peak_sustained_elapsed) -# - Warp stall reasons (smsp__warp_issue_stalled_*.avg.pct_of_peak_sustained_active) -# - L1 cache hit rate (l1tex__t_requests_pipe_lsu_mem_global_op_ld.sum) -``` - -### Common Performance Issues - -1. **Low occupancy**: Too many registers or shared memory per block - - Check: `--ptxas-options=-v` in `extra_cuda_cflags` to see register count - - Fix: Reduce `--maxrregcount=N`; use smaller block size - -2. **Memory bound, low bandwidth**: Achieved < 30% of 3.35 TB/s - - Check: `dram__throughput.avg.pct_of_peak_sustained_elapsed` - - Fix: Switch to `AlignedVector` for 128-bit vector loads - -3. **Shared memory bank conflicts**: `l1tex__data_bank_conflicts_pipe_lmem_op_st.sum` is high - - Fix: Add padding — `__shared__ float data[32][33]` - -4. **Warp divergence**: Conditional branches splitting warps - - Check: `smsp__warp_issue_stalled_branch.avg.pct_of_peak_sustained_active` - - Fix: Restructure so elements with identical branches are in the same warp - -5. **Too many small kernels**: High kernel launch overhead - - Fix: Fuse operations (e.g., norm + scale + shift → AdaLN in one kernel) - ---- - -## JIT Compilation Notes - -SGLang's JIT compiles kernels on first use via `load_jit`. For H100-specific flags: - -```python -return load_jit( - "my_kernel", - *args, - cuda_files=["diffusion/my_kernel.cuh"], - cuda_wrappers=[("my_kernel", f"my_kernel<{args}>")], - extra_cuda_cflags=[ - "-O3", - "--use_fast_math", - "-arch=sm_90", # H100 only; omit for multi-arch - "--ptxas-options=-v", # Remove after tuning - ], -) -``` - -For multi-arch (H100 + A100): - -```python -extra_cuda_cflags=[ - "-O3", - "--use_fast_math", - "-gencode=arch=compute_80,code=sm_80", # A100 - "-gencode=arch=compute_90,code=sm_90", # H100 -] -``` - ---- - -## Best Practices Summary - -1. **Memory access**: Coalesce writes, align to 128-byte boundaries -2. **Vectorization**: Use `AlignedVector` for all element-wise loads/stores -3. **Reductions**: Use `warp::reduce_sum/max`, then shared memory pattern above -4. **Precision**: BF16 for I/O, FP32 for accumulation; use `static_cast` -5. **Block size**: 256 threads default; 512 for reductions; tune with `runtime::get_blocks_per_sm` -6. **Grid sizing**: Multiples of 132 SMs; use persistent kernel pattern for small N -7. **Shared memory**: Add padding (`[32][33]`) to avoid bank conflicts -8. **Profile**: Run `ncu` before claiming a speedup; check dram throughput % -9. **Fuse**: Combine norm + scale + shift into a single pass to reduce memory traffic -10. **Abstractions**: Always use `TensorMatcher`, `AlignedVector`, `LaunchKernel` — never raw CUDA - -## Reference Benchmark Results (H100 80GB, BF16) - -| Kernel | Shape | Time (ms) | -|--------|-------|-----------| -| RMSNorm | [2, 1024, 2048] | 0.054 | -| GEGLU | [2, 1024, 4096] → [2, 1024, 2048] | 0.030 | -| RoPE 3D | [2, 480, 8, 64] | 1.670 | -| RMSNorm vectorized | [1, 1024, 2048] | 0.019 | -| RMSNorm vectorized | [4, 4096, 3072] | 0.157 | - -> See `kernel-templates.md` for copy-paste ready sglang JIT kernel implementations. diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/kernel-templates.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/kernel-templates.md deleted file mode 100644 index dec2acf52..000000000 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/kernel-templates.md +++ /dev/null @@ -1,570 +0,0 @@ -# CUDA Kernel Templates — SGLang Diffusion JIT Style - -Copy-paste ready templates for JIT CUDA kernels in `python/sglang/jit_kernel/csrc/diffusion/`. -All templates use SGLang's internal abstractions; no raw CUDA headers needed. - -> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) - ---- - -## Prerequisite: Standard Includes - -Every kernel file in `csrc/diffusion/` starts with: - -```cpp -#include // TensorMatcher, SymbolicSize, SymbolicDevice -#include // fp16_t, bf16_t, fp32_t, dtype_trait, packed_t -#include // RuntimeCheck, Panic, div_ceil -#include // LaunchKernel, SGL_DEVICE, type aliases -#include // AlignedVector -#include // warp::reduce_sum, warp::reduce_max -#include // device::math::rsqrt, sqrt, ... -#include // tile::Memory (strided access pattern) - -#include -#include -``` - -**Key type aliases** (from `utils.cuh`): -- `fp16_t` = `__half`, `fp16x2_t` = `__half2` -- `bf16_t` = `__nv_bfloat16`, `bf16x2_t` = `__nv_bfloat162` -- `fp32_t` = `float`, `fp32x2_t` = `float2` -- `SGL_DEVICE` = `__forceinline__ __device__` - ---- - -## Template 1: Element-wise Operation - -Use for ops that process elements independently: RoPE, SiLU, GEGLU, scale+bias. - -### `.cuh` file: `csrc/diffusion/silu_gate.cuh` - -```cpp -#include -#include -#include -#include -#include -#include - -#include -#include - -namespace { - -// SiLU gate: out[i] = x[i] * sigmoid(x[i]) -// Input layout: [B, L, hidden] -template -__global__ void silu_gate_kernel( - T* __restrict__ dst, - const T* __restrict__ src, - uint32_t n_vecs, - uint32_t n_remainder, - uint32_t n_total) -{ - using vec_t = device::AlignedVector; - - const uint32_t stride = blockDim.x * gridDim.x; - - // --- vectorized body --- - for (uint32_t vi = blockIdx.x * blockDim.x + threadIdx.x; vi < n_vecs; vi += stride) { - vec_t v; - v.load(src, vi); - #pragma unroll - for (int i = 0; i < kVecN; ++i) { - float val = static_cast(v[i]); - float sig = 1.f / (1.f + device::math::exp(-val)); - v[i] = static_cast(val * sig); - } - v.store(dst, vi); - } - - // --- scalar tail (for sizes not divisible by kVecN) --- - const uint32_t base = n_vecs * kVecN; - for (uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n_remainder; i += stride) { - float val = static_cast(src[base + i]); - float sig = 1.f / (1.f + device::math::exp(-val)); - dst[base + i] = static_cast(val * sig); - } -} - -template -void silu_gate(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { - using namespace host; - - SymbolicSize N{"num_elements"}; - SymbolicDevice device; - device.set_options(); - - TensorMatcher({N}) - .with_dtype() - .with_device(device) - .verify(dst) - .verify(src); - - const uint32_t n = static_cast(N.unwrap()); - const DLDevice dev = device.unwrap(); - RuntimeCheck(n > 0, "silu_gate: num_elements must be > 0"); - - constexpr int kVecN = 16 / sizeof(T); // 128-bit vector load - const uint32_t n_vecs = n / kVecN; - const uint32_t n_rem = n % kVecN; - - constexpr uint32_t kBlock = 256; - const uint32_t grid = div_ceil(std::max(n_vecs, n_rem), kBlock); - - LaunchKernel(grid, kBlock, dev)( - silu_gate_kernel, - static_cast(dst.data_ptr()), - static_cast(src.data_ptr()), - n_vecs, n_rem, n); -} - -} // namespace -``` - -### Python wrapper: `diffusion/silu_gate.py` - -```python -from __future__ import annotations -import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args - -@cache_once -def _jit_silu_gate_module(dtype: torch.dtype): - args = make_cpp_args(dtype) - return load_jit( - "diffusion_silu_gate", - *args, - cuda_files=["diffusion/silu_gate.cuh"], - cuda_wrappers=[("silu_gate", f"silu_gate<{args}>")], - extra_cuda_cflags=["-O3", "--use_fast_math"], - ) - -def diffusion_silu_gate(src: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor: - assert src.is_cuda and src.dtype in (torch.float16, torch.bfloat16, torch.float32) - if out is None: - out = torch.empty_like(src) - module = _jit_silu_gate_module(src.dtype) - module.silu_gate(out, src) - return out -``` - ---- - -## Template 2: Row-wise Reduction (RMSNorm / LayerNorm) - -Use for ops that reduce across the last dimension of each row. - -### `.cuh` file: `csrc/diffusion/rmsnorm.cuh` - -```cpp -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -namespace { - -// RMSNorm: y = x / rms(x) * weight -// One block per row; vectorized loads/stores; warp + shared-mem reduction -template -__global__ void rmsnorm_kernel( - T* __restrict__ dst, - const T* __restrict__ src, - const T* __restrict__ weight, // nullptr if no affine weight - uint32_t hidden, - uint32_t n_vecs, - float eps) -{ - using vec_t = device::AlignedVector; - - const uint32_t row = blockIdx.x; - const T* row_src = src + row * hidden; - T* row_dst = dst + row * hidden; - - // Pass 1: sum of squares - float sum_sq = 0.f; - for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { - vec_t v; - v.load(row_src, vi); - #pragma unroll - for (int i = 0; i < kVecN; ++i) { - float val = static_cast(v[i]); - sum_sq += val * val; - } - } - - // Warp + block reduction - sum_sq = device::warp::reduce_sum(sum_sq); - __shared__ float smem[32]; - if (threadIdx.x % 32 == 0) smem[threadIdx.x / 32] = sum_sq; - __syncthreads(); - if (threadIdx.x < 32) { - sum_sq = (threadIdx.x < blockDim.x / 32) ? smem[threadIdx.x] : 0.f; - sum_sq = device::warp::reduce_sum(sum_sq); - } - __syncthreads(); - - const float rms_inv = device::math::rsqrt(sum_sq / static_cast(hidden) + eps); - - // Pass 2: normalize + optional weight - for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { - vec_t v_in, v_out; - v_in.load(row_src, vi); - if (weight != nullptr) { - vec_t v_w; - v_w.load(weight, vi); - #pragma unroll - for (int i = 0; i < kVecN; ++i) - v_out[i] = static_cast(static_cast(v_in[i]) * rms_inv - * static_cast(v_w[i])); - } else { - #pragma unroll - for (int i = 0; i < kVecN; ++i) - v_out[i] = static_cast(static_cast(v_in[i]) * rms_inv); - } - v_out.store(row_dst, vi); - } -} - -template -void rmsnorm( - tvm::ffi::TensorView dst, - tvm::ffi::TensorView src, - tvm::ffi::TensorView weight, // data_ptr == nullptr → no weight - float eps) -{ - using namespace host; - - SymbolicSize B{"batch_tokens"}, H{"hidden_size"}; - SymbolicDevice device; - device.set_options(); - - TensorMatcher({B, H}) - .with_dtype() - .with_device(device) - .verify(dst) - .verify(src); - - const uint32_t num_rows = static_cast(B.unwrap()); - const uint32_t hidden = static_cast(H.unwrap()); - const DLDevice dev = device.unwrap(); - - constexpr int kVecN = 16 / sizeof(T); - RuntimeCheck(hidden % kVecN == 0, - "rmsnorm: hidden_size (", hidden, ") must be divisible by ", kVecN); - const uint32_t n_vecs = hidden / kVecN; - - uint32_t threads = std::min(n_vecs, 512u); - threads = (threads + 31) / 32 * 32; - - const T* w_ptr = (weight.data_ptr() != nullptr) - ? static_cast(weight.data_ptr()) : nullptr; - - LaunchKernel(num_rows, threads, dev)( - rmsnorm_kernel, - static_cast(dst.data_ptr()), - static_cast(src.data_ptr()), - w_ptr, hidden, n_vecs, eps); -} - -} // namespace -``` - ---- - -## Template 3: Fused Row-Reduction + Element-wise (AdaLN) - -Combines RMSNorm + AdaLN modulation into one pass: `y = norm(x) * (1 + scale) + shift`. - -### `.cuh` file: `csrc/diffusion/adaln.cuh` - -```cpp -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -namespace { - -// AdaLN: y = norm(x) * (1 + scale) + shift -// scale, shift: [batch, hidden] (one per row) -template -__global__ void adaln_kernel( - T* __restrict__ dst, - const T* __restrict__ src, - const T* __restrict__ weight, - const T* __restrict__ scale, - const T* __restrict__ shift, - uint32_t hidden, - uint32_t n_vecs, - float eps) -{ - using vec_t = device::AlignedVector; - - const uint32_t row = blockIdx.x; - const T* row_src = src + row * hidden; - const T* row_scale = scale + row * hidden; - const T* row_shift = shift + row * hidden; - T* row_dst = dst + row * hidden; - - // Pass 1: compute RMS - float sum_sq = 0.f; - for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { - vec_t v; - v.load(row_src, vi); - #pragma unroll - for (int i = 0; i < kVecN; ++i) { - float val = static_cast(v[i]); - sum_sq += val * val; - } - } - sum_sq = device::warp::reduce_sum(sum_sq); - __shared__ float smem[32]; - if (threadIdx.x % 32 == 0) smem[threadIdx.x / 32] = sum_sq; - __syncthreads(); - if (threadIdx.x < 32) { - sum_sq = (threadIdx.x < blockDim.x / 32) ? smem[threadIdx.x] : 0.f; - sum_sq = device::warp::reduce_sum(sum_sq); - } - __syncthreads(); - const float rms_inv = device::math::rsqrt(sum_sq / static_cast(hidden) + eps); - - // Pass 2: normalize + modulate - for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { - vec_t v_in, v_w, v_sc, v_sh, v_out; - v_in.load(row_src, vi); - v_w.load(weight, vi); - v_sc.load(row_scale, vi); - v_sh.load(row_shift, vi); - #pragma unroll - for (int i = 0; i < kVecN; ++i) { - float x = static_cast(v_in[i]) * rms_inv * static_cast(v_w[i]); - float sc = static_cast(v_sc[i]); - float sh = static_cast(v_sh[i]); - v_out[i] = static_cast(x * (1.f + sc) + sh); - } - v_out.store(row_dst, vi); - } -} - -template -void adaln( - tvm::ffi::TensorView dst, - tvm::ffi::TensorView src, - tvm::ffi::TensorView weight, - tvm::ffi::TensorView scale, - tvm::ffi::TensorView shift, - float eps) -{ - using namespace host; - - SymbolicSize B{"batch_tokens"}, H{"hidden_size"}; - SymbolicDevice device; - device.set_options(); - - TensorMatcher({B, H}) - .with_dtype() - .with_device(device) - .verify(dst).verify(src).verify(weight).verify(scale).verify(shift); - - const uint32_t num_rows = static_cast(B.unwrap()); - const uint32_t hidden = static_cast(H.unwrap()); - const DLDevice dev = device.unwrap(); - - constexpr int kVecN = 16 / sizeof(T); - RuntimeCheck(hidden % kVecN == 0, "adaln: hidden_size must be divisible by ", kVecN); - const uint32_t n_vecs = hidden / kVecN; - - uint32_t threads = std::min(n_vecs, 512u); - threads = (threads + 31) / 32 * 32; - - LaunchKernel(num_rows, threads, dev)( - adaln_kernel, - static_cast(dst.data_ptr()), - static_cast(src.data_ptr()), - static_cast(weight.data_ptr()), - static_cast(scale.data_ptr()), - static_cast(shift.data_ptr()), - hidden, n_vecs, eps); -} - -} // namespace -``` - ---- - -## Template 4: Python Wrapper (generic pattern) - -File location: `python/sglang/jit_kernel/diffusion/.py` - -```python -from __future__ import annotations -from typing import TYPE_CHECKING - -import torch - -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args - -if TYPE_CHECKING: - from tvm_ffi.module import Module - - -@cache_once -def _jit_module(dtype: torch.dtype) -> Module: - """Cache key: dtype (and any other template params you need).""" - args = make_cpp_args(dtype) - return load_jit( - "diffusion_your_op", # unique build cache key - *args, - cuda_files=["diffusion/your_op.cuh"], # relative to csrc/ - cuda_wrappers=[("your_op", f"your_op<{args}>")], - extra_cuda_cflags=["-O3", "--use_fast_math"], - ) - - -def diffusion_your_op( - src: torch.Tensor, - out: torch.Tensor | None = None, -) -> torch.Tensor: - """ - Your op description. - - Supported dtypes: float16, bfloat16, float32. - """ - assert src.is_cuda, "src must be a CUDA tensor" - assert src.dtype in (torch.float16, torch.bfloat16, torch.float32), ( - f"Unsupported dtype {src.dtype}" - ) - if out is None: - out = torch.empty_like(src) - - module = _jit_module(src.dtype) - module.your_op(out, src) - return out -``` - -**`make_cpp_args` conversion table:** - -| `torch.dtype` | C++ type | -|---------------|----------| -| `torch.float16` | `fp16_t` | -| `torch.bfloat16` | `bf16_t` | -| `torch.float32` | `fp32_t` | - ---- - -## Template 5: Correctness Test - -```python -# python/sglang/jit_kernel/tests/test_diffusion_.py -import pytest -import torch -from sglang.jit_kernel.diffusion. import diffusion_ - - -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("shape", [(1, 2048), (4, 3072), (16, 4096)]) -def test__correctness(dtype, shape): - src = torch.randn(*shape, dtype=dtype, device="cuda") - - out_jit = diffusion_(src) - ref = reference_(src.float()).to(dtype) # reference in fp32 - - tol = {"rtol": 1e-2, "atol": 1e-2} if dtype != torch.float32 else {"rtol": 1e-5, "atol": 1e-6} - torch.testing.assert_close(out_jit, ref, **tol) - - -def test__out_param(): - src = torch.randn(1024, 2048, dtype=torch.bfloat16, device="cuda") - out = torch.empty_like(src) - result = diffusion_(src, out=out) - assert result is out - - -def test__cpu_error(): - src = torch.randn(128, dtype=torch.float16) # CPU tensor - with pytest.raises(AssertionError): - diffusion_(src) - - -if __name__ == "__main__": - import sys - sys.exit(pytest.main([__file__, "-v", "-s"])) -``` - ---- - -## Template 6: Benchmark - -```python -# python/sglang/jit_kernel/benchmark/bench_diffusion_.py -import torch -import triton.testing - -from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE, run_benchmark -from sglang.jit_kernel.diffusion. import diffusion_ - -SHAPES = [(4096, 2048), (4096, 3072), (4096, 4096)] - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["hidden"], - x_vals=[s[1] for s in SHAPES], - line_arg="provider", - line_vals=["jit_cuda", "torch"], - line_names=["SGLang JIT CUDA", "PyTorch"], - styles=[("blue", "-"), ("red", "--")], - ylabel="us", - plot_name="diffusion-", - args={}, - ) -) -def benchmark(hidden: int, provider: str): - src = torch.randn(4096, hidden, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) - - if provider == "jit_cuda": - fn = lambda: diffusion_(src) - else: - fn = lambda: reference_(src) # torch baseline - - return run_benchmark(fn) - - -if __name__ == "__main__": - benchmark.run(print_data=True) -``` - ---- - -## Summary of New Files per Kernel - -``` -python/sglang/jit_kernel/csrc/diffusion/ -└── .cuh # CUDA kernel + launcher - -python/sglang/jit_kernel/diffusion/ -└── .py # Python wrapper (load_jit + cache_once) - -python/sglang/jit_kernel/tests/ -└── test_diffusion_.py # correctness tests - -python/sglang/jit_kernel/benchmark/ -└── bench_diffusion_.py # triton.testing benchmark -``` - -> See `scripts/bench_diffusion_rmsnorm.py` and `scripts/bench_diffusion_denoise.py` for full runnable examples. diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/t4-optimization-guide.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/t4-optimization-guide.md deleted file mode 100644 index d080a6523..000000000 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/t4-optimization-guide.md +++ /dev/null @@ -1,340 +0,0 @@ -# T4 GPU Optimization Guide — SGLang Diffusion JIT Kernels - -T4 is a Turing architecture GPU (GCP n1+T4, AWS g4dn) commonly used for cloud inference. -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. - ---- - -## T4 Turing Architecture Overview - -| Component | T4 | A100 | H100 | -|-----------|-----|------|------| -| Compute Capability | sm_75 | sm_80 | sm_90 | -| SMs | 40 | 108 | 132 | -| Shared Memory/SM | **64 KB** | 164 KB | 192 KB | -| L2 Cache | 4 MB | 40 MB | 50 MB | -| Memory Bandwidth | **320 GB/s** | 2.0 TB/s | 3.35 TB/s | -| Memory | 16 GB GDDR6 | 40–80 GB HBM2e | 80 GB HBM3 | -| Max Threads/SM | **1024** | 2048 | 2048 | -| BF16 Support | **No** | Yes | Yes | - -### Critical T4 Constraints - -1. **No BFloat16** — must use FP16 everywhere -2. **320 GB/s bandwidth** — ~10x lower than H100; vectorization is critical -3. **16 GB memory** — limits model size; use offloading -4. **64 KB shared memory/SM** — smaller attention tiles -5. **Max 1024 threads/SM** — half of A100/H100; affects occupancy calculations - ---- - -## No BF16: Always Use FP16 - -This is the most impactful constraint. **Never use `bf16_t` or `__nv_bfloat16` on T4.** - -**Python wrapper guard:** - -```python -import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args - -@cache_once -def _jit_rmsnorm_module(dtype: torch.dtype): - # T4 (sm_75) does not support BF16 - cap = torch.cuda.get_device_capability() - if cap < (8, 0) and dtype == torch.bfloat16: - raise RuntimeError( - f"T4 (sm_75) does not support BF16. Use torch.float16 instead. " - f"Got dtype={dtype}" - ) - args = make_cpp_args(dtype) - return load_jit( - "diffusion_rmsnorm", - *args, - cuda_files=["diffusion/rmsnorm.cuh"], - cuda_wrappers=[("rmsnorm", f"rmsnorm<{args}>")], - ) -``` - -**Conditional type in kernel:** - -```cuda -#if __CUDA_ARCH__ >= 800 - // A100/H100: BF16 available - using DefaultHalf = bf16_t; -#else - // T4/Turing: FP16 only - using DefaultHalf = fp16_t; -#endif -``` - -**Runtime detection helper:** - -```python -def get_diffusion_dtype() -> torch.dtype: - """Return the appropriate half-precision dtype for the current GPU.""" - cap = torch.cuda.get_device_capability() - if cap >= (8, 0): - return torch.bfloat16 # A100/H100: prefer BF16 - else: - return torch.float16 # T4/older: FP16 only -``` - ---- - -## Memory Access Optimization - -With only 320 GB/s, **vectorization is more critical on T4 than on A100/H100**. - -### `AlignedVector` (same abstraction, FP16 only) - -```cpp -#include - -// On T4, T must be fp16_t or fp32_t (NOT bf16_t) -constexpr int kVecN = 16 / sizeof(T); // 8 for fp16, 4 for fp32 -using vec_t = device::AlignedVector; -``` - -**Target bandwidth**: 40–50% of T4's 320 GB/s = 128–160 GB/s. - -### Increase Arithmetic Intensity - -With low bandwidth, fusing ops saves more on T4 than on H100: - -```cpp -// BAD on T4: separate passes → 2× memory traffic -output1[i] = input[i] * scale; // pass 1 -output2[i] = output1[i] + bias; // pass 2 - -// GOOD: fuse → single memory read, single write -float val = static_cast(v[i]); -val = val * scale + bias; -val = device::math::max(val, 0.f); // ReLU -v[i] = static_cast(val); -``` - -### Expected T4 Performance - -| Kernel | T4 (ms) | A100 (ms) | H100 (ms) | T4 vs H100 | -|--------|---------|-----------|-----------|------------| -| RMSNorm [2, 1024, 2048] | ~0.5 | ~0.08 | 0.054 | ~9x slower | -| GEGLU [2, 1024, 4096] | ~0.3 | ~0.05 | 0.030 | ~10x slower | - ---- - -## Shared Memory Configuration - -T4 max: **64 KB/SM**. Use smaller tiles vs A100/H100. - -```cpp -// T4: request max shared memory (64 KB) -cudaFuncSetAttribute( - your_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - 64 * 1024 -); -``` - -**Attention tile sizes for T4** (halved vs H100): - -``` -H100/A100: BLOCK_SIZE_M = 128, BLOCK_SIZE_N = 64 -T4: BLOCK_SIZE_M = 64, BLOCK_SIZE_N = 32 ← reduced for 64 KB limit -``` - ---- - -## Occupancy Tuning - -T4 max: **1024 threads/SM** (vs 2048 on A100/H100). This halves max occupancy for a given block size. - -**Block sizes for T4:** - -| Kernel Type | Threads/Block | Notes | -|-------------|---------------|-------| -| Element-wise | 256 | Same as H100 | -| Row reduction | 256–512 | Avoid > 512 to fit multiple blocks/SM | -| Tiled/attention | 128–256 | Small tiles due to 64 KB shared mem | - -**Grid sizing for T4 (40 SMs)** — `runtime::get_sm_count` handles this automatically: - -```cpp -// get_sm_count() returns 40 on T4, 108 on A100, 132 on H100 -const uint32_t num_sm = host::runtime::get_sm_count(device.unwrap().device_id); -``` - ---- - -## Numerical Stability with FP16 - -FP16 has a smaller dynamic range (±65504) vs BF16 (±3.39×10³⁸). Watch for overflow in attention: - -```cuda -// Scale attention scores to prevent FP16 overflow -float scale_factor = 1.0f / sqrtf(static_cast(head_dim)); -// For very long sequences on T4, may need additional scaling: -// if (score * scale_factor > 65000.f) { /* clamp */ } -``` - -Always accumulate in FP32: - -```cpp -float acc = 0.f; // FP32 accumulation -for (uint32_t vi = threadIdx.x; vi < n_vecs; vi += blockDim.x) { - vec_t v; - v.load(src, vi); - #pragma unroll - for (int i = 0; i < kVecN; ++i) { - float val = static_cast(v[i]); // fp16 → fp32 - acc += val * val; - } -} -``` - ---- - -## Memory Management for 16 GB - -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= - -# Enable CPU offloading to fit within 16 GB -sglang generate \ - --model-path=black-forest-labs/FLUX.1-dev \ - --dit-cpu-offload true \ # DiT weights to CPU - --text-encoder-cpu-offload true \ - --vae-cpu-offload true \ - --width=512 --height=512 \ # Reduce resolution - --num-inference-steps=20 \ # Fewer steps - --seed=42 -``` - -**Resolution recommendations for T4:** - -| Model | H100/A100 | T4 | -|-------|-----------|-----| -| FLUX.1-dev | 1024×1024 | 512×512 | -| Wan2.2-TI2V-5B | 720P | 480P | -| FLUX.2-dev | 1024×1024 | 512×512 | - ---- - -## JIT Compilation for T4 - -```python -return load_jit( - "my_kernel", - *args, - cuda_files=["diffusion/my_kernel.cuh"], - cuda_wrappers=[("my_kernel", f"my_kernel<{args}>")], - extra_cuda_cflags=[ - "-O3", - "--use_fast_math", - "-arch=sm_75", # T4 only; omit for multi-arch - ], -) -``` - -**Multi-arch (T4 + A100 + H100):** - -```python -extra_cuda_cflags=[ - "-O3", "--use_fast_math", - "-gencode=arch=compute_75,code=sm_75", # T4 - "-gencode=arch=compute_80,code=sm_80", # A100 - "-gencode=arch=compute_90,code=sm_90", # H100 -] -``` - ---- - -## H100/A100 → T4 Migration Checklist - -| Item | H100/A100 | T4 | Action | -|------|-----------|-----|--------| -| BF16 | Available | **Not available** | Replace `bf16_t` with `fp16_t`; guard in Python wrapper | -| Shared memory | 164–192 KB | **64 KB** | Halve tile sizes | -| Grid sizing | ×108/132 SMs | ×40 SMs | `get_sm_count()` auto-handles | -| Max threads/SM | 2048 | **1024** | Don't exceed 512 threads/block | -| Memory | 40–80 GB | **16 GB** | Enable CPU offloading | -| cp.async | Available | No (Turing has limited async) | Remove async copy patterns | -| `AlignedVector` | Same | Same | No changes | -| `warp::reduce_sum` | Same | Same | No changes | - ---- - -## Performance Profiling - -### NVIDIA Nsight Systems (nsys) - -```bash -nsys profile -o t4_profile python scripts/bench_diffusion_rmsnorm.py - -# Key metrics to watch: -# - Kernel duration -# - Memory transfer time -# - GPU idle time -# - Stream utilization -``` - -### NVIDIA Nsight Compute (ncu) - -```bash -# Full metrics -ncu --set full -o t4_metrics.ncu-rep \ - python scripts/bench_diffusion_rmsnorm.py - -# Specific metrics — T4 is memory-bound; focus on dram throughput -ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\ -dram__throughput.avg.pct_of_peak_sustained_elapsed \ - python scripts/bench_diffusion_rmsnorm.py - -# Key metrics for T4 diffusion kernels: -# - Memory throughput (dram__throughput.avg.pct_of_peak_sustained_elapsed) -# → Target: 40–50% of 320 GB/s (128–160 GB/s) for vectorized kernels -# - SM utilization (sm__throughput.avg.pct_of_peak_sustained_elapsed) -# → Target high with only 40 SMs -# - Achieved occupancy (sm__warps_active.avg.pct_of_peak_sustained_active) -# → Max 1024 threads/SM on T4 — block size ≤ 512 for decent occupancy -# - Warp stall reasons (smsp__warp_issue_stalled_*.avg.pct_of_peak_sustained_active) -``` - -### Common T4 Bottlenecks - -1. **Memory Bandwidth** — 320 GB/s is the primary limit; if `dram__throughput` < 40% → use `AlignedVector` -2. **Limited Memory** — 16 GB; enable `--dit-cpu-offload`/`--vae-cpu-offload` as needed -3. **No BF16** — guard in Python wrapper; FP16 overflow risk in long-sequence attention -4. **Smaller tiles** — 64 KB shared memory; reduce `BLOCK_SIZE_M/N` vs H100 - ---- - -## Best Practices Summary (T4) - -1. **No BF16**: Guard in Python wrapper, raise clear error -2. **Vectorization**: Even more critical at 320 GB/s — always use `AlignedVector` -3. **Tile sizes**: 64 KB shared memory limit → halve BLOCK_SIZE vs H100 -4. **Block size**: Max 512 threads/block for decent occupancy (max 1024 threads/SM) -5. **Grid sizing**: 40 SMs — `runtime::get_sm_count()` auto-handles -6. **FP32 accumulation**: Always accumulate in FP32 to avoid FP16 overflow -7. **Memory**: Plan for 16 GB; use `--dit-cpu-offload`/`--vae-cpu-offload` as needed -8. **Fuse more**: Low bandwidth makes kernel fusion more impactful than on H100 -9. **Multi-arch build**: Always build for `sm_75,sm_80,sm_90` together - -## T4 Cloud Instance Quick Reference - -| Provider | Instance | Notes | -|----------|----------|-------| -| GCP | n1-standard-4 + T4 | Most common inference setup | -| AWS | g4dn.xlarge | 1× T4, 16 GB | -| AWS | g4dn.12xlarge | 4× T4, 64 GB total | -| Azure | NC4as T4 v3 | 1× T4 | diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/troubleshooting.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/troubleshooting.md deleted file mode 100644 index 61db55b98..000000000 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cuda-kernel/references/troubleshooting.md +++ /dev/null @@ -1,330 +0,0 @@ -# Troubleshooting Guide — SGLang Diffusion JIT CUDA Kernels - -Common issues and solutions when writing and integrating JIT CUDA kernels for SGLang Diffusion. - -> **Adapted from**: [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels) - ---- - -## Build / Compile Issues - -### 1. JIT compilation fails: "No such file or directory" - -**Problem:** `load_jit` cannot find your `.cuh` file. - -``` -FileNotFoundError: .../jit_kernel/csrc/diffusion/your_op.cuh not found -``` - -**Fix:** Ensure the file is under `python/sglang/jit_kernel/csrc/diffusion/`. The path passed to `cuda_files` is relative to `csrc/`: - -```python -# CORRECT — file lives at csrc/diffusion/your_op.cuh -load_jit(..., cuda_files=["diffusion/your_op.cuh"]) -# resolves to: python/sglang/jit_kernel/csrc/diffusion/your_op.cuh - -# ALSO CORRECT — absolute path (pathlib replaces the csrc/ prefix) -load_jit(..., cuda_files=["/full/absolute/path/to/your_op.cuh"]) -``` - -### 2. Type conversion errors (FP16/BF16) - -**Problem:** Implicit FP16/BF16 conversion fails because PyTorch compiles with `-D__CUDA_NO_HALF_OPERATORS__`: - -``` -error: no suitable conversion function from "__half" to "float" exists -``` - -**Fix:** SGLang's `static_cast` works because `fp16_t` and `bf16_t` are typedef'd with proper conversion operators. Always use explicit casts: - -```cpp -// CORRECT — explicit cast -float val = static_cast(v[i]); // fp16_t / bf16_t → float -v[i] = static_cast(fp32_result); // float → T - -// WRONG — implicit conversion (disabled by PyTorch build flags) -float val = v[i]; // compile error -v[i] = fp32_result; // compile error -``` - -If you need the raw intrinsics for packed types: -```cpp -// bf16x2_t → two floats -bf16x2_t packed = ...; -float v0 = __bfloat162float(packed.x); -float v1 = __bfloat162float(packed.y); -``` - -### 3. Template instantiation explodes / slow first compile - -**Problem:** Many template combinations makes the first JIT compile very slow. - -**Fix:** Reduce template argument combinations. Move compile-time constants to runtime if they don't affect performance critically: - -```cpp -// Fewer template args = fewer instantiations -template // only dtype varies -void my_op(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, int block_size); -``` - -### 4. SM check: kernel requires sm_90 but device is sm_80 - -**Problem:** Kernel uses H100-only features on A100. - -**Fix:** Add a Python guard before calling `load_jit`: - -```python -cap = torch.cuda.get_device_capability() -if cap[0] < 9: - raise RuntimeError( - f"This kernel requires H100 (sm_90+). " - f"Got compute capability {cap[0]}.{cap[1]}. " - f"Use the Triton fallback instead: diffusion_triton_()" - ) -``` - ---- - -## Performance Issues - -### 5. Kernel is slower than Triton / PyTorch baseline - -**Steps to diagnose:** - -1. Check dtype: are you using `bf16_t` on T4? (T4 has no BF16 — silently falls back to slow emulation) -2. Check vectorization: is `hidden_size` divisible by `kVecN = 16/sizeof(T)` (8 for bf16, 4 for fp32)? -3. Profile with `ncu`: - ```bash - ncu --set full --csv -o metrics.csv \ - python -c "from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm; ..." - ``` - Look at `dram__throughput.avg.pct_of_peak_sustained_elapsed` — if < 30%, check coalescing. - -4. Check occupancy: run with `--ptxas-options=-v` in `extra_cuda_cflags` to see register usage. - -### 6. Shared memory bank conflicts - -**Problem:** `ncu` reports high `l1tex__data_bank_conflicts_pipe_lmem_op_st.sum`. - -**Fix:** Add padding to shared memory arrays: - -```cpp -// Conflict (all threads hit same bank when stride=32) -__shared__ float data[32][32]; - -// Fixed with padding -__shared__ float data[32][33]; // 33 instead of 32 -``` - -### 7. Low occupancy from too many registers - -**Problem:** `nvcc --ptxas-options=-v` shows high register count; occupancy < 25%. - -**Fix:** Add `--maxrregcount=N` to limit registers: - -```python -extra_cuda_cflags=["-O3", "--use_fast_math", "--maxrregcount=64"] -``` - -Reduces registers per thread at the cost of possible register spilling to local memory. - ---- - -## Integration Issues - -### 8. RMSNorm weight is None (`elementwise_affine=False`) - -**Problem:** -``` -AttributeError: 'NoneType' object has no attribute 'data_ptr' -``` - -**Root Cause:** DiT transformer blocks often use `RMSNorm(dim, elementwise_affine=False)` — no learnable weight. - -**Fix in Python wrapper:** pass an empty tensor when weight is absent; the kernel launcher checks `data_ptr == nullptr`: - -```python -w = weight if weight is not None else torch.empty(0, dtype=src.dtype, device=src.device) -module.rmsnorm(out, src, w, eps) -``` - -**Fix in `.cuh` launcher:** - -```cpp -const T* w_ptr = (weight.data_ptr() != nullptr) - ? static_cast(weight.data_ptr()) : nullptr; -// ... pass w_ptr to kernel ... -``` - -**Fix in module patching:** - -```python -has_weight = hasattr(module, "weight") and module.weight is not None -if has_weight: - def _fwd(mod, eps): - def forward(x): return diffusion_rmsnorm(x, weight=mod.weight, eps=eps) - return forward - module.forward = _fwd(module, module.eps) -else: - def _fwd_noweight(eps): - def forward(x): return diffusion_rmsnorm(x, weight=None, eps=eps) - return forward - module.forward = _fwd_noweight(module.eps) -``` - -### 9. `isinstance(module, torch.nn.RMSNorm)` misses diffusion variants - -**Problem:** Patching doesn't apply because diffusers / sglang diffusion models define their own `RMSNorm` class that is **not** a subclass of `torch.nn.RMSNorm`. - -**Fix:** Match by class name string: - -```python -# WRONG — misses diffusers/sglang RMSNorm -if isinstance(module, torch.nn.RMSNorm): - -# CORRECT — catches all variants -if type(module).__name__ == "RMSNorm": -# or for broader matching: -if "RMSNorm" in type(module).__name__: -``` - -### 10. Kernel patching doesn't persist after CPU offloading - -**Problem:** After calling `pipe.enable_model_cpu_offload()`, patched modules revert. - -**Fix:** Always inject **after** moving to CUDA, **before** enabling any offloading: - -```python -pipe = load_pipeline(...) -pipe.to("cuda") # 1. Move to CUDA -inject_optimized_kernels(pipe) # 2. Patch modules -pipe.enable_model_cpu_offload() # 3. Now safe to enable offloading -``` - -### 11. Kernel patched after `torch.compile` - -**Problem:** Module is already compiled; patching its `forward` after compilation has no effect. - -**Fix:** Apply patches **before** any `torch.compile` call: - -```python -inject_optimized_kernels(pipe) # FIRST: patch -pipe.transformer = torch.compile(...) # SECOND: compile -``` - ---- - -## `torch.compile` Compatibility - -### 12. Custom CUDA kernel causes graph break - -**Problem:** -``` -torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped -``` -or: -``` -torch._dynamo.exc.TorchRuntimeError: Cannot access data pointer of Tensor (FakeTensor) -``` - -**Root Cause:** `torch.compile` traces with "fake tensors" that have no real data. Any kernel that calls `.data_ptr()` during tracing fails. - -**Options:** - -**Option A (simplest):** Don't use `torch.compile` with CUDA JIT kernels — use Triton instead: -```python -# Triton kernels are torch.compile compatible -from sglang.jit_kernel.diffusion.triton.norm import fused_rmsnorm -``` - -**Option B:** Register as a `@torch.library.custom_op` (advanced): -```python -import torch - -@torch.library.custom_op("diffusion_jit::rmsnorm", mutates_args={"out"}) -def _rmsnorm_op(out: torch.Tensor, src: torch.Tensor, - weight: torch.Tensor, eps: float) -> None: - module = _jit_rmsnorm_module(src.dtype) - module.rmsnorm(out, src, weight, eps) - -@_rmsnorm_op.register_fake -def _(out, src, weight, eps): - pass # no shape changes; output already allocated in 'out' -``` - -**Performance trade-off:** - -| Approach | Speedup (denoise) | torch.compile | Notes | -|----------|-------------------|---------------|-------| -| CUDA JIT kernel | best | Yes (via `torch.library.custom_op`) | Performance-optimal regardless of whether `torch.compile` is enabled; use `custom_op` + `register_fake` for compile compatibility | -| Triton kernel | good | Yes | Use when you need faster iteration/portability, or when you do not have a well-tuned CUDA kernel yet | -| Triton + compile | good | Yes | Use for end-to-end `torch.compile` integration convenience; typically slower than a well-tuned CUDA kernel | - -### 13. Unstable benchmark results from JIT timing - -**Problem:** First few runs are slow due to JIT compilation; timing is noisy. - -**Fix:** Use `triton.testing.do_bench` / `run_benchmark` which use CUDA-graph-based timing automatically. Always do a warmup run first: - -```python -# Pre-compile by running once before timing -diffusion_rmsnorm(dummy_src, weight=dummy_w, eps=1e-6) -torch.cuda.synchronize() -# Now time -result = run_benchmark(lambda: diffusion_rmsnorm(src, weight=w, eps=1e-6)) -``` - ---- - -## Debugging Checklist - -```bash -# 1. Verify CUDA device and compute capability -python -c "import torch; print(torch.cuda.get_device_name(), torch.cuda.get_device_capability())" - -# 2. Force synchronous CUDA execution to get real error location -CUDA_LAUNCH_BLOCKING=1 python scripts/bench_diffusion_rmsnorm.py - -# 3. Run memory sanitizer to catch illegal accesses -compute-sanitizer --tool memcheck python scripts/bench_diffusion_rmsnorm.py - -# 4. Check register and shared memory usage -# Add to extra_cuda_cflags: "--ptxas-options=-v" - -# 5a. Kernel-level profiling — full metrics -ncu --set full -o metrics.ncu-rep \ - python scripts/bench_diffusion_rmsnorm.py - -# 5b. Kernel-level profiling — targeted bandwidth + occupancy check -ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\ -dram__throughput.avg.pct_of_peak_sustained_elapsed \ - python scripts/bench_diffusion_rmsnorm.py - -# Key metrics to interpret: -# - sm__throughput : compute utilization % of peak -# - dram__throughput: memory bandwidth % of peak (target ≥ 30% on H100/A100) -# - 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= -nsys profile -o denoise_profile \ - sglang generate --model-path=black-forest-labs/FLUX.1-dev \ - --width=1024 --height=1024 --num-inference-steps=50 \ - --seed=42 --enable-torch-compile --warmup - -# 7. Verify a patched module produces correct output -python - << 'EOF' -import torch -from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm - -x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") -w = torch.ones(2048, dtype=torch.bfloat16, device="cuda") - -out_jit = diffusion_rmsnorm(x, weight=w, eps=1e-6) -out_ref = torch.nn.functional.rms_norm(x.float(), (2048,), w.float(), eps=1e-6).to(torch.bfloat16) - -max_diff = (out_jit - out_ref).abs().max().item() -print(f"Max diff: {max_diff:.2e} ({'PASS' if max_diff < 0.02 else 'FAIL'})") -EOF -``` diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md index a94f9ed06..e81449004 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md @@ -26,8 +26,8 @@ These options are intended to preserve output quality. In practice, some paths ( | **torch.compile** | `--enable-torch-compile` | Applies `torch.compile` to the DiT forward pass, fusing ops and reducing kernel launch overhead. | ~1.2–1.5x on denoising | First request is slow (compilation). May cause minor precision drifts due to [PyTorch issue #145213](https://github.com/pytorch/pytorch/issues/145213). Pair with `--warmup` for best results. | | **Warmup** | `--warmup` | Runs dummy forward passes to warm up CUDA caches, JIT, and `torch.compile`. Eliminates cold-start penalty. | Removes first-request latency spike | Adds startup time. Without `--warmup-resolutions`, warmup happens on first request. | | **Warmup Resolutions** | `--warmup-resolutions 256x256 720x720` | Pre-compiles and warms up specific resolutions at server startup (instead of lazily on first request). | Faster first request per resolution | Each resolution adds to startup time. Serving mode only; useful when you know your target resolutions in advance. | -| **Multi-GPU (SP)** | `--num-gpus N --ulysses-degree N` | Sequence parallelism across GPUs. Shards sequence tokens (not frames) to minimize padding. | Near-linear scaling with N GPUs | Requires NCCL; inter-GPU bandwidth matters. `ulysses_degree * ring_degree = sp_degree`. | -| **CFG Parallel** | `--enable-cfg-parallel` | Runs conditional and unconditional CFG branches in parallel across GPUs. For CFG models on multi-GPU, benchmark this against pure Ulysses on your topology instead of assuming one always wins. | Often faster than pure SP for CFG models | Requires `num_gpus >= 2`. Halves the Ulysses group size (e.g. 8 GPU → two 4-GPU groups). Only for models that use CFG. | +| **Multi-GPU (SP)** | `--num-gpus N --ulysses-degree N` | Sequence parallelism across GPUs. Shards sequence tokens (not frames) to minimize padding. | Near-linear scaling with N GPUs | Requires NCCL; inter-GPU bandwidth matters. `ulysses_degree * ring_degree = sp_degree`. For Wan2.2 video, start by benchmarking pure Ulysses before assuming a mixed Ulysses/Ring layout is fastest. | +| **CFG Parallel** | `--enable-cfg-parallel` | Runs conditional and unconditional CFG branches in parallel across GPUs. For CFG models on multi-GPU, benchmark this against pure Ulysses on your topology instead of assuming one always wins. | Often faster than pure SP for CFG models | Requires `num_gpus >= 2`. Halves the Ulysses group size (e.g. 8 GPU → two 4-GPU groups). Only for models that use CFG. Nightly coverage configs may intentionally use smaller Ulysses groups to keep ring behavior exercised; that does not automatically make them the lowest-latency choice. | | **Layerwise Offload** | `--dit-layerwise-offload` | Async layer-by-layer H2D prefetch with compute overlap. Only ~2 DiT layers reside on GPU at a time, dramatically reducing VRAM. For some video models the copy stream can be almost fully hidden behind compute ([PR #15511](https://github.com/sgl-project/sglang/pull/15511)). | Saves VRAM (40 GB → ~11 GB for Wan A14B); can be near-zero speed cost on the right workload | Enabled by default for Wan/MOVA video models. Incompatible with Cache-DiT. For **image models** or highly parallelized setups (many GPUs, small per-GPU compute), the copy stream may not be fully hidden and can cause slowdown. | | **Offload Prefetch Size** | `--dit-offload-prefetch-size F` | Fine-grained control over layerwise offload: how many layers to prefetch ahead. `0.0` = 1 layer (min VRAM), `0.1` = 10% of layers, `≥1` = absolute layer count. | Tune for cases where default offload has copy stream interference (e.g. image models). 0.05–0.1 is a good starting point. | Values ≥ 0.5 approach no-offload VRAM with worse performance. See [PR #17693](https://github.com/sgl-project/sglang/pull/17693) for benchmarks on image models. | | **FSDP Inference** | `--use-fsdp-inference` | Uses PyTorch FSDP to shard model weights across GPUs with prefetch. Low latency, low VRAM. | Reduces per-GPU VRAM | Mutually exclusive with `--dit-layerwise-offload`. More overhead than SP on high-bandwidth interconnects. | @@ -67,6 +67,27 @@ sglang generate --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ Note: `--dit-layerwise-offload` is enabled by default for Wan/MOVA video models and is often a good default, but still benchmark it on your exact workload if latency matters. +For Wan2.2 specifically: +- the nightly-aligned 4-GPU benchmark may use `--enable-cfg-parallel --ulysses-degree=2` to keep CFG and ring behavior covered +- that is a **coverage** choice, not a guaranteed best-performance choice +- for pure latency tuning, benchmark pure Ulysses too, for example `--ulysses-degree=4 --ring-degree=1` on 4 GPUs +- on 8 GPUs, compare pure `--ulysses-degree=8` against `--enable-cfg-parallel --ulysses-degree=4` + +### Nightly-aligned model, single GPU: LTX-2 two-stage + +```bash +sglang generate --model-path Lightricks/LTX-2 \ + --pipeline-class-name LTX2TwoStagePipeline \ + --prompt "A beautiful sunset over the ocean" \ + --negative-prompt "shaky, glitchy, low quality, worst quality, deformed, distorted, disfigured, motion smear, motion artifacts, fused fingers, bad anatomy, weird hand, ugly, transition, static." \ + --width 1536 --height 1024 \ + --num-frames 121 --fps 24 \ + --seed 1234 --num-gpus 1 \ + --enable-torch-compile --warmup --save-output +``` + +Note: this generate recipe is aligned with the nightly comparison case `ltx2_twostage_t2v`. After [PR #20707](https://github.com/sgl-project/sglang/pull/20707), `LTX2TwoStagePipeline` is a native path and auto-resolves the spatial upsampler plus distilled LoRA from the same model snapshot unless you override them. + ### Maximum speed, image model, single GPU, lossless ```bash @@ -107,3 +128,7 @@ SGLANG_CACHE_DIT_ENABLED=true sglang generate --model-path \ - **Perf dump**: use `--perf-dump-path result.json` to save structured metrics, then compare with `python python/sglang/multimodal_gen/benchmarks/compare_perf.py baseline.json result.json`. - **Offload tuning**: after the first request, the runtime logs peak GPU memory and which components could stay resident. Use this to decide which `--*-cpu-offload` flags to disable. - **Backend selection**: `--backend sglang` (default, auto-detected) enables all native optimizations (fused kernels, SP, etc.). `--backend diffusers` falls back to vanilla Diffusers pipelines but supports `--cache-dit-config` and diffusers attention backends. +- **Wan2.2-I2V sizing**: after [PR #21390](https://github.com/sgl-project/sglang/pull/21390), explicit `--width/--height` on `Wan2.2-I2V-A14B` control the target area while preserving the condition-image aspect ratio. +- **Merged diffusion fast paths**: before proposing a new kernel or overlap scheme, check `sglang-diffusion-benchmark-profile/existing-fast-paths.md`. It now covers merged Z-Image residual-form modulation, fused diffusion `QK norm + RoPE`, and existing multi-GPU overlap families such as Ulysses / USP and turbo-layer async all-to-all. +- **NVFP4 trace interpretation**: on FLUX.2 NVFP4 and Nunchaku-style checkpoints, packed QKV is expected. SGLang intentionally uses fused projection modules such as `to_qkv` / `to_added_qkv` instead of separate `to_q` / `to_k` / `to_v`, so a split-QKV trace usually means the quantized path did not engage rather than a brand new fusion opportunity. +- **Hotspot workflow split**: use `sglang-diffusion-benchmark-profile` to prove and classify a slowdown with perf dumps plus `torch.profiler`; hand concrete kernel work to `sglang-diffusion-ako4all-kernel` or another specialized optimization skill instead of expanding the benchmark skill. diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-triton-kernel/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-triton-kernel/SKILL.md deleted file mode 100644 index 43ff87d7d..000000000 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-triton-kernel/SKILL.md +++ /dev/null @@ -1,515 +0,0 @@ ---- -name: sglang-diffusion-triton-kernel -description: Use when writing or tuning a Triton diffusion kernel in SGLang. ---- - -# Adding a Triton Kernel to SGLang Diffusion - -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 `../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. - ---- - -## Directory Layout - -``` -python/sglang/jit_kernel/diffusion/ -├── triton/ -│ ├── scale_shift.py # AdaLN scale/shift fused kernels -│ ├── norm.py # LayerNorm / RMSNorm fused kernels -│ ├── rmsnorm_onepass.py # One-pass RMSNorm for small hidden size -│ └── rotary.py # RoPE kernel -└── cutedsl/ - └── ... # CuTe DSL kernels (see existing-fast-paths.md in the benchmark/profile skill) -``` - -New Triton kernels go into `triton/.py`. - ---- - -## Step 1: Write the Triton Kernel - -Create `python/sglang/jit_kernel/diffusion/triton/.py`. - -### 1a. Imports - -```python -import torch -import triton # type: ignore -import triton.language as tl # type: ignore -``` - -Always use `# type: ignore` on triton imports — the stubs are incomplete. - -### 1b. The `@triton.jit` Kernel Function - -Follow the naming convention `__kernel` (private, underscore prefix). - -```python -@triton.autotune( - configs=[ - triton.Config({"BLOCK_C": 64}, num_warps=2), - triton.Config({"BLOCK_C": 128}, num_warps=4), - triton.Config({"BLOCK_C": 256}, num_warps=4), - triton.Config({"BLOCK_C": 512}, num_warps=8), - ], - key=["C"], # re-tune when hidden dim changes -) -@triton.jit -def _fused_scale_shift_kernel( - # Pointers — always pass raw tensors; Triton takes .data_ptr() internally - x_ptr, - scale_ptr, - shift_ptr, - y_ptr, - # Dimensions - B, # batch size - L, # sequence length - C, # hidden / channel dim - # Strides — pass every stride separately; do NOT assume contiguous - stride_xb, stride_xl, stride_xc, - stride_sb, stride_sc, - stride_yb, stride_yl, stride_yc, - # Compile-time constants (tl.constexpr) - BLOCK_C: tl.constexpr, -): - # Grid: (cdiv(L, 1), B) — one program per (batch, token) - pid_l = tl.program_id(0) - pid_b = tl.program_id(1) - - c_offs = tl.arange(0, BLOCK_C) - mask = c_offs < C - - x_row = pid_b * stride_xb + pid_l * stride_xl - y_row = pid_b * stride_yb + pid_l * stride_yl - s_row = pid_b * stride_sb - - x = tl.load(x_ptr + x_row + c_offs * stride_xc, mask=mask, other=0.0) - scale = tl.load(scale_ptr + s_row + c_offs * stride_sc, mask=mask, other=0.0) - shift = tl.load(shift_ptr + s_row + c_offs * stride_sc, mask=mask, other=0.0) - - y = x * (1.0 + scale) + shift - tl.store(y_ptr + y_row + c_offs * stride_yc, y, mask=mask) -``` - -**Rules:** -- All pointer arguments are raw (Triton extracts `.data_ptr()` internally when called via `kernel[grid](...)`). -- Pass every stride as a separate scalar — never assume a tensor is contiguous inside the kernel. -- Use `tl.constexpr` for block sizes and boolean flags (`HAS_RESIDUAL`, `IS_RMS_NORM`, etc.). -- Use `mask=mask, other=0.0` on every `tl.load` to avoid out-of-bounds reads. -- Compute in `tl.float32` when precision matters (`x.to(tl.float32)`), then cast back to output dtype before `tl.store`. -- Use `tl.fma(a, b, c)` (`a*b + c`) for fused multiply-add — avoids rounding errors and maps to a single instruction. - -### 1c. `@triton.autotune` Guidelines - -| `key` entry | When to include | -|-------------|-----------------| -| `"C"` / `"hidden_dim"` | Always — block tile size depends on C | -| `"IS_RMS_NORM"` | When the kernel has a `constexpr` boolean flag that changes code paths | -| `"HAS_RESIDUAL"` | Same — constexpr path branching | -| Shape / batch / seq | Usually NOT — autotune cost outweighs benefit | - -Keep configs in ascending `BLOCK_C` order with matching `num_warps` (warp × 32 threads ≤ 1024). - -### 1d. `torch.compile` Compatibility - -When the kernel is called inside a `torch.compile`-d region, wrap the launch with `torch.library.wrap_triton`: - -```python -with torch.get_device_module().device(x.device): - torch.library.wrap_triton(_fused_scale_shift_kernel)[grid]( - x, scale, shift, y, - B, L, C, - x.stride(0), x.stride(1), x.stride(2), - scale.stride(0), scale.stride(1), - y.stride(0), y.stride(1), y.stride(2), - ) -``` - -Use `wrap_triton` when the kernel is called from a layer that runs under `torch.compile`. -Skip it for utility kernels called only at Python graph boundaries. - ---- - -## Step 2: Write the Python Launcher - -The launcher is a regular Python function (public, no underscore) in the same file. - -```python -def fused_scale_shift( - x: torch.Tensor, - scale: torch.Tensor, - shift: torch.Tensor, -) -> torch.Tensor: - """ - Fused AdaLN modulation: y = x * (1 + scale) + shift. - - Args: - x: [B, L, C], CUDA, contiguous - scale: [B, C], CUDA - shift: [B, C], CUDA (same shape as scale) - - Returns: - y: same shape and dtype as x - """ - # --- Precondition checks --- - assert x.is_cuda, "x must be on CUDA" - assert x.is_contiguous(), "x must be contiguous" - assert scale.is_cuda and shift.is_cuda - assert x.ndim == 3, f"x must be 3D [B, L, C], got {x.shape}" - assert scale.shape == shift.shape - B, L, C = x.shape - - # Allocate output - y = torch.empty_like(x) - - # Grid: one program per token - grid = (L, B) - - _fused_scale_shift_kernel[grid]( - x, scale, shift, y, - B, L, C, - x.stride(0), x.stride(1), x.stride(2), - scale.stride(0), scale.stride(1), - y.stride(0), y.stride(1), y.stride(2), - ) - return y -``` - -**Rules:** -- Validate CUDA placement and shape/dtype **before** launching — use `assert` with a helpful message. -- Call `.contiguous()` on inputs that the kernel requires contiguous **before** the launch, not inside it. -- Allocate the output with `torch.empty_like(x)` — never reuse input buffers unless the op is explicitly in-place. -- The `grid` is a tuple or a lambda `(META)` when block sizes are auto-tuned: - -```python -# Static grid (block size fixed) -grid = (triton.cdiv(L, BLOCK_L), triton.cdiv(C, BLOCK_C), B) - -# Dynamic grid (block size comes from autotune) -grid = lambda META: (triton.cdiv(L, META["BLOCK_C"]), B) -``` - -### Handling Non-Contiguous Inputs - -Never call `.contiguous()` silently — it copies data. Instead, pass strides to the kernel and let it handle arbitrary layouts. Only call `.contiguous()` when the kernel genuinely requires it (e.g., after a reshape): - -```python -# OK: reshape + contiguous needed for 2D view trick -x_2d = x.view(B * L, C) # view only works on contiguous -if not x.is_contiguous(): - x = x.contiguous() - x_2d = x.view(B * L, C) -``` - ---- - -## Step 3: Integrate into the Layer - -Call the new kernel from the appropriate layer file in -`python/sglang/multimodal_gen/runtime/layers/` (typically `layernorm.py` or `elementwise.py`). - -```python -# In layernorm.py or elementwise.py -import torch - -def apply_scale_shift(x, scale, shift): - if x.is_cuda: - from sglang.jit_kernel.diffusion.triton.my_op import fused_scale_shift - return fused_scale_shift(x, scale, shift) - # Pure-PyTorch fallback for non-CUDA execution - return x * (1.0 + scale) + shift -``` - -**Rules:** -- Gate on `x.is_cuda` — the Triton kernel only runs on CUDA; the fallback handles everything else. -- The launcher raises `AssertionError` on invalid inputs (wrong shape, CPU tensor, etc.) — do **not** silently catch these. Let them propagate so bugs are visible during development. -- Add `logger.warning_once(...)` only when falling back due to a **known hardware limitation** (e.g., unsupported SM compute capability), not for wrong-input errors. - ---- - -## Step 4: Write Tests - -Create `python/sglang/jit_kernel/tests/test_.py`. - -```python -import pytest -import torch - -from sglang.jit_kernel.diffusion.triton.my_op import fused_scale_shift - - -def _ref_fused_scale_shift(x, scale, shift): - """PyTorch reference implementation.""" - # Broadcast scale/shift from [B, C] to [B, L, C] - return x * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) - - -@pytest.fixture(autouse=True) -def require_cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA required") - - -@pytest.mark.parametrize("B,L,C", [ - (1, 6, 3072), # Qwen (small batch) - (1, 1024, 1536), # Wan - (2, 512, 3072), # typical training shape - (1, 1, 256), # edge: L=1 -]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) -def test_fused_scale_shift_correctness(B, L, C, dtype): - torch.manual_seed(0) - x = torch.randn(B, L, C, dtype=dtype, device="cuda") - scale = torch.randn(B, C, dtype=dtype, device="cuda") * 0.1 - shift = torch.randn(B, C, dtype=dtype, device="cuda") * 0.1 - - out = fused_scale_shift(x, scale, shift) - ref = _ref_fused_scale_shift(x.float(), scale.float(), shift.float()).to(dtype) - - atol = 1e-5 if dtype == torch.float32 else 1e-2 - torch.testing.assert_close(out, ref, atol=atol, rtol=atol, - msg=f"Mismatch at B={B} L={L} C={C} dtype={dtype}") - - -def test_fused_scale_shift_non_cuda_raises(): - x = torch.randn(1, 4, 64) - scale = torch.randn(1, 64) - shift = torch.randn(1, 64) - with pytest.raises(AssertionError, match="CUDA"): - fused_scale_shift(x, scale, shift) - - -def test_fused_scale_shift_output_dtype_preserved(): - x = torch.randn(1, 8, 128, dtype=torch.bfloat16, device="cuda") - scale = torch.randn(1, 128, dtype=torch.bfloat16, device="cuda") - shift = torch.zeros(1, 128, dtype=torch.bfloat16, device="cuda") - out = fused_scale_shift(x, scale, shift) - assert out.dtype == torch.bfloat16 - assert out.shape == x.shape - - -if __name__ == "__main__": - import sys - sys.exit(pytest.main([__file__, "-v"])) -``` - -Run: - -```bash -pytest python/sglang/jit_kernel/tests/test_.py -v -``` - -**Test coverage requirements:** -1. Reference comparison against pure-PyTorch for all supported dtypes (fp16, bf16, fp32). -2. Edge shapes: `L=1`, `C` not a multiple of the largest BLOCK_C, large `B`. -3. Error cases: CPU tensor, wrong shape. -4. Output dtype and shape preservation. - ---- - -## Step 5: Add a Benchmark (required) - -Create `python/sglang/jit_kernel/benchmark/bench_.py`. - -```python -import torch -import triton.testing - -from sglang.jit_kernel.diffusion.triton.my_op import fused_scale_shift - - -SHAPES = [ - # (B, L, C) — representative diffusion shapes - (1, 6, 3072), # Qwen image - (1, 1024, 1536), # Wan video - (1, 4096, 3072), # FLUX double-stream -] - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["B", "L", "C"], - x_vals=SHAPES, - line_arg="provider", - line_vals=["triton", "torch"], - line_names=["Triton Fused", "PyTorch"], - styles=[("blue", "-"), ("red", "--")], - ylabel="µs (median)", - plot_name="fused-scale-shift", - args={}, - ) -) -def benchmark(B, L, C, provider): - dtype = torch.bfloat16 - x = torch.randn(B, L, C, dtype=dtype, device="cuda") - scale = torch.randn(B, C, dtype=dtype, device="cuda") - shift = torch.randn(B, C, dtype=dtype, device="cuda") - - if provider == "triton": - fn = lambda: fused_scale_shift(x, scale, shift) - else: - fn = lambda: x * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) - - ms, *_ = triton.testing.do_bench_cudagraph(fn, quantiles=[0.5, 0.2, 0.8]) - return ms * 1000 # µs - - -if __name__ == "__main__": - benchmark.run(print_data=True) -``` - -Run: - -```bash -python python/sglang/jit_kernel/benchmark/bench_.py -``` - ---- - -## Step 6: Profile with Nsight Compute (required for optimization work) - -After correctness tests, you must use **ncu (Nsight Compute)** to validate hardware efficiency (bandwidth/throughput/occupancy/bottleneck type). - -To avoid duplicating ncu CLI details across multiple skills, this skill does not repeat command flags. Follow the canonical docs: - -- `../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) - ---- - -## Common Patterns Reference - -### Pattern 1: Autotune over a 2D tile (L × C) - -Used in `scale_shift.py` (`fuse_scale_shift_kernel_blc_opt`): - -```python -@triton.jit -def _kernel(..., BLOCK_L: tl.constexpr, BLOCK_C: tl.constexpr): - pid_l = tl.program_id(0) - pid_c = tl.program_id(1) - pid_b = tl.program_id(2) - l_offs = pid_l * BLOCK_L + tl.arange(0, BLOCK_L) - c_offs = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) - mask = (l_offs[:, None] < L) & (c_offs[None, :] < C) - ... - -# Launch: -grid = (triton.cdiv(L, BLOCK_L), triton.cdiv(C, BLOCK_C), B) -_kernel[grid](..., BLOCK_L=block_l, BLOCK_C=block_c, num_warps=4, num_stages=2) -``` - -### Pattern 2: One-pass RMSNorm for small hidden size - -Used in `rmsnorm_onepass.py`: - -```python -@triton.jit -def _rms_norm_tiled_onepass(y_ptr, x_ptr, w_ptr, - SEQ: tl.constexpr, DIM: tl.constexpr, EPS: tl.constexpr, - BLOCK_SIZE_SEQ: tl.constexpr, BLOCK_SIZE_DIM: tl.constexpr): - seq_blk_id = tl.program_id(0) - seq_id = seq_blk_id * BLOCK_SIZE_SEQ - seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None] - d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :] - ... - x = tl.load(x_ptr + seq_offset * DIM + d_offset, mask=..., other=0.0).to(tl.float32) - mean_sq = tl.sum(x * x, axis=1, keep_dims=True) / DIM - rstd = tl.math.rsqrt(mean_sq + EPS) - tl.store(y_ptr + ..., x * rstd * w, mask=...) - -# Launch with wrap_triton for torch.compile compat: -with torch.get_device_module().device(x.device): - torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid]( - y_view, x_view, w, - S, D, eps, - BLOCK_SIZE_DIM=triton.next_power_of_2(D), - BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ, - ) -``` - -### Pattern 3: `tl.constexpr` boolean flags for conditional paths - -Used in `norm.py` and `scale_shift.py`: - -```python -@triton.jit -def _kernel(..., - IS_RMS_NORM: tl.constexpr, - HAS_RESIDUAL: tl.constexpr, - SCALE_IS_SCALAR: tl.constexpr): - ... - if IS_RMS_NORM: - var = tl.sum(x * x, axis=0) / N - else: - mean = tl.sum(x, axis=0) / N - var = tl.sum((x - mean) ** 2, axis=0) / N - - if HAS_RESIDUAL: - x = x + tl.load(residual_ptr + ...) - - if SCALE_IS_SCALAR: - scale_val = tl.load(scale_ptr) - scale = tl.full([BLOCK_N], scale_val, dtype=scale_val.dtype) - else: - scale = tl.load(scale_ptr + col_offsets, mask=mask, other=0.0) -``` - -Autotune key must include these booleans so the compiler generates separate specializations. - -### Pattern 4: Computing in fp32, storing in original dtype - -Always up-cast to `tl.float32` for reductions and math, then down-cast before storing: - -```python -x_f32 = x.to(tl.float32) -scale_f32 = scale.to(tl.float32) -y_f32 = x_f32 * (1.0 + scale_f32) + shift_f32 -tl.store(y_ptr + offsets, y_f32.to(x.dtype), mask=mask) -``` - ---- - -## Checklist Before Submitting - -### Prerequisites -- [ ] `ncu --version` prints a valid Nsight Compute version (required for Step 7 profiling) - -### Implementation -- [ ] Kernel file at `python/sglang/jit_kernel/diffusion/triton/.py` -- [ ] All pointer arguments passed with separate stride scalars -- [ ] Every `tl.load` uses `mask=` and `other=` -- [ ] Autotune `key` includes all `constexpr` flags that change code paths -- [ ] `torch.library.wrap_triton` used if kernel runs inside `torch.compile` region -- [ ] PyTorch fallback path in the layer integration (see Step 4) - -### Validation -- [ ] Tests pass: `pytest python/sglang/jit_kernel/tests/test_.py -v` -- [ ] Benchmark runs: `python python/sglang/jit_kernel/benchmark/bench_.py` -- [ ] **Correctness verified**: Triton output matches PyTorch reference within tolerance -- [ ] Nsight Compute profile collected (`ncu --set full`); achieved occupancy ≥ 50% and memory throughput ≥ 70% of peak (or bottleneck documented) - ---- - -## Summary of Files Created/Modified - -``` -python/sglang/jit_kernel/diffusion/triton/.py # NEW: Triton kernel + launcher -python/sglang/jit_kernel/tests/test_.py # NEW: correctness tests -python/sglang/jit_kernel/benchmark/bench_.py # NEW: performance benchmark -python/sglang/multimodal_gen/runtime/layers/layernorm.py # MODIFIED: integrate into layer - (or elementwise.py, depending on op type) -``` - -## References - -- `python/sglang/jit_kernel/diffusion/triton/scale_shift.py` — 2D tile pattern, scalar broadcast, 4D shape handling -- `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 -- `../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 diff --git a/python/sglang/multimodal_gen/runtime/utils/profiler.py b/python/sglang/multimodal_gen/runtime/utils/profiler.py index 19a0ebfd8..f9eebc2c4 100644 --- a/python/sglang/multimodal_gen/runtime/utils/profiler.py +++ b/python/sglang/multimodal_gen/runtime/utils/profiler.py @@ -18,6 +18,17 @@ if current_platform.is_npu(): logger = init_logger(__name__) +def _resolve_profiler_log_dir(log_dir: str | None) -> str: + if log_dir is not None: + return log_dir + + diffusion_profiler_dir = os.getenv("SGLANG_DIFFUSION_TORCH_PROFILER_DIR") + if diffusion_profiler_dir: + return diffusion_profiler_dir + + return os.getenv("SGLANG_TORCH_PROFILER_DIR", "./logs") + + class SGLDiffusionProfiler: """ A wrapper around torch.profiler to simplify usage in pipelines. @@ -43,11 +54,7 @@ class SGLDiffusionProfiler: self.rank = rank self.full_profile = full_profile - self.log_dir = ( - log_dir - if log_dir is not None - else os.getenv("SGLANG_TORCH_PROFILER_DIR", "./logs") - ) + self.log_dir = _resolve_profiler_log_dir(log_dir) try: os.makedirs(self.log_dir, exist_ok=True) diff --git a/scripts/ci/utils/diffusion/run_comparison.py b/scripts/ci/utils/diffusion/run_comparison.py index d24b64b52..7bdadfec3 100644 --- a/scripts/ci/utils/diffusion/run_comparison.py +++ b/scripts/ci/utils/diffusion/run_comparison.py @@ -317,7 +317,14 @@ def _build_sglang_payload(case: dict) -> dict: "n": 1, "response_format": "b64_json", } - for key in ("num_inference_steps", "guidance_scale", "seed", "num_frames"): + for key in ( + "num_inference_steps", + "guidance_scale", + "seed", + "num_frames", + "fps", + "negative_prompt", + ): if key in case: payload[key] = case[key] return payload @@ -447,7 +454,14 @@ def send_image_conditioned_request_sglang( "n": "1", "response_format": "b64_json", } - for key in ("num_inference_steps", "guidance_scale", "seed", "num_frames"): + for key in ( + "num_inference_steps", + "guidance_scale", + "seed", + "num_frames", + "fps", + "negative_prompt", + ): if key in case: data[key] = str(case[key]) if perf_dump_path: @@ -521,6 +535,10 @@ def send_request_vllm_omni(base_url: str, case: dict, config: dict) -> float: } if "num_frames" in case: extra_body["num_frames"] = case["num_frames"] + if "fps" in case: + extra_body["fps"] = case["fps"] + if "negative_prompt" in case: + extra_body["negative_prompt"] = case["negative_prompt"] # Build message content (text or text+image) content: list[dict] | str = case["prompt"] @@ -583,6 +601,10 @@ def send_request_lightx2v(base_url: str, case: dict, config: dict) -> float: payload["width"] = case["width"] if "guidance_scale" in case: payload["guidance_scale"] = case["guidance_scale"] + if "fps" in case: + payload["fps"] = case["fps"] + if "negative_prompt" in case: + payload["negative_prompt"] = case["negative_prompt"] # Image-conditioned: LightX2V accepts image_path (URL or local path) if case.get("reference_image"): payload["image_path"] = config.get("test_image_url", "")