minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/speculative/eagle_utils.cu
+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);
sgl-kernel/csrc/speculative/speculative_sampling.cu
+2 -3
View File
@@ -42,8 +42,7 @@ void tree_speculative_sampling_target_only(
at::Tensor draft_probs,
double threshold_single,
double threshold_acc,
bool deterministic = true,
int64_t cuda_stream = 0) {
bool deterministic = true) {
CHECK_INPUT(candidates);
CHECK_INPUT(retrive_index);
CHECK_INPUT(retrive_next_token);
@@ -124,7 +123,7 @@ void tree_speculative_sampling_target_only(
CHECK_GE(threshold_acc, 0);
CHECK_GE(1, threshold_acc);
cudaStream_t stream = reinterpret_cast<cudaStream_t>(cuda_stream);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
cudaError_t status = sampling::TreeSpeculativeSamplingTargetOnly<float, int32_t, int64_t>(
static_cast<int32_t*>(predicts.data_ptr()),
static_cast<int32_t*>(accept_index.data_ptr()),