forward_context.py 15.8 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import time
5
from collections import defaultdict
6
from contextlib import contextmanager
7
from dataclasses import dataclass, field
8
from typing import Any, NamedTuple
9

10
11
import torch

12
import vllm.envs as envs
13
from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
14
from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
from vllm.v1.attention.backend import AttentionMetadata
17
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
18
from vllm.v1.worker.ubatch_utils import UBatchSlices
19
20
21
22
23

logger = init_logger(__name__)

track_batchsize: bool = envs.VLLM_LOG_BATCHSIZE_INTERVAL >= 0
last_logging_time: float = 0
24
forward_start_time: float = 0
25
batchsize_logging_interval: float = envs.VLLM_LOG_BATCHSIZE_INTERVAL
26
batchsize_forward_time: defaultdict = defaultdict(list)
27
28


29
30
31
32
33
34
class BatchDescriptor(NamedTuple):
    """
    Batch descriptor for cudagraph dispatching. We should keep the num of
    items as minimal as possible to properly and uniquely describe the padded
    batch for cudagraph.
    """
35

36
    num_tokens: int
37
    num_reqs: int | None = None
38
    """
39
40
41
42
43
44
    Number of requests in the batch. Can be None for PIECEWISE cudagraphs where
    the cudagraphs can handle any number of requests.
    """
    uniform: bool = False
    """
    True if all the requests in the batch have the same number of tokens.
45
    """
46
47
48
49
    has_lora: bool = False
    """
    Whether this batch has active LoRA adapters.
    """
50

51
    def relax_for_mixed_batch_cudagraphs(self) -> "BatchDescriptor":
52
        """
53
54
        Return a relaxed version of current batch descriptor that is still compatible
        with PIECEWISE cudagraphs (or mixed prefill-decode FA cudagraphs).
55
        """
56
        return BatchDescriptor(
57
            self.num_tokens, num_reqs=None, uniform=False, has_lora=self.has_lora
58
        )
59
60


61
62
63
64
65
66
def _compute_sp_num_tokens(
    num_tokens_across_dp_cpu: torch.Tensor, sequence_parallel_size: int
) -> list[int]:
    sp_tokens = (
        num_tokens_across_dp_cpu + sequence_parallel_size - 1
    ) // sequence_parallel_size
67
68
69
70
71

    sp_tokens = sp_tokens.repeat_interleave(sequence_parallel_size)
    return sp_tokens.tolist()


72
73
74
75
76
77
78
def _compute_chunked_local_num_tokens(
    num_tokens_across_dp_cpu: torch.Tensor,
    sequence_parallel_size: int,
    max_num_tokens: int,
    chunk_idx: int,
) -> list[int]:
    sp_tokens = _compute_sp_num_tokens(num_tokens_across_dp_cpu, sequence_parallel_size)
79
80
81
82
83
    sp_size = len(sp_tokens)

    local_size = [-1] * sp_size
    for i in range(sp_size):
        # Take into account sharding if MoE activation is sequence parallel.
84
        local_size[i] = min(max_num_tokens, sp_tokens[i] - (max_num_tokens * chunk_idx))
85
86
87
88
89
        if local_size[i] <= 0:
            local_size[i] = 1  # ensure lockstep even if done
    return local_size


90
91
@dataclass
class DPMetadata:
92
    max_tokens_across_dp_cpu: torch.Tensor
93
94
95
    num_tokens_across_dp_cpu: torch.Tensor

    # NOTE: local_sizes should only be set by the chunked_sizes context manager
96
    local_sizes: list[int] | None = None
97

98
    @staticmethod
99
    def make(
100
101
        parallel_config: ParallelConfig,
        num_tokens: int,
102
        num_tokens_across_dp_cpu: torch.Tensor,
103
    ) -> "DPMetadata":
104
        assert num_tokens_across_dp_cpu is not None
105
        assert parallel_config.data_parallel_size > 1
106
        assert parallel_config.is_moe_model is not False
107
        dp_rank = parallel_config.data_parallel_rank
108
        batchsize = num_tokens
109

110
111
        # If num_tokens_across_dp is None, it will be computed by all_reduce
        # Otherwise, num_tokens_across_dp[dp_rank] should be equal to batchsize
