[Diffusion] Fix torch.compile RMSNorm fallback for Z-Image (#20962)

Co-authored-by: Mick <mickjagger19@icloud.com>

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),
         ],
     ),
     (

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():

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`

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/<model>.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
 

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 <file>.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 <preset> --label <name>`. 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
 

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
 

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)

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__":

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.

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
 

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

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)