From 5acc1c4e4b8730298cbff7a4a7c68c814dc24fd7 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Wed, 26 Aug 2026 20:59:35 -0700 Subject: [PATCH] [Model Runner V2][Spec Decode] Skip DP sync before EAGLE/MTP draft prefill (#53694) --- .../worker/test_gpu_model_runner_v2_eplb.py | 2 +- vllm/v1/worker/gpu/dp_utils.py | 124 +++++++++++++++--- vllm/v1/worker/gpu/model_runner.py | 14 +- .../spec_decode/autoregressive/speculator.py | 19 ++- .../gpu/spec_decode/dflash/speculator.py | 13 +- .../gpu/spec_decode/extract_hidden_states.py | 7 +- .../multi_module_mtp/speculator.py | 10 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 3 +- 8 files changed, 156 insertions(+), 36 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 6a4603db38..356b313515 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -179,10 +179,10 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): slot_mappings_by_layer=None, hidden_states=None, aux_hidden_states=None, + dp_sync=None, finished_req_ids=set(), ec_connector_output=None, routed_experts=None, - num_tokens_across_dp=None, ) runner.req_states = SimpleNamespace() diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index fcdb4e28cf..d1a935684f 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations +from dataclasses import dataclass, replace + import torch import torch.distributed as dist @@ -13,6 +15,26 @@ from vllm.v1.worker.gpu.cudagraph_utils import ( ) +@dataclass(frozen=True) +class DPSyncState: + """What a `dispatch_cg_and_sync_dp` call agreed across DP ranks. + + Every field is identical on every rank, so callers can branch on them + without disagreeing about whether to run a collective. Never add a per-rank + value here. + """ + + # Per-rank token counts, for the forward context. All entries hold the + # agreed padded count, unless `eager`, where nothing was padded. + num_tokens_across_dp: torch.Tensor + # Agreed uniform decode length. None means the ranks disagreed, which is an + # answer, not "unknown". + uniform_token_count: int | None + # Whether the ranks agreed to run eager. A dispatch reusing this must run + # eager too; no shape was agreed, so picking a graph per rank would diverge. + eager: bool + + def sync_cudagraph_and_dp_padding( cudagraph_manager: CudaGraphManager | None, desired_batch_desc: BatchExecutionDescriptor, @@ -23,11 +45,11 @@ def sync_cudagraph_and_dp_padding( dp_rank: int, max_query_len: int | None = None, num_active_loras: int = 0, -) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]: +) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: """ Coordinates the batch descriptor and DP padding across all ranks. - Returns (synced_batch_desc, num_tokens_across_dp). + Returns (synced_batch_desc, sync). `sync` is None when no rank has work. """ assert dp_size > 1, "DP size must be greater than 1" group = get_dp_group().cpu_group @@ -43,6 +65,13 @@ def sync_cudagraph_and_dp_padding( uniform_token_counts_across_dp = tensor[2] max_query_lens_across_dp = tensor[3] + # If ranks disagree on the uniform token count, or its 0 (means None) set to None + synced_uniform_token_count: int | None = int(uniform_token_counts_across_dp[0]) + if synced_uniform_token_count == 0 or not torch.all( + uniform_token_counts_across_dp == synced_uniform_token_count + ): + synced_uniform_token_count = None + if torch.all(num_tokens_across_dp == 0).item(): synced_desc = BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=0, num_reqs=0 @@ -53,24 +82,25 @@ def sync_cudagraph_and_dp_padding( # If any rank wants to run eager, all ranks run eager if synced_cg_mode == CUDAGraphMode.NONE: - return BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=num_tokens, - num_reqs=num_reqs, - num_active_loras=desired_batch_desc.num_active_loras, - ), num_tokens_across_dp + return ( + BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=num_reqs, + num_active_loras=desired_batch_desc.num_active_loras, + ), + DPSyncState( + num_tokens_across_dp=num_tokens_across_dp, + uniform_token_count=synced_uniform_token_count, + eager=True, + ), + ) assert cudagraph_manager is not None, ( "cudagraph_manager should only be None during profile run, " "where synced_cg_mode must be NONE across all DP ranks" ) synced_num_tokens = int(num_tokens_across_dp.max().item()) - synced_uniform_token_count = uniform_token_counts_across_dp[0] - # If ranks disagree on the uniform token count, or its 0 (means None) set to None - if synced_uniform_token_count == 0 or not torch.all( - uniform_token_counts_across_dp == synced_uniform_token_count - ): - synced_uniform_token_count = None # Varlen decode graphs are selected by the query-length bound, so ranks must agree # on it or they pad to different token counts below. @@ -92,7 +122,11 @@ def sync_cudagraph_and_dp_padding( # Update num_tokens_across_dp to reflect padded size. num_tokens_across_dp[:] = synced_desc.num_tokens - return synced_desc, num_tokens_across_dp + return synced_desc, DPSyncState( + num_tokens_across_dp=num_tokens_across_dp, + uniform_token_count=synced_uniform_token_count, + eager=False, + ) def dispatch_cg_and_sync_dp( @@ -105,8 +139,43 @@ def dispatch_cg_and_sync_dp( max_query_len: int | None = None, need_eager: bool = False, num_active_loras: int = 0, -) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]: - if need_eager: + dp_sync: DPSyncState | None = None, +) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: + """Pick a cudagraph descriptor for this batch, agreeing it across DP ranks. + + Runs a collective when dp_size > 1 so every rank dispatches to the same + shape. Pass `dp_sync` from a dispatch already made over this same batch (a + drafter's prefill runs the target's batch shape) to reuse that agreement + instead, with no collective. + + Args: + cudagraph_manager: Manager to dispatch against. May be None only when + `need_eager` is True (profile run). + num_reqs: Requests in this rank's batch. + num_tokens: Tokens in this rank's batch, already padded by the caller. + uniform_token_count: Per-request token count if this rank's batch is a + uniform decode, else None. `dp_sync.uniform_token_count` takes its + place when a sync is reused, since that one is agreed across ranks. + dp_size: Data-parallel world size. 1 skips all cross-rank work. + dp_rank: This rank's index in the DP group. + max_query_len: Upper bound on per-request query length, for selecting + varlen decode graphs. None means the graph must not constrain it. + need_eager: Force `CUDAGraphMode.NONE` instead of dispatching. + num_active_loras: Active LoRA count for this rank. Does not need + cross-rank agreement; it never changes a bucket's token count. + dp_sync: Agreement from a prior dispatch over this same batch, to reuse. + Must come from a batch with this same padded `num_tokens` and the + same `uniform_token_count`; `num_reqs` may differ, as neither + depends on it. Passing a sync from a different batch is a caller + error and trips an assert. + + Returns: + (batch_desc, sync), where `sync` is this batch's agreement for a later + dispatch to reuse. It is None when `dp_size` is 1 or no rank has work. + """ + reuse_eager = dp_sync is not None and dp_sync.eager + + if need_eager or reuse_eager: batch_desc = BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens, @@ -121,7 +190,7 @@ def dispatch_cg_and_sync_dp( batch_desc = cudagraph_manager.dispatch( num_reqs, num_tokens, - uniform_token_count, + dp_sync.uniform_token_count if dp_sync is not None else uniform_token_count, num_active_loras=num_active_loras, max_query_len=max_query_len, ) @@ -129,6 +198,25 @@ def dispatch_cg_and_sync_dp( if dp_size == 1: return batch_desc, None + if dp_sync is not None: + assert dp_sync.num_tokens_across_dp[dp_rank] == num_tokens, ( + "reusing a DP sync taken over a different batch" + ) + assert ( + dp_sync.uniform_token_count is None + or uniform_token_count == dp_sync.uniform_token_count + ), "reusing a DP sync taken over a different batch" + if not dp_sync.eager and batch_desc.num_tokens != num_tokens: + # Capture sizes can differ between managers, so this one may + # pad further. Every rank pads alike, so report what will run. + dp_sync = replace( + dp_sync, + num_tokens_across_dp=torch.full_like( + dp_sync.num_tokens_across_dp, batch_desc.num_tokens + ), + ) + return batch_desc, dp_sync + return sync_cudagraph_and_dp_padding( cudagraph_manager, batch_desc, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index a58c86166d..cde043d1cb 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -97,7 +97,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import ( from vllm.v1.worker.gpu.cudagraph_utils import ( profile_cudagraph_memory as _profile_cudagraph_memory, ) -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp +from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.ec_connector import get_ec_connector from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after from vllm.v1.worker.gpu.input_batch import ( @@ -752,6 +752,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states + dp_sync = self.execute_model_state.dp_sync self.execute_model_state = None self.step_timing.forward_end() @@ -796,6 +797,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): next_prefill_tokens=self.req_states.next_prefill_tokens, temperature=self.sampler.sampling_states.temperature.gpu, seeds=self.sampler.sampling_states.seeds.gpu, + dp_sync=dp_sync, dummy_run=True, skip_attn_for_dummy_run=skip_attn, mm_inputs=mm_inputs, @@ -1544,7 +1546,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): # cross-attention cache with dynamic encoder outputs. skip_compiled = True - batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + batch_desc, dp_sync = dispatch_cg_and_sync_dp( self.cudagraph_manager, num_reqs, num_toks, @@ -1731,7 +1733,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.vllm_config, num_tokens=input_batch.num_tokens_after_padding, cudagraph_runtime_mode=batch_desc.cg_mode, - num_tokens_across_dp=num_tokens_across_dp, + num_tokens_across_dp=( + dp_sync.num_tokens_across_dp if dp_sync is not None else None + ), batch_descriptor=batch_descriptor, slot_mapping=slot_mappings_by_layer, skip_compiled=skip_compiled, @@ -1777,6 +1781,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): slot_mappings_by_layer=slot_mappings_by_layer, hidden_states=hidden_states, aux_hidden_states=aux_hidden_states, + dp_sync=dp_sync, finished_req_ids=finished_req_ids, ec_connector_output=ec_connector_output, routed_experts=routed_experts, @@ -1801,6 +1806,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states + dp_sync = self.execute_model_state.dp_sync finished_req_ids = self.execute_model_state.finished_req_ids ec_connector_output = self.execute_model_state.ec_connector_output routed_experts = self.execute_model_state.routed_experts @@ -1922,6 +1928,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.req_states.next_prefill_tokens, self.sampler.sampling_states.temperature.gpu, self.sampler.sampling_states.seeds.gpu, + dp_sync=dp_sync, mm_inputs=mm_inputs, ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens @@ -2066,6 +2073,7 @@ class ExecuteModelState(NamedTuple): slot_mappings_by_layer: dict[str, torch.Tensor] | None hidden_states: torch.Tensor | None aux_hidden_states: list[torch.Tensor] | None + dp_sync: DPSyncState | None finished_req_ids: set[str] ec_connector_output: ECConnectorOutput | None routed_experts: RoutedExpertsTensors | None diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 57430529f6..dc5f115c70 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -14,7 +14,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp +from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp 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.spec_decode.autoregressive.cudagraph_utils import ( @@ -219,7 +219,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - num_tokens_across_dp: torch.Tensor | None = None, + dp_sync: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -279,7 +279,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): max_query_len, input_batch.has_prefill, ) - prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + prefill_batch_desc, prefill_batch_sync = dispatch_cg_and_sync_dp( self.prefill_cudagraph_manager, num_reqs, num_tokens_padded, @@ -287,6 +287,12 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, + dp_sync=dp_sync, + ) + num_tokens_across_dp = ( + prefill_batch_sync.num_tokens_across_dp + if prefill_batch_sync is not None + else None ) self._prepare_eplb_forward(num_tokens) @@ -328,7 +334,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): # Each request produces exactly 1 token per draft generation step, # enabling FULL graph replay. - decode_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + decode_batch_desc, decode_batch_sync = dispatch_cg_and_sync_dp( self.decode_cudagraph_manager, num_reqs, num_reqs, @@ -337,6 +343,11 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): dp_rank=self.dp_rank, need_eager=is_profile, ) + num_tokens_across_dp = ( + decode_batch_sync.num_tokens_across_dp + if decode_batch_sync is not None + else None + ) self.on_multi_step_decode_begin(num_reqs) # Generate the remaining num_speculative_steps - 1 draft tokens. diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 684a3d3a44..2ff7b22d06 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -19,7 +19,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp +from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp 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.spec_decode.dflash.cudagraph import DFlashCudaGraphManager @@ -335,7 +335,7 @@ class DFlashSpeculator(DraftModelSpeculator): temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - num_tokens_across_dp: torch.Tensor | None = None, + dp_sync: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -377,7 +377,9 @@ class DFlashSpeculator(DraftModelSpeculator): num_query_tokens, attn_metadata=None, slot_mappings=None, - num_tokens_across_dp=num_tokens_across_dp, + num_tokens_across_dp=( + dp_sync.num_tokens_across_dp if dp_sync is not None else None + ), cudagraph_runtime_mode=CUDAGraphMode.NONE, ) return self.draft_tokens[:num_reqs] @@ -438,7 +440,7 @@ class DFlashSpeculator(DraftModelSpeculator): ) # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs - batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + batch_desc, batch_sync = dispatch_cg_and_sync_dp( self.query_cudagraph_manager, num_reqs, num_query_tokens, @@ -450,6 +452,9 @@ class DFlashSpeculator(DraftModelSpeculator): num_reqs_padded = batch_desc.num_reqs or num_reqs num_tokens_padded = batch_desc.num_tokens + num_tokens_across_dp = ( + batch_sync.num_tokens_across_dp if batch_sync is not None else None + ) # Rebuild the draft attention metadata even when replaying the FULL # graph so that any attention metadata builder state is updated. diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py index efcdec8018..c411eeae55 100644 --- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py +++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py @@ -10,6 +10,7 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import set_forward_context from vllm.model_executor.model_loader import get_model +from vllm.v1.worker.gpu.dp_utils import DPSyncState from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator @@ -97,7 +98,7 @@ class ExtractHiddenStatesSpeculator(DraftModelSpeculator): next_prefill_tokens: torch.Tensor, temperature: torch.Tensor, seeds: torch.Tensor, - num_tokens_across_dp: torch.Tensor | None = None, + dp_sync: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -143,7 +144,9 @@ class ExtractHiddenStatesSpeculator(DraftModelSpeculator): draft_attn_metadata, self.vllm_config, num_tokens=num_tokens, - num_tokens_across_dp=num_tokens_across_dp, + num_tokens_across_dp=( + dp_sync.num_tokens_across_dp if dp_sync is not None else None + ), cudagraph_runtime_mode=CUDAGraphMode.NONE, slot_mapping=draft_slot_mappings, is_padding=input_batch.is_padding[:num_tokens], diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index fc147b3f49..82835c7b35 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -12,7 +12,7 @@ from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp +from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import ( SpeculatorCudaGraphManager, @@ -150,7 +150,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator): temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - num_tokens_across_dp: torch.Tensor | None = None, + dp_sync: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, @@ -191,7 +191,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator): uniform_token_count = get_uniform_decode_token_count( num_reqs, num_tokens, max_query_len, input_batch.has_prefill ) - batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + batch_desc, batch_sync = dispatch_cg_and_sync_dp( self.cudagraph_manager, num_reqs, input_batch.num_tokens_after_padding, @@ -199,6 +199,10 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator): dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, + dp_sync=dp_sync, + ) + num_tokens_across_dp = ( + batch_sync.num_tokens_across_dp if batch_sync is not None else None ) # Rebuild the slot mappings and attention metadata. diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 8427bf3d00..303f16b79f 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -21,6 +21,7 @@ from vllm.v1.worker.gpu.attn_utils import ( init_attn_backend, ) from vllm.v1.worker.gpu.block_table import BlockTables +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.gumbel import gumbel_sample @@ -60,7 +61,7 @@ class BaseSpeculator(ABC): temperature: torch.Tensor, # [max_num_reqs] seeds: torch.Tensor, - num_tokens_across_dp: torch.Tensor | None = None, + dp_sync: DPSyncState | None = None, dummy_run: bool = False, skip_attn_for_dummy_run: bool = False, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,