@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5acc1c4e4b

[Model Runner V2][Spec Decode] Skip DP sync before EAGLE/MTP draft prefill (#53694)

Giancarlo Delfin committed Aug 26, 2026 at 20:59 UTC 5acc1c4e4b8730298cbff7a4a7c68c814dc24fd7

8 files changed +156 -36

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

+1 -1

@@ -179,10 +179,10 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):

179 slot_mappings_by_layer=None,

180 hidden_states=None,

181 aux_hidden_states=None,

182 + dp_sync=None,

183 finished_req_ids=set(),

184 ec_connector_output=None,

185 routed_experts=None,

185 - num_tokens_across_dp=None,

186 )

187 runner.req_states = SimpleNamespace()

188

vllm/v1/worker/gpu/dp_utils.py

+106 -18

@@ -2,6 +2,8 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from __future__ import annotations

5 +from dataclasses import dataclass, replace

6 +

7 import torch

8 import torch.distributed as dist

@@ -13,6 +15,26 @@ from vllm.v1.worker.gpu.cudagraph_utils import (

15 )

18 +@dataclass(frozen=True)

19 +class DPSyncState:

20 + """What a `dispatch_cg_and_sync_dp` call agreed across DP ranks.

21 +

22 + Every field is identical on every rank, so callers can branch on them

23 + without disagreeing about whether to run a collective. Never add a per-rank

24 + value here.

25 + """

26 +

27 + # Per-rank token counts, for the forward context. All entries hold the

28 + # agreed padded count, unless `eager`, where nothing was padded.

29 + num_tokens_across_dp: torch.Tensor

30 + # Agreed uniform decode length. None means the ranks disagreed, which is an

31 + # answer, not "unknown".

32 + uniform_token_count: int | None

33 + # Whether the ranks agreed to run eager. A dispatch reusing this must run

34 + # eager too; no shape was agreed, so picking a graph per rank would diverge.

35 + eager: bool

36 +

37 +

38 def sync_cudagraph_and_dp_padding(

39 cudagraph_manager: CudaGraphManager | None,

40 desired_batch_desc: BatchExecutionDescriptor,

@@ -23,11 +45,11 @@ def sync_cudagraph_and_dp_padding(

45 dp_rank: int,

46 max_query_len: int | None = None,

47 num_active_loras: int = 0,

26 -) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:

48 +) -> tuple[BatchExecutionDescriptor, DPSyncState | None]:

49 """

50 Coordinates the batch descriptor and DP padding across all ranks.

30 - Returns (synced_batch_desc, num_tokens_across_dp).

52 + Returns (synced_batch_desc, sync). `sync` is None when no rank has work.

53 """

54 assert dp_size > 1, "DP size must be greater than 1"

55 group = get_dp_group().cpu_group