112
113
114
        assert num_tokens_across_dp_cpu[dp_rank] == batchsize, (
            f"{num_tokens_across_dp_cpu[dp_rank]} {batchsize}"
        )
115
116
        max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp_cpu)
        return DPMetadata(max_tokens_across_dp_cpu, num_tokens_across_dp_cpu)
117

118
    @contextmanager
119
120
121
    def chunked_sizes(
        self, sequence_parallel_size: int, max_chunk_size_per_rank: int, chunk_idx: int
    ):
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
        """
        Context manager to compute and temporarily set the per-rank local token
        sizes for a specific chunk during chunked forward execution.

        This is necessary to ensure each DP (data parallel) rank processes its
        designated portion of tokens in lockstep with others, even when the
        token counts are uneven or some ranks have completed their input early.

        For chunked execution, we break up the total tokens on each rank into
        multiple chunks (of at most `max_chunk_size_per_rank`), and for a given
        `chunk_idx`, this context manager sets `self.local_sizes` to the number
        of tokens to process in that chunk on each rank.

        `self.local_sizes` is only valid inside the context.

        Args:
138
139
140
141
            sequence_parallel_size: When Attn is TP and MoE layers are EP,
                                    we use SP between the layers to avoid
                                    redundant ops. We need this value to
                                    compute the chunked sizes.
142
            max_chunk_size_per_rank: The max number of tokens each rank is
143
144
145
146
                                     allowed to process in this chunk.
            chunk_idx: The index of the chunk to compute sizes for.
        """
        self.local_sizes = _compute_chunked_local_num_tokens(
147
148
149
150
151
            self.num_tokens_across_dp_cpu,
            sequence_parallel_size,
            max_chunk_size_per_rank,
            chunk_idx,
        )
152
153
154
155
156
157
158
159
        try:
            yield self.local_sizes
        finally:
            self.local_sizes = None

    @contextmanager
    def sp_local_sizes(self, sequence_parallel_size: int):
        """
160
        Context manager for setting self.local_sizes. Same as self.chunked_sizes
161
162
163
        but without any chunking.
        """
        self.local_sizes = _compute_sp_num_tokens(
164
165
            self.num_tokens_across_dp_cpu, sequence_parallel_size
        )
166
167
168
169
170
        try:
            yield self.local_sizes
        finally:
            self.local_sizes = None

171
    def get_chunk_sizes_across_dp_rank(self) -> list[int] | None:
172
        assert self.local_sizes is not None
173
174
        return self.local_sizes

175
176
177
178
179
180
181
182
183
184
185
    # Get the cumulative tokens across sequence parallel ranks.
    # In this case the input to the MoEs will be distributed w.r.t both
    # DP and TP rank.
    # When sp_size==1, this is just the cummulative num tokens across DP.
    def cu_tokens_across_sp(self, sp_size: int) -> torch.Tensor:
        num_tokens_across_sp_cpu = (
            self.num_tokens_across_dp_cpu - 1 + sp_size
        ) // sp_size
        num_tokens_across_sp_cpu = num_tokens_across_sp_cpu.repeat_interleave(sp_size)
        return torch.cumsum(num_tokens_across_sp_cpu, dim=0)

186

187
188
@dataclass
class ForwardContext:
189
    # copy from vllm_config.compilation_config.static_forward_context
190
    no_compile_layers: dict[str, Any]
191
    attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]]
192
    slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]]
193
194
195
    """
    Type Dict[str, AttentionMetadata] for v1, map from layer_name of each 
    attention layer to its attention metadata
196
197
198
    Type List[Dict[str, AttentionMetadata]] for DBO. List of size two, one
    for each microbatch.
    Set dynamically for each forward pass
199
    """
200
201
    # TODO: remove after making all virtual_engines share the same kv cache
    virtual_engine: int  # set dynamically for each forward pass
202
    # set dynamically for each forward pass
203
    dp_metadata: DPMetadata | None = None
204
205
206
    # determine the cudagraph style at runtime to be FULL, PIECEWISE, or NONE.
    # by default NONE, no cudagraph is used.
    cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE
207
    batch_descriptor: BatchDescriptor | None = None
