model_runner.py 39.2 KB
Newer Older
1
import contextlib
2
import time
3
from typing import Dict, List, Optional, Tuple, Set
4

5
import numpy as np
6
import torch
7
import torch.nn as nn
8

9
from vllm.attention import AttentionMetadata, get_attn_backend
Woosuk Kwon's avatar
Woosuk Kwon committed
10
11
from vllm.config import (DeviceConfig, ModelConfig, LoRAConfig, ParallelConfig,
                         SchedulerConfig)
12
from vllm.logger import init_logger
13
from vllm.model_executor import SamplingMetadata
14
from vllm.model_executor.model_loader import get_model
15
from vllm.model_executor.parallel_utils import cupy_utils
16
from vllm.model_executor.parallel_utils.communication_op import (
17
    broadcast_tensor_dict)
Woosuk Kwon's avatar
Woosuk Kwon committed
18
19
from vllm.model_executor.parallel_utils.parallel_state import (
    with_cupy_nccl_for_all_reduce)
20
from vllm.model_executor.parallel_utils import custom_all_reduce
21
22
from vllm.sampling_params import SamplingParams, SamplingType
from vllm.sequence import SamplerOutput, SequenceData, SequenceGroupMetadata
23
24
25
from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
from vllm.lora.layers import LoRAMapping
from vllm.lora.request import LoRARequest
26
27
28
from vllm.utils import (async_tensor_h2d, CudaMemoryProfiler,
                        is_pin_memory_available, make_tensor_with_pad,
                        maybe_expand_dim)
29
30
31
32

logger = init_logger(__name__)

_PAD_SLOT_ID = -1
33
LORA_WARMUP_RANK = 8
34
35
_BATCH_SIZE_ALIGNMENT = 8
# Capture graphs for token size 1, 2, 4, 8, 16, 24, 32, 40, ..., 256.
36
# NOTE: _get_graph_batch_size needs to be updated if this list is changed.
37
38
39
_BATCH_SIZES_TO_CAPTURE = [1, 2, 4] + [
    _BATCH_SIZE_ALIGNMENT * i for i in range(1, 33)
]
40
41
42
43
44
45
46
47
48


class ModelRunner:

    def __init__(
        self,
        model_config: ModelConfig,
        parallel_config: ParallelConfig,
        scheduler_config: SchedulerConfig,
49
        device_config: DeviceConfig,
50
        lora_config: Optional[LoRAConfig],
51
        kv_cache_dtype: Optional[str] = "auto",
52
        is_driver_worker: bool = False,
53
54
55
56
    ):
        self.model_config = model_config
        self.parallel_config = parallel_config
        self.scheduler_config = scheduler_config
57
        self.lora_config = lora_config
58
        self.is_driver_worker = is_driver_worker
59

Woosuk Kwon's avatar
Woosuk Kwon committed
60
61
62
63
        # model_config can be None in tests/samplers/test_sampler.py.
        # FIXME(woosuk): This is a hack to make the tests work. Refactor this.
        self.sliding_window = (model_config.get_sliding_window()
                               if model_config is not None else None)
64
65
66
67
        self.device_config = (device_config
                              if device_config is not None else DeviceConfig())
        self.device = self.device_config.device

68
69
        self.model = None
        self.block_size = None  # Set after initial profiling.
70
        self.lora_manager = None
71

72
73
74
75
76
77
78
79
80
81
82
83
84
        self.graph_runners: Dict[int, CUDAGraphRunner] = {}
        self.graph_memory_pool = None  # Set during graph capture.

        self.max_context_len_to_capture = (
            self.model_config.max_context_len_to_capture
            if self.model_config is not None else 0)
        # When using CUDA graph, the input block tables must be padded to
        # max_context_len_to_capture. However, creating the block table in
        # Python can be expensive. To optimize this, we cache the block table
        # in numpy and only copy the actual input content at every iteration.
        # The shape of the cached block table will be
        # (max batch size to capture, max context len to capture / block size).
        self.graph_block_tables = None  # Set after initial profiling.
85
        self.pin_memory = is_pin_memory_available()
86
        self.kv_cache_dtype = kv_cache_dtype
87

88
89
90
        self.attn_backend = get_attn_backend(
            self.model_config.dtype if model_config is not None else None)

91
    def load_model(self) -> None:
92
        with CudaMemoryProfiler() as m:
93
94
95
96
97
98
99
            self.model = get_model(self.model_config,
                                   self.device_config,
                                   lora_config=self.lora_config,
                                   parallel_config=self.parallel_config,
                                   scheduler_config=self.scheduler_config)

        self.model_memory_usage = m.consumed_memory
100
101
        logger.info(f"Loading model weights took "
                    f"{self.model_memory_usage / float(2**30):.4f} GB")
102
103

        if self.lora_config:
104
105
106
            assert hasattr(self.model, "supported_lora_modules"
                           ) and self.model.supported_lora_modules, (
                               "Model does not support LoRA")
