[Model Runner V2][Spec Decode] Skip DP sync for all speculator uniform decodes (#54856)
Giancarlo Delfin committed
Sep 3, 2026 at 03:32 UTC
facd9a74a1cd1b9fed324cdc2cceb8d54fdad3d0
4 files changed
+48
-7
vllm/v1/worker/gpu/dp_utils.py
+13
-1
@@ -33,6 +33,9 @@ class DPSyncState:
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
+ # Agreed upper bound on any rank's request count. Holds the padded count when
37
+ # a FULL descriptor imposed one, else the most any rank scheduled.
38
+ num_reqs: int
39
40
41
def sync_cudagraph_and_dp_padding(
@@ -53,17 +56,19 @@ def sync_cudagraph_and_dp_padding(
56
"""
57
assert dp_size > 1, "DP size must be greater than 1"
58
group = get_dp_group().cpu_group
56
- tensor = torch.zeros(4, dp_size, dtype=torch.int32, device="cpu")
59
+ tensor = torch.zeros(5, dp_size, dtype=torch.int32, device="cpu")
60
tensor[0][dp_rank] = num_tokens
61
tensor[1][dp_rank] = desired_batch_desc.cg_mode.value
62
tensor[2][dp_rank] = uniform_token_count or 0 # (0 means None)
63
tensor[3][dp_rank] = max_query_len or -1 # (-1 means None)
64
+ tensor[4][dp_rank] = num_reqs
65
dist.all_reduce(tensor, group=group)
66
67
num_tokens_across_dp = tensor[0]
68
cg_mode_across_dp = tensor[1]
69
uniform_token_counts_across_dp = tensor[2]
70
max_query_lens_across_dp = tensor[3]
71
+ num_reqs_across_dp = tensor[4]
72
73
# If ranks disagree on the uniform token count, or its 0 (means None) set to None
74
synced_uniform_token_count: int | None = int(uniform_token_counts_across_dp[0])
@@ -93,6 +98,7 @@ def sync_cudagraph_and_dp_padding(
98
num_tokens_across_dp=num_tokens_across_dp,
99
uniform_token_count=synced_uniform_token_count,
100
eager=True,
101
+ num_reqs=int(num_reqs_across_dp.max().item()),
102
),
103
)
104
@@ -126,6 +132,12 @@ def sync_cudagraph_and_dp_padding(
132
num_tokens_across_dp=num_tokens_across_dp,
133
uniform_token_count=synced_uniform_token_count,
134
eager=False,
135
+ num_reqs=(
136
+ synced_desc.num_reqs
137
+ if synced_desc.cg_mode == CUDAGraphMode.FULL
138
+ and synced_desc.num_reqs is not None
139
+ else int(num_reqs_across_dp.max().item())
140
+ ),
141
)
142
143
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+7
-1
@@ -336,16 +336,22 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
336
advance_draft_positions=self.advance_draft_positions,
337
)
338
339
+ decode_batch_sync, num_batch_tokens = (
340
+ self._build_uniform_batch_dp_sync(dp_sync, num_reqs, num_query_per_req=1)
341
+ if dp_sync is not None
342
+ else (None, num_reqs)
343
+ )
344
# Each request produces exactly 1 token per draft generation step,
345
# enabling FULL graph replay.
346
decode_batch_desc, decode_batch_sync = dispatch_cg_and_sync_dp(
347
self.decode_cudagraph_manager,
348
num_reqs,
344
- num_reqs,
349
+ num_batch_tokens,
350
uniform_token_count=1,
351
dp_size=self.dp_size,
352
dp_rank=self.dp_rank,
353
need_eager=is_profile,
354
+ dp_sync=decode_batch_sync,
355
)
356
num_tokens_across_dp = (
357
decode_batch_sync.num_tokens_across_dp
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+8
-5
@@ -367,9 +367,7 @@ class DFlashSpeculator(DraftModelSpeculator):
367
num_query_tokens,
368
attn_metadata=None,
369
slot_mappings=None,
370
- num_tokens_across_dp=(
371
- dp_sync.num_tokens_across_dp if dp_sync is not None else None
372
- ),
370
+ num_tokens_across_dp=None,
371
cudagraph_runtime_mode=CUDAGraphMode.NONE,
372
)
373
return self.draft_tokens[:num_reqs]
@@ -429,17 +427,22 @@ class DFlashSpeculator(DraftModelSpeculator):
427
context_slots,
428
)
429
430
+ batch_sync, num_batch_tokens = (
431
+ self._build_uniform_batch_dp_sync(dp_sync, num_reqs, self.num_query_per_req)
432
+ if dp_sync is not None
433
+ else (None, num_query_tokens)
434
+ )
435
# Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs
436
batch_desc, batch_sync = dispatch_cg_and_sync_dp(
437
self.query_cudagraph_manager,
438
num_reqs,
436
- num_query_tokens,
439
+ num_batch_tokens,
440
uniform_token_count=self.num_query_per_req,
441
dp_size=self.dp_size,
442
dp_rank=self.dp_rank,
443
need_eager=is_profile,
444
+ dp_sync=batch_sync,
445
)
442
-
446
num_reqs_padded = batch_desc.num_reqs or num_reqs
447
num_tokens_padded = batch_desc.num_tokens
448
num_tokens_across_dp = (
vllm/v1/worker/gpu/spec_decode/speculator.py
+20
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from abc import ABC, abstractmethod
4
from collections.abc import Mapping
5
+from dataclasses import replace
6
from typing import Any
7
8
import numpy as np
@@ -409,3 +410,22 @@ class DraftModelSpeculator(BaseSpeculator):
410
# idx_mapping for CG padded requests points to -1, which is ignored
411
# during sampling to prevent writing stale values to draft logits.
412
self.idx_mapping[num_reqs:].fill_(-1)
413
+
414
+ def _build_uniform_batch_dp_sync(
415
+ self,
416
+ target_dp_sync: DPSyncState,
417
+ num_reqs: int,
418
+ num_query_per_req: int = 1,
419
+ ) -> tuple[DPSyncState, int]:
420
+ num_batch_tokens = target_dp_sync.num_reqs * num_query_per_req
421
+ assert num_reqs * num_query_per_req <= num_batch_tokens, (
422
+ "reusing a DP sync that does not cover this batch's requests"
423
+ )
424
+ return replace(
425
+ target_dp_sync,
426
+ num_tokens_across_dp=torch.full_like(
427
+ target_dp_sync.num_tokens_across_dp, num_batch_tokens
428
+ ),
429
+ uniform_token_count=num_query_per_req,
430
+ eager=False,
431
+ ), num_batch_tokens