diff --git a/.claude/skills/add-jit-kernel/SKILL.md b/.claude/skills/add-jit-kernel/SKILL.md index d7f944bf5..232c9eb99 100644 --- a/.claude/skills/add-jit-kernel/SKILL.md +++ b/.claude/skills/add-jit-kernel/SKILL.md @@ -17,8 +17,9 @@ Add a new operation that scales each element of a tensor by a scalar factor: ## 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 +- **JIT (`jit_kernel`)**: prefer this first for kernels that do **not** depend on CUTLASS or another large C++ project. It is the default choice for lightweight kernels that benefit from rapid iteration and first-use compilation. +- **AOT (`sgl-kernel`)**: prefer this when the kernel **does** depend on CUTLASS or another large C++ project, or when it should live in `sgl-kernel/` and participate in the wheel build / torch op registration flow. +- **Exception**: kernels that depend on `flashinfer`, or on CUTLASS that is already provided through `flashinfer`, can still be implemented as `jit_kernel`. --- @@ -26,6 +27,8 @@ Add a new operation that scales each element of a tensor by a scalar factor: **Always prefer these abstractions over raw CUDA primitives.** They provide safety, readability, and consistency with the rest of the codebase. +**Important include rule:** for every `#include ` line, add a short trailing comment explaining why that header is included (for example `// For TensorMatcher, SymbolicSize, SymbolicDevice`). This matches the current JIT kernel style and keeps include usage self-documenting. + ### `utils.h` — Host-side utilities ```cpp @@ -80,7 +83,7 @@ auto device = SymbolicDevice{}; device.set_options(); TensorMatcher({N}) // .with_dtype() - .with_device(device) + .with_device(device) .verify(dst) .verify(src); // same shape, dtype, device as dst const size_t n = N.unwrap(); @@ -95,7 +98,8 @@ const DLDevice dev = device.unwrap(); - **`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`) + - `dtype_trait::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math (primarily for `fp32_t`) + - `dtype_trait::max/min(x, y)` — type-dispatched binary math (primarily 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)`. @@ -105,7 +109,7 @@ const DLDevice dev = device.unwrap(); #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. +- **`device::AlignedVector`** — Aligned storage for N elements of type T. N must be a power of two, `sizeof(T)*N <= 32`. Enables vectorized loads/stores for bandwidth efficiency. In terms of API/codegen constraints, the upper bound is 256-bit; in practice, 128-bit is the portable default, while 256-bit vectorization is typically only viable on `SM100+` and should be gated by an architecture check when needed. - `.load(ptr, offset)` — vectorized load from `ptr[offset]` - `.store(ptr, offset)` — vectorized store to `ptr[offset]` - `.fill(value)` — fill all lanes @@ -117,19 +121,22 @@ const DLDevice dev = device.unwrap(); #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]` +- `tile::Memory` is fundamentally a **1D cooperative accessor** over a contiguous region. +- **`device::tile::Memory::cta(blockDim.x)`** — Creates a tile accessor where each thread handles `tid = threadIdx.x` with stride `tsize` (for `cta(blockDim.x)`, this is `blockDim.x`). Common for loops over a 1D array. +- **`.load(ptr, offset)`** — loads `ptr[tid + offset * tsize]` +- **`.store(ptr, val, offset)`** — stores to `ptr[tid + offset * tsize]` - **`.in_bound(n, offset)`** — boundary check +For a **2D tile**, either flatten `(row, col)` into a linear tile index first, or compute the address manually with `ptr[row * stride + col]` using your thread/block coordinates. + ### `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 +- `device::math::max/min(a, b)` — type-dispatched binary math via `dtype_trait` +- `device::math::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math via `dtype_trait` ### `warp.cuh` — Warp-level primitives @@ -191,11 +198,11 @@ Create `python/sglang/jit_kernel/csrc/elementwise/scale.cuh`. The implementation fully uses the project abstractions described above: ```cpp -#include // TensorMatcher, SymbolicSize, SymbolicDevice -#include // dtype_trait, fp16_t, bf16_t, fp32_t -#include // RuntimeCheck, div_ceil -#include // LaunchKernel, SGL_DEVICE -#include // AlignedVector +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For dtype_trait, fp16_t, bf16_t, fp32_t +#include // For RuntimeCheck, div_ceil +#include // For LaunchKernel, SGL_DEVICE +#include // For AlignedVector #include #include @@ -257,7 +264,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { TensorMatcher({N}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(dst) .verify(src); // same shape / dtype / device as dst @@ -292,6 +299,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) { **Key points:** - Include headers from `sgl_kernel/` — **not** raw CUDA headers for anything already covered +- Add a short trailing `// For ...` explanation to every `#include ` line - 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 diff --git a/.claude/skills/add-sgl-kernel/SKILL.md b/.claude/skills/add-sgl-kernel/SKILL.md index b31187613..07767c16b 100644 --- a/.claude/skills/add-sgl-kernel/SKILL.md +++ b/.claude/skills/add-sgl-kernel/SKILL.md @@ -18,8 +18,9 @@ Add a new operation that scales each element of a tensor by a scalar factor: ## 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 `python/sglang/jit_kernel`.** If it is small, has few dependencies, and benefits from rapid iteration, implement it as a JIT kernel instead. +1. **Prefer `python/sglang/jit_kernel` 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: @@ -156,40 +157,51 @@ csrc/elementwise/scale.cu ## Step 5: Expose a Python API under `sgl-kernel/python/sgl_kernel/` -In `sgl-kernel/python/sgl_kernel/__init__.py`, add: +Prefer following the existing module organization first. For elementwise kernels, the usual pattern is: + +- implement the Python wrapper in `sgl-kernel/python/sgl_kernel/elementwise.py` +- then re-export it from `sgl-kernel/python/sgl_kernel/__init__.py` + +For example, in `sgl-kernel/python/sgl_kernel/elementwise.py`, add: ```python -from torch.ops import sgl_kernel as _ops +import torch -def scale(out: torch.Tensor, input: torch.Tensor, factor: float) -> None: +def scale( + input: torch.Tensor, + factor: float, + out: torch.Tensor | None = None, +) -> torch.Tensor: """ - Element-wise scale: out = input * factor (in-place into out). + Element-wise scale: out = input * factor. 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) + out : optional pre-allocated CUDA output tensor (same shape/dtype as input) """ - _ops.scale(out, input, factor) + if out is None: + out = torch.empty_like(input) + torch.ops.sgl_kernel.scale.default(out, input, factor) + return out ``` -Or export it from the existing module organisation — follow the pattern already used by similar ops in `__init__.py`. +Then re-export it from `sgl-kernel/python/sgl_kernel/__init__.py` following the existing import style used by other kernels. --- ## Step 6: Write tests (required) Create `sgl-kernel/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]) @@ -197,7 +209,8 @@ def test_scale_correctness(dtype, size, factor): input = torch.randn(size, dtype=dtype, device="cuda") out = torch.empty_like(input) - sgl_kernel.scale(out, input, factor) + 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) @@ -208,26 +221,20 @@ 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) + 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(out, input, 2.0) + sgl_kernel.scale(input, 2.0, out=out) if __name__ == "__main__": pytest.main([__file__, "-q"]) ``` -Run: - -```bash -pytest sgl-kernel/tests/test_scale.py -q -``` - --- ## Step 7: Add a benchmark (required) @@ -279,7 +286,7 @@ def benchmark(dtype, size, provider): factor = 2.0 if provider == "sglang": - fn = lambda: sgl_kernel.scale(out, input, factor) + fn = lambda: sgl_kernel.scale(input, factor, out=out) else: fn = lambda: torch_scale(input, factor) @@ -293,15 +300,9 @@ if __name__ == "__main__": benchmark.run(print_data=True) ``` -Run: - -```bash -python sgl-kernel/benchmark/bench_scale.py -``` - --- -## Step 8: Build and validate +## Step 8: Build Build: @@ -317,7 +318,11 @@ cd sgl-kernel make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1" ``` -Validate: +--- + +## Step 9: Validate + +After building successfully, run the test and benchmark: ```bash pytest sgl-kernel/tests/test_scale.py -q @@ -352,7 +357,8 @@ 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/python/sgl_kernel/elementwise.py # MODIFIED: Python wrapper +sgl-kernel/python/sgl_kernel/__init__.py # MODIFIED: re-export Python API sgl-kernel/tests/test_scale.py # NEW: tests sgl-kernel/benchmark/bench_scale.py # NEW: benchmark ```