minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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/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()),