[Diffusion] diffusion profile and opt skills (#19540)

2026-03-02 15:06:29 +08:00
parent 07ef5f7be1
commit e42fa009d4
3 changed files with 1552 additions and 0 deletions
--- a/.claude/skills/diffusion/add-triton-kernel.md
+++ b/.claude/skills/diffusion/add-triton-kernel.md
@@ -0,0 +1,599 @@
+---
+name: add-triton-kernel
+description: Step-by-step guide for adding a new Triton kernel to SGLang Diffusion's jit_kernel module. Use when implementing fused elementwise ops, norm variants, RoPE variants, or any other lightweight GPU kernel for diffusion models using Triton JIT. Covers kernel authoring, autotune, torch.compile compatibility, layer integration, and tests.
+---
+
+# Adding a Triton Kernel to SGLang Diffusion
+
+This guide walks through adding a Triton kernel to `python/sglang/jit_kernel/diffusion/triton/`.
+We use a fused elementwise operation as the running example: `y = x * (1 + scale) + shift` (AdaLN modulation).
+
+---
+
+## 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 use-efficient-diffusion-kernels.md)
+```
+
+New Triton kernels go into `triton/<op_name>.py`.
+
+---
+
+## Step 1: Write the Triton Kernel
+
+Create `python/sglang/jit_kernel/diffusion/triton/<op_name>.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 `_<op_name>_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_<op_name>.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__":
+    pytest.main([__file__, "-v"])
+```
+
+Run:
+
+```bash
+pytest python/sglang/jit_kernel/tests/test_<op_name>.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_<op_name>.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_<op_name>.py
+```
+
+---
+
+## Step 6: Profile with Nsight Compute (optional but recommended)
+
+After the kernel passes correctness tests, use **`nsight-profiler.md`** to measure its hardware-level efficiency. This step requires **Nsight Compute (`ncu`)** to be installed.
+
+### Dependency Check
+
+Before profiling, verify `ncu` is available:
+
+```bash
+ncu --version   # must print a version string, e.g. "NVIDIA Nsight Compute 2024.1.0"
+```
+
+If `ncu` is missing, install it via the CUDA Toolkit package or the standalone [Nsight Compute installer](https://developer.nvidia.com/nsight-compute).
+
+### Quick Kernel Profile
+
+```bash
+# Profile the Triton kernel during the benchmark script
+# --kernel-name: match the Triton-mangled name (check with nsys first if unsure)
+ncu --set full \
+    -o /tmp/triton_<op_name> \
+    python python/sglang/jit_kernel/benchmark/bench_<op_name>.py
+```
+
+To profile only the Triton kernel (skip PyTorch reference and warmup launches), add `--launch-skip N --launch-count M`:
+
+```bash
+# Skip first 2 launches (warmup), capture 3 kernel invocations
+ncu --set full --launch-skip 2 --launch-count 3 \
+    -o /tmp/triton_<op_name> \
+    python python/sglang/jit_kernel/benchmark/bench_<op_name>.py
+```
+
+### Key Metrics to Check for a Triton Kernel
+
+| Metric | Healthy Range | Action if Low |
+|--------|--------------|---------------|
+| **Achieved Occupancy** | ≥ 50% | Reduce register usage or shared memory; try smaller block sizes |
+| **Memory Throughput** | ≥ 70% of peak BW | Check for non-coalesced access (pass contiguous strides) |
+| **Compute Throughput** | ≥ 50% of peak | Increase arithmetic intensity; fuse more ops per load |
+| **Warp Efficiency (No Stall)** | ≥ 60% | Reduce branch divergence; avoid `tl.atomic_*` on hot paths |
+| **L1/L2 Hit Rate** | L2 ≥ 40% | Reorder loads for locality; check broadcast patterns |
+
+### CLI-Only Analysis Workflow
+
+```bash
+# 1. Collect profile (--csv is NOT valid here; only -o to save .ncu-rep)
+ncu --set full \
+    --launch-skip 2 --launch-count 1 \
+    -o /tmp/prof_<op_name> \
+    python python/sglang/jit_kernel/benchmark/bench_<op_name>.py
+
+# 2. Export key metrics to CSV from the saved .ncu-rep (--csv is valid here)
+ncu --import /tmp/prof_<op_name>.ncu-rep \
+    --page details --csv \
+    > /tmp/prof_<op_name>_details.csv
+
+# 3. Quick summary — top bottleneck sections
+python3 - << 'EOF'
+import csv, sys
+rows = list(csv.DictReader(open("/tmp/prof_<op_name>_details.csv")))
+# print section names and their achieved % of peak
+for r in rows:
+    name = r.get("Metric Name", "")
+    val  = r.get("Metric Value", "")
+    if any(k in name for k in ["sol", "Occupancy", "Throughput"]):
+        print(f"{name:60s} {val}")
+EOF
+```
+
+### Compare Two Kernel Versions
+
+```bash
+# Profile baseline
+ncu --set full --launch-skip 2 --launch-count 1 \
+    -o /tmp/baseline python .../bench_<op_name>.py
+
+# Profile optimized version (after your changes)
+ncu --set full --launch-skip 2 --launch-count 1 \
+    -o /tmp/optimized python .../bench_<op_name>.py
+
+# Diff (CSV, no GUI)
+ncu --import /tmp/baseline.ncu-rep \
+    --import /tmp/optimized.ncu-rep \
+    --page diff --csv > /tmp/diff_<op_name>.csv
+
+python3 -c "
+import csv
+rows = list(csv.DictReader(open('/tmp/diff_<op_name>.csv')))
+for r in rows[:30]:
+    print(r.get('Metric Name','')[:55], r.get('Baseline',''), '->', r.get('Comparison',''))
+"
+```
+
+> For a complete guide to Nsight Compute metrics, occupancy analysis, roofline model interpretation, and warp efficiency optimization, refer to **`nsight-profiler.md`** in this directory.
+
+---
+
+## 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/<op_name>.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_<op_name>.py -v`
+- [ ] Benchmark runs: `python python/sglang/jit_kernel/benchmark/bench_<op_name>.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/<op_name>.py      # NEW: Triton kernel + launcher
+python/sglang/jit_kernel/tests/test_<op_name>.py            # NEW: correctness tests
+python/sglang/jit_kernel/benchmark/bench_<op_name>.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
+- `nsight-profiler.md` — full Nsight Compute guide: occupancy analysis, roofline model, warp efficiency, kernel comparison
+- `diffusion-benchmark-and-profile.md` — how to verify the kernel's impact on denoise latency
+- `use-efficient-diffusion-kernels.md` — overview of existing fused kernel entry points
--- a/.claude/skills/diffusion/diffusion-benchmark-and-profile.md
+++ b/.claude/skills/diffusion/diffusion-benchmark-and-profile.md
@@ -0,0 +1,686 @@
+---
+name: diffusion-benchmark-and-profile
+description: End-to-end benchmark and per-layer kernel profiling guide for SGLang Diffusion models. Use when measuring SGLang Diffusion generation performance, running latency benchmarks across Qwen-Image, FLUX, Z-Image-Turbo, and Wan2.2 models, profiling DiT layer kernel breakdown with torch.profiler or nsys+gputrc2graph.py, investigating performance bottlenecks, or tracking performance regressions. Always verify output correctness before and after any optimization.
+---
+
+# SGLang Diffusion Benchmark and Profile Guide
+
+**Overview**
+This skill covers how to run end-to-end benchmarks for SGLang Diffusion (`sglang generate`) across a standard set of models, profile per-layer kernel execution inside the DiT, and use those results to continuously improve performance. These metrics collectively reflect the overall performance of the current SGLang Diffusion release.
+
+> **Correctness First**: Any performance optimization must be validated for correctness before being considered complete. Faster but incorrect output is not an improvement. Always compare generated images/videos against a reference baseline before and after any change.
+
+**Primary Metric: Denoise Latency**
+The most important latency signal is the **denoising loop latency** — the total time spent running the DiT forward pass across all inference steps. This is the dominant cost in every diffusion model and the main target for optimization. End-to-end latency (including VAE decode and text encoding) is also recorded as a secondary metric, but denoising latency is the key indicator of DiT model performance.
+
+---
+
+## Prerequisites
+
+### Tool Dependency Check
+
+Before running any benchmark or profiling command, verify that all required tools are available. Run the following check script:
+
+```bash
+#!/usr/bin/env bash
+set -euo pipefail
+PASS=0; FAIL=0
+
+check() {
+  local label=$1; shift
+  if "$@" &>/dev/null; then
+    echo "  [OK]  $label"
+    ((PASS++)) || true
+  else
+    echo "  [MISS] $label"
+    ((FAIL++)) || true
+  fi
+}
+
+echo "=== SGLang Diffusion Benchmark Prerequisites ==="
+
+# Core runtime
+check "sglang CLI"             python3 -c "import sglang"
+check "torch"                  python3 -c "import torch; assert torch.cuda.is_available()"
+check "CUDA available"         python3 -c "import torch; torch.zeros(1).cuda()"
+
+# Profiling (torch.profiler is built into torch — no extra install needed)
+check "torch.profiler"         python3 -c "import torch.profiler"
+
+# nsys (optional, for Level 2 profiling)
+check "nsys in PATH"           which nsys
+
+# gputrc2graph.py dependencies
+# The tool lives in the sglang repo at examples/profiler/nsys_profile_tools/gputrc2graph.py
+# Set SGLANG_REPO to your local sglang repo root, e.g.:
+#   export SGLANG_REPO=/workspace/sglang
+SGLANG_REPO="${SGLANG_REPO:-$(python3 -c "import sglang, os; print(os.path.abspath(os.path.join(os.path.dirname(sglang.__file__), '../../..')))" 2>/dev/null || echo "")}"
+GPUTRC="${SGLANG_REPO}/examples/profiler/nsys_profile_tools/gputrc2graph.py"
+check "gputrc2graph.py exists (set SGLANG_REPO if missing)" test -f "$GPUTRC"
+check "pandas"                 python3 -c "import pandas"
+check "regex"                  python3 -c "import regex"
+check "plotly (optional)"      python3 -c "import plotly"
+
+echo ""
+echo "Result: $PASS passed, $FAIL missing"
+if [ "$FAIL" -gt 0 ]; then
+  echo ""
+  echo "Install missing dependencies:"
+  echo "  pip install pandas regex plotly   # for gputrc2graph.py"
+  echo "  # nsys: install NVIDIA Nsight Systems from https://developer.nvidia.com/nsight-systems"
+fi
+```
+
+**Minimum required for benchmarking**: `sglang`, `torch` with CUDA.
+**Additional for Level 1 profiling**: `torch.profiler` (bundled with torch, always available).
+**Additional for Level 2 profiling**: `nsys` on PATH, `gputrc2graph.py` present, `pandas`, `regex`.
+`plotly` is only needed to generate the HTML chart; `result.csv` is generated regardless.
+
+---
+
+### Download Required Input Images
+
+Some models (image editing / image-guided video generation) require input images:
+
+```bash
+mkdir -p /workspace/gen_benchmark/figs
+cd /workspace/gen_benchmark/figs
+
+wget -O cat.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png
+wget -O astronaut.jpg https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
+```
+
+---
+
+## Standard Benchmark Model Suite
+
+All commands include `--warmup` (pre-warm torch.compile) and `--enable-torch-compile` to reflect real production deployment performance. Timing is measured after warmup completes.
+
+### 1. Qwen/Qwen-Image-2512 (Text-to-Image, single GPU)
+
+**Task**: Text-to-Image
+**Resolution**: 1024×1024, 50 steps
+
+```bash
+sglang generate \
+  --model-path=Qwen/Qwen-Image-2512 \
+  --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \
+  '--negative-prompt= ' \
+  --width=1024 \
+  --height=1024 \
+  --num-inference-steps=50 \
+  --guidance-scale=4.0 \
+  --seed=42 \
+  --save-output \
+  --enable-torch-compile \
+  --warmup \
+  --dit-cpu-offload false \
+  --text-encoder-cpu-offload false
+```
+
+**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
+
+---
+
+### 2. Qwen/Qwen-Image-Edit-2511 (Image Editing, single GPU)
+
+**Task**: Text-guided Image Editing
+**Prerequisite**: `cat.png` (see Prerequisites)
+**Resolution**: 1024×1024, 50 steps
+
+```bash
+sglang generate \
+  --model-path=Qwen/Qwen-Image-Edit-2511 \
+  '--prompt=Transform into anime style' \
+  '--negative-prompt= ' \
+  --image-path=/workspace/gen_benchmark/figs/cat.png \
+  --width=1024 \
+  --height=1024 \
+  --num-inference-steps=50 \
+  --guidance-scale=4.0 \
+  --seed=42 \
+  --save-output \
+  --enable-torch-compile \
+  --warmup \
+  --dit-cpu-offload false \
+  --text-encoder-cpu-offload false
+```
+
+**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
+
+---
+
+### 3. black-forest-labs/FLUX.1-dev (Text-to-Image, single GPU)
+
+**Task**: Text-to-Image
+**Resolution**: 1024×1024, 50 steps
+
+```bash
+sglang generate \
+  --model-path=black-forest-labs/FLUX.1-dev \
+  --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \
+  --width=1024 \
+  --height=1024 \
+  --num-inference-steps=50 \
+  --guidance-scale=4.0 \
+  --seed=42 \
+  --save-output \
+  --warmup \
+  --enable-torch-compile
+```
+
+**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
+
+---
+
+### 4. black-forest-labs/FLUX.2-dev (Text-to-Image, single GPU)
+
+**Task**: Text-to-Image
+**Resolution**: 1024×1024
+
+```bash
+sglang generate \
+  --model-path black-forest-labs/FLUX.2-dev \
+  --prompt "A Logo With Bold Large Text: SGL Diffusion" \
+  --width=1024 \
+  --height=1024 \
+  --dit-layerwise-offload false \
+  --enable-torch-compile \
+  --warmup \
+  --dit-cpu-offload false \
+  --text-encoder-cpu-offload true \
+  --vae-cpu-offload false
+```
+
+**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
+
+---
+
+### 5. Tongyi-MAI/Z-Image-Turbo (Turbo Text-to-Image, single GPU)
+
+**Task**: Text-to-Image (few-step turbo mode, guidance=0)
+**Resolution**: 1024×1024, **9 steps**
+
+```bash
+sglang generate \
+  --model-path=Tongyi-MAI/Z-Image-Turbo \
+  --log-level=info \
+  --prompt='A fantasy landscape with mountains and a river, detailed, vibrant colors' \
+  --width=1024 \
+  --height=1024 \
+  --num-inference-steps=9 \
+  --guidance-scale=0.0 \
+  --seed=42 \
+  --save-output \
+  --enable-torch-compile \
+  --warmup \
+  --dit-cpu-offload false \
+  --text-encoder-cpu-offload false
+```
+
+**Key metrics**: denoise latency (s, primary), end-to-end latency (s/image), peak GPU memory (GB)
+
+---
+
+### 6. Wan-AI/Wan2.2-T2V-A14B-Diffusers 720P (Text-to-Video, 8 GPUs)
+
+**Task**: Text-to-Video
+**Resolution**: 720P, 81 frames, 40 steps
+**Parallelism**: 8 GPUs, Ulysses degree=4, CFG parallel + layerwise offload
+
+```bash
+sglang generate \
+  --model-path=Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --log-level=info \
+  --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." \
+  --negative-prompt=" " \
+  --720p \
+  --num-inference-steps=40 \
+  --num-frames=81 \
+  --guidance-scale=5.0 \
+  --seed=42 \
+  --save-output \
+  --num-gpus=8 \
+  --enable-cfg-parallel \
+  --ulysses-degree=4 \
+  --dit-layerwise-offload true \
+  --dit-cpu-offload false \
+  --vae-cpu-offload false \
+  --text-encoder-cpu-offload true \
+  --warmup \
+  --enable-torch-compile
+```
+
+**Key metrics**: denoise latency (s, primary), total end-to-end latency (s/video), peak GPU memory per device (GB)
+
+---
+
+### 7. Wan-AI/Wan2.2-TI2V-5B-Diffusers 720P (Text-Image-to-Video, single GPU)
+
+**Task**: Text-Image-to-Video
+**Prerequisite**: `astronaut.jpg` (see Prerequisites)
+**Resolution**: 720P, 81 frames, 50 steps
+
+```bash
+sglang generate \
+  --model-path Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+  --log-level info \
+  --warmup \
+  --dit-layerwise-offload false \
+  --dit-cpu-offload false \
+  --vae-cpu-offload false \
+  --text-encoder-cpu-offload false \
+  --enable-torch-compile \
+  --prompt "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot." \
+  --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
+  --image-path=/workspace/gen_benchmark/figs/astronaut.jpg \
+  --num-frames 81 \
+  --720p \
+  --num-inference-steps 50 \
+  --guidance-scale 5.0 \
+  --seed 42 \
+  --save-output
+```
+
+**Key metrics**: denoise latency (s, primary), total end-to-end latency (s/video), peak GPU memory (GB)
+
+---
+
+## Result Recording Format
+
+Record results in the following format for cross-version comparison:
+
+Denoise latency is the primary optimization target. End-to-end latency includes text encoding and VAE decode on top of denoise.
+
+```
+| Model                             | Task       | Steps | Resolution  | Denoise (s) ★ | E2E (s) | Peak Mem (GB) | SGLang Ver | GPU |
+|-----------------------------------|------------|-------|-------------|---------------|---------|---------------|------------|-----|
+| Qwen-Image-2512                   | T2I        |  50   | 1024×1024   |               |         |               |            |     |
+| Qwen-Image-Edit-2511              | Img Edit   |  50   | 1024×1024   |               |         |               |            |     |
+| FLUX.1-dev                        | T2I        |  50   | 1024×1024   |               |         |               |            |     |
+| FLUX.2-dev                        | T2I        |   -   | 1024×1024   |               |         |               |            |     |
+| Z-Image-Turbo                     | T2I Turbo  |   9   | 1024×1024   |               |         |               |            |     |
+| Wan2.2-T2V-A14B 720P (8 GPU)      | T2V        |  40   | 720P 81fr   |               |         |               |            |     |
+| Wan2.2-TI2V-5B 720P               | TI2V       |  50   | 720P 81fr   |               |         |               |            |     |
+```
+
+★ Denoise latency = total time of the denoising loop (all DiT forward passes). Reported separately from text encoding and VAE decode.
+
+---
+
+## Performance Bottleneck Investigation Workflow
+
+### Step 1: Identify the Slow Stage
+
+Add `--log-level=info` (or `debug`) and observe the following stages in order of priority:
+
+- **Denoise loop latency** ★ — total time across all DiT forward passes; this is the primary optimization target and usually accounts for >80% of end-to-end latency
+- **Per-step DiT forward latency** — denoise latency divided by number of steps; useful for pinpointing per-step overhead
+- VAE decode latency — runs once after the denoising loop; significant for video models
+- Text encoder encoding latency — runs once before the denoising loop
+- Warmup / torch.compile compilation time — excluded from all reported latency numbers
+
+Focus optimization effort on denoise latency first. Improvements to VAE or text encoder only matter after denoise is already well-optimized.
+
+### Step 2: Cross-reference Kernel Optimization Skill
+
+After identifying the slow stage, refer to `use-efficient-diffusion-kernels.md`:
+- Slow attention → check whether attention backend is FlashAttention (FA3/FA4)
+- Slow AdaLN modulation → verify `LayerNormScaleShift` / `fuse_scale_shift_kernel` is active
+- Slow RMSNorm → verify `sgl_kernel.rmsnorm` / `fused_add_rmsnorm` is hit
+- Slow RoPE → check FlashInfer inplace RoPE or Triton RoPE fallback
+- Slow QK Norm → verify `fused_inplace_qknorm` path; confirm `head_dim` is in the supported list (`64, 128, 256, 512, 1024`)
+
+If no existing fused kernel covers the slow operation (e.g., a new elementwise fusion opportunity, a norm variant, or a custom DiT sub-op), implement a new Triton kernel using **`add-triton-kernel.md`**. That skill covers the full workflow: kernel authoring, autotune, `torch.compile` compatibility, NPU fallback, layer integration, and tests.
+
+### Step 3: Check torch.compile Coverage
+
+```bash
+TORCH_COMPILE_DEBUG=1 sglang generate ...
+```
+
+Key things to watch:
+- Dynamic shape changes trigger recompilation — fix resolution and frame count when benchmarking
+- Conditional branches containing `tensor.item()` cause graph breaks and must be rewritten
+
+### Step 4: Multi-GPU Parallel Efficiency (Wan2.2-T2V-A14B)
+
+- Verify `--ulysses-degree` evenly divides `--num-gpus`
+- Confirm `--enable-cfg-parallel` is active (requires `guidance_scale > 1`)
+- Use `torch.distributed` profiling or `nsys` to identify communication bottlenecks
+- `--dit-layerwise-offload true` introduces CPU↔GPU transfer overhead; only enable when memory-constrained
+
+### Step 5: Offload Strategy Trade-offs
+
+| Offload Flag                   | Effect                         | Cost                            |
+|--------------------------------|--------------------------------|---------------------------------|
+| `--dit-cpu-offload true`       | Move all DiT weights to CPU    | Significantly increases per-step latency |
+| `--dit-layerwise-offload true` | Load DiT layer-by-layer on demand | Moderate latency, saves GPU memory |
+| `--text-encoder-cpu-offload`   | Move text encoder to CPU       | Only affects encoding phase (once per run) |
+| `--vae-cpu-offload`            | Move VAE to CPU                | Affects decode phase only       |
+
+When establishing a GPU performance baseline, disable all offloading (`false`).
+
+---
+
+## Per-Layer Kernel Profiling
+
+When denoise latency is higher than expected, the next step is to understand **which layer and which kernel** inside the DiT is responsible. Two profiling levels are used, from coarse to fine:
+
+```
+Level 1 — torch.profiler (built-in --profile flag)  →  per-PyTorch-op breakdown, per DiT layer
+Level 2 — nsys + gputrc2graph.py                    →  CUDA kernel category breakdown, CPU vs GPU time
+```
+
+Use Level 1 to find the slow DiT layer and sub-component; use Level 2 to get the full CUDA kernel category distribution and confirm where GPU time is spent.
+
+---
+
+### Level 1: torch.profiler — Built-in `--profile` Flag
+
+SGLang Diffusion has a built-in profiler (`SGLDiffusionProfiler` in `python/sglang/multimodal_gen/runtime/utils/profiler.py`) that is activated via CLI flags. No driver script needed.
+
+**Key flags:**
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--profile` | off | Enable torch profiler for the denoising stage |
+| `--num-profiled-timesteps N` | 5 | Profile N denoising steps (use `-1` for all steps) |
+| `--profile-all-stages` | off | Also profile text encoding and VAE decode stages |
+
+**Output:** gzipped Chrome trace JSON at `$SGLANG_TORCH_PROFILER_DIR/{uuid4}-{mode}-global-rank0.trace.json.gz` (default dir: `./logs`).
+The `{uuid4}` part is a randomly generated request ID. Use `ls -t` or `ls -1` to find the latest trace file.
+
+**Example — profile 3 denoising steps of FLUX.1-dev:**
+
+```bash
+SGLANG_TORCH_PROFILER_DIR=/workspace/gen_benchmark/profiles \
+sglang generate \
+  --model-path=black-forest-labs/FLUX.1-dev \
+  --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \
+  --width=1024 --height=1024 \
+  --num-inference-steps=50 \
+  --guidance-scale=4.0 \
+  --seed=42 \
+  --enable-torch-compile \
+  --warmup \
+  --profile \
+  --num-profiled-timesteps 3
+```
+
+**Example — profile all pipeline stages (text encoder + denoise + VAE):**
+
+```bash
+SGLANG_TORCH_PROFILER_DIR=/workspace/gen_benchmark/profiles \
+sglang generate \
+  --model-path=black-forest-labs/FLUX.1-dev \
+  --prompt="A futuristic cyberpunk city at night" \
+  --width=1024 --height=1024 \
+  --num-inference-steps=5 \
+  --seed=42 \
+  --enable-torch-compile \
+  --warmup \
+  --profile \
+  --profile-all-stages
+```
+
+**Reading the trace without a UI:**
+
+The output is a `.trace.json.gz` file (Chrome trace format). Parse it with Python to get a ranked op table — no browser needed:
+
+```python
+import gzip, json, collections, glob, os
+
+# The request_id is a uuid4, so the filename is e.g.:
+#   550e8400-e29b-41d4-a716-446655440000-3_steps-global-rank0.trace.json.gz
+# Find the latest trace file automatically:
+log_dir = os.environ.get("SGLANG_TORCH_PROFILER_DIR", "./logs")
+traces = sorted(glob.glob(f"{log_dir}/*.trace.json.gz"), key=os.path.getmtime, reverse=True)
+assert traces, f"No trace files found in {log_dir}"
+trace_path = traces[0]
+print(f"Reading: {trace_path}")
+with gzip.open(trace_path, "rb") as f:
+    data = json.loads(f.read())
+
+events = data.get("traceEvents", [])
+
+# Collect CUDA kernel durations per op name
+cuda_ops = collections.defaultdict(lambda: {"total_us": 0, "count": 0})
+for e in events:
+    if e.get("cat") in ("kernel", "gpu_memcpy") and "dur" in e:
+        name = e.get("name", "unknown")
+        cuda_ops[name]["total_us"] += e["dur"]
+        cuda_ops[name]["count"] += 1
+
+# Print top 40 by total CUDA time
+sorted_ops = sorted(cuda_ops.items(), key=lambda x: x[1]["total_us"], reverse=True)
+print(f"{'Kernel Name':<80} {'Total (ms)':>12} {'Count':>8}")
+print("-" * 102)
+for name, stats in sorted_ops[:40]:
+    print(f"{name:<80} {stats['total_us']/1000:>12.3f} {stats['count']:>8}")
+```
+
+**Add `record_function` scopes for per-layer attribution:**
+
+The built-in profiler captures all PyTorch ops, but adding named scopes makes it easy to attribute costs to specific DiT layers. Add these directly in the DiT block `forward()`:
+
+```python
+import torch
+
+# Inside DiT transformer block forward() — e.g., FluxTransformerBlock, QwenDiTBlock
+with torch.profiler.record_function(f"dit_block_{block_idx}.norm"):
+    x = self.norm1(x, scale, shift)
+with torch.profiler.record_function(f"dit_block_{block_idx}.attn"):
+    x = self.attn(x)
+with torch.profiler.record_function(f"dit_block_{block_idx}.mlp"):
+    x = self.mlp(x)
+```
+
+These scopes appear as named spans in the trace and in `key_averages()` output.
+
+**Key ops to watch per DiT sub-component:**
+
+| Sub-component        | Expected dominant kernel                                    |
+|----------------------|-------------------------------------------------------------|
+| QKV projection       | `cutlass_gemm` / `ampere_*_gemm`                            |
+| Attention compute    | `flash_attn_fwd` / `fmha_*` (FA3/FA4)                      |
+| Output projection    | `cutlass_gemm`                                              |
+| AdaLN modulation     | `fuse_scale_shift_kernel` / `fuse_scale_shift_gate_*`       |
+| RMSNorm / LayerNorm  | `sgl_kernel_rmsnorm` / `fused_add_rmsnorm` / Triton norm    |
+| MLP fc1 / fc2        | `cutlass_gemm`                                              |
+| SiLU gate            | `vectorized_elementwise_kernel` / `silu_and_mul`            |
+| RoPE                 | `apply_rotary_embedding` (Triton) / FlashInfer inplace      |
+| QK Norm              | `fused_inplace_qknorm` (JIT) or fallback `rmsnorm`          |
+
+---
+
+### Level 2: nsys + gputrc2graph.py — CUDA Kernel Category Breakdown
+
+SGLang ships a CLI analysis tool at `examples/profiler/nsys_profile_tools/gputrc2graph.py` ([PR #9314](https://github.com/sgl-project/sglang/pull/9314)). It processes `.nsys-rep` files and outputs:
+- `result.csv` — kernel-to-category mapping with elapsed time per category (read directly on server)
+- `result.html` — stacked bar chart (for environments with a browser)
+
+**Step 1: Install dependencies**
+
+```bash
+pip install pandas regex plotly
+```
+
+**Step 2: Collect nsys trace**
+
+Run two passes: one with profiling (to get the trace) and one without (to measure the true runtime for CPU time calculation):
+
+```bash
+# Pass A — with profiling (skip warmup using --delay)
+# Estimate DELAY as: warmup time + a few seconds buffer
+nsys profile \
+  -t cuda \
+  -o /workspace/gen_benchmark/profiles/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, neon lights reflecting on wet streets, highly detailed, 8k" \
+    --width=1024 --height=1024 \
+    --num-inference-steps=50 \
+    --guidance-scale=4.0 \
+    --seed=42 \
+    --enable-torch-compile \
+    --warmup
+
+# Pass B — without profiling, record total wall-clock time
+time sglang generate \
+  --model-path=black-forest-labs/FLUX.1-dev \
+  --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets, highly detailed, 8k" \
+  --width=1024 --height=1024 --num-inference-steps=50 \
+  --guidance-scale=4.0 --seed=42 \
+  --enable-torch-compile --warmup
+# Record the elapsed seconds as ELAPSED_SEC
+```
+
+**Step 3: Add a diffusion kernel classification JSON**
+
+`gputrc2graph.py` loads all `.json` files in its directory. Create a diffusion-specific one 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"
+    }
+  }
+}
+```
+
+**Step 4: Run analysis (CLI-only, no UI)**
+
+```bash
+cd examples/profiler/nsys_profile_tools
+
+python3 gputrc2graph.py \
+  --in_file /workspace/gen_benchmark/profiles/flux_dev.nsys-rep,sglang,diffusion,ELAPSED_SEC \
+  --out_dir /workspace/gen_benchmark/profiles/analysis \
+  --title "FLUX.1-dev denoise kernel breakdown"
+```
+
+Replace `ELAPSED_SEC` with the wall-clock seconds measured in Pass B (e.g. `132` if `time sglang generate ...` reported `real 2m12s`). The tool uses this value to compute CPU (non-GPU) idle time = `ELAPSED_SEC - total_GPU_sec`. Passing `0` is also valid but will use the nsys-measured elapsed time, which may inflate non-GPU time if profiling overhead is significant.
+
+**Step 5: Read results without a UI**
+
+```bash
+# Read the CSV directly — sorted by elapsed time, largest first
+cat /workspace/gen_benchmark/profiles/analysis/result.csv | column -t -s,
+
+# Or with Python for a clean summary
+python3 - << 'EOF'
+import pandas as pd
+df = pd.read_csv("/workspace/gen_benchmark/profiles/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 to look for in the category breakdown:**
+
+| Category shows high time | Investigation |
+|--------------------------|---------------|
+| `gemm` dominant          | Check QKV / output / MLP projections; verify tensor parallelism is active |
+| `attn` dominant          | Verify FA3/FA4 is active; check sequence length and head_dim |
+| `adaln_modulation` unexpectedly high | Verify fused `fuse_scale_shift_kernel` path is used |
+| `norm` high              | Verify `sgl_kernel_rmsnorm` / CuTe DSL fused path; check D alignment |
+| `nccl_comm` high         | Multi-GPU: check Ulysses degree; consider reducing TP degree |
+| `triton_kernel` high     | Identify which Triton kernel; check if a CuTe DSL or sgl-kernel replacement exists |
+| `non-gpu-H_D_memops` high | CPU↔GPU copy detected; check for accidental offload or `.cpu()` calls mid-denoising |
+| `CPU(non-GPU)` high      | Python dispatch overhead; check for graph breaks in torch.compile |
+
+**Comparing two versions:**
+
+The 4th field in `--in_file` is `elapsed_nonprofiled_sec` — the **total wall-clock time (seconds) of the same run without profiling** (measured via `time sglang generate ...` in Pass B above). The tool uses this to calculate how much time was spent on CPU (non-GPU) work. Replace `120` and `118` with the actual measured seconds for each run:
+
+```bash
+python3 gputrc2graph.py \
+  --in_file \
+    before.nsys-rep,sglang,diffusion,120 \   # replace 120 with Pass B elapsed seconds for "before"
+    after.nsys-rep,sglang,diffusion,118 \    # replace 118 with Pass B elapsed seconds for "after"
+  --out_dir /workspace/gen_benchmark/profiles/compare \
+  --title "FLUX.1-dev before vs after optimization"
+
+# Then read the comparison CSV
+python3 - << 'EOF'
+import pandas as pd
+df = pd.read_csv("/workspace/gen_benchmark/profiles/compare/result.csv")
+pivot = df.pivot_table(values="Elapsed Time (sec)", index="Category",
+                       columns="Model_Engine", aggfunc="sum").round(3)
+pivot["delta_s"] = pivot.iloc[:, 1] - pivot.iloc[:, 0]
+pivot["delta_%"] = (pivot["delta_s"] / pivot.iloc[:, 0] * 100).round(1)
+print(pivot.sort_values("delta_s"))
+EOF
+```
+
+---
+
+### Profiling Workflow Summary
+
+```
+0. CORRECTNESS BASELINE
+   sglang generate --seed=42 --save-output → save reference images/videos
+   (do this before any optimization; compare against it after every change)
+       ↓
+1. Run benchmark → identify slow model (denoise latency > baseline)
+       ↓
+2. sglang generate --profile --num-profiled-timesteps 3
+   → parse .trace.json.gz → rank ops by self_cuda_time_total
+   → identify slow DiT layer / sub-component (norm / attn / mlp / rope)
+       ↓
+3. nsys profile + gputrc2graph.py → result.csv
+   → read category breakdown (gemm / attn / adaln / norm / nccl / cpu-overhead)
+   → confirm where GPU time is concentrated
+       ↓
+4. Cross-reference use-efficient-diffusion-kernels.md
+   → apply existing fused kernel if available
+   → if no existing kernel covers the case, use add-triton-kernel.md
+     to implement and integrate a new Triton kernel
+       ↓
+5. VERIFY CORRECTNESS FIRST
+   sglang generate --seed=42 --save-output → diff against reference baseline
+   If output differs beyond tolerance → reject the optimization regardless of speedup
+       ↓
+6. Re-run benchmark → verify denoise latency improvement; no regression elsewhere
+```
+
+---
+
+## Continuous Improvement Checklist
+
+Before merging any PR that affects diffusion performance, run the full benchmark suite and compare.
+
+> **Rule: correctness gates performance.** A PR that improves latency but changes output is not acceptable. Correctness checks must pass before performance numbers are even considered.
+
+### Correctness Checks (must pass first)
+
+- [ ] Reference outputs collected with `--seed=42 --save-output` **before** any change
+- [ ] After change: regenerate with identical args and compare against reference
+- [ ] No visible quality degradation in generated images / videos
+- [ ] For numerical changes: pixel-level diff or PSNR/SSIM within agreed tolerance
+- [ ] Correctness verified on **all 7 benchmark models**, not just the model being optimized
+
+### Performance Checks (only after correctness passes)
+
+- [ ] All 7 model benchmarks executed; denoise latency (★), end-to-end latency, and peak memory recorded
+- [ ] No regression in denoise latency vs. previous baseline (allow ±2% variance)
+- [ ] New optimization shows measurable improvement in denoise latency on at least 2 models
+- [ ] No new graph breaks introduced (verify via `torch._dynamo` logs)
+- [ ] Results reproducible with all offloads disabled and fixed `--seed=42`
--- a/.claude/skills/diffusion/nsight-profiler.md
+++ b/.claude/skills/diffusion/nsight-profiler.md
@@ -0,0 +1,267 @@
+---
+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
+nsys export -t sqlite,json report.nsys-rep
+
+# Generate summary statistics
+nsys stats report.nsys-rep
+```
+
+### 2. Nsight Compute Profiling
+
+Detailed kernel analysis:
+
+```bash
+# Profile all kernels
+ncu -o profile ./cuda_program
+
+# Profile specific kernel
+ncu --kernel-name myKernel -o profile ./cuda_program
+
+# Full metric collection
+ncu --set full -o profile ./cuda_program
+
+# Roofline analysis
+ncu --set roofline -o profile ./cuda_program
+
+# Memory analysis
+ncu --section MemoryWorkloadAnalysis -o profile ./cuda_program
+
+# Compare two runs
+ncu --import baseline.ncu-rep --diff ./cuda_program
+```
+
+### 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 --set full -o baseline ./program_v1
+
+# Step 2: Profile optimized version
+ncu --set full -o optimized ./program_v2
+
+# Step 3: Generate comparison report (CLI, no GUI needed)
+# Both --import flags required; --page diff generates a side-by-side diff
+ncu --import baseline.ncu-rep \
+    --import optimized.ncu-rep \
+    --page diff --csv > comparison.csv
+```
+
+> **Note**: `ncu --diff` (the old single-flag syntax) was removed in Nsight Compute 2022.x. Always use two `--import` flags with `--page diff` for comparisons.
+
+### 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