diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 924dfb73aba0..5279a0524987 100644 --- a/tests/distributed/test_pynccl.py +++ b/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( diff --git a/tests/v1/e2e/spec_decode/test_sharded_sampling.py b/tests/v1/e2e/spec_decode/test_sharded_sampling.py index ea323810bd95..ab0a5f59103e 100644 --- a/tests/v1/e2e/spec_decode/test_sharded_sampling.py +++ b/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" + ) diff --git a/tests/v1/worker/test_gpu_batch_shard.py b/tests/v1/worker/test_gpu_batch_shard.py index 72d3affed239..1454f410ee9e 100644 --- a/tests/v1/worker/test_gpu_batch_shard.py +++ b/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) diff --git a/vllm/distributed/device_communicators/pynccl.py b/vllm/distributed/device_communicators/pynccl.py index bb508c303927..8981594affda 100644 --- a/vllm/distributed/device_communicators/pynccl.py +++ b/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 diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index 31d1a094ebef..bc6e1a9cd7a8 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/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) diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index cfa06c63cb95..641d286efc95 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/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: diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index dda1026a53ef..341aa7e0811e 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/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 diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index c0fbfa553c30..de8538a46d0d 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index ebfa9da5235c..eaa48a21b145 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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] diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py index 9c2da8d018c7..fa6355e020f3 100644 --- a/vllm/v1/worker/gpu/sample/batch_shard.py +++ b/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) ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index a40aa9fab0b6..00cbaf26a33e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -21,6 +21,10 @@ CUDA graphs (FULL, mirroring DFlash) cover the whole draft step: the parallel backbone forward AND the sequential Markov sampling. + +Under batch-sharded sampling, draft steps also shard the sequential sampling +across TP ranks: each rank runs the Markov loop only for the requests it owns +(the same owners as verification), see ``_sample_sequential_sharded``. """ from typing import Any @@ -29,7 +33,15 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator +from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv +from vllm.v1.worker.gpu.sample.batch_shard import ( + all_to_all_vocab_shards, + batch_sharded_sampling_enabled, + get_tp_pynccl_comm, +) 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.utils import load_dspark_model @@ -81,6 +93,53 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.use_confidence_head: bool = False + # Set when the sequential sampling can be batch-sharded across TP. + self._sharded_sequential_comm: PyNcclCommunicator | None = None + + def load_model(self, target_model: torch.nn.Module) -> None: + super().load_model(target_model) + tp_size = get_tp_group().world_size + if tp_size > 1 and batch_sharded_sampling_enabled( + self.vllm_config, target_model + ): + self._try_enable_sharded_sequential_sampling() + + def _try_enable_sharded_sequential_sampling(self) -> None: + tp_size = get_tp_group().world_size + lm_head = getattr(self.model, "lm_head", None) + pynccl_comm = get_tp_pynccl_comm() + reason = None + if self.draft_logits is not None and self.draft_logits.row_mapping is None: + reason = "the draft logits cache is replicated" + elif self.use_acceptance_estimator: + reason = ( + "the online acceptance estimator fits on every request's " + "draft logits on every rank" + ) + elif self._draft_topk is not None: + reason = "top-k Markov sampling is not supported" + elif hasattr(self.model, "sample_sequential"): + reason = "the draft model fuses its own sequential sampling" + elif not hasattr(self.model, "compute_draft_logits_local") or ( + getattr(lm_head, "tp_size", None) != tp_size + ): + reason = "the draft model cannot compute vocab-sharded draft logits" + elif pynccl_comm is None: + reason = "PyNCCL is unavailable" + + if reason is not None: + logger.info( + "DSpark sequential sampling stays replicated across TP ranks: %s.", + reason, + ) + return + + self._sharded_sequential_comm = pynccl_comm + logger.info( + "DSpark sequential sampling is batch-sharded across %d TP ranks.", + tp_size, + ) + def load_draft_model( self, target_model: torch.nn.Module, @@ -189,19 +248,42 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: vocab_size = base_logits.shape[-1] base_logits = base_logits.view(num_reqs, n_spec, vocab_size) - idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec) - sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec) - confidence_markov_embeds = [] - # Anchor (bonus) token per request = the input id at query offset 0, # read via the precomputed persistent index (fixed buffer for capture). - prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]] + confidence = self._sample_markov_chain( + base_logits, + sample_hidden, + self.input_buffers.input_ids[self._anchor_idx[:num_reqs]], + self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec), + self.sample_pos[:num_sample].view(num_reqs, n_spec), + self.draft_tokens[:num_reqs], + ) + if confidence is not None: + self.draft_token_confidence_probs[:num_reqs] = confidence - for i in range(n_spec): + def _sample_markov_chain( + self, + base_logits: torch.Tensor, + hidden: torch.Tensor, + prev: torch.Tensor, + idx_map: torch.Tensor, + sample_pos: torch.Tensor, + out_tokens: torch.Tensor, + ) -> torch.Tensor | None: + """Run the sequential Markov steps for a block of requests, writing the + drafts to `out_tokens`. + + Returns: + [num_rows, num_speculative_steps] draft confidences, or None + without a confidence head. + + """ + markov_embeds = [] + for i in range(self.num_speculative_steps): # Sequential stage: Markov bias from the previously sampled token. markov_embed = self.model.markov_embed(prev) if self.use_confidence_head: - confidence_markov_embeds.append(markov_embed) + markov_embeds.append(markov_embed) bias = self.model.markov_bias(markov_embed) logits_i = base_logits[:, i] + bias if hasattr(self.model, "apply_output_logit_transform"): @@ -209,16 +291,88 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: draft_sampled_i = self._sample_logits( logits_i, idx_map[:, i], sample_pos[:, i], i ) - self.draft_tokens[:num_reqs, i] = draft_sampled_i + out_tokens[:, i] = draft_sampled_i prev = draft_sampled_i + confidence = None if self.use_confidence_head: confidence = self.model.compute_confidence( - sample_hidden, - torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1), + hidden, torch.stack(markov_embeds, dim=1).flatten(0, 1) + ).view(out_tokens.shape) + return confidence + + def _sample_sequential_sharded( + self, num_reqs: int, head_hidden: torch.Tensor + ) -> None: + """_sample_sequential with each TP rank sampling only its own requests. + + Request i is owned by rank shard_owner(i), as in DraftLogitsCache, so no + rank owns more than cdiv(num_reqs, tp_size) requests. Each rank computes + vocab-shard base logits for every request, in owner-sorted order. An + all-to-all then gives each rank the full-vocab rows of the requests it + owns, the Markov loop runs on just those rows, and an all-gather returns + the drafts in batch order. All shapes depend only on num_reqs, so the + step can be captured. + """ + assert self._sharded_sequential_comm is not None + tp_group = get_tp_group() + tp_size, tp_rank = tp_group.world_size, tp_group.rank_in_group + n_spec = self.num_speculative_steps + capacity = cdiv(num_reqs, tp_size) + # Rank r owns requests r, r + tp_size, r + 2 * tp_size, ... + send_reqs = torch.arange(capacity, device=self.device) * tp_size + torch.arange( + tp_size, device=self.device + ).unsqueeze(-1) + # Mark padding slots with -1. + send_reqs = torch.where(send_reqs < num_reqs, send_reqs, -1) + req_idxs = torch.arange(num_reqs, device=self.device) + # Inverse of send_reqs (ignoring padding), used for restoring the + # all-gathered drafts/confidences into the original request order. + gathered_src_idxs = (req_idxs % tp_size) * capacity + req_idxs // tp_size + + # Re-shard the draft logits from the vocab dimension to the batch dimension. + # Afterwards, each rank holds the full-vocab logits of its own requests. + # Padding slots read request 0's rows, and their results are discarded. + send_rows = send_reqs.clamp(min=0).unsqueeze(-1) * n_spec + self._step_cols + send_hidden = head_hidden[self.sample_indices[send_rows.flatten()]] + local_logits = self.model.compute_draft_logits_local(send_hidden) + num_local_rows = capacity * n_spec + base_logits = all_to_all_vocab_shards( + local_logits, [num_local_rows] * tp_size, num_local_rows + ).view(capacity, n_spec, -1)[..., : self.model.logits_processor.org_vocab_size] + + # Sample draft tokens for the locally assigned requests. + local_reqs = send_reqs[tp_rank] + local_rows = send_rows[tp_rank] + local_tokens = self.draft_tokens.new_empty(capacity, n_spec) + local_hidden = send_hidden.view(tp_size, capacity * n_spec, -1)[tp_rank] + confidence = self._sample_markov_chain( + base_logits, + local_hidden, + self.input_buffers.input_ids[self._anchor_idx[local_reqs.clamp(min=0)]], + # Empty slots are inert: no draft logits cache writes. + torch.where( + local_reqs.view(-1, 1) >= 0, self.sample_idx_mapping[local_rows], -1 + ), + self.sample_pos[local_rows], + local_tokens, + ) + + # Gather the draft tokens and confidences sampled by each rank. + packed = local_tokens + if confidence is not None: + packed = torch.cat( + [local_tokens, confidence.float().view(torch.int32).long()], dim=1 ) - self.draft_token_confidence_probs[:num_reqs] = confidence.view( - num_reqs, n_spec + gathered = packed.new_empty(tp_size * capacity, packed.shape[1]) + self._sharded_sequential_comm.all_gather(gathered, packed.contiguous()) + + # Write the output draft tokens and confidences. + packed = gathered[gathered_src_idxs] + self.draft_tokens[:num_reqs] = packed[:, :n_spec] + if confidence is not None: + self.draft_token_confidence_probs[:num_reqs] = ( + packed[:, n_spec:].int().view(torch.float32) ) def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> None: @@ -288,4 +442,7 @@ def _generate_draft( num_tokens_across_dp, cudagraph_runtime_mode, ) - self._sample_sequential(num_reqs, head_hidden) + if self._sharded_sequential_comm is not None: + self._sample_sequential_sharded(num_reqs, head_hidden) + else: + self._sample_sequential(num_reqs, head_hidden)