@@ -43,6 +65,13 @@ def sync_cudagraph_and_dp_padding(

65 uniform_token_counts_across_dp = tensor[2]

66 max_query_lens_across_dp = tensor[3]

68 + # If ranks disagree on the uniform token count, or its 0 (means None) set to None

69 + synced_uniform_token_count: int | None = int(uniform_token_counts_across_dp[0])

70 + if synced_uniform_token_count == 0 or not torch.all(

71 + uniform_token_counts_across_dp == synced_uniform_token_count

72 + ):

73 + synced_uniform_token_count = None

74 +

75 if torch.all(num_tokens_across_dp == 0).item():

76 synced_desc = BatchExecutionDescriptor(

77 cg_mode=CUDAGraphMode.NONE, num_tokens=0, num_reqs=0

@@ -53,24 +82,25 @@ def sync_cudagraph_and_dp_padding(

83 # If any rank wants to run eager, all ranks run eager

84 if synced_cg_mode == CUDAGraphMode.NONE:

56 - return BatchExecutionDescriptor(

57 - cg_mode=CUDAGraphMode.NONE,

58 - num_tokens=num_tokens,

59 - num_reqs=num_reqs,

60 - num_active_loras=desired_batch_desc.num_active_loras,

61 - ), num_tokens_across_dp

85 + return (

86 + BatchExecutionDescriptor(

87 + cg_mode=CUDAGraphMode.NONE,

88 + num_tokens=num_tokens,

89 + num_reqs=num_reqs,

90 + num_active_loras=desired_batch_desc.num_active_loras,

91 + ),

92 + DPSyncState(

93 + num_tokens_across_dp=num_tokens_across_dp,

94 + uniform_token_count=synced_uniform_token_count,

95 + eager=True,

96 + ),

97 + )

99 assert cudagraph_manager is not None, (

100 "cudagraph_manager should only be None during profile run, "

101 "where synced_cg_mode must be NONE across all DP ranks"

102 )

103 synced_num_tokens = int(num_tokens_across_dp.max().item())

68 - synced_uniform_token_count = uniform_token_counts_across_dp[0]

69 - # If ranks disagree on the uniform token count, or its 0 (means None) set to None

70 - if synced_uniform_token_count == 0 or not torch.all(

71 - uniform_token_counts_across_dp == synced_uniform_token_count

72 - ):

73 - synced_uniform_token_count = None

104

105 # Varlen decode graphs are selected by the query-length bound, so ranks must agree

106 # on it or they pad to different token counts below.

@@ -92,7 +122,11 @@ def sync_cudagraph_and_dp_padding(

122 # Update num_tokens_across_dp to reflect padded size.

123 num_tokens_across_dp[:] = synced_desc.num_tokens

124

95 - return synced_desc, num_tokens_across_dp

125 + return synced_desc, DPSyncState(

126 + num_tokens_across_dp=num_tokens_across_dp,

127 + uniform_token_count=synced_uniform_token_count,

128 + eager=False,

129 + )

130

131

132 def dispatch_cg_and_sync_dp(

@@ -105,8 +139,43 @@ def dispatch_cg_and_sync_dp(

139 max_query_len: int | None = None,

140 need_eager: bool = False,

141 num_active_loras: int = 0,

108 -) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:

109 - if need_eager:

142 + dp_sync: DPSyncState | None = None,

143 +) -> tuple[BatchExecutionDescriptor, DPSyncState | None]:

144 + """Pick a cudagraph descriptor for this batch, agreeing it across DP ranks.

145 +

146 + Runs a collective when dp_size > 1 so every rank dispatches to the same

147 + shape. Pass `dp_sync` from a dispatch already made over this same batch (a

148 + drafter's prefill runs the target's batch shape) to reuse that agreement

149 + instead, with no collective.

150 +

151 + Args:

152 + cudagraph_manager: Manager to dispatch against. May be None only when

153 + `need_eager` is True (profile run).

154 + num_reqs: Requests in this rank's batch.

155 + num_tokens: Tokens in this rank's batch, already padded by the caller.

156 + uniform_token_count: Per-request token count if this rank's batch is a

157 + uniform decode, else None. `dp_sync.uniform_token_count` takes its

158 + place when a sync is reused, since that one is agreed across ranks.

159 + dp_size: Data-parallel world size. 1 skips all cross-rank work.

160 + dp_rank: This rank's index in the DP group.

161 + max_query_len: Upper bound on per-request query length, for selecting

162 + varlen decode graphs. None means the graph must not constrain it.

163 + need_eager: Force `CUDAGraphMode.NONE` instead of dispatching.

164 + num_active_loras: Active LoRA count for this rank. Does not need

165 + cross-rank agreement; it never changes a bucket's token count.

166 + dp_sync: Agreement from a prior dispatch over this same batch, to reuse.

167 + Must come from a batch with this same padded `num_tokens` and the

168 + same `uniform_token_count`; `num_reqs` may differ, as neither

169 + depends on it. Passing a sync from a different batch is a caller

170 + error and trips an assert.

171 +

172 + Returns:

173 + (batch_desc, sync), where `sync` is this batch's agreement for a later

174 + dispatch to reuse. It is None when `dp_size` is 1 or no rank has work.

175 + """

176 + reuse_eager = dp_sync is not None and dp_sync.eager

177 +

178 + if need_eager or reuse_eager:

179 batch_desc = BatchExecutionDescriptor(

180 cg_mode=CUDAGraphMode.NONE,

181 num_tokens=num_tokens,

@@ -121,7 +190,7 @@ def dispatch_cg_and_sync_dp(

190 batch_desc = cudagraph_manager.dispatch(

191 num_reqs,

192 num_tokens,

124 - uniform_token_count,

193 + dp_sync.uniform_token_count if dp_sync is not None else uniform_token_count,

194 num_active_loras=num_active_loras,

195 max_query_len=max_query_len,

196 )

@@ -129,6 +198,25 @@ def dispatch_cg_and_sync_dp(

198 if dp_size == 1:

199 return batch_desc, None

200

201 + if dp_sync is not None:

202 + assert dp_sync.num_tokens_across_dp[dp_rank] == num_tokens, (

203 + "reusing a DP sync taken over a different batch"

204 + )

205 + assert (

206 + dp_sync.uniform_token_count is None

207 + or uniform_token_count == dp_sync.uniform_token_count

208 + ), "reusing a DP sync taken over a different batch"

209 + if not dp_sync.eager and batch_desc.num_tokens != num_tokens:

210 + # Capture sizes can differ between managers, so this one may

211 + # pad further. Every rank pads alike, so report what will run.

212 + dp_sync = replace(

213 + dp_sync,

214 + num_tokens_across_dp=torch.full_like(

215 + dp_sync.num_tokens_across_dp, batch_desc.num_tokens

216 + ),

217 + )

218 + return batch_desc, dp_sync

219 +

220 return sync_cudagraph_and_dp_padding(

221 cudagraph_manager,

222 batch_desc,

vllm/v1/worker/gpu/model_runner.py

+11 -3

@@ -97,7 +97,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import (

97 from vllm.v1.worker.gpu.cudagraph_utils import (

98 profile_cudagraph_memory as _profile_cudagraph_memory,

99 )

100 -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp

100 +from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp

101 from vllm.v1.worker.gpu.ec_connector import get_ec_connector

102 from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after

103 from vllm.v1.worker.gpu.input_batch import (

@@ -752,6 +752,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

752 slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer

753 hidden_states = self.execute_model_state.hidden_states

754 aux_hidden_states = self.execute_model_state.aux_hidden_states

755 + dp_sync = self.execute_model_state.dp_sync

756 self.execute_model_state = None

757

758 self.step_timing.forward_end()

@@ -796,6 +797,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

797 next_prefill_tokens=self.req_states.next_prefill_tokens,

798 temperature=self.sampler.sampling_states.temperature.gpu,

799 seeds=self.sampler.sampling_states.seeds.gpu,

800 + dp_sync=dp_sync,

801 dummy_run=True,

802 skip_attn_for_dummy_run=skip_attn,

803 mm_inputs=mm_inputs,

@@ -1544,7 +1546,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1546 # cross-attention cache with dynamic encoder outputs.

1547 skip_compiled = True

1548

1547 - batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(

1549 + batch_desc, dp_sync = dispatch_cg_and_sync_dp(

1550 self.cudagraph_manager,

1551 num_reqs,

1552 num_toks,

@@ -1731,7 +1733,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1733 self.vllm_config,

1734 num_tokens=input_batch.num_tokens_after_padding,

1735 cudagraph_runtime_mode=batch_desc.cg_mode,

1734 - num_tokens_across_dp=num_tokens_across_dp,

1736 + num_tokens_across_dp=(

1737 + dp_sync.num_tokens_across_dp if dp_sync is not None else None

1738 + ),

1739 batch_descriptor=batch_descriptor,

1740 slot_mapping=slot_mappings_by_layer,

1741 skip_compiled=skip_compiled,

@@ -1777,6 +1781,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1781 slot_mappings_by_layer=slot_mappings_by_layer,

1782 hidden_states=hidden_states,

1783 aux_hidden_states=aux_hidden_states,

1784 + dp_sync=dp_sync,

1785 finished_req_ids=finished_req_ids,

1786 ec_connector_output=ec_connector_output,

1787 routed_experts=routed_experts,

@@ -1801,6 +1806,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1806 slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer

1807 hidden_states = self.execute_model_state.hidden_states

1808 aux_hidden_states = self.execute_model_state.aux_hidden_states

1809 + dp_sync = self.execute_model_state.dp_sync

1810 finished_req_ids = self.execute_model_state.finished_req_ids

1811 ec_connector_output = self.execute_model_state.ec_connector_output

1812 routed_experts = self.execute_model_state.routed_experts

@@ -1922,6 +1928,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1928 self.req_states.next_prefill_tokens,

1929 self.sampler.sampling_states.temperature.gpu,

1930 self.sampler.sampling_states.seeds.gpu,

1931 + dp_sync=dp_sync,

1932 mm_inputs=mm_inputs,

1933 )

1934 self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens

@@ -2066,6 +2073,7 @@ class ExecuteModelState(NamedTuple):

2073 slot_mappings_by_layer: dict[str, torch.Tensor] | None

2074 hidden_states: torch.Tensor | None

2075 aux_hidden_states: list[torch.Tensor] | None

2076 + dp_sync: DPSyncState | None

2077 finished_req_ids: set[str]

2078 ec_connector_output: ECConnectorOutput | None

2079 routed_experts: RoutedExpertsTensors | None

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

+15 -4

@@ -14,7 +14,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig

14 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer

15 from vllm.v1.worker.gpu.block_table import BlockTables

16 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor

17 -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp

17 +from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp

18 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

19 from vllm.v1.worker.gpu.model_states.interface import ModelState

20 from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (

@@ -219,7 +219,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

219 temperature: torch.Tensor,

220 # [max_num_reqs]

221 seeds: torch.Tensor,

222 - num_tokens_across_dp: torch.Tensor | None = None,

222 + dp_sync: DPSyncState | None = None,

223 dummy_run: bool = False,

224 skip_attn_for_dummy_run: bool = False,

225 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,

@@ -279,7 +279,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

279 max_query_len,

280 input_batch.has_prefill,

281 )

282 - prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(

282 + prefill_batch_desc, prefill_batch_sync = dispatch_cg_and_sync_dp(

283 self.prefill_cudagraph_manager,

284 num_reqs,

285 num_tokens_padded,

@@ -287,6 +287,12 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

287 dp_size=self.dp_size,

288 dp_rank=self.dp_rank,

289 need_eager=is_profile,

290 + dp_sync=dp_sync,

291 + )

292 + num_tokens_across_dp = (

293 + prefill_batch_sync.num_tokens_across_dp

294 + if prefill_batch_sync is not None

295 + else None

296 )

297

298 self._prepare_eplb_forward(num_tokens)

@@ -328,7 +334,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

334

335 # Each request produces exactly 1 token per draft generation step,

336 # enabling FULL graph replay.

331 - decode_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(

337 + decode_batch_desc, decode_batch_sync = dispatch_cg_and_sync_dp(

338 self.decode_cudagraph_manager,

339 num_reqs,

340 num_reqs,

@@ -337,6 +343,11 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

343 dp_rank=self.dp_rank,

344 need_eager=is_profile,

345 )

346 + num_tokens_across_dp = (

347 + decode_batch_sync.num_tokens_across_dp

348 + if decode_batch_sync is not None

349 + else None

350 + )

351

352 self.on_multi_step_decode_begin(num_reqs)

353 # Generate the remaining num_speculative_steps - 1 draft tokens.

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

+9 -4

@@ -19,7 +19,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig

19 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer

20 from vllm.v1.worker.gpu.block_table import BlockTables

21 from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens

22 -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp

22 +from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp

23 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

24 from vllm.v1.worker.gpu.model_states.interface import ModelState

25 from vllm.v1.worker.gpu.spec_decode.dflash.cudagraph import DFlashCudaGraphManager

@@ -335,7 +335,7 @@ class DFlashSpeculator(DraftModelSpeculator):

335 temperature: torch.Tensor,

336 # [max_num_reqs]

337 seeds: torch.Tensor,

338 - num_tokens_across_dp: torch.Tensor | None = None,

338 + dp_sync: DPSyncState | None = None,

339 dummy_run: bool = False,

340 skip_attn_for_dummy_run: bool = False,

341 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,

@@ -377,7 +377,9 @@ class DFlashSpeculator(DraftModelSpeculator):

377 num_query_tokens,

378 attn_metadata=None,

379 slot_mappings=None,

380 - num_tokens_across_dp=num_tokens_across_dp,

380 + num_tokens_across_dp=(

381 + dp_sync.num_tokens_across_dp if dp_sync is not None else None

382 + ),

383 cudagraph_runtime_mode=CUDAGraphMode.NONE,

384 )

385 return self.draft_tokens[:num_reqs]

@@ -438,7 +440,7 @@ class DFlashSpeculator(DraftModelSpeculator):

440 )

441

442 # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs

441 - batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(

443 + batch_desc, batch_sync = dispatch_cg_and_sync_dp(

444 self.query_cudagraph_manager,

445 num_reqs,

446 num_query_tokens,

@@ -450,6 +452,9 @@ class DFlashSpeculator(DraftModelSpeculator):

452

453 num_reqs_padded = batch_desc.num_reqs or num_reqs

454 num_tokens_padded = batch_desc.num_tokens

455 + num_tokens_across_dp = (

456 + batch_sync.num_tokens_across_dp if batch_sync is not None else None

457 + )

458

459 # Rebuild the draft attention metadata even when replaying the FULL

460 # graph so that any attention metadata builder state is updated.

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

+5 -2

@@ -10,6 +10,7 @@ from vllm.config import VllmConfig

10 from vllm.config.compilation import CUDAGraphMode

11 from vllm.forward_context import set_forward_context

12 from vllm.model_executor.model_loader import get_model

13 +from vllm.v1.worker.gpu.dp_utils import DPSyncState

14 from vllm.v1.worker.gpu.input_batch import InputBatch

15 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator

@@ -97,7 +98,7 @@ class ExtractHiddenStatesSpeculator(DraftModelSpeculator):

98 next_prefill_tokens: torch.Tensor,

99 temperature: torch.Tensor,

100 seeds: torch.Tensor,

100 - num_tokens_across_dp: torch.Tensor | None = None,

101 + dp_sync: DPSyncState | None = None,

102 dummy_run: bool = False,

103 skip_attn_for_dummy_run: bool = False,

104 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,

@@ -143,7 +144,9 @@ class ExtractHiddenStatesSpeculator(DraftModelSpeculator):

144 draft_attn_metadata,

145 self.vllm_config,

146 num_tokens=num_tokens,

146 - num_tokens_across_dp=num_tokens_across_dp,

147 + num_tokens_across_dp=(

148 + dp_sync.num_tokens_across_dp if dp_sync is not None else None

149 + ),

150 cudagraph_runtime_mode=CUDAGraphMode.NONE,

151 slot_mapping=draft_slot_mappings,

152 is_padding=input_batch.is_padding[:num_tokens],

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

+7 -3

@@ -12,7 +12,7 @@ from vllm.logger import init_logger

12 from vllm.triton_utils import tl, triton

13 from vllm.v1.attention.backends.utils import PAD_SLOT_ID

14 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer

15 -from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp

15 +from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp

16 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

17 from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (

18 SpeculatorCudaGraphManager,

@@ -150,7 +150,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):

150 temperature: torch.Tensor,

151 # [max_num_reqs]

152 seeds: torch.Tensor,

153 - num_tokens_across_dp: torch.Tensor | None = None,

153 + dp_sync: DPSyncState | None = None,

154 dummy_run: bool = False,

155 skip_attn_for_dummy_run: bool = False,

156 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,

@@ -191,7 +191,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):

191 uniform_token_count = get_uniform_decode_token_count(

192 num_reqs, num_tokens, max_query_len, input_batch.has_prefill

193 )

194 - batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(

194 + batch_desc, batch_sync = dispatch_cg_and_sync_dp(

195 self.cudagraph_manager,

196 num_reqs,

197 input_batch.num_tokens_after_padding,

@@ -199,6 +199,10 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):

199 dp_size=self.dp_size,

200 dp_rank=self.dp_rank,

201 need_eager=is_profile,

202 + dp_sync=dp_sync,

203 + )

204 + num_tokens_across_dp = (

205 + batch_sync.num_tokens_across_dp if batch_sync is not None else None

206 )

207

208 # Rebuild the slot mappings and attention metadata.

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

+2 -1

@@ -21,6 +21,7 @@ from vllm.v1.worker.gpu.attn_utils import (

21 init_attn_backend,

22 )

23 from vllm.v1.worker.gpu.block_table import BlockTables

24 +from vllm.v1.worker.gpu.dp_utils import DPSyncState

25 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

26 from vllm.v1.worker.gpu.model_states.interface import ModelState

27 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

@@ -60,7 +61,7 @@ class BaseSpeculator(ABC):

61 temperature: torch.Tensor,

62 # [max_num_reqs]

63 seeds: torch.Tensor,

63 - num_tokens_across_dp: torch.Tensor | None = None,

64 + dp_sync: DPSyncState | None = None,

65 dummy_run: bool = False,

66 skip_attn_for_dummy_run: bool = False,

67 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,