[Spec Decode][DSpark] Batch-sharded sequential block

tests/distributed/test_pynccl.py

@@ -246,6 +246,52 @@ def test_cuda_communicator_all_gather_dim_not_zero():
     distributed_run(cuda_communicator_all_gather_dim_worker_fn, 2)
 
 
+@worker_fn_wrapper
+def all_to_all_single_worker_fn():
+    pynccl_comm = PyNcclCommunicator(
+        get_world_group().cpu_group, device=get_world_group().device
+    )
+    rank = pynccl_comm.rank
+    world_size = pynccl_comm.world_size
+    device = f"cuda:{rank}"
+
+    # Rank src sends (src + dst) % 3 rows to rank dst, so some splits are
+    # empty. Row i of the src -> dst chunk carries src * 100 + dst * 10 + i.
+    def num_rows(src: int, dst: int) -> int:
+        return (src + dst) % 3
+
+    def chunk(src: int, dst: int) -> torch.Tensor:
+        values = src * 100 + dst * 10 + torch.arange(num_rows(src, dst))
+        return values.to(torch.float32).view(-1, 1).expand(-1, 8)
+
+    send = torch.cat([chunk(rank, dst) for dst in range(world_size)]).to(device)
+    expected = torch.cat([chunk(src, rank) for src in range(world_size)]).to(device)
+    recv = torch.empty_like(expected)
+    output_split_sizes = [num_rows(src, rank) for src in range(world_size)]
+    input_split_sizes = [num_rows(rank, dst) for dst in range(world_size)]
+
+    pynccl_comm.all_to_all_single(recv, send, output_split_sizes, input_split_sizes)
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(recv, expected, rtol=0, atol=0)
+
+    # Split sizes fixed at capture, so the collective replays from a CUDA
+    # graph with fresh inputs.
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        pynccl_comm.all_to_all_single(recv, send, output_split_sizes, input_split_sizes)
+    send.add_(1000)
+    graph.replay()
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(recv, expected + 1000, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
+)
+def test_pynccl_all_to_all_single():
+    distributed_run(all_to_all_single_worker_fn, 2)
+
+
 @worker_fn_wrapper
 def all_gatherv_worker_fn():
     pynccl_comm = PyNcclCommunicator(

tests/v1/e2e/spec_decode/test_sharded_sampling.py

@@ -101,3 +101,73 @@ def run(disable_sharding: bool):
                 f"{name}: {len(divergent)}/{len(ref_outputs)} prompts diverged "
                 f"at {divergent}, beyond near-tie noise"
             )
+
+
+def _dspark_sharding_state(worker) -> tuple[bool, bool]:
+    speculator = worker.model_runner.speculator
+    draft_logits = speculator.draft_logits
+    sharded = draft_logits is not None and draft_logits.row_mapping is not None
+    return sharded, speculator._sharded_sequential_comm is not None
+
+
+def test_dspark_sharded_sampling_equivalence(monkeypatch: pytest.MonkeyPatch):
+    """Batch-sharded sampling, including the sharded draft logits cache and
+    the batch-sharded DSpark sequential sampling loop, must match replicated
+    sampling up to near-tie noise."""
+    tp_size = 2
+    _skip_if_insufficient_gpus_for_tp(tp_size)
+    test_prompts = get_test_prompts(mm_enabled=False, num_prompts=32)
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+        m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+        m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+        def run(disable_sharding: bool):
+            llm = LLM(
+                model="Qwen/Qwen3-4B-FP8",
+                tensor_parallel_size=tp_size,
+                enable_batch_sharded_sampling=not disable_sharding,
+                speculative_config={
+                    "method": "dspark",
+                    "model": "deepseek-ai/dspark_qwen3_4b_block7",
+                    "num_speculative_tokens": 7,
+                    "draft_sample_method": "probabilistic",
+                },
+                max_model_len=2048,
+                max_num_seqs=len(test_prompts),
+                enable_flashinfer_autotune=False,
+            )
+            try:
+                states = llm.llm_engine.collective_rpc(_dspark_sharding_state)
+                for state in states:
+                    # (draft logits cache sharded, sequential loop sharded)
+                    assert state == (not disable_sharding, not disable_sharding)
+                greedy = llm.chat(test_prompts, greedy_sampling())
+                seeded = llm.chat(
+                    test_prompts,
+                    SamplingParams(temperature=1.0, seed=33, max_tokens=32),
+                )
+                return greedy, seeded
+            finally:
+                del llm
+                torch.accelerator.empty_cache()
+                cleanup_dist_env_and_memory()
+                wait_for_memory_to_settle()
+
+        ref_greedy, ref_seeded = run(disable_sharding=True)
+        shard_greedy, shard_seeded = run(disable_sharding=False)
+
+        for name, ref_outputs, shard_outputs in (
+            ("greedy", ref_greedy, shard_greedy),
+            ("seeded", ref_seeded, shard_seeded),
+        ):
+            divergent = [
+                i
+                for i, (ref, out) in enumerate(zip(ref_outputs, shard_outputs))
+                if list(ref.outputs[0].token_ids) != list(out.outputs[0].token_ids)
+            ]
+            assert len(divergent) <= MAX_DIVERGENT_PROMPTS, (
+                f"{name}: {len(divergent)}/{len(ref_outputs)} prompts diverged "
+                f"at {divergent}, beyond near-tie noise"
+            )

