cpu_model_runner.py 27.6 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
import dataclasses
import weakref
5
from collections import defaultdict
6
from dataclasses import dataclass
7
8
from typing import (TYPE_CHECKING, Any, Dict, List, Optional, Set, Type,
                    TypeVar, Union)
9
10

import torch
11
from torch import nn
12
13

from vllm.attention import AttentionMetadata, get_attn_backend
14
from vllm.config import VllmConfig
youkaichao's avatar
youkaichao committed
15
from vllm.forward_context import set_forward_context
16
from vllm.logger import init_logger
17
18
19
from vllm.lora.layers import LoRAMapping
from vllm.lora.request import LoRARequest
from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
20
from vllm.model_executor import SamplingMetadata
21
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
22
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
23
from vllm.model_executor.model_loader import get_model
24
from vllm.model_executor.models import supports_lora, supports_multimodal
25
26
from vllm.multimodal import (BatchedTensorInputs, MultiModalKwargs,
                             MultiModalPlaceholderMap)
27
28
from vllm.sequence import (IntermediateTensors, SequenceData,
                           SequenceGroupMetadata)
29
from vllm.worker.model_runner_base import (
30
    ModelRunnerBase, ModelRunnerInputBase, ModelRunnerInputBuilderBase,
31
32
33
34
35
36
37
    _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
38
39
40

logger = init_logger(__name__)

41
TModelInputForCPU = TypeVar('TModelInputForCPU', bound="ModelInputForCPU")
42
43
44
_PAD_SLOT_ID = -1


45
@dataclass(frozen=True)
46
class ModelInputForCPU(ModelRunnerInputBase):
47
    """
48
    Base class contains metadata needed for the base model forward pass on CPU
49
50
51
    """
    input_tokens: Optional[torch.Tensor] = None
    input_positions: Optional[torch.Tensor] = None
52
    token_type_ids: Optional[torch.Tensor] = None
53
    attn_metadata: Optional["AttentionMetadata"] = None
54
    multi_modal_kwargs: Optional[BatchedTensorInputs] = None
55
    virtual_engine: Optional[int] = None
56
57
    seq_lens: Optional[List[int]] = None
    query_lens: Optional[List[int]] = None
58
59
    lora_mapping: Optional["LoRAMapping"] = None
    lora_requests: Optional[Set[LoRARequest]] = None
60
61
62
63
64
65

    def as_broadcastable_tensor_dict(
            self) -> Dict[str, Union[int, torch.Tensor]]:
        tensor_dict = {
            "input_tokens": self.input_tokens,
            "input_positions": self.input_positions,
66
            "token_type_ids": self.token_type_ids,
67
            "multi_modal_kwargs": self.multi_modal_kwargs,
68
69
            "lora_requests": self.lora_requests,
            "lora_mapping": self.lora_mapping,
70
71
        }
        _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
72

73
74
75
76
        return tensor_dict

    @classmethod
    def from_broadcasted_tensor_dict(
77
        cls: Type[TModelInputForCPU],
78
79
        tensor_dict: Dict[str, Any],
        attn_backend: Optional["AttentionBackend"] = None
80
    ) -> TModelInputForCPU:
81
82
83
84
85
86
        if attn_backend is not None:
            tensor_dict = _init_attn_metadata_from_tensor_dict(
                attn_backend, tensor_dict)
        return cls(**tensor_dict)


87
88
89
90
91
92
@dataclass(frozen=True)
class ModelInputForCPUWithSamplingMetadata(ModelInputForCPU):
    """
    Used by the ModelRunner.
    """
    sampling_metadata: Optional["SamplingMetadata"] = None
93
    is_prompt: Optional[bool] = None
94

95
96
97
98
    def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:
        tensor_dict = {
            "input_tokens": self.input_tokens,
            "input_positions": self.input_positions,
99
            "token_type_ids": self.token_type_ids,
100
            "multi_modal_kwargs": self.multi_modal_kwargs,
101
102
103
104
105
        }
        _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
        _add_sampling_metadata_broadcastable_dict(tensor_dict,
                                                  self.sampling_metadata)
        return tensor_dict
106

107
108
109
110
111
112
113
114
115
116
117
    @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)
