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

python/sglang/jit_kernel/csrc/cuda_wait_value.cuh

@@ -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

python/sglang/jit_kernel/cuda_wait_value.py

@@ -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()

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:
+    needs_pad = dim_og % pad_multiple != 0
-        x = torch.nn.functional.pad(x, (0, 8 - dim_og % 8))
+    if needs_pad:
-    dim = x.size(1)
+        x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
 
     out = torch.empty_like(x)
-    module = _jit_hadamard_module(x.dtype)
+    kernel_fn(x, out, scale)
-    module.hadamard_transform(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)
+    return _hadamard_transform_impl(x, scale, 4 * 12, module.hadamard_transform_12n)
-
-    if dim_og % pad_multiple != 0:
-        out = out[:, :dim_og]
-    return out.reshape(shapes_og)
 
 
 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)
+    return _hadamard_transform_impl(x, scale, 4 * 20, module.hadamard_transform_20n)
-
-    if dim_og % pad_multiple != 0:
-        out = out[:, :dim_og]
-    return out.reshape(shapes_og)
 
 
 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)
+    return _hadamard_transform_impl(x, scale, 4 * 28, module.hadamard_transform_28n)
-
-    if dim_og % pad_multiple != 0:
-        out = out[:, :dim_og]
-    return out.reshape(shapes_og)
 
 
 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)
+    return _hadamard_transform_impl(x, scale, 4 * 40, module.hadamard_transform_40n)
-
-    if dim_og % pad_multiple != 0:
-        out = out[:, :dim_og]
-    return out.reshape(shapes_og)