From 740f57a02c0b0b615f39513ff2f273fbfe071126 Mon Sep 17 00:00:00 2001
From: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com>
Date: Thu, 17 Sep 2026 19:41:52 -0700
Subject: [PATCH] [Spec] Fix CDF boundary handling in
 `TreeSpeculativeSamplingTargetOnly` (#35798)

---
 .../sgl_kernel/speculative/sampling.cuh       |  3 +-
 .../speculative/test_speculative_sampling.py  | 86 ++++++++++++++++++-
 2 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh b/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh
index c4136bd5b..ac46f636b 100644
--- a/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh
+++ b/python/sglang/kernels/jit/include/sgl_kernel/speculative/sampling.cuh
@@ -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;
diff --git a/test/registered/kernels/ops/speculative/test_speculative_sampling.py b/test/registered/kernels/ops/speculative/test_speculative_sampling.py
index b94de4f45..77e16ade1 100644
--- a/test/registered/kernels/ops/speculative/test_speculative_sampling.py
+++ b/test/registered/kernels/ops/speculative/test_speculative_sampling.py
@@ -26,9 +26,9 @@ test_cases = [
     (
         0,  # threshold_single
         0,  # threshold_acc
-        [1, 2, 18, -1, -1, -1, 11, -1, -1, -1, 12, 18],
-        [[0, 1, 2, -1], [6, 10, 11, -1]],
-        [2, 2],
+        [3, -1, -1, 4, 5, 18, 11, -1, -1, -1, 12, 18],
+        [[0, 3, 4, 5], [6, 10, 11, -1]],
+        [3, 2],
     ),
 ]
 
@@ -137,5 +137,85 @@ def test_tree_speculative_sampling_target_only(
     )
 
 
+@pytest.mark.parametrize(
+    "candidate_tokens,target_probabilities,coin,threshold_single,expected_token,expected_accept_token_num",
+    [
+        ([2], [0.0, 1.0, 0.0], 0.0, 0.0, 1, 0),
+        ([2], [0.0, 1.0, 0.0], 0.0, 1.0, 1, 0),
+        ([1, 2, 3], [0.0, 0.25, 0.0, 0.75], 0.25, 1.0, 3, 1),
+        ([1], [0.0, 0.25, 0.75], 0.0, 1.0, 1, 1),
+        ([1], [0.0, 1.0], 1.0 - 2**-24, 1.0, 1, 1),
+    ],
+    ids=[
+        "zero-mass-zero-threshold",
+        "zero-mass-default-threshold",
+        "interior-boundary",
+        "lower-endpoint",
+        "upper-endpoint",
+    ],
+)
+def test_target_only_sampling_cdf_boundaries(
+    candidate_tokens,
+    target_probabilities,
+    coin,
+    threshold_single,
+    expected_token,
+    expected_accept_token_num,
+):
+    num_candidates = len(candidate_tokens)
+    candidates = torch.tensor(
+        [[0, *candidate_tokens]], dtype=torch.int64, device="cuda"
+    )
+    num_draft_tokens = candidates.shape[1]
+    retrive_index = torch.arange(
+        num_draft_tokens, dtype=torch.int64, device="cuda"
+    ).unsqueeze(0)
+    retrive_next_token = torch.full_like(retrive_index, -1)
+    retrive_next_token[0, 0] = 1
+    retrive_next_sibling = torch.full_like(retrive_index, -1)
+    if num_candidates > 1:
+        retrive_next_sibling[0, 1:num_candidates] = torch.arange(
+            2, num_candidates + 1, dtype=torch.int64, device="cuda"
+        )
+
+    target_probs = torch.zeros(
+        (1, num_draft_tokens, len(target_probabilities)),
+        dtype=torch.float32,
+        device="cuda",
+    )
+    target_probs[0, 0] = torch.tensor(
+        target_probabilities, dtype=torch.float32, device="cuda"
+    )
+    target_probs[0, 1:, 0] = 1.0
+    draft_probs = torch.zeros_like(target_probs)
+    predicts = torch.full((num_draft_tokens,), -1, dtype=torch.int32, device="cuda")
+    accept_index = torch.full((1, 2), -1, dtype=torch.int32, device="cuda")
+    accept_token_num = torch.zeros((1,), dtype=torch.int32, device="cuda")
+    coins = torch.zeros((1, num_draft_tokens), dtype=torch.float32, device="cuda")
+    coins[0, 0] = coin
+
+    tree_speculative_sampling_target_only(
+        predicts=predicts,
+        accept_index=accept_index,
+        accept_token_num=accept_token_num,
+        candidates=candidates,
+        retrive_index=retrive_index,
+        retrive_next_token=retrive_next_token,
+        retrive_next_sibling=retrive_next_sibling,
+        uniform_samples=coins,
+        uniform_samples_for_final_sampling=torch.zeros(
+            (1,), dtype=torch.float32, device="cuda"
+        ),
+        target_probs=target_probs,
+        draft_probs=draft_probs,
+        threshold_single=threshold_single,
+        threshold_acc=1.0,
+        deterministic=True,
+    )
+
+    assert predicts[0].item() == expected_token
+    assert accept_token_num.item() == expected_accept_token_num
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))