From b425ccc1121b6a19fb20e13ec2453951b5ed4946 Mon Sep 17 00:00:00 2001
From: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com>
Date: Fri, 25 Sep 2026 03:50:29 +0000
Subject: [PATCH] fix(speculative): keep final chain draws on positive support

Mask zero-mass and padded CDF lanes. If endpoint rounding leaves a finite positive distribution without a match, select its last positive token. Keep malformed inputs outside the supported sampling contract.

Add analytic synthetic regressions for endpoint gaps, zero tails, partial blocks, uniform boundaries, tiny mass, malformed no-match inputs, and graph replay.
---
 .../ops/speculative/reject_sampling.py        |  51 ++++-
 .../speculative/test_chain_sampling_cdf.py    | 198 ++++++++++++++++++
 2 files changed, 241 insertions(+), 8 deletions(-)
 create mode 100644 test/registered/kernel/speculative/test_chain_sampling_cdf.py

diff --git a/python/sglang/kernels/ops/speculative/reject_sampling.py b/python/sglang/kernels/ops/speculative/reject_sampling.py
index fd32ca3ea1b9..ae58d89db365 100644
--- a/python/sglang/kernels/ops/speculative/reject_sampling.py
+++ b/python/sglang/kernels/ops/speculative/reject_sampling.py
@@ -118,9 +118,7 @@ def speculative_sampling_classic_kernel(
 
         norm_sum += tl.sum(val)
 
-    # Pass 2: CDF. Degenerate residual (norm_sum == 0, i.e. p == q everywhere on
-    # rejection) leaves the cumsum at 0 <= target_u, so final_token falls back to
-    # VOCAB_SIZE - 1; acceptable since this case is numerically near-impossible.
+    # Reduction and prefix scan can round differently at the endpoint.
     target_u = coin_final * norm_sum
     cum_sum = 0.0
     final_token = VOCAB_SIZE - 1
@@ -147,16 +145,50 @@ def speculative_sampling_classic_kernel(
             block_cumsum = tl.cumsum(val, axis=0)
             total_cumsum = cum_sum + block_cumsum
 
-            candidates_mask = total_cumsum > target_u
-            has_match = tl.max(candidates_mask, axis=0)
+            candidates_mask = mask & (val > 0.0) & (total_cumsum > target_u)
+            match_idx = tl.min(tl.where(candidates_mask, v_offsets, VOCAB_SIZE), axis=0)
 
-            if has_match:
-                match_idx = tl.argmax(candidates_mask.to(tl.int32), axis=0)
-                final_token = v_start + match_idx
+            if match_idx < VOCAB_SIZE:
+                final_token = match_idx
                 found = 1
 
             cum_sum += tl.sum(val)
 
+    # Assign a rounding gap only to positive support of a valid distribution.
+    if (
+        (found == 0)
+        & (norm_sum > 0.0)
+        & (norm_sum < float("inf"))
+        & (coin_final >= 0.0)
+        & (coin_final < 1.0)
+    ):
+        last_positive = -1
+        valid_probs = 1
+        for block_idx in range(tl.cdiv(VOCAB_SIZE, BLOCK_V) - 1, -1, -1):
+            v_offsets = block_idx * BLOCK_V + tl.arange(0, BLOCK_V)
+            mask = v_offsets < VOCAB_SIZE
+            p_val = tl.load(tp_base_ptr + v_offsets * stride_tp_v, mask=mask, other=0.0)
+            valid = (p_val >= 0.0) & (p_val < float("inf"))
+            if all_drafts_accepted:
+                val = p_val
+            else:
+                q_val = tl.load(
+                    dp_base_ptr_safe + v_offsets * stride_dp_v, mask=mask, other=0.0
+                )
+                # Keep the residual identical to both forward passes.
+                q_val = tl.where(q_val == q_val, q_val, 0.0)
+                diff = p_val - q_val
+                val = tl.where(diff > 0.0, diff, 0.0)
+            valid_probs = valid_probs & tl.min(
+                tl.where(mask, valid, True).to(tl.int32), axis=0
+            )
+            if last_positive < 0:
+                last_positive = tl.max(
+                    tl.where(mask & (val > 0.0), v_offsets, -1), axis=0
+                )
+        if valid_probs & (last_positive >= 0):
+            final_token = last_positive
+
     tl.store(Predicts + last_accepted_global_idx, final_token)
 
 
@@ -176,6 +208,9 @@ def chain_speculative_sampling_triton(
     threshold_acc,
     deterministic,  # not used
 ):
+    """Final draws require finite nonnegative p, positive finite sampling mass,
+    and a uniform in [0, 1). Malformed inputs are unsupported.
+    """
     batch_size, num_slots = candidates.shape
     vocab_size = target_probs.shape[-1]
 
diff --git a/test/registered/kernel/speculative/test_chain_sampling_cdf.py b/test/registered/kernel/speculative/test_chain_sampling_cdf.py
new file mode 100644
index 000000000000..048fccba4092
--- /dev/null
+++ b/test/registered/kernel/speculative/test_chain_sampling_cdf.py
@@ -0,0 +1,198 @@
+"""Final chain draws must remain on positive in-vocabulary probability support."""
+
+import os
+import unittest
+
+import torch
+
+from sglang.kernels.ops.speculative.reject_sampling import (
+    chain_speculative_sampling_triton,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+INTERPRET = os.environ.get("TRITON_INTERPRET") == "1"
+DEVICE = "cpu" if INTERPRET else "cuda"
+
+
+def _endpoint_mass(kind, vocab_size, offset=0):
+    p = torch.zeros(vocab_size, dtype=torch.float32, device=DEVICE)
+    if kind == "equal":
+        # Twelve rounded copies of 1/12 can have distinct reduction/scan sums.
+        p[offset : offset + 12] = 1.0 / 12.0
+        last_positive = offset + 11
+    else:
+        # The small terms sum to one ULP of 0.5, but each ties down when added
+        # to 0.5 alone. The four binary fractions sum to exactly 1 in float64.
+        p[offset] = p[offset + 8] = 2.0**-25
+        p[offset + 4] = 0.5
+        p[offset + 128] = 0.5 - 2.0**-24
+        last_positive = offset + 128
+    return p, last_positive
+
+
+def _chain_inputs(p, uniform, mode):
+    num_slots = 1 if mode == "target" else 3
+    proposal = p.numel() - 1
+    target = p[None, None].repeat(1, num_slots, 1)
+    draft = torch.zeros((1, num_slots - 1, p.numel()), device=DEVICE)
+    if mode != "target":
+        draft[:, :, proposal] = 1.0
+    if mode == "accept":
+        target[:, :-1] = draft
+    return dict(
+        predicts=torch.full((num_slots,), -1, dtype=torch.int32, device=DEVICE),
+        accept_index=torch.full((1, num_slots), -1, dtype=torch.int64, device=DEVICE),
+        accept_token_num=torch.full((1,), -1, dtype=torch.int32, device=DEVICE),
+        candidates=torch.full(
+            (1, num_slots), proposal, dtype=torch.int64, device=DEVICE
+        ),
+        retrive_index=torch.arange(num_slots, dtype=torch.int64, device=DEVICE)[None],
+        retrive_next_token=None,
+        retrive_next_sibling=None,
+        uniform_samples=torch.full((1, num_slots - 1), 0.5, device=DEVICE),
+        uniform_samples_for_final_sampling=torch.tensor([uniform], device=DEVICE),
+        target_probs=target,
+        draft_probs=draft,
+        threshold_single=1.0,
+        threshold_acc=1.0,
+        deterministic=True,
+    )
+
+
+@unittest.skipUnless(INTERPRET or torch.cuda.is_available(), "CUDA is required")
+class TestChainSamplingCDF(CustomTestCase):
+    def _check_draw(self, kwargs, expected, mode):
+        num_correct = 2 if mode == "accept" else 0
+        self.assertEqual(kwargs["accept_token_num"].item(), num_correct)
+        bonus = kwargs["predicts"][num_correct].item()
+        p = kwargs["target_probs"][0, num_correct]
+        mass = p if mode != "reject" else (p - kwargs["draft_probs"][0, 0]).clamp_min(0)
+        self.assertTrue(0 <= bonus < p.numel(), bonus)
+        self.assertGreater(mass[bonus].item(), 0.0, (bonus, mass[bonus].item()))
+        self.assertEqual(bonus, expected)
+        self.assertEqual(
+            kwargs["accept_index"][0, : num_correct + 1].tolist(),
+            list(range(num_correct + 1)),
+        )
+        if num_correct:
+            self.assertEqual(
+                kwargs["predicts"][:num_correct].tolist(),
+                [p.numel() - 1] * num_correct,
+            )
+
+    def test_rounding_gap_selects_last_positive(self):
+        """Endpoint roundoff must not emit a zero-mass tail or vocabulary end."""
+        uniform = float.fromhex("0x1.fffffep-1")
+        for kind in ("equal", "binary"):
+            for vocab_size, offset in ((4096, 0), (8193, 0), (4226, 4096)):
+                for mode in ("target", "accept", "reject"):
+                    with self.subTest(
+                        kind=kind, vocab=vocab_size, offset=offset, mode=mode
+                    ):
+                        p, expected = _endpoint_mass(kind, vocab_size, offset)
+                        self.assertLess(expected, vocab_size - 1)
+                        kwargs = _chain_inputs(p, uniform, mode)
+                        chain_speculative_sampling_triton(**kwargs)
+                        self._check_draw(kwargs, expected, mode)
+
+    def test_uniform_boundaries_and_partial_block(self):
+        """A CDF tie advances to the next positive lane across block boundaries."""
+        p = torch.zeros(4101, device=DEVICE)
+        p[17], p[4095], p[4096] = 0.25, 0.5, 0.25
+        for uniform, expected in (
+            (0.0, 17),
+            (float.fromhex("0x1.fffffep-3"), 17),
+            (0.25, 4095),
+            (0.75, 4096),
+            (float.fromhex("0x1.fffffep-1"), 4096),
+        ):
+            for mode in ("target", "accept", "reject"):
+                with self.subTest(uniform=uniform, mode=mode):
+                    kwargs = _chain_inputs(p, uniform, mode)
+                    chain_speculative_sampling_triton(**kwargs)
+                    self._check_draw(kwargs, expected, mode)
+
+    def test_tiny_positive_tail(self):
+        """Support is strictly positive mass, without an absolute epsilon cutoff."""
+        p = torch.zeros(4101, device=DEVICE)
+        p[4096] = p[4099] = 2.0**-100
+        for uniform, expected in ((0.0, 4096), (0.5, 4099)):
+            for mode in ("target", "accept", "reject"):
+                with self.subTest(uniform=uniform, mode=mode):
+                    kwargs = _chain_inputs(p, uniform, mode)
+                    chain_speculative_sampling_triton(**kwargs)
+                    self._check_draw(kwargs, expected, mode)
+
+    def test_no_match_malformed_inputs_do_not_use_support_fallback(self):
+        """Unsupported no-match inputs must not be turned into a support draw."""
+        for case in (
+            "zero",
+            "zero_residual",
+            "nan",
+            "inf",
+            "negative_total",
+            "one",
+            "nan_u",
+            "inf_u",
+        ):
+            with self.subTest(case=case):
+                p = torch.zeros(4097, device=DEVICE)
+                p[17] = 1.0
+                uniform = 0.5
+                if case == "zero":
+                    p.zero_()
+                elif case == "nan":
+                    p[0] = float("nan")
+                elif case == "inf":
+                    p[0] = float("inf")
+                elif case == "negative_total":
+                    p[0], p[17] = -0.2, 0.1
+                elif case == "one":
+                    uniform = 1.0
+                elif case == "nan_u":
+                    uniform = float("nan")
+                elif case == "inf_u":
+                    uniform = float("inf")
+                mode = "reject" if case == "zero_residual" else "target"
+                kwargs = _chain_inputs(p, uniform, mode)
+                if case == "zero_residual":
+                    kwargs["draft_probs"].copy_(p[None, None])
+                chain_speculative_sampling_triton(**kwargs)
+                self.assertEqual(kwargs["predicts"][0].item(), p.numel() - 1)
+
+    def test_fallback_validates_blocks_before_last_support(self):
+        """Finding late support must not skip malformed target mass before it."""
+        for invalid in (-1.0, float("nan")):
+            with self.subTest(invalid=invalid):
+                p, _ = _endpoint_mass("binary", 4226, 4096)
+                p[7] = invalid
+                kwargs = _chain_inputs(p, float.fromhex("0x1.fffffep-1"), "reject")
+                chain_speculative_sampling_triton(**kwargs)
+                self.assertEqual(kwargs["accept_token_num"].item(), 0)
+                self.assertEqual(kwargs["predicts"][0].item(), p.numel() - 1)
+
+    @unittest.skipIf(INTERPRET, "CUDA graph replay requires CUDA execution")
+    def test_graph_replay_uses_current_support(self):
+        p, expected = _endpoint_mass("binary", 8193)
+        kwargs = _chain_inputs(p, float.fromhex("0x1.fffffep-1"), "reject")
+        warmup = torch.cuda.Stream()
+        warmup.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(warmup):
+            chain_speculative_sampling_triton(**kwargs)
+        torch.cuda.current_stream().wait_stream(warmup)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            chain_speculative_sampling_triton(**kwargs)
+        for offset in (0, 4096):
+            p, expected = _endpoint_mass("binary", 8193, offset)
+            kwargs["target_probs"].copy_(p[None, None])
+            kwargs["predicts"].fill_(-1)
+            graph.replay()
+            self._check_draw(kwargs, expected, "reject")
+
+
+if __name__ == "__main__":
+    unittest.main()