From facd9a74a1cd1b9fed324cdc2cceb8d54fdad3d0 Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Thu, 3 Sep 2026 03:32:33 -0700 Subject: [PATCH] [Model Runner V2][Spec Decode] Skip DP sync for all speculator uniform decodes (#54856) --- vllm/v1/worker/gpu/dp_utils.py | 14 ++++++++++++- .../spec_decode/autoregressive/speculator.py | 8 +++++++- .../gpu/spec_decode/dflash/speculator.py | 13 +++++++----- vllm/v1/worker/gpu/spec_decode/speculator.py | 20 +++++++++++++++++++ 4 files changed, 48 insertions(+), 7 deletions(-) diff --git a/vllm/v1/worker/gpu/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index d1a935684f..d720d50ca0 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -33,6 +33,9 @@ class DPSyncState: # 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 + # Agreed upper bound on any rank's request count. Holds the padded count when + # a FULL descriptor imposed one, else the most any rank scheduled. + num_reqs: int def sync_cudagraph_and_dp_padding( @@ -53,17 +56,19 @@ def sync_cudagraph_and_dp_padding( """ assert dp_size > 1, "DP size must be greater than 1" group = get_dp_group().cpu_group - tensor = torch.zeros(4, dp_size, dtype=torch.int32, device="cpu") + tensor = torch.zeros(5, dp_size, dtype=torch.int32, device="cpu") tensor[0][dp_rank] = num_tokens tensor[1][dp_rank] = desired_batch_desc.cg_mode.value tensor[2][dp_rank] = uniform_token_count or 0 # (0 means None) tensor[3][dp_rank] = max_query_len or -1 # (-1 means None) + tensor[4][dp_rank] = num_reqs dist.all_reduce(tensor, group=group) num_tokens_across_dp = tensor[0] cg_mode_across_dp = tensor[1] uniform_token_counts_across_dp = tensor[2] max_query_lens_across_dp = tensor[3] + num_reqs_across_dp = tensor[4] # 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]) @@ -93,6 +98,7 @@ def sync_cudagraph_and_dp_padding( num_tokens_across_dp=num_tokens_across_dp, uniform_token_count=synced_uniform_token_count, eager=True, + num_reqs=int(num_reqs_across_dp.max().item()), ), ) @@ -126,6 +132,12 @@ def sync_cudagraph_and_dp_padding( num_tokens_across_dp=num_tokens_across_dp, uniform_token_count=synced_uniform_token_count, eager=False, + num_reqs=( + synced_desc.num_reqs + if synced_desc.cg_mode == CUDAGraphMode.FULL + and synced_desc.num_reqs is not None + else int(num_reqs_across_dp.max().item()) + ), ) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 8cdbc7322c..aba53e53c3 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -336,16 +336,22 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): advance_draft_positions=self.advance_draft_positions, ) + decode_batch_sync, num_batch_tokens = ( + self._build_uniform_batch_dp_sync(dp_sync, num_reqs, num_query_per_req=1) + if dp_sync is not None + else (None, num_reqs) + ) # Each request produces exactly 1 token per draft generation step, # enabling FULL graph replay. decode_batch_desc, decode_batch_sync = dispatch_cg_and_sync_dp( self.decode_cudagraph_manager, num_reqs, - num_reqs, + num_batch_tokens, uniform_token_count=1, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, + dp_sync=decode_batch_sync, ) num_tokens_across_dp = ( decode_batch_sync.num_tokens_across_dp diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index dd1bfdd5e7..f926bf87cc 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -367,9 +367,7 @@ class DFlashSpeculator(DraftModelSpeculator): num_query_tokens, attn_metadata=None, slot_mappings=None, - num_tokens_across_dp=( - dp_sync.num_tokens_across_dp if dp_sync is not None else None - ), + num_tokens_across_dp=None, cudagraph_runtime_mode=CUDAGraphMode.NONE, ) return self.draft_tokens[:num_reqs] @@ -429,17 +427,22 @@ class DFlashSpeculator(DraftModelSpeculator): context_slots, ) + batch_sync, num_batch_tokens = ( + self._build_uniform_batch_dp_sync(dp_sync, num_reqs, self.num_query_per_req) + if dp_sync is not None + else (None, num_query_tokens) + ) # Every DFlash step has exactly num_query_per_req tokens, so we can use FULL CGs batch_desc, batch_sync = dispatch_cg_and_sync_dp( self.query_cudagraph_manager, num_reqs, - num_query_tokens, + num_batch_tokens, uniform_token_count=self.num_query_per_req, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, + dp_sync=batch_sync, ) - num_reqs_padded = batch_desc.num_reqs or num_reqs num_tokens_padded = batch_desc.num_tokens num_tokens_across_dp = ( diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index d6e51aad7f..ecf36ce517 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod from collections.abc import Mapping +from dataclasses import replace from typing import Any import numpy as np @@ -409,3 +410,22 @@ class DraftModelSpeculator(BaseSpeculator): # idx_mapping for CG padded requests points to -1, which is ignored # during sampling to prevent writing stale values to draft logits. self.idx_mapping[num_reqs:].fill_(-1) + + def _build_uniform_batch_dp_sync( + self, + target_dp_sync: DPSyncState, + num_reqs: int, + num_query_per_req: int = 1, + ) -> tuple[DPSyncState, int]: + num_batch_tokens = target_dp_sync.num_reqs * num_query_per_req + assert num_reqs * num_query_per_req <= num_batch_tokens, ( + "reusing a DP sync that does not cover this batch's requests" + ) + return replace( + target_dp_sync, + num_tokens_across_dp=torch.full_like( + target_dp_sync.num_tokens_across_dp, num_batch_tokens + ), + uniform_token_count=num_query_per_req, + eager=False, + ), num_batch_tokens