118
119


120
class ModelInputForCPUBuilder(ModelRunnerInputBuilderBase[ModelInputForCPU]):
121

122
123
124
125
126
    class ModelInputData:

        def __init__(self, use_mrope: bool):
            self.use_mrope = use_mrope
            self.input_tokens: List[int] = []
127
            self.input_positions: List[int] = []
128
            self.token_type_ids: Optional[List[int]] = []
129
130
131
132
133
134
135
136
137
138
139
140
141
            self.seq_lens: List[int] = []
            self.query_lens: List[int] = []
            self.prefill_block_tables: List[List[int]] = []
            self.decode_block_tables: List[List[int]] = []
            self.max_decode_seq_len: int = 0
            self.num_prefills: int = 0
            self.num_prefill_tokens: int = 0
            self.num_decode_tokens: int = 0
            self.slot_mapping: List[int] = []
            self.multi_modal_inputs_list: List[MultiModalKwargs] = []
            self.multi_modal_placeholder_maps: Dict[
                str, MultiModalPlaceholderMap] = defaultdict(
                    MultiModalPlaceholderMap)
142
143
            self.input_mrope_positions: List[List[int]] = [[]
                                                           for _ in range(3)]
144

145
146
147
148
149
    def __init__(self,
                 runner: "CPUModelRunner",
                 finished_requests_ids: Optional[List[str]] = None) -> None:
        super().__init__()
        self.runner = runner
150
151
        self.chunked_prefill = (runner.scheduler_config.chunked_prefill_enabled
                                or runner.cache_config.enable_prefix_caching)
152
153
        self.model_input_cls = self.runner._model_input_cls
        self.attn_backend = self.runner.attn_backend
154
155
156
157
        self.sliding_window = self.runner.sliding_window
        self.block_size = self.runner.block_size
        self.device = self.runner.device
        self.enable_lora = self.runner.lora_config is not None
158
159
160
161
162
163
164
165
        if self.runner.attn_backend is not None:
            # spec decode (e.g. Medusa) does not have atten backend
            attn_backend = self.runner.attn_backend
            self.att_metadata_builder = attn_backend.get_builder_cls()(self)

    def prepare(self,
                finished_requests_ids: Optional[List[str]] = None) -> None:
        self.seq_group_metadata_list: List[SequenceGroupMetadata] = []
166
167
        self.input_data = ModelInputForCPUBuilder.ModelInputData(
            self.runner.model_config.uses_mrope)
168
        self.att_metadata_builder.prepare()
169

170
171
    def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
        self.seq_group_metadata_list.append(seq_group_metadata)
172

173
174
175
176
    def set_seq_group_list(
            self, seq_group_metadata_list: List[SequenceGroupMetadata]):
        self.seq_group_metadata_list = seq_group_metadata_list

177
    def build(self) -> ModelInputForCPU:
178
179
180
181
182
183
184
185
        self._build_input_data()

        input_data = self.input_data
        input_tokens = torch.tensor(input_data.input_tokens,
                                    dtype=torch.long,
                                    device="cpu")
        input_positions = torch.tensor(
            input_data.input_positions
186
187
            if not any(input_data.input_mrope_positions) else
            input_data.input_mrope_positions,
188
189
            dtype=torch.long,
            device="cpu")
190
191
192
193
        token_type_ids = torch.tensor(input_data.token_type_ids,
                                    dtype=torch.long,
                                    device="cpu") \
                                    if input_data.token_type_ids else None
194
195

        # For multi-modal models
196
        multi_modal_kwargs = None
197
198
199
200
201
202
        if len(input_data.multi_modal_inputs_list) != 0:
            multi_modal_kwargs = MultiModalKwargs.batch(
                input_data.multi_modal_inputs_list)

        attn_metadata = self.att_metadata_builder.build(
            input_data.seq_lens, input_data.query_lens, -1, -1)
203

