243 lines
9.4 KiB
Python
243 lines
9.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, replace
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from vllm.config.compilation import CUDAGraphMode
|
|
from vllm.distributed.parallel_state import get_dp_group
|
|
from vllm.v1.worker.gpu.cudagraph_utils import (
|
|
BatchExecutionDescriptor,
|
|
CudaGraphManager,
|
|
)
|
|
|
|
|
|
@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
|
|
# 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(
|
|
cudagraph_manager: CudaGraphManager | None,
|
|
desired_batch_desc: BatchExecutionDescriptor,
|
|
num_tokens: int,
|
|
num_reqs: int,
|
|
uniform_token_count: int | None,
|
|
dp_size: int,
|
|
dp_rank: int,
|
|
max_query_len: int | None = None,
|
|
num_active_loras: int = 0,
|
|
) -> tuple[BatchExecutionDescriptor, DPSyncState | None]:
|
|
"""
|
|
Coordinates the batch descriptor and DP padding across all ranks.
|
|
|
|
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
|
|
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])
|
|
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
|
|
)
|
|
return synced_desc, None
|
|
|
|
synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item()))
|
|
|
|
# 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,
|
|
),
|
|
DPSyncState(
|
|
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()),
|
|
),
|
|
)
|
|
|
|
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())
|
|
|
|
# Varlen decode graphs are selected by the query-length bound, so ranks must agree
|
|
# on it or they pad to different token counts below.
|
|
synced_max_query_len: int | None = None
|
|
if bool(torch.all(max_query_lens_across_dp != -1).item()):
|
|
synced_max_query_len = int(max_query_lens_across_dp.max().item())
|
|
|
|
# Dispatch for the final synced values, use num_reqs instead of synced_num_reqs
|
|
# so we don't perform request padding for PIECEWISE graphs.
|
|
# num_active_loras is per-rank and doesn't need cross-rank agreement.
|
|
synced_desc = cudagraph_manager.dispatch(
|
|
num_reqs,
|
|
synced_num_tokens,
|
|
synced_uniform_token_count,
|
|
num_active_loras=num_active_loras,
|
|
max_query_len=synced_max_query_len,
|
|
)
|
|
|
|
# Update num_tokens_across_dp to reflect padded size.
|
|
num_tokens_across_dp[:] = synced_desc.num_tokens
|
|
|
|
return synced_desc, DPSyncState(
|
|
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())
|
|
),
|
|
)
|
|
|
|
|
|
def dispatch_cg_and_sync_dp(
|
|
cudagraph_manager: CudaGraphManager | None,
|
|
num_reqs: int,
|
|
num_tokens: int,
|
|
uniform_token_count: int | None,
|
|
dp_size: int,
|
|
dp_rank: int,
|
|
max_query_len: int | None = None,
|
|
need_eager: bool = False,
|
|
num_active_loras: int = 0,
|
|
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,
|
|
num_reqs=num_reqs,
|
|
num_active_loras=num_active_loras,
|
|
)
|
|
else:
|
|
assert cudagraph_manager is not None, (
|
|
"cudagraph_manager should only be None during profile run, "
|
|
"where need_eager must be True"
|
|
)
|
|
batch_desc = cudagraph_manager.dispatch(
|
|
num_reqs,
|
|
num_tokens,
|
|
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,
|
|
)
|
|
|
|
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,
|
|
num_tokens,
|
|
num_reqs,
|
|
uniform_token_count,
|
|
dp_size,
|
|
dp_rank,
|
|
max_query_len=max_query_len,
|
|
num_active_loras=num_active_loras,
|
|
)
|