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

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()
 

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,

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

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.

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.

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],

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.

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,