diff --git a/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh b/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh new file mode 100644 index 000000000..5874ce1ef --- /dev/null +++ b/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh @@ -0,0 +1,38 @@ +#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 new file mode 100644 index 000000000..27665dbb9 --- /dev/null +++ b/python/sglang/jit_kernel/cuda_wait_value.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from functools import lru_cache +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import load_jit + +if TYPE_CHECKING: + import torch + from tvm_ffi.module import Module + + +@lru_cache(maxsize=1) +def _jit_stream_wait_value_module() -> Module: + return load_jit( + "cuda_wait_value", + cuda_files=["cuda_wait_value.cuh"], + cuda_wrappers=[("stream_wait_value", "stream_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()