[GG] fix(spec-decode): centralize probabilistic draft sampling

tests/v1/worker/test_gpu_autoregressive_speculator.py

@@ -7,10 +7,12 @@
 import torch
 
 from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module
 from vllm.v1.worker.gpu.spec_decode.autoregressive import speculator as spec_module
 from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
     AutoRegressiveSpeculator,
 )
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
 
 
 class _TestSpeculator(AutoRegressiveSpeculator):
@@ -80,3 +82,87 @@ def test_run_model_reuses_tensor_return_for_mtp(monkeypatch):
 
     assert actual_logits_hidden is hidden
     assert actual_feedback_hidden is hidden
+
+
+def test_probabilistic_draft_sampler_owns_disjoint_philox_offset(monkeypatch):
+    captured = {}
+
+    def fake_gumbel_sample(
+        logits,
+        idx_mapping,
+        temperature,
+        seeds,
+        positions,
+        **kwargs,
+    ):
+        captured["positions"] = positions
+        captured.update(kwargs)
+        return torch.zeros(logits.shape[0], dtype=torch.int64)
+
+    monkeypatch.setattr(base_spec_module, "gumbel_sample", fake_gumbel_sample)
+    positions = torch.tensor([12, 99], dtype=torch.int64)
+    active_rows = torch.tensor(2, dtype=torch.int32)
+    speculator = object.__new__(_TestSpeculator)
+    speculator.use_fp64_gumbel = False
+
+    speculator._sample_probabilistic_draft(
+        logits=torch.zeros(2, 5),
+        positions=positions,
+        idx_mapping=torch.arange(2),
+        temperature=torch.ones(2),
+        seeds=torch.tensor([7, 11], dtype=torch.int64),
+        draft_step=torch.tensor(0, dtype=torch.int64),
+        draft_logits=torch.empty(2, 5),
+        active_rows=active_rows,
+    )
+
+    torch.testing.assert_close(
+        captured["positions"], positions + 1 + (1 << 30), rtol=0, atol=0
+    )
+    assert captured["apply_temperature"] is True
+    assert captured["output_processed_logits_active_rows"] is active_rows
+
+
+def test_ar_probabilistic_draft_uses_shared_sampler(monkeypatch):
+    captured = {}
+
+    def fake_sample_probabilistic_draft(
+        self,
+        logits,
+        positions,
+        idx_mapping,
+        temperature,
+        seeds,
+        draft_step,
+        draft_logits,
+        active_rows=None,
+    ):
+        captured["positions"] = positions
+        captured["active_rows"] = active_rows
+        return torch.zeros(logits.shape[0], dtype=torch.int64)
+
+    monkeypatch.setattr(
+        DraftModelSpeculator,
+        "_sample_probabilistic_draft",
+        fake_sample_probabilistic_draft,
+    )
+    speculator = object.__new__(_TestSpeculator)
+    speculator.model = SimpleNamespace(
+        compute_logits=lambda hidden_states: torch.zeros(hidden_states.shape[0], 5)
+    )
+    speculator.active_num_reqs = torch.tensor(2, dtype=torch.int32)
+    speculator.use_fp64_gumbel = False
+    positions = torch.tensor([12, 99], dtype=torch.int64)
+
+    speculator.sample_draft(
+        hidden_states=torch.zeros(2, 3),
+        positions=positions,
+        idx_mapping=torch.arange(2),
+        temperature=torch.ones(2),
+        seeds=torch.tensor([7, 11], dtype=torch.int64),
+        draft_step=torch.tensor(0, dtype=torch.int64),
+        draft_logits=torch.empty(2, 5),
+    )
+
+    assert captured["positions"] is positions
+    assert captured["active_rows"] is speculator.active_num_reqs

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

@@ -18,7 +18,6 @@
 )
 from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
     SpeculatorCudaGraphManager,
 )
@@ -369,19 +368,15 @@ def sample_draft(
     ) -> torch.Tensor:
         logits = self.model.compute_logits(hidden_states)
         if draft_logits is not None:
-            # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
-            # used for draft and target sampling.
-            return gumbel_sample(
-                logits,
-                idx_mapping,
-                temperature,
-                seeds,
-                positions + 1,
-                apply_temperature=True,
-                output_processed_logits=draft_logits,
-                output_processed_logits_col=draft_step,
-                output_processed_logits_active_rows=self.active_num_reqs,
-                use_fp64=self.use_fp64_gumbel,
+            return self._sample_probabilistic_draft(
+                logits=logits,
+                positions=positions,
+                idx_mapping=idx_mapping,
+                temperature=temperature,
+                seeds=seeds,
+                draft_step=draft_step,
+                draft_logits=draft_logits,
+                active_rows=self.active_num_reqs,
             )
         else:
             return logits.argmax(dim=-1)

vllm/v1/worker/gpu/spec_decode/dspark/speculator.py

@@ -32,15 +32,13 @@
 from vllm.config.compilation import CUDAGraphMode
 from vllm.triton_utils import triton
 from vllm.v1.worker.gpu.input_batch import InputBatch
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
 from vllm.v1.worker.gpu.spec_decode.dspark.capacity import (
     build_sps_table,
     compute_draft_token_capacity_from_confidence,
 )
 from vllm.v1.worker.gpu.spec_decode.dspark.online_sts import DSparkOnlineSTS
 from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model
-from vllm.v1.worker.gpu.spec_decode.utils import draft_gumbel_pos
 
 
 class DSparkSpeculator(DFlashSpeculator):
@@ -252,21 +250,17 @@ def _sample_sequential(
                     buf = self._draft_scatter_buf[:num_reqs]
                     buf.index_copy_(1, self._d2t_scatter_index, logits_i.to(buf.dtype))
                     logits_i = buf
-                # sample_pos is the predicted token's position Q;
-                # draft_gumbel_pos keys the (salted) draft Gumbel stream by
-                # positions + 1, so pass Q-2 to get a key unique per
-                # predicted position and disjoint from the rejection
-                # sampler's acceptance/recovery keys.
-                draft_sampled_i = gumbel_sample(
-                    logits_i,
-                    idx_map[:, i],
-                    self.temperature,
-                    self.seeds,
-                    draft_gumbel_pos(sample_pos[:, i] - 2),
-                    apply_temperature=True,
-                    output_processed_logits=self.draft_logits,
-                    output_processed_logits_col=self._step_cols[i],
-                    use_fp64=self.use_fp64_gumbel,
+                # sample_pos is the predicted token's position Q. The shared
+                # sampler adds one before salting, so Q-2 produces a unique
+                # draft key for each predicted position.
+                draft_sampled_i = self._sample_probabilistic_draft(
+                    logits=logits_i,
+                    positions=sample_pos[:, i] - 2,
+                    idx_mapping=idx_map[:, i],
+                    temperature=self.temperature,
+                    seeds=self.seeds,
+                    draft_step=self._step_cols[i],
+                    draft_logits=self.draft_logits,
                 )
             else:
                 draft_sampled_i = self.model.map_draft_to_target(

vllm/v1/worker/gpu/spec_decode/speculator.py

@@ -353,6 +353,31 @@ def _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor:
         logits = self.model.compute_logits(hidden_states)
         return logits.argmax(dim=-1)
 
+    def _sample_probabilistic_draft(
+        self,
+        logits: torch.Tensor,
+        positions: torch.Tensor,
+        idx_mapping: torch.Tensor,
+        temperature: torch.Tensor,
+        seeds: torch.Tensor,
+        draft_step: torch.Tensor,
+        draft_logits: torch.Tensor,
+        active_rows: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Sample a draft from a stream disjoint from verifier recovery."""
+        return gumbel_sample(
+            logits,
+            idx_mapping,
+            temperature,
+            seeds,
+            draft_gumbel_pos(positions),
+            apply_temperature=True,
+            output_processed_logits=draft_logits,
+            output_processed_logits_col=draft_step,
+            output_processed_logits_active_rows=active_rows,
+            use_fp64=self.use_fp64_gumbel,
+        )
+
     def sample_draft(
         self,
         hidden_states: torch.Tensor,
@@ -365,16 +390,14 @@ def sample_draft(
     ) -> torch.Tensor:
         if draft_logits is not None:
             logits = self.model.compute_logits(hidden_states)
-            return gumbel_sample(
-                logits,
-                idx_mapping,
-                temperature,
-                seeds,
-                draft_gumbel_pos(positions),
-                apply_temperature=True,
-                output_processed_logits=draft_logits,
-                output_processed_logits_col=draft_step,
-                use_fp64=self.use_fp64_gumbel,
+            return self._sample_probabilistic_draft(
+                logits=logits,
+                positions=positions,
+                idx_mapping=idx_mapping,
+                temperature=temperature,
+                seeds=seeds,
+                draft_step=draft_step,
+                draft_logits=draft_logits,
             )
         return self._greedy_sample_draft(hidden_states)