208

209
    ubatch_slices: UBatchSlices | None = None
210

211
212
213
    # If True, bypass the compiled model call, e.g. by using .forward() directly
    skip_compiled: bool = False

214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
    # For torch.compile cold start times, we need to avoid hard-coding
    # any strings into the graph. Right now, the vllm.moe_forward
    # and vllm.moe_forward_shared custom operators hard-code strings into
    # the graph.
    #
    # The workaround is to store a list of the strings that each of those
    # custom ops needs, in reverse order, in the ForwardContext.
    # The ForwardContext object is alive for the duration of the forward pass.
    # When the custom op needs the string, pop the string from this list.
    #
    # This assumes that the custom operators will always be executed in
    # order and that torch.compile will not try to reorder these
    # operations with respect to each other.
    #
    # TODO(https://github.com/vllm-project/vllm/issues/31985):
    # There are longer-term solutions, like unwrapping the moe custom operator,
    # that aren't ready yet.
    # We could also treat the string as a "symbolic input" to the graph but
    # the PyTorch-side bits for that aren't ready yet either.
    #
    # If this value is None (like in some tests), then we end up baking the string
    # into the graph. Otherwise, the moe custom ops will pop a string from this list.
    remaining_moe_layers: list[str] | None = None

238
239
    additional_kwargs: dict[str, Any] = field(default_factory=dict)

240
    def __post_init__(self):
241
        assert self.cudagraph_runtime_mode.valid_runtime_modes(), (
242
            f"Invalid cudagraph runtime mode: {self.cudagraph_runtime_mode}"
243
        )
244
245


246
_forward_context: ForwardContext | None = None
247
248
249


def get_forward_context() -> ForwardContext:
250
    """Get the current forward context."""
251
252
    assert _forward_context is not None, (
        "Forward context is not set. "
253
254
        "Please use `set_forward_context` to set the forward context."
    )
255
256
257
    return _forward_context


258
259
260
261
def is_forward_context_available() -> bool:
    return _forward_context is not None


262
def create_forward_context(
263
264
265
    attn_metadata: Any,
    vllm_config: VllmConfig,
    virtual_engine: int = 0,
266
    dp_metadata: DPMetadata | None = None,
267
    cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
268
269
    batch_descriptor: BatchDescriptor | None = None,
    ubatch_slices: UBatchSlices | None = None,
270
    slot_mapping: dict[str, torch.Tensor] | None = None,
271
    additional_kwargs: dict[str, Any] | None = None,
272
    skip_compiled: bool = False,
273
):
274
275
276
277
278
279
280
281
    no_compile_layers = vllm_config.compilation_config.static_forward_context
    from vllm.model_executor.layers.fused_moe.layer import FusedMoE

    remaining_moe_layers = [
        name for name, layer in no_compile_layers.items() if isinstance(layer, FusedMoE)
    ]
    remaining_moe_layers.reverse()

282
    return ForwardContext(
283
284
        no_compile_layers=no_compile_layers,
        remaining_moe_layers=remaining_moe_layers,
285
286
        virtual_engine=virtual_engine,
        attn_metadata=attn_metadata,
287
        slot_mapping=slot_mapping or {},
288
289
290
291
        dp_metadata=dp_metadata,
        cudagraph_runtime_mode=cudagraph_runtime_mode,
        batch_descriptor=batch_descriptor,
        ubatch_slices=ubatch_slices,
292
        skip_compiled=skip_compiled,
293
        additional_kwargs=additional_kwargs or {},
294
    )
295
296
297


@contextmanager
298
def override_forward_context(forward_context: ForwardContext | None):
299
300
301
302
303
304
305
306
307
308
309
310
311
    """A context manager that overrides the current forward context.
    This is used to override the forward context for a specific
    forward pass.
    """
    global _forward_context
    prev_context = _forward_context
    _forward_context = forward_context
    try:
        yield
    finally:
        _forward_context = prev_context


