[Spec Decode][Sampling] Batch-sharded draft logits (#63554)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Co-authored-by: vllm-agent <claw@inferact.ai>

tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py

@@ -99,6 +99,8 @@ def _speculator(temp: float, seed: int):
     s._draft_scatter_buf = None
     s.use_fp64_gumbel = False
     s.draft_logits = torch.zeros(MAX_REQS, N_SPEC, VOCAB, dtype=bf16, device=dev)
+    s.draft_logits_num_shards = 1
+    s.draft_logits_shard_idx = 0
     s.temperature = torch.full((MAX_REQS,), temp, dtype=torch.float32, device=dev)
     s.seeds = torch.randint(0, 2**62, (MAX_REQS,), generator=g, device=dev)
     s._step_cols = torch.arange(N_SPEC, dtype=torch.int32, device=dev)

tests/v1/spec_decode/test_dspark_topk.py

@@ -87,6 +87,8 @@ def test_dspark_revisited_position_uses_step_seed(use_fp64):
     spec.seeds = torch.arange(batch, dtype=torch.int64, device=device)
     spec._step_cols = torch.arange(2, dtype=torch.int32, device=device)
     spec.draft_logits = torch.empty(batch, 2, 2, device=device)
+    spec.draft_logits_num_shards = 1
+    spec.draft_logits_shard_idx = 0
     logits = torch.tensor([0.7, 0.3], device=device).log().expand(batch, -1)
     mapping = torch.arange(batch, device=device)
     positions = torch.full((batch,), 22, dtype=torch.int64, device=device)

tests/v1/spec_decode/test_lilicorr.py

@@ -364,6 +364,8 @@ def test_candidate_generation_routes_scores_and_adaptive_inputs(monkeypatch, lil
     spec.seeds = torch.arange(3)
     spec.draft_tokens = torch.empty(3, 2, dtype=torch.long)
     spec.draft_logits = None
+    spec.draft_logits_num_shards = 1
+    spec.draft_logits_shard_idx = 0
     spec.use_fp64_gumbel = False
     spec.enable_adaptive_verification = True
     spec.input_buffers = SimpleNamespace(input_ids=torch.arange(9))

tests/v1/spec_decode/test_rejection_sampler_utils.py

@@ -1105,6 +1105,7 @@ def test_verify_rejects_unproposed_drafts():
     rejection_sampler.fly_window_size = 0
     rejection_sampler.fly_entropy_threshold = 0.3
     rejection_sampler.fly_entropy_top_k = 3
+    rejection_sampler.draft_logits_num_shards = 1
     rejection_sampler.watermark_key = None
 
     def verify() -> tuple[torch.Tensor, torch.Tensor]:

tests/v1/worker/test_gpu_rejection_sampler_chunking.py

@@ -56,6 +56,7 @@ def _make_rejection_sampler(
     rejection_sampler.fly_window_size = 0
     rejection_sampler.fly_entropy_threshold = 0.3
     rejection_sampler.fly_entropy_top_k = 3
+    rejection_sampler.draft_logits_num_shards = 1
     rejection_sampler.watermark_key = None
     return rejection_sampler
 

tests/v1/worker/test_gpu_target_dependent_ar_speculator.py

@@ -232,6 +232,7 @@ def test_mm_support_configured_after_model_load(monkeypatch):
         speculator.draft_model_config = draft_model_config
         speculator.supports_mm_inputs = False
         speculator.use_acceptance_estimator = False
+        speculator.use_draft_logits = False
 
     monkeypatch.setattr(DraftModelSpeculator, "__init__", init_base)
     _mock_base_model_load(monkeypatch)
@@ -256,6 +257,7 @@ def test_load_model_keeps_mm_support_for_capable_drafter(monkeypatch):
     speculator.supports_mm_inputs = False
     speculator.inputs_embeds = None
     speculator.use_acceptance_estimator = False
+    speculator.use_draft_logits = False
     speculator.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(supports_multimodal_inputs=True)
     )
@@ -278,6 +280,7 @@ def test_load_model_disables_mm_support_for_text_only_drafter(monkeypatch):
     speculator.supports_mm_inputs = False
     speculator.inputs_embeds = None
     speculator.use_acceptance_estimator = False
+    speculator.use_draft_logits = False
     speculator.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(supports_multimodal_inputs=True)
     )
@@ -308,6 +311,7 @@ def test_multi_module_mm_support_configured_after_model_load(monkeypatch):
     speculator.supports_mm_inputs = False
     speculator.inputs_embeds = None
     speculator.use_acceptance_estimator = False
+    speculator.use_draft_logits = False
     speculator.cached_draft_input_embeds = None
     speculator.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(supports_multimodal_inputs=True)

tests/watermarking/test_watermarking.py

@@ -1861,6 +1861,8 @@ def test_dflash_draft_sampler_matches_step_loop_on_cuda(deduplicate_contexts, dt
 def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):
     speculator = object.__new__(DSparkSpeculator)
     speculator.draft_logits = torch.zeros(2, 1, 8)
+    speculator.draft_logits_num_shards = 1
+    speculator.draft_logits_shard_idx = 0
     speculator._d2t_scatter_index = torch.tensor([1, 5])
     speculator._draft_scatter_buf = torch.full((2, 8), float("-inf"))
     speculator.temperature = torch.ones(2)
@@ -1905,6 +1907,8 @@ def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):
 def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch):
     speculator = object.__new__(DSparkSpeculator)
     speculator.draft_logits = torch.zeros(2, 1, 8)
