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

[sgl] update specdec sampling kernel to return valid token ID (#22643)

Browse Source
This commit is contained in:
Tarushii Goel
2026-04-21 20:28:19 -07:00
committed by GitHub
parent f41f1a74a4
commit 3ebf066d13
1 changed files with 15 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/speculative/speculative_sampling.cuh
+15 -3
View File
@@ -125,8 +125,9 @@ __global__ void TreeSpeculativeSamplingTargetOnly(
if (tx == 0) { if (tx == 0) {
temp_storage.block_aggregate.value = sum_relu_q_minus_p; temp_storage.block_aggregate.value = sum_relu_q_minus_p;
} }
// init the first rejected token to (d - 1)
temp_storage.sampled_id = d - 1; temp_storage.sampled_id = d;
temp_storage.last_valid_id = -1;
__syncthreads(); __syncthreads();
sum_relu_q_minus_p = temp_storage.block_aggregate.value; sum_relu_q_minus_p = temp_storage.block_aggregate.value;
DType u = coin * sum_relu_q_minus_p; DType u = coin * sum_relu_q_minus_p;
@@ -156,8 +157,19 @@ __global__ void TreeSpeculativeSamplingTargetOnly(
} }
} }
__syncthreads(); __syncthreads();
// This would happen when u is very close to 1
// and the sum of probabilities is smaller than u
// In this case, we use the last valid index as the sampled id
int sampled_id = temp_storage.sampled_id;
if (sampled_id == d) {
if (temp_storage.last_valid_id == -1) {
sampled_id = d - 1;
} else {
sampled_id = temp_storage.last_valid_id;
}
}
// set the first rejected token // set the first rejected token
predicts[last_accepted_retrive_idx] = temp_storage.sampled_id; predicts[last_accepted_retrive_idx] = sampled_id;
// value at not used indices are undefined // value at not used indices are undefined
} }