[Diffusion] Fix and validate diffusion skills benchmarking/profiling workflow (#20528)

2026-03-13 21:11:37 +08:00
parent b1246c50f8
commit be7a0311a0
10 changed files with 318 additions and 78 deletions
--- a/examples/profiler/nsys_profile_tools/gputrc2graph.py
+++ b/examples/profiler/nsys_profile_tools/gputrc2graph.py
@@ -25,7 +25,9 @@ def load_engine_model():
    json_files = glob.glob(os.path.join(os.path.dirname(__file__) or ".", "*.json"))
    for fname in json_files:
        with open(fname, encoding="utf-8") as f:
-            engine_model.update(json.load(f))
+            file_engine_model = json.load(f)
+        for engine, models in file_engine_model.items():
+            engine_model.setdefault(engine, {}).update(models)
    return engine_model


@@ -38,7 +40,6 @@ class GPUTrace2Graph:
        import pandas as pd  # avoid importing till needed

        self.pd = pd
-        self.pd.options.mode.copy_on_write = True

    # helper functions for generating trace->summary csvs
    def gen_nonoverlapped_sum_from_gputrace(self, in_file, out_file):
--- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/SKILL.md
+++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/SKILL.md
@@ -32,6 +32,15 @@ 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.
+
+- [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py)
+
+  Shared preflight helper for all diffusion skill commands. Use it to print the repo root, create benchmark/profile output directories, and choose idle GPUs before running `sglang generate`, torch profiler, nsys, or ncu.
+
 - [add-triton-kernel.md](./add-triton-kernel.md)

  Step-by-step guide for adding a new Triton kernel to SGLang Diffusion's `jit_kernel/diffusion/triton/` module, including authoring, autotune, `torch.compile` compatibility, integration, and tests. Use for fused elementwise ops, norm variants, RoPE variants, or when NPU/CPU fallback is needed.
@@ -64,5 +73,6 @@ Loaded by `add-cuda-kernel.md`. Adapted from [HuggingFace kernels cuda-kernels s

 ## Scripts (runnable benchmarks)

+- [scripts/diffusion_skill_env.py](scripts/diffusion_skill_env.py) — preflight helper: repo root discovery via `sglang.__file__`, write-access probe, benchmark/profile output directories, idle GPU selection
 - [scripts/bench_diffusion_rmsnorm.py](scripts/bench_diffusion_rmsnorm.py) — RMSNorm micro-benchmark: JIT CUDA vs PyTorch, correctness check, bandwidth efficiency analysis
 - [scripts/bench_diffusion_denoise.py](scripts/bench_diffusion_denoise.py) — end-to-end denoise benchmark via `sglang generate`, baseline vs custom kernels comparison table
--- 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
@@ -7,6 +7,8 @@ description: Step-by-step guide for adding a new JIT CUDA kernel to SGLang Diffu

 > **Origin**: This skill is adapted from the [HuggingFace kernels cuda-kernels skill](https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels), rewritten to follow SGLang's JIT compilation system and internal abstractions.
 >
+> **Run environment first**: before compiling, benchmarking, or profiling any kernel from this guide, use `scripts/diffusion_skill_env.py` (or the setup block in `diffusion-benchmark-and-profile.md`) to `cd` to the repo root resolved from `sglang.__file__`, verify write access, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and pick an idle GPU.
+>
 > **Extended references** (in this directory's `references/` and `scripts/`):
 > - [references/kernel-templates.md](references/kernel-templates.md) — copy-paste ready templates for element-wise, row-reduction (RMSNorm), fused AdaLN
 > - [references/troubleshooting.md](references/troubleshooting.md) — build errors, perf issues, integration pitfalls
--- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/add-triton-kernel.md
+++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/add-triton-kernel.md
@@ -8,6 +8,8 @@ description: Step-by-step guide for adding a new Triton kernel to SGLang Diffusi
 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).

+Before compiling, benchmarking, or profiling any Triton kernel from this guide, use `scripts/diffusion_skill_env.py` (or the setup block in `diffusion-benchmark-and-profile.md`) to `cd` to the repo root resolved from `sglang.__file__`, verify write access, export `FLASHINFER_DISABLE_VERSION_CHECK=1`, and choose an idle GPU.
+
 ---

 ## Directory Layout
--- 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
@@ -15,27 +15,45 @@ The denoising loop latency — total DiT forward pass time across all inference
 ## Prerequisites

 ```bash
-#!/usr/bin/env bash
-# Quick dependency check
-check() { "$@" &>/dev/null && echo "[OK]  $1" || echo "[MISS] $1"; }
-check "sglang"         python3 -c "import sglang"
-check "torch+CUDA"     python3 -c "import torch; assert torch.cuda.is_available()"
+ENV_PY=python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py
+ROOT=$(python3 "$ENV_PY" print-root)
+cd "$ROOT"
+python3 "$ENV_PY" check-write-access >/dev/null
+
+export FLASHINFER_DISABLE_VERSION_CHECK=1
+export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 1)
+
+ASSET_DIR=$(python3 "$ENV_PY" print-assets-dir --mkdir)
+BENCH_DIR=$(python3 "$ENV_PY" print-output-dir --kind benchmarks --mkdir)
+PROFILE_DIR=$(python3 "$ENV_PY" print-output-dir --kind profiles --mkdir)
+NCU_DIR=$(python3 "$ENV_PY" print-output-dir --kind ncu --mkdir)
+export PROFILE_DIR
+
+check() {
+  local label="$1"
+  shift
+  "$@" &>/dev/null && echo "[OK]  $label" || echo "[MISS] $label"
+}
+
+check "sglang" python3 -c "import sglang"
+check "torch+CUDA" python3 -c "import torch; assert torch.cuda.is_available()"
 check "torch.profiler" python3 -c "import torch.profiler"
 check "nsys (Level 2)" which nsys
-check "pandas"         python3 -c "import pandas"
-check "plotly"         python3 -c "import plotly"
+check "ncu (Level 3)" which ncu
+check "pandas" python3 -c "import pandas"
+check "plotly" python3 -c "import plotly"
 ```

 **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. For 8-GPU commands, request eight idle GPUs and export the comma-separated list to `CUDA_VISIBLE_DEVICES` first.

 Download input images required by some models:
 ```bash
-mkdir -p /workspace/gen_benchmark/figs
-wget -O /workspace/gen_benchmark/figs/cat.png \
+wget -O "${ASSET_DIR}/cat.png" \
  https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png
-wget -O /workspace/gen_benchmark/figs/astronaut.jpg \
+wget -O "${ASSET_DIR}/astronaut.jpg" \
  https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
 ```

@@ -50,19 +68,20 @@ All commands include `--warmup` and `--enable-torch-compile` for real production
 For every benchmark run, always write a perf dump JSON:

 ```bash
-sglang generate ... --warmup --perf-dump-path <result>.json
+sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/<result>.json"
 ```

 Before/after comparison (outputs a Markdown table suitable for PR descriptions):

 ```bash
 # Baseline (on main branch or before changes)
-sglang generate ... --warmup --perf-dump-path baseline.json
+sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/baseline.json"

 # New (after changes)
-sglang generate ... --warmup --perf-dump-path new.json
+sglang generate ... --warmup --perf-dump-path "${BENCH_DIR}/new.json"

-python python/sglang/multimodal_gen/benchmarks/compare_perf.py baseline.json new.json
+python3 python/sglang/multimodal_gen/benchmarks/compare_perf.py \
+  "${BENCH_DIR}/baseline.json" "${BENCH_DIR}/new.json"
 ```

 ### Qwen-Image-2512 (1024×1024, 50 steps)
@@ -81,7 +100,7 @@ sglang generate \
 sglang generate \
  --model-path=Qwen/Qwen-Image-Edit-2511 \
  '--prompt=Transform into anime style' '--negative-prompt= ' \
-  --image-path=/workspace/gen_benchmark/figs/cat.png \
+  --image-path="${ASSET_DIR}/cat.png" \
  --width=1024 --height=1024 --num-inference-steps=50 --guidance-scale=4.0 \
  --seed=42 --save-output --enable-torch-compile --warmup \
  --dit-cpu-offload false --text-encoder-cpu-offload false
@@ -118,6 +137,8 @@ sglang generate \

 ### Wan2.2-T2V-A14B 720P (8 GPUs, 81 frames, 40 steps)
 ```bash
+# Select eight idle GPUs first:
+# export CUDA_VISIBLE_DEVICES=$(python3 "$ENV_PY" print-idle-gpus --count 8)
 sglang generate \
  --model-path=Wan-AI/Wan2.2-T2V-A14B-Diffusers \
  --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." \
@@ -135,7 +156,7 @@ sglang generate \
  --model-path Wan-AI/Wan2.2-TI2V-5B-Diffusers \
  --prompt "An astronaut hatching from an egg, on the surface of the moon..." \
  --negative-prompt "Bright tones, overexposed, static, blurred details..." \
-  --image-path=/workspace/gen_benchmark/figs/astronaut.jpg \
+  --image-path="${ASSET_DIR}/astronaut.jpg" \
  --num-frames 81 --720p --num-inference-steps 50 --guidance-scale 5.0 \
  --seed 42 --save-output \
  --dit-layerwise-offload false --dit-cpu-offload false \
@@ -158,7 +179,7 @@ Add `--log-level=info` and observe:
 ### Step 2: Profile with torch.profiler (Level 1)

 ```bash
-SGLANG_TORCH_PROFILER_DIR=/workspace/profiles \
+SGLANG_TORCH_PROFILER_DIR="${PROFILE_DIR}/torch" \
 sglang generate \
  --model-path=black-forest-labs/FLUX.1-dev \
  --prompt="A futuristic cyberpunk city at night" \
@@ -211,7 +232,7 @@ with torch.profiler.record_function(f"dit_block_{idx}.norm"):

 ```bash
 # Pass A — collect nsys trace (skip warmup with --delay)
-nsys profile -t cuda -o /workspace/profiles/flux_dev -f true \
+nsys profile -t cuda -o "${PROFILE_DIR}/flux_dev" -f true \
  --trace-fork-before-exec=true --delay 120 --duration 60 \
  sglang generate \
    --model-path=black-forest-labs/FLUX.1-dev \
@@ -248,16 +269,16 @@ Create classification JSON at `examples/profiler/nsys_profile_tools/sglang_diffu

 Run analysis:
 ```bash
-cd examples/profiler/nsys_profile_tools
+cd "$ROOT/examples/profiler/nsys_profile_tools"
 python3 gputrc2graph.py \
-  --in_file /workspace/profiles/flux_dev.nsys-rep,sglang,diffusion,ELAPSED_SEC \
-  --out_dir /workspace/profiles/analysis \
+  --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 pandas as pd
-df = pd.read_csv("/workspace/profiles/analysis/result.csv")
+df = pd.read_csv(f"{os.environ['PROFILE_DIR']}/analysis/result.csv")
 summary = df.groupby("Category")["Elapsed Time (sec)"].sum().sort_values(ascending=False)
 total = summary.sum()
 for cat, sec in summary.items():
@@ -296,7 +317,7 @@ EOF
 ncu --kernel-name "_fused_gated_residual_add_kernel" \
    --launch-skip 10 --launch-count 3 \
    --set full \
-    -o /workspace/ncu_reports/gated_residual \
+    -o "${NCU_DIR}/gated_residual" \
    sglang generate \
      --model-path=black-forest-labs/FLUX.1-dev \
      --prompt="test" --width=1024 --height=1024 \
@@ -305,48 +326,49 @@ ncu --kernel-name "_fused_gated_residual_add_kernel" \
 # 2. Profile all kernels in a short run (use few steps to limit time)
 ncu --launch-skip 50 --launch-count 200 \
    --set full \
-    -o /workspace/ncu_reports/all_kernels \
+    -o "${NCU_DIR}/all_kernels" \
    sglang generate \
      --model-path=black-forest-labs/FLUX.1-dev \
      --prompt="test" --width=1024 --height=1024 \
      --num-inference-steps=3 --seed=42

-# 3. For CUDA graph mode, use --graph-profiling=node to profile inside the graph
+# 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 \
    --kernel-name "_fused_gated_residual_add_kernel" \
    --launch-skip 5 --launch-count 3 \
    --set full \
-    -o /workspace/ncu_reports/gated_residual_cudagraph \
+    -o "${NCU_DIR}/gated_residual_cudagraph" \
    sglang generate \
      --model-path=black-forest-labs/FLUX.1-dev \
      --prompt="test" --width=1024 --height=1024 \
-      --num-inference-steps=5 --seed=42 \
-      --enable-piecewise-cuda-graph
+      --num-inference-steps=5 --seed=42
 ```

 #### Reading ncu results (CLI, no GUI needed)

 ```bash
 # Summary of all profiled kernels
-ncu --import /workspace/ncu_reports/gated_residual.ncu-rep --page raw --csv 2>/dev/null | head -50
+ncu --import "${NCU_DIR}/gated_residual.ncu-rep" --page raw --csv 2>/dev/null | head -50

 # Key metrics to extract:
-ncu --import /workspace/ncu_reports/gated_residual.ncu-rep \
+ncu --import "${NCU_DIR}/gated_residual.ncu-rep" \
    --page details --csv 2>/dev/null | python3 -c "
 import csv, sys
 reader = csv.DictReader(sys.stdin)
-key_metrics = [
-    'gpu__time_duration.avg',           # Kernel duration
-    'sm__throughput.avg.pct_of_peak_sustained_elapsed',  # SM utilization
-    'dram__throughput.avg.pct_of_peak_sustained_elapsed', # DRAM bandwidth util
-    'l1tex__throughput.avg.pct_of_peak_sustained_elapsed', # L1 throughput
-    'sm__warps_active.avg.pct_of_peak_sustained_active',  # Achieved occupancy
-    'launch__occupancy_limit_registers',   # Occupancy limiter
-    'launch__occupancy_limit_shared_mem',
-]
+key_metrics = {
+    'gpu__time_duration.avg': 'Duration',
+    'sm__throughput.avg.pct_of_peak_sustained_elapsed': 'Compute (SM) Throughput',
+    'dram__throughput.avg.pct_of_peak_sustained_elapsed': 'DRAM Throughput',
+    'l1tex__throughput.avg.pct_of_peak_sustained_elapsed': 'L1/TEX Cache Throughput',
+    'sm__warps_active.avg.pct_of_peak_sustained_active': 'Achieved Occupancy',
+    'launch__occupancy_limit_registers': 'Block Limit Registers',
+    'launch__occupancy_limit_shared_mem': 'Block Limit Shared Mem',
+}
 for row in reader:
    name = row.get('Metric Name', '')
-    if any(m in name for m in key_metrics):
+    if any(alias in name or metric in name for metric, alias in key_metrics.items()):
        print(f'{name:<60} {row.get(\"Metric Value\",\"\")}')
 "
 ```
@@ -368,17 +390,17 @@ for row in reader:
 # Profile baseline kernel
 ncu --kernel-name "vectorized_elementwise_kernel" \
    --launch-skip 10 --launch-count 3 --set full \
-    -o /workspace/ncu_reports/baseline ./program
+    -o "${NCU_DIR}/baseline" ./program

 # Profile optimized kernel
 ncu --kernel-name "_fused_gated_residual_add_kernel" \
    --launch-skip 10 --launch-count 3 --set full \
-    -o /workspace/ncu_reports/optimized ./program
+    -o "${NCU_DIR}/optimized" ./program

 # Compare key metrics
 for report in baseline optimized; do
  echo "=== $report ==="
-  ncu --import /workspace/ncu_reports/${report}.ncu-rep \
+  ncu --import "${NCU_DIR}/${report}.ncu-rep" \
      --page details --csv 2>/dev/null | grep -E "time_duration|throughput.*pct|occupancy"
 done
 ```
--- 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
@@ -8,22 +8,23 @@ Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kerne

 Usage:
    # Baseline — single model
-    python scripts/bench_diffusion_denoise.py --model flux
+    cd /data/bbuf/sglang
+    python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux

    # With custom JIT CUDA kernels
-    python scripts/bench_diffusion_denoise.py --model flux --custom-kernels
+    python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux --custom-kernels

    # Side-by-side comparison
-    python scripts/bench_diffusion_denoise.py --model flux --compare
+    python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --model flux --compare

    # All 7 models, comparison
-    python scripts/bench_diffusion_denoise.py --all --compare
+    python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_denoise.py --all --compare

 Input images required for image-guided models:
-    mkdir -p /workspace/gen_benchmark/figs
-    wget -O /workspace/gen_benchmark/figs/cat.png \
+    ASSET_DIR=$(python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py print-assets-dir --mkdir)
+    wget -O "${ASSET_DIR}/cat.png" \
      https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png
-    wget -O /workspace/gen_benchmark/figs/astronaut.jpg \
+    wget -O "${ASSET_DIR}/astronaut.jpg" \
      https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
 """

@@ -31,10 +32,26 @@ import argparse
 import json
 import os
 import subprocess
+import sys
 import time
 from pathlib import Path
 from typing import Optional

+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPT_DIR))
+
+from diffusion_skill_env import (
+    ensure_dir,
+    get_assets_dir,
+    get_output_dir,
+    get_repo_root,
+    pick_idle_gpus,
+)
+
+REPO_ROOT = get_repo_root()
+ASSET_DIR = ensure_dir(get_assets_dir(REPO_ROOT))
+
 # ---------------------------------------------------------------------------
 # Model configs — kept in exact sync with diffusion-benchmark-and-profile.md
 # Each entry produces the same `sglang generate` command as shown in that doc.
@@ -57,12 +74,12 @@ MODELS = {
        ],
    },
    # 2. Qwen/Qwen-Image-Edit-2511 — Image Editing, 1024×1024, 50 steps
-    # Requires: /workspace/gen_benchmark/figs/cat.png
+    # Requires: <repo>/inputs/diffusion_benchmark/figs/cat.png
    "qwen-edit": {
        "path": "Qwen/Qwen-Image-Edit-2511",
        "prompt": "Transform into anime style",
        "negative_prompt": " ",
-        "image_path": "/workspace/gen_benchmark/figs/cat.png",
+        "image_path": str(ASSET_DIR / "cat.png"),
        "extra_args": [
            "--width=1024",
            "--height=1024",
@@ -141,12 +158,12 @@ MODELS = {
        ],
    },
    # 7. Wan-AI/Wan2.2-TI2V-5B-Diffusers — Text-Image-to-Video, 720P, 1 GPU, 81 frames, 50 steps
-    # Requires: /workspace/gen_benchmark/figs/astronaut.jpg
+    # Requires: <repo>/inputs/diffusion_benchmark/figs/astronaut.jpg
    "wan-ti2v": {
        "path": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
        "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",
+        "image_path": str(ASSET_DIR / "astronaut.jpg"),
        "extra_args": [
            "--num-frames",
            "81",
@@ -168,6 +185,10 @@ MODELS = {
 }


+def required_gpus_for_model(model_key: str) -> int:
+    return 8 if model_key == "wan-t2v" else 1
+
+
 def build_sglang_cmd(
    model_key: str,
    use_custom_kernels: bool,
@@ -230,6 +251,11 @@ def run_benchmark_once(
    )

    env = os.environ.copy()
+    env.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
+    if not env.get("CUDA_VISIBLE_DEVICES"):
+        env["CUDA_VISIBLE_DEVICES"] = ",".join(
+            str(index) for index in pick_idle_gpus(required_gpus_for_model(model_key))
+        )
    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
@@ -238,6 +264,7 @@ def run_benchmark_once(

    print(f"\n{'=' * 64}")
    print(f"[{label.upper()}] {model_key}")
+    print(f"  CUDA_VISIBLE_DEVICES={env.get('CUDA_VISIBLE_DEVICES', '<unset>')}")
    print("  " + " \\\n  ".join(cmd))
    print()

@@ -425,7 +452,7 @@ def main():
    parser.add_argument(
        "--output-dir",
        type=str,
-        default="/workspace/gen_benchmark/bench_results",
+        default=str(get_output_dir("benchmarks", REPO_ROOT)),
        help="Directory for perf dump JSON files",
    )
    parser.add_argument("--no-warmup", action="store_true", help="Skip warmup")
--- 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
@@ -8,16 +8,28 @@ Compares:
 Adapted from: https://github.com/huggingface/kernels/tree/main/skills/cuda-kernels

 Usage:
-    python scripts/bench_diffusion_rmsnorm.py
+    cd /data/bbuf/sglang
+    python3 python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/bench_diffusion_rmsnorm.py

 Requirements:
-    pip install triton  # for triton.testing timing utilities
-    # SGLang must be installed and CUDA available
+    # Run inside the configured SGLang diffusion container shell.
+    # This script auto-selects an idle GPU when CUDA_VISIBLE_DEVICES is unset.
 """

+import os
+import sys
 import time
+from pathlib import Path
 from typing import Tuple

+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPT_DIR))
+
+from diffusion_skill_env import configure_runtime_env
+
+configure_runtime_env(required_gpus=1)
+
 import torch

 # ---------------------------------------------------------------------------
@@ -75,6 +87,7 @@ def run_benchmark():
    print("=" * 72)
    print("SGLang Diffusion RMSNorm Micro-Benchmark: JIT CUDA vs PyTorch")
    print("=" * 72)
+    print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', '<unset>')}")
    print(f"Device: {torch.cuda.get_device_name(0)}")
    cap = torch.cuda.get_device_capability()
    print(f"Compute Capability: sm_{cap[0]}{cap[1]}")
--- a/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py
+++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-kernel/scripts/diffusion_skill_env.py
@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import os
+import subprocess
+from pathlib import Path
+
+OUTPUT_DIR_NAMES = {
+    "benchmarks": Path("outputs/diffusion_benchmarks"),
+    "profiles": Path("outputs/diffusion_profiles"),
+    "ncu": Path("outputs/ncu_reports"),
+}
+
+
+def get_repo_root() -> Path:
+    import sglang
+
+    return Path(sglang.__file__).resolve().parents[2]
+
+
+def get_assets_dir(repo_root: Path | None = None) -> Path:
+    root = repo_root or get_repo_root()
+    return root / "inputs" / "diffusion_benchmark" / "figs"
+
+
+def get_output_dir(name: str, repo_root: Path | None = None) -> Path:
+    if name not in OUTPUT_DIR_NAMES:
+        raise KeyError(f"Unknown output dir name: {name}")
+    root = repo_root or get_repo_root()
+    return root / OUTPUT_DIR_NAMES[name]
+
+
+def ensure_dir(path: Path) -> Path:
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def check_write_access(repo_root: Path | None = None) -> Path:
+    root = repo_root or get_repo_root()
+    probe_dir = ensure_dir(root / ".cache" / "diffusion_skill_write_test")
+    probe_file = probe_dir / "probe.txt"
+    probe_file.write_text("ok", encoding="utf-8")
+    return probe_file
+
+
+def _run_nvidia_smi(query: str) -> list[list[str]]:
+    command = [
+        "nvidia-smi",
+        f"--query-{query}",
+        "--format=csv,noheader,nounits",
+    ]
+    result = subprocess.run(command, check=True, capture_output=True, text=True)
+    rows: list[list[str]] = []
+    for raw_line in result.stdout.splitlines():
+        line = raw_line.strip()
+        if not line:
+            continue
+        rows.append([field.strip() for field in csv.reader([line]).__next__()])
+    return rows
+
+
+def get_gpu_inventory() -> list[dict[str, int | str]]:
+    rows = _run_nvidia_smi("gpu=index,uuid,memory.used,memory.total,utilization.gpu")
+    inventory = []
+    for index, uuid, memory_used, memory_total, utilization_gpu in rows:
+        inventory.append(
+            {
+                "index": int(index),
+                "uuid": uuid,
+                "memory_used_mib": int(memory_used),
+                "memory_total_mib": int(memory_total),
+                "utilization_gpu_pct": int(utilization_gpu),
+            }
+        )
+    return inventory
+
+
+def get_busy_gpu_uuids() -> set[str]:
+    rows = _run_nvidia_smi("compute-apps=gpu_uuid,pid,process_name,used_gpu_memory")
+    return {gpu_uuid for gpu_uuid, *_ in rows}
+
+
+def pick_idle_gpus(
+    required_gpus: int,
+    max_memory_used_mib: int = 32,
+    max_utilization_gpu_pct: int = 5,
+) -> list[int]:
+    inventory = get_gpu_inventory()
+    busy_uuids = get_busy_gpu_uuids()
+
+    idle = [
+        int(gpu["index"])
+        for gpu in inventory
+        if gpu["uuid"] not in busy_uuids
+        and int(gpu["memory_used_mib"]) <= max_memory_used_mib
+        and int(gpu["utilization_gpu_pct"]) <= max_utilization_gpu_pct
+    ]
+    if len(idle) < required_gpus:
+        raise RuntimeError(
+            "Not enough idle GPUs. "
+            f"required={required_gpus}, idle={idle}, inventory={inventory}, busy={sorted(busy_uuids)}"
+        )
+    return idle[:required_gpus]
+
+
+def configure_runtime_env(required_gpus: int = 1) -> str | None:
+    os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
+    if os.environ.get("CUDA_VISIBLE_DEVICES"):
+        return None
+    selected = ",".join(str(index) for index in pick_idle_gpus(required_gpus))
+    os.environ["CUDA_VISIBLE_DEVICES"] = selected
+    return selected
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Resolve SGLang diffusion skill paths and idle GPUs."
+    )
+    parser.add_argument(
+        "command",
+        choices=[
+            "print-root",
+            "print-assets-dir",
+            "print-output-dir",
+            "print-idle-gpus",
+            "check-write-access",
+        ],
+    )
+    parser.add_argument(
+        "--kind",
+        choices=sorted(OUTPUT_DIR_NAMES),
+        help="Output directory kind for print-output-dir.",
+    )
+    parser.add_argument(
+        "--count",
+        type=int,
+        default=1,
+        help="Number of idle GPUs to print.",
+    )
+    parser.add_argument(
+        "--mkdir",
+        action="store_true",
+        help="Create the requested directory before printing it.",
+    )
+    args = parser.parse_args()
+
+    if args.command == "print-root":
+        print(get_repo_root())
+        return
+    if args.command == "print-assets-dir":
+        path = get_assets_dir()
+        if args.mkdir:
+            ensure_dir(path)
+        print(path)
+        return
+    if args.command == "print-output-dir":
+        if not args.kind:
+            raise SystemExit("--kind is required for print-output-dir")
+        path = get_output_dir(args.kind)
+        if args.mkdir:
+            ensure_dir(path)
+        print(path)
+        return
+    if args.command == "print-idle-gpus":
+        print(",".join(str(index) for index in pick_idle_gpus(args.count)))
+        return
+    if args.command == "check-write-access":
+        print(check_write_access())
+        return
+    raise SystemExit(f"Unhandled command: {args.command}")
+
+
+if __name__ == "__main__":
+    main()
--- a/python/sglang/multimodal_gen/.claude/skills/diffusion-optimal-perf/SKILL.md
+++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-optimal-perf/SKILL.md
@@ -7,6 +7,12 @@ 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
+- 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
+
 Reference: [SGLang-Diffusion Advanced Optimizations Blog](https://lmsys.org/blog/2026-02-16-sglang-diffusion-advanced-optimizations/)

 ---
--- a/python/sglang/multimodal_gen/.claude/skills/diffusion-perf/SKILL.md
+++ b/python/sglang/multimodal_gen/.claude/skills/diffusion-perf/SKILL.md
@@ -1,18 +0,0 @@
---
-name: diffusion-perf
-description: Deprecated alias (merged into diffusion-kernel).
-user-invocable: false
-allowed-tools: Bash, Read
-argument-hint: <model-path> [--prompt "..."] [--baseline baseline.json]
---
-
-# Diffusion Performance Measurement
-
-This skill has been merged into the canonical docs under `diffusion-kernel`:
-
- `../diffusion-kernel/diffusion-benchmark-and-profile.md` → **Perf dump & before/after compare**
-
-Follow that document as the single source of truth:
-
- Always run `sglang generate ... --warmup --perf-dump-path <file>.json`
- Use `python python/sglang/multimodal_gen/benchmarks/compare_perf.py <baseline.json> <new.json>` to generate a PR-ready comparison table