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

[Spec] Fix CDF boundary handling in TreeSpeculativeSamplingTargetOnly (#35798)

Browse Source
This commit is contained in:
Nan Jiang
2026-09-17 19:41:52 -07:00
committed by GitHub
parent 0214954f26
commit 740f57a02c
2 changed files with 85 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh
+2 -1
View File
@@ -77,7 +77,8 @@ __global__ void TreeSpeculativeSamplingTargetOnly(
DType target_prob_single = target_probs[cur_prob_offset + draft_token_id];
prob_acc += target_prob_single;
if (coin <= prob_acc / threshold_acc || target_prob_single >= threshold_single) {
const bool has_target_mass = target_prob_single > DType(0);
if (has_target_mass && (coin < prob_acc / threshold_acc || target_prob_single >= threshold_single)) {
// accept token
prob_acc = 0.;
cur_prob_offset = (bx * num_draft_tokens + cur_index) * d;