[Spec Decode][Sampling] Batch-sharded draft logits

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 init_base(self, config, device):
     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

@@ -229,6 +229,7 @@ def init_base(speculator, vllm_config, device):
         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)
@@ -253,6 +254,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)
     )
@@ -275,6 +277,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)
     )
@@ -305,6 +308,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 run(fused: bool):
 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 sample(
 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

@@ -147,6 +147,7 @@
 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
@@ -445,7 +446,7 @@ def load_model(
         )
 
         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,
@@ -511,7 +512,15 @@ def load_model(
                 )
             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:
@@ -522,6 +531,7 @@ def load_model(
                     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 @@
 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 @@ def prepare_context_anchor(
         """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 @@
 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 @@ def sample(
         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 @@ def sample(
 
         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 @@ def _cache_draft_logits(
         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 @@ def _cache_draft_logits(
             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 @@ def _generate_draft(
             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(