cpu_model_runner.py 22.1 KB
Newer Older
1
2
import dataclasses
import weakref
3
from dataclasses import dataclass
4
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
5
6

import torch
7
from torch import nn
8
9

from vllm.attention import AttentionMetadata, get_attn_backend
10
from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig,
11
12
                         ModelConfig, ParallelConfig, PromptAdapterConfig,
                         SchedulerConfig)
13
14
from vllm.logger import init_logger
from vllm.model_executor import SamplingMetadata
15
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
16
from vllm.model_executor.layers.sampler import SamplerOutput
17
from vllm.model_executor.model_loader import get_model
18
from vllm.multimodal import (MULTIMODAL_REGISTRY, BatchedTensorInputs,
19
                             MultiModalInputs)
20
21
from vllm.sequence import (IntermediateTensors, SequenceData,
                           SequenceGroupMetadata)
22
from vllm.utils import make_tensor_with_pad
23
from vllm.worker.model_runner_base import (
24
    ModelRunnerBase, ModelRunnerInputBase, ModelRunnerInputBuilderBase,
25
26
27
28
29
30
31
    _add_attn_metadata_broadcastable_dict,
    _add_sampling_metadata_broadcastable_dict,
    _init_attn_metadata_from_tensor_dict,
    _init_sampling_metadata_from_tensor_dict)

if TYPE_CHECKING:
    from vllm.attention.backends.abstract import AttentionBackend
32
33
34
35
36
37

logger = init_logger(__name__)

_PAD_SLOT_ID = -1


38
@dataclass(frozen=True)
39
class ModelInputForCPU(ModelRunnerInputBase):
40
    """
41
    Base class contains metadata needed for the base model forward pass on CPU
42
43
44
45
    """
    input_tokens: Optional[torch.Tensor] = None
    input_positions: Optional[torch.Tensor] = None
    attn_metadata: Optional["AttentionMetadata"] = None
46
    multi_modal_kwargs: Optional[BatchedTensorInputs] = None
47
    virtual_engine: Optional[int] = None
48
49
    seq_lens: Optional[List[int]] = None
    query_lens: Optional[List[int]] = None
50
51
52
53
54
55
56
57
58

    def as_broadcastable_tensor_dict(
            self) -> Dict[str, Union[int, torch.Tensor]]:
        tensor_dict = {
            "input_tokens": self.input_tokens,
            "input_positions": self.input_positions,
            "multi_modal_kwargs": self.multi_modal_kwargs,
        }
        _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
59

60
61
62
63
        return tensor_dict

    @classmethod
    def from_broadcasted_tensor_dict(
64
65
66
67
        cls: Type["ModelInputForCPU"],
        tensor_dict: Dict[str, Any],
        attn_backend: Optional["AttentionBackend"] = None
    ) -> "ModelInputForCPU":
68
69
70
71
72
73
        if attn_backend is not None:
            tensor_dict = _init_attn_metadata_from_tensor_dict(
                attn_backend, tensor_dict)
        return cls(**tensor_dict)


74
75
76
77
78
79
@dataclass(frozen=True)
class ModelInputForCPUWithSamplingMetadata(ModelInputForCPU):
    """
    Used by the ModelRunner.
    """
    sampling_metadata: Optional["SamplingMetadata"] = None
80

81
82
83
84
85
86
87
88
89
    def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
        tensor_dict = {
            "input_tokens": self.input_tokens,
            "input_positions": self.input_positions,
        }
        _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
        _add_sampling_metadata_broadcastable_dict(tensor_dict,
                                                  self.sampling_metadata)
        return tensor_dict
90

91
92
93
94
95
96
97
98
99
100
101
    @classmethod
    def from_broadcasted_tensor_dict(
        cls,
        tensor_dict: Dict[str, Any],
        attn_backend: Optional["AttentionBackend"] = None,
    ) -> "ModelInputForCPUWithSamplingMetadata":
        tensor_dict = _init_sampling_metadata_from_tensor_dict(tensor_dict)
        if attn_backend is not None:
            tensor_dict = _init_attn_metadata_from_tensor_dict(
                attn_backend, tensor_dict)
        return cls(**tensor_dict)