312
@contextmanager
313
def set_forward_context(
314
315
316
    attn_metadata: Any,
    vllm_config: VllmConfig,
    virtual_engine: int = 0,
317
318
    num_tokens: int | None = None,
    num_tokens_across_dp: torch.Tensor | None = None,
319
    cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
320
321
    batch_descriptor: BatchDescriptor | None = None,
    ubatch_slices: UBatchSlices | None = None,
322
    slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None,
323
    skip_compiled: bool = False,
324
):
325
    """A context manager that stores the current forward context,
326
327
328
    can be attention metadata, etc.
    Here we can inject common logic for every model forward pass.
    """
329
    global forward_start_time
330
    need_to_track_batchsize = track_batchsize and attn_metadata is not None
331
332
    if need_to_track_batchsize:
        forward_start_time = time.perf_counter()
333

334
    dp_metadata: DPMetadata | None = None
335
    if vllm_config.parallel_config.data_parallel_size > 1 and (
336
337
        attn_metadata is not None or num_tokens is not None
    ):
338
339
340
341
342
343
        # If num_tokens_across_dp hasn't already been initialized, then
        # initialize it here. Both DP padding and Microbatching will be
        # disabled.
        if num_tokens_across_dp is None:
            assert ubatch_slices is None
            assert num_tokens is not None
344
            _, num_tokens_across_dp, _ = coordinate_batch_across_dp(
345
346
347
348
349
350
                num_tokens_unpadded=num_tokens,
                parallel_config=vllm_config.parallel_config,
                allow_microbatching=False,
                allow_dp_padding=False,
            )
            assert num_tokens_across_dp is not None
351
        dp_metadata = DPMetadata.make(
352
            vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp
353
354
        )

355
356
357
358
359
360
    # Convenience: if cudagraph is used and num_tokens is given, we can just
    # create a batch descriptor here if not given (there's no harm since if it
    # doesn't match in the wrapper it'll fall through).
    if cudagraph_runtime_mode != CUDAGraphMode.NONE and num_tokens is not None:
        batch_descriptor = batch_descriptor or BatchDescriptor(num_tokens=num_tokens)

361
362
363
364
    additional_kwargs = current_platform.set_additional_forward_context(
        attn_metadata=attn_metadata,
        vllm_config=vllm_config,
        virtual_engine=virtual_engine,
365
        dp_metadata=dp_metadata,
366
367
368
369
370
371
372
        num_tokens=num_tokens,
        num_tokens_across_dp=num_tokens_across_dp,
        cudagraph_runtime_mode=cudagraph_runtime_mode,
        batch_descriptor=batch_descriptor,
        ubatch_slices=ubatch_slices,
    )

373
374
375
376
377
378
379
380
    forward_context = create_forward_context(
        attn_metadata,
        vllm_config,
        virtual_engine,
        dp_metadata,
        cudagraph_runtime_mode,
        batch_descriptor,
        ubatch_slices,
381
        slot_mapping,
382
        additional_kwargs,
383
        skip_compiled,
384
    )
385

386
    try:
387
388
        with override_forward_context(forward_context):
            yield
389
    finally:
390
391
        global last_logging_time, batchsize_logging_interval
        if need_to_track_batchsize:
392
            batchsize = num_tokens
393
394
395
            # we use synchronous scheduling right now,
            # adding a sync point here should not affect
            # scheduling of the next batch
396
397
398
            synchronize = current_platform.synchronize
            if synchronize is not None:
                synchronize()
399
400
            now = time.perf_counter()
            # time measurement is in milliseconds
401
            batchsize_forward_time[batchsize].append((now - forward_start_time) * 1000)
402
403
404
405
406
407
408
409
410
411
412
413
            if now - last_logging_time > batchsize_logging_interval:
                last_logging_time = now
                forward_stats = []
                for bs, times in batchsize_forward_time.items():
                    if len(times) <= 1:
                        # can be cudagraph / profiling run
                        continue
                    medium = torch.quantile(torch.tensor(times), q=0.5).item()
                    medium = round(medium, 2)
                    forward_stats.append((bs, len(times), medium))
                forward_stats.sort(key=lambda x: x[1], reverse=True)
                if forward_stats:
414
415
416
417
418
419
420
                    logger.info(
                        (
                            "Batchsize forward time stats "
                            "(batchsize, count, median_time(ms)): %s"
                        ),
                        forward_stats,
                    )