Terry's avatar
Terry committed
107
108
109
110
111
            assert hasattr(
                self.model,
                "embedding_modules"), "Model does not have embedding_modules"
            assert hasattr(self.model, "embedding_padding_modules"
                           ), "Model does not have embedding_padding_modules"
112
113
            self.lora_manager = LRUCacheWorkerLoRAManager(
                self.scheduler_config.max_num_seqs,
114
                self.scheduler_config.max_num_batched_tokens, self.vocab_size,
Terry's avatar
Terry committed
115
116
                self.lora_config, self.device, self.model.embedding_modules,
                self.model.embedding_padding_modules)
117
            self.model = self.lora_manager.create_lora_manager(self.model)
118
119
120
121

    def set_block_size(self, block_size: int) -> None:
        self.block_size = block_size

122
        self.graph_block_tables = np.zeros(
123
124
125
126
127
128
            (max(_BATCH_SIZES_TO_CAPTURE), self.get_max_block_per_batch()),
            dtype=np.int32)

    def get_max_block_per_batch(self) -> int:
        block_size = self.block_size
        return (self.max_context_len_to_capture + block_size - 1) // block_size
129

130
131
132
    def _prepare_prompt(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
133
134
    ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, List[int],
               List[int], List[int], List[int], Set[LoRARequest]]:
135
        assert len(seq_group_metadata_list) > 0
136
137
138
        input_tokens: List[int] = []
        input_positions: List[int] = []
        slot_mapping: List[int] = []
139
140
141
        lora_index_mapping: List[int] = []
        lora_prompt_mapping: List[int] = []
        lora_requests: Set[LoRARequest] = set()
142
143

        prompt_lens: List[int] = []
144
145
146
        context_lens: List[int] = []
        subquery_lens: List[int] = []
        prefix_block_tables: List[List[int]] = []
147
148
149
150
151
152
153
154
155
156
        for seq_group_metadata in seq_group_metadata_list:
            assert seq_group_metadata.is_prompt
            seq_ids = list(seq_group_metadata.seq_data.keys())
            assert len(seq_ids) == 1
            seq_id = seq_ids[0]

            seq_data = seq_group_metadata.seq_data[seq_id]
            prompt_tokens = seq_data.get_token_ids()
            prompt_len = len(prompt_tokens)
            prompt_lens.append(prompt_len)
157
158
159
160
161
162
163
164
165
166
            computed_len = 0

            # NOTE: This only works for oooooooxxx style attention.
            computed_block_nums = seq_group_metadata.computed_block_nums
            if computed_block_nums is not None and len(
                    computed_block_nums) > 0 and self.sliding_window is None:
                # Prefix is not supported with sliding_window
                computed_len = len(computed_block_nums) * self.block_size
                prompt_tokens = prompt_tokens[computed_len:]
                prefix_block_tables.append(computed_block_nums)
167
                context_len = computed_len
168
169
            else:
                prefix_block_tables.append([])
170
                context_len = 0
171
            # actual prompt lens
172
            context_lens.append(context_len)
173
            subquery_lens.append(prompt_len - computed_len)
174

175
            input_tokens.extend(prompt_tokens)
176
177
            # NOTE(woosuk): Here we assume that the first token in the prompt
            # is always the first token in the sequence.
178
            input_positions.extend(
179
                list(range(computed_len, computed_len + len(prompt_tokens))))
180

181
182
183
184
185
            lora_id = seq_group_metadata.lora_int_id

            if lora_id > 0:
                lora_requests.add(seq_group_metadata.lora_request)

186
            lora_index_mapping += [lora_id] * (prompt_len - computed_len)
187
188
            lora_prompt_mapping.extend(
                [lora_id] *
189
                (prompt_len - computed_len
190
191
                 if seq_group_metadata.sampling_params.prompt_logprobs else 1))

192
193
194
            if seq_group_metadata.block_tables is None:
                # During memory profiling, the block tables are not initialized
                # yet. In this case, we just use a dummy slot mapping.
195
                slot_mapping.extend([_PAD_SLOT_ID] * prompt_len)
196
197
198
199
200
201
202
203
204
205
206
                continue

            # Compute the slot mapping.
            block_table = seq_group_metadata.block_tables[seq_id]
            # Mask the [0, start_idx) tokens of the prompt with _PAD_SLOT_ID,
            # where start_idx is max(0, prompt_len - sliding_window).
            # For example, if the prompt len is 10, sliding window is 8, and
            # block size is 4, the first two tokens are masked and the slot
            # mapping will be [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].
            start_idx = 0
            if self.sliding_window is not None:
207
                assert computed_len == 0, (
208
209
                    "Prefix caching is currently not supported with "
                    "sliding window attention")
210
                start_idx = max(0, prompt_len - self.sliding_window)
211
            for i in range(computed_len, prompt_len):
212
                if i < start_idx:
213
                    slot_mapping.append(_PAD_SLOT_ID)
