---
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).

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

```
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 (required for optimization work)

After correctness tests, you must use **ncu (Nsight Compute)** to validate hardware efficiency (bandwidth/throughput/occupancy/bottleneck type).

To avoid duplicating ncu CLI details across multiple skills, this skill does not repeat command flags. Follow the canonical docs:

- `diffusion-benchmark-and-profile.md` → Step 3.5 (ncu workflow, including CUDA graph profiling)
- `nsight-profiler.md` (metrics interpretation: bandwidth / occupancy / roofline / warp stalls)

---

## Common Patterns Reference

### Pattern 1: Autotune over a 2D tile (L × C)

Used in `scale_shift.py` (`fuse_scale_shift_kernel_blc_opt`):

```python
@triton.jit
def _kernel(..., BLOCK_L: tl.constexpr, BLOCK_C: tl.constexpr):
    pid_l = tl.program_id(0)
    pid_c = tl.program_id(1)
    pid_b = tl.program_id(2)
    l_offs = pid_l * BLOCK_L + tl.arange(0, BLOCK_L)
    c_offs = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
    mask = (l_offs[:, None] < L) & (c_offs[None, :] < C)
    ...

# Launch:
grid = (triton.cdiv(L, BLOCK_L), triton.cdiv(C, BLOCK_C), B)
_kernel[grid](..., BLOCK_L=block_l, BLOCK_C=block_c, num_warps=4, num_stages=2)
```

### Pattern 2: One-pass RMSNorm for small hidden size

Used in `rmsnorm_onepass.py`:

```python
@triton.jit
def _rms_norm_tiled_onepass(y_ptr, x_ptr, w_ptr,
                              SEQ: tl.constexpr, DIM: tl.constexpr, EPS: tl.constexpr,
                              BLOCK_SIZE_SEQ: tl.constexpr, BLOCK_SIZE_DIM: tl.constexpr):
    seq_blk_id = tl.program_id(0)
    seq_id     = seq_blk_id * BLOCK_SIZE_SEQ
    seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None]
    d_offset   = tl.arange(0, BLOCK_SIZE_DIM)[None, :]
    ...
    x = tl.load(x_ptr + seq_offset * DIM + d_offset, mask=..., other=0.0).to(tl.float32)
    mean_sq = tl.sum(x * x, axis=1, keep_dims=True) / DIM
    rstd    = tl.math.rsqrt(mean_sq + EPS)
    tl.store(y_ptr + ..., x * rstd * w, mask=...)

# Launch with wrap_triton for torch.compile compat:
with torch.get_device_module().device(x.device):
    torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid](
        y_view, x_view, w,
        S, D, eps,
        BLOCK_SIZE_DIM=triton.next_power_of_2(D),
        BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ,
    )
```

### Pattern 3: `tl.constexpr` boolean flags for conditional paths

Used in `norm.py` and `scale_shift.py`:

```python
@triton.jit
def _kernel(...,
            IS_RMS_NORM:   tl.constexpr,
            HAS_RESIDUAL:  tl.constexpr,
            SCALE_IS_SCALAR: tl.constexpr):
    ...
    if IS_RMS_NORM:
        var = tl.sum(x * x, axis=0) / N
    else:
        mean = tl.sum(x, axis=0) / N
        var  = tl.sum((x - mean) ** 2, axis=0) / N

    if HAS_RESIDUAL:
        x = x + tl.load(residual_ptr + ...)

    if SCALE_IS_SCALAR:
        scale_val = tl.load(scale_ptr)
        scale = tl.full([BLOCK_N], scale_val, dtype=scale_val.dtype)
    else:
        scale = tl.load(scale_ptr + col_offsets, mask=mask, other=0.0)
```

Autotune key must include these booleans so the compiler generates separate specializations.

### Pattern 4: Computing in fp32, storing in original dtype

Always up-cast to `tl.float32` for reductions and math, then down-cast before storing:

```python
x_f32    = x.to(tl.float32)
scale_f32 = scale.to(tl.float32)
y_f32    = x_f32 * (1.0 + scale_f32) + shift_f32
tl.store(y_ptr + offsets, y_f32.to(x.dtype), mask=mask)
```

---

## Checklist Before Submitting

### Prerequisites
- [ ] `ncu --version` prints a valid Nsight Compute version (required for Step 7 profiling)

### Implementation
- [ ] Kernel file at `python/sglang/jit_kernel/diffusion/triton/<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