204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
        is_prompt = (self.seq_group_metadata_list[0].is_prompt
                     if self.seq_group_metadata_list else None)
        # LoRA data.
        lora_requests = set()
        lora_mapping = None
        if self.enable_lora:
            lora_requests = set(seq.lora_request
                                for seq in self.seq_group_metadata_list
                                if seq.lora_request is not None)

            lora_mapping = self._prepare_lora_input(
                self.seq_group_metadata_list, is_prompt)

        return self.model_input_cls(input_tokens=input_tokens,
                                    input_positions=input_positions,
                                    token_type_ids=token_type_ids,
                                    seq_lens=input_data.seq_lens,
                                    query_lens=input_data.query_lens,
                                    attn_metadata=attn_metadata,
                                    multi_modal_kwargs=multi_modal_kwargs,
                                    lora_mapping=lora_mapping,
                                    lora_requests=lora_requests)
226

227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
    def _build_input_data(self):
        for seq_group_metadata in self.seq_group_metadata_list:
            for seq_id, seq_data in seq_group_metadata.seq_data.items():
                if seq_group_metadata.is_prompt:
                    self._compute_prompt_input_tokens(self.input_data,
                                                      seq_group_metadata,
                                                      seq_data, seq_id)
                    if seq_group_metadata.multi_modal_data:
                        self._compute_multi_modal_input(
                            seq_group_metadata, seq_data)
                else:
                    self._compute_decode_input_tokens(self.input_data,
                                                      seq_group_metadata,
                                                      seq_data, seq_id)

    def _compute_decode_input_tokens(self, data: ModelInputData,
                                     seq_group_metadata: SequenceGroupMetadata,
                                     seq_data: SequenceData, seq_id: int):
        """
        Compute decode input tokens, positions, block table and slot mapping.
        """
        block_size = self.runner.block_size

        block_table = seq_group_metadata.block_tables[seq_id]
        seq_len = seq_data.get_len()
        context_len = seq_data.get_num_computed_tokens()

        tokens = seq_data.get_last_token_id()
        token_positions = seq_len - 1
        block_number = block_table[token_positions // block_size]
        block_offset = token_positions % block_size
        slot = block_number * block_size + block_offset

        # For paged_attention kernel
        if self.runner.sliding_window:
            start_idx = max(0, seq_len - self.runner.sliding_window)
            start_block = start_idx // block_size
            start_idx = start_block * block_size
            seq_len = seq_len - start_idx
            block_table = block_table[start_block:]

        # For MRotaryEmbedding
269
        if seq_data.mrope_position_delta is not None:
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
            next_pos = MRotaryEmbedding.get_next_input_positions(
                seq_data.mrope_position_delta,
                context_len,
                seq_len,
            )
            for idx in range(3):
                data.input_mrope_positions[idx].extend(  # type: ignore
                    next_pos[idx])
        else:
            data.input_positions.append(token_positions)  # type: ignore

        # Update fields
        data.input_tokens.append(tokens)
        data.max_decode_seq_len = max(data.max_decode_seq_len, seq_len)
        data.num_decode_tokens += 1
        data.slot_mapping.append(slot)
        data.decode_block_tables.append(block_table)
        data.query_lens.append(1)
        data.seq_lens.append(seq_len)

    def _compute_prompt_input_tokens(self, data: ModelInputData,
                                     seq_group_metadata: SequenceGroupMetadata,
                                     seq_data: SequenceData, seq_id: int):
        """
        Compute prompt input tokens, positions, block table and slot mapping.
        """
        token_chunk_size = seq_group_metadata.token_chunk_size
        block_size = self.runner.block_size

        block_table = seq_group_metadata.block_tables[seq_id]
        seq_len = seq_data.get_len()
        context_len = seq_data.get_num_computed_tokens()
        seq_len = min(seq_len, context_len + token_chunk_size)

        # For prefix caching
        prefix_cache_block_num = len(seq_group_metadata.computed_block_nums)
        if prefix_cache_block_num > 0:
            prefix_cache_len = (prefix_cache_block_num *
                                self.runner.block_size)
            if prefix_cache_len <= context_len:
                # We already passed the cache hit region,
                # so do normal computation.
                pass
            elif context_len < prefix_cache_len < seq_len:
                # Partial hit. Compute the missing part.
                context_len = prefix_cache_len
                token_chunk_size = seq_len - context_len
            elif seq_len <= prefix_cache_len:
                # Full hit. Only compute the last token to avoid
                # erroneous behavior. FIXME: Ideally we should directly
                # mark all tokens as computed in the scheduler and do not
                # schedule this sequence, so this case should not happen.
                context_len = seq_len - 1
                token_chunk_size = 1

        tokens = seq_data.get_token_ids()
        tokens = tokens[context_len:seq_len]
        token_positions = range(context_len, seq_len)
328
        token_types = seq_group_metadata.token_type_ids
329
330
331
332
333
334
335
336
337
338
339
340
341

        # For encoder-only models, the block_table is None,
        # and there is no need to initialize the slot_mapping.
        if block_table is not None:
            slot_mapping = [_PAD_SLOT_ID] * len(token_positions)
            for i, pos in enumerate(token_positions):
                block_number = block_table[pos // block_size]
                block_offset = pos % block_size
                slot = block_number * block_size + block_offset
                slot_mapping[i] = slot
            data.slot_mapping.extend(slot_mapping)

        # The MROPE positions are prepared in _compute_multi_modal_input
342
        data.input_positions.extend(token_positions)
343

344
345
346
        if data.token_type_ids is not None:
            data.token_type_ids.extend(token_types if token_types else [])

347
348
349
350
351
352
353
354
355
356
357
358
359
360
        # Update fields
        data.input_tokens.extend(tokens)
        data.num_prefills += 1
        data.num_prefill_tokens += len(tokens)
        data.query_lens.append(len(tokens))
        data.prefill_block_tables.append(block_table)
        data.seq_lens.append(seq_len)

    def _compute_multi_modal_input(self,
                                   seq_group_metadata: SequenceGroupMetadata,
                                   seq_data: SequenceData):
        computed_len = seq_data.get_num_computed_tokens()
        seq_len = self.input_data.seq_lens[-1]

361
        # NOTE: mm_kwargs only includes the subset of multi-modal items that
362
        # intersect with the current prefill positions.
363
        mm_kwargs, placeholder_maps = MultiModalPlaceholderMap.from_seq_group(
364
            seq_group_metadata, range(computed_len, seq_len))
365

366
        if not mm_kwargs:
367
            return
368

369
        # special processing for mrope position deltas.
370
        if self.runner.model_config.uses_mrope:
371
372
373
            assert not self.chunked_prefill, \
                "MROPE on CPU does not support chunked-prefill."

374
375
            image_grid_thw = mm_kwargs.get("image_grid_thw", None)
            video_grid_thw = mm_kwargs.get("video_grid_thw", None)
376
377
378
379
380
381
382
383
            audio_feature_lengths = mm_kwargs.get("audio_feature_lengths",
                                                  None)
            assert (
                image_grid_thw is not None or video_grid_thw is not None
                or audio_feature_lengths is not None), (
                    "mrope embedding type requires multi-modal input mapper "
                    "returns 'image_grid_thw' or 'video_grid_thw' or "
                    "'audio_feature_lengths'.")
384

Roger Wang's avatar
Roger Wang committed
385
            second_per_grid_ts = mm_kwargs.get("second_per_grid_ts", None)
386
            use_audio_in_video = mm_kwargs.get("use_audio_in_video", False)
387
388
389
390
391
392
            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,
Roger Wang's avatar
Roger Wang committed
393
                    hf_config=hf_config,
394
395
                    image_grid_thw=image_grid_thw,
                    video_grid_thw=video_grid_thw,
Roger Wang's avatar
Roger Wang committed
396
                    second_per_grid_ts=second_per_grid_ts,
397
                    context_len=computed_len,
398
399
                    audio_feature_lengths=audio_feature_lengths,
                    use_audio_in_video=use_audio_in_video,
400
401
402
                )
            seq_data.mrope_position_delta = mrope_position_delta

403
404
405
            for i in range(3):
                self.input_data.input_mrope_positions[  # type: ignore
                    i].extend(mrope_positions[i])
406

407
408
409
410
        self.input_data.multi_modal_inputs_list.append(mm_kwargs)
        for modality, placeholder_map in placeholder_maps.items():
            self.input_data.multi_modal_placeholder_maps[modality].extend(
                placeholder_map)
411

412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
    def _prepare_lora_input(
            self, seq_group_metadata_list: List[SequenceGroupMetadata],
            is_prefill: bool) -> LoRAMapping:
        index_mapping = []
        prompt_mapping = []
        for seq in seq_group_metadata_list:
            lora_id = seq.lora_int_id
            query_len = seq.token_chunk_size

            index_mapping += [lora_id] * query_len
            prompt_mapping += [lora_id] * (
                query_len if seq.sampling_params
                and seq.sampling_params.prompt_logprobs is not None else 1)

        return LoRAMapping(index_mapping=tuple(index_mapping),
                           prompt_mapping=tuple(prompt_mapping),
                           is_prefill=is_prefill)

430

431
432
433
434
435
436
class CPUModelRunnerBase(ModelRunnerBase[TModelInputForCPU]):
    """
    Helper class for shared methods between CPU model runners.
    """
    _model_input_cls: Type[TModelInputForCPU]
    _builder_cls: Type[ModelInputForCPUBuilder]
437
    builder: ModelInputForCPUBuilder
438
439
440

    def __init__(
        self,
441
        vllm_config: VllmConfig,
442
443
        kv_cache_dtype: Optional[str] = "auto",
        is_driver_worker: bool = False,
444
        return_hidden_states: bool = False,
445
446
447
        *args,
        **kwargs,
    ):
448
449
450
451
        ModelRunnerBase.__init__(self, vllm_config)
        model_config = self.model_config
        cache_config = self.cache_config

452
        self.is_driver_worker = is_driver_worker
453
        self.return_hidden_states = return_hidden_states
454
455

        self.device = self.device_config.device
456
        self.pin_memory = False
457
458
459
460

        self.kv_cache_dtype = kv_cache_dtype
        self.sliding_window = model_config.get_sliding_window()
        self.block_size = cache_config.block_size
461
462
463
464
        num_attn_heads = self.model_config.get_num_attention_heads(
            self.parallel_config)
        needs_attn_backend = (num_attn_heads != 0
                              or self.model_config.is_attention_free)
465
466
467
468
469
        self.attn_backend = get_attn_backend(
            self.model_config.get_head_size(),
            self.model_config.dtype,
            self.kv_cache_dtype,
            self.block_size,
470
            self.model_config.is_attention_free,
Thien Tran's avatar
Thien Tran committed
471
            use_mla=self.model_config.use_mla,
472
        ) if needs_attn_backend else None
473
474
475

        # Lazy initialization.
        self.model: nn.Module  # Set after init_Model
476
477
        # Set after load_model.
        self.lora_manager: Optional[LRUCacheWorkerLoRAManager] = None
478
        self.sampler = get_sampler()
479

480
481
482
483
        if hasattr(self, "_builder_cls"):
            # multi-step model runner does not have `_builder_cls`
            self.builder = self._builder_cls(weakref.proxy(self))

484
    def load_model(self) -> None:
485
        self.model = get_model(vllm_config=self.vllm_config)
486

487
488
489
490
491
492
493
494
495
        if self.lora_config:
            assert supports_lora(
                self.model
            ), f"{self.model.__class__.__name__} does not support LoRA yet."

            if supports_multimodal(self.model):
                logger.warning("Regarding multimodal models, vLLM currently "
                               "only supports adding LoRA to language model.")

496
497
            # Use get_text_config() in case of multimodal models
            text_config = self.model_config.hf_config.get_text_config()
498
499
500
501
502
503
504
505
506

            self.lora_manager = LRUCacheWorkerLoRAManager(
                self.scheduler_config.max_num_seqs,
                self.scheduler_config.max_num_batched_tokens,
                self.vocab_size,
                self.lora_config,
                self.device,
                self.model.embedding_modules,
                self.model.embedding_padding_modules,
507
                max_position_embeddings=text_config.max_position_embeddings,
508
509
510
            )
            self.model = self.lora_manager.create_lora_manager(self.model)

511
512
513
    def get_model(self) -> nn.Module:
        return self.model

514
515
516
517
    def _prepare_model_input_tensors(
        self,
        seq_group_metadata_list: List[SequenceGroupMetadata],
        finished_requests_ids: Optional[List[str]] = None
518
    ) -> TModelInputForCPU:
519
520
521
522
523
        """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.

        """
524
525
        self.builder.prepare(finished_requests_ids)
        self.builder.set_seq_group_list(seq_group_metadata_list)
526

527
        return self.builder.build()  # type: ignore
528

529
530
531
532
    @property
    def vocab_size(self) -> int:
        return self.model_config.get_vocab_size()

533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
    def remove_all_loras(self):
        if not self.lora_manager:
            raise RuntimeError("LoRA is not enabled.")
        self.lora_manager.remove_all_adapters()

    def set_active_loras(self, lora_requests: Set[LoRARequest],
                         lora_mapping: LoRAMapping) -> None:
        if not self.lora_manager:
            raise RuntimeError("LoRA is not enabled.")
        self.lora_manager.set_active_adapters(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_adapter(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_adapter(lora_id)

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

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

564
565
566
567
568
569
570
571
572
573
574
575
576
577
578

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

    def make_model_input_from_broadcasted_tensor_dict(
        self,
        tensor_dict: Dict[str, Any],
    ) -> ModelInputForCPUWithSamplingMetadata:
        return ModelInputForCPUWithSamplingMetadata.from_broadcasted_tensor_dict(  # noqa: E501
            tensor_dict,
            attn_backend=self.attn_backend,
        )

579
    def prepare_model_input(
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
        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)

600
601
        is_prompt = (seq_group_metadata_list[0].is_prompt
                     if seq_group_metadata_list else None)
602
603
        return dataclasses.replace(model_input,
                                   sampling_metadata=sampling_metadata,
604
605
                                   virtual_engine=virtual_engine,
                                   is_prompt=is_prompt)
606

607
    @torch.no_grad()
608
609
    def execute_model(
        self,
610
        model_input: ModelInputForCPUWithSamplingMetadata,
611
        kv_caches: List[torch.Tensor],
612
        intermediate_tensors: Optional[IntermediateTensors] = None,
613
        num_steps: int = 1,
614
        previous_hidden_states: Optional[torch.Tensor] = None,
615
616
617
618
619
    ) -> Optional[List[SamplerOutput]]:
        if num_steps > 1:
            raise ValueError(
                "CPU worker does not support multi-step execution.")

620
621
622
623
624
625
        if self.lora_config:
            assert model_input.lora_requests is not None
            assert model_input.lora_mapping is not None
            self.set_active_loras(model_input.lora_requests,
                                  model_input.lora_mapping)

626
        model_executable = self.model
627

628
629
630
631
        multimodal_kwargs = {}
        if model_input.multi_modal_kwargs is not None:
            multimodal_kwargs = MultiModalKwargs.as_kwargs(
                model_input.multi_modal_kwargs, device=self.device)
632
633
634
635
        execute_model_kwargs = {}
        if previous_hidden_states is not None:
            execute_model_kwargs.update(
                {"previous_hidden_states": previous_hidden_states})
636

637
638
        with set_forward_context(model_input.attn_metadata, self.vllm_config,
                                 model_input.virtual_engine):
youkaichao's avatar
youkaichao committed
639
640
641
642
            hidden_states = model_executable(
                input_ids=model_input.input_tokens,
                positions=model_input.input_positions,
                intermediate_tensors=intermediate_tensors,
643
                **execute_model_kwargs,
youkaichao's avatar
youkaichao committed
644
645
                **multimodal_kwargs,
            )
646
647

        # Compute the logits.
648
649
        logits = self.model.compute_logits(hidden_states,
                                           model_input.sampling_metadata)
650
651

        # Only perform sampling in the driver worker.
652
        if not self.is_driver_worker:
653
            return []
654
655

        # Sample the next token.
656
        output = self.sampler(
657
            logits=logits,
658
            sampling_metadata=model_input.sampling_metadata,
659
        )
660
661
662
663
664
        if self.return_hidden_states:
            # we only need to pass hidden states of most recent token
            if model_input.is_prompt:
                output.prefill_hidden_states = hidden_states
            output.hidden_states = hidden_states
665
        return [output]
666
667
668

    def generate_proposals(self, *args, **kwargs):
        return self.model.generate_proposals(*args, **kwargs)