Add cuda event based on waiting value (#14214)

python/sglang/jit_kernel/csrc/cuda_wait_value.cuh

@@ -0,0 +1,38 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.cuh>
+
+#include <cuda_runtime_api.h>
+
+#include <cstdint>
+
+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

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