102
103


104
class ModelInputForCPUBuilder(ModelRunnerInputBuilderBase[ModelInputForCPU]):
105

106
107
108
109
110
111
112
113
114
115
116
117
    def __init__(self,
                 runner: "CPUModelRunner",
                 finished_requests_ids: Optional[List[str]] = None) -> None:
        super().__init__()
        self.seq_group_metadata_list: List[SequenceGroupMetadata] = []
        self.runner = runner
        self.model_input_cls = self.runner._model_input_cls
        self.attn_backend = self.runner.attn_backend
        self.sliding_window = self.runner.sliding_window
        self.block_size = self.runner.block_size
        self.device = self.runner.device
        self.multi_modal_input_mapper = self.runner.multi_modal_input_mapper
118

119
120
    def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
        self.seq_group_metadata_list.append(seq_group_metadata)
121

122
123
124
125
126
127
128
129
130
131
132
133
134
135
    def build(self) -> ModelInputForCPU:
        multi_modal_kwargs = None
        # NOTE: We assume that all sequences in the group are all prompts or
        # all decodes.
        is_prompt = self.seq_group_metadata_list[0].is_prompt
        # Prepare input tensors.
        if is_prompt:
            (input_tokens, input_positions, attn_metadata, seq_lens,
             multi_modal_kwargs) = self._prepare_prompt(
                 self.seq_group_metadata_list)
        else:
            (input_tokens, input_positions,
             attn_metadata) = self._prepare_decode(
                 self.seq_group_metadata_list)
136
            seq_lens = None
137
138
139
140
141
142
143
144
145
146
147
148

        return self.model_input_cls(
            input_tokens=input_tokens,
            input_positions=input_positions,
            attn_metadata=attn_metadata,
            multi_modal_kwargs=multi_modal_kwargs,
            # query_lens is not needed if chunked prefill is not
            # supported. Since CPU worker doesn't support chunked prefill
            # just use seq_lens instead.
            seq_lens=seq_lens,
            query_lens=seq_lens,
        )
149

150
    def _compute_multi_modal_input(self, seq_data: SequenceData, mm_data,
151
152
153
                                   computed_len: int,
                                   mm_processor_kwargs: Dict[str, Any]):
        mm_kwargs = self.multi_modal_input_mapper(mm_data, mm_processor_kwargs)
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182

        # special processing for mrope position deltas.
        mrope_positions = None
        if self.runner.model_is_mrope:
            image_grid_thw = mm_kwargs.get("image_grid_thw", None)
            video_grid_thw = mm_kwargs.get("video_grid_thw", None)
            assert image_grid_thw is not None or video_grid_thw is not None, (
                "mrope embedding type requires multi-modal input mapper "
                "returns 'image_grid_thw' or 'video_grid_thw'.")

            hf_config = self.runner.model_config.hf_config
            token_ids = seq_data.get_token_ids()

            mrope_positions, mrope_position_delta = \
                MRotaryEmbedding.get_input_positions(
                    token_ids,
                    image_grid_thw=image_grid_thw,
                    video_grid_thw=video_grid_thw,
                    image_token_id=hf_config.image_token_id,
                    video_token_id=hf_config.video_token_id,
                    vision_start_token_id=hf_config.vision_start_token_id,
                    vision_end_token_id=hf_config.vision_end_token_id,
                    spatial_merge_size=hf_config.vision_config.
                    spatial_merge_size,
                    context_len=computed_len,
                )
            seq_data.mrope_position_delta = mrope_position_delta
        return mm_kwargs, mrope_positions

183
184
185
    def _prepare_prompt(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
186
    ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, List[int],
187
               BatchedTensorInputs]:
188
189
190
        assert len(seq_group_metadata_list) > 0
        input_tokens: List[int] = []
        input_positions: List[int] = []
191
192
        input_mrope_positions: List[List[int]] = [[] for _ in range(3)]

193
        slot_mapping: List[int] = []
194
        seq_lens: List[int] = []
195
        multi_modal_inputs_list: List[MultiModalInputs] = []
196
197
198
199
200
201
202
203
204
205

        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()
            computed_len = seq_data.get_num_computed_tokens()
