forward_context.py 16.5 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


zhuwenwen's avatar
zhuwenwen committed
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

zhuwenwen's avatar
zhuwenwen committed
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
    has_lora: bool = False
zhuwenwen's avatar
zhuwenwen committed
47
    """
48
    Whether this batch has active LoRA adapters.
zhuwenwen's avatar
zhuwenwen committed
49
    """
luopl's avatar
luopl committed
50
51
52
53
54
55
56
57
    num_active_loras: int = 0
    """
    Number of distinct active LoRA adapters in this batch.
    When cudagraph_specialize_lora_count is enabled, separate CUDA graphs
    are captured for each num_active_loras value. This allows kernels
    (like fused_moe_lora) whose grid size depends on num_active_loras
    to be properly captured.
    """
zhuwenwen's avatar
zhuwenwen committed
58

59
    def relax_for_mixed_batch_cudagraphs(self) -> "BatchDescriptor":
zhuwenwen's avatar
zhuwenwen committed
60
        """
61
62
        Return a relaxed version of current batch descriptor that is still compatible
        with PIECEWISE cudagraphs (or mixed prefill-decode FA cudagraphs).
zhuwenwen's avatar
zhuwenwen committed
63
        """
64
        return BatchDescriptor(
65
            self.num_tokens, num_reqs=None, uniform=False, has_lora=self.has_lora
66
        )
zhuwenwen's avatar
zhuwenwen committed
67
68


69
70
71
72
73
74
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
75
76
77
78
79

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


80
81
82
83
84
85
86
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)
87
88
89
90
91
    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.
92
        local_size[i] = min(max_num_tokens, sp_tokens[i] - (max_num_tokens * chunk_idx))
zhuwenwen's avatar
zhuwenwen committed
93
94
95
96
97
        if local_size[i] <= 0:
            local_size[i] = 1  # ensure lockstep even if done
    return local_size


98
99
@dataclass
class DPMetadata:
100
    max_tokens_across_dp_cpu: torch.Tensor
101
102
103
    num_tokens_across_dp_cpu: torch.Tensor

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

    @staticmethod
107
    def make(
108
109
        parallel_config: ParallelConfig,
        num_tokens: int,
110
        num_tokens_across_dp_cpu: torch.Tensor,
111
    ) -> "DPMetadata":
112
        assert num_tokens_across_dp_cpu is not None
113
        assert parallel_config.data_parallel_size > 1
114
        assert parallel_config.is_moe_model is not False
115
        dp_rank = parallel_config.data_parallel_rank
116
        batchsize = num_tokens
117

118
119
        # 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
120
121
122
        assert num_tokens_across_dp_cpu[dp_rank] == batchsize, (
            f"{num_tokens_across_dp_cpu[dp_rank]} {batchsize}"
        )
123
124
        max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp_cpu)
        return DPMetadata(max_tokens_across_dp_cpu, num_tokens_across_dp_cpu)
125

126
    @contextmanager
127
128
129
    def chunked_sizes(
        self, sequence_parallel_size: int, max_chunk_size_per_rank: int, chunk_idx: int
    ):
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
        """
        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:
146
147
148
149
            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.
150
            max_chunk_size_per_rank: The max number of tokens each rank is
151
152
153
154
                                     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(
155
156
157
158
159
            self.num_tokens_across_dp_cpu,
            sequence_parallel_size,
            max_chunk_size_per_rank,
            chunk_idx,
        )
160
161
162
163
164
165
166
167
        try:
            yield self.local_sizes
        finally:
            self.local_sizes = None

    @contextmanager
    def sp_local_sizes(self, sequence_parallel_size: int):
        """
168
        Context manager for setting self.local_sizes. Same as self.chunked_sizes
169
170
171
        but without any chunking.
        """
        self.local_sizes = _compute_sp_num_tokens(
172
173
            self.num_tokens_across_dp_cpu, sequence_parallel_size
        )
174
175
176
177
178
        try:
            yield self.local_sizes
        finally:
            self.local_sizes = None

179
    def get_chunk_sizes_across_dp_rank(self) -> list[int] | None:
180
        assert self.local_sizes is not None
181
182
        return self.local_sizes

183
184
185
186
187
188
189
190
191
192
193
    # 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)

194

195
196
@dataclass
class ForwardContext:
197
    # copy from vllm_config.compilation_config.static_forward_context
198
    no_compile_layers: dict[str, Any]
199
    attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]]
200
    slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]]
201
    """
luopl's avatar
luopl committed
202
    Type Dict[str, AttentionMetadata] for v1, map from layer_name of each
203
    attention layer to its attention metadata
204
205
206
    Type List[Dict[str, AttentionMetadata]] for DBO. List of size two, one
    for each microbatch.
    Set dynamically for each forward pass
