chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

refactor(jit_kernel): reduce duplication and separate test code (#19323)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-02-26 18:30:49 +08:00
committed by GitHub
parent a7152df2e3
commit 74c8e7b215
3 changed files with 22 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
python/sglang/jit_kernel/csrc/cuda_wait_value.cuh
View File

@@ -1,38 +0,0 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.cuh>
#include <cstdint>
#include <cuda_runtime_api.h>
namespace {
__global__ void wait_flag_kernel(const int32_t* flag, int32_t target) {
const volatile int32_t* vflag = (volatile const int32_t*)flag;
while (*vflag != target) {
#if __CUDA_ARCH__ >= 700
__nanosleep(100);
#else
// Note: This falls back to an inefficient busy-wait on pre-Volta architectures.
#endif
}
}
auto stream_wait_value(const tvm::ffi::TensorView flag, std::int32_t value) -> void {
using namespace host;
auto length = SymbolicSize{"length"};
TensorMatcher({length}).with_dtype<int32_t>().with_device<kDLCUDA>().verify(flag);
RuntimeCheck(length.unwrap() >= 1, "wait_flag expects a non-empty tensor.");
auto* ptr = static_cast<std::int32_t*>(flag.data_ptr());
const auto stream = LaunchKernel::resolve_device(flag.device());
constexpr int blocks = 1;
constexpr int threads = 1;
wait_flag_kernel<<<blocks, threads, 0, stream>>>(ptr, value);
RuntimeDeviceCheck(cudaGetLastError());
}
} // namespace

78
python/sglang/jit_kernel/cuda_wait_value.py
View File

@@ -1,78 +0,0 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from sglang.jit_kernel.utils import cache_once, load_jit
if TYPE_CHECKING:
import torch
from tvm_ffi.module import Module
@cache_once
def _jit_stream_wait_value_module() -> Module:
return load_jit(
"cuda_wait_value",
cuda_files=["cuda_wait_value.cuh"],
cuda_wrappers=[("stream_wait_value", "cuda_wait_value")],
)
def stream_wait_value(flag: torch.Tensor, value: int) -> None:
module = _jit_stream_wait_value_module()
module.stream_wait_value(flag, value)
class Event:
def __init__(self) -> None:
self.flag = torch.zeros(1, dtype=torch.int32, device="cuda")
def record(self, value: int = 1) -> None:
self.flag[0] = value
def wait(self, value: int = 1) -> None:
stream_wait_value(self.flag, value)
def test_wait_before_record(event: Event | torch.cuda.Event):
stream_a = torch.cuda.Stream()
stream_b = torch.cuda.Stream()
with torch.cuda.stream(stream_a):
event.wait()
stream_a.synchronize()
with torch.cuda.stream(stream_b):
event.record()
def main():
import threading
import time
block_thead = threading.Thread(
target=test_wait_before_record, args=(Event(),), daemon=True
)
block_thead.start()
non_block_thread = threading.Thread(
target=test_wait_before_record, args=(torch.cuda.Event(),)
)
non_block_thread.start()
print("Checking if custom Event blocks the stream...", flush=True)
for _ in range(5):
print(f"{block_thead.is_alive()=}, {non_block_thread.is_alive()=}", flush=True)
time.sleep(1)
assert block_thead.is_alive(), "Custom Event did not block as expected"
assert not non_block_thread.is_alive(), "torch.cuda.Event should not block"
print("=" * 40)
print("Test completed successfully.")
if __name__ == "__main__":
main()

107
python/sglang/jit_kernel/hadamard.py
View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Callable
import torch
@@ -29,9 +29,14 @@ def _jit_hadamard_module(dtype: torch.dtype) -> Module:
)
def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
def _hadamard_transform_impl(
x: torch.Tensor,
scale: float,
pad_multiple: int,
kernel_fn: Callable,
) -> torch.Tensor:
if not x.is_cuda:
raise RuntimeError("hadamard_transform only supports CUDA tensors")
raise RuntimeError(f"{kernel_fn.__name__} only supports CUDA tensors")
shapes_og = x.size()
dim_og = x.size(-1)
@@ -39,106 +44,38 @@ def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
if x.stride(-1) != 1:
x = x.contiguous()
if dim_og % 8 != 0:
x = torch.nn.functional.pad(x, (0, 8 - dim_og % 8))
dim = x.size(1)
needs_pad = dim_og % pad_multiple != 0
if needs_pad:
x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
out = torch.empty_like(x)
module = _jit_hadamard_module(x.dtype)
module.hadamard_transform(x, out, scale)
kernel_fn(x, out, scale)
if dim_og % 8 != 0:
if needs_pad:
out = out[:, :dim_og]
return out.reshape(shapes_og)
def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
module = _jit_hadamard_module(x.dtype)
return _hadamard_transform_impl(x, scale, 8, module.hadamard_transform)
def hadamard_transform_12n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
if not x.is_cuda:
raise RuntimeError("hadamard_transform_12n only supports CUDA tensors")
shapes_og = x.size()
dim_og = x.size(-1)
x = x.reshape(-1, dim_og)
if x.stride(-1) != 1:
x = x.contiguous()
pad_multiple = 4 * 12
if dim_og % pad_multiple != 0:
x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
out = torch.empty_like(x)
module = _jit_hadamard_module(x.dtype)
module.hadamard_transform_12n(x, out, scale)
if dim_og % pad_multiple != 0:
out = out[:, :dim_og]
return out.reshape(shapes_og)
return _hadamard_transform_impl(x, scale, 4 * 12, module.hadamard_transform_12n)
def hadamard_transform_20n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
if not x.is_cuda:
raise RuntimeError("hadamard_transform_20n only supports CUDA tensors")
shapes_og = x.size()
dim_og = x.size(-1)
x = x.reshape(-1, dim_og)
if x.stride(-1) != 1:
x = x.contiguous()
pad_multiple = 4 * 20
if dim_og % pad_multiple != 0:
x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
out = torch.empty_like(x)
module = _jit_hadamard_module(x.dtype)
module.hadamard_transform_20n(x, out, scale)
if dim_og % pad_multiple != 0:
out = out[:, :dim_og]
return out.reshape(shapes_og)
return _hadamard_transform_impl(x, scale, 4 * 20, module.hadamard_transform_20n)
def hadamard_transform_28n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
if not x.is_cuda:
raise RuntimeError("hadamard_transform_28n only supports CUDA tensors")
shapes_og = x.size()
dim_og = x.size(-1)
x = x.reshape(-1, dim_og)
if x.stride(-1) != 1:
x = x.contiguous()
pad_multiple = 4 * 28
if dim_og % pad_multiple != 0:
x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
out = torch.empty_like(x)
module = _jit_hadamard_module(x.dtype)
module.hadamard_transform_28n(x, out, scale)
if dim_og % pad_multiple != 0:
out = out[:, :dim_og]
return out.reshape(shapes_og)
return _hadamard_transform_impl(x, scale, 4 * 28, module.hadamard_transform_28n)
def hadamard_transform_40n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
if not x.is_cuda:
raise RuntimeError("hadamard_transform_40n only supports CUDA tensors")
shapes_og = x.size()
dim_og = x.size(-1)
x = x.reshape(-1, dim_og)
if x.stride(-1) != 1:
x = x.contiguous()
pad_multiple = 4 * 40
if dim_og % pad_multiple != 0:
x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
out = torch.empty_like(x)
module = _jit_hadamard_module(x.dtype)
module.hadamard_transform_40n(x, out, scale)
if dim_og % pad_multiple != 0:
out = out[:, :dim_og]
return out.reshape(shapes_og)
return _hadamard_transform_impl(x, scale, 4 * 40, module.hadamard_transform_40n)