206
            seq_len = len(prompt_tokens)
207

208
            seq_lens.append(seq_len)  # Prompt token num
209
210
            input_tokens.extend(prompt_tokens)  # Token ids

211
212
213
            mrope_positions = None
            if (mm_data := seq_group_metadata.multi_modal_data):
                mm_kwargs, mrope_positions = self._compute_multi_modal_input(
214
215
                    seq_data, mm_data, computed_len,
                    seq_group_metadata.mm_processor_kwargs)
216
217
                multi_modal_inputs_list.append(mm_kwargs)

218
219
220
            # Token position ids
            # NOTE(woosuk): Here we assume that the first token in the prompt
            # is always the first token in the sequence.
221
222
223
224
225
            if mrope_positions:
                for idx in range(3):
                    input_mrope_positions[idx].extend(mrope_positions[idx])
            else:
                input_positions.extend(list(range(computed_len, seq_len)))
226

227
228
229
            # 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,
230
            # where start_idx is max(0, seq_len - sliding_window).
231
232
233
234
235
            # 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:
236
                start_idx = max(0, seq_len - self.sliding_window)
237

238
            for i in range(computed_len, seq_len):
239
240
241
242
243
244
245
246
247
248
                if i < start_idx:
                    slot_mapping.append(_PAD_SLOT_ID)
                    continue

                block_number = block_table[i //
                                           self.block_size]  # type: ignore
                block_offset = i % self.block_size  # type: ignore
                slot = block_number * self.block_size + block_offset
                slot_mapping.append(slot)

249
250
251
252
253
        if any(input_mrope_positions):
            input_positions = None  # type: ignore
        else:
            input_mrope_positions = None  # type: ignore

254
255
256
257
258
        num_prompt_tokens = len(input_tokens)

        input_tokens = torch.tensor(input_tokens,
                                    dtype=torch.long,
                                    device=self.device)  # type: ignore
259
260
        input_positions = torch.tensor(input_positions
                                       or input_mrope_positions,
261
262
263
264
265
266
267
268
                                       dtype=torch.long,
                                       device=self.device)  # type: ignore
        slot_mapping = torch.tensor(slot_mapping,
                                    dtype=torch.long,
                                    device=self.device)  # type: ignore

        attn_metadata = self.attn_backend.make_metadata(
            is_prompt=True,
269
            seq_lens=seq_lens,
270
271
            seq_lens_tensor=torch.tensor([]),
            max_decode_seq_len=0,
272
            num_prefills=len(seq_lens),
273
274
275
276
277
            num_prefill_tokens=num_prompt_tokens,
            num_decode_tokens=0,
            block_tables=torch.tensor([]),
            slot_mapping=slot_mapping,
        )
278

279
        multi_modal_kwargs = MultiModalInputs.batch(multi_modal_inputs_list)
280

281
        return (input_tokens, input_positions, attn_metadata, seq_lens,
282
                multi_modal_kwargs)
283
284
285
286
287
288
289
290

    def _prepare_decode(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
    ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata]:
        assert len(seq_group_metadata_list) > 0
        input_tokens: List[int] = []
        input_positions: List[int] = []
291
        input_mrope_positions: List[List[int]] = [[] for _ in range(3)]
292
        slot_mapping: List[int] = []
293
        seq_lens: List[int] = []
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
        block_tables: List[List[int]] = []

        for seq_group_metadata in seq_group_metadata_list:
            assert not seq_group_metadata.is_prompt
            assert seq_group_metadata.token_chunk_size == 1

            seq_ids = list(seq_group_metadata.seq_data.keys())

            for seq_id in seq_ids:
                seq_data = seq_group_metadata.seq_data[seq_id]
                generation_token = seq_data.get_last_token_id()
                input_tokens.append(generation_token)

                seq_len = seq_data.get_len()
                position = seq_len - 1
309
310
311
312
313
314
315
316
317
318
319
                if seq_data.mrope_position_delta is not None:
                    context_len = seq_data.get_num_computed_tokens()
                    next_pos = MRotaryEmbedding.get_next_input_positions(
                        seq_data.mrope_position_delta,
                        context_len,
                        seq_len,
                    )
                    for idx in range(3):
                        input_mrope_positions[idx].extend(next_pos[idx])
                else:
                    input_positions.append(position)
