chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Skills] Refine jit_kernel and sgl-kernel skills (#20095)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-07 22:54:01 +08:00
committed by GitHub
parent 5297b02c88
commit fd79cd8d9c
2 changed files with 60 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
.claude/skills/add-jit-kernel/SKILL.md
View File

@@ -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 <sgl_kernel/...>` 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<kDLCUDA>();
TensorMatcher({N}) //
.with_dtype<fp16_t>()
.with_device(device)
.with_device<kDLCUDA>(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<T>`** — Static trait struct for each scalar type. Provides:
- `dtype_trait<T>::from(value)` — convert from another type (e.g. `fp32_t` → `fp16_t`)
- `dtype_trait<T>::abs/sqrt/rsqrt/max/min(x)` — type-dispatched math (for `fp32_t`)
- `dtype_trait<T>::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math (primarily for `fp32_t`)
- `dtype_trait<T>::max/min(x, y)` — type-dispatched binary math (primarily for `fp32_t`)
- **`packed_t<T>`** — Two-element packed alias: `packed_t<fp16_t>` = `fp16x2_t`, `packed_t<bf16_t>` = `bf16x2_t`, `packed_t<fp32_t>` = `fp32x2_t`. Use for vectorized loads/stores.
- **`device::cast<To, From>(value)`** — Type-safe cast using `dtype_trait`, e.g. `cast<fp32x2_t, fp16x2_t>(v)`.
@@ -105,7 +109,7 @@ const DLDevice dev = device.unwrap();
#include <sgl_kernel/vec.cuh>
```
- **`device::AlignedVector<T, N>`** — 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<T, N>`** — 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 <sgl_kernel/tile.cuh>
```
- **`device::tile::Memory<T>::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<T>` is fundamentally a **1D cooperative accessor** over a contiguous region.
- **`device::tile::Memory<T>::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 <sgl_kernel/math.cuh>
```
- `device::math::max/min/abs/sqrt/rsqrt<T>(a, b)` — type-dispatched math via `dtype_trait`
- `device::math::exp/sin/cos(float)` — fast float math wrappers
- `device::math::max/min<T>(a, b)` — type-dispatched binary math via `dtype_trait`
- `device::math::abs/sqrt/rsqrt/exp/sin/cos<T>(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 <sgl_kernel/tensor.h> // TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/type.cuh> // dtype_trait, fp16_t, bf16_t, fp32_t
#include <sgl_kernel/utils.h> // RuntimeCheck, div_ceil
#include <sgl_kernel/utils.cuh> // LaunchKernel, SGL_DEVICE
#include <sgl_kernel/vec.cuh> // AlignedVector
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
#include <sgl_kernel/type.cuh> // For dtype_trait, fp16_t, bf16_t, fp32_t
#include <sgl_kernel/utils.h> // For RuntimeCheck, div_ceil
#include <sgl_kernel/utils.cuh> // For LaunchKernel, SGL_DEVICE
#include <sgl_kernel/vec.cuh> // For AlignedVector
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
@@ -257,7 +264,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
TensorMatcher({N}) //
.with_dtype<T>()
.with_device(device_)
.with_device<kDLCUDA>(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 <sgl_kernel/...>` 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

66
.claude/skills/add-sgl-kernel/SKILL.md
View File

@@ -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
```