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-jit-kernel/SKILL.md
+1 -1
View File
@@ -18,7 +18,7 @@ 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`)**: 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.
- **AOT (`sgl-kernel`)**: prefer this when the kernel **does** depend on CUTLASS or another large C++ project, or when it should live in `python/sglang/kernels/aot/` 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`.
---
.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
```
.claude/skills/llm-torch-profiler-analysis/references/fuse-overlap-catalog.md
+1 -1
View File
@@ -145,7 +145,7 @@ Stable entries should be folded into the mainline family rows above.
| PR `#21491` FlashInfer TRTLLM FP8 MoE with fused shared experts | `num_fused_shared_experts`<br>`trtllm_fp8_block_scale_moe` | `PR #21491`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py`<br>`python/sglang/srt/models/deepseek_v2.py` | FlashInfer TRTLLM FP8 MoE path can fuse shared experts inside the routed MoE kernel | On FP8 TRTLLM MoE discussions, treat fused shared experts as an upstream pattern that already has a concrete PR. |
| PR `#22005` fused add + RMSNorm + per-token FP8 quant | `fused_add_rmsnorm_per_token_quant`<br>`per_token_quant_fp8` | `PR #22005`<br>`python/sglang/kernels/jit/csrc/elementwise/fused_add_rmsnorm_per_token_quant.cuh`<br>`python/sglang/kernels/jit/fused_add_rmsnorm_per_token_quant.py` | CUDA JIT kernel keeps normed values in registers and emits BF16 + FP8 outputs plus per-token scales | If FP8 online-quant traces show add+norm followed by per-token quant, treat this as an in-flight upstream CUDA fuse family. |
| PR `#20667` Qwen3.5 fused QK norm + RoPE + KV cache write | `fused_qk_norm_rope_cache_pts_quant_shuffle`<br>`fused_qk_norm_mrope_3d_cache_pts_quant_shuffle`<br>`rotary_dim` | `PR #20667`<br>`python/sglang/srt/models/qwen3_5.py`<br>`python/sglang/srt/models/utils.py` | ROCm / AITER path fuses Q / K RMSNorm, partial or 3D RoPE, and direct KV cache write for Qwen3.5 attention | Treat split QK-norm + RoPE + cache-store on Qwen3.5 as a concrete in-flight upstream family, not a novel idea. |
| PR `#22392` CUTLASS FP8 GEMM replacing nvjet | `cutlass_scaled_mm`<br>`fp8_scaled_mm`<br>`nvjet`<br>`cudaMemsetAsync` | `PR #22392`<br>`sgl-kernel/python/sgl_kernel/gemm.py`<br>`python/sglang/srt/layers/quantization/fp8_utils.py` | Runtime replacement swaps nvjet FP8 GEMMs for CUTLASS kernels, removing per-launch memset bubbles and extra output-copy kernels | Treat nvjet GEMM + memset bubble ladders as an in-flight SGLang linear-kernel family before calling them novel. |
| PR `#22392` CUTLASS FP8 GEMM replacing nvjet | `cutlass_scaled_mm`<br>`fp8_scaled_mm`<br>`nvjet`<br>`cudaMemsetAsync` | `PR #22392`<br>`python/sglang/kernels/aot/python/sgl_kernel/gemm.py`<br>`python/sglang/srt/layers/quantization/fp8_utils.py` | Runtime replacement swaps nvjet FP8 GEMMs for CUTLASS kernels, removing per-launch memset bubbles and extra output-copy kernels | Treat nvjet GEMM + memset bubble ladders as an in-flight SGLang linear-kernel family before calling them novel. |
| PR `#18612` NVFP4 CUTLASS MoE fused SiLU+Mul+quant | `silu_and_mul_scaled_nvfp4`<br>`nvfp4 expert quant`<br>`cutlass moe` | `PR #18612`<br>`python/sglang/srt/layers/moe/cutlass_w4a8_moe.py`<br>`python/sglang/kernels/ops/quantization/nvfp4_gemm_swiglu_nvfp4_quant.py` | Fuses MoE activation epilogue and NVFP4 expert quantization before the CUTLASS MoE second GEMM | Treat split SiLU+Mul then NVFP4 expert quant in CUTLASS MoE traces as an in-flight upstream SGLang family. |
| PR `#22918` FlashInfer per-token NVFP4 MoE | `per_token_nvfp4`<br>`trtllm_fp4_block_scale_moe`<br>`FlashInfer MoE` | `PR #22918`<br>`python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py` | Adds FlashInfer-backed per-token NVFP4 MoE execution so expert quant/dequant work can move into the fused MoE backend | Treat standalone per-token NVFP4 MoE support kernels as a candidate missing backend-selection path, not an automatically novel kernel idea. |
| PR `#22851` NSA top-k backend and FlashInfer / PyTorch top-k split | `nsa topk`<br>`flashinfer_topk`<br>`pytorch_topk`<br>`fast_topk_transform` | `PR #22851`<br>`python/sglang/srt/layers/attention/nsa_backend.py` | Makes NSA top-k backend selection explicit and aligns fused top-k transform with FlashInfer / PyTorch fallbacks | When NSA top-k dominates decode, first classify it as backend selection or fused-transform eligibility work. |
.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py
+2 -2
View File
@@ -840,7 +840,7 @@ FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
pattern="PR #22392 CUTLASS FP8 scaled MM replacing nvjet",
candidate_path=(
"PR #22392"
"<br>sgl-kernel/python/sgl_kernel/gemm.py"
"<br>python/sglang/kernels/aot/python/sgl_kernel/gemm.py"
"<br>python/sglang/srt/layers/quantization/fp8_utils.py"
),
active_keywords=("cutlass_scaled_mm", "fp8_scaled_mm"),
@@ -2632,7 +2632,7 @@ def fusion_framework_hints(spec: FusionPatternSpec) -> set[str]:
hints.add("tokenspeed")
if "tensorrt_llm/" in text:
hints.add("trtllm")
if any(token in text for token in ("python/sglang/", "sgl-kernel/", "sgl_kernel/")):
if any(token in text for token in ("python/sglang/", "sgl_kernel/")):
hints.add("sglang")
return hints