320

321
                seq_len = seq_len if self.sliding_window is None else min(
322
                    seq_len, self.sliding_window)
323
                seq_lens.append(seq_len)
324
325
326
327
328
329
330
331
332
333
334
335
336

                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
                slot_mapping.append(slot)

                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)

337
338
339
340
341
        if any(input_mrope_positions):
            input_positions = None  # type: ignore
        else:
            input_mrope_positions = None  # type: ignore

342
        max_decode_seq_len = max(seq_lens)
343
344
345
346

        input_tokens = torch.tensor(input_tokens,
                                    dtype=torch.long,
                                    device=self.device)
347
348
        input_positions = torch.tensor(input_positions
                                       or input_mrope_positions,
349
350
351
352
353
                                       dtype=torch.long,
                                       device=self.device)
        slot_mapping = torch.tensor(slot_mapping,
                                    dtype=torch.long,
                                    device=self.device)
354
355
356
        seq_lens_tensor = torch.tensor(seq_lens,
                                       dtype=torch.int,
                                       device=self.device)
357
358
359
360
361
362
363
364
365
366
367

        block_tables = make_tensor_with_pad(
            block_tables,
            pad=0,
            dtype=torch.int,
            device=self.device,
        )

        attn_metadata = self.attn_backend.make_metadata(
            is_prompt=False,
            slot_mapping=slot_mapping,
368
369
            seq_lens=seq_lens,
            seq_lens_tensor=seq_lens_tensor,
370
            max_decode_seq_len=max_decode_seq_len,
371
372
373
374
375
376
377
378
379
380
381
            num_prefill_tokens=0,
            num_decode_tokens=len(input_tokens),
            num_prefills=0,
            block_tables=block_tables,
        )
        return (
            input_tokens,
            input_positions,
            attn_metadata,
        )

382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438

class CPUModelRunner(ModelRunnerBase[ModelInputForCPU]):
    _model_input_cls: Type[ModelInputForCPUWithSamplingMetadata] = (
        ModelInputForCPUWithSamplingMetadata)
    _builder_cls: Type[ModelInputForCPUBuilder] = ModelInputForCPUBuilder

    def __init__(
        self,
        model_config: ModelConfig,
        parallel_config: ParallelConfig,
        scheduler_config: SchedulerConfig,
        device_config: DeviceConfig,
        cache_config: CacheConfig,
        load_config: LoadConfig,
        lora_config: Optional[LoRAConfig],
        kv_cache_dtype: Optional[str] = "auto",
        prompt_adapter_config: Optional[PromptAdapterConfig] = None,
        is_driver_worker: bool = False,
        *args,
        **kwargs,
    ):
        self.model_config = model_config
        self.parallel_config = parallel_config
        self.scheduler_config = scheduler_config
        # Currently, CPU worker doesn't support chunked prefill.
        assert self.scheduler_config.chunked_prefill_enabled is False
        self.device_config = device_config
        self.cache_config = cache_config
        self.lora_config = lora_config
        self.prompt_adapter_config = prompt_adapter_config
        self.load_config = load_config
        self.is_driver_worker = is_driver_worker

        self.device = self.device_config.device

        self.kv_cache_dtype = kv_cache_dtype
        self.sliding_window = model_config.get_sliding_window()
        self.block_size = cache_config.block_size
        self.attn_backend = get_attn_backend(
            self.model_config.get_num_attention_heads(self.parallel_config),
            self.model_config.get_head_size(),
            self.model_config.get_num_kv_heads(self.parallel_config),
            self.model_config.get_sliding_window(),
            self.model_config.dtype,
            self.kv_cache_dtype,
            self.block_size,
        )

        # Multi-modal data support
        self.mm_registry = MULTIMODAL_REGISTRY
        self.multi_modal_input_mapper = self.mm_registry \
            .create_input_mapper(self.model_config)
        self.mm_registry.init_mm_limits_per_prompt(self.model_config)

        # Lazy initialization.
        self.model: nn.Module  # Set after init_Model

