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

[JIT] Refactor dtype traits into DTypeTrait and unify warp reductions (#30838)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: BBuf <xiaoyu.zhang@radixark.ai>
Co-authored-by: jessiewei7 <jessiewei747@gmail.com>
Co-authored-by: root <root@GPUC5A6.maas>
This commit is contained in:
DarkSharpness
2026-07-18 10:07:18 +08:00
committed by GitHub
co-authored by Claude Fable 5 BBuf jessiewei7 root
parent e48eabbeee
commit 67e7f8d13a
19 changed files with 930 additions and 507 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs/developer_guide/development_jit_kernel_guide.mdx
+15 -10
View File
@@ -61,19 +61,24 @@ void test() {
#### 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.
`CHECK_HOST` is the preferred runtime check: stream-style, and zero overhead when the
check passes — the message expressions are only evaluated on failure.
`RuntimeCheck` is the function-style alternative; note its message arguments are always
evaluated, even when the check passes.
`RuntimeDeviceCheck` verifies the status of the last kernel launch, and `CHECK_CUDA`
is its stream-style equivalent for checking a `cudaError_t` with extra context.
```C++ Example
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
void test() {
CHECK_HOST(1 + 1 == 2) << 1 + 1 << " != " << 2; // preferred
host::RuntimeCheck(1 + 1 == 2, 1 + 1, " != ", 2);
host::RuntimeDeviceCheck();
// check the provided `cudaError_t`
host::RuntimeDeviceCheck(cudaGetLastError());
CHECK_CUDA(cudaGetLastError()) << "after my_kernel launch";
}
```
@@ -161,7 +166,7 @@ Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](https://github.com/
```cpp Example
#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 <sgl_kernel/utils.h> // For div_ceil, CHECK_HOST
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
@@ -199,8 +204,8 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
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);
// some extra runtime checks using CHECK_HOST
CHECK_HOST(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*/)(
@@ -289,12 +294,12 @@ and its key APIs.
<tr>
<td><code>utils.h</code></td>
<td><code>host</code></td>
<td>Host-side essentials: <code>RuntimeCheck</code>, <code>Panic</code>, <code>div_ceil</code>, <code>irange</code></td>
<td>Host-side essentials: <code>RuntimeCheck</code>, <code>CHECK_HOST(cond) &lt;&lt; ...</code>, <code>Panic</code>, <code>div_ceil</code>, <code>irange</code></td>
</tr>
<tr>
<td><code>utils.cuh</code></td>
<td><code>device</code> / <code>host</code></td>
<td>Type aliases (<code>fp16_t</code>, <code>bf16_t</code>, ...), <code>SGL_DEVICE</code> macro, PDL helpers, <code>LaunchKernel</code>, <code>RuntimeDeviceCheck</code></td>
<td>Type aliases (<code>fp16_t</code>, <code>bf16_t</code>, ...), <code>SGL_DEVICE</code> macro, PDL helpers, <code>LaunchKernel</code>, <code>RuntimeDeviceCheck</code>, <code>CHECK_CUDA(expr) &lt;&lt; ...</code></td>
</tr>
<tr>
<td><code>source_location.h</code></td>
@@ -347,7 +352,7 @@ and its key APIs.
<tr>
<td><code>type.cuh</code></td>
<td>(global) / <code>device</code></td>
<td><code>dtype_trait&lt;T&gt;</code>, <code>packed_t&lt;T&gt;</code>, <code>device::cast&lt;To&gt;(from)</code></td>
<td><code>DTypeTrait&lt;T&gt;</code>, <code>packed_t&lt;T&gt;</code>, <code>device::cast&lt;To&gt;(from)</code></td>
</tr>
</tbody>
</table>
@@ -390,7 +395,7 @@ and its key APIs.
<tr>
<td><code>warp.cuh</code></td>
<td><code>device::warp</code></td>
<td><code>reduce_sum</code>, <code>reduce_max</code> via <code>__shfl_xor_sync</code></td>
<td><code>reduce&lt;Op, kNumThreads, kInner&gt;</code> (SUM/MAX/MIN, grouped or inter-group) and <code>reduce_sum</code> / <code>reduce_max</code> / <code>reduce_min</code> wrappers via <code>__shfl_xor_sync</code></td>
</tr>
<tr>
<td><code>cta.cuh</code></td>