214
215
216
217
218
                    continue

                block_number = block_table[i // self.block_size]
                block_offset = i % self.block_size
                slot = block_number * self.block_size + block_offset
219
220
221
                slot_mapping.append(slot)

        max_subquery_len = max(subquery_lens)
222
        max_prompt_len = max(prompt_lens)
223
224
225
226
227
228
229
230
231
232
233
234
235
236
        num_prompt_tokens = len(input_tokens)
        assert max_subquery_len > 0

        input_tokens = torch.tensor(input_tokens,
                                    dtype=torch.long,
                                    device=self.device)
        input_positions = torch.tensor(input_positions,
                                       dtype=torch.long,
                                       device=self.device)
        slot_mapping = torch.tensor(slot_mapping,
                                    dtype=torch.long,
                                    device=self.device)
        lora_index_mapping = lora_index_mapping

237
238
        context_lens_tensor = torch.tensor(context_lens,
                                           dtype=torch.int,
239
                                           device=self.device)
240
241
        # Prepare prefix block tables
        max_prompt_block_table_len = max(len(t) for t in prefix_block_tables)
242
        block_tables = make_tensor_with_pad(
243
244
245
246
            prefix_block_tables,
            max_len=max_prompt_block_table_len,
            pad=0,
            dtype=torch.int,
247
            device=self.device,
248
        )
249
250
251
252
253
254
255
256
257
258

        # Query length can be shorter than key (i.e., prompt) when prefill
        # is chunked or prefix cached.
        subquery_lens_tensor = torch.tensor(subquery_lens,
                                            dtype=torch.long,
                                            device=self.device)
        subquery_start_loc = torch.zeros(subquery_lens_tensor.shape[0] + 1,
                                         dtype=torch.int32,
                                         device=self.device)

259
260
        prompt_lens_tensor = torch.tensor(prompt_lens,
                                          dtype=torch.long,
261
                                          device=self.device)
262
263
264
265
266
267
268
269
270
271
272
273
274
        seq_start_loc = torch.zeros(prompt_lens_tensor.shape[0] + 1,
                                    dtype=torch.int32,
                                    device=self.device)

        torch.cumsum(subquery_lens_tensor,
                     dim=0,
                     dtype=subquery_start_loc.dtype,
                     out=subquery_start_loc[1:])

        torch.cumsum(prompt_lens_tensor,
                     dim=0,
                     dtype=seq_start_loc.dtype,
                     out=seq_start_loc[1:])
275

276
        attn_metadata = self.attn_backend.make_metadata(
277
            is_prompt=True,
278
            slot_mapping=slot_mapping,
279
280
281
282
283
            prompt_lens=prompt_lens,
            prompt_lens_tensor=prompt_lens_tensor,
            num_prompt_tokens=num_prompt_tokens,
            num_generation_tokens=0,
            max_subquery_len=max_subquery_len,
284
            max_context_len=None,
285
            max_prompt_len=max_prompt_len,
286
287
            subquery_start_loc=subquery_start_loc,
            seq_start_loc=seq_start_loc,
288
289
            context_lens=context_lens_tensor,
            block_tables=block_tables,
290
            use_cuda_graph=False,
291
            kv_cache_dtype=self.kv_cache_dtype,
292
        )
293
        return (input_tokens, input_positions, attn_metadata, prompt_lens,
294
295
                subquery_lens, lora_index_mapping, lora_prompt_mapping,
                lora_requests)
296
297
298
299

    def _prepare_decode(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
300
301
    ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, List[int],
               List[int], Set[LoRARequest]]:
302
        assert len(seq_group_metadata_list) > 0
303
304
305
        input_tokens: List[int] = []
        input_positions: List[int] = []
        slot_mapping: List[int] = []
306
307
        context_lens: List[int] = []
        block_tables: List[List[int]] = []
308
309
310
        lora_index_mapping: List[int] = []
        lora_prompt_mapping: List[int] = []
        lora_requests: Set[LoRARequest] = set()
311
312
313
314
315

        for seq_group_metadata in seq_group_metadata_list:
            assert not seq_group_metadata.is_prompt

            seq_ids = list(seq_group_metadata.seq_data.keys())
316
317
318
319
320
            lora_id = seq_group_metadata.lora_int_id

            if lora_id > 0:
                lora_requests.add(seq_group_metadata.lora_request)

321
322
323
            for seq_id in seq_ids:
                seq_data = seq_group_metadata.seq_data[seq_id]
                generation_token = seq_data.get_last_token_id()
324
                input_tokens.append(generation_token)
325

326
327
                seq_len = seq_data.get_len()
                position = seq_len - 1
328
                input_positions.append(position)
329

330
331
332
333
                context_len = seq_len if self.sliding_window is None else min(
                    seq_len, self.sliding_window)
                context_lens.append(context_len)

334
335
336
337
                block_table = seq_group_metadata.block_tables[seq_id]
                block_number = block_table[position // self.block_size]
                block_offset = position % self.block_size
                slot = block_number * self.block_size + block_offset
338
339
                slot_mapping.append(slot)
                lora_index_mapping.append(lora_id)
340
                lora_prompt_mapping.append(lora_id)
341
342
343
344
345
346
347

                if self.sliding_window is not None:
                    sliding_window_blocks = (self.sliding_window //
                                             self.block_size)
                    block_table = block_table[-sliding_window_blocks:]
                block_tables.append(block_table)

348
349
350
        # vLLM uses cuda graph only for decoding requests.
        # See `capture_model` API for more details.
        # For decoding requests, batch_size == input_tokens.
351
352
353
354
355
356
357
358
359
360
        batch_size = len(input_tokens)
        max_context_len = max(context_lens)
        use_captured_graph = (
            not self.model_config.enforce_eager
            and batch_size <= _BATCH_SIZES_TO_CAPTURE[-1]
            and max_context_len <= self.max_context_len_to_capture)
        if use_captured_graph:
            graph_batch_size = _get_graph_batch_size(batch_size)
            assert graph_batch_size >= batch_size
            for _ in range(graph_batch_size - batch_size):
361
362
363
                input_tokens.append(0)
                input_positions.append(0)
                slot_mapping.append(_PAD_SLOT_ID)
364
365
                context_lens.append(1)
                block_tables.append([])
366
                lora_index_mapping.append(0)
367
368
            batch_size = graph_batch_size

369
370
371
372
373
374
375
376
377
        input_tokens = torch.tensor(input_tokens,
                                    dtype=torch.long,
                                    device=self.device)
        input_positions = torch.tensor(input_positions,
                                       dtype=torch.long,
                                       device=self.device)
        slot_mapping = torch.tensor(slot_mapping,
                                    dtype=torch.long,
                                    device=self.device)
378
379
        context_lens = torch.tensor(context_lens,
                                    dtype=torch.int,
380
                                    device=self.device)
381
382

        if use_captured_graph:
383
384
385
386
387
388
            # When using cuda-graph all these tensors should be
            # padded.
            assert context_lens.shape[0] == input_tokens.shape[0]
            assert context_lens.shape[0] == input_positions.shape[0]
            assert context_lens.shape[0] == slot_mapping.shape[0]

389
390
391
392
393
394
            # The shape of graph_block_tables is
            # [max batch size, max context len // block size].
            input_block_tables = self.graph_block_tables[:batch_size]
            for i, block_table in enumerate(block_tables):
                if block_table:
                    input_block_tables[i, :len(block_table)] = block_table
395
            block_tables = torch.tensor(input_block_tables, device=self.device)
396
        else:
397
398
            max_block_table_len = max(
                len(block_table) for block_table in block_tables)
399
            block_tables = make_tensor_with_pad(
400
                block_tables,
401
                max_len=max_block_table_len,
402
403
                pad=0,
                dtype=torch.int,
404
                device=self.device,
405
            )
406

407
        attn_metadata = self.attn_backend.make_metadata(
408
            is_prompt=False,
409
            slot_mapping=slot_mapping,
410
            prompt_lens=None,
411
412
413
414
            prompt_lens_tensor=None,
            num_prompt_tokens=0,
            num_generation_tokens=len(input_tokens),
            max_subquery_len=None,
415
            max_context_len=max_context_len,
416
            max_prompt_len=None,
417
418
            subquery_start_loc=None,
            seq_start_loc=None,
419
420
            context_lens=context_lens,
            block_tables=block_tables,
421
            use_cuda_graph=use_captured_graph,
422
            kv_cache_dtype=self.kv_cache_dtype,
423
        )
424
        return (input_tokens, input_positions, attn_metadata,
425
                lora_index_mapping, lora_prompt_mapping, lora_requests)
426
427
428
429
430

    def _prepare_sample(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
        prompt_lens: List[int],
431
        subquery_lens: Optional[List[int]],
432
433
434
    ) -> SamplingMetadata:
        seq_groups: List[Tuple[List[int], SamplingParams]] = []
        selected_token_indices: List[int] = []
Nick Hill's avatar
Nick Hill committed
435
        generators: List[torch.Generator] = []
436
437
438
        selected_token_start_idx = 0
        categorized_sample_indices = {t: [] for t in SamplingType}
        categorized_sample_indices_start_idx = 0
439
        categorized_sampled_token_indices_start_idx = 0
440
441
442
443
444
445
446
447

        for i, seq_group_metadata in enumerate(seq_group_metadata_list):
            seq_ids = list(seq_group_metadata.seq_data.keys())
            sampling_params = seq_group_metadata.sampling_params
            seq_groups.append((seq_ids, sampling_params))

            if seq_group_metadata.is_prompt:
                assert len(seq_ids) == 1
448
449
                assert subquery_lens is not None
                subquery_len = subquery_lens[i]
450
451
                if sampling_params.prompt_logprobs is not None:
                    # NOTE: prompt token positions do not need sample, skip
452
                    categorized_sample_indices_start_idx += subquery_len - 1
453
454

                categorized_sample_indices[
455
456
457
458
                    sampling_params.sampling_type].append([
                        categorized_sample_indices_start_idx,
                        categorized_sampled_token_indices_start_idx
                    ])
459
                categorized_sample_indices_start_idx += 1
460
                categorized_sampled_token_indices_start_idx += 1
461
462
463
464

                if sampling_params.prompt_logprobs is not None:
                    selected_token_indices.extend(
                        range(selected_token_start_idx,
465
                              selected_token_start_idx + subquery_len - 1))
466
                selected_token_indices.append(selected_token_start_idx +
467
                                              subquery_len - 1)
468
                selected_token_start_idx += subquery_len
Nick Hill's avatar
Nick Hill committed
469
470
471

                if sampling_params.seed is not None:
                    seq_group_metadata.state.generator = torch.Generator(
472
                        device=self.device).manual_seed(sampling_params.seed)
473
474
475
476
477
478
479
480
481
            else:
                num_seqs = len(seq_ids)
                selected_token_indices.extend(
                    range(selected_token_start_idx,
                          selected_token_start_idx + num_seqs))
                selected_token_start_idx += num_seqs

                categorized_sample_indices[
                    sampling_params.sampling_type].extend(
482
483
484
485
486
487
488
489
490
                        zip(
                            range(
                                categorized_sample_indices_start_idx,
                                categorized_sample_indices_start_idx +
                                num_seqs),
                            range(
                                categorized_sampled_token_indices_start_idx,
                                categorized_sampled_token_indices_start_idx +
                                num_seqs)))
491
                categorized_sample_indices_start_idx += num_seqs
492
                categorized_sampled_token_indices_start_idx += num_seqs
493

Nick Hill's avatar
Nick Hill committed
494
495
496
            if sampling_params.seed is not None:
                generators.append(seq_group_metadata.state.generator)

497
498
499
500
        selected_token_indices = async_tensor_h2d(selected_token_indices,
                                                  dtype=torch.long,
                                                  target_device=self.device,
                                                  pin_memory=self.pin_memory)
501

502
        categorized_sample_indices = {
503
504
505
506
507
            t: maybe_expand_dim(
                async_tensor_h2d(seq_ids,
                                 dtype=torch.int,
                                 target_device=self.device,
                                 pin_memory=self.pin_memory), 2, 2)
508
509
510
511
512
513
514
515
516
517
518
519
520
            for t, seq_ids in categorized_sample_indices.items()
        }

        seq_data: Dict[int, SequenceData] = {}
        for seq_group_metadata in seq_group_metadata_list:
            seq_data.update(seq_group_metadata.seq_data)

        sampling_metadata = SamplingMetadata(
            seq_groups=seq_groups,
            seq_data=seq_data,
            prompt_lens=prompt_lens,
            selected_token_indices=selected_token_indices,
            categorized_sample_indices=categorized_sample_indices,
Nick Hill's avatar
Nick Hill committed
521
            generators=generators,
522
523
524
        )
        return sampling_metadata

525
526
527
    def prepare_input_tensors(
        self,
        seq_group_metadata_list: Optional[List[SequenceGroupMetadata]],
528
    ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, SamplingMetadata,
529
               Set[int], LoRAMapping]:
530
531
532
533
534
535
        if self.is_driver_worker:
            # NOTE: We assume that all sequences in the group are all prompts or
            # all decodes.
            is_prompt = seq_group_metadata_list[0].is_prompt
            # Prepare input tensors.
            if is_prompt:
536
                (input_tokens, input_positions, attn_metadata, prompt_lens,
537
538
                 subquery_lens, lora_index_mapping, lora_prompt_mapping,
                 lora_requests) = self._prepare_prompt(seq_group_metadata_list)
539
            else:
540
                (input_tokens, input_positions, attn_metadata,
541
542
                 lora_index_mapping, lora_prompt_mapping,
                 lora_requests) = self._prepare_decode(seq_group_metadata_list)
543
                prompt_lens = []
544
                subquery_lens = None
545
            sampling_metadata = self._prepare_sample(seq_group_metadata_list,
546
547
                                                     prompt_lens,
                                                     subquery_lens)
548

549
550
            if self.lora_config:
                lora_mapping = LoRAMapping(
551
                    lora_index_mapping,
552
553
554
555
556
                    lora_prompt_mapping,
                )
            else:
                lora_mapping = None

557
558
559
560
561
562
            # Broadcast the metadata.
            metadata_dict = {
                "input_tokens": input_tokens,
                "input_positions": input_positions,
                "selected_token_indices":
                sampling_metadata.selected_token_indices,
563
564
                "lora_requests": lora_requests,
                "lora_mapping": lora_mapping,
565
            }
566
            metadata_dict.update(attn_metadata.asdict_zerocopy())
567
            broadcast_tensor_dict(metadata_dict, src=0)
568
        else:
569
            metadata_dict = broadcast_tensor_dict(src=0)
570
571
572
573
574
575
            input_tokens = metadata_dict.pop("input_tokens")
            input_positions = metadata_dict.pop("input_positions")
            selected_token_indices = metadata_dict.pop(
                "selected_token_indices")
            lora_mapping = metadata_dict.pop("lora_mapping")
            lora_requests = metadata_dict.pop("lora_requests")
576
            attn_metadata = self.attn_backend.make_metadata(**metadata_dict)
577
578
579
580
            sampling_metadata = SamplingMetadata(
                seq_groups=None,
                seq_data=None,
                prompt_lens=None,
581
                selected_token_indices=selected_token_indices,
582
                categorized_sample_indices=None,
Nick Hill's avatar
Nick Hill committed
583
                generators=None,
584
585
586
                perform_sampling=False,
            )

587
        return (input_tokens, input_positions, attn_metadata,
588
                sampling_metadata, lora_requests, lora_mapping)
589

590
591
592
    @torch.inference_mode()
    def execute_model(
        self,
593
        seq_group_metadata_list: Optional[List[SequenceGroupMetadata]],
594
        kv_caches: List[torch.Tensor],
595
    ) -> Optional[SamplerOutput]:
596
        (input_tokens, input_positions, attn_metadata, sampling_metadata,
597
598
         lora_requests,
         lora_mapping) = self.prepare_input_tensors(seq_group_metadata_list)
599
600
601
602

        if self.lora_config:
            self.set_active_loras(lora_requests, lora_mapping)

603
        # Execute the model.
604
        if attn_metadata.use_cuda_graph:
605
606
607
608
609
            graph_batch_size = input_tokens.shape[0]
            model_executable = self.graph_runners[graph_batch_size]
        else:
            model_executable = self.model
        hidden_states = model_executable(
610
611
612
            input_ids=input_tokens,
            positions=input_positions,
            kv_caches=kv_caches,
613
            attn_metadata=attn_metadata,
614
615
        )

616
617
618
619
620
621
622
        # Compute the logits.
        logits = self.model.compute_logits(hidden_states, sampling_metadata)

        # Only perform sampling in the driver worker.
        if not sampling_metadata.perform_sampling:
            return None

623
624
        # Sample the next token.
        output = self.model.sample(
625
            logits=logits,
626
627
628
629
630
631
632
            sampling_metadata=sampling_metadata,
        )
        return output

    @torch.inference_mode()
    def profile_run(self) -> None:
        # Enable top-k sampling to reflect the accurate memory usage.
633
        sampling_params = SamplingParams(top_p=0.99, top_k=self.vocab_size - 1)
634
635
636
        max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
        max_num_seqs = self.scheduler_config.max_num_seqs

637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
        # This represents the maximum number of different requests
        # that will have unique loras, an therefore the max amount of memory
        # consumption create dummy lora request copies from the lora request
        # passed in, which contains a lora from the lora warmup path.
        dummy_lora_requests = []
        dummy_lora_requests_per_seq = []
        if self.lora_config:
            for idx in range(self.lora_config.max_loras):
                lora_id = idx + 1
                dummy_lora_request = LoRARequest(
                    lora_name=f"warmup_{lora_id}",
                    lora_int_id=lora_id,
                    lora_local_path="/not/a/real/path",
                )
                self.lora_manager.add_dummy_lora(dummy_lora_request,
                                                 rank=LORA_WARMUP_RANK)
                dummy_lora_requests.append(dummy_lora_request)
            dummy_lora_requests_per_seq = [
                dummy_lora_requests[idx % len(dummy_lora_requests)]
                for idx in range(max_num_seqs)
            ]

659
660
661
662
663
664
665
666
667
668
669
670
671
        # Profile memory usage with max_num_sequences sequences and the total
        # number of tokens equal to max_num_batched_tokens.
        seqs: List[SequenceGroupMetadata] = []
        for group_id in range(max_num_seqs):
            seq_len = (max_num_batched_tokens // max_num_seqs +
                       (group_id < max_num_batched_tokens % max_num_seqs))
            seq_data = SequenceData([0] * seq_len)
            seq = SequenceGroupMetadata(
                request_id=str(group_id),
                is_prompt=True,
                seq_data={group_id: seq_data},
                sampling_params=sampling_params,
                block_tables=None,
672
673
                lora_request=dummy_lora_requests_per_seq[group_id]
                if dummy_lora_requests_per_seq else None,
674
675
676
677
678
            )
            seqs.append(seq)

        # Run the model with the dummy inputs.
        num_layers = self.model_config.get_num_layers(self.parallel_config)
679
        kv_caches = [None] * num_layers
680
        self.execute_model(seqs, kv_caches)
681
        torch.cuda.synchronize()
682
683
        return

684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
    def remove_all_loras(self) -> bool:
        if not self.lora_manager:
            raise RuntimeError("LoRA is not enabled.")
        return self.lora_manager.remove_all_loras()

    def set_active_loras(self, lora_requests: List[LoRARequest],
                         lora_mapping: LoRAMapping) -> None:
        if not self.lora_manager:
            raise RuntimeError("LoRA is not enabled.")
        self.lora_manager.set_active_loras(lora_requests, lora_mapping)

    def add_lora(self, lora_request: LoRARequest) -> bool:
        if not self.lora_manager:
            raise RuntimeError("LoRA is not enabled.")
        return self.lora_manager.add_lora(lora_request)

    def remove_lora(self, lora_id: int) -> bool:
        if not self.lora_manager:
            raise RuntimeError("LoRA is not enabled.")
        return self.lora_manager.remove_lora(lora_id)

    def list_loras(self) -> Set[int]:
        if not self.lora_manager:
            raise RuntimeError("LoRA is not enabled.")
        return self.lora_manager.list_loras()

710
    @torch.inference_mode()
711
    def capture_model(self, kv_caches: List[torch.Tensor]) -> None:
712
713
714
715
716
717
718
719
720
721
722
723
        """Cuda graph capture a model.

        Note that CUDA graph's performance gain is negligible if number
        of batched tokens are larger than 200. And since CUDA graph
        requires fixed sized tensors, supporting large/variable batch
        size requires high GPU memory overhead. Thus, vLLM only captures
        decoding requests. Mixed batch (chunked prefill + decoding) or
        prefill requests are not captured.

        Since it is used for decoding-only, it assumes there's only 1 token
        per sequence in the batch.
        """
Woosuk Kwon's avatar
Woosuk Kwon committed
724
725
        # NOTE(woosuk): This is a hack to ensure that the NCCL backend is never
        # deleted before the CUDA graphs.
726
        self.cupy_nccl_backend = cupy_utils.get_nccl_backend()
Woosuk Kwon's avatar
Woosuk Kwon committed
727

728
729
730
731
732
        assert not self.model_config.enforce_eager
        logger.info("Capturing the model for CUDA graphs. This may lead to "
                    "unexpected consequences if the model is not static. To "
                    "run the model in eager mode, set 'enforce_eager=True' or "
                    "use '--enforce-eager' in the CLI.")
733
734
        logger.info("CUDA graphs can take additional 1~3 GiB memory per GPU. "
                    "If you are running out of memory, consider decreasing "
735
736
737
                    "`gpu_memory_utilization` or enforcing eager mode. "
                    "You can also reduce the `max_num_seqs` as needed "
                    "to decrease memory usage.")
738
739
740
741
        start_time = time.perf_counter()

        # Prepare dummy inputs. These will be reused for all batch sizes.
        max_batch_size = max(_BATCH_SIZES_TO_CAPTURE)
742
743
744
        input_tokens = torch.zeros(max_batch_size, dtype=torch.long).cuda()
        input_positions = torch.zeros(max_batch_size, dtype=torch.long).cuda()
        slot_mapping = torch.empty(max_batch_size, dtype=torch.long).cuda()
745
746
747
748
        slot_mapping.fill_(_PAD_SLOT_ID)
        context_lens = torch.ones(max_batch_size, dtype=torch.int32).cuda()
        block_tables = torch.from_numpy(self.graph_block_tables).cuda()

749
750
751
752
753
754
        graph_batch_size = _get_graph_batch_size(
            self.scheduler_config.max_num_seqs)
        batch_size_capture_list = [
            bs for bs in _BATCH_SIZES_TO_CAPTURE if bs <= graph_batch_size
        ]

Woosuk Kwon's avatar
Woosuk Kwon committed
755
756
757
758
759
760
        # NOTE(woosuk): There are 3 backends for all-reduce: custom all-reduce
        # kernel, CuPy NCCL, and PyTorch NCCL. When using CUDA graph, we use
        # either custom all-reduce kernel or CuPy NCCL. When not using CUDA
        # graph, we use either custom all-reduce kernel or PyTorch NCCL.
        # We always prioritize using custom all-reduce kernel but fall back
        # to PyTorch or CuPy NCCL if it is disabled or not supported.
761
        with custom_all_reduce.capture():
762
763
            # NOTE: Capturing the largest batch size first may help reduce the
            # memory usage of CUDA graph.
764
            for batch_size in reversed(batch_size_capture_list):
765
766
                # Create dummy attn_metadata.
                attn_metadata = self.attn_backend.make_metadata(
767
768
769
                    is_prompt=False,
                    slot_mapping=slot_mapping[:batch_size],
                    prompt_lens=None,
770
771
772
773
                    prompt_lens_tensor=None,
                    num_prompt_tokens=0,
                    num_generation_tokens=batch_size,
                    max_subquery_len=None,
774
                    max_context_len=self.max_context_len_to_capture,
775
                    max_prompt_len=None,
776
777
                    subquery_start_loc=None,
                    seq_start_loc=None,
778
779
780
                    context_lens=context_lens[:batch_size],
                    block_tables=block_tables[:batch_size],
                    use_cuda_graph=True,
781
                    kv_cache_dtype=self.kv_cache_dtype,
782
                )
783

784
785
786
787
788
789
790
791
792
793
794
795
                if self.lora_config:
                    lora_mapping = LoRAMapping(
                        [0] * batch_size,
                        [0] * batch_size,
                    )
                    self.set_active_loras(set(), lora_mapping)

                graph_runner = CUDAGraphRunner(self.model)
                graph_runner.capture(
                    input_tokens[:batch_size],
                    input_positions[:batch_size],
                    kv_caches,
796
                    attn_metadata,
797
                    memory_pool=self.graph_memory_pool,
798
                )
799
800
                self.graph_memory_pool = graph_runner.graph.pool()
                self.graph_runners[batch_size] = graph_runner
801
802
803
804
805
806

        end_time = time.perf_counter()
        elapsed_time = end_time - start_time
        # This usually takes < 10 seconds.
        logger.info(f"Graph capturing finished in {elapsed_time:.0f} secs.")

Woosuk Kwon's avatar
Woosuk Kwon committed
807
808
809
810
811
812
813
814
    def __del__(self) -> None:
        # Delete the CUDA graphs before deleting the CuPy NCCL communicator.
        # NOTE(woosuk): This is necessary because otherwise deadlocks can
        # happen.
        # FIXME(woosuk): This is a bit hacky. Find a more robust solution.
        self.graph_runners.clear()
        self.cupy_nccl_backend = None

815
816
817
818
    @property
    def vocab_size(self) -> int:
        return self.model_config.get_vocab_size()

819
820
821
822
823
824
825
826
827
828
829
830
831

class CUDAGraphRunner:

    def __init__(self, model: nn.Module):
        self.model = model
        self.graph = None
        self.input_buffers: Dict[str, torch.Tensor] = {}
        self.output_buffers: Dict[str, torch.Tensor] = {}

    def capture(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
832
833
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
834
835
836
837
838
839
        memory_pool,
    ) -> None:
        assert self.graph is None
        # Run the model once without capturing the graph.
        # This is to make sure that the captured graph does not include the
        # kernel launches for initial benchmarking (e.g., Triton autotune).
840
        with _maybe_cupy_nccl():
Woosuk Kwon's avatar
Woosuk Kwon committed
841
            self.model(
842
843
844
                input_ids,
                positions,
                kv_caches,
845
                attn_metadata,
846
847
848
            )
        torch.cuda.synchronize()

Woosuk Kwon's avatar
Woosuk Kwon committed
849
850
851
852
853
        # Capture the graph.
        # NOTE(woosuk): Python 3.8 does not support multi-line with statements.
        # https://stackoverflow.com/questions/31039022/python-multi-line-with-statement
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph, pool=memory_pool):  # noqa: SIM117
854
            with _maybe_cupy_nccl():
Woosuk Kwon's avatar
Woosuk Kwon committed
855
856
857
858
                hidden_states = self.model(
                    input_ids,
                    positions,
                    kv_caches,
859
                    attn_metadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
860
861
862
                )
        torch.cuda.synchronize()

863
864
865
866
867
        # Save the input and output buffers.
        self.input_buffers = {
            "input_ids": input_ids,
            "positions": positions,
            "kv_caches": kv_caches,
868
869
870
            "slot_mapping": attn_metadata.slot_mapping,
            "context_lens": attn_metadata.context_lens,
            "block_tables": attn_metadata.block_tables,
871
872
873
874
875
876
877
878
        }
        self.output_buffers = {"hidden_states": hidden_states}
        return

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
879
880
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
881
882
883
884
885
    ) -> torch.Tensor:
        # KV caches are fixed tensors, so we don't need to copy them.
        del kv_caches

        # Copy the input tensors to the input buffers.
