From 1b65c0d2598036468261647ec93135d9c2f3154e Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sun, 22 Mar 2026 15:38:22 +0800 Subject: [PATCH] [Diffusion] Fix torch.compile RMSNorm fallback for Z-Image (#20962) Co-authored-by: Mick --- .../jit_kernel/benchmark/bench_norm_impls.py | 12 +- .../diffusion/triton/rmsnorm_onepass.py | 19 +- .../.claude/skills/diffusion-kernel/SKILL.md | 11 +- .../diffusion-kernel/add-cuda-kernel.md | 87 +++------- .../diffusion-benchmark-and-profile.md | 50 ++++-- .../diffusion-kernel/nsight-profiler.md | 27 ++- .../scripts/bench_diffusion_denoise.py | 162 +++--------------- .../scripts/bench_diffusion_rmsnorm.py | 20 ++- .../use-efficient-diffusion-kernels.md | 9 +- .../skills/diffusion-optimal-perf/SKILL.md | 15 +- .../.claude/skills/support-new-model/SKILL.md | 2 +- .../runtime/layers/layernorm.py | 8 +- 12 files changed, 170 insertions(+), 252 deletions(-) diff --git a/python/sglang/jit_kernel/benchmark/bench_norm_impls.py b/python/sglang/jit_kernel/benchmark/bench_norm_impls.py index 94b029597..9642635b1 100644 --- a/python/sglang/jit_kernel/benchmark/bench_norm_impls.py +++ b/python/sglang/jit_kernel/benchmark/bench_norm_impls.py @@ -57,11 +57,11 @@ ACTUAL_DIFFUSION_GROUPS: list[ "qwen-edit", "1 GPU", [ - ("qwen_edit_ln_189x3072", "layernorm", (1, 189, 3072), SGL_LN_PAIR), - ("qwen_edit_ln_192x3072", "layernorm", (1, 192, 3072), SGL_LN_PAIR), + ("qwen_edit_ln_200x3072", "layernorm", (1, 200, 3072), SGL_LN_PAIR), + ("qwen_edit_ln_203x3072", "layernorm", (1, 203, 3072), SGL_LN_PAIR), ("qwen_edit_ln_8308x3072", "layernorm", (1, 8308, 3072), TORCH_LN), - ("qwen_edit_rms_189x3584", "rmsnorm", (1, 189, 3584), SGL_RMS), - ("qwen_edit_rms_192x3584", "rmsnorm", (1, 192, 3584), SGL_RMS), + ("qwen_edit_rms_200x3584", "rmsnorm", (1, 200, 3584), SGL_RMS), + ("qwen_edit_rms_203x3584", "rmsnorm", (1, 203, 3584), SGL_RMS), ], ), ( @@ -93,9 +93,7 @@ ACTUAL_DIFFUSION_GROUPS: list[ ("zimage_rms_32x3840", "rmsnorm", (1, 32, 3840), SGL_RMS), ("zimage_rms_4096x3840", "rmsnorm", (1, 4096, 3840), SGL_RMS), ("zimage_rms_4128x3840", "rmsnorm", (1, 4128, 3840), SGL_RMS), - ("zimage_rms_512x2560", "rmsnorm", (1, 512, 2560), SGL_RMS), - ("zimage_rms_512x32x128", "rmsnorm", (1, 512, 32, 128), SGL_RMS), - ("zimage_rms_512x8x128", "rmsnorm", (1, 512, 8, 128), SGL_RMS), + ("zimage_rms_32x2560", "rmsnorm", (32, 2560), SGL_RMS), ], ), ( diff --git a/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py index c6992ac72..963c59f55 100644 --- a/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py +++ b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py @@ -2,6 +2,8 @@ import torch import triton # type: ignore import triton.language as tl # type: ignore +from sglang.srt.utils.custom_op import register_custom_op + # Adapted from https://github.com/ModelTC/LightX2V/blob/main/lightx2v/common/ops/norm/triton_ops.py#L905-L956 @triton.jit @@ -33,7 +35,10 @@ def _rms_norm_tiled_onepass( tl.store(y_blk, x * rstd * w, mask=mask) -def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6): +@register_custom_op(op_name="triton_one_pass_rms_norm_cuda", out_shape="x") +def _triton_one_pass_rms_norm_cuda( + x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6 +) -> torch.Tensor: shape = x.shape x = x.contiguous() y = torch.empty_like(x) @@ -41,11 +46,11 @@ def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6 y_view = y.reshape(-1, shape[-1]) S, D = x_view.shape - BLOCK_SIZE_SEQ = min(16, triton.next_power_of_2(max(1, S // 512))) - grid = (triton.cdiv(S, BLOCK_SIZE_SEQ),) + block_size_seq = min(16, triton.next_power_of_2(max(1, S // 512))) + grid = (triton.cdiv(S, block_size_seq),) with torch.get_device_module().device(x.device): - torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid]( + _rms_norm_tiled_onepass[grid]( y_view, x_view, w, @@ -53,11 +58,15 @@ def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6 D, eps, BLOCK_SIZE_DIM=triton.next_power_of_2(D), - BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ, + BLOCK_SIZE_SEQ=block_size_seq, ) return y +def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6): + return _triton_one_pass_rms_norm_cuda(x, w, eps) + + from sglang.multimodal_gen.runtime.platforms import current_platform if current_platform.is_mps(): diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/SKILL.md index b70c603ad..45bd678f2 100644 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/SKILL.md @@ -32,10 +32,11 @@ python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/ ## Index -Before running any benchmark, profiler, or kernel-validation command, use -`scripts/diffusion_skill_env.py` to derive the repo root from `sglang.__file__`, -verify the repo is writable, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and -choose idle GPU(s) before starting perf work. +Before running any benchmark, profiler, or kernel-validation command: +- use `scripts/diffusion_skill_env.py` to derive the repo root from `sglang.__file__` +- verify the repo is writable +- export `FLASHINFER_DISABLE_VERSION_CHECK=1` +- choose idle GPU(s) before starting perf work - [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py) @@ -75,4 +76,4 @@ Loaded by `add-cuda-kernel.md`. Adapted from [HuggingFace kernels cuda-kernels s - [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 via `sglang generate`, baseline vs custom kernels comparison table +- [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` diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/add-cuda-kernel.md b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/add-cuda-kernel.md index 38cd3b0be..48290a7dc 100644 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/add-cuda-kernel.md +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/add-cuda-kernel.md @@ -16,7 +16,7 @@ description: Step-by-step guide for adding a new JIT CUDA kernel to SGLang Diffu > - [references/a100-optimization-guide.md](references/a100-optimization-guide.md) — A100 (sm_80) deep dive > - [references/t4-optimization-guide.md](references/t4-optimization-guide.md) — T4 (sm_75, FP16 only) deep dive > - [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark vs PyTorch -> - [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark with/without kernels +> - [scripts/bench_diffusion_denoise.py](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 @@ -224,21 +224,25 @@ from typing import TYPE_CHECKING import torch -from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args +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(dtype: torch.dtype) -> Module: - args = make_cpp_args(dtype) +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"rmsnorm<{args}>")], - extra_cuda_cflags=["-O3", "--use_fast_math"], + cuda_files=["diffusion/rmsnorm.cuh"], # relative to csrc/ + cuda_wrappers=[("rmsnorm", f"RMSNormKernel<{args}>::run")], ) @@ -253,32 +257,33 @@ def diffusion_rmsnorm( y = x / rms(x) * weight (weight=None → no affine scaling) - Supported dtypes: float16, bfloat16, float32. - hidden_size must be divisible by 8 (fp16/bf16) or 4 (fp32). + 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) - # Pass a zero-sized tensor when weight is absent (launcher checks data_ptr == nullptr) - w = weight if weight is not None else torch.empty(0, dtype=src.dtype, device=src.device) + w = weight if weight is not None else torch.ones(hidden_size, dtype=src.dtype, device=src.device) - module = _jit_rmsnorm_module(src.dtype) - module.rmsnorm(out, src, w, eps) + 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`) -- First arg(s) to `load_jit` form the unique build cache key +- 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 Denoising Stage +## 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: @@ -287,7 +292,7 @@ python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py python/sglang/multimodal_gen/runtime/models/dits/.py ``` -**Pattern — monkey-patch the DiT block's RMSNorm:** +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 @@ -387,7 +392,7 @@ if torch.cuda.get_device_capability()[0] < 9: ## Step 6: Tests -Create `python/sglang/jit_kernel/tests/test_diffusion_rmsnorm.py`: +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 @@ -422,46 +427,7 @@ if __name__ == "__main__": ## Step 7: Benchmark -Create `python/sglang/jit_kernel/benchmark/bench_diffusion_rmsnorm.py`: - -```python -import torch -import triton.testing - -from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE, run_benchmark -from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm - -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 rms_norm"], - styles=[("blue", "-"), ("red", "--")], - ylabel="us", - plot_name="diffusion-rmsnorm", - args={}, - ) -) -def benchmark(hidden: int, provider: str): - src = torch.randn(4096, hidden, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) - w = torch.ones(hidden, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) - - if provider == "jit_cuda": - fn = lambda: diffusion_rmsnorm(src, weight=w, eps=1e-6) - else: - fn = lambda: torch.nn.functional.rms_norm(src, (hidden,), w, eps=1e-6) - - return run_benchmark(fn) - - -if __name__ == "__main__": - benchmark.run(print_data=True) -``` +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`. --- @@ -505,8 +471,9 @@ python/sglang/jit_kernel/diffusion/ python/sglang/jit_kernel/tests/ └── test_diffusion_rmsnorm.py # NEW: correctness tests -python/sglang/jit_kernel/benchmark/ -└── bench_diffusion_rmsnorm.py # NEW: benchmark +python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/ +├── bench_diffusion_rmsnorm.py # Validated micro-benchmark used by this skill +└── bench_diffusion_denoise.py # Preset runner for end-to-end perf dumps ``` --- @@ -523,7 +490,7 @@ python/sglang/jit_kernel/benchmark/ | [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: `sglang generate` baseline vs custom kernels, comparison table | +| [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 diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/diffusion-benchmark-and-profile.md b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/diffusion-benchmark-and-profile.md index aa693087b..7b71ef351 100644 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/diffusion-benchmark-and-profile.md +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/diffusion-benchmark-and-profile.md @@ -6,7 +6,10 @@ description: Denoise-stage benchmark and per-layer kernel profiling guide for SG # SGLang Diffusion Benchmark and Profile Guide **Primary Metric: Denoise Latency** -The denoising loop latency — total DiT forward pass time across all inference steps — is the dominant cost (>80% of end-to-end) and the **sole optimization target** for kernel work. End-to-end latency is recorded as a secondary check only. +- 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. > **Correctness First**: Faster but incorrect output is not an improvement. Always compare generated images/videos against a reference baseline before and after any change. @@ -42,12 +45,17 @@ 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" ``` -**Minimum for benchmarking**: `sglang`, `torch` with CUDA. -**Level 1 profiling**: `torch.profiler` (bundled with torch). -**Level 2 profiling**: `nsys`, `pandas`, `plotly` + `gputrc2graph.py` from the sglang repo. -All commands below assume you are inside the configured diffusion container shell, already `cd`'d to the repo root derived from `sglang.__file__`, with `FLASHINFER_DISABLE_VERSION_CHECK=1` exported. Re-run `print-idle-gpus` before each perf command if GPU availability may have changed. Keep benchmark commands within 4 GPUs or fewer. +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 `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. Download input images required by some models: ```bash @@ -65,6 +73,8 @@ wget -O "${ASSET_DIR}/mova_single_person.jpg" \ All commands include `--warmup` and `--enable-torch-compile` for real production performance. Add `--perf-dump-path .json` for machine-readable output. +If you want a checked-in preset runner instead of copying commands manually, use `scripts/bench_diffusion_denoise.py --model --label `. It writes the same perf dump JSONs used by `compare_perf.py`. + ### Perf dump & before/after compare For every benchmark run, always write a perf dump JSON: @@ -264,6 +274,11 @@ with torch.profiler.record_function(f"dit_block_{idx}.norm"): ### Step 3: Deep CUDA Kernel Breakdown (Level 2 — nsys) +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`. + ```bash # Pass A — collect nsys trace (skip warmup with --delay) nsys profile -t cuda -o "${PROFILE_DIR}/flux_dev" -f true \ @@ -301,17 +316,24 @@ Create classification JSON at `examples/profiler/nsys_profile_tools/sglang_diffu } ``` +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. + 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" \ + --in_file "${PROFILE_DIR}/flux_dev.nsys-rep,sglang,diffusion,${ELAPSED_SEC}" \ --out_dir "${PROFILE_DIR}/analysis" \ --title "FLUX.1-dev denoise kernel breakdown" # 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() @@ -343,12 +365,14 @@ EOF - 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 --kernel-name "_fused_gated_residual_add_kernel" \ +ncu --target-processes all \ + --kernel-name "_fused_gated_residual_add_kernel" \ --launch-skip 10 --launch-count 3 \ --set full \ -o "${NCU_DIR}/gated_residual" \ @@ -358,7 +382,8 @@ ncu --kernel-name "_fused_gated_residual_add_kernel" \ --num-inference-steps=5 --seed=42 # 2. Profile all kernels in a short run (use few steps to limit time) -ncu --launch-skip 50 --launch-count 200 \ +ncu --target-processes all \ + --launch-skip 50 --launch-count 200 \ --set full \ -o "${NCU_DIR}/all_kernels" \ sglang generate \ @@ -369,7 +394,8 @@ ncu --launch-skip 50 --launch-count 200 \ # 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 --graph-profiling node \ +ncu --target-processes all \ + --graph-profiling node \ --kernel-name "_fused_gated_residual_add_kernel" \ --launch-skip 5 --launch-count 3 \ --set full \ @@ -422,12 +448,14 @@ for row in reader: ```bash # Profile baseline kernel -ncu --kernel-name "vectorized_elementwise_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 --kernel-name "_fused_gated_residual_add_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 diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/nsight-profiler.md b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/nsight-profiler.md index 0c0b9e857..5d15cb399 100644 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/nsight-profiler.md +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/nsight-profiler.md @@ -62,30 +62,38 @@ nsys export -t json report.nsys-rep 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 -o profile ./cuda_program +ncu --target-processes all -o profile ./cuda_program # Profile specific kernel -ncu --kernel-name myKernel -o profile ./cuda_program +ncu --target-processes all --kernel-name myKernel -o profile ./cuda_program # Full metric collection -ncu --set full -o profile ./cuda_program +ncu --target-processes all --set full -o profile ./cuda_program # Roofline analysis -ncu --set roofline -o profile ./cuda_program +ncu --target-processes all --set roofline -o profile ./cuda_program # Memory analysis -ncu --section MemoryWorkloadAnalysis -o profile ./cuda_program +ncu --target-processes all --section MemoryWorkloadAnalysis -o profile ./cuda_program -# Compare two runs -ncu --import baseline.ncu-rep --diff ./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: @@ -175,10 +183,10 @@ Compare kernel variants: ```bash # Step 1: Profile baseline -ncu --set full -o baseline ./program_v1 +ncu --target-processes all --set full -o baseline ./program_v1 # Step 2: Profile optimized version -ncu --set full -o optimized ./program_v2 +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. @@ -197,6 +205,7 @@ for k in sorted(set(b) | set(o)): if bv != ov: print(f'{k[:55]:<55} {bv} -> {ov}') " +``` ### 8. Performance Recommendations diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py index c95c7ef98..2d67c15e2 100755 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py @@ -1,24 +1,19 @@ """ -End-to-end denoise-stage benchmark for SGLang Diffusion with/without custom JIT CUDA kernels. +End-to-end denoise-stage benchmark presets for SGLang Diffusion. Measures denoise latency (primary metric ★) and peak GPU memory. All model configs are kept in exact sync with diffusion-benchmark-and-profile.md. -Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels - Usage: - # Baseline — single model + # Single model cd /path/to/sglang python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux - # With custom JIT CUDA kernels - python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux --custom-kernels + # Tag the run for later compare_perf.py usage + python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux --label tuned - # Side-by-side comparison - python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux --compare - - # All 10 models, comparison - python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --all --compare + # All 10 preset models + python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --all Input images required for image-guided models: ASSET_DIR=$(python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py print-assets-dir --mkdir) @@ -237,7 +232,6 @@ def required_gpus_for_model(model_key: str) -> int: def build_sglang_cmd( model_key: str, - use_custom_kernels: bool, perf_dump_path: Optional[str] = None, warmup: bool = True, torch_compile: bool = True, @@ -283,17 +277,15 @@ def build_sglang_cmd( def run_benchmark_once( model_key: str, - use_custom_kernels: bool, + label: str, output_dir: Path, warmup: bool = True, ) -> dict: """Run a single benchmark pass and return results dict.""" - label = "custom" if use_custom_kernels else "baseline" perf_path = output_dir / f"{model_key}_{label}.json" cmd = build_sglang_cmd( model_key, - use_custom_kernels=use_custom_kernels, perf_dump_path=str(perf_path), warmup=warmup, ) @@ -304,11 +296,6 @@ def run_benchmark_once( env["CUDA_VISIBLE_DEVICES"] = ",".join( str(index) for index in pick_idle_gpus(required_gpus_for_model(model_key)) ) - if use_custom_kernels: - # NOTE: This env var is a convention for user-implemented kernel injection - # logic. SGLang runtime does not read it by default — you must add handling - # in your denoising stage or model code to check this var and apply patches. - env["SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS"] = "1" print(f"\n{'=' * 64}") print(f"[{label.upper()}] {model_key}") @@ -375,127 +362,52 @@ def run_benchmark_once( def print_results_table(results: list[dict]): - """Print baseline vs custom kernel comparison table.""" + """Print a compact table for one or more benchmark runs.""" print() print("=" * 80) print("BENCHMARK RESULTS — Denoise Latency (primary metric ★)") print("(Models and params match diffusion-benchmark-and-profile.md)") print("=" * 80) - by_model: dict[str, dict] = {} - for r in results: - by_model.setdefault(r["model"], {})[r["label"]] = r - print( - f"{'Model':<16} {'Baseline(s)':>12} {'Custom(s)':>10} {'Speedup':>9} {'Peak Mem(GB)':>14}" + f"{'Model':<16} {'Label':<12} {'Denoise(s)':>12} {'E2E(s)':>10} {'Peak Mem(GB)':>14}" ) print("-" * 64) - for model_key in MODELS: # preserve order - if model_key not in by_model: - continue - runs = by_model[model_key] - base = runs.get("baseline", {}) - custom = runs.get("custom", {}) - - base_lat = base.get("denoise_latency_s") - custom_lat = custom.get("denoise_latency_s") - peak_mem = base.get("peak_memory_gb") or custom.get("peak_memory_gb") - - speedup = f"{base_lat / custom_lat:.2f}x" if base_lat and custom_lat else "n/a" - base_s = f"{base_lat:.2f}" if base_lat else "n/a" - custom_s = f"{custom_lat:.2f}" if custom_lat else "n/a" - mem_s = f"{peak_mem:.1f}" if isinstance(peak_mem, float) else "n/a" - - print(f"{model_key:<16} {base_s:>12} {custom_s:>10} {speedup:>9} {mem_s:>14}") + for result in results: + denoise_s = result.get("denoise_latency_s") + e2e_s = result.get("e2e_latency_s") + peak_mem = result.get("peak_memory_gb") + denoise_text = f"{denoise_s:.2f}" if isinstance(denoise_s, float) else "n/a" + e2e_text = f"{e2e_s:.2f}" if isinstance(e2e_s, float) else "n/a" + mem_text = f"{peak_mem:.1f}" if isinstance(peak_mem, float) else "n/a" + print( + f"{result['model']:<16} {result['label']:<12} {denoise_text:>12} {e2e_text:>10} {mem_text:>14}" + ) print("-" * 64) print() print("★ Denoise latency = total DiT forward pass time across all inference steps.") print( - " See diffusion-benchmark-and-profile.md for full Level 1/2 profiling workflow." + " Compare two runs with python/sglang/multimodal_gen/benchmarks/compare_perf.py." ) -def inject_kernels_example(): - """ - Show the kernel injection pattern used when SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS=1. - After implementing add-cuda-kernel.md, this logic lives in denoising.py or - the model's transformer.py — NOT in this script. - - Call patch_rmsnorm(dit_model) BEFORE torch.compile and BEFORE any CPU offloading. - """ - import torch.nn as nn - - try: - from sglang.jit_kernel.diffusion.rmsnorm import diffusion_rmsnorm - except ImportError: - print( - "diffusion.rmsnorm JIT kernel not available. " - "Implement add-cuda-kernel.md first." - ) - return - - def patch_rmsnorm(model: nn.Module, verbose: bool = False) -> int: - """Monkey-patch all RMSNorm variants to use the JIT CUDA kernel.""" - patched = 0 - for name, module in model.named_modules(): - if "RMSNorm" not in type(module).__name__: - continue - 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(mod, ep): - def fwd(x): - return diffusion_rmsnorm(x, weight=mod.weight, eps=ep) - - return fwd - - module.forward = _make(module, eps) - else: - - def _make_no_w(ep): - def fwd(x): - return diffusion_rmsnorm(x, weight=None, eps=ep) - - return fwd - - module.forward = _make_no_w(eps) - - patched += 1 - if verbose: - print(f" Patched: {name} (weight={has_weight})") - return patched - - return patch_rmsnorm - - def main(): parser = argparse.ArgumentParser( - description="SGLang Diffusion denoise benchmark — baseline vs JIT CUDA kernels" + description="SGLang Diffusion denoise benchmark preset runner" ) parser.add_argument( "--model", choices=list(MODELS.keys()), help="Model to benchmark (default: flux)", ) - parser.add_argument("--all", action="store_true", help="Benchmark all 7 models") + parser.add_argument("--all", action="store_true", help="Benchmark all 10 models") parser.add_argument( - "--custom-kernels", - action="store_true", - help="Run with custom JIT CUDA kernels (SGLANG_DIFFUSION_CUSTOM_CUDA_KERNELS=1)", - ) - parser.add_argument( - "--no-custom-kernels", - action="store_true", - help="Run baseline (no custom kernels)", - ) - parser.add_argument( - "--compare", - action="store_true", - help="Run both baseline and custom, print comparison table", + "--label", + type=str, + default="baseline", + help="Result label and perf dump suffix (e.g. baseline, tuned, pr20962).", ) parser.add_argument( "--output-dir", @@ -504,23 +416,9 @@ def main(): help="Directory for perf dump JSON files", ) parser.add_argument("--no-warmup", action="store_true", help="Skip warmup") - parser.add_argument( - "--show-injection-example", - action="store_true", - help="Print kernel injection pattern and exit", - ) args = parser.parse_args() - if args.show_injection_example: - patch_fn = inject_kernels_example() - if patch_fn: - print( - "patch_rmsnorm function defined. " - "Call it on the DiT model before torch.compile and CPU offloading." - ) - return - output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) warmup = not args.no_warmup @@ -529,13 +427,7 @@ def main(): results = [] for model_key in models_to_run: - if args.compare: - results.append(run_benchmark_once(model_key, False, output_dir, warmup)) - results.append(run_benchmark_once(model_key, True, output_dir, warmup)) - elif args.custom_kernels: - results.append(run_benchmark_once(model_key, True, output_dir, warmup)) - else: - results.append(run_benchmark_once(model_key, False, output_dir, warmup)) + results.append(run_benchmark_once(model_key, args.label, output_dir, warmup)) if results: print_results_table(results) diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py index 4bdadcf7c..8a68eb74c 100755 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py @@ -182,21 +182,25 @@ def run_benchmark(): jit_avg, _ = benchmark_kernel(diffusion_rmsnorm, (x, weight, 1e-6)) bandwidth_gbps = (total_bytes / 1e9) / (jit_avg / 1000) - theoretical_bw = { - (9, 0): 3350, # H100: 3.35 TB/s - (8, 0): 2000, # A100 80GB - }.get( - cap, 320 - ) # T4: 320 GB/s + 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 ({torch.cuda.get_device_name(0)}): {theoretical_bw} GB/s") + print(f" Theoretical ({device_name}): {theoretical_bw} GB/s") print(f" Bandwidth efficiency: {efficiency:.1f}%") print() - print("Target: ≥ 30% efficiency (H100/A100), ≥ 40% (T4)") + print("Target: ≥ 30% efficiency (H200/H100/A100), ≥ 40% (T4)") if __name__ == "__main__": diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/use-efficient-diffusion-kernels.md b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/use-efficient-diffusion-kernels.md index c77216ce4..9e6f4402f 100644 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/use-efficient-diffusion-kernels.md +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/use-efficient-diffusion-kernels.md @@ -50,6 +50,7 @@ This skill focuses on SGLang Diffusion (`sglang.multimodal_gen`) kernel fusion p - 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 - Kernel: `apply_rotary_embedding` @@ -62,7 +63,11 @@ This skill focuses on SGLang Diffusion (`sglang.multimodal_gen`) kernel fusion p 1. sgl-kernel RMSNorm and fused add RMSNorm - Location: `layernorm.py` -- Behavior: CUDA uses `sgl_kernel.fused_add_rmsnorm` and `sgl_kernel.rmsnorm`. `hidden_size <= 128` uses Triton one-pass. ROCm falls back to native. +- Behavior: +- Standard `bf16`/`fp16` CUDA paths use `sgl_kernel.fused_add_rmsnorm` and `sgl_kernel.rmsnorm`. +- The Z-Image `fp32` `32x2560` path under `torch.compile` avoids `wrap_triton` and uses the native fp32 path. +- `hidden_size <= 128` uses Triton one-pass. +- ROCm falls back to native. 2. Attention backend selection (FlashAttention, Sage, SDPA) - Locations: `platforms/cuda.py`, `attention/selector.py`, `docs/diffusion/performance/attention_backends.md` @@ -109,4 +114,4 @@ This skill focuses on SGLang Diffusion (`sglang.multimodal_gen`) kernel fusion p - Keep CuTe compile cache keys aligned to `(dtype, ndim, D)`. - Avoid implicit broadcasts that force hidden `contiguous()` copies. - Preserve NPU and ROCm fallback paths. -- **Always verify with ncu** (`ncu --set full`) that the kernel achieves adequate memory bandwidth utilization (>70% of peak for bandwidth-bound ops) and occupancy (>50%). See `diffusion-benchmark-and-profile.md` Step 3.5 for the ncu workflow. +- **Always verify with ncu** (`ncu --set full`) and compare against both the unfused baseline and the hardware roofline. Do not rely on a single universal bandwidth/occupancy threshold; the right target depends on whether the kernel is memory-bound, compute-bound, or launch-limited. See `diffusion-benchmark-and-profile.md` Step 3.5 for the ncu workflow. diff --git a/python/sglang/multimodal_gen/.claude/skills/diffusion-optimal-perf/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/diffusion-optimal-perf/SKILL.md index 917c20edf..efcb8fd24 100644 --- a/python/sglang/multimodal_gen/.claude/skills/diffusion-optimal-perf/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-optimal-perf/SKILL.md @@ -8,10 +8,9 @@ description: Guide for achieving optimal performance with SGLang-Diffusion. Cove Use this guide when a user asks how to speed up diffusion inference, reduce latency, lower VRAM usage, or tune SGLang-Diffusion for production. Before running any `sglang generate` command below inside the diffusion container: -- derive the repo root from `python3 -c "import os, sglang; print(os.path.abspath(os.path.join(os.path.dirname(sglang.__file__), '..', '..')))"` and `cd` there +- use `python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py` to derive the repo root, verify write access, and choose idle GPU(s) - export `FLASHINFER_DISABLE_VERSION_CHECK=1` -- verify the repo is writable if you expect perf dumps or outputs -- choose idle GPU(s) first; reuse `diffusion-kernel/scripts/diffusion_skill_env.py` when doing perf work +- `cd` to the repo root resolved from `sglang.__file__` Reference: [SGLang-Diffusion Advanced Optimizations Blog](https://lmsys.org/blog/2026-02-16-sglang-diffusion-advanced-optimizations/) @@ -19,7 +18,7 @@ Reference: [SGLang-Diffusion Advanced Optimizations Blog](https://lmsys.org/blog ## Section 1: Lossless Optimizations -These options **do not** affect output quality. The generated images/videos are numerically identical (or within floating-point rounding) to the baseline. +These options are intended to preserve output quality. In practice, some paths (most notably `torch.compile`) can still introduce small floating-point drift, so validate on your target model when numerical parity matters. | Option | CLI Flag / Env Var | What It Does | Speedup | Limitations / Notes | |---|---|---|---|---| @@ -27,13 +26,13 @@ These options **do not** affect output quality. The generated images/videos are | **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 with multi-GPU, always prefer `--enable-cfg-parallel` + Ulysses over pure Ulysses** — it is generally faster at the same GPU count due to better compute-to-communication ratio and elimination of sequential branch execution. | Typically 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. | -| **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 **video models** (where per-layer compute >> H2D transfer), the memcpy is completely hidden behind computation — **zero-cost offload** that saves VRAM without speed penalty ([PR #15511](https://github.com/sgl-project/sglang/pull/15511)). | Saves VRAM (40 GB → ~11 GB for Wan A14B); zero or near-zero speed cost for video models | 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. | +| **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. | +| **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. | | **CPU Offload (components)** | `--text-encoder-cpu-offload`, `--image-encoder-cpu-offload`, `--vae-cpu-offload`, `--dit-cpu-offload` | Offloads specific pipeline components to CPU when not in use. | Reduces peak VRAM | Adds H2D transfer latency when the component is needed. Auto-enabled for low-VRAM GPUs (<30 GB). **Tip:** after the first request completes, the console prints a peak VRAM analysis with suggestions on which offload flags can be safely disabled — look for the `"Components that could stay resident"` log line. | | **Pin CPU Memory** | `--pin-cpu-memory` | Uses pinned (page-locked) memory for CPU offload transfers. | Faster H2D transfers | Slightly higher host memory usage. Enabled by default; disable only as workaround for CUDA errors. | -| **Attention Backend (lossless)** | `--attention-backend fa` | Selects lossless attention kernel: `fa` (FlashAttention 2/3/4), `torch_sdpa`. FA is the fastest lossless option. | FA >> SDPA for long sequences | FA requires compatible GPU (Ampere+). `fa3`/`fa4` are aliased to `fa`. Ring attention only works with `fa` or `sage_attn`. | +| **Attention Backend (lossless)** | `--attention-backend fa` | Selects a lossless attention kernel for SGLang-native pipelines: `fa` (FlashAttention 2/3/4 alias) or `torch_sdpa`. | FA is usually faster than SDPA on long sequences | FA requires compatible GPU (Ampere+). For `--backend diffusers`, valid backend names differ; use the names documented in `docs/diffusion/performance/attention_backends.md`. | | **Parallel Folding** | *(automatic when SP > 1)* | Reuses the SP process group as TP for the T5 text encoder, so text encoding is parallelized "for free". | Faster text encoding on multi-GPU | Automatic; no user action needed. Only applies to T5-based pipelines. | --- @@ -65,7 +64,7 @@ sglang generate --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers \ --prompt "..." --save-output ``` -Note: `--dit-layerwise-offload` is enabled by default for Wan/MOVA video models and is zero-cost (H2D fully overlapped with compute). No need to disable it. +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. ### Maximum speed, image model, single GPU, lossless diff --git a/python/sglang/multimodal_gen/.claude/skills/support-new-model/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/support-new-model/SKILL.md index 0d5b0a397..f75c7a3d5 100644 --- a/python/sglang/multimodal_gen/.claude/skills/support-new-model/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/support-new-model/SKILL.md @@ -71,7 +71,7 @@ See existing Modular examples: `QwenImagePipeline` (uses `add_standard_t2i_stage ### Step 1: Obtain and Study the Reference Implementation -**Before writing any code, ask the user to provide the model's original implementation or Diffusers pipeline code.** You need the actual source code to work from — do not guess or assume the model's architecture. If the user has not provided it, request: +**Before writing any code, obtain the model's reference implementation or Diffusers pipeline code.** You need the actual source code to work from — do not guess or assume the model's architecture. If the user already gave a HuggingFace model ID or repo, inspect that yourself first. Ask the user only when the reference implementation is private, ambiguous, or otherwise unavailable. Typical sources are: - The model's Diffusers pipeline source (e.g., the `pipeline_*.py` file from the `diffusers` library or HuggingFace repo) - Or the model's official reference implementation (e.g., from the model author's GitHub repo) - Or the HuggingFace model ID so you can look up `model_index.json` and the associated pipeline class diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 66f3f1f61..175f48201 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -68,6 +68,11 @@ class RMSNorm(CustomOp): x, self.weight, bias=None, residual=residual, eps=self.variance_epsilon ) + def _forward_cuda_fp32_rmsnorm(self, x: torch.Tensor) -> torch.Tensor: + # Avoid wrap_triton in torch.compile: it specializes on a fresh + # constant_args_idx every call and eventually falls back to eager. + return self.forward_native(x) + def forward_cuda( self, x: torch.Tensor, @@ -81,7 +86,8 @@ class RMSNorm(CustomOp): residual = residual.view(-1, shape[-1]) if x.dtype == torch.float: - # fp32 + if residual is None and self.variance_size_override is None: + return self._forward_cuda_fp32_rmsnorm(x).view(shape) out = self.forward_triton(x, residual) if residual is not None: return out[0].view(shape), out[1].view(residual_shape)