207
    """
208
209
    # TODO: remove after making all virtual_engines share the same kv cache
    virtual_engine: int  # set dynamically for each forward pass
210
    # set dynamically for each forward pass
211
    dp_metadata: DPMetadata | None = None
212
213
214
    # 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
215
    batch_descriptor: BatchDescriptor | None = None
216

217
    ubatch_slices: UBatchSlices | None = None
218

219
220
221
    # If True, bypass the compiled model call, e.g. by using .forward() directly
    skip_compiled: bool = False

222
223
224
225
226
227
    # 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
228
229
    # custom ops needs in the ForwardContext (all_moe_layers)
    # as well as a counter (moe_layer_index).
230
    # The ForwardContext object is alive for the duration of the forward pass.
231
232
    # When the custom op needs a layer string, get the next string
    # from all_moe_layers and increment the counter.
233
234
235
236
237
238
239
240
241
242
243
244
245
    #
    # 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.
246
247
    all_moe_layers: list[str] | None = None
    moe_layer_index: int = 0
248

249
250
    additional_kwargs: dict[str, Any] = field(default_factory=dict)

251
    def __post_init__(self):
252
        assert self.cudagraph_runtime_mode.valid_runtime_modes(), (
253
            f"Invalid cudagraph runtime mode: {self.cudagraph_runtime_mode}"
254
        )
255
256


257
_forward_context: ForwardContext | None = None
258
259
260


def get_forward_context() -> ForwardContext:
261
    """Get the current forward context."""
262
263
    assert _forward_context is not None, (
        "Forward context is not set. "
264
265
        "Please use `set_forward_context` to set the forward context."
    )
266
267
268
    return _forward_context


269
270
271
272
def is_forward_context_available() -> bool:
    return _forward_context is not None


273
def create_forward_context(
274
275
276
    attn_metadata: Any,
    vllm_config: VllmConfig,
    virtual_engine: int = 0,
277
    dp_metadata: DPMetadata | None = None,
278
    cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
279
280
    batch_descriptor: BatchDescriptor | None = None,
    ubatch_slices: UBatchSlices | None = None,
281
    slot_mapping: dict[str, torch.Tensor] | None = None,
282
    additional_kwargs: dict[str, Any] | None = None,
283
    skip_compiled: bool = False,
284
):
285
286
287
288
289
290
291
292
293
294
295
296
    if vllm_config.compilation_config.fast_moe_cold_start:
        if vllm_config.speculative_config is None:
            all_moe_layers = vllm_config.compilation_config.static_all_moe_layers
        else:
            logger.warning_once(
                "vllm_config.compilation_config.fast_moe_cold_start is not "
                "compatible with speculative decoding so we are ignoring "
                "fast_moe_cold_start."
            )
            all_moe_layers = None
    else:
        all_moe_layers = None
297

298
    return ForwardContext(
299
        no_compile_layers=vllm_config.compilation_config.static_forward_context,
300
        all_moe_layers=all_moe_layers,
301
302
        virtual_engine=virtual_engine,
        attn_metadata=attn_metadata,
303
        slot_mapping=slot_mapping or {},
304
305
306
307
        dp_metadata=dp_metadata,
        cudagraph_runtime_mode=cudagraph_runtime_mode,
        batch_descriptor=batch_descriptor,
        ubatch_slices=ubatch_slices,
308
        skip_compiled=skip_compiled,
309
        additional_kwargs=additional_kwargs or {},
310
    )
311
312
313


@contextmanager
314
def override_forward_context(forward_context: ForwardContext | None):
315
316
317
318
319
320
321
322
323
324
325
326
327
    """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


328
@contextmanager
329
def set_forward_context(
330
331
332
    attn_metadata: Any,
    vllm_config: VllmConfig,
    virtual_engine: int = 0,
333
334
    num_tokens: int | None = None,
    num_tokens_across_dp: torch.Tensor | None = None,
335
    cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
336
337
    batch_descriptor: BatchDescriptor | None = None,
    ubatch_slices: UBatchSlices | None = None,
338
    slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None,
339
    skip_compiled: bool = False,
340
):
341
    """A context manager that stores the current forward context,
342
343
344
    can be attention metadata, etc.
    Here we can inject common logic for every model forward pass.
    """
345
    global forward_start_time
346
    need_to_track_batchsize = track_batchsize and attn_metadata is not None
347
348
    if need_to_track_batchsize:
        forward_start_time = time.perf_counter()
349

350
    dp_metadata: DPMetadata | None = None
351
    if vllm_config.parallel_config.data_parallel_size > 1 and (
352
353
        attn_metadata is not None or num_tokens is not None
    ):
354
355
356
357
358
359
        # 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
360
            _, num_tokens_across_dp, _ = coordinate_batch_across_dp(
361
362
363
364
365
366
                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
367
        dp_metadata = DPMetadata.make(
368
            vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp
369
370
        )

371
372
373
374
375
376
    # 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)

377
378
379
380
    additional_kwargs = current_platform.set_additional_forward_context(
        attn_metadata=attn_metadata,
        vllm_config=vllm_config,
        virtual_engine=virtual_engine,
381
        dp_metadata=dp_metadata,
382
383
384
385
386
387
388
        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,
    )

389
390
391
392
393
394
395
396
    forward_context = create_forward_context(
        attn_metadata,
        vllm_config,
        virtual_engine,
        dp_metadata,
        cudagraph_runtime_mode,
        batch_descriptor,
        ubatch_slices,
397
        slot_mapping,
398
        additional_kwargs,
399
        skip_compiled,
400
    )
401

402
    try:
403
404
        with override_forward_context(forward_context):
            yield
405
    finally:
406
407
        global last_logging_time, batchsize_logging_interval
        if need_to_track_batchsize:
408
            batchsize = num_tokens
409
410
411
            # we use synchronous scheduling right now,
            # adding a sync point here should not affect
            # scheduling of the next batch
412
413
414
            synchronize = current_platform.synchronize
            if synchronize is not None:
                synchronize()
415
416
            now = time.perf_counter()
            # time measurement is in milliseconds
417
            batchsize_forward_time[batchsize].append((now - forward_start_time) * 1000)
418
419
420
421
422
423
424
425
426
427
428
429
            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:
430
431
432
433
434
435
436
                    logger.info(
                        (
                            "Batchsize forward time stats "
                            "(batchsize, count, median_time(ms)): %s"
                        ),
                        forward_stats,
                    )