tests/v1/worker/test_gpu_batch_shard.py

@@ -323,7 +323,7 @@ def test_all_to_all_reassembly_simulated(tp_size: int):
             [sends[src][bounds[rank] : bounds[rank + 1]] for src in range(tp_size)]
         )
         # Reassemble source-rank-major chunks into full-vocab rows, mirroring
-        # all_to_all_logits (source-rank order is vocab-shard order).
+        # all_to_all_vocab_shards (source-rank order is vocab-shard order).
         reassembled = (
             recv.view(tp_size, metadata.num_local_logits, shard_w)
             .permute(1, 0, 2)

vllm/distributed/device_communicators/pynccl.py

@@ -486,6 +486,60 @@ def scatter(
             self.recv(output_tensor, root, stream)
         self.nccl.ncclGroupEnd()
 
+    def all_to_all_single(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        output_split_sizes: list[int],
+        input_split_sizes: list[int],
+        stream=None,
+    ):
+        """All-to-all along dim 0, like `torch.distributed.all_to_all_single`:
+        the next `input_split_sizes[i]` rows of `input_tensor` go to rank i, and
+        the next `output_split_sizes[i]` rows of `output_tensor` come from it.
+        CUDA graph capturable for split sizes fixed at capture.
+        """
+        if self.disabled:
+            return
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        assert input_tensor.is_contiguous() and output_tensor.is_contiguous()
+        if stream is None:
+            stream = current_stream()
+        nccl_stream = cudaStream_t(stream.cuda_stream)
+        nccl_dtype = ncclDataTypeEnum.from_torch(input_tensor.dtype)
+        row_numel = input_tensor[0].numel() if input_tensor.dim() > 1 else 1
+        row_bytes = row_numel * input_tensor.element_size()
+        send_ptr = input_tensor.data_ptr()
+        recv_ptr = output_tensor.data_ptr()
+        self.nccl.ncclGroupStart()
+        for peer in range(self.world_size):
+            send_rows = input_split_sizes[peer]
+            recv_rows = output_split_sizes[peer]
+            if send_rows > 0:
+                self.nccl.ncclSend(
+                    buffer_type(send_ptr),
+                    send_rows * row_numel,
+                    nccl_dtype,
+                    peer,
+                    self.comm,
+                    nccl_stream,
+                )
+            if recv_rows > 0:
+                self.nccl.ncclRecv(
+                    buffer_type(recv_ptr),
+                    recv_rows * row_numel,
+                    nccl_dtype,
+                    peer,
+                    self.comm,
+                    nccl_stream,
+                )
+            send_ptr += send_rows * row_bytes
+            recv_ptr += recv_rows * row_bytes
+        self.nccl.ncclGroupEnd()
+
     def send(self, tensor: torch.Tensor, dst: int, stream=None):
         if self.disabled:
             return

vllm/model_executor/models/qwen3.py

@@ -343,6 +343,12 @@ def compute_logits(
         logits = self.logits_processor(self.lm_head, hidden_states)
         return logits
 
+    def compute_logits_local(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)

vllm/model_executor/models/qwen3_dspark.py

@@ -277,6 +277,10 @@ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
         # Markov bias in draft space, then remaps via map_draft_to_target.
         return self.logits_processor(self.lm_head, hidden_states)
 
+    def compute_draft_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """This TP rank's vocab shard of compute_draft_logits, not gathered."""
+        return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
+
     def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
         # Map draft-vocab ids to target ids (identity for full-vocab drafts).
         if self.draft_id_to_target_id is None:

vllm/models/deepseek_v4/nvidia/dspark.py

@@ -423,6 +423,12 @@ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
         # Full-vocab draft: base logits, no d2t scatter.
         return self.compute_logits(hidden_states)
 
+    def compute_draft_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """This TP rank's vocab shard of compute_draft_logits, not gathered."""
+        return self.logits_processor(
+            self.lm_head, self.model.norm(hidden_states), skip_gather=True
+        )
+
     def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
         return draft_ids  # full-vocab: draft ids are target ids
 

vllm/models/deepseek_v41/nvidia/dspark.py

@@ -358,6 +358,12 @@ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
         # Full-vocab draft: base logits, no d2t scatter.
         return self.compute_logits(hidden_states)
 
+    def compute_draft_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """This TP rank's vocab shard of compute_draft_logits, not gathered."""
+        return self.logits_processor(
+            self.lm_head, self.model.norm(hidden_states), skip_gather=True
+        )
+
     def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
         return draft_ids  # full-vocab: draft ids are target ids
 

vllm/v1/worker/gpu/model_runner.py

@@ -148,7 +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,
+    all_to_all_vocab_shards,
     batch_sharded_sampling_enabled,
     gather_sampler_output,
 )
@@ -1677,7 +1677,11 @@ def sample(
             # produces full-vocab logits for just the locally-owned requests.
             sample_hidden_states = hidden_states[sorted_logits_indices]
             local_logits = self.model.compute_logits_local(sample_hidden_states)
-            logits = all_to_all_logits(local_logits, shard_metadata)
+            logits = all_to_all_vocab_shards(
+                local_logits,
+                shard_metadata.num_logits_per_rank,
+                shard_metadata.num_local_logits,
+            )
             logits = logits[:, : self.vocab_size]
         else:
             sample_hidden_states = hidden_states[input_batch.logits_indices]

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

@@ -19,6 +19,7 @@
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
 
 if TYPE_CHECKING:
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
     from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
 
 
@@ -64,6 +65,14 @@ def batch_sharded_sampling_enabled(
     )
 
 
+def get_tp_pynccl_comm() -> "PyNcclCommunicator | None":
+    """The TP group's PyNCCL communicator, or None if it is unavailable."""
+    pynccl_comm = getattr(get_tp_group().device_communicator, "pynccl_comm", None)
+    if pynccl_comm is None or pynccl_comm.disabled:
+        return None
+    return pynccl_comm
+
+
 @triton.jit
 def shard_owner(req_idx, NUM_SHARDS: tl.constexpr):
     """TP rank that samples the request at a batch position.
@@ -374,25 +383,39 @@ def shard_sampler_inputs(
         return local_batch, sorted_logits_indices, local_grammar_output, metadata
 
 
-def all_to_all_logits(
-    logits_shard: torch.Tensor,
-    metadata: BatchShardMetadata,
+def all_to_all_vocab_shards(
+    logits_shard: torch.Tensor, send_splits: list[int], num_local_rows: int
 ) -> torch.Tensor:
-    recv = logits_shard.new_empty(
-        metadata.tp_size * metadata.num_local_logits, logits_shard.shape[-1]
-    )
-    torch.distributed.all_to_all_single(
-        recv,
-        logits_shard.contiguous(),
-        output_split_sizes=[metadata.num_local_logits] * metadata.tp_size,
-        input_split_sizes=metadata.num_logits_per_rank,
-        group=get_tp_group().device_group,
-    )
+    """Exchange vocab-shard logits so each rank gets full-vocab logits for its
+    own rows. `send_splits[r]` rows of `logits_shard` go to rank r.
+
+    Returns:
+        [num_local_rows, tp_size * shard_width] logits of this rank's rows.
+
+    """
+    tp_size = len(send_splits)
+    recv = logits_shard.new_empty(tp_size * num_local_rows, logits_shard.shape[-1])
+    recv_splits = [num_local_rows] * tp_size
+    pynccl_comm = get_tp_pynccl_comm()
+    if pynccl_comm is not None:
+        # The TP collectives' communicator: a c10d all-to-all would make NCCL
+        # allocate a second set of point-to-point buffers.
+        pynccl_comm.all_to_all_single(
+            recv, logits_shard.contiguous(), recv_splits, send_splits
+        )
+    else:
+        torch.distributed.all_to_all_single(
+            recv,
+            logits_shard.contiguous(),
+            output_split_sizes=recv_splits,
+            input_split_sizes=send_splits,
+            group=get_tp_group().device_group,
+        )
     shard_width = recv.shape[-1]
     return (
-        recv.view(metadata.tp_size, metadata.num_local_logits, shard_width)
+        recv.view(tp_size, num_local_rows, shard_width)
         .permute(1, 0, 2)
-        .reshape(metadata.num_local_logits, metadata.tp_size * shard_width)
+        .reshape(num_local_rows, tp_size * shard_width)
     )