add-sgl-kernel

---
name: add-sgl-kernel
description: Step-by-step tutorial for adding a heavyweight AOT CUDA/C++ kernel to sgl-kernel (including tests & benchmarks)
---

# Tutorial: Adding a New Kernel to `sgl-kernel` (AOT / Heavyweight)

Apply [kernel-organization](../kernel-organization/SKILL.md) for the public
operator namespace, logical grouping, lazy registry metadata, and test placement.
The implementation tutorial below does not replace that API contract.

This tutorial walks through adding a simple element-wise scale operation as an AOT kernel. We'll implement `scale(x, factor) = x * factor` to demonstrate the complete workflow.

## Goal

Add a new operation that scales each element of a tensor by a scalar factor:

- Input: tensor `x` (CUDA) and scalar `factor` (float)
- Output: `x * factor` (element-wise, in-place or into pre-allocated `out`)
- Supported dtypes: **FP16 (`torch.float16`), BF16 (`torch.bfloat16`), FP32 (`torch.float32`)**
  - Dispatched via `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` macro (defined in `python/sglang/kernels/aot/include/utils.h`)

## Two rules of thumb (must follow)

1. **Prefer `python/sglang/kernels/jit` first** when the kernel does **not** depend on CUTLASS or another large C++ project. This is the default path for lightweight kernels that benefit from rapid iteration.
2. **Prefer `sgl-kernel`** when the kernel **does** depend on CUTLASS or another large C++ project, or when it should be part of the AOT wheel / torch op registration flow.
3. **Exception**: if the dependency is `flashinfer`, or CUTLASS that is already provided through `flashinfer`, the kernel can still be implemented as `jit_kernel`.

In addition, every new kernel must ship with:

- **Tests** (pytest)
- **A benchmark script** (`marker.do_bench` from `sglang.kernels.jit.benchmark` — it works for any callable, not only JIT kernels)

---

## Repository integration map

You will typically touch these files/areas:

- Implementation: `python/sglang/kernels/aot/csrc/elementwise/scale.cu` (pick the right subdirectory)
- Public declarations: `python/sglang/kernels/aot/include/sgl_kernel_ops.h`
- Torch extension registration: `python/sglang/kernels/aot/csrc/common_extension.cc`
- Build: `python/sglang/kernels/aot/CMakeLists.txt` (`set(SOURCES ...)`)
- Python API: `python/sglang/kernels/aot/python/sgl_kernel/` and `python/sglang/kernels/aot/python/sgl_kernel/__init__.py`
- Tests: `python/sglang/kernels/aot/tests/test_scale.py`
- Benchmarks: `python/sglang/kernels/aot/benchmark/bench_scale.py`

---

## Step 1: Implement the kernel in `csrc/`

Pick the right subdirectory:

- `csrc/elementwise/` — for element-wise ops (our example)
- `csrc/gemm/`, `csrc/attention/`, `csrc/moe/` — for other categories

Create `python/sglang/kernels/aot/csrc/elementwise/scale.cu`:

```cpp
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>

#include "utils.h"  // DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16

// scale_kernel: out[i] = input[i] * factor
// Supports float, half (__half), __nv_bfloat16 via template T
template <typename T>
__global__ void scale_kernel(T* __restrict__ out,
                              const T* __restrict__ input,
                              float factor,
                              int64_t n) {
  int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
  if (idx < n) {
    out[idx] = static_cast<T>(static_cast<float>(input[idx]) * factor);
  }
}

void scale(at::Tensor& out, const at::Tensor& input, double factor) {
  TORCH_CHECK(input.is_cuda(),       "input must be a CUDA tensor");
  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
  TORCH_CHECK(out.is_cuda(),         "out must be a CUDA tensor");
  TORCH_CHECK(out.is_contiguous(),   "out must be contiguous");
  TORCH_CHECK(out.sizes() == input.sizes(),  "out and input must have the same shape");
  TORCH_CHECK(out.scalar_type() == input.scalar_type(),
              "out and input must have the same dtype");

  const int64_t n = input.numel();
  const int threads = 256;
  const int blocks  = (n + threads - 1) / threads;

  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));

  // Dispatches over float, float16, bfloat16
  DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), c_type, [&] {
    scale_kernel<c_type><<<blocks, threads, 0, stream>>>(
        static_cast<c_type*>(out.data_ptr()),
        static_cast<const c_type*>(input.data_ptr()),
        static_cast<float>(factor),
        n);
    cudaError_t status = cudaGetLastError();
    TORCH_CHECK(status == cudaSuccess,
                "scale_kernel launch failed: ", cudaGetErrorString(status));
    return true;
  });
}
```

**Key points:**