886
887
        self.input_buffers["input_ids"].copy_(input_ids, non_blocking=True)
        self.input_buffers["positions"].copy_(positions, non_blocking=True)
888
        self.input_buffers["slot_mapping"].copy_(attn_metadata.slot_mapping,
889
                                                 non_blocking=True)
890
        self.input_buffers["context_lens"].copy_(attn_metadata.context_lens,
891
                                                 non_blocking=True)
892
        self.input_buffers["block_tables"].copy_(attn_metadata.block_tables,
893
                                                 non_blocking=True)
894
895
896
897
898
899
900
901
902
        # Run the graph.
        self.graph.replay()

        # Return the output tensor.
        return self.output_buffers["hidden_states"]

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

903

904
905
906
907
908
909
910
911
912
@contextlib.contextmanager
def _maybe_cupy_nccl():
    if cupy_utils.is_initialized() and not custom_all_reduce.is_initialized():
        with with_cupy_nccl_for_all_reduce():
            yield
    else:
        yield


913
def _get_graph_batch_size(batch_size: int) -> int:
914
915
916
917
918
    """Returns the padded batch size given actual batch size.

    Batch sizes are 1, 2, 4, _BATCH_SIZE_ALIGNMENT,
    2*_BATCH_SIZE_ALIGNMENT, 3*_BATCH_SIZE_ALIGNMENT...
    """
919
920
921
922
923
    if batch_size <= 2:
        return batch_size
    elif batch_size <= 4:
        return 4
    else:
924
925
        return ((batch_size + _BATCH_SIZE_ALIGNMENT - 1) //
                _BATCH_SIZE_ALIGNMENT * _BATCH_SIZE_ALIGNMENT)