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