diff --git a/.claude/skills/add-jit-kernel/SKILL.md b/.claude/skills/add-jit-kernel/SKILL.md index eed313866..3175e761c 100644 --- a/.claude/skills/add-jit-kernel/SKILL.md +++ b/.claude/skills/add-jit-kernel/SKILL.md @@ -1,276 +1,559 @@ --- name: add-jit-kernel -description: Step-by-step tutorial for adding a lightweight JIT CUDA/C++ kernel to python/sglang/jit_kernel (including tests & benchmarks) +description: Step-by-step tutorial for adding a new lightweight JIT CUDA kernel to sglang's jit_kernel module --- -# Tutorial: Adding a New Kernel to `python/sglang/jit_kernel` (JIT / Lightweight) +# Tutorial: Adding a New JIT Kernel to SGLang -This SKILL is a step-by-step guide for adding a **lightweight** CUDA/C++ kernel to `python/sglang/jit_kernel/`. - -Typical characteristics: - -- Few dependencies (usually tvm-ffi + a small subset of `sgl_kernel` utility headers) -- Compiled at runtime (JIT), optimized for fast iteration -- Avoids pulling heavyweight third-party/template code into AOT builds - -## Two rules of thumb (must follow) - -1. **Heavyweight kernels go to `sgl-kernel`.** If it depends on CUTLASS / FlashInfer / DeepGEMM (or similarly heavy stacks), implement it in `sgl-kernel/`. -2. **Lightweight kernels go to `jit_kernel`.** If it is small and can be compiled independently, implement it here. - -## Stop and use `sgl-kernel` instead (important) - -Do **not** add a new kernel under `jit_kernel` if any of the following applies: - -- It directly depends on CUTLASS / FlashInfer (or other heavyweight third-party stacks) -- It requires complex link-time integration, large template instantiations, or AOT-style packaging - -In addition, every new JIT kernel must ship with: - -- **Tests** (pytest) -- **A benchmark script** (triton.testing) - ---- +This tutorial walks through adding a simple element-wise scale operation as a JIT kernel. We'll implement `scale(x, factor) = x * factor` to demonstrate the complete workflow. ## Goal -Add a new JIT kernel end-to-end, including: +Add a new operation that scales each element of a tensor by a scalar factor: -- CUDA/C++ implementation in `jit_kernel/csrc` -- A Python wrapper that compiles + loads the JIT module via tvm-ffi -- Correctness tests -- A reproducible benchmark (with CI-friendly ranges) +- Input: tensor `x` (CUDA) and scalar `factor` (float, passed as C++ template argument) +- Output: `x * factor` (element-wise), allocated internally +- Supported dtypes: **FP16 (`torch.float16`), BF16 (`torch.bfloat16`), FP32 (`torch.float32`)** + +## When to use JIT vs AOT (`sgl-kernel`) + +- **JIT (`jit_kernel`)**: lightweight, few dependencies, rapid iteration, compiled on first use +- **AOT (`sgl-kernel`)**: depends on CUTLASS / FlashInfer / DeepGEMM, needs pre-built wheel --- -## Repository integration map +## Common Abstractions in `python/sglang/jit_kernel/include/sgl_kernel/` -You will typically touch these files/areas: +**Always prefer these abstractions over raw CUDA primitives.** They provide safety, readability, and consistency with the rest of the codebase. -- Implementation: `python/sglang/jit_kernel/csrc/` -- Reusable headers: `python/sglang/jit_kernel/include/` -- Python API: `python/sglang/jit_kernel/.py` -- JIT build + cache utilities: `python/sglang/jit_kernel/utils.py` -- Tests: `python/sglang/jit_kernel/tests/test_.py` -- Benchmarks: `python/sglang/jit_kernel/benchmark/bench_.py` -- Benchmark helpers: `python/sglang/jit_kernel/benchmark/utils.py` +### `utils.h` — Host-side utilities ---- +```cpp +#include +``` -## tvm-ffi primer (practical, as used in this repo) +- **`host::RuntimeCheck(cond, args...)`** — Assert a condition at runtime; throws `PanicError` with file/line info on failure. Prefer this over bare `assert`. +- **`host::Panic(args...)`** — Unconditionally throw a `PanicError` with a descriptive message. +- **`host::div_ceil(a, b)`** — Integer ceiling division `(a + b - 1) / b`. +- **`host::irange(n)`** / **`host::irange(start, end)`** — Range views for cleaner loops. +- **`host::pointer::offset(ptr, offsets...)`** — Byte-safe pointer arithmetic on `void*`. Use this instead of raw casts. -This repository uses tvm-ffi primarily as a **stable C++ ABI** and a set of **lightweight container types** to move data between Python and C++ with minimal overhead. +### `utils.cuh` — Device-side utilities + `LaunchKernel` -### Core types you will see in JIT kernels +```cpp +#include +``` -- `tvm::ffi::TensorView` - - A **non-owning view** of a tensor (backed by DLPack) that enables zero-copy interop. - - Use it for most tensor arguments in kernel entrypoints. - - You typically inspect/validate: - - Shape/strides: `dim()`, `shape()`, `strides()`, `size(i)`, `stride(i)`, `is_contiguous()` - - Dtype/device: `dtype()`, `device()` - - Raw pointer: `data_ptr()` (then cast after dtype checks) +- **Type aliases**: `fp16_t`, `bf16_t`, `fp32_t`, `fp8_e4m3_t`, `fp8_e5m2_t` and their packed variants `fp16x2_t`, `bf16x2_t`, `fp32x2_t`, etc. +- **`SGL_DEVICE`** — Expands to `__forceinline__ __device__`. Use on all device functions. +- **`device::kWarpThreads`** — Constant `32`. +- **`device::load_as(ptr, offset)`** / **`device::store_as(ptr, val, offset)`** — Type-safe loads/stores from `void*`. +- **`device::pointer::offset(ptr, offsets...)`** — Pointer arithmetic on device. +- **`host::LaunchKernel(grid, block, device_or_stream [, smem])`** — RAII kernel launcher that: + - Resolves the CUDA stream from a `DLDevice` via TVM-FFI automatically. + - Checks the CUDA error with file/line info after launch via `operator()(kernel, args...)`. + - Supports `.enable_pdl(bool)` for PDL (Programmatic Dependent Launch, SM90+). +- **`host::RuntimeDeviceCheck(cudaError_t)`** — Check a CUDA error; throw on failure. -- `tvm::ffi::Optional` - - Used for optional tensor arguments, e.g. `tvm::ffi::Optional`. - - Always check `has_value()` before using it. +### `tensor.h` — Tensor validation (`TensorMatcher`, Symbolic types) -### Containers you may want (even if not widely used here yet) +```cpp +#include +``` -- `tvm::ffi::Array`, `tvm::ffi::Tuple<...>` - - Useful for passing small structured metadata without inventing ad-hoc pointer conventions. +This is the **primary validation API** for all kernel launchers. Use it to validate every `tvm::ffi::TensorView` argument. -### STL support +- **`host::SymbolicSize{"name"}`** — A named symbolic dimension. Call `.set_value(n)` to pin it, `.unwrap()` to extract after verification. +- **`host::SymbolicDType`** — Symbolic dtype. Use `.set_options()` to restrict allowed types. +- **`host::SymbolicDevice`** — Symbolic device. Use `.set_options()` to restrict to CUDA. +- **`host::TensorMatcher({dims...})`** — Fluent builder for tensor validation: + - `.with_dtype()` — require a specific C++ type (e.g. `fp16_t`) + - `.with_dtype()` — allow a set of types + - `.with_device(device_sym)` — require CUDA, bind device to symbol + - `.with_strides({strides...})` — validate strides (omit to require contiguous) + - `.verify(tensor_view)` — execute the check; throws `PanicError` with full context on failure; **chainable** (`verify(a).verify(b)` to check multiple tensors with the same shape) -tvm-ffi has optional headers to interop with parts of the C++ standard library (review mentions `extra/stl.h`). This repo currently mostly relies on `TensorView` + `Optional` for kernel interfaces. +**Typical pattern:** +```cpp +auto N = SymbolicSize{"num_elements"}; +auto device = SymbolicDevice{}; +device.set_options(); +TensorMatcher({N}) // + .with_dtype() + .with_device(device) + .verify(dst) + .verify(src); // same shape, dtype, device as dst +const size_t n = N.unwrap(); +const DLDevice dev = device.unwrap(); +``` -### Source of truth in `sglang` +### `type.cuh` — `dtype_trait` and `packed_t` -The most reliable documentation for how tvm-ffi is used in `sglang` is the code under: +```cpp +#include +``` -- `python/sglang/jit_kernel/include/` +- **`dtype_trait`** — Static trait struct for each scalar type. Provides: + - `dtype_trait::from(value)` — convert from another type (e.g. `fp32_t` → `fp16_t`) + - `dtype_trait::abs/sqrt/rsqrt/max/min(x)` — type-dispatched math (for `fp32_t`) +- **`packed_t`** — Two-element packed alias: `packed_t` = `fp16x2_t`, `packed_t` = `bf16x2_t`, `packed_t` = `fp32x2_t`. Use for vectorized loads/stores. +- **`device::cast(value)`** — Type-safe cast using `dtype_trait`, e.g. `cast(v)`. -In particular: +### `vec.cuh` — Vectorized memory access (`AlignedVector`) -- `python/sglang/jit_kernel/include/sgl_kernel/tensor.h` - - `host::TensorMatcher` for validating shapes/strides/dtypes/devices - - Symbolic helper types used across many kernels: - - `host::SymbolicSize` / `host::SymbolicDType` / `host::SymbolicDevice` - - Typical pattern: declare symbols, validate with `TensorMatcher(...).verify(...)`, then `unwrap()` the resolved values for launch configuration +```cpp +#include +``` + +- **`device::AlignedVector`** — Aligned storage for N elements of type T. N must be a power of two, `sizeof(T)*N <= 32`. Enables 128-bit vector loads/stores for bandwidth efficiency. + - `.load(ptr, offset)` — vectorized load from `ptr[offset]` + - `.store(ptr, offset)` — vectorized store to `ptr[offset]` + - `.fill(value)` — fill all lanes + - `operator[](i)` — element access + +### `tile.cuh` — `tile::Memory` (strided memory access pattern) + +```cpp +#include +``` + +- **`device::tile::Memory::cta(blockDim.x)`** — Creates a tile accessor where each thread handles `tid = threadIdx.x` with stride `blockDim.x`. Common for loops over a 1D array. +- **`.load(ptr, offset)`** — loads `ptr[tid + offset * blockDim.x]` +- **`.store(ptr, val, offset)`** — stores to `ptr[tid + offset * blockDim.x]` +- **`.in_bound(n, offset)`** — boundary check + +### `math.cuh` — Device math (`device::math::`) + +```cpp +#include +``` + +- `device::math::max/min/abs/sqrt/rsqrt(a, b)` — type-dispatched math via `dtype_trait` +- `device::math::exp/sin/cos(float)` — fast float math wrappers + +### `warp.cuh` — Warp-level primitives + +```cpp +#include +``` + +- `device::warp::reduce_sum(value)` — warp-level sum reduction via `__shfl_xor_sync` +- `device::warp::reduce_max(value)` — warp-level max reduction + +### `cta.cuh` — CTA-level primitives + +```cpp +#include +``` + +- `device::cta::reduce_max(value, smem, min_value)` — CTA-wide max using shared memory + warp reduction. Caller is responsible for a `__syncthreads()` after if the result in `smem[0]` is needed. + +### `atomic.cuh` — Atomic operations + +```cpp +#include +``` + +- `device::atomic::max(float* addr, float value)` — float atomic max (handles negative values correctly via bit tricks). + +### `runtime.cuh` — Occupancy and device info + +```cpp +#include +``` + +- `host::runtime::get_blocks_per_sm(kernel, block_dim)` — max active blocks per SM (occupancy) +- `host::runtime::get_sm_count(device_id)` — number of SMs on the device +- `host::runtime::get_cc_major(device_id)` — compute capability major version + +**Persistent kernel pattern** (cap blocks to SM count × occupancy): +```cpp +static const uint32_t max_occ = runtime::get_blocks_per_sm(kernel, kBlockSize); +static const uint32_t num_sm = runtime::get_sm_count(device.unwrap().device_id); +const auto num_blocks = std::min(num_sm * max_occ, div_ceil(n, kBlockSize)); +LaunchKernel(num_blocks, kBlockSize, device.unwrap())(kernel, params); +``` --- ## Step 0 (optional): Generate a `.clangd` config for better IDE support -Because JIT kernels compile at runtime, there is no static `compile_commands.json`. - -Run from your working directory (typically the repository root): - ```bash python -m sglang.jit_kernel ``` -This will generate a `.clangd` file (and will not overwrite an existing one). - --- -## Step 1: Implement the CUDA/C++ kernel in `jit_kernel/csrc/` +## Step 1: Implement the CUDA kernel in `jit_kernel/csrc/` -1. Create a new source file: +Create `python/sglang/jit_kernel/csrc/elementwise/scale.cuh`. -- `python/sglang/jit_kernel/csrc/.cuh` (common pattern) +The implementation fully uses the project abstractions described above: -2. Use the project’s recommended utilities. +```cpp +#include // TensorMatcher, SymbolicSize, SymbolicDevice +#include // dtype_trait, fp16_t, bf16_t, fp32_t +#include // RuntimeCheck, div_ceil +#include // LaunchKernel, SGL_DEVICE +#include // AlignedVector -Notes: +#include +#include -- Prefer reading and reusing the actual helper code in `python/sglang/jit_kernel/include/`. -- If you find a missing helper that would be reusable across kernels, add it under `python/sglang/jit_kernel/include/`. +namespace { -- Use `tvm::ffi::TensorView` for tensor arguments (PyTorch tensors are passed through tvm-ffi) -- Validate inputs with `TensorMatcher` (shape/stride/dtype/device) -- Use `RuntimeCheck` / `RuntimeDeviceCheck` for readable runtime validation -- Launch kernels via `LaunchKernel` (stream/device resolution) +// ---------------------------------------------------------------- +// Kernel: element-wise scale using vectorized 128-bit loads/stores +// T = fp16_t | bf16_t | fp32_t +// kVecN = number of elements per vector load (e.g. 8 for fp16) +// kFactor = scale factor encoded as kFactorNumer / kFactorDenom +// ---------------------------------------------------------------- +template +__global__ void scale_kernel(T* __restrict__ dst, + const T* __restrict__ src, + uint32_t n_vecs, + uint32_t n_remainder, + uint32_t n_total) { + constexpr float kFactor = static_cast(kFactorNumer) + / static_cast(kFactorDenom); + + using vec_t = device::AlignedVector; + + // --- vectorised body --- + const uint32_t vec_stride = blockDim.x * gridDim.x; + for (uint32_t vi = blockIdx.x * blockDim.x + threadIdx.x; + vi < n_vecs; + vi += vec_stride) { + vec_t v; + v.load(src, vi); +#pragma unroll + for (int i = 0; i < kVecN; ++i) { + v[i] = static_cast(static_cast(v[i]) * kFactor); + } + v.store(dst, vi); + } + + // --- scalar tail --- + const uint32_t base = n_vecs * kVecN; + const uint32_t scalar_stride = blockDim.x * gridDim.x; + for (uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; + i < n_remainder; + i += scalar_stride) { + dst[base + i] = static_cast(static_cast(src[base + i]) * kFactor); + } +} + +// ---------------------------------------------------------------- +// Launcher: validates tensors, selects vector width, launches kernel +// ---------------------------------------------------------------- +template +void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { + using namespace host; + + // 1. Validate input tensors with TensorMatcher + SymbolicSize N = {"num_elements"}; + SymbolicDevice device_; + device_.set_options(); + + TensorMatcher({N}) // + .with_dtype() + .with_device(device_) + .verify(dst) + .verify(src); // same shape / dtype / device as dst + + const uint32_t n = static_cast(N.unwrap()); + const DLDevice device = device_.unwrap(); + + RuntimeCheck(n > 0, "scale: num_elements must be > 0, got ", n); + + // 2. Choose vector width for 128-bit loads (16 bytes) + // fp16/bf16: 8 elements × 2 bytes = 16 bytes + // fp32: 4 elements × 4 bytes = 16 bytes + constexpr int kVecN = 16 / sizeof(T); + const uint32_t n_vecs = n / kVecN; + const uint32_t n_remainder = n % kVecN; + + // 3. Launch + constexpr uint32_t kBlockSize = 256; + const uint32_t grid = div_ceil(std::max(n_vecs, n_remainder), kBlockSize); + + LaunchKernel(grid, kBlockSize, device)( + scale_kernel, + static_cast(dst.data_ptr()), + static_cast(src.data_ptr()), + n_vecs, + n_remainder, + n); +} + +} // namespace +``` **Key points:** -- Be explicit about contiguity/stride assumptions. -- Make failures readable. A crash is not an error message. +- Include headers from `sgl_kernel/` — **not** raw CUDA headers for anything already covered +- Use `TensorMatcher` for all tensor validation; never manually check shape/dtype/device +- Use `AlignedVector` for vectorised 128-bit loads/stores — significant bandwidth win +- Use `LaunchKernel` — it resolves the stream and checks errors automatically +- Use `RuntimeCheck` for runtime assertions with useful error messages +- `fp16_t` / `bf16_t` / `fp32_t` are the project's type aliases (from `utils.cuh`) +- `device::cast` or `dtype_trait::from(val)` for cross-type conversions +- `device::math::` functions for device math instead of bare `__` intrinsics --- -## Step 2: Add the Python wrapper (compile + load with `load_jit`) +## Step 2: Add the Python wrapper in `jit_kernel/` -Create: +Create `python/sglang/jit_kernel/scale.py`: -- `python/sglang/jit_kernel/.py` +```python +from __future__ import annotations -### 2.1 Use `cache_once` for module caching +from typing import TYPE_CHECKING -Use `sglang.jit_kernel.utils.cache_once` (do **not** use `functools.lru_cache`). +import torch -Reason: `functools.lru_cache` is not compatible with `torch.compile` in this codebase. +from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args -### 2.2 Build and load the module with `load_jit` +if TYPE_CHECKING: + from tvm_ffi.module import Module -`load_jit` compiles a tvm-ffi module from C++/CUDA sources and returns a module object. -Key fields: +@cache_once +def _jit_scale_module(dtype: torch.dtype, factor_numer: int, factor_denom: int) -> Module: + """Compile and cache the JIT scale module for a given dtype and factor.""" + args = make_cpp_args(dtype, factor_numer, factor_denom) + return load_jit( + "scale", + *args, + cuda_files=["elementwise/scale.cuh"], + cuda_wrappers=[("scale", f"scale<{args}>")], + ) -- `*args: str`: a unique marker for the build (different kernels / different template args must produce different markers) -- `cpp_files` / `cuda_files`: filenames under `jit_kernel/csrc/` -- `cpp_wrappers` / `cuda_wrappers`: list of `(export_name, kernel_symbol)` - - `export_name` is how you call it from Python: `module.export_name(...)` - - `kernel_symbol` is the C++ symbol name (can include template args) -### 2.3 Template arguments (if needed) +def scale(src: torch.Tensor, factor: float, out: torch.Tensor | None = None) -> torch.Tensor: + """ + Element-wise scale: dst = src * factor. -Use `make_cpp_args(...)` to convert Python values (int/float/bool/torch.dtype) into C++ template arguments. + Supported dtypes: torch.float16, torch.bfloat16, torch.float32. -### 2.4 Destination-passing style (recommended) + Parameters + ---------- + src : CUDA tensor (FP16 / BF16 / FP32) + factor : scale factor + out : optional pre-allocated output tensor (same shape/dtype as src) -Prefer APIs that accept preallocated outputs (e.g. `out=` / `output=`) to avoid allocations in hot paths. + Returns + ------- + Scaled tensor (dst = src * factor). + """ + assert src.is_cuda, "src must be a CUDA tensor" + assert src.dtype in (torch.float16, torch.bfloat16, torch.float32), ( + f"Unsupported dtype {src.dtype}. Supported: float16, bfloat16, float32" + ) + if out is None: + out = torch.empty_like(src) + else: + assert out.shape == src.shape, "out shape must match src" + assert out.dtype == src.dtype, "out dtype must match src" + + # Encode factor as integer ratio; denom=1000 gives 3 decimal places of precision + factor_denom = 1000 + factor_numer = round(factor * factor_denom) + + module = _jit_scale_module(src.dtype, factor_numer, factor_denom) + module.scale(out, src) + return out +``` + +**Key points:** + +- Use `cache_once` — **not** `functools.lru_cache` (incompatible with `torch.compile`) +- `load_jit` first arg(s) form the unique build marker; same marker = same cached binary +- `cuda_wrappers`: `(export_name, kernel_symbol)` — `export_name` is called from Python +- `make_cpp_args(dtype, ...)` converts `torch.dtype` to C++ type alias: + +| `torch.dtype` | C++ type | +|--------------------|------------| +| `torch.float16` | `fp16_t` | +| `torch.bfloat16` | `bf16_t` | +| `torch.float32` | `fp32_t` | --- ## Step 3 (optional): Tune JIT build flags -`load_jit` supports: +```python +return load_jit( + "scale", + *args, + cuda_files=["elementwise/scale.cuh"], + cuda_wrappers=[("scale", f"scale<{args}>")], + extra_cuda_cflags=["-O3", "--use_fast_math"], +) +``` -- `extra_cflags`, `extra_cuda_cflags`, `extra_ldflags` -- `extra_include_paths` -- `build_directory` +If your kernel requires SM90+, raise a clear Python error before calling `load_jit`: -**CUDA arch list:** - -`load_jit` sets `TVM_FFI_CUDA_ARCH_LIST` automatically if it is not already present. - -If your kernel has hard arch requirements (e.g. SM90+ only), enforce that: - -- In Python wrapper (raise a clear error) -- In tests/benchmarks (skip or return NaN for unsupported providers) +```python +if torch.cuda.get_device_capability()[0] < 9: + raise RuntimeError("This kernel requires SM90 (Hopper) or later") +``` --- ## Step 4: Write tests (required) -Create: +Create `python/sglang/jit_kernel/tests/test_scale.py`: -- `python/sglang/jit_kernel/tests/test_.py` +```python +import pytest +import torch +from sglang.jit_kernel.scale import scale -**Recommended test patterns:** -- Compare against a reference implementation (PyTorch or math definition) -- If a corresponding op exists in `sgl-kernel` (AOT) or FlashInfer, add a cross-implementation equivalence test +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("size", [1, 127, 128, 1024, 4097]) # cover tail remainder +@pytest.mark.parametrize("factor", [0.5, 1.0, 2.0, 3.0]) +def test_scale_correctness(dtype, size, factor): + src = torch.randn(size, dtype=dtype, device="cuda") + out = scale(src, factor) + expected = src * factor -**Minimum coverage:** + 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) -- Shapes: typical + edge cases -- Dtypes: the dtypes you claim to support -- Correctness: `torch.testing.assert_close` with appropriate tolerances -- Failure modes: invalid dtype/shape/device should fail clearly (or be skipped) + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_scale_out_param(dtype): + src = torch.randn(1024, dtype=dtype, device="cuda") + out = torch.empty_like(src) + result = scale(src, 2.0, out=out) + assert result is out + torch.testing.assert_close(out, src * 2.0, rtol=1e-2, atol=1e-2) + + +def test_scale_cpu_error(): + src = torch.randn(128, dtype=torch.float16) # CPU tensor + with pytest.raises(AssertionError, match="CUDA"): + scale(src, 2.0) + + +def test_scale_unsupported_dtype(): + src = torch.randint(0, 10, (128,), dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="Unsupported dtype"): + scale(src, 2.0) + + +if __name__ == "__main__": + pytest.main([__file__, "-q"]) +``` Run: ```bash -pytest python/sglang/jit_kernel/tests/test_.py -q +pytest python/sglang/jit_kernel/tests/test_scale.py -q ``` --- ## Step 5: Add a benchmark (required) -Create: +Create `python/sglang/jit_kernel/benchmark/bench_scale.py`: -- `python/sglang/jit_kernel/benchmark/bench_.py` +```python +import itertools -Use the shared helpers: +import torch +import triton +import triton.testing -- `python/sglang/jit_kernel/benchmark/utils.py` - - `is_in_ci()` - - `get_benchmark_range(...)` - - `run_benchmark(fn)` (uses `triton.testing.do_bench_cudagraph` and returns microseconds) +from sglang.jit_kernel.benchmark.utils import ( + DEFAULT_DEVICE, + DEFAULT_DTYPE, + get_benchmark_range, + run_benchmark, +) +from sglang.jit_kernel.scale import scale as jit_scale -**Minimum benchmark requirements:** -- At least two providers/variants: - - Your JIT kernel - - A baseline (FlashInfer / `sgl-kernel` AOT / PyTorch / `torch.compile`) -- CI-friendly reduced ranges (guard with `is_in_ci()` or env vars) -- Use `triton.testing.Benchmark` + `triton.testing.perf_report` +SIZE_LIST = get_benchmark_range( + full_range=[2**n for n in range(10, 20)], # 1K … 512K elements + ci_range=[4096, 65536], +) + +configs = list(itertools.product(SIZE_LIST)) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["size"], + x_vals=configs, + line_arg="provider", + line_vals=["jit", "torch"], + line_names=["SGL JIT Kernel", "PyTorch"], + styles=[("blue", "-"), ("red", "--")], + ylabel="us", + plot_name="scale-performance", + args={}, + ) +) +def benchmark(size: int, provider: str): + src = torch.randn(size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) + factor = 2.0 + + if provider == "jit": + fn = lambda: jit_scale(src, factor) + else: + fn = lambda: src * factor + + return run_benchmark(fn) + + +if __name__ == "__main__": + benchmark.run(print_data=True) +``` Run: ```bash -python python/sglang/jit_kernel/benchmark/bench_.py +python python/sglang/jit_kernel/benchmark/bench_scale.py ``` --- ## Troubleshooting -- **JIT compilation fails**: - - Ensure the file is under `python/sglang/jit_kernel/csrc/` - - Reduce template argument combinations to minimize compilation scope - -- **CUDA crash / illegal memory access**: - - `CUDA_LAUNCH_BLOCKING=1` - - `compute-sanitizer --tool memcheck python ...` - -- **Unstable benchmark results**: - - Use CUDA-graph-based benchmarking (`run_benchmark` does this by default) - - Fix input distributions and shapes +- **JIT compilation fails**: ensure the `.cuh` file is under `python/sglang/jit_kernel/csrc/`; reduce template argument combinations +- **CUDA crash / illegal memory access**: `CUDA_LAUNCH_BLOCKING=1`; `compute-sanitizer --tool memcheck python ...` +- **Unstable benchmark results**: `run_benchmark` uses CUDA-graph-based timing by default --- ## References - `docs/developer_guide/development_jit_kernel_guide.md` -- `python/sglang/jit_kernel/utils.py` (`cache_once`, `load_jit`, wrappers, CUDA arch list) -- `python/sglang/jit_kernel/include/sgl_kernel/tensor.h` (`TensorMatcher` and symbolic size/dtype/device helpers) -- Existing kernels that are good references for utility usage: - - `python/sglang/jit_kernel/per_tensor_quant_fp8.py` + `python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh` - - `python/sglang/jit_kernel/norm.py` + `python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh` - - `python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh` -- `python/sglang/jit_kernel/tests/test_add_constant.py` (minimal runnable example) -- `python/sglang/jit_kernel/benchmark/utils.py` (benchmark helpers) +- `python/sglang/jit_kernel/utils.py` — `cache_once`, `load_jit`, `make_cpp_args` +- `python/sglang/jit_kernel/include/sgl_kernel/tensor.h` — `TensorMatcher`, `SymbolicSize/DType/Device` +- `python/sglang/jit_kernel/include/sgl_kernel/utils.cuh` — type aliases, `LaunchKernel`, `SGL_DEVICE` +- `python/sglang/jit_kernel/include/sgl_kernel/vec.cuh` — `AlignedVector` +- `python/sglang/jit_kernel/include/sgl_kernel/tile.cuh` — `tile::Memory` +- `python/sglang/jit_kernel/include/sgl_kernel/type.cuh` — `dtype_trait`, `packed_t`, `device::cast` +- `python/sglang/jit_kernel/include/sgl_kernel/math.cuh` — `device::math::` +- `python/sglang/jit_kernel/include/sgl_kernel/warp.cuh` — `warp::reduce_sum/max` +- `python/sglang/jit_kernel/include/sgl_kernel/cta.cuh` — `cta::reduce_max` +- `python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh` — `atomic::max` +- `python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh` — occupancy / SM count helpers +- `python/sglang/jit_kernel/csrc/add_constant.cuh` — minimal runnable reference +- `python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh` — real example using `TensorMatcher` + `LaunchKernel` + `tile::Memory` +- `python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh` — real example using `runtime::get_blocks_per_sm` + persistent kernel pattern +- `python/sglang/jit_kernel/benchmark/utils.py` — benchmark helpers + +## Summary of Files Created + +``` +python/sglang/jit_kernel/csrc/elementwise/scale.cuh # NEW: CUDA kernel +python/sglang/jit_kernel/scale.py # NEW: Python wrapper +python/sglang/jit_kernel/tests/test_scale.py # NEW: Tests +python/sglang/jit_kernel/benchmark/bench_scale.py # NEW: Benchmark +``` diff --git a/.claude/skills/add-sgl-kernel/SKILL.md b/.claude/skills/add-sgl-kernel/SKILL.md index f53900845..b31187613 100644 --- a/.claude/skills/add-sgl-kernel/SKILL.md +++ b/.claude/skills/add-sgl-kernel/SKILL.md @@ -5,17 +5,20 @@ description: Step-by-step tutorial for adding a heavyweight AOT CUDA/C++ kernel # Tutorial: Adding a New Kernel to `sgl-kernel` (AOT / Heavyweight) -This SKILL is a step-by-step guide for adding a **heavyweight** CUDA/C++ kernel to `sgl-kernel/`. +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. -Typical characteristics: +## Goal -- Depends on heavyweight components such as CUTLASS / FlashInfer / DeepGEMM / sgl-attn -- Needs AOT build and distribution (wheel / torch extension), so build time, link flags, CUDA arch targets, and binary size matter -- Exposed as a stable `sgl_kernel` API and used by higher-level code (including `torch.compile`) +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 `sgl-kernel/include/utils.h`) ## Two rules of thumb (must follow) -1. **Heavyweight kernels go to `sgl-kernel`.** If it depends on CUTLASS/FlashInfer/DeepGEMM (or similarly heavy stacks), implement it in `sgl-kernel/`. +1. **Heavyweight kernels go to `sgl-kernel`.** If it depends on CUTLASS / FlashInfer / DeepGEMM (or similarly heavy stacks), implement it in `sgl-kernel/`. 2. **Lightweight kernels go to `python/sglang/jit_kernel`.** If it is small, has few dependencies, and benefits from rapid iteration, implement it as a JIT kernel instead. In addition, every new kernel must ship with: @@ -25,154 +28,275 @@ In addition, every new kernel must ship with: --- -## Goal - -Add a new kernel end-to-end, including: - -- CUDA/C++ implementation -- Torch library registration (`m.def` schema + `m.impl` dispatch) -- Build system integration (CMake sources list) -- Python-facing API -- Correctness tests and performance benchmarks - ---- - ## Repository integration map You will typically touch these files/areas: -- Implementation: `sgl-kernel/csrc/...` +- Implementation: `sgl-kernel/csrc/elementwise/scale.cu` (pick the right subdirectory) - Public declarations: `sgl-kernel/include/sgl_kernel_ops.h` - Torch extension registration: `sgl-kernel/csrc/common_extension.cc` - Build: `sgl-kernel/CMakeLists.txt` (`set(SOURCES ...)`) -- Python API: `sgl-kernel/python/sgl_kernel/...` and `sgl-kernel/python/sgl_kernel/__init__.py` -- Tests: `sgl-kernel/tests/test_.py` -- Benchmarks: `sgl-kernel/benchmark/bench_.py` +- Python API: `sgl-kernel/python/sgl_kernel/` and `sgl-kernel/python/sgl_kernel/__init__.py` +- Tests: `sgl-kernel/tests/test_scale.py` +- Benchmarks: `sgl-kernel/benchmark/bench_scale.py` --- ## Step 1: Implement the kernel in `csrc/` -1. Pick the right subdirectory: +Pick the right subdirectory: -- `csrc/elementwise/` -- `csrc/gemm/` -- `csrc/attention/` -- `csrc/moe/` +- `csrc/elementwise/` — for element-wise ops (our example) +- `csrc/gemm/`, `csrc/attention/`, `csrc/moe/` — for other categories -2. Implementation requirements: +Create `sgl-kernel/csrc/elementwise/scale.cu`: -- Clearly define dtype/shape/stride/contiguity assumptions -- If assumptions are violated, fail fast with a readable error (e.g. `TORCH_CHECK(...)`) -- After kernel launch, perform device error checking (follow existing project conventions) +```cpp +#include +#include +#include + +#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 +__global__ void scale_kernel(T* __restrict__ out, + const T* __restrict__ input, + float factor, + int64_t n) { + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx < n) { + out[idx] = static_cast(static_cast(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<<>>( + static_cast(out.data_ptr()), + static_cast(input.data_ptr()), + static_cast(factor), + n); + cudaError_t status = cudaGetLastError(); + TORCH_CHECK(status == cudaSuccess, + "scale_kernel launch failed: ", cudaGetErrorString(status)); + return true; + }); +} +``` **Key points:** -- Prefer explicit validation over "it probably works". -- If a kernel only works on certain architectures, make that restriction explicit (error/skip behavior). +- Use `at::Tensor` (PyTorch tensors), `TORCH_CHECK` for validation, `at::cuda::getCurrentCUDAStream()` for stream +- `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: +Edit `sgl-kernel/include/sgl_kernel_ops.h`, add to the elementwise section: -- `sgl-kernel/include/sgl_kernel_ops.h` - -Add your function declaration in the appropriate section. +```cpp +void scale(at::Tensor& out, const at::Tensor& input, double factor); +``` --- -## Step 3: Register the op in `csrc/common_extension.cc` (schema + dispatch) +## Step 3: Register the op in `csrc/common_extension.cc` -Edit: +Edit `sgl-kernel/csrc/common_extension.cc`, inside `TORCH_LIBRARY_FRAGMENT(sgl_kernel, m)`: -- `sgl-kernel/csrc/common_extension.cc` - -Inside `TORCH_LIBRARY_FRAGMENT(sgl_kernel, m)`: - -1. Add `m.def(...)` with a **schema**. -2. Add `m.impl(...)` for CUDA dispatch. +```cpp +// From csrc/elementwise +m.def("scale(Tensor! out, Tensor input, float factor) -> ()"); +m.impl("scale", torch::kCUDA, &scale); +``` **Key points:** -- The schema is important for `torch.compile` and for consistent call signatures. -- If your underlying C++ API uses native types (e.g. `int`, `float`), but PyTorch bindings expect `int64_t` / `double`, use the project’s recommended shim approach (see `sgl-kernel/README.md`). +- `Tensor!` means in-place / mutable output argument +- The schema is important for `torch.compile` and for consistent call signatures +- If your underlying C++ API uses `float` but PyTorch bindings expect `double`, the implicit cast is fine for scalars; use shims if needed for other types --- ## Step 4: Add the new source file to `CMakeLists.txt` -Edit: +Edit `sgl-kernel/CMakeLists.txt`, add to `set(SOURCES ...)`: -- `sgl-kernel/CMakeLists.txt` - -Add your new `.cu` / `.cc` file to the `set(SOURCES ...)` list. +```cmake +csrc/elementwise/scale.cu +``` **Key points:** -- Keep the list **alphabetically sorted** (the file explicitly requires this). -- If your kernel has arch constraints, reflect that in tests/benchmarks via skip logic. +- 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 `sgl-kernel/python/sgl_kernel/` -Goal: users can call `sgl_kernel.(...)`. +In `sgl-kernel/python/sgl_kernel/__init__.py`, add: -- Add/extend a Python wrapper under `sgl-kernel/python/sgl_kernel/` (follow existing module organization). -- Export it from `sgl-kernel/python/sgl_kernel/__init__.py`. +```python +from torch.ops import sgl_kernel as _ops + +def scale(out: torch.Tensor, input: torch.Tensor, factor: float) -> None: + """ + Element-wise scale: out = input * factor (in-place into out). + + Supported dtypes: torch.float16, torch.bfloat16, torch.float32. + + Parameters + ---------- + out : pre-allocated CUDA output tensor (same shape/dtype as input) + input : CUDA input tensor + factor : scale factor (float) + """ + _ops.scale(out, input, factor) +``` + +Or export it from the existing module organisation — follow the pattern already used by similar ops in `__init__.py`. --- ## Step 6: Write tests (required) -Create: +Create `sgl-kernel/tests/test_scale.py`: -- `sgl-kernel/tests/test_.py` +```python +import pytest +import torch +import sgl_kernel -**Minimum coverage:** -- **Shapes**: typical + edge cases -- **Dtypes**: whatever the kernel claims to support -- **Correctness**: compare with a reference implementation (PyTorch / FlashInfer / another stable backend) -- **Negative cases**: unsupported dtype/shape/arch should either raise a clear error or be explicitly skipped +@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) -**Skipping by architecture:** + sgl_kernel.scale(out, input, factor) -- Use `@pytest.mark.skipif(..., reason="...")` when compute capability requirements apply. + 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(out, input, 2.0) + + +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(out, input, 2.0) + + +if __name__ == "__main__": + pytest.main([__file__, "-q"]) +``` Run: ```bash -pytest sgl-kernel/tests/test_.py -q +pytest sgl-kernel/tests/test_scale.py -q ``` --- ## Step 7: Add a benchmark (required) -Create: +Create `sgl-kernel/benchmark/bench_scale.py`: -- `sgl-kernel/benchmark/bench_.py` +```python +import itertools +import os -Follow the repository convention: +import torch +import triton +import triton.testing -- Use `triton.testing.Benchmark` + `triton.testing.perf_report` -- Prefer `triton.testing.do_bench_cudagraph` for timing +import sgl_kernel -**Minimum benchmark requirements:** +IS_CI = ( + os.getenv("CI", "false").lower() == "true" + or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" +) -- At least two providers/variants: - - Your `sgl_kernel` implementation - - A baseline (PyTorch / `torch.compile` / Triton / FlashInfer) -- Quantiles output (median/min/max) -- CI-friendly ranges controlled by `CI` / `GITHUB_ACTIONS` +dtypes = [torch.float16] if IS_CI else [torch.float16, torch.bfloat16, torch.float32] +sizes = [4096] if IS_CI else [2**n for n in range(10, 20)] # 1K … 512K +factors = [2.0] + +configs = list(itertools.product(dtypes, sizes)) + + +def torch_scale(input: torch.Tensor, factor: float) -> torch.Tensor: + return input * factor + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["dtype", "size"], + x_vals=configs, + line_arg="provider", + line_vals=["sglang", "torch"], + line_names=["SGL Kernel", "PyTorch"], + styles=[("green", "-"), ("red", "--")], + ylabel="µs (median)", + plot_name="scale-performance", + args={}, + ) +) +def benchmark(dtype, size, provider): + input = torch.randn(size, dtype=dtype, device="cuda") + out = torch.empty_like(input) + factor = 2.0 + + if provider == "sglang": + fn = lambda: sgl_kernel.scale(out, input, factor) + else: + fn = lambda: torch_scale(input, factor) + + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + fn, quantiles=[0.5, 0.2, 0.8] + ) + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +if __name__ == "__main__": + benchmark.run(print_data=True) +``` Run: ```bash -python sgl-kernel/benchmark/bench_.py +python sgl-kernel/benchmark/bench_scale.py ``` --- @@ -195,8 +319,10 @@ make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1" Validate: -- Tests: `pytest sgl-kernel/tests/test_.py -q` -- Benchmark: `python sgl-kernel/benchmark/bench_.py` +```bash +pytest sgl-kernel/tests/test_scale.py -q +python sgl-kernel/benchmark/bench_scale.py +``` --- @@ -206,6 +332,7 @@ Validate: - **Memory errors**: `compute-sanitizer --tool memcheck python ...` - **Build is too slow / OOM**: reduce `MAX_JOBS` and `SGL_KERNEL_COMPILE_THREADS` - **Binary bloat**: use `sgl-kernel/analyze_whl_kernel_sizes.py` +- **CMake sources list**: if your `.cu` file is missing from `SOURCES`, the symbol will be undefined at link time --- @@ -215,3 +342,17 @@ Validate: - `sgl-kernel/include/sgl_kernel_ops.h` - `sgl-kernel/csrc/common_extension.cc` - `sgl-kernel/CMakeLists.txt` +- `sgl-kernel/include/utils.h` — `DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16` macro and friends +- `sgl-kernel/csrc/elementwise/activation.cu` — reference for the FP16/BF16/FP32 dispatch pattern + +## Summary of Files Created/Modified + +``` +sgl-kernel/csrc/elementwise/scale.cu # NEW: CUDA kernel + launcher +sgl-kernel/include/sgl_kernel_ops.h # MODIFIED: C++ declaration +sgl-kernel/csrc/common_extension.cc # MODIFIED: schema + dispatch registration +sgl-kernel/CMakeLists.txt # MODIFIED: add source file (alphabetical) +sgl-kernel/python/sgl_kernel/__init__.py # MODIFIED: export Python API +sgl-kernel/tests/test_scale.py # NEW: tests +sgl-kernel/benchmark/bench_scale.py # NEW: benchmark +```