From 3ebf066d131de7b2e0fdaa0c5b4c7c779dcf61c6 Mon Sep 17 00:00:00 2001 From: Tarushii Goel Date: Tue, 21 Apr 2026 20:28:19 -0700 Subject: [PATCH] [sgl] update specdec sampling kernel to return valid token ID (#22643) --- .../csrc/speculative/speculative_sampling.cuh | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/sgl-kernel/csrc/speculative/speculative_sampling.cuh b/sgl-kernel/csrc/speculative/speculative_sampling.cuh index 59f18bc2f..b30b18eeb 100644 --- a/sgl-kernel/csrc/speculative/speculative_sampling.cuh +++ b/sgl-kernel/csrc/speculative/speculative_sampling.cuh @@ -125,8 +125,9 @@ __global__ void TreeSpeculativeSamplingTargetOnly( if (tx == 0) { 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(); sum_relu_q_minus_p = temp_storage.block_aggregate.value; DType u = coin * sum_relu_q_minus_p; @@ -156,8 +157,19 @@ __global__ void TreeSpeculativeSamplingTargetOnly( } } __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 - predicts[last_accepted_retrive_idx] = temp_storage.sampled_id; + predicts[last_accepted_retrive_idx] = sampled_id; // value at not used indices are undefined }