diff --git a/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh b/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh deleted file mode 100644 index 49b274fb1..000000000 --- a/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh +++ /dev/null @@ -1,38 +0,0 @@ -#include - -#include - -#include -#include - -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().with_device().verify(flag); - RuntimeCheck(length.unwrap() >= 1, "wait_flag expects a non-empty tensor."); - - auto* ptr = static_cast(flag.data_ptr()); - const auto stream = LaunchKernel::resolve_device(flag.device()); - - constexpr int blocks = 1; - constexpr int threads = 1; - wait_flag_kernel<<>>(ptr, value); - RuntimeDeviceCheck(cudaGetLastError()); -} - -} // namespace diff --git a/python/sglang/jit_kernel/cuda_wait_value.py b/python/sglang/jit_kernel/cuda_wait_value.py deleted file mode 100644 index e93185e97..000000000 --- a/python/sglang/jit_kernel/cuda_wait_value.py +++ /dev/null @@ -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() diff --git a/python/sglang/jit_kernel/hadamard.py b/python/sglang/jit_kernel/hadamard.py index afe9629be..25930ce94 100644 --- a/python/sglang/jit_kernel/hadamard.py +++ b/python/sglang/jit_kernel/hadamard.py @@ -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)