+    speculator.draft_logits_num_shards = 1
+    speculator.draft_logits_shard_idx = 0
     speculator._d2t_scatter_index = None
     speculator.temperature = torch.ones(2)
     speculator.seeds = torch.zeros(2, dtype=torch.int64)

vllm/v1/worker/gpu/model_runner.py

@@ -148,6 +148,7 @@ from vllm.v1.worker.gpu.pp_utils import PPHandler
 from vllm.v1.worker.gpu.sample.batch_shard import (
     BatchSharder,
     all_to_all_logits,
+    batch_sharded_sampling_enabled,
     gather_sampler_output,
 )
 from vllm.v1.worker.gpu.sample.logits_processor import build_custom_logits_processors
@@ -446,7 +447,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         )
 
         if self.parallel_config.enable_batch_sharded_sampling:
-            if hasattr(self.model, "compute_logits_local"):
+            if batch_sharded_sampling_enabled(self.vllm_config, self.model):
                 self.batch_sharder = BatchSharder(
                     max_num_reqs=self.max_num_reqs,
                     max_num_logits_per_req=self.decode_query_len,
@@ -512,7 +513,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 )
             custom = self.model_state.custom_sampler(self.sampler)
 
+            draft_logits_num_shards = (
+                self.speculator.draft_logits_num_shards
+                if isinstance(self.speculator, DraftModelSpeculator)
+                else 1
+            )
             if custom:
+                assert draft_logits_num_shards == 1, (
+                    "A custom rejection sampler cannot read batch-sharded draft logits"
+                )
                 self.vllm_config._check_supports_watermarking(custom_sampler=True)
                 self.sampler, self.rejection_sampler = custom
             elif self.speculative_config is not None:
@@ -523,6 +532,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     watermark_key=speculative_target_watermark_key(
                         self.vllm_config.watermark_config
                     ),
+                    draft_logits_num_shards=draft_logits_num_shards,
                 )
             self.prompt_logprobs_worker = PromptLogprobsWorker(
                 self.max_num_reqs,

vllm/v1/worker/gpu/sample/batch_shard.py

@@ -6,7 +6,9 @@ from dataclasses import dataclass, replace
 import numpy as np
 import torch
 import torch.distributed
+import torch.nn as nn
 
+from vllm.config import VllmConfig
 from vllm.distributed.communication_op import tensor_model_parallel_all_gather
 from vllm.distributed.parallel_state import get_tp_group
 from vllm.triton_utils import tl, triton
@@ -48,6 +50,31 @@ class BatchShardMetadata:
     adaptive_verification_enabled: bool
 
 
+def batch_sharded_sampling_enabled(
+    vllm_config: VllmConfig, target_model: nn.Module
+) -> bool:
+    """Whether the target's sampling and verification are batch-sharded."""
+    return bool(
+        vllm_config.parallel_config.enable_batch_sharded_sampling
+        and hasattr(target_model, "compute_logits_local")
+    )
+
+
+@triton.jit
+def shard_owner(req_state_idx, NUM_SHARDS: tl.constexpr):
+    """TP rank that samples a request and holds its draft logits.
+
+    Mirrored on the host in `BatchSharder.shard_sampler_inputs`.
+    """
+    return req_state_idx % NUM_SHARDS
+
+
+@triton.jit
+def shard_row(req_state_idx, NUM_SHARDS: tl.constexpr):
+    """Row of a request in its owner's batch-sharded draft logits cache."""
+    return req_state_idx // NUM_SHARDS
+
+
 @triton.jit(
     do_not_specialize=["num_reqs", "local_logits_start", "max_num_reqs_per_rank"]
 )
@@ -74,14 +101,16 @@ def _build_shard_plan_kernel(
 ):
     req_idx = tl.program_id(0)
     req_state_idx = tl.load(idx_mapping_ptr + req_idx)
-    owner = req_state_idx % TP_SIZE
+    owner = shard_owner(req_state_idx, TP_SIZE)
     num_logits = tl.load(cu_num_logits_ptr + req_idx + 1) - tl.load(
         cu_num_logits_ptr + req_idx
     )
 
     req_block = tl.arange(0, PADDED_NUM_REQS)
     req_mask = req_block < num_reqs
-    owners = tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0) % TP_SIZE
+    owners = shard_owner(
+        tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0), TP_SIZE
+    )
     num_logits_per_req = tl.load(
         cu_num_logits_ptr + req_block + 1, mask=req_mask, other=0
     ) - tl.load(cu_num_logits_ptr + req_block, mask=req_mask, other=0)