439
440
441
442
443
444
445
446
447
    @property
    def model_is_mrope(self) -> bool:
        """Detect if the model has "mrope" rope_scaling type.
        mrope requires keep "rope_deltas" between prompt and decoding phases."""
        rope_scaling = getattr(self.model_config.hf_config, "rope_scaling", {})
        if rope_scaling is None:
            return False
        return rope_scaling.get("type", None) == "mrope"

448
449
450
451
452
453
454
455
456
    def load_model(self) -> None:
        self.model = get_model(model_config=self.model_config,
                               load_config=self.load_config,
                               device_config=self.device_config,
                               lora_config=self.lora_config,
                               parallel_config=self.parallel_config,
                               scheduler_config=self.scheduler_config,
                               cache_config=self.cache_config)

457
458
459
    def make_model_input_from_broadcasted_tensor_dict(
        self,
        tensor_dict: Dict[str, Any],
460
461
    ) -> ModelInputForCPUWithSamplingMetadata:
        return ModelInputForCPUWithSamplingMetadata.from_broadcasted_tensor_dict(  # noqa: E501
462
463
464
465
            tensor_dict,
            attn_backend=self.attn_backend,
        )

466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
    def _prepare_model_input_tensors(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
        finished_requests_ids: Optional[List[str]] = None
    ) -> ModelInputForCPUWithSamplingMetadata:
        """Helper method to prepare the model input based on a given sequence
        group. Prepares metadata needed for the base model forward pass but not
        metadata for possible additional steps, e.g., sampling.

        """
        builder = self._builder_cls(weakref.proxy(self), finished_requests_ids)
        for seq_group_metadata in seq_group_metadata_list:
            builder.add_seq_group(seq_group_metadata)

        return builder.build()  # type: ignore

482
    def prepare_model_input(
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
        virtual_engine: int = 0,
        finished_requests_ids: Optional[List[str]] = None
    ) -> ModelInputForCPUWithSamplingMetadata:
        """Prepare the model input based on a given sequence group, including
        metadata for the sampling step.

        """
        model_input = self._prepare_model_input_tensors(
            seq_group_metadata_list, finished_requests_ids)
        # Sampling metadata is only required for the final pp group
        generators = self.get_generators(finished_requests_ids)
        sampling_metadata = SamplingMetadata.prepare(seq_group_metadata_list,
                                                     model_input.seq_lens,
                                                     model_input.query_lens,
                                                     self.device,
                                                     pin_memory=False,
                                                     generators=generators)

        return dataclasses.replace(model_input,
                                   sampling_metadata=sampling_metadata,
                                   virtual_engine=virtual_engine)
506

507
    @torch.no_grad()
508
509
    def execute_model(
        self,
510
        model_input: ModelInputForCPUWithSamplingMetadata,
511
        kv_caches: List[torch.Tensor],
512
        intermediate_tensors: Optional[IntermediateTensors] = None,
513
514
515
516
517
518
        num_steps: int = 1,
    ) -> Optional[List[SamplerOutput]]:
        if num_steps > 1:
            raise ValueError(
                "CPU worker does not support multi-step execution.")

519
520
        model_executable = self.model
        execute_model_kwargs = {
521
522
523
524
525
526
527
528
529
530
            "input_ids":
            model_input.input_tokens,
            "positions":
            model_input.input_positions,
            "kv_caches":
            kv_caches,
            "attn_metadata":
            model_input.attn_metadata,
            **MultiModalInputs.as_kwargs(model_input.multi_modal_kwargs or {},
                                         device=self.device),
531
532
            "intermediate_tensors":
            intermediate_tensors,
533
534
535
536
537
        }

        hidden_states = model_executable(**execute_model_kwargs)

        # Compute the logits.
538
539
        logits = self.model.compute_logits(hidden_states,
                                           model_input.sampling_metadata)
540
541

        # Only perform sampling in the driver worker.
542
        if not self.is_driver_worker:
543
            return []
544
545
546
547

        # Sample the next token.
        output = self.model.sample(
            logits=logits,
548
            sampling_metadata=model_input.sampling_metadata,
549
        )
550
        return [output]