A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Core] Simplify the Dp padding/should ubatch coordination logic (#25768)
Signed-off-by: Sage Moore <sage@neuralmagic.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Sage Moore committed
Oct 6, 2025 at 18:57 UTC
2111b4643c1039894dc730a32eec78e937767532
10 files changed
+297
-462
tests/v1/attention/test_attention_splitting.py
+1
-1
@@ -13,7 +13,7 @@ from vllm.v1.attention.backends.utils import (
13
split_attn_metadata,
14
split_decodes_and_prefills,
15
)
16
-from vllm.v1.worker.ubatch_splitting import create_ubatch_slices
16
+from vllm.v1.worker.ubatch_utils import create_ubatch_slices
17
18
19
@pytest.fixture
vllm/config/parallel.py
+4
@@ -152,6 +152,10 @@ class ParallelConfig:
152
threshold, microbatching will be used. Otherwise, the request will be
153
processed in a single batch."""
154
155
+ disable_nccl_for_dp_synchronization: bool = False
156
+ """Forces the dp synchronization logic in vllm/v1/worker/dp_utils.py
157
+ to use Gloo instead of NCCL for its all reduce"""
158
+
159
ray_workers_use_nsight: bool = False
160
"""Whether to profile Ray workers with nsight, see https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html#profiling-nsight-profiler."""
161
vllm/engine/arg_utils.py
+8
@@ -365,6 +365,9 @@ class EngineArgs:
365
enable_dbo: bool = ParallelConfig.enable_dbo
366
dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold
367
dbo_prefill_token_threshold: int = ParallelConfig.dbo_prefill_token_threshold
368
+ disable_nccl_for_dp_synchronization: bool = (
369
+ ParallelConfig.disable_nccl_for_dp_synchronization
370
+ )
371
eplb_config: EPLBConfig = get_field(ParallelConfig, "eplb_config")
372
enable_eplb: bool = ParallelConfig.enable_eplb
373
expert_placement_strategy: ExpertPlacementStrategy = (
@@ -760,6 +763,10 @@ class EngineArgs:
763
"--dbo-prefill-token-threshold",
764
**parallel_kwargs["dbo_prefill_token_threshold"],
765
)
766
+ parallel_group.add_argument(
767
+ "--disable-nccl-for-dp-synchronization",
768
+ **parallel_kwargs["disable_nccl_for_dp_synchronization"],
769
+ )
770
parallel_group.add_argument("--enable-eplb", **parallel_kwargs["enable_eplb"])
771
parallel_group.add_argument("--eplb-config", **parallel_kwargs["eplb_config"])
772
parallel_group.add_argument(
@@ -1437,6 +1444,7 @@ class EngineArgs:
1444
enable_dbo=self.enable_dbo,
1445
dbo_decode_token_threshold=self.dbo_decode_token_threshold,
1446
dbo_prefill_token_threshold=self.dbo_prefill_token_threshold,
1447
+ disable_nccl_for_dp_synchronization=self.disable_nccl_for_dp_synchronization,
1448
enable_eplb=self.enable_eplb,
1449
eplb_config=self.eplb_config,
1450
expert_placement_strategy=self.expert_placement_strategy,
vllm/envs.py
-7
@@ -95,7 +95,6 @@ if TYPE_CHECKING:
95
VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
96
VLLM_SKIP_P2P_CHECK: bool = False
97
VLLM_DISABLED_KERNELS: list[str] = []
98
- VLLM_DISABLE_NCCL_FOR_DP_SYNCHRONIZATION: bool = False
98
VLLM_DISABLE_PYNCCL: bool = False
99
VLLM_USE_V1: bool = True
100
VLLM_ROCM_USE_AITER: bool = False
@@ -830,12 +829,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
829
"VLLM_DISABLED_KERNELS": lambda: []
830
if "VLLM_DISABLED_KERNELS" not in os.environ
831
else os.environ["VLLM_DISABLED_KERNELS"].split(","),
833
- # Swaps the all reduce backend that we use to coordinate the DP padding
834
- # information from NCCL to gloo.
835
- "VLLM_DISABLE_NCCL_FOR_DP_SYNCHRONIZATION": lambda: (
836
- os.getenv("VLLM_DISABLE_NCCL_FOR_DP_SYNCHRONIZATION", "False").lower()
837
- in ("true", "1")
838
- ),
832
# Disable pynccl (using torch.distributed instead)
833
"VLLM_DISABLE_PYNCCL": lambda: (
834
os.getenv("VLLM_DISABLE_PYNCCL", "False").lower() in ("true", "1")
vllm/forward_context.py
+9
-120
@@ -8,13 +8,11 @@ from dataclasses import dataclass
8
from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union
9
10
import torch
11
-import torch.distributed as dist
11
12
import vllm.envs as envs
13
from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
14
from vllm.logger import init_logger
16
-from vllm.platforms import current_platform
17
-from vllm.v1.worker.ubatch_utils import UBatchSlices, is_second_ubatch_empty
15
+from vllm.v1.worker.ubatch_utils import UBatchSlices
16
17
if TYPE_CHECKING:
18
from vllm.attention.backends.abstract import AttentionMetadata
@@ -87,129 +85,22 @@ class DPMetadata:
85
# NOTE: local_sizes should only be set by the chunked_sizes context manager
86
local_sizes: Optional[list[int]] = None
87
90
- @staticmethod
91
- def num_tokens_across_dp(
92
- num_tokens: int, dp_size: int, dp_rank: int
93
- ) -> torch.Tensor:
94
- """
95
- Gather the num_tokens across all DP ranks and return results in a
96
- CPU tensor of size dp_size.
97
- """
98
- from vllm.distributed.parallel_state import get_dp_group
99
-
100
- device = current_platform.device_type
101
- group = get_dp_group().device_group
102
-
103
- # Transfering this tensor from GPU to CPU will introduce a GPU sync
104
- # point that could adversely affect performance of vllm with asynch
105
- # scheduling. This environment variable exists to quickly disable
106
- # this optimization if we run into this case.
107
- if envs.VLLM_DISABLE_NCCL_FOR_DP_SYNCHRONIZATION:
108
- logger.info_once(
109
- "Using CPU all reduce to syncronize DP padding between ranks."
110
- )
111
- device = "cpu"
112
- group = get_dp_group().cpu_group
113
- num_tokens_across_dp = [0] * dp_size
114
- num_tokens_across_dp[dp_rank] = num_tokens
115
- num_tokens_tensor = torch.tensor(
116
- num_tokens_across_dp, device=device, dtype=torch.int32
117
- )
118
- dist.all_reduce(num_tokens_tensor, group=group)
119
- return num_tokens_tensor.cpu()
120
-
121
- # Get the cumulative tokens across sequence parallel ranks.
122
- # In this case the input to the MoEs will be distributed w.r.t both
123
- # DP and TP rank.
124
- # When sp_size==1, this is just the cummulative num tokens across DP.
125
- def cu_tokens_across_sp(self, sp_size: int) -> torch.Tensor:
126
- num_tokens_across_sp_cpu = (
127
- self.num_tokens_across_dp_cpu - 1 + sp_size
128
- ) // sp_size
129
- num_tokens_across_sp_cpu = num_tokens_across_sp_cpu.repeat_interleave(sp_size)
130
- return torch.cumsum(num_tokens_across_sp_cpu, dim=0)
131
-
132
- @staticmethod
133
- def should_ubatch_across_dp(
134
- should_ubatch: bool,
135
- orig_num_tokens_per_ubatch: int,
136
- padded_num_tokens_per_ubatch: int,
137
- dp_size: int,
138
- dp_rank: int,
139
- ) -> tuple[bool, Optional[torch.Tensor]]:
140
- """
141
- 1. Decides if each DP rank is going to microbatch. Either all ranks
142
- run with microbatching or none of them do. If this function decides
143
- not to run with microbatching. It will "abort" meaning that no padding
144
- information will be returned to the caller. It will return (False, None)
145
-
146
- 2. Determines the total number of tokens that each rank will run.
147
- All ranks will be padded out so that the run with the same number
148
- of tokens
149
-
150
- Returns: tuple[
151
- should_ubatch: Are all DP ranks going to microbatch
152
- num_tokens_after_padding: A tensor containing the total number of
153
- tokens per-microbatch for each DP rank including padding. Will be
154
- None if should_ubatch if False
155
- ]
156
- """
157
-
158
- device = current_platform.device_type
159
- tensor = torch.zeros(3, dp_size, device=device, dtype=torch.int32)
160
- tensor[0][dp_rank] = orig_num_tokens_per_ubatch
161
- tensor[1][dp_rank] = padded_num_tokens_per_ubatch
162
- tensor[2][dp_rank] = 1 if should_ubatch else 0
163
-
164
- from vllm.distributed.parallel_state import get_dp_group
165
-
166
- dist.all_reduce(tensor, group=get_dp_group().device_group)
167
-
168
- result: bool = bool(torch.all(tensor[2] == 1).item())
169
- if not result:
170
- return result, None
171
-
172
- orig_num_tokens_tensor = tensor[0, :]
173
- padded_num_tokens_tensor = tensor[1, :]
174
-
175
- orig_min_num_tokens = int(orig_num_tokens_tensor.min().item())
176
- padded_max_num_tokens = int(padded_num_tokens_tensor.max().item())
177
- if is_second_ubatch_empty(orig_min_num_tokens, padded_max_num_tokens):
178
- logger.debug(
179
- "Aborting ubatching %s %s", orig_min_num_tokens, padded_max_num_tokens
180
- )
181
- return False, None
182
- return result, padded_num_tokens_tensor.cpu()
183
-
88
@staticmethod
89
def make(
90
parallel_config: ParallelConfig,
187
- attn_metadata: Any,
91
num_tokens: int,
189
- num_tokens_across_dp_cpu: Optional[torch.Tensor] = None,
92
+ num_tokens_across_dp_cpu: torch.Tensor,
93
) -> "DPMetadata":
94
+ assert num_tokens_across_dp_cpu is not None
95
assert parallel_config.data_parallel_size > 1
192
- dp_size = parallel_config.data_parallel_size
96
dp_rank = parallel_config.data_parallel_rank
194
- if attn_metadata is not None and hasattr(attn_metadata, "num_prefill_tokens"):
195
- # for v0 attention backends
196
- batchsize = (
197
- attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens
198
- )
199
- else:
200
- # for v1 attention backends or no attn_metadata
201
- batchsize = num_tokens
97
+ batchsize = num_tokens
98
99
# If num_tokens_across_dp is None, it will be computed by all_reduce
100
# Otherwise, num_tokens_across_dp[dp_rank] should be equal to batchsize
205
- assert (
206
- num_tokens_across_dp_cpu is None
207
- or num_tokens_across_dp_cpu[dp_rank] == batchsize
208
- ), f"{num_tokens_across_dp_cpu[dp_rank]} {batchsize}"
209
- if num_tokens_across_dp_cpu is None:
210
- num_tokens_across_dp_cpu = DPMetadata.num_tokens_across_dp(
211
- batchsize, dp_size, dp_rank
212
- )
101
+ assert num_tokens_across_dp_cpu[dp_rank] == batchsize, (
102
+ f"{num_tokens_across_dp_cpu[dp_rank]} {batchsize}"
103
+ )
104
max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp_cpu)
105
return DPMetadata(max_tokens_across_dp_cpu, num_tokens_across_dp_cpu)
106
@@ -376,11 +267,9 @@ def set_forward_context(
267
if vllm_config.parallel_config.data_parallel_size > 1 and (
268
attn_metadata is not None or num_tokens is not None
269
):
270
+ assert num_tokens_across_dp is not None
271
dp_metadata = DPMetadata.make(
380
- vllm_config.parallel_config,
381
- attn_metadata,
382
- num_tokens or 0,
383
- num_tokens_across_dp,
272
+ vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp
273
)
274
275
forward_context = create_forward_context(
vllm/v1/worker/dp_utils.py
new
+177
@@ -0,0 +1,177 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Optional
4
+
5
+import numpy as np
6
+import torch
7
+import torch.distributed as dist
8
+
9
+from vllm.config import ParallelConfig
10
+from vllm.distributed.parallel_state import get_dp_group
11
+from vllm.logger import init_logger
12
+from vllm.platforms import current_platform
13
+from vllm.v1.worker.ubatch_utils import (
14
+ UBatchSlices,
15
+ check_ubatch_thresholds,
16
+ create_ubatch_slices,
17
+ is_second_ubatch_empty,
18
+)
19
+
20
+logger = init_logger(__name__)
21
+
22
+
23
+def _get_device_and_group(parallel_config: ParallelConfig):
24
+ device = current_platform.device_type
25
+ group = get_dp_group().device_group
26
+
27
+ # Transfering this tensor from GPU to CPU will introduce a GPU sync
28
+ # point that could adversely affect performance of vllm with asynch
29
+ # scheduling. This environment variable exists to quickly disable
30
+ # this optimization if we run into this case.
31
+ if parallel_config.disable_nccl_for_dp_synchronization:
32
+ logger.info_once("Using CPU all reduce to syncronize DP padding between ranks.")
33
+ device = "cpu"
34
+ group = get_dp_group().cpu_group
35
+ return device, group
36
+
37
+
38
+def _run_ar(
39
+ should_ubatch: bool,
40
+ orig_num_tokens_per_ubatch: int,
41
+ padded_num_tokens_per_ubatch: int,
42
+ parallel_config: ParallelConfig,
43
+) -> torch.Tensor:
44
+ dp_size = parallel_config.data_parallel_size
45
+ dp_rank = parallel_config.data_parallel_rank
46
+ device, group = _get_device_and_group(parallel_config)
47
+ tensor = torch.zeros(3, dp_size, device=device, dtype=torch.int32)
48
+ tensor[0][dp_rank] = orig_num_tokens_per_ubatch
49
+ tensor[1][dp_rank] = padded_num_tokens_per_ubatch
50
+ tensor[2][dp_rank] = 1 if should_ubatch else 0
51
+ dist.all_reduce(tensor, group=group)
52
+ return tensor
53
+
54
+
55
+def _post_process_ubatch(tensor: torch.Tensor) -> bool:
56
+ orig_num_tokens_tensor = tensor[0, :]
57
+ padded_num_tokens_tensor = tensor[1, :]
58
+
59
+ # First determine if we are going to be ubatching.
60
+ should_ubatch: bool = bool(torch.all(tensor[2] == 1).item())
61
+ if not should_ubatch:
62
+ return False
63
+ # If the DP ranks are planning to ubatch, make sure that
64
+ # there are no "empty" second ubatches
65
+ orig_min_num_tokens = int(orig_num_tokens_tensor.min().item())
66
+ padded_max_num_tokens = int(padded_num_tokens_tensor.max().item())
67
+ if is_second_ubatch_empty(orig_min_num_tokens, padded_max_num_tokens):
68
+ logger.debug(
69
+ "Aborting ubatching %s %s", orig_min_num_tokens, padded_max_num_tokens
70
+ )
71
+ should_ubatch = False
72
+ return should_ubatch
73
+
74
+
75
+def _synchronize_dp_ranks(
76
+ num_tokens_unpadded: int,
77
+ num_tokens_padded: int,
78
+ should_attempt_ubatching: bool,
79
+ parallel_config: ParallelConfig,
80
+) -> tuple[bool, Optional[torch.Tensor]]:
81
+ """
82
+ 1. Decides if each DP rank is going to microbatch. Either all ranks
83
+ run with microbatching or none of them do.
84
+
85
+ 2. Determines the total number of tokens that each rank will run.
86
+ All ranks will be padded out so that the run with the same number
87
+ of tokens
88
+
89
+ Returns: tuple[
90
+ should_ubatch: Are all DP ranks going to microbatch
91
+ num_tokens_after_padding: A tensor containing the total number of
92
+ tokens per-microbatch for each DP rank including padding.
93
+ ]
94
+
95
+ """
96
+ assert num_tokens_padded >= num_tokens_unpadded
97
+
98
+ # First we coordinate between the DP ranks via an All Reduce
99
+ # to determine the total number of tokens that each rank
100
+ # will run and if we are using ubatching or not.
101
+ tensor = _run_ar(
102
+ should_ubatch=should_attempt_ubatching,
103
+ orig_num_tokens_per_ubatch=num_tokens_unpadded,
104
+ padded_num_tokens_per_ubatch=num_tokens_padded,
105
+ parallel_config=parallel_config,
106
+ )
107
+
108
+ # Ensure that each rank is processing the same nuber of tokens
109
+ num_tokens_across_dp = tensor[1, :]
110
+ max_num_tokens = int(num_tokens_across_dp.max().item())
111
+ num_tokens_after_padding = torch.tensor(
112
+ [max_num_tokens] * len(num_tokens_across_dp), device="cpu", dtype=torch.int32
113
+ )
114
+
115
+ should_ubatch = _post_process_ubatch(tensor)
116
+
117
+ return should_ubatch, num_tokens_after_padding
118
+
119
+
120
+def coordinate_batch_across_dp(
121
+ num_scheduled_tokens_per_request: np.ndarray,
122
+ num_tokens_unpadded: int,
123
+ num_tokens_padded: int,
124
+ parallel_config: ParallelConfig,
125
+ allow_microbatching: bool,
126
+ uniform_decode: bool,
127
+) -> tuple[Optional[UBatchSlices], Optional[torch.Tensor]]:
128
+ """
129
+ Coordinates amongst all DP ranks to determine if and how the full batch
130
+ should be split into microbatches.
131
+
132
+ Returns: tuple[
133
+ ubatch_slices: if this is set then all DP ranks have agreed to
134
+ microbatch
135
+ num_tokens_after_padding: A tensor containing the total number of
136
+ tokens per-microbatch for each DP rank including padding.
137
+ ]
138
+
139
+ """
140
+ if parallel_config.data_parallel_size == 1:
141
+ # Early exit.
142
+ return None, None
143
+
144
+ # Check preconditions for microbatching
145
+ should_attempt_ubatching = check_ubatch_thresholds(
146
+ parallel_config,
147
+ num_tokens_unpadded,
148
+ uniform_decode=uniform_decode,
149
+ )
150
+
151
+ # If the caller has explicitly disabled microbatching.
152
+ if not allow_microbatching:
153
+ should_attempt_ubatching = False
154
+
155
+ (should_ubatch, num_tokens_after_padding) = _synchronize_dp_ranks(
156
+ num_tokens_unpadded,
157
+ num_tokens_padded,
158
+ should_attempt_ubatching,
159
+ parallel_config,
160
+ )
161
+
162
+ # Don't microbatch unless every other DP worker is also microbatching
163
+ if not should_ubatch:
164
+ return (None, num_tokens_after_padding)
165
+
166
+ # This doesn't actually pad the ubatch slices. It just initializes the
167
+ # split point to the padded value so that padding can be applied
168
+ # to the second ubatch in pad_out_ubatch_slice after attention
169
+ # metadata creation
170
+ assert num_tokens_after_padding is not None
171
+ token_split_point = int(num_tokens_after_padding[0].item()) // 2
172
+
173
+ ubatch_slices = create_ubatch_slices(
174
+ num_scheduled_tokens_per_request, token_split_point
175
+ )
176
+
177
+ return (ubatch_slices, num_tokens_after_padding)
vllm/v1/worker/gpu_model_runner.py
+38
-123
@@ -41,7 +41,7 @@ from vllm.distributed.parallel_state import (
41
is_global_first_rank,
42
prepare_communication_buffer_for_model,
43
)
44
-from vllm.forward_context import BatchDescriptor, DPMetadata, set_forward_context
44
+from vllm.forward_context import BatchDescriptor, set_forward_context
45
from vllm.logger import init_logger
46
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
47
from vllm.model_executor.layers.mamba.abstract import MambaBase
@@ -131,12 +131,16 @@ from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
131
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
132
from vllm.v1.structured_output.utils import apply_grammar_bitmask
133
from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext
134
+from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
135
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
136
from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
137
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
138
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
138
-from vllm.v1.worker.ubatch_splitting import check_ubatch_thresholds, ubatch_split
139
-from vllm.v1.worker.ubatch_utils import UBatchSlice, UBatchSlices
139
+from vllm.v1.worker.ubatch_utils import (
140
+ UBatchSlice,
141
+ UBatchSlices,
142
+ check_ubatch_thresholds,
143
+)
144
from vllm.v1.worker.utils import is_residual_scattered_for_sp
145
146
from .utils import (
@@ -1161,18 +1165,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1165
query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]
1166
1167
num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
1164
- num_tokens_padded = num_tokens_unpadded + self.get_local_padding(
1165
- num_tokens_unpadded
1166
- )
1168
+ num_tokens_padded = self._get_num_input_tokens(num_tokens_unpadded)
1169
uniform_decode = (
1170
max_num_scheduled_tokens == self.uniform_decode_query_len
1171
) and (total_num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)
1170
- ubatch_slices, num_tokens_after_padding = ubatch_split(
1172
+ ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(
1173
num_scheduled_tokens,
1174
num_tokens_unpadded,
1175
num_tokens_padded,
1174
- uniform_decode=uniform_decode,
1175
- vllm_config=self.vllm_config,
1176
+ self.parallel_config,
1177
+ True,
1178
+ uniform_decode,
1179
)
1180
1181
self.seq_lens.np[:num_reqs] = (
@@ -1405,7 +1408,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1408
spec_decode_common_attn_metadata,
1409
max_num_scheduled_tokens,
1410
ubatch_slices,
1408
- num_tokens_after_padding,
1411
+ num_tokens_across_dp,
1412
use_cascade_attn,
1413
)
1414
@@ -1986,65 +1989,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1989
log_stats=self.parallel_config.eplb_config.log_balancedness,
1990
)
1991
1989
- def get_dp_padding(self, num_tokens: int) -> tuple[int, Optional[torch.Tensor]]:
1990
- """
1991
- Determines the total number of tokens that each rank will run.
1992
- All ranks will be padded out so that they run with the same number
1993
- of tokens
1994
-
1995
- Returns: tuple[
1996
- num_pad_tokens: The number of tokens that will be added to the batch
1997
- num_tokens_after_padding: A tensor containing the total number of
1998
- tokens for each DP rank including padding.
1999
- ]
2000
- """
2001
- dp_size = self.vllm_config.parallel_config.data_parallel_size
2002
- dp_rank = self.vllm_config.parallel_config.data_parallel_rank
2003
-
2004
- # For DP: Don't pad when setting enforce_eager.
2005
- # This lets us set enforce_eager on the prefiller in a P/D setup and
2006
- # still use CUDA graphs (enabled by this padding) on the decoder.
2007
- #
2008
- # TODO(tms) : There are many cases where padding is enabled for
2009
- # prefills, causing unnecessary and excessive padding of activations.
2010
-
2011
- if dp_size == 1 or self.vllm_config.model_config.enforce_eager:
2012
- # Early exit.
2013
- return 0, None
2014
-
2015
- num_tokens_across_dp = DPMetadata.num_tokens_across_dp(
2016
- num_tokens, dp_size, dp_rank
2017
- )
2018
- max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp).item()
2019
- num_tokens_after_padding = torch.tensor(
2020
- [max_tokens_across_dp_cpu] * dp_size, device="cpu", dtype=torch.int32
2021
- )
2022
- return max_tokens_across_dp_cpu - num_tokens, num_tokens_after_padding
2023
-
2024
- def get_local_padding(self, num_tokens_unpadded: int) -> int:
2025
- num_tokens_padded = num_tokens_unpadded
2026
-
2027
- if (
2028
- self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
2029
- and num_tokens_unpadded <= self.cudagraph_batch_sizes[-1]
2030
- ):
2031
- # Use piecewise CUDA graphs.
2032
- # Add padding to the batch size.
2033
- num_tokens_padded = self.vllm_config.pad_for_cudagraph(num_tokens_unpadded)
2034
- else:
2035
- # Eager mode.
2036
- # Pad tokens to multiple of tensor_parallel_size when
2037
- # enabled collective fusion for SP
2038
- tp_size = self.vllm_config.parallel_config.tensor_parallel_size
2039
- if (
2040
- self.vllm_config.compilation_config.pass_config.enable_sequence_parallelism
2041
- and tp_size > 1
2042
- ):
2043
- num_tokens_padded = round_up(num_tokens_unpadded, tp_size)
2044
-
2045
- num_pad_tokens = num_tokens_padded - num_tokens_unpadded
2046
- return num_pad_tokens
2047
-
1992
# This is where the second ubatch is adjusted to account for the padding.
1993
# Should be called after attention metadata creation. This just pads
1994
# the second ubatch slice out to the total number of tokens
@@ -2127,13 +2071,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2071
def _preprocess(
2072
self,
2073
scheduler_output: "SchedulerOutput",
2074
+ num_input_tokens: int, # Padded
2075
intermediate_tensors: Optional[IntermediateTensors] = None,
2131
- ubatch_slices: Optional[UBatchSlices] = None,
2132
- num_tokens_after_padding: Optional[torch.Tensor] = None,
2076
) -> tuple[
2077
int,
2135
- int,
2136
- Optional[torch.Tensor],
2078
Optional[torch.Tensor],
2079
Optional[torch.Tensor],
2080
torch.Tensor,
@@ -2141,14 +2082,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2082
dict[str, Any],
2083
]:
2084
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2144
- if ubatch_slices:
2145
- assert num_tokens_after_padding is not None
2146
- num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
2147
- self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
2148
- elif ubatch_slices is None:
2149
- num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
2150
- num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
2151
- num_input_tokens += num_pad
2085
2086
# _prepare_inputs may reorder the batch, so we must gather multi
2087
# modal outputs after that to ensure the correct order
@@ -2235,8 +2168,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2168
2169
return (
2170
num_scheduled_tokens,
2238
- num_input_tokens,
2239
- num_tokens_after_padding,
2171
input_ids,
2172
inputs_embeds,
2173
positions,
@@ -2506,24 +2437,30 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2437
spec_decode_common_attn_metadata,
2438
max_query_len,
2439
ubatch_slices,
2509
- num_tokens_after_padding,
2440
+ num_tokens_across_dp,
2441
use_cascade_attn,
2442
) = self._prepare_inputs(scheduler_output)
2443
2444
+ if ubatch_slices:
2445
+ assert num_tokens_across_dp is not None
2446
+ num_input_tokens = int(num_tokens_across_dp[0].item())
2447
+ self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
2448
+ elif num_tokens_across_dp is not None:
2449
+ num_input_tokens = int(num_tokens_across_dp[0].item())
2450
+ else:
2451
+ num_input_tokens = self._get_num_input_tokens(
2452
+ scheduler_output.total_num_scheduled_tokens
2453
+ )
2454
+
2455
(
2456
num_scheduled_tokens,
2515
- num_input_tokens,
2516
- num_tokens_across_dp,
2457
input_ids,
2458
inputs_embeds,
2459
positions,
2460
intermediate_tensors,
2461
model_kwargs,
2462
) = self._preprocess(
2523
- scheduler_output,
2524
- intermediate_tensors,
2525
- ubatch_slices,
2526
- num_tokens_after_padding,
2463
+ scheduler_output, num_input_tokens, intermediate_tensors
2464
)
2465
2466
uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
@@ -2548,11 +2485,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2485
):
2486
cudagraph_runtime_mode = CUDAGraphMode.NONE
2487
2551
- # This is currently to get around the assert in the DPMetadata
2552
- # where it wants `num_tokens_across_dp` to align with `num_tokens`
2553
- if ubatch_slices is not None:
2554
- num_input_tokens = ubatch_slices[0].num_tokens
2555
-
2488
# Run the model.
2489
# Use persistent buffers for CUDA graphs.
2490
with (
@@ -3329,36 +3261,19 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3261
num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
3262
total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
3263
3332
- ubatch_slices = None
3333
- num_tokens_after_padding = None
3334
-
3264
# We currently only microbatch if the number of tokens is
3265
# over a certain threshold.
3337
- if self.parallel_config.enable_dbo and allow_microbatching:
3338
- ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split(
3339
- num_scheduled_tokens,
3340
- total_num_scheduled_tokens,
3341
- total_num_scheduled_tokens,
3342
- uniform_decode=uniform_decode,
3343
- vllm_config=self.vllm_config,
3344
- )
3345
- # Currently when DBO is enabled `ubatch_split` returns
3346
- # the num_tokens_after_padding for a single ubatch, but we have 2
3347
- # TODO(sage,lucas): this is cruft that should be addressed in the
3348
- # padding refactor.
3349
- if ubatch_num_tokens_after_padding is not None:
3350
- num_tokens_after_padding = ubatch_num_tokens_after_padding * 2
3351
-
3352
- # If we failed to microbatch, currently need to resynchronize
3353
- # TODO(lucas,sage): we should be able to avoid this second sync by
3354
- # refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
3355
- # a single `coordinate_batch_across_dp` function.
3356
- if num_tokens_after_padding is None:
3357
- num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
3358
- num_tokens_after_padding = num_tokens + num_pad
3359
- else:
3360
- num_tokens_across_dp = num_tokens_after_padding
3361
- num_tokens_after_padding = int(num_tokens_after_padding[0].item())
3266
+ ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(
3267
+ num_scheduled_tokens,
3268
+ total_num_scheduled_tokens,
3269
+ total_num_scheduled_tokens,
3270
+ self.vllm_config.parallel_config,
3271
+ allow_microbatching,
3272
+ uniform_decode,
3273
+ )
3274
+ num_tokens_after_padding = num_tokens
3275
+ if num_tokens_across_dp is not None:
3276
+ num_tokens_after_padding = int(num_tokens_across_dp[0])
3277
3278
attn_metadata: Optional[PerLayerAttnMetadata] = None
3279
vllm/v1/worker/gpu_ubatch_wrapper.py
+14
-1
@@ -13,6 +13,7 @@ from vllm.config import CUDAGraphMode, VllmConfig
13
from vllm.distributed import get_ep_group
14
from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
15
from vllm.forward_context import (
16
+ DPMetadata,
17
create_forward_context,
18
get_forward_context,
19
override_forward_context,
@@ -409,6 +410,18 @@ class UBatchWrapper:
410
411
# We shouldn't be here unless we are running with multiple DP ranks
412
assert dp_metadata is not None
413
+ num_tokens_per_ubatch = (
414
+ ubatch_slices[0].token_slice.stop - ubatch_slices[0].token_slice.start
415
+ )
416
+ dp_size = self.vllm_config.parallel_config.data_parallel_size
417
+ ubatch_num_tokens_across_dp = torch.tensor(
418
+ [num_tokens_per_ubatch] * dp_size, device="cpu", dtype=torch.int32
419
+ )
420
+ ubatch_dp_metadata = DPMetadata.make(
421
+ self.vllm_config.parallel_config,
422
+ num_tokens_per_ubatch,
423
+ ubatch_num_tokens_across_dp,
424
+ )
425
426
if (
427
num_tokens not in self.cudagraphs
@@ -422,7 +435,7 @@ class UBatchWrapper:
435
intermediate_tensors=intermediate_tensors,
436
inputs_embeds=inputs_embeds,
437
compute_stream=compute_stream,
425
- dp_metadata=dp_metadata,
438
+ dp_metadata=ubatch_dp_metadata,
439
batch_descriptor=batch_descriptor,
440
cudagraph_runtime_mode=CUDAGraphMode.NONE,
441
)
vllm/v1/worker/ubatch_splitting.py
deleted
-207
@@ -1,207 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import Optional
5
-
6
-import numpy as np
7
-import torch
8
-
9
-from vllm.config import ParallelConfig, VllmConfig
10
-from vllm.forward_context import DPMetadata
11
-from vllm.logger import init_logger
12
-from vllm.utils import round_up
13
-from vllm.v1.worker.ubatch_utils import (
14
- UBatchSlice,
15
- UBatchSlices,
16
- is_second_ubatch_empty,
17
-)
18
-
19
-logger = init_logger(__name__)
20
-
21
-
22
-def should_ubatch_with_num_tokens(
23
- should_ubatch: bool,
24
- orig_num_tokens_per_ubatch: int,
25
- padded_num_tokens_per_ubatch: int,
26
- vllm_config: VllmConfig,
27
-) -> tuple[bool, Optional[torch.Tensor]]:
28
- dp_size = vllm_config.parallel_config.data_parallel_size
29
- dp_rank = vllm_config.parallel_config.data_parallel_rank
30
- return DPMetadata.should_ubatch_across_dp(
31
- should_ubatch,
32
- orig_num_tokens_per_ubatch,
33
- padded_num_tokens_per_ubatch,
34
- dp_size,
35
- dp_rank,
36
- )
37
-
38
-
39
-def check_ubatch_thresholds(
40
- config: ParallelConfig, num_tokens: int, uniform_decode: bool
41
-) -> bool:
42
- if not config.enable_dbo:
43
- return False
44
- if uniform_decode:
45
- return num_tokens >= config.dbo_decode_token_threshold
46
- else:
47
- return num_tokens >= config.dbo_prefill_token_threshold
48
-
49
-
50
-def get_dp_padding_ubatch(
51
- num_tokens_unpadded: int,
52
- num_tokens_padded: int,
53
- should_attempt_ubatching: bool,
54
- vllm_config: VllmConfig,
55
-) -> tuple[bool, Optional[torch.Tensor]]:
56
- """
57
- 1. Decides if each DP rank is going to microbatch. Either all ranks
58
- run with microbatching or none of them do. If this function decides
59
- not to run with microbatching. It will "abort" meaning that no padding
60
- information will be returned to the caller. It will return (False, None)
61
-
62
- 2. Determines the total number of tokens that each rank will run.
63
- All ranks will be padded out so that the run with the same number
64
- of tokens
65
-
66
- Returns: tuple[
67
- should_ubatch: Are all DP ranks going to microbatch
68
- num_tokens_after_padding: A tensor containing the total number of
69
- tokens per-microbatch for each DP rank including padding. Will be
70
- None if should_ubatch if False
71
- ]
72
-
73
- """
74
- assert num_tokens_padded >= num_tokens_unpadded
75
- dp_size = vllm_config.parallel_config.data_parallel_size
76
- if dp_size == 1:
77
- # Early exit.
78
- return False, None
79
-
80
- # If this DP rank doesn't want to attempt microbatching
81
- if not should_attempt_ubatching:
82
- (should_ubatch, num_tokens_across_dp) = should_ubatch_with_num_tokens(
83
- False, 0, 0, vllm_config
84
- )
85
- assert should_ubatch is False
86
- assert num_tokens_across_dp is None
87
- return should_ubatch, num_tokens_across_dp
88
-
89
- # Round up to the next multiple of two for even divisibility
90
- num_tokens_padded = round_up(num_tokens_padded, 2)
91
- num_tokens_per_ubatch = num_tokens_padded // 2
92
- should_ubatch = True
93
-
94
- # Sanity Check that the existing padding isn't giving us an empty second
95
- # ubatch. Abort if so
96
- if is_second_ubatch_empty(num_tokens_unpadded, num_tokens_padded):
97
- logger.debug(
98
- "Empty second µbatch detected: unpadded tokens: %s, padded tokens: %s",
99
- num_tokens_unpadded,
100
- num_tokens_padded,
101
- )
102
- should_ubatch = False
103
-
104
- # Note that we compute the number of padded tokens per ubatch
105
- (should_ubatch, num_tokens_across_dp) = should_ubatch_with_num_tokens(
106
- should_ubatch, num_tokens_unpadded // 2, num_tokens_per_ubatch, vllm_config
107
- )
108
- if not should_ubatch:
109
- assert num_tokens_across_dp is None
110
- return should_ubatch, num_tokens_across_dp
111
-
112
- assert num_tokens_across_dp is not None
113
-
114
- max_tokens_across_dp_cpu = int(torch.max(num_tokens_across_dp).item())
115
- num_tokens_after_padding = torch.tensor(
116
- [max_tokens_across_dp_cpu] * dp_size, device="cpu", dtype=torch.int32
117
- )
118
- return should_ubatch, num_tokens_after_padding
119
-
120
-
121
-def create_ubatch_slices(
122
- num_scheduled_tokens: np.ndarray, split_point: int
123
-) -> UBatchSlices:
124
- # TODO(lucas): Refactor the gpu_model_runner.py so we can pass
125
- # in cu_num_tokens directly (i.e. query_start_loc)
126
- cu_num_tokens = np.zeros(len(num_scheduled_tokens) + 1, dtype=np.int32)
127
- np.cumsum(num_scheduled_tokens, dtype=np.int32, out=cu_num_tokens[1:])
128
-
129
- first_ubatch_token_slice = slice(0, split_point)
130
- second_ubatch_token_slice = slice(split_point, cu_num_tokens[-1])
131
-
132
- # Determine request slices using exclusive stop semantics
133
- # First ubatch includes requests whose tokens overlap [0, split_point)
134
- first_ubatch_req_stop = int(
135
- np.searchsorted(cu_num_tokens, split_point, side="left")
136
- )
137
- first_ubatch_req_slice = slice(0, first_ubatch_req_stop)
138
-
139
- # Second ubatch starts at the request that contains the split_point
140
- # or the request starting exactly at split_point (if on boundary)
141
- second_ubatch_req_start = int(
142
- np.searchsorted(cu_num_tokens, split_point, side="right") - 1
143
- )
144
- second_ubatch_req_slice = slice(second_ubatch_req_start, len(cu_num_tokens) - 1)
145
-
146
- return [
147
- UBatchSlice(first_ubatch_req_slice, first_ubatch_token_slice),
148
- UBatchSlice(second_ubatch_req_slice, second_ubatch_token_slice),
149
- ]
150
-
151
-
152
-def ubatch_split(
153
- num_scheduled_tokens_per_request: np.ndarray,
154
- num_tokens_unpadded: int,
155
- num_tokens_padded: int,
156
- uniform_decode: bool,
157
- vllm_config: VllmConfig,
158
-) -> tuple[Optional[UBatchSlices], Optional[torch.Tensor]]:
159
- """
160
- Coordinates amongst all DP ranks to determine if and how the full batch
161
- should be split into microbatches.
162
-
163
- Returns: tuple[
164
- ubatch_slices: if this is set then all DP ranks have agreed to
165
- microbatch
166
- num_tokens_after_padding: A tensor containing the total number of
167
- tokens per-microbatch for each DP rank including padding. Will be
168
- None if ubatch_slices is None
169
- ]
170
-
171
- """
172
- parallel_config = vllm_config.parallel_config
173
- # Don't bother with the should_ubatch handshaking unless microbatching
174
- # is enabled
175
- if not parallel_config.enable_dbo:
176
- return (None, None)
177
-
178
- # Check preconditions for microbatching
179
- should_attempt_ubatching = check_ubatch_thresholds(
180
- parallel_config,
181
- num_tokens_unpadded,
182
- uniform_decode=uniform_decode,
183
- )
184
-
185
- # Don't microbatch unless every other DP worker is also microbatching
186
- should_ubatch, num_tokens_after_padding = get_dp_padding_ubatch(
187
- num_tokens_unpadded,
188
- num_tokens_padded,
189
- should_attempt_ubatching,
190
- vllm_config,
191
- )
192
-
193
- if not should_ubatch:
194
- return (None, None)
195
-
196
- # This doesn't actually pad the ubatch slices. It just initializes the
197
- # split point to the padded value so that padding can be applied
198
- # to the second ubatch in pad_out_ubatch_slice after attention
199
- # metadata creation
200
- assert num_tokens_after_padding is not None
201
- token_split_point = int(num_tokens_after_padding[0].item())
202
-
203
- ubatch_slices = create_ubatch_slices(
204
- num_scheduled_tokens_per_request, token_split_point
205
- )
206
-
207
- return (ubatch_slices, num_tokens_after_padding)
vllm/v1/worker/ubatch_utils.py
+46
-3
@@ -2,8 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
5
+import numpy as np
6
from typing_extensions import TypeAlias
7
8
+from vllm.config import ParallelConfig
9
+
10
11
@dataclass
12
class UBatchSlice:
@@ -24,7 +27,47 @@ class UBatchSlice:
27
UBatchSlices: TypeAlias = list[UBatchSlice]
28
29
27
-def is_second_ubatch_empty(
28
- orig_num_tokens_per_ubatch: int, padded_num_tokens_per_ubatch: int
30
+def is_second_ubatch_empty(orig_num_tokens: int, padded_num_tokens: int) -> bool:
31
+ return (padded_num_tokens // 2) >= orig_num_tokens
32
+
33
+
34
+def check_ubatch_thresholds(
35
+ config: ParallelConfig, num_tokens: int, uniform_decode: bool
36
) -> bool:
30
- return padded_num_tokens_per_ubatch >= 2 * orig_num_tokens_per_ubatch
37
+ if not config.enable_dbo:
38
+ return False
39
+ if uniform_decode:
40
+ return num_tokens >= config.dbo_decode_token_threshold
41
+ else:
42
+ return num_tokens >= config.dbo_prefill_token_threshold
43
+
44
+
45
+def create_ubatch_slices(
46
+ num_scheduled_tokens: np.ndarray, split_point: int
47
+) -> UBatchSlices:
48
+ # TODO(lucas): Refactor the gpu_model_runner.py so we can pass
49
+ # in cu_num_tokens directly (i.e. query_start_loc)
50
+ cu_num_tokens = np.zeros(len(num_scheduled_tokens) + 1, dtype=np.int32)
51
+ np.cumsum(num_scheduled_tokens, dtype=np.int32, out=cu_num_tokens[1:])
52
+
53
+ first_ubatch_token_slice = slice(0, split_point)
54
+ second_ubatch_token_slice = slice(split_point, cu_num_tokens[-1])
55
+
56
+ # Determine request slices using exclusive stop semantics
57
+ # First ubatch includes requests whose tokens overlap [0, split_point)
58
+ first_ubatch_req_stop = int(
59
+ np.searchsorted(cu_num_tokens, split_point, side="left")
60
+ )
61
+ first_ubatch_req_slice = slice(0, first_ubatch_req_stop)
62
+
63
+ # Second ubatch starts at the request that contains the split_point
64
+ # or the request starting exactly at split_point (if on boundary)
65
+ second_ubatch_req_start = int(
66
+ np.searchsorted(cu_num_tokens, split_point, side="right") - 1
67
+ )
68
+ second_ubatch_req_slice = slice(second_ubatch_req_start, len(cu_num_tokens) - 1)
69
+
70
+ return [
71
+ UBatchSlice(first_ubatch_req_slice, first_ubatch_token_slice),
72
+ UBatchSlice(second_ubatch_req_slice, second_ubatch_token_slice),
73
+ ]