vllm/v1/worker/gpu/sample/gumbel.py

@@ -3,6 +3,7 @@
 import torch
 
 from vllm.triton_utils import tl, triton
+from vllm.v1.worker.gpu.sample.batch_shard import shard_owner, shard_row
 
 # Smallest positive value produced by Triton's fp32 `tl.rand`. Used by the
 # retained Philox helper for rejection sampling.
@@ -233,6 +234,8 @@ def gumbel_block_argmax(
     APPLY_TEMPERATURE: tl.constexpr,
     USE_FP64: tl.constexpr,
     PER_TOKEN_COL: tl.constexpr = False,
+    CACHE_NUM_SHARDS: tl.constexpr = 1,
+    CACHE_SHARD_IDX: tl.constexpr = 0,
 ):
     req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64)
     is_valid_req = req_state_idx >= 0
@@ -254,11 +257,14 @@ def gumbel_block_argmax(
         )
         tl.store(
             logits_cache_ptr
-            + req_state_idx * logits_cache_stride_0
+            + shard_row(req_state_idx, CACHE_NUM_SHARDS) * logits_cache_stride_0
             + col * logits_cache_stride_1
             + block,
             cached_logits,
-            mask=mask & is_valid_req,
+            mask=mask
+            & is_valid_req
+            # A batch-sharded cache holds only this rank's requests, packed by row.
+            & (shard_owner(req_state_idx, CACHE_NUM_SHARDS) == CACHE_SHARD_IDX),
         )
 
     seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0)
@@ -303,6 +309,8 @@ def _gumbel_sample_kernel(
     APPLY_TEMPERATURE: tl.constexpr,
     USE_FP64: tl.constexpr,
     PER_TOKEN_COL: tl.constexpr,
+    CACHE_NUM_SHARDS: tl.constexpr,
+    CACHE_SHARD_IDX: tl.constexpr,
 ):
     token_idx = tl.program_id(0).to(tl.int64)
     block_idx = tl.program_id(1)
@@ -335,6 +343,8 @@ def _gumbel_sample_kernel(
         APPLY_TEMPERATURE=APPLY_TEMPERATURE,
         USE_FP64=USE_FP64,
         PER_TOKEN_COL=PER_TOKEN_COL,
+        CACHE_NUM_SHARDS=CACHE_NUM_SHARDS,
+        CACHE_SHARD_IDX=CACHE_SHARD_IDX,
     )
     token_id = block_idx * BLOCK_SIZE + idx
     tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id)
@@ -353,6 +363,8 @@ def gumbel_sample(
     logits_cache_col: torch.Tensor | None = None,  # scalar or [num_tokens]
     use_fp64: bool = False,
     logits_cache_source: torch.Tensor | None = None,
+    logits_cache_num_shards: int = 1,
+    logits_cache_shard_idx: int = 0,
 ) -> torch.Tensor:
     # Enforce contiguity on non-strided input tensors
     expanded_idx_mapping = expanded_idx_mapping.contiguous()
