minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Kernel] Move sgl-kernel under sglang.kernels.aot (#32648)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-07-29 17:25:00 +08:00
committed by GitHub
parent 1b9dfa14e6
commit c32c4ef79c
370 changed files with 300 additions and 269 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/skills/add-sgl-kernel/SKILL.md
+38 -38
View File
@@ -14,7 +14,7 @@ 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`)
- 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)
@@ -33,13 +33,13 @@ In addition, every new kernel must ship with:
You will typically touch these files/areas:
- 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_scale.py`
- Benchmarks: `sgl-kernel/benchmark/bench_scale.py`
- 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`
---
@@ -50,7 +50,7 @@ Pick the right subdirectory:
- `csrc/elementwise/` — for element-wise ops (our example)
- `csrc/gemm/`, `csrc/attention/`, `csrc/moe/` — for other categories
Create `sgl-kernel/csrc/elementwise/scale.cu`:
Create `python/sglang/kernels/aot/csrc/elementwise/scale.cu`:
```cpp
#include <ATen/cuda/CUDAContext.h>
@@ -115,7 +115,7 @@ void scale(at::Tensor& out, const at::Tensor& input, double factor) {
## Step 2: Add a C++ declaration in `include/sgl_kernel_ops.h`
Edit `sgl-kernel/include/sgl_kernel_ops.h`, add to the elementwise section:
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);
@@ -125,7 +125,7 @@ void scale(at::Tensor& out, const at::Tensor& input, double factor);
## Step 3: Register the op in `csrc/common_extension.cc`
Edit `sgl-kernel/csrc/common_extension.cc`, inside `TORCH_LIBRARY_FRAGMENT(sgl_kernel, m)`:
Edit `python/sglang/kernels/aot/csrc/common_extension.cc`, inside `TORCH_LIBRARY_FRAGMENT(sgl_kernel, m)`:
```cpp
// From csrc/elementwise
@@ -143,7 +143,7 @@ m.impl("scale", torch::kCUDA, &scale);
## Step 4: Add the new source file to `CMakeLists.txt`
Edit `sgl-kernel/CMakeLists.txt`, add to `set(SOURCES ...)`:
Edit `python/sglang/kernels/aot/CMakeLists.txt`, add to `set(SOURCES ...)`:
```cmake
csrc/elementwise/scale.cu
@@ -156,14 +156,14 @@ csrc/elementwise/scale.cu
---
## Step 5: Expose a Python API under `sgl-kernel/python/sgl_kernel/`
## 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 `sgl-kernel/python/sgl_kernel/elementwise.py`
- then re-export it from `sgl-kernel/python/sgl_kernel/__init__.py`
- 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 `sgl-kernel/python/sgl_kernel/elementwise.py`, add:
For example, in `python/sglang/kernels/aot/python/sgl_kernel/elementwise.py`, add:
```python
import torch
@@ -190,13 +190,13 @@ def scale(
return out
```
Then re-export it from `sgl-kernel/python/sgl_kernel/__init__.py` following the existing import style used by other kernels.
Then re-export it from `python/sglang/kernels/aot/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`:
Create `python/sglang/kernels/aot/tests/test_scale.py`:
```python
import pytest
@@ -241,7 +241,7 @@ if __name__ == "__main__":
## Step 7: Add a benchmark (required)
Create `sgl-kernel/benchmark/bench_scale.py`:
Create `python/sglang/kernels/aot/benchmark/bench_scale.py`:
```python
import itertools
@@ -306,14 +306,14 @@ if __name__ == "__main__":
Build:
```bash
cd sgl-kernel
cd python/sglang/kernels/aot
make build -j16
```
If you need to limit host resource usage:
```bash
cd sgl-kernel
cd python/sglang/kernels/aot
make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"
```
@@ -324,8 +324,8 @@ make build -j1 MAX_JOBS=2 CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=1"
After building successfully, run the test and benchmark:
```bash
pytest sgl-kernel/tests/test_scale.py -q
python sgl-kernel/benchmark/bench_scale.py
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
@@ -339,29 +339,29 @@ Blackwell coverage signal for AOT `sgl-kernel` changes.
- **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 `sgl-kernel/analyze_whl_kernel_sizes.py`
- **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
- `sgl-kernel/README.md`
- `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
- `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
```
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/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
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
```