[jit_kernel] Move JIT kernels into namespace sglang (#33400)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5
parent
5fdf6cd18f
commit
4ad5bb5d9a
@@ -23,6 +23,11 @@ After generating the file, restart the clangd language server. It should now rec
|
||||
C++ source code is located in `python/sglang/kernels/jit/csrc`.
|
||||
Reusable functions should be placed in `python/sglang/kernels/jit/include`.
|
||||
|
||||
JIT C++ lives in `namespace sglang`: open it after the include block and close it at the
|
||||
end of the file, with the device kernels and the host wrapper both inside.
|
||||
The shared `host::` and `device::` helpers are nested in it as well, so they resolve unqualified
|
||||
and need no `sglang::` prefix.
|
||||
|
||||
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.
|
||||
Typically, `tvm::ffi::TensorView` is sufficient for passing PyTorch Tensors from Python.
|
||||
@@ -33,6 +38,8 @@ 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`.
|
||||
`load_jit` emits the export wrapper inside `namespace sglang`, so write `cpp_func` without a
|
||||
`sglang::` prefix.
|
||||
|
||||
For caching compiled modules, prefer `sglang.kernels.jit.utils.cache_once` over `functools.lru_cache`.
|
||||
`functools.lru_cache` is not compatible with `torch.compile`.
|
||||
@@ -174,7 +181,7 @@ Write your CUDA kernel in [kernels/jit/csrc/add_constant.cuh](https://github.com
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
namespace sglang {
|
||||
|
||||
template <int32_t kConstant>
|
||||
__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
|
||||
@@ -217,7 +224,7 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
|
||||
num_elements);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace sglang
|
||||
|
||||
```
|
||||
|
||||
|
||||
Reference in New Issue
Block a user