@@ -410,6 +422,8 @@ def gumbel_sample(
         APPLY_TEMPERATURE=apply_temperature,
         USE_FP64=use_fp64,
         PER_TOKEN_COL=per_token_col,
+        CACHE_NUM_SHARDS=logits_cache_num_shards,
+        CACHE_SHARD_IDX=logits_cache_shard_idx,
     )
     # NOTE(woosuk): Use int64 for later indexing.
     max_block_idx = local_max.argmax(dim=-1, keepdim=True)

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

@@ -431,6 +431,10 @@ class DFlashSpeculator(DraftModelSpeculator):
         """Publish context features required by a draft's candidate head. Runs
         inside the FULL graph, so it must not sync with the host."""
 
+    def supports_batch_sharded_draft_logits(self) -> bool:
+        # Every draft logits write goes through sample_draft.
+        return True
+
     @torch.inference_mode()
     def propose(
         self,

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

@@ -8,6 +8,7 @@ import torch
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.triton_utils import tl, triton
+from vllm.v1.worker.gpu.sample.batch_shard import shard_owner, shard_row
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax, spec_step_seed
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
 
@@ -87,18 +88,25 @@ def _cache_draft_logits_kernel(
     num_steps: tl.constexpr,
     top_k: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    CACHE_NUM_SHARDS: tl.constexpr,
+    CACHE_SHARD_IDX: tl.constexpr,
 ):
     flat = tl.program_id(0)
     req_state = tl.load(req_state_ptr + flat)
     step = flat % num_steps
     offsets = tl.arange(0, BLOCK_K)
-    mask = (req_state >= 0) & (offsets < top_k)
+    mask = (
+        (req_state >= 0)
+        & (offsets < top_k)
+        # A batch-sharded cache holds only this rank's requests, packed by row.
+        & (shard_owner(req_state, CACHE_NUM_SHARDS) == CACHE_SHARD_IDX)
+    )
     candidate_base = flat * top_k
     cache_base = (req_state * num_steps + step) * top_k
     old_token_ids = tl.load(cached_candidate_ptr + cache_base + offsets, mask=mask)
     logits_base = (
         draft_logits_ptr
-        + req_state * draft_logits_stride_0
+        + shard_row(req_state, CACHE_NUM_SHARDS) * draft_logits_stride_0
         + step * draft_logits_stride_1
     )
     tl.store(logits_base + old_token_ids, -float("inf"), mask=mask)
@@ -135,6 +143,8 @@ class CandidateSampler:
         draft_tokens: torch.Tensor,
         draft_logits: torch.Tensor | None,
         use_fp64: bool,
+        draft_logits_num_shards: int = 1,
+        draft_logits_shard_idx: int = 0,
     ) -> None:
         candidate_ids = candidate_ids.contiguous()
         sample_pos = sample_pos.contiguous()
@@ -159,7 +169,12 @@ class CandidateSampler:
 
         if draft_logits is not None:
             self._cache_draft_logits(
-                candidate_ids, num_reqs * self.num_steps, idx_mapping, draft_logits
+                candidate_ids,
+                num_reqs * self.num_steps,
+                idx_mapping,
+                draft_logits,
+                draft_logits_num_shards,
+                draft_logits_shard_idx,
             )
 
     def _cache_draft_logits(
@@ -168,6 +183,8 @@ class CandidateSampler:
         num_sample: int,
         idx_mapping: torch.Tensor,
         draft_logits: torch.Tensor,
+        draft_logits_num_shards: int,
+        draft_logits_shard_idx: int,
     ) -> None:
         block_k = triton.next_power_of_2(self.top_k)
         _cache_draft_logits_kernel[(num_sample,)](
@@ -181,6 +198,8 @@ class CandidateSampler:
             num_steps=self.num_steps,
             top_k=self.top_k,
             BLOCK_K=block_k,
+            CACHE_NUM_SHARDS=draft_logits_num_shards,
+            CACHE_SHARD_IDX=draft_logits_shard_idx,
             num_warps=1,
         )
 
@@ -242,6 +261,8 @@ class DFlash2Speculator(DFlashSpeculator):
             self.draft_tokens,
             self.draft_logits,
             self.use_fp64_gumbel,