- Use `at::Tensor` (PyTorch tensors), `TORCH_CHECK` for validation, `at::cuda::getCurrentCUDAStream()` for stream
- Keep Python wrappers thin; do shape/dtype/device validation in C++ right around the launch path
- `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` covers `float`, `half` (FP16), `__nv_bfloat16` (BF16)
- Add device error checking after every kernel launch
- If a kernel only works on certain architectures, enforce that with `TORCH_CHECK` and skip logic in tests

---

## Step 2: Add a C++ declaration in `include/sgl_kernel_ops.h`

Edit `python/sglang/kernels/aot/include/sgl_kernel_ops.h`, add to the elementwise section:

```cpp
void scale(at::Tensor& out, const at::Tensor& input, double factor);
```

---

## Step 3: Register the op in `csrc/common_extension.cc`

Edit `python/sglang/kernels/aot/csrc/common_extension.cc`, inside `TORCH_LIBRARY_FRAGMENT(sgl_kernel, m)`:

```cpp
// From csrc/elementwise
m.def("scale(Tensor! out, Tensor input, float factor) -> ()");
m.impl("scale", torch::kCUDA, &scale);
```

**Key points:**

- `Tensor!` means in-place / mutable output argument
- The schema is important for `torch.compile` and for consistent call signatures
- Keep the torch schema in PyTorch scalar types (`float` here), but note that the C++ launcher signature still needs `double` for scalar arguments accepted by `torch::Library`

---

## Step 4: Add the new source file to `CMakeLists.txt`

Edit `python/sglang/kernels/aot/CMakeLists.txt`, add to `set(SOURCES ...)`:

```cmake
csrc/elementwise/scale.cu
```

**Key points:**

- Keep the list **alphabetically sorted** (the file explicitly requires this)
- If the kernel has arch constraints, reflect that in tests/benchmarks via skip logic

---

## Step 5: Expose a Python API under `python/sglang/kernels/aot/python/sgl_kernel/`

Prefer following the existing module organization first. For elementwise kernels, the usual pattern is:

- implement the Python wrapper in `python/sglang/kernels/aot/python/sgl_kernel/elementwise.py`
- then re-export it from `python/sglang/kernels/aot/python/sgl_kernel/__init__.py`

For example, in `python/sglang/kernels/aot/python/sgl_kernel/elementwise.py`, add:

```python
import torch

def scale(
    input: torch.Tensor,
    factor: float,
    out: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Element-wise scale: out = input * factor.

    Supported dtypes: torch.float16, torch.bfloat16, torch.float32.

    Parameters
    ----------
    input  : CUDA input tensor
    factor : scale factor (float)
    out    : optional pre-allocated CUDA output tensor (same shape/dtype as input)
    """
    if out is None:
        out = torch.empty_like(input)
    torch.ops.sgl_kernel.scale.default(out, input, factor)
    return out
```

Then re-export it from `python/sglang/kernels/aot/python/sgl_kernel/__init__.py` following the existing import style used by other kernels.

---

## SGLang integration (required for runtime use)

After exposing the AOT wheel symbol, add a lazy wrapper and a `KernelSpec` under
`python/sglang/kernels/ops/<group>/`. Set `backend=KernelBackend.AOT`, record
actual device/architecture support, and keep `sgl_kernel` imports inside the
implementation path. SGLang runtime and integration tests import that wrapper.
For this example, use `sglang.kernels.ops.elementwise.scale`.

The wheel-level tests below validate its standalone API/build. They do not
replace CI-registered SGLang correctness tests in
`test/registered/kernels/ops/elementwise/test_scale.py` and benchmarks in
`test/registered/kernels/benchmark/elementwise/bench_scale.py`. Follow
`write-sglang-test` for their registration and CI budget; do not register tests
under the shipped `python/sglang/` package.

## Step 6: Write tests (required)

Create `python/sglang/kernels/aot/tests/test_scale.py`:
```python
import pytest

import torch
import sgl_kernel

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("size", [128, 1024, 4096, 65536])
@pytest.mark.parametrize("factor", [0.5, 1.0, 2.0])
def test_scale_correctness(dtype, size, factor):
    input = torch.randn(size, dtype=dtype, device="cuda")
    out   = torch.empty_like(input)

    result = sgl_kernel.scale(input, factor, out=out)
    assert result is out

    expected = input * factor
    rtol, atol = (1e-5, 1e-6) if dtype == torch.float32 else (1e-2, 1e-2)
    torch.testing.assert_close(out, expected, rtol=rtol, atol=atol)


def test_scale_shape_mismatch():
    input = torch.randn(128, dtype=torch.float16, device="cuda")
    out   = torch.empty(256, dtype=torch.float16, device="cuda")
    with pytest.raises(RuntimeError, match="same shape"):
        sgl_kernel.scale(input, 2.0, out=out)


def test_scale_cpu_input():
    input = torch.randn(128, dtype=torch.float16)  # CPU
    out   = torch.empty_like(input)
    with pytest.raises(RuntimeError, match="CUDA"):
        sgl_kernel.scale(input, 2.0, out=out)


if __name__ == "__main__":
    import sys
    sys.exit(pytest.main([__file__, "-q"]))
```

