[Minor] Enhance JIT kernel and add dev docs (#14570)
This commit is contained in:
DarkSharpness
6
.gitignore
vendored
6
.gitignore
vendored
@@ -252,3 +252,9 @@ outputs/
|
||||
|
||||
# Eval Cache
|
||||
.longbench_cache/
|
||||
|
||||
# CUDA kernel develop, profile and debug
|
||||
.clangd
|
||||
*.nsys-rep
|
||||
*.ncu-rep
|
||||
*.nvcudmp
|
||||
|
||||
258
docs/developer_guide/JIT_kernels.md
Normal file
258
docs/developer_guide/JIT_kernels.md
Normal file
@@ -0,0 +1,258 @@
|
||||
# Development Guide for JIT Kernels
|
||||
|
||||
## Environment Setup
|
||||
|
||||
We strongly recommend using `clangd` as the language server for JIT kernel development.
|
||||
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`.
|
||||
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.
|
||||
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`.
|
||||
|
||||
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.
|
||||
|
||||
### 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.
|
||||
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`.
|
||||
|
||||
### C++ Utilities
|
||||
|
||||
The following C++ utilities are available:
|
||||
|
||||
#### Integer Range
|
||||
|
||||
Similar to PyTorch, we provide an `irange` function to represent an integer range.
|
||||
|
||||
```C++
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
void test() {
|
||||
for (auto i : host::irange(100)) { // [0, 100)
|
||||
// do something
|
||||
}
|
||||
for (auto i : host::irange(0, 100)) { // [0, 100)
|
||||
// do something
|
||||
}
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
#### Runtime Checking
|
||||
|
||||
`RuntimeCheck` validates conditions at runtime. It accepts optional arguments for error reporting.
|
||||
If the check fails, these arguments are output to aid debugging.
|
||||
`RuntimeDeviceCheck` verifies the status of the last kernel launch.
|
||||
|
||||
```C++
|
||||
#include <sgl_kernel/utils.h>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
void test() {
|
||||
host::RuntimeCheck(1 + 1 == 2, 1 + 1, " != ", 2);
|
||||
host::RuntimeDeviceCheck();
|
||||
// check the provided `cudaError_t`
|
||||
host::RuntimeDeviceCheck(cudaGetLastError());
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
#### Tensor Checking
|
||||
|
||||
`TensorMatcher` provides a readable way to validate and extract tensor shape information.
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h>
|
||||
|
||||
void test(const tvm::ffi::TensorView k_cache, const tvm::ffi::TensorView v_cache) {
|
||||
using namespace host;
|
||||
|
||||
auto D = SymbolicSize{"D"}; // cache dimension
|
||||
auto N = SymbolicSize{"N"}; // kvcache stride
|
||||
auto dtype = SymbolicDType{};
|
||||
auto device = SymbolicDevice{};
|
||||
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_strides({N, 1})
|
||||
.with_dtype<int32_t, int64_t>(dtype)
|
||||
.with_device<kDLCUDA, kDLCPU>(device)
|
||||
.verify(k_cache)
|
||||
.verify(v_cache);
|
||||
}
|
||||
```
|
||||
|
||||
Configure the `TensorMatcher` with expected stride, dtype, and device properties before verification.
|
||||
- If `with_strides` is omitted, the tensor is expected to be contiguous.
|
||||
- Template arguments in `with_dtype` restrict the allowed data types.
|
||||
- Template arguments in `with_device` restrict the allowed devices.
|
||||
- Values passed to `with_xxx` methods enforce equality checks.
|
||||
- Passing `-1` for size or stride allows matching any value.
|
||||
|
||||
A `Symbolic` variable must resolve to the same value across all verifications.
|
||||
Use `.unwrap()` to retrieve the matched value after verification.
|
||||
|
||||
> Note: `TensorMatcher` is a temporary expression and should not be stored in a variable.
|
||||
|
||||
> Tip: Add `//` at the end of the `TensorMatcher` chain to enforce proper indentation.
|
||||
|
||||
#### Kernel Launching
|
||||
|
||||
`LaunchKernel::resolve_device` retrieves the current `cudaStream` from PyTorch.
|
||||
Kernels can also be launched directly using `LaunchKernel`.
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
|
||||
__global__ void kernel() {}
|
||||
|
||||
void test() {
|
||||
const auto num_blocks = 1;
|
||||
const auto num_threads = 32;
|
||||
const auto dynamic_smem = 0;
|
||||
|
||||
DLDevice dev; // suppose this is initialized properly
|
||||
host::LaunchKernel(num_blocks, num_threads, dev)(kernel);
|
||||
|
||||
cudaStream_t stream = host::LaunchKernel::resolve_device(dev);
|
||||
host::LaunchKernel(num_blocks, num_threads, stream, dynamic_smem)(kernel);
|
||||
}
|
||||
|
||||
```
|
||||
|
||||
## Add new kernels
|
||||
|
||||
This section walks through a complete, end-to-end example of adding a new JIT kernel to the system.
|
||||
We use a simple add_constant kernel as a running example, which adds a constant integer value to every element of an input tensor.
|
||||
|
||||
Conceptually, the Python interface looks like this:
|
||||
|
||||
```python
|
||||
def add_constant(src: torch.Tensor, c: int):
|
||||
return src + c
|
||||
```
|
||||
|
||||
### STEP 1: Write the C++ kernel
|
||||
|
||||
Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](../../python/sglang/jit_kernel/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter.
|
||||
|
||||
```cpp
|
||||
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/utils.cuh> // For LaunchKernel
|
||||
#include <sgl_kernel/utils.h> // For div_ceil, RuntimeCheck
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
template <int32_t kConstant>
|
||||
__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
|
||||
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < length) {
|
||||
dst[idx] = src[idx] + kConstant;
|
||||
}
|
||||
}
|
||||
|
||||
constexpr size_t kBlockSize = 256;
|
||||
|
||||
// You can also use struct with static method as an alternative
|
||||
template <int32_t kConstant>
|
||||
void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
|
||||
using namespace host;
|
||||
|
||||
// 1. Validate input tensors
|
||||
SymbolicSize N = {"num_elements"};
|
||||
SymbolicDevice device_;
|
||||
TensorMatcher({N}) // 1D tensor, must be contiguous
|
||||
.with_dtype<int32_t>() // must be int32
|
||||
.with_device<kDLCUDA>(device_) // must be on CUDA device
|
||||
.verify(dst) // check tensor dst
|
||||
.verify(src); // check tensor src
|
||||
|
||||
// 2. Extract required parameters, prepare for kernel launch
|
||||
const size_t num_elements = N.unwrap();
|
||||
const size_t grid_size = div_ceil(num_elements, kBlockSize);
|
||||
const DLDevice device = device_.unwrap();
|
||||
// some extra runtime checks using host::RuntimeCheck
|
||||
RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements);
|
||||
|
||||
// 3. Launch the kernel. Error code will be automatically checked.
|
||||
LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)(
|
||||
// kernel function
|
||||
add_constant_kernel<kConstant>,
|
||||
// kernel arguments
|
||||
static_cast<int32_t*>(dst.data_ptr()),
|
||||
static_cast<int32_t*>(src.data_ptr()),
|
||||
num_elements);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
```
|
||||
|
||||
### STEP 2: Create Python Interfaces
|
||||
|
||||
Next, expose the kernel through a Python wrapper.
|
||||
Create a new file at [jit_kernel/add_constant.py](../../python/sglang/jit_kernel/add_constant.py) and expose the needed interfaces.
|
||||
|
||||
```python
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import load_jit, make_cpp_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _jit_add_constant_module(constant: int) -> Module:
|
||||
args = make_cpp_args(constant) # pass all the template argument
|
||||
return load_jit(
|
||||
"add_constant",
|
||||
*args,
|
||||
cuda_files=["add_constant.cuh"],
|
||||
cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
|
||||
)
|
||||
|
||||
|
||||
def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor:
|
||||
dst = torch.empty_like(src)
|
||||
module = _jit_add_constant_module(constant)
|
||||
module.add_constant(dst, src)
|
||||
return dst
|
||||
|
||||
```
|
||||
|
||||
### STEP 3: Use your kernel
|
||||
|
||||
Finally, import and use the kernel like a regular Python function:
|
||||
|
||||
```python
|
||||
from sglang.jit_kernel.add_constant import add_constant
|
||||
```
|
||||
|
||||
For a complete, runnable example, refer to [test_add_constant.py](../../python/sglang/jit_kernel/test_add_constant.py).
|
||||
48
python/sglang/jit_kernel/__main__.py
Normal file
48
python/sglang/jit_kernel/__main__.py
Normal file
@@ -0,0 +1,48 @@
|
||||
assert __name__ == "__main__"
|
||||
|
||||
|
||||
def generate_clangd():
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
|
||||
|
||||
from sglang.jit_kernel.utils import DEFAULT_INCLUDE
|
||||
|
||||
logger = logging.getLogger()
|
||||
logger.info("Generating .clangd file...")
|
||||
include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE
|
||||
status = subprocess.run(
|
||||
args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
|
||||
capture_output=True,
|
||||
check=True,
|
||||
)
|
||||
compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
|
||||
major, minor = compute_cap.split(".")
|
||||
compile_flags = ",\n ".join(
|
||||
[
|
||||
"-xcuda",
|
||||
f"--cuda-gpu-arch=sm_{major}{minor}",
|
||||
"-std=c++20",
|
||||
"-Wall",
|
||||
"-Wextra",
|
||||
]
|
||||
+ [f"-isystem{path}" for path in include_paths]
|
||||
)
|
||||
clangd_content = f"""
|
||||
CompileFlags:
|
||||
Add: [
|
||||
{compile_flags}
|
||||
]
|
||||
"""
|
||||
if os.path.exists(".clangd"):
|
||||
logger.warning(".clangd file already exists, nothing done.")
|
||||
logger.warning(f"suggested content: {clangd_content}")
|
||||
else:
|
||||
with open(".clangd", "w") as f:
|
||||
f.write(clangd_content)
|
||||
logger.info(".clangd file generated.")
|
||||
|
||||
|
||||
generate_clangd()
|
||||
29
python/sglang/jit_kernel/add_constant.py
Normal file
29
python/sglang/jit_kernel/add_constant.py
Normal file
@@ -0,0 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import load_jit, make_cpp_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _jit_add_constant_module(constant: int) -> Module:
|
||||
args = make_cpp_args(constant) # pass all the template argument
|
||||
return load_jit(
|
||||
"add_constant",
|
||||
*args,
|
||||
cuda_files=["add_constant.cuh"],
|
||||
cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
|
||||
)
|
||||
|
||||
|
||||
def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor:
|
||||
dst = torch.empty_like(src)
|
||||
module = _jit_add_constant_module(constant)
|
||||
module.add_constant(dst, src)
|
||||
return dst
|
||||
60
python/sglang/jit_kernel/csrc/add_constant.cuh
Normal file
60
python/sglang/jit_kernel/csrc/add_constant.cuh
Normal file
@@ -0,0 +1,60 @@
|
||||
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
|
||||
#include <sgl_kernel/utils.cuh> // For LaunchKernel
|
||||
#include <sgl_kernel/utils.h> // For div_ceil, RuntimeCheck
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
template <int32_t kConstant>
|
||||
__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
|
||||
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < length) {
|
||||
dst[idx] = src[idx] + kConstant;
|
||||
}
|
||||
}
|
||||
|
||||
constexpr size_t kBlockSize = 256;
|
||||
|
||||
// You can also use struct with static method as an alternative
|
||||
template <int32_t kConstant>
|
||||
void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
|
||||
using namespace host;
|
||||
|
||||
// 1. Validate input tensors
|
||||
SymbolicSize N = {"num_elements"};
|
||||
SymbolicDevice device_;
|
||||
TensorMatcher({N}) // 1D tensor, must be contiguous
|
||||
.with_dtype<int32_t>() // must be int32
|
||||
.with_device<kDLCUDA>(device_) // must be on CUDA device
|
||||
.verify(dst) // check tensor dst
|
||||
.verify(src); // check tensor src
|
||||
|
||||
// 2. Extract required parameters, prepare for kernel launch
|
||||
const size_t num_elements = N.unwrap();
|
||||
const size_t grid_size = div_ceil(num_elements, kBlockSize);
|
||||
const DLDevice device = device_.unwrap();
|
||||
[[maybe_unused]] // optional, can be omitted
|
||||
const size_t dynamic_smem = 0;
|
||||
[[maybe_unused]] // optional, LaunchKernel can auto determine stream from device
|
||||
const cudaStream_t stream = LaunchKernel::resolve_device(device);
|
||||
// some extra runtime checks using host::RuntimeCheck
|
||||
RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements);
|
||||
|
||||
// 3. Launch the kernel. Error code will be automatically checked.
|
||||
LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)(
|
||||
// kernel function
|
||||
add_constant_kernel<kConstant>,
|
||||
// kernel arguments
|
||||
static_cast<int32_t*>(dst.data_ptr()),
|
||||
static_cast<int32_t*>(src.data_ptr()),
|
||||
num_elements);
|
||||
// You can also manually check the last CUDA error code via:
|
||||
// RuntimeDeviceCheck();
|
||||
}
|
||||
|
||||
} // namespace
|
||||
@@ -1,7 +1,6 @@
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/utils.h>
|
||||
#include <sgl_kernel/warp.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
|
||||
@@ -11,6 +10,105 @@
|
||||
#include <cstdint>
|
||||
#include <type_traits>
|
||||
|
||||
namespace device::warp {
|
||||
|
||||
namespace details {
|
||||
|
||||
template <std::size_t kUnit>
|
||||
inline constexpr auto get_mem_package() {
|
||||
if constexpr (kUnit == 16) {
|
||||
return uint4{};
|
||||
} else if constexpr (kUnit == 8) {
|
||||
return uint2{};
|
||||
} else if constexpr (kUnit == 4) {
|
||||
return uint1{};
|
||||
} else {
|
||||
static_assert(kUnit == 16 || kUnit == 8 || kUnit == 4, "Unsupported memory package size");
|
||||
}
|
||||
}
|
||||
|
||||
template <std::size_t kBytes, std::size_t kUnit>
|
||||
using mem_package_t = decltype(get_mem_package<kUnit>());
|
||||
|
||||
__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
|
||||
uint32_t tmp;
|
||||
asm volatile("ld.global.cs.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
|
||||
return uint1{tmp};
|
||||
}
|
||||
|
||||
__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 {
|
||||
uint32_t tmp0, tmp1;
|
||||
asm volatile("ld.global.cs.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
|
||||
return uint2{tmp0, tmp1};
|
||||
}
|
||||
|
||||
__always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 {
|
||||
uint32_t tmp0, tmp1, tmp2, tmp3;
|
||||
asm volatile("ld.global.cs.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src));
|
||||
return uint4{tmp0, tmp1, tmp2, tmp3};
|
||||
}
|
||||
|
||||
__always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) {
|
||||
uint32_t tmp = value.x;
|
||||
asm volatile("st.global.cs.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
|
||||
}
|
||||
|
||||
__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) {
|
||||
uint32_t tmp0 = value.x;
|
||||
uint32_t tmp1 = value.y;
|
||||
asm volatile("st.global.cs.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
|
||||
}
|
||||
|
||||
__always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) {
|
||||
uint32_t tmp0 = value.x;
|
||||
uint32_t tmp1 = value.y;
|
||||
uint32_t tmp2 = value.z;
|
||||
uint32_t tmp3 = value.w;
|
||||
asm volatile("st.global.cs.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
|
||||
}
|
||||
|
||||
} // namespace details
|
||||
|
||||
template <std::size_t kBytes, std::size_t kUnit, std::size_t kThreads>
|
||||
__always_inline __device__ auto load_vec(const void* __restrict__ src) {
|
||||
using Package = details::mem_package_t<kBytes, kUnit>;
|
||||
constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
|
||||
constexpr auto kLoopCount = kBytes / kBytesPerLoop;
|
||||
static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
|
||||
|
||||
const auto src_packed = static_cast<const Package*>(src);
|
||||
const auto lane_id = threadIdx.x % kThreads;
|
||||
device_vec<Package, kLoopCount> vec;
|
||||
|
||||
#pragma unroll kLoopCount
|
||||
for (std::size_t i = 0; i < kLoopCount; ++i) {
|
||||
const auto j = i * kThreads + lane_id;
|
||||
vec.data[i] = details::load_nc(src_packed + j);
|
||||
}
|
||||
|
||||
return vec;
|
||||
}
|
||||
|
||||
template <std::size_t kBytes, std::size_t kUnit, std::size_t kThreads, typename Tp>
|
||||
__always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) {
|
||||
using Package = details::mem_package_t<kBytes, kUnit>;
|
||||
constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
|
||||
constexpr auto kLoopCount = kBytes / kBytesPerLoop;
|
||||
static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
|
||||
static_assert(std::is_same_v<Tp, device_vec<Package, kLoopCount>>);
|
||||
|
||||
const auto dst_packed = static_cast<Package*>(dst);
|
||||
const auto lane_id = threadIdx.x % kThreads;
|
||||
|
||||
#pragma unroll kLoopCount
|
||||
for (std::size_t i = 0; i < kLoopCount; ++i) {
|
||||
const auto j = i * kThreads + lane_id;
|
||||
details::store_nc(dst_packed + j, vec.data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace device::warp
|
||||
|
||||
namespace {
|
||||
|
||||
struct HicacheKernelParams {
|
||||
@@ -142,10 +240,10 @@ struct HiCacheKernel {
|
||||
const tvm::ffi::TensorView indices_src) {
|
||||
using namespace host;
|
||||
|
||||
auto D = SymbolicSize{"D"}; // cache dimension
|
||||
auto N = SymbolicSize{"N"}; // src kv stride
|
||||
auto M = SymbolicSize{"M"}; // dst kv stride
|
||||
auto L = SymbolicSize{"L"}; // indices length
|
||||
auto D = SymbolicSize{"head dimension"};
|
||||
auto N = SymbolicSize{"src kv stride"};
|
||||
auto M = SymbolicSize{"dst kv stride"};
|
||||
auto L = SymbolicSize{"indices length"};
|
||||
auto cache_dtype = SymbolicDType{};
|
||||
auto indices_dtype = SymbolicDType{};
|
||||
auto indices_device = SymbolicDevice{};
|
||||
@@ -213,8 +311,8 @@ struct HiCacheKernel {
|
||||
const std::size_t kv_dst_stride) {
|
||||
using namespace host;
|
||||
|
||||
auto N = SymbolicSize{"N"}; // num layers
|
||||
auto L = SymbolicSize{"L"}; // indices length
|
||||
auto N = SymbolicSize{"num_layers"};
|
||||
auto L = SymbolicSize{"indices length"};
|
||||
auto dtype_ = SymbolicDType{};
|
||||
auto device_ = SymbolicDevice{};
|
||||
|
||||
|
||||
22
python/sglang/jit_kernel/csrc/test_utils.h
Normal file
22
python/sglang/jit_kernel/csrc/test_utils.h
Normal file
@@ -0,0 +1,22 @@
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
namespace {
|
||||
|
||||
[[maybe_unused]]
|
||||
void assert_same_shape(tvm::ffi::TensorView a, tvm::ffi::TensorView b) {
|
||||
using namespace host;
|
||||
auto N = SymbolicSize{"N"};
|
||||
auto D = SymbolicSize{"D"};
|
||||
TensorMatcher({N, D}) //
|
||||
.with_dtype<float>()
|
||||
.with_device<kDLCUDA>()
|
||||
.verify(a)
|
||||
.verify(b);
|
||||
RuntimeCheck(N.unwrap() > 0 && D.unwrap() > 0);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
@@ -13,7 +13,6 @@
|
||||
#include <initializer_list>
|
||||
#include <optional>
|
||||
#include <ranges>
|
||||
#include <source_location>
|
||||
#include <span>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
@@ -21,13 +20,21 @@
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
|
||||
#ifdef __CUDACC__
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#endif
|
||||
|
||||
namespace host {
|
||||
|
||||
namespace stdr = std::ranges;
|
||||
namespace stdv = std::views;
|
||||
|
||||
namespace details {
|
||||
|
||||
inline constexpr auto kAnyDeviceID = -1;
|
||||
inline constexpr auto kAnySize = static_cast<int64_t>(-1);
|
||||
inline constexpr auto kNullSize = static_cast<int64_t>(-1);
|
||||
inline constexpr auto kNullDType = static_cast<DLDataTypeCode>(18u);
|
||||
inline constexpr auto kNullDevice = static_cast<DLDeviceType>(-1);
|
||||
|
||||
struct SizeRef;
|
||||
struct DTypeRef;
|
||||
struct DeviceRef;
|
||||
@@ -37,7 +44,7 @@ struct dtype_trait {};
|
||||
|
||||
template <std::integral T>
|
||||
struct dtype_trait<T> {
|
||||
inline static constexpr auto value = DLDataType{
|
||||
inline static constexpr DLDataType value = {
|
||||
.code = std::is_signed_v<T> ? DLDataTypeCode::kDLInt : DLDataTypeCode::kDLUInt,
|
||||
.bits = static_cast<std::uint8_t>(sizeof(T) * 8),
|
||||
.lanes = 1};
|
||||
@@ -45,22 +52,31 @@ struct dtype_trait<T> {
|
||||
|
||||
template <std::floating_point T>
|
||||
struct dtype_trait<T> {
|
||||
inline static constexpr auto value =
|
||||
DLDataType{.code = DLDataTypeCode::kDLFloat, .bits = static_cast<std::uint8_t>(sizeof(T) * 8), .lanes = 1};
|
||||
inline static constexpr DLDataType value = {
|
||||
.code = DLDataTypeCode::kDLFloat, .bits = static_cast<std::uint8_t>(sizeof(T) * 8), .lanes = 1};
|
||||
};
|
||||
|
||||
inline constexpr auto kAnyDeviceID = -1;
|
||||
inline constexpr auto kAnySize = static_cast<int64_t>(-1);
|
||||
inline constexpr auto kNullSize = static_cast<int64_t>(-1);
|
||||
inline constexpr auto kNullDType = static_cast<DLDataTypeCode>(18u);
|
||||
inline constexpr auto kNullDevice = static_cast<DLDeviceType>(-1);
|
||||
#ifdef __CUDACC__
|
||||
template <>
|
||||
struct dtype_trait<__half> {
|
||||
inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1};
|
||||
};
|
||||
template <>
|
||||
struct dtype_trait<__nv_bfloat16> {
|
||||
inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1};
|
||||
};
|
||||
#endif
|
||||
|
||||
template <DLDeviceType Code>
|
||||
struct device_trait {
|
||||
inline static constexpr DLDevice value = {.device_type = Code, .device_id = kAnyDeviceID};
|
||||
};
|
||||
|
||||
template <typename... Ts>
|
||||
inline constexpr auto kDTypeList = std::array<DLDataType, sizeof...(Ts)>{dtype_trait<Ts>::value...};
|
||||
|
||||
template <DLDeviceType... Codes>
|
||||
inline constexpr auto kDeviceList = std::array<DLDevice, sizeof...(Codes)>{
|
||||
DLDevice{.device_type = static_cast<DLDeviceType>(Codes), .device_id = kAnyDeviceID}...};
|
||||
inline constexpr auto kDeviceList = std::array<DLDevice, sizeof...(Codes)>{device_trait<Codes>::value...};
|
||||
|
||||
template <typename T>
|
||||
struct PrintAbleSpan {
|
||||
@@ -103,11 +119,13 @@ struct PrintableDevice {
|
||||
inline auto& operator<<(std::ostream& os, DLDevice device) {
|
||||
const auto& mapping = kDeviceStringMap;
|
||||
const auto entry = static_cast<std::size_t>(device.device_type);
|
||||
host::RuntimeCheck(entry < mapping.size());
|
||||
RuntimeCheck(entry < mapping.size());
|
||||
const auto name = mapping[entry];
|
||||
host::RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type));
|
||||
RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type));
|
||||
os << name;
|
||||
if (device.device_id != kAnyDeviceID) os << "[" << device.device_id << "]";
|
||||
if (device.device_id != kAnyDeviceID && device.device_type != DLDeviceType::kDLCPU) {
|
||||
os << ":" << device.device_id;
|
||||
}
|
||||
return os;
|
||||
}
|
||||
|
||||
@@ -118,7 +136,7 @@ inline auto& operator<<(std::ostream& os, PrintableDevice pd) {
|
||||
template <typename T>
|
||||
inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
|
||||
os << "[";
|
||||
for (const auto i : stdv::iota(std::size_t{0}, span.data.size())) {
|
||||
for (const auto i : irange(span.data.size())) {
|
||||
if (i > 0) {
|
||||
os << ", ";
|
||||
}
|
||||
@@ -133,37 +151,58 @@ inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
|
||||
struct SymbolicSize {
|
||||
public:
|
||||
SymbolicSize(std::string_view annotation = {}) : m_value(details::kNullSize), m_annotation(annotation) {}
|
||||
SymbolicSize(const SymbolicSize&) = delete;
|
||||
SymbolicSize& operator=(const SymbolicSize&) = delete;
|
||||
|
||||
auto get_name() const -> std::string_view {
|
||||
return m_annotation;
|
||||
}
|
||||
|
||||
auto set_value(int64_t value) -> void {
|
||||
host::RuntimeCheck(!this->has_value(), "Size value already set");
|
||||
RuntimeCheck(!this->has_value(), "Size value already set");
|
||||
m_value = value;
|
||||
}
|
||||
|
||||
auto has_value() const -> bool {
|
||||
return m_value != details::kNullSize;
|
||||
}
|
||||
|
||||
auto get_value() const -> std::optional<int64_t> {
|
||||
return this->has_value() ? std::optional{m_value} : std::nullopt;
|
||||
}
|
||||
auto unwrap() const -> int64_t {
|
||||
host::RuntimeCheck(this->has_value(), "Size value is not set");
|
||||
|
||||
auto unwrap(DebugInfo info = {}) const -> int64_t {
|
||||
RuntimeCheck(info, this->has_value(), "Size value is not set");
|
||||
return m_value;
|
||||
}
|
||||
|
||||
SymbolicSize(const SymbolicSize&) = delete;
|
||||
SymbolicSize& operator=(const SymbolicSize&) = delete;
|
||||
|
||||
auto verify(int64_t dim) -> void {
|
||||
auto verify(int64_t value, const char* prefix, int64_t dim) -> void {
|
||||
if (this->has_value()) {
|
||||
host::RuntimeCheck(m_value == dim, "Size mismatch: expected ", m_value, " but got ", dim);
|
||||
if (m_value != value) {
|
||||
[[unlikely]];
|
||||
Panic("Size mismatch for ", m_name_str(prefix, dim), ": expected ", m_value, " but got ", value);
|
||||
}
|
||||
} else {
|
||||
this->set_value(dim);
|
||||
this->set_value(value);
|
||||
}
|
||||
}
|
||||
|
||||
auto value_or_name(const char* prefix, int64_t dim) const -> std::string {
|
||||
if (const auto value = this->get_value()) {
|
||||
return std::to_string(*value);
|
||||
} else {
|
||||
return m_name_str(prefix, dim);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
auto m_name_str(const char* prefix, int64_t dim) const -> std::string {
|
||||
std::ostringstream os;
|
||||
os << prefix << '#' << dim;
|
||||
if (!m_annotation.empty()) os << "('" << m_annotation << "')";
|
||||
return std::move(os).str();
|
||||
}
|
||||
|
||||
std::int64_t m_value;
|
||||
std::string_view m_annotation;
|
||||
};
|
||||
@@ -175,27 +214,33 @@ inline auto operator==(DLDevice lhs, DLDevice rhs) -> bool {
|
||||
struct SymbolicDType {
|
||||
public:
|
||||
SymbolicDType() : m_value({details::kNullDType, 0, 0}) {}
|
||||
SymbolicDType(const SymbolicDType&) = delete;
|
||||
SymbolicDType& operator=(const SymbolicDType&) = delete;
|
||||
|
||||
auto set_value(DLDataType value) -> void {
|
||||
host::RuntimeCheck(!this->has_value(), "Dtype value already set");
|
||||
host::RuntimeCheck(
|
||||
RuntimeCheck(!this->has_value(), "Dtype value already set");
|
||||
RuntimeCheck(
|
||||
m_check(value), "Dtype value [", value, "] not in the allowed options: ", details::PrintAbleSpan{m_options});
|
||||
m_value = value;
|
||||
}
|
||||
|
||||
auto has_value() const -> bool {
|
||||
return m_value.code != details::kNullDType;
|
||||
}
|
||||
|
||||
auto get_value() const -> std::optional<DLDataType> {
|
||||
return this->has_value() ? std::optional{m_value} : std::nullopt;
|
||||
}
|
||||
auto unwrap() const -> DLDataType {
|
||||
host::RuntimeCheck(this->has_value(), "Dtype value is not set");
|
||||
|
||||
auto unwrap(DebugInfo info = {}) const -> DLDataType {
|
||||
RuntimeCheck(info, this->has_value(), "Dtype value is not set");
|
||||
return m_value;
|
||||
}
|
||||
|
||||
auto set_options(std::span<const DLDataType> options) -> void {
|
||||
m_options = options;
|
||||
}
|
||||
|
||||
template <typename... Ts>
|
||||
auto set_options() -> void {
|
||||
m_options = details::kDTypeList<Ts...>;
|
||||
@@ -203,7 +248,7 @@ struct SymbolicDType {
|
||||
|
||||
auto verify(DLDataType dtype) -> void {
|
||||
if (this->has_value()) {
|
||||
host::RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " but got ", dtype);
|
||||
RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " but got ", dtype);
|
||||
} else {
|
||||
this->set_value(dtype);
|
||||
}
|
||||
@@ -221,10 +266,12 @@ struct SymbolicDType {
|
||||
struct SymbolicDevice {
|
||||
public:
|
||||
SymbolicDevice() : m_value({details::kNullDevice, details::kAnyDeviceID}) {}
|
||||
SymbolicDevice(const SymbolicDevice&) = delete;
|
||||
SymbolicDevice& operator=(const SymbolicDevice&) = delete;
|
||||
|
||||
auto set_value(DLDevice value) -> void {
|
||||
host::RuntimeCheck(!this->has_value(), "Device value already set");
|
||||
host::RuntimeCheck(
|
||||
RuntimeCheck(!this->has_value(), "Device value already set");
|
||||
RuntimeCheck(
|
||||
m_check(value),
|
||||
"Device value [",
|
||||
details::PrintableDevice{value},
|
||||
@@ -232,20 +279,24 @@ struct SymbolicDevice {
|
||||
details::PrintAbleSpan{m_options});
|
||||
m_value = value;
|
||||
}
|
||||
|
||||
auto has_value() const -> bool {
|
||||
return m_value.device_type != details::kNullDevice;
|
||||
}
|
||||
|
||||
auto get_value() const -> std::optional<DLDevice> {
|
||||
return this->has_value() ? std::optional{m_value} : std::nullopt;
|
||||
}
|
||||
auto unwrap() const -> DLDevice {
|
||||
host::RuntimeCheck(this->has_value(), "Device value is not set");
|
||||
|
||||
auto unwrap(DebugInfo info = {}) const -> DLDevice {
|
||||
RuntimeCheck(info, this->has_value(), "Device value is not set");
|
||||
return m_value;
|
||||
}
|
||||
|
||||
auto set_options(std::span<const DLDevice> options) -> void {
|
||||
m_options = options;
|
||||
}
|
||||
|
||||
template <DLDeviceType... Codes>
|
||||
auto set_options() -> void {
|
||||
m_options = details::kDeviceList<Codes...>;
|
||||
@@ -253,7 +304,7 @@ struct SymbolicDevice {
|
||||
|
||||
auto verify(DLDevice device) -> void {
|
||||
if (this->has_value()) {
|
||||
host::RuntimeCheck(
|
||||
RuntimeCheck(
|
||||
m_value == device,
|
||||
"Device mismatch: expected ",
|
||||
details::PrintableDevice{m_value},
|
||||
@@ -313,19 +364,6 @@ struct SizeRef : BaseRef<SymbolicSize> {
|
||||
// otherwise, we can match any size
|
||||
}
|
||||
}
|
||||
|
||||
auto value_or_name(std::size_t dim) const -> std::string {
|
||||
if (const auto value = (**this).get_value()) {
|
||||
return std::to_string(*value);
|
||||
} else {
|
||||
const auto annotation = (**this).get_name();
|
||||
if (annotation.empty()) {
|
||||
return "dim#" + std::to_string(dim);
|
||||
} else {
|
||||
return static_cast<std::string>(annotation);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct DTypeRef : BaseRef<SymbolicDType> {
|
||||
@@ -361,7 +399,6 @@ struct TensorMatcher {
|
||||
using SizeRef = details::SizeRef;
|
||||
using DTypeRef = details::DTypeRef;
|
||||
using DeviceRef = details::DeviceRef;
|
||||
using Loc_t = std::source_location;
|
||||
|
||||
public:
|
||||
TensorMatcher(const TensorMatcher&) = delete;
|
||||
@@ -371,8 +408,8 @@ struct TensorMatcher {
|
||||
|
||||
auto with_strides(std::initializer_list<SizeRef> strides) && -> TensorMatcher&& {
|
||||
// no partial update allowed
|
||||
host::RuntimeCheck(m_strides.size() == 0, "Strides already specified");
|
||||
host::RuntimeCheck(m_shape.size() == strides.size(), "Strides size must match shape size");
|
||||
RuntimeCheck(m_strides.size() == 0, "Strides already specified");
|
||||
RuntimeCheck(m_shape.size() == strides.size(), "Strides size must match shape size");
|
||||
m_strides = strides;
|
||||
return std::move(*this);
|
||||
}
|
||||
@@ -381,6 +418,7 @@ struct TensorMatcher {
|
||||
auto with_dtype(DTypeRef&& dtype) && -> TensorMatcher&& {
|
||||
m_init_dtype();
|
||||
m_dtype.rebind(*dtype);
|
||||
m_dtype->set_options<Ts...>();
|
||||
return std::move(*this);
|
||||
}
|
||||
|
||||
@@ -396,6 +434,7 @@ struct TensorMatcher {
|
||||
auto with_device(DeviceRef&& device) && -> TensorMatcher&& {
|
||||
m_init_device();
|
||||
m_device.rebind(*device);
|
||||
m_device->set_options<Codes...>();
|
||||
return std::move(*this);
|
||||
}
|
||||
|
||||
@@ -408,70 +447,70 @@ struct TensorMatcher {
|
||||
}
|
||||
|
||||
// once we start verification, we cannot modify anymore
|
||||
auto verify(tvm::ffi::TensorView view, Loc_t loc = Loc_t::current()) const&& -> const TensorMatcher&& {
|
||||
auto verify(tvm::ffi::TensorView view, DebugInfo info = {}) const&& -> const TensorMatcher&& {
|
||||
try {
|
||||
this->m_verify_impl(view);
|
||||
m_verify_impl(view);
|
||||
} catch (PanicError& e) {
|
||||
auto oss = std::ostringstream{};
|
||||
oss << "Tensor match failed for " << this->debug_str() << " at " << loc.file_name() << ":" << loc.line()
|
||||
<< "\n- Root cause: " << e.detail();
|
||||
oss << "Tensor match failed for ";
|
||||
s_print_tensor(oss, view);
|
||||
oss << " at " << info.file_name() << ":" << info.line() << "\n- Root cause: " << e.root_cause();
|
||||
throw PanicError(std::move(oss).str());
|
||||
}
|
||||
return std::move(*this);
|
||||
}
|
||||
|
||||
auto debug_str() const -> std::string {
|
||||
auto oss = std::ostringstream{};
|
||||
private:
|
||||
static auto s_print_tensor(std::ostringstream& oss, tvm::ffi::TensorView view) -> void {
|
||||
oss << "Tensor<";
|
||||
std::size_t dim = 0;
|
||||
for (const auto& size_ref : m_shape) {
|
||||
if (dim > 0) {
|
||||
int64_t dim = 0;
|
||||
for (const auto& size : view.shape()) {
|
||||
if (dim++ > 0) oss << ", ";
|
||||
oss << size;
|
||||
}
|
||||
oss << ">[strides=<";
|
||||
dim = 0;
|
||||
for (const auto& stride : view.strides()) {
|
||||
if (dim++ > 0) {
|
||||
oss << ", ";
|
||||
}
|
||||
oss << size_ref.value_or_name(dim++);
|
||||
oss << stride;
|
||||
}
|
||||
oss << ">";
|
||||
if (m_strides.size() > 0) {
|
||||
oss << " [strides=<";
|
||||
dim = 0;
|
||||
for (const auto& stride_ref : m_strides) {
|
||||
if (dim > 0) {
|
||||
oss << ", ";
|
||||
}
|
||||
oss << stride_ref.value_or_name(dim++);
|
||||
}
|
||||
oss << ">]";
|
||||
}
|
||||
return std::move(oss).str();
|
||||
oss << ">, dtype=" << view.dtype();
|
||||
oss << ", device=" << details::PrintableDevice{view.device()} << "]";
|
||||
}
|
||||
|
||||
private:
|
||||
auto m_verify_impl(tvm::ffi::TensorView view) const -> void {
|
||||
const auto dim = static_cast<std::size_t>(view.dim());
|
||||
host::RuntimeCheck(dim == m_shape.size(), "Tensor dimension mismatch: expected ", m_shape.size(), " but got ", dim);
|
||||
for (const auto i : stdv::iota(std::size_t{0}, dim)) {
|
||||
m_shape[i]->verify(view.size(i));
|
||||
RuntimeCheck(dim == m_shape.size(), "Tensor dimension mismatch: expected ", m_shape.size(), " but got ", dim);
|
||||
for (const auto i : irange(dim)) {
|
||||
m_shape[i]->verify(view.size(i), "shape", i);
|
||||
}
|
||||
if (this->m_has_strides()) {
|
||||
for (const auto i : stdv::iota(std::size_t{0}, dim)) {
|
||||
m_strides[i]->verify(view.stride(i));
|
||||
if (m_has_strides()) {
|
||||
for (const auto i : irange(dim)) {
|
||||
if (view.size(i) != 1 || !m_strides[i]->has_value()) {
|
||||
// skip stride check for size 1 dimension
|
||||
m_strides[i]->verify(view.stride(i), "stride", i);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
host::RuntimeCheck(view.is_contiguous(), "Tensor is not contiguous as expected");
|
||||
RuntimeCheck(view.is_contiguous(), "Tensor is not contiguous as expected");
|
||||
}
|
||||
// since we may use the same matcher to verify again, we will force to check
|
||||
// since we may double verify, we will force to check
|
||||
m_dtype->verify(view.dtype());
|
||||
m_device->verify(view.device());
|
||||
}
|
||||
|
||||
auto m_init_dtype() -> void {
|
||||
host::RuntimeCheck(!m_has_dtype, "DType already specified");
|
||||
RuntimeCheck(!m_has_dtype, "DType already specified");
|
||||
m_has_dtype = true;
|
||||
}
|
||||
|
||||
auto m_init_device() -> void {
|
||||
host::RuntimeCheck(!m_has_device, "Device already specified");
|
||||
RuntimeCheck(!m_has_device, "Device already specified");
|
||||
m_has_device = true;
|
||||
}
|
||||
|
||||
auto m_has_strides() const -> bool {
|
||||
return !m_strides.empty();
|
||||
}
|
||||
|
||||
@@ -7,7 +7,6 @@
|
||||
|
||||
#include <concepts>
|
||||
#include <cstddef>
|
||||
#include <source_location>
|
||||
#include <type_traits>
|
||||
|
||||
namespace device {
|
||||
@@ -32,60 +31,63 @@ __always_inline __device__ auto offset(const T* ptr, U... offset) -> const void*
|
||||
|
||||
} // namespace pointer
|
||||
|
||||
template <typename T, std::size_t N>
|
||||
struct device_vec {
|
||||
T data[N];
|
||||
};
|
||||
|
||||
} // namespace device
|
||||
|
||||
namespace host {
|
||||
|
||||
inline auto
|
||||
RuntimeDeviceCheck(::cudaError_t error, std::source_location location = std::source_location::current()) -> void {
|
||||
inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) {
|
||||
if (error != ::cudaSuccess) {
|
||||
[[unlikely]];
|
||||
::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error));
|
||||
}
|
||||
}
|
||||
|
||||
inline auto RuntimeCudaCheck(std::source_location location = std::source_location::current()) -> void {
|
||||
inline void RuntimeDeviceCheck(DebugInfo location = {}) {
|
||||
return RuntimeDeviceCheck(::cudaGetLastError(), location);
|
||||
}
|
||||
|
||||
template <auto F>
|
||||
inline void set_smem_once(std::size_t smem_size) {
|
||||
static const auto last_smem_size = [&] {
|
||||
RuntimeDeviceCheck(::cudaFuncSetAttribute(F, ::cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
|
||||
return smem_size;
|
||||
}();
|
||||
RuntimeCheck(
|
||||
smem_size <= last_smem_size,
|
||||
"Dynamic shared memory size exceeds the previously set maximum size: ",
|
||||
last_smem_size,
|
||||
" bytes");
|
||||
}
|
||||
|
||||
struct LaunchKernel {
|
||||
public:
|
||||
explicit LaunchKernel(
|
||||
dim3 grid_dim, dim3 block_dim, DLDevice device, std::size_t dynamic_shared_mem_bytes = 0) noexcept
|
||||
: m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)) {}
|
||||
dim3 grid_dim,
|
||||
dim3 block_dim,
|
||||
DLDevice device,
|
||||
std::size_t dynamic_shared_mem_bytes = 0,
|
||||
DebugInfo location = {}) noexcept
|
||||
: m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)),
|
||||
m_location(location) {}
|
||||
|
||||
explicit LaunchKernel(
|
||||
dim3 grid_dim, dim3 block_dim, cudaStream_t stream, std::size_t dynamic_shared_mem_bytes = 0) noexcept
|
||||
: m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)) {}
|
||||
dim3 grid_dim,
|
||||
dim3 block_dim,
|
||||
cudaStream_t stream,
|
||||
std::size_t dynamic_shared_mem_bytes = 0,
|
||||
DebugInfo location = {}) noexcept
|
||||
: m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)), m_location(location) {}
|
||||
|
||||
LaunchKernel(const LaunchKernel&) = delete;
|
||||
LaunchKernel& operator=(const LaunchKernel&) = delete;
|
||||
|
||||
static auto resolve_device(DLDevice device) -> cudaStream_t {
|
||||
return static_cast<cudaStream_t>(::TVMFFIEnvGetStream(device.device_type, device.device_id));
|
||||
}
|
||||
|
||||
LaunchKernel(const LaunchKernel&) = delete;
|
||||
LaunchKernel& operator=(const LaunchKernel&) = delete;
|
||||
|
||||
template <typename T, typename... Args>
|
||||
auto operator()(T&& kernel, Args&&... args) const -> void {
|
||||
host::RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...));
|
||||
RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...), m_location);
|
||||
}
|
||||
|
||||
private:
|
||||
static auto
|
||||
s_make_config(dim3 grid_dim, dim3 block_dim, cudaStream_t stream, std::size_t smem) -> cudaLaunchConfig_t {
|
||||
static auto s_make_config( // Make a config for kernel launch
|
||||
dim3 grid_dim,
|
||||
dim3 block_dim,
|
||||
cudaStream_t stream,
|
||||
std::size_t smem) -> cudaLaunchConfig_t {
|
||||
auto config = ::cudaLaunchConfig_t{};
|
||||
config.gridDim = grid_dim;
|
||||
config.blockDim = block_dim;
|
||||
@@ -94,8 +96,10 @@ struct LaunchKernel {
|
||||
config.numAttrs = 0;
|
||||
return config;
|
||||
}
|
||||
|
||||
cudaLaunchConfig_t m_config;
|
||||
/// TODO: We can add a queue to store the attributes if needed in the future.
|
||||
const DebugInfo m_location;
|
||||
/// TODO: We can add a queue to store the attributes (e.g. for PDL) if needed in the future.
|
||||
};
|
||||
|
||||
} // namespace host
|
||||
|
||||
@@ -1,23 +1,55 @@
|
||||
#pragma once
|
||||
|
||||
// ref: https://forums.developer.nvidia.com/t/c-20s-source-location-compilation-error-when-using-nvcc-12-1/258026/3
|
||||
#ifdef __CUDACC__
|
||||
#pragma push_macro("__cpp_consteval")
|
||||
#pragma push_macro("_NODISCARD")
|
||||
#pragma push_macro("__builtin_LINE")
|
||||
|
||||
#pragma clang diagnostic push
|
||||
#pragma clang diagnostic ignored "-Wbuiltin-macro-redefined"
|
||||
#define __cpp_consteval 201811L
|
||||
#pragma clang diagnostic pop
|
||||
|
||||
#ifdef _NODISCARD
|
||||
#undef _NODISCARD
|
||||
#define _NODISCARD
|
||||
#endif
|
||||
|
||||
#define consteval constexpr
|
||||
|
||||
#include <source_location>
|
||||
|
||||
#undef consteval
|
||||
#pragma pop_macro("__cpp_consteval")
|
||||
#pragma pop_macro("_NODISCARD")
|
||||
#else
|
||||
#include <source_location>
|
||||
#endif
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
|
||||
#include <concepts>
|
||||
#include <cstddef>
|
||||
#include <ostream>
|
||||
#include <ranges>
|
||||
#include <source_location>
|
||||
#include <sstream>
|
||||
#include <utility>
|
||||
|
||||
namespace host {
|
||||
|
||||
struct DebugInfo : public std::source_location {
|
||||
DebugInfo(std::source_location loc = std::source_location::current()) : std::source_location(loc) {}
|
||||
};
|
||||
|
||||
struct PanicError : public std::runtime_error {
|
||||
public:
|
||||
// copy and move constructors
|
||||
explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {}
|
||||
auto detail() const -> std::string_view {
|
||||
const auto sv = std::string_view{m_message};
|
||||
const auto pos = sv.find(": ");
|
||||
return pos == std::string_view::npos ? sv : sv.substr(pos + 2);
|
||||
auto root_cause() const -> std::string_view {
|
||||
const auto str = std::string_view{m_message};
|
||||
const auto pos = str.find(": ");
|
||||
return pos == std::string_view::npos ? str : str.substr(pos + 2);
|
||||
}
|
||||
|
||||
private:
|
||||
@@ -26,7 +58,7 @@ struct PanicError : public std::runtime_error {
|
||||
|
||||
template <typename... Args>
|
||||
[[noreturn]]
|
||||
inline auto panic(std::source_location location, Args&&... args) -> void {
|
||||
inline auto panic(DebugInfo location, Args&&... args) -> void {
|
||||
std::ostringstream os;
|
||||
os << "Runtime check failed at " << location.file_name() << ":" << location.line();
|
||||
if constexpr (sizeof...(args) > 0) {
|
||||
@@ -40,32 +72,42 @@ inline auto panic(std::source_location location, Args&&... args) -> void {
|
||||
|
||||
template <typename... Args>
|
||||
struct RuntimeCheck {
|
||||
using Loc_t = std::source_location;
|
||||
template <typename Cond>
|
||||
explicit RuntimeCheck(Cond&& condition, Args&&... args, Loc_t location = Loc_t::current()) {
|
||||
if (!condition) {
|
||||
[[unlikely]];
|
||||
::host::panic(location, std::forward<Args>(args)...);
|
||||
}
|
||||
explicit RuntimeCheck(Cond&& condition, Args&&... args, DebugInfo location = {}) {
|
||||
if (condition) return;
|
||||
[[unlikely]] ::host::panic(location, std::forward<Args>(args)...);
|
||||
}
|
||||
template <typename Cond>
|
||||
explicit RuntimeCheck(DebugInfo location, Cond&& condition, Args&&... args) {
|
||||
if (condition) return;
|
||||
[[unlikely]] ::host::panic(location, std::forward<Args>(args)...);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename... Args>
|
||||
struct Panic {
|
||||
explicit Panic(Args&&... args, DebugInfo location = {}) {
|
||||
::host::panic(location, std::forward<Args>(args)...);
|
||||
}
|
||||
explicit Panic(DebugInfo location, Args&&... args) {
|
||||
::host::panic(location, std::forward<Args>(args)...);
|
||||
}
|
||||
[[noreturn]] ~Panic() {
|
||||
std::terminate();
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Cond, typename... Args>
|
||||
explicit RuntimeCheck(Cond&&, Args&&...) -> RuntimeCheck<Args...>;
|
||||
|
||||
template <std::signed_integral T, std::signed_integral U>
|
||||
inline constexpr auto div_ceil(T a, U b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
template <typename Cond, typename... Args>
|
||||
explicit RuntimeCheck(DebugInfo, Cond&&, Args&&...) -> RuntimeCheck<Args...>;
|
||||
|
||||
template <std::unsigned_integral T, std::unsigned_integral U>
|
||||
inline constexpr auto div_ceil(T a, U b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
template <typename... Args>
|
||||
explicit Panic(Args&&...) -> Panic<Args...>;
|
||||
|
||||
inline auto dtype_bytes(DLDataType dtype) -> std::size_t {
|
||||
return static_cast<std::size_t>(dtype.bits / 8);
|
||||
}
|
||||
template <typename... Args>
|
||||
explicit Panic(DebugInfo, Args&&...) -> Panic<Args...>;
|
||||
|
||||
namespace pointer {
|
||||
|
||||
@@ -85,4 +127,26 @@ inline auto offset(const T* ptr, U... offset) -> const void* {
|
||||
|
||||
} // namespace pointer
|
||||
|
||||
template <std::integral T, std::integral U>
|
||||
inline constexpr auto div_ceil(T a, U b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
inline auto dtype_bytes(DLDataType dtype) -> std::size_t {
|
||||
return static_cast<std::size_t>(dtype.bits / 8);
|
||||
}
|
||||
|
||||
namespace stdr = std::ranges;
|
||||
namespace stdv = stdr::views;
|
||||
|
||||
template <std::integral T>
|
||||
inline auto irange(T end) {
|
||||
return stdv::iota(static_cast<T>(0), end);
|
||||
}
|
||||
|
||||
template <std::integral T>
|
||||
inline auto irange(T start, T end) {
|
||||
return stdv::iota(start, end);
|
||||
}
|
||||
|
||||
} // namespace host
|
||||
|
||||
@@ -1,145 +0,0 @@
|
||||
#pragma once
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <type_traits>
|
||||
|
||||
namespace device::warp {
|
||||
|
||||
namespace details {
|
||||
|
||||
template <std::size_t kUnit>
|
||||
inline constexpr auto get_mem_package() {
|
||||
if constexpr (kUnit == 16) {
|
||||
return uint4{};
|
||||
} else if constexpr (kUnit == 8) {
|
||||
return uint2{};
|
||||
} else if constexpr (kUnit == 4) {
|
||||
return uint1{};
|
||||
} else {
|
||||
static_assert(kUnit == 16 || kUnit == 8 || kUnit == 4, "Unsupported memory package size");
|
||||
}
|
||||
}
|
||||
|
||||
inline constexpr auto default_unit_size(std::size_t x) -> std::size_t {
|
||||
if (x % (16 * kWarpThreads) == 0) return 16;
|
||||
if (x % (8 * kWarpThreads) == 0) return 8;
|
||||
if (x % (4 * kWarpThreads) == 0) return 4;
|
||||
return 0; // trigger static assert in _get_mem_package
|
||||
}
|
||||
|
||||
template <std::size_t kBytes, std::size_t kUnit>
|
||||
using mem_package_t = decltype(get_mem_package<kUnit>());
|
||||
|
||||
template <typename T, std::size_t N>
|
||||
struct storage_vec {
|
||||
T data[N];
|
||||
};
|
||||
|
||||
__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
|
||||
uint32_t tmp;
|
||||
asm volatile("ld.global.cs.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
|
||||
return uint1{tmp};
|
||||
}
|
||||
|
||||
__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 {
|
||||
uint32_t tmp0, tmp1;
|
||||
asm volatile("ld.global.cs.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
|
||||
return uint2{tmp0, tmp1};
|
||||
}
|
||||
|
||||
__always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 {
|
||||
uint32_t tmp0, tmp1, tmp2, tmp3;
|
||||
asm volatile("ld.global.cs.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src));
|
||||
return uint4{tmp0, tmp1, tmp2, tmp3};
|
||||
}
|
||||
|
||||
__always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) {
|
||||
uint32_t tmp = value.x;
|
||||
asm volatile("st.global.cs.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
|
||||
}
|
||||
|
||||
__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) {
|
||||
uint32_t tmp0 = value.x;
|
||||
uint32_t tmp1 = value.y;
|
||||
asm volatile("st.global.cs.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
|
||||
}
|
||||
|
||||
__always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) {
|
||||
uint32_t tmp0 = value.x;
|
||||
uint32_t tmp1 = value.y;
|
||||
uint32_t tmp2 = value.z;
|
||||
uint32_t tmp3 = value.w;
|
||||
asm volatile("st.global.cs.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
|
||||
}
|
||||
|
||||
} // namespace details
|
||||
|
||||
template <
|
||||
std::size_t kBytes,
|
||||
std::size_t kUnit = details::default_unit_size(kBytes),
|
||||
std::size_t kThreads = ::device::kWarpThreads>
|
||||
__always_inline __device__ void copy(void* __restrict__ dst, const void* __restrict__ src) {
|
||||
using Package = details::mem_package_t<kBytes, kUnit>;
|
||||
constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
|
||||
constexpr auto kLoopCount = kBytes / kBytesPerLoop;
|
||||
static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
|
||||
|
||||
const auto dst_packed = static_cast<Package*>(dst);
|
||||
const auto src_packed = static_cast<const Package*>(src);
|
||||
const auto lane_id = threadIdx.x % kThreads;
|
||||
|
||||
#pragma unroll kLoopCount
|
||||
for (std::size_t i = 0; i < kLoopCount; ++i) {
|
||||
const auto j = i * kThreads + lane_id;
|
||||
dst_packed[j] = src_packed[j];
|
||||
}
|
||||
}
|
||||
|
||||
template <
|
||||
std::size_t kBytes,
|
||||
std::size_t kUnit = details::default_unit_size(kBytes),
|
||||
std::size_t kThreads = ::device::kWarpThreads>
|
||||
__always_inline __device__ auto load_vec(const void* __restrict__ src) {
|
||||
using Package = details::mem_package_t<kBytes, kUnit>;
|
||||
constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
|
||||
constexpr auto kLoopCount = kBytes / kBytesPerLoop;
|
||||
static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
|
||||
|
||||
const auto src_packed = static_cast<const Package*>(src);
|
||||
const auto lane_id = threadIdx.x % kThreads;
|
||||
details::storage_vec<Package, kLoopCount> vec;
|
||||
|
||||
#pragma unroll kLoopCount
|
||||
for (std::size_t i = 0; i < kLoopCount; ++i) {
|
||||
const auto j = i * kThreads + lane_id;
|
||||
vec.data[i] = details::load_nc(src_packed + j);
|
||||
}
|
||||
|
||||
return vec;
|
||||
}
|
||||
|
||||
template <
|
||||
std::size_t kBytes,
|
||||
std::size_t kUnit = details::default_unit_size(kBytes),
|
||||
std::size_t kThreads = ::device::kWarpThreads,
|
||||
typename Tp>
|
||||
__always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) {
|
||||
using Package = details::mem_package_t<kBytes, kUnit>;
|
||||
constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
|
||||
constexpr auto kLoopCount = kBytes / kBytesPerLoop;
|
||||
static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
|
||||
static_assert(std::is_same_v<Tp, details::storage_vec<Package, kLoopCount>>);
|
||||
|
||||
const auto dst_packed = static_cast<Package*>(dst);
|
||||
const auto lane_id = threadIdx.x % kThreads;
|
||||
|
||||
#pragma unroll kLoopCount
|
||||
for (std::size_t i = 0; i < kLoopCount; ++i) {
|
||||
const auto j = i * kThreads + lane_id;
|
||||
details::store_nc(dst_packed + j, vec.data[i]);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace device::warp
|
||||
14
python/sglang/jit_kernel/test_add_constant.py
Normal file
14
python/sglang/jit_kernel/test_add_constant.py
Normal file
@@ -0,0 +1,14 @@
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.add_constant import add_constant
|
||||
|
||||
|
||||
def main():
|
||||
c = 1024
|
||||
src = torch.arange(0, 1024 + 1, dtype=torch.int32).cuda()
|
||||
dst = add_constant(src, c)
|
||||
assert torch.all(dst == src + c)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -70,6 +70,36 @@ def load_jit(
|
||||
extra_include_paths: List[str] | None = None,
|
||||
build_directory: str | None = None,
|
||||
) -> Module:
|
||||
"""
|
||||
Loading a JIT module from C++/CUDA source files.
|
||||
We define a wrapper as a tuple of (export_name, kernel_name),
|
||||
where `export_name` is the name used to called from Python,
|
||||
and `kernel_name` is the name of the kernel class in C++/CUDA source.
|
||||
|
||||
:param args: Unique marker of the JIT module. Must be distinct for different kernels.
|
||||
:type args: str
|
||||
:param cpp_files: A list of C++ source files.
|
||||
:type cpp_files: List[str] | None
|
||||
:param cuda_files: A list of CUDA source files.
|
||||
:type cuda_files: List[str] | None
|
||||
:param cpp_wrappers: A list of C++ wrappers, defining the export name and kernel name.
|
||||
:type cpp_wrappers: List[Tuple[str, str]] | None
|
||||
:param cuda_wrappers: A list of CUDA wrappers, defining the export name and kernel name.
|
||||
:type cuda_wrappers: List[Tuple[str, str]] | None
|
||||
:param extra_cflags: Extra C++ compiler flags.
|
||||
:type extra_cflags: List[str] | None
|
||||
:param extra_cuda_cflags: Extra CUDA compiler flags.
|
||||
:type extra_cuda_cflags: List[str] | None
|
||||
:param extra_ldflags: Extra linker flags.
|
||||
:type extra_ldflags: List[str] | None
|
||||
:param extra_include_paths: Extra include paths.
|
||||
:type extra_include_paths: List[str] | None
|
||||
:param build_directory: The build directory for JIT compilation.
|
||||
:type build_directory: str | None
|
||||
:return: A just-in-time(JIT) compiled module.
|
||||
:rtype: Module
|
||||
"""
|
||||
|
||||
from tvm_ffi.cpp import load_inline
|
||||
|
||||
cpp_files = cpp_files or []
|
||||
|
||||
Reference in New Issue
Block a user