+            self.draft_logits_num_shards,
+            self.draft_logits_shard_idx,
         )
         if self.enable_adaptive_verification:
             self._maybe_predict_acceptance(

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

@@ -90,7 +90,7 @@ class DSparkSpeculator(DFlashSpeculator):
         # Reduced draft vocab: probabilistic rejection sampling indexes draft
         # logits by target id, so precompute the draft->target column map and a
         # scratch buffer to scatter logits into target vocab before sampling.
-        if self.draft_logits is not None and model.draft_id_to_target_id is not None:
+        if self.use_draft_logits and model.draft_id_to_target_id is not None:
             d2t = model.draft_id_to_target_id
             self._d2t_scatter_index = (
                 torch.arange(d2t.shape[0], device=d2t.device) + d2t
@@ -100,7 +100,7 @@ class DSparkSpeculator(DFlashSpeculator):
             self._draft_scatter_buf = torch.full(
                 (self.max_num_reqs, self.vocab_size),
                 float("-inf"),
-                dtype=self.draft_logits.dtype,
+                dtype=self.draft_logits_spec(self.vllm_config)[0],
                 device=self.device,
             )
         self.use_confidence_head = (
@@ -113,6 +113,10 @@ class DSparkSpeculator(DFlashSpeculator):
             self.use_acceptance_estimator = False
         return model
 
+    def supports_batch_sharded_draft_logits(self) -> bool:
+        # A model's own sample_sequential kernels write the full-size cache.
+        return not hasattr(self.model, "sample_sequential")
+
     def _sample_logits(
         self,
         logits: torch.Tensor,
@@ -135,12 +139,20 @@ class DSparkSpeculator(DFlashSpeculator):
 
         # sample_pos is the predicted token's position P. Sampling keys a draw
         # by the position before the sampled token, P-1.
-        sampler = (
-            gumbel_sample
-            if self.draft_watermarker is None
-            else self.draft_watermarker.sample
-        )
-        return sampler(
+        if self.draft_watermarker is not None:
+            return self.draft_watermarker.sample(
+                logits,
+                idx_map,
+                self.temperature,
+                self.seeds,
+                sample_pos - 1,
+                apply_temperature=True,
+                is_drafting=True,
+                logits_cache=self.draft_logits,
+                logits_cache_col=self._step_cols[step],
+                use_fp64=self.use_fp64_gumbel,
+            )
+        return gumbel_sample(
             logits,
             idx_map,
             self.temperature,
@@ -151,6 +163,8 @@ class DSparkSpeculator(DFlashSpeculator):
             logits_cache=self.draft_logits,
             logits_cache_col=self._step_cols[step],
             use_fp64=self.use_fp64_gumbel,
+            logits_cache_num_shards=self.draft_logits_num_shards,
+            logits_cache_shard_idx=self.draft_logits_shard_idx,
         )
 
     def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:

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

@@ -136,6 +136,10 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
             progress_bar_desc="Capturing multi-module MTP CUDA graphs",
         )
 
+    def supports_batch_sharded_draft_logits(self) -> bool:
+        # Every draft logits write goes through sample_draft.
+        return True
+
     @torch.inference_mode()
     def propose(
         self,

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

@@ -152,6 +152,7 @@ class RejectionSampler:
         device: torch.device,
         *,
         watermark_key: int | None = None,
+        draft_logits_num_shards: int = 1,
     ):
         self.sampler = sampler
         self.watermark_key = watermark_key
@@ -164,6 +165,9 @@ class RejectionSampler:
         self.fly_window_size = 0
         self.fly_entropy_threshold = spec_config.fly_entropy_threshold
         self.fly_entropy_top_k = spec_config.fly_entropy_top_k
+        # Above 1 when each TP rank caches only the draft logits of the
+        # requests it owns (see DraftModelSpeculator).
+        self.draft_logits_num_shards = draft_logits_num_shards
         self.synthetic_conditional_rates: torch.Tensor | None = None
         if rejection_sample_method == "synthetic":
             assert spec_config.synthetic_acceptance_rates is not None
@@ -289,6 +293,7 @@ class RejectionSampler:
             fly_window_size=self.fly_window_size,
             fly_entropy_threshold=self.fly_entropy_threshold,
             fly_entropy_top_k=self.fly_entropy_top_k,
+            draft_logits_num_shards=self.draft_logits_num_shards,
             **self._watermarking_kwargs(
                 draft_sampled, expanded_idx_mapping, expanded_local_pos
             ),
@@ -344,7 +349,7 @@ class RejectionSampler:
             hi = int(cu_num_logits_np[end])
             chunk_cu_num_logits_np = cu_num_logits_np[start : end + 1] - lo
             chunk_cu_num_logits = input_batch.cu_num_logits[start : end + 1] - lo
-            # draft_logits uses persistent request-state indices and stays global.
+            # draft_logits uses persistent request-state indices, not batch rows.
             processed_logits, sampled, num_sampled = self._verify(
                 logits[lo:hi],
                 draft_logits,

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

@@ -4,6 +4,7 @@ import torch
 
 from vllm.triton_utils import tl, tldevice, triton
 from vllm.v1.spec_decode.fly import compute_fly_entropy
+from vllm.v1.worker.gpu.sample.batch_shard import shard_row
 from vllm.v1.worker.gpu.sample.gumbel import (
     gumbel_block_argmax,
     spec_step_seed,
@@ -149,6 +150,7 @@ def _compute_global_logprobs_and_logsumexp(
     vocab_num_blocks,
     PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
+    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
 ):
     target_logit = tl.load(
         target_logits_ptr + logit_idx * target_logits_stride + token,
@@ -170,7 +172,8 @@ def _compute_global_logprobs_and_logsumexp(
         draft_logit = (
             tl.load(
                 draft_logits_ptr
-                + req_state_idx * draft_logits_stride_0
+                + shard_row(req_state_idx, DRAFT_LOGITS_NUM_SHARDS)
+                * draft_logits_stride_0
                 + draft_step * draft_logits_stride_1
                 + token,
                 mask=mask,
@@ -239,6 +242,7 @@ def _compute_local_logits_stats_kernel(
     num_logits,
     BLOCK_SIZE: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
+    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
 ):
     logit_idx = tl.program_id(0).to(tl.int64)
 
@@ -297,7 +301,8 @@ def _compute_local_logits_stats_kernel(
             draft_logits = (
                 tl.load(
                     draft_logits_ptr
-                    + req_state_idx * draft_logits_stride_0
+                    + shard_row(req_state_idx, DRAFT_LOGITS_NUM_SHARDS)
+                    * draft_logits_stride_0
                     + draft_step_idx * draft_logits_stride_1
                     + block_offsets,
                     mask=mask,
@@ -352,6 +357,7 @@ def _compute_cumulative_log_p_kernel(
     vocab_num_blocks,
     PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
+    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
 ):
     req_idx = tl.program_id(0)
     req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64)
@@ -393,6 +399,7 @@ def _compute_cumulative_log_p_kernel(
                     vocab_num_blocks,
                     PADDED_VOCAB_NUM_BLOCKS,
                     HAS_DRAFT_LOGITS,
+                    DRAFT_LOGITS_NUM_SHARDS,
                 )
             )
             log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0)
@@ -438,6 +445,7 @@ def _compute_local_residual_mass_kernel(
     vocab_num_blocks,
     BLOCK_SIZE: tl.constexpr,
     PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
+    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
 ):
     logit_idx = tl.program_id(0).to(tl.int64)
     # Adaptive verification can end before num_speculative_steps.
@@ -486,6 +494,7 @@ def _compute_local_residual_mass_kernel(
         vocab_num_blocks,
         PADDED_VOCAB_NUM_BLOCKS,
         True,  # HAS_DRAFT_LOGITS
+        DRAFT_LOGITS_NUM_SHARDS,
     )
 
     # Compute the residual mass: max(p_i * M_b(x|x_{<i}) - M_s(x|x_{<i}), 0)
@@ -557,6 +566,7 @@ def _rejection_kernel(
     vocab_num_blocks,
     PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
+    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
     SYNTHETIC_MODE: tl.constexpr,
     USE_BLOCK_VERIFICATION: tl.constexpr,
     FLY_WINDOW_SIZE: tl.constexpr,
@@ -683,6 +693,7 @@ def _rejection_kernel(
                         vocab_num_blocks,
                         PADDED_VOCAB_NUM_BLOCKS,
                         HAS_DRAFT_LOGITS,
+                        DRAFT_LOGITS_NUM_SHARDS,
                     )
                 )
                 if FLY_WINDOW_SIZE == 0 or pending_rejection < 0:
@@ -846,6 +857,7 @@ def _resample_kernel(
     vocab_size,
     BLOCK_SIZE: tl.constexpr,
     HAS_DRAFT_LOGITS: tl.constexpr,
+    DRAFT_LOGITS_NUM_SHARDS: tl.constexpr,
     USE_FP64: tl.constexpr,
     USE_BLOCK_VERIFICATION: tl.constexpr,
     CONTEXT_WIDTH: tl.constexpr,
@@ -911,7 +923,8 @@ def _resample_kernel(
         draft_logits = (
             tl.load(
                 draft_logits_ptr
-                + req_state_idx * draft_logits_stride_0
+                + shard_row(req_state_idx, DRAFT_LOGITS_NUM_SHARDS)
+                * draft_logits_stride_0
                 + resample_idx * draft_logits_stride_1
                 + block,
                 mask=mask,
@@ -1085,7 +1098,9 @@ def _insert_resampled_kernel(
 def rejection_sample(
     # [num_logits, V]
     target_logits: torch.Tensor,
-    # [max_num_reqs, num_speculative_steps, V]
+    # [max_num_reqs, num_speculative_steps, V], or under batch-sharded
+    # sampling, [cdiv(max_num_reqs, draft_logits_num_shards), ...] holding this
+    # rank's requests at row req_state_idx // draft_logits_num_shards.
     draft_logits: torch.Tensor | None,
     # [num_logits]
     draft_sampled: torch.Tensor,
@@ -1115,6 +1130,7 @@ def rejection_sample(
     fly_window_size: int = 0,
     fly_entropy_threshold: float = 0.3,
     fly_entropy_top_k: int = 3,
+    draft_logits_num_shards: int = 1,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert target_logits.ndim == 2 and target_logits.stride(-1) == 1
     assert draft_logits is None or (
@@ -1216,6 +1232,7 @@ def rejection_sample(
         num_logits,
         BLOCK_SIZE=VOCAB_BLOCK_SIZE,
         HAS_DRAFT_LOGITS=has_draft_logits,
+        DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
     )
 
     # Precompute the running joint ratio and residual mass for block
@@ -1251,6 +1268,7 @@ def rejection_sample(
             vocab_num_blocks,
             PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
             HAS_DRAFT_LOGITS=has_draft_logits,
+            DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
             num_warps=1,
         )
 
@@ -1288,6 +1306,7 @@ def rejection_sample(
                 vocab_num_blocks,
                 BLOCK_SIZE=VOCAB_BLOCK_SIZE,
                 PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
+                DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
             )
         else:
             local_residual_mass = None
@@ -1339,6 +1358,7 @@ def rejection_sample(
         vocab_num_blocks,
         PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
         HAS_DRAFT_LOGITS=has_draft_logits,
+        DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
         SYNTHETIC_MODE=synthetic_conditional_rates is not None,
         USE_BLOCK_VERIFICATION=use_block_verification,
         FLY_WINDOW_SIZE=fly_window_size,
@@ -1386,6 +1406,7 @@ def rejection_sample(
         vocab_size,
         BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
         HAS_DRAFT_LOGITS=has_draft_logits,
+        DRAFT_LOGITS_NUM_SHARDS=draft_logits_num_shards,
         USE_FP64=use_fp64,
         USE_BLOCK_VERIFICATION=use_block_verification,
         CONTEXT_WIDTH=context_width,

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

@@ -3,7 +3,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
 from dataclasses import replace
-from functools import partial
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -13,9 +12,11 @@ import torch.nn as nn
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.config.compilation import CUDAGraphMode
 from vllm.distributed.eplb.eplb_state import EplbState
+from vllm.distributed.parallel_state import get_tp_group
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.models import supports_multimodal_embeddings
+from vllm.utils.math_utils import cdiv
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.watermarking import create_watermarker
 from vllm.v1.watermarking.spec_decode import (
@@ -32,6 +33,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.dp_utils import DPSyncState
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.model_states.interface import ModelState
+from vllm.v1.worker.gpu.sample.batch_shard import batch_sharded_sampling_enabled
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (
     OnlineAcceptanceEstimator,
@@ -176,20 +178,14 @@ class DraftModelSpeculator(BaseSpeculator):
         self.use_acceptance_estimator = self.enable_adaptive_verification
         self.acceptance_estimator: OnlineAcceptanceEstimator | None = None
 
+        # Pre-temperature logits, cached from the previous decode step.
+        # Allocated in load_model. See _init_draft_logits.
         self.draft_logits: torch.Tensor | None = None
-        if self.speculative_config.draft_sample_method == "probabilistic":
-            # Pre-temperature logits, cached from the previous decode step.
-            dtype, fill = self.draft_logits_spec(vllm_config)
-            self.draft_logits = torch.full(
-                (
-                    self.max_num_reqs,
-                    self.num_speculative_steps,
-                    self.vocab_size,
-                ),
-                fill,
-                dtype=dtype,
-                device=device,
-            )
+        self.use_draft_logits = (
+            self.speculative_config.draft_sample_method == "probabilistic"
+        )
+        self.draft_logits_num_shards = 1
+        self.draft_logits_shard_idx = 0
 
         self.draft_watermarker: DraftWatermarker | None = None
         watermark_config = getattr(vllm_config, "watermark_config", None)
@@ -232,6 +228,8 @@ class DraftModelSpeculator(BaseSpeculator):
 
         self.model = self.load_draft_model(target_model, target_attn_layer_names)
         self._validate_local_argmax_reduction()
+        if self.use_draft_logits:
+            self._init_draft_logits(target_model)
 
         all_attn_layers = get_layers_from_vllm_config(
             self.vllm_config,
@@ -266,6 +264,39 @@ class DraftModelSpeculator(BaseSpeculator):
                 self.device,
             )
 
+    def supports_batch_sharded_draft_logits(self) -> bool:
+        """Whether every draft logits write on this speculator's path honors
+        draft_logits_num_shards and draft_logits_shard_idx."""
+        return False
+
+    def _init_draft_logits(self, target_model: nn.Module) -> None:
+        """Allocate the draft logits cache.
+
+        Under batch-sharded sampling, the rank that verifies a request is the
+        only one that reads its draft logits, so each TP rank keeps just the
+        rows of the requests it owns, at row
+        req_state_idx // draft_logits_num_shards.
+        """
+        if (
+            batch_sharded_sampling_enabled(self.vllm_config, target_model)
+            and self.draft_watermarker is None
+            and self.supports_batch_sharded_draft_logits()
+        ):
+            tp_group = get_tp_group()
+            self.draft_logits_num_shards = tp_group.world_size
+            self.draft_logits_shard_idx = tp_group.rank_in_group
+        dtype, fill = self.draft_logits_spec(self.vllm_config)
+        self.draft_logits = torch.full(
+            (
+                cdiv(self.max_num_reqs, self.draft_logits_num_shards),
+                self.num_speculative_steps,
+                self.vocab_size,
+            ),
+            fill,
+            dtype=dtype,
+            device=self.device,
+        )
+
     def set_eplb_state(self, eplb_state: EplbState) -> None:
         """Inject EPLB state after construction."""
         self.eplb_state = eplb_state
@@ -457,13 +488,8 @@ class DraftModelSpeculator(BaseSpeculator):
             return self.get_draft_top_tokens(hidden_states, spec_step_idx)
 
         logits = self.compute_draft_logits(hidden_states, spec_step_idx)
-        if draft_logits is not None:
-            sampler = (
-                gumbel_sample
-                if self.draft_watermarker is None
-                else partial(self.draft_watermarker.sample, num_steps=num_steps)
-            )
-            sampled = sampler(
+        if draft_logits is not None and self.draft_watermarker is not None:
+            sampled = self.draft_watermarker.sample(
                 logits,
                 idx_mapping,
                 temperature,
@@ -474,6 +500,22 @@ class DraftModelSpeculator(BaseSpeculator):
                 logits_cache=draft_logits,
                 logits_cache_col=draft_step,
                 use_fp64=self.use_fp64_gumbel,
+                num_steps=num_steps,
+            )
+        elif draft_logits is not None:
+            sampled = gumbel_sample(
+                logits,
+                idx_mapping,
+                temperature,
+                seeds,
+                sample_src_positions,
+                apply_temperature=True,
+                is_drafting=True,
+                logits_cache=draft_logits,
+                logits_cache_col=draft_step,
+                use_fp64=self.use_fp64_gumbel,
+                logits_cache_num_shards=self.draft_logits_num_shards,
+                logits_cache_shard_idx=self.draft_logits_shard_idx,
             )
         else:
             sampled = logits.argmax(dim=-1)

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

@@ -351,6 +351,10 @@ class StandaloneARSpeculator(DraftModelSpeculator):
                 self.draft_logits,
             )
 
+    def supports_batch_sharded_draft_logits(self) -> bool:
+        # Every draft logits write goes through sample_draft.
+        return True
+
     @torch.inference_mode()
     def propose(
         self,

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

@@ -161,6 +161,10 @@ class TargetDependentARSpeculator(DraftModelSpeculator):
         )
         self.on_multi_step_decode_end(self.max_num_reqs)
 
+    def supports_batch_sharded_draft_logits(self) -> bool:
+        # Every draft logits write goes through sample_draft.
+        return True
+
     @torch.inference_mode()
     def propose(
         self,