move all get_stream in sgl_kernel to c++ to reduce the launch overhead (#12521)

This commit is contained in:
Lianmin Zheng
2025-11-02 13:15:05 -08:00
committed by GitHub
parent c9db79117f
commit 20315697f4
20 changed files with 61 additions and 94 deletions
+2 -3
View File
@@ -328,8 +328,7 @@ void verify_tree_greedy(
at::Tensor retrive_index,
at::Tensor retrive_next_token,
at::Tensor retrive_next_sibling,
at::Tensor target_predict,
int64_t cuda_stream = 0) {
at::Tensor target_predict) {
CHECK_INPUT(candidates);
CHECK_INPUT(retrive_index);
CHECK_INPUT(retrive_next_token);
@@ -389,7 +388,7 @@ void verify_tree_greedy(
throw std::runtime_error("Expected 'target_predict' to be of type long (torch.int64).");
}
cudaStream_t stream = reinterpret_cast<cudaStream_t>(cuda_stream);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
dim3 grid(batch_size);
dim3 block(1);