---

## Step 7: Add a benchmark (required)

Every benchmark must account for L2 cache reuse — see [`rules/kernel-benchmark.md`](../../rules/kernel-benchmark.md).

Create `python/sglang/kernels/aot/benchmark/bench_scale.py`:

```python
import torch

import sgl_kernel
from sglang.kernels.jit.benchmark import marker


def sglang_scale(input: torch.Tensor, factor: float, out: torch.Tensor) -> None:
    sgl_kernel.scale(input, factor, out=out)


def torch_scale(input: torch.Tensor, factor: float, out: torch.Tensor) -> None:
    torch.mul(input, factor, out=out)


FN_MAP = {"sglang": sglang_scale, "torch": torch_scale}


@marker.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], [torch.float16])
@marker.parametrize("size", [2**n for n in range(10, 20)], [4096])  # 1K .. 512K
@marker.benchmark("provider", ["sglang", "torch"])
def benchmark(dtype: torch.dtype, size: int, provider: str):
    input = torch.randn(size, dtype=dtype, device="cuda")
    out = torch.empty_like(input)
    return marker.do_bench(
        FN_MAP[provider],
        # Pass every tensor through input_args (not a closure) so marker rotates
        # them across CUDA-graph calls to defeat L2 reuse.
        input_args=(input, 2.0, out),
        # Bandwidth = bytes(input) + bytes(out), both already in input_args.
        memory_output=None,
    )


if __name__ == "__main__":
    benchmark.run()
```

---

## Step 8: Build

Build:

```bash
cd python/sglang/kernels/aot
make build -j16
```

If you need to limit host resource usage:

```bash
cd python/sglang/kernels/aot
make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"
```

---

## Step 9: Validate

After building successfully, run the test and benchmark:

```bash
pytest python/sglang/kernels/aot/tests/test_scale.py -q
python python/sglang/kernels/aot/benchmark/bench_scale.py
```

PR CI also runs `pr-test-sgl-kernel.yml`, including the B200 job
`sgl-kernel-b200-test` when kernel changes are detected. Use that job as the
Blackwell coverage signal for AOT `sgl-kernel` changes.

---

## Troubleshooting

- **Async CUDA errors**: `CUDA_LAUNCH_BLOCKING=1`
- **Memory errors**: `compute-sanitizer --tool memcheck python ...`
- **Build is too slow / OOM**: reduce `MAX_JOBS` and `SGL_KERNEL_COMPILE_THREADS`
- **Binary bloat**: use `python/sglang/kernels/aot/analyze_whl_kernel_sizes.py`
- **CMake sources list**: if your `.cu` file is missing from `SOURCES`, the symbol will be undefined at link time

---

## References

- `python/sglang/kernels/aot/README.md`
- `python/sglang/kernels/aot/include/sgl_kernel_ops.h`
- `python/sglang/kernels/aot/csrc/common_extension.cc`
- `python/sglang/kernels/aot/CMakeLists.txt`
- `python/sglang/kernels/aot/include/utils.h` — `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` macro and friends
- `python/sglang/kernels/aot/csrc/elementwise/activation.cu` — reference for the FP16/BF16/FP32 dispatch pattern

## Summary of Files Created/Modified

```
python/sglang/kernels/aot/csrc/elementwise/scale.cu          # NEW: CUDA kernel + launcher
python/sglang/kernels/aot/include/sgl_kernel_ops.h           # MODIFIED: C++ declaration
python/sglang/kernels/aot/csrc/common_extension.cc           # MODIFIED: schema + dispatch registration
python/sglang/kernels/aot/CMakeLists.txt                     # MODIFIED: add source file (alphabetical)
python/sglang/kernels/aot/python/sgl_kernel/elementwise.py   # MODIFIED: Python wrapper
python/sglang/kernels/aot/python/sgl_kernel/__init__.py      # MODIFIED: re-export Python API
python/sglang/kernels/aot/tests/test_scale.py                # NEW: tests
python/sglang/kernels/aot/benchmark/bench_scale.py           # NEW: benchmark
```

More agent context in sgl-project/sglang

27 other files this repository gives its agents.