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

[Kernel] RFC #29630 finale: retire sglang.jit_kernel into sglang.kernels (#32072)

Browse Source 
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
2026-07-23 08:35:09 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 8ce68370b5
commit 99f636a86f
354 changed files with 889 additions and 875 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs/developer_guide/development_jit_kernel_guide.mdx
+12 -12
View File
@@ -10,18 +10,18 @@ We strongly recommend using `clangd` as the language server for JIT kernel devel
For Ubuntu/Debian, you can download clangd from [apt.llvm.org](https://apt.llvm.org/).
If you are using VS Code, we recommend installing the `clangd` extension for better IDE integration.
All JIT-related files are located in `python/sglang/jit_kernel`.
All JIT-related files are located in `python/sglang/kernels/jit`.
Unlike `sgl-kernel`, which compiles CUDA/C++ binaries ahead of time (AOT), just-in-time (JIT) kernels are compiled at runtime.
Consequently, a static `compile_commands.json` cannot be generated.
To enable code completion with `clangd`, run `python -m sglang.jit_kernel` to generate a `.clangd` configuration file in your current directory.
To enable code completion with `clangd`, run `python -m sglang.kernels.jit` to generate a `.clangd` configuration file in your current directory.
After generating the file, restart the clangd language server. It should now recognize all JIT kernel files.
## Code Structure
### C++ Implementation
C++ source code is located in `python/sglang/jit_kernel/csrc`.
Reusable functions should be placed in `python/sglang/jit_kernel/include`.
C++ source code is located in `python/sglang/kernels/jit/csrc`.
Reusable functions should be placed in `python/sglang/kernels/jit/include`.
We use [tvm-ffi](https://github.com/apache/tvm-ffi) for efficient foreign language bindings.
Refer to the [documentation](https://tvm.apache.org/ffi/) for advanced usage, such as exporting C++ objects.
@@ -29,12 +29,12 @@ Typically, `tvm::ffi::TensorView` is sufficient for passing PyTorch Tensors from
### Python Interface
Python interfaces are defined in `python/sglang/jit_kernel`.
The `load_jit` utility function in `python/sglang/jit_kernel/utils.py` loads and returns the compiled module.
Python interfaces are defined in `python/sglang/kernels/jit`.
The `load_jit` utility function in `python/sglang/kernels/jit/utils/compile.py` loads and returns the compiled module.
To export a C++ function (e.g., `cpp_func`), pass `cuda_wrappers=[("func", "cpp_func")]` to `load_jit`.
The function can then be called in Python as `module.func`.
For caching compiled modules, prefer `sglang.jit_kernel.utils.cache_once` over `functools.lru_cache`.
For caching compiled modules, prefer `sglang.kernels.jit.utils.cache_once` over `functools.lru_cache`.
`functools.lru_cache` is not compatible with `torch.compile`.
### C++ Utilities
@@ -161,7 +161,7 @@ def add_constant(src: torch.Tensor, c: int):
### STEP 1: Write the C++ kernel
Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter.
Write your CUDA kernel in [kernels/jit/csrc/add_constant.cuh](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/jit/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter.
```cpp Example
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
@@ -224,7 +224,7 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
### STEP 2: Create Python Interfaces
Next, expose the kernel through a Python wrapper.
Create a new file at [jit_kernel/add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/add_constant.py) and expose the needed interfaces.
Create a new file at [kernels/ops/attention/add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/ops/attention/add_constant.py) and expose the needed interfaces.
```python Example
from __future__ import annotations
@@ -232,7 +232,7 @@ from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
if TYPE_CHECKING:
from tvm_ffi.module import Module
@@ -268,7 +268,7 @@ Keep the Python wrapper thin, but still validate the basic invariants such as de
Finally, import and use the kernel like a regular Python function:
```python Example
from sglang.jit_kernel.add_constant import add_constant
from sglang.kernels.jit.add_constant import add_constant
```
For a complete, runnable example, refer to [test_add_constant.py](https://github.com/sgl-project/sglang/blob/main/test/registered/jit/test_add_constant.py).
@@ -276,7 +276,7 @@ For a complete, runnable example, refer to [test_add_constant.py](https://github
## C++ Include Library Reference
The JIT kernel framework provides a set of reusable C++ headers in
`python/sglang/jit_kernel/include/sgl_kernel/`. Each header is designed
`python/sglang/kernels/jit/include/sgl_kernel/`. Each header is designed
to be lightweight and self-contained. Below is a summary of each header
and its key APIs.