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

4
import asyncio
5
from collections.abc import Mapping
6
from typing import Any, Optional, Union, cast
7
8
9
10
11
12

from typing_extensions import assert_never

from vllm.config import ModelConfig
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
13
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
14
15
from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalEncDecInputs,
                                    MultiModalInputs)
16
from vllm.transformers_utils.tokenizer import AnyTokenizer
17
from vllm.transformers_utils.tokenizer_group import TokenizerGroup
18

19
20
21
22
from .data import (DecoderOnlyInputs, EmbedsInputs, EmbedsPrompt,
                   EncoderDecoderInputs, ProcessorInputs, PromptType,
                   SingletonInputs, SingletonPrompt, TextPrompt, TokenInputs,
                   TokensPrompt, embeds_inputs, token_inputs)
23
24
25
26
27
28
29
30
31
32
from .parse import is_explicit_encoder_decoder_prompt, parse_singleton_prompt

logger = init_logger(__name__)


class InputPreprocessor:

    def __init__(
        self,
        model_config: ModelConfig,
33
        tokenizer: Optional[TokenizerGroup],
34
        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
35
36
37
38
39
    ) -> None:
        super().__init__()

        self.model_config = model_config
        self.tokenizer = tokenizer
40
        self.mm_registry = mm_registry
41

42
    def get_tokenizer_group(self) -> TokenizerGroup:
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
        if self.tokenizer is None:
            raise ValueError("You cannot pass text prompts when "
                             "`skip_tokenizer_init` is True")

        return self.tokenizer

    def get_bos_token_id(self,
                         lora_request: Optional[LoRARequest] = None
                         ) -> Optional[int]:
        if self.tokenizer is None:
            logger.warning("Using None for BOS token id because tokenizer "
                           "is not initialized")
            return None

        return self.tokenizer.get_lora_tokenizer(lora_request).bos_token_id

    def get_eos_token_id(self,
                         lora_request: Optional[LoRARequest] = None
                         ) -> Optional[int]:
        if self.tokenizer is None:
            logger.warning("Using None for EOS token id because tokenizer "
                           "is not initialized")
            return None

        return self.tokenizer.get_lora_tokenizer(lora_request).eos_token_id

    def get_decoder_start_token_id(self) -> Optional[int]:
70
        """
71
72
73
        Obtain the decoder start token id employed by an encoder/decoder
        model. Returns None for non-encoder/decoder models or if the
        model config is unavailable.
74
        """
75

76
        if not self.model_config.is_encoder_decoder:
77
78
79
            logger.warning_once(
                "Using None for decoder start token id because "
                "this is not an encoder/decoder model.")
80
81
            return None

82
        if self.model_config is None or self.model_config.hf_config is None:
83
84
85
            logger.warning_once(
                "Using None for decoder start token id because "
                "model config is not available.")
86
87
88
            return None

        dec_start_token_id = getattr(self.model_config.hf_config,
89
                                     "decoder_start_token_id", None)
90
        if dec_start_token_id is None:
91
92
93
94
            logger.warning_once(
                "Falling back on <BOS> for decoder start token "
                "id because decoder start token id is not "
                "available.")
95
96
97
98
            dec_start_token_id = self.get_bos_token_id()

        return dec_start_token_id

99
    def _get_default_enc_dec_decoder_prompt(self) -> list[int]:
100
        """
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
        Specifically for encoder/decoder models:
        generate a default decoder prompt for when
        the user specifies only the encoder prompt.

        Encoder/decoder models utilize the decoder
        prompt in different ways; as new models are
        added, it is intended that this function
        will be extended to produce differing
        default decoder prompts, depending on the
        model variety.

        Absent a special case, the default behavior
        of this method is to mirror the behavior of
        the HuggingFace (HF) GenerationMixin for a None
        decoder prompt, which is to employ a logit processor
        setting to force the first decoded token to be <BOS>.
        Here, this behavior is approximated by having the
        "default" decoder prompt be <BOS>.

        However, it is possible that in the future
121
        other models may have different or more
122
123
124
125
126
127
128
        complex logic for the default decoder prompt.
        This motivates having a special helper method
        for default decoder prompts.

        Returns:

        * prompt_token_ids
129
        """
130
131
132
133
134
135
136

        bos_token_id = self.get_bos_token_id()
        assert bos_token_id is not None
        return [bos_token_id]

    def _prepare_decoder_input_ids_for_generation(
        self,
137
138
        decoder_input_ids: Optional[list[int]],
    ) -> list[int]:
139
140
141
        """
        Prepares `decoder_input_ids` for generation with encoder-decoder models.

142
143
144
145
        Based on:
        https://github.com/huggingface/transformers/blob/4037a2b5b1278736e566aec12e169100275545ea/src/transformers/generation/utils.py
        specifically,
        `GenerationMixin._prepare_decoder_input_ids_for_generation()`.
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163

        Arguments:

        * decoder_input_ids: input token ids to preprocess

        Returns:

        * Processed token list
        """

        decoder_start_token_id = self.get_decoder_start_token_id()
        assert decoder_start_token_id is not None

        if decoder_input_ids is None:
            # no decoder prompt input ->
            # use decoder_start_token_id as decoder_input_ids
            decoder_input_ids = self._get_default_enc_dec_decoder_prompt()

164
165
        if (len(decoder_input_ids) == 0
                or decoder_input_ids[0] != decoder_start_token_id):
166
167
168
169
            decoder_input_ids = [decoder_start_token_id] + decoder_input_ids

        return decoder_input_ids

170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
    def _get_tokenization_kw(
        self,
        overrides: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        kwargs = dict[str, Any]()

        if self.model_config.hf_config.model_type == "whisper":
            # For Whisper, special tokens should be provided by the user based
            # on the task and language of their request. Also needed to avoid
            # appending an EOS token to the prompt which disrupts generation.
            kwargs["add_special_tokens"] = False

        if overrides:
            kwargs.update(overrides)

        return kwargs

187
188
189
190
    def _tokenize_prompt(
        self,
        prompt: str,
        lora_request: Optional[LoRARequest],
191
        tokenization_kwargs: Optional[dict[str, Any]] = None,
192
    ) -> list[int]:
193
194
195
196
197
        """
        Apply the model's tokenizer to a text prompt, returning the
        corresponding token IDs.
        """
        tokenizer = self.get_tokenizer_group()
198
        tokenization_kwargs = self._get_tokenization_kw(tokenization_kwargs)
199

200
        encoder_config = self.model_config.encoder_config
201

202
        if encoder_config and encoder_config.get("do_lower_case", False):
203
204
            prompt = prompt.lower()

205
206
207
208
209
        if self.model_config.tokenizer_mode == "cpm":
                return [tokenizer.bos_id] + tokenizer.encode(prompt)
        else:
            return tokenizer.encode(prompt=prompt,
                                    lora_request=lora_request,
zhuwenwen's avatar
zhuwenwen committed
210
                                    **tokenization_kwargs)
211
212
213
214
215

    async def _tokenize_prompt_async(
        self,
        prompt: str,
        lora_request: Optional[LoRARequest],
216
        tokenization_kwargs: Optional[dict[str, Any]] = None,
217
    ) -> list[int]:
218
219
220
221
        """
        Async version of
        [`_tokenize_prompt`][vllm.inputs.preprocess.InputPreprocessor._tokenize_prompt].
        """
222
        tokenizer = self.get_tokenizer_group()
223
        tokenization_kwargs = self._get_tokenization_kw(tokenization_kwargs)
224
225
226
227

        return await tokenizer.encode_async(prompt=prompt,
                                            lora_request=lora_request,
                                            **tokenization_kwargs)
228

229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
    def _get_mm_tokenizer(
        self,
        lora_request: Optional[LoRARequest],
    ) -> AnyTokenizer:
        # PrithviGeoSpatialMAE needs to be initialized without a tokenizer
        # while using also multi-modal input
        if not self.tokenizer:
            return cast(AnyTokenizer, object())  # Dummy

        tokenizer_group = self.get_tokenizer_group()
        return tokenizer_group.get_lora_tokenizer(lora_request)

    async def _get_mm_tokenizer_async(
        self,
        lora_request: Optional[LoRARequest],
    ) -> AnyTokenizer:
        # PrithviGeoSpatialMAE needs to be initialized without a tokenizer
        # while using also multi-modal input
        if not self.tokenizer:
            return cast(AnyTokenizer, object())  # Dummy

        tokenizer_group = self.get_tokenizer_group()
        return await tokenizer_group.get_lora_tokenizer_async(lora_request)
252

253
254
    def _process_multimodal(
        self,
255
        prompt: Union[str, list[int]],
256
257
        mm_data: MultiModalDataDict,
        mm_processor_kwargs: Optional[Mapping[str, object]],
258
259
        tokenization_kwargs: Optional[dict[str, Any]] = None,
        lora_request: Optional[LoRARequest] = None,
260
        return_mm_hashes: bool = False,
261
    ) -> MultiModalInputs:
262
263
264
265
        """
        Apply the model's multi-modal processor to a multi-modal prompt,
        returning the corresponding token IDs and metadata.
        """
266
        tokenizer = self._get_mm_tokenizer(lora_request)
267

268
269
        mm_processor = self.mm_registry.create_processor(self.model_config,
                                                         tokenizer=tokenizer)
270
271
272
273

        if mm_processor_kwargs is None:
            mm_processor_kwargs = {}

274
275
276
277
278
        return mm_processor.apply(prompt,
                                  mm_data,
                                  hf_processor_mm_kwargs=mm_processor_kwargs,
                                  tokenization_kwargs=tokenization_kwargs,
                                  return_mm_hashes=return_mm_hashes)
279
280
281

    async def _process_multimodal_async(
        self,
282
        prompt: Union[str, list[int]],
283
284
        mm_data: MultiModalDataDict,
        mm_processor_kwargs: Optional[Mapping[str, object]],
285
286
        tokenization_kwargs: Optional[dict[str, Any]] = None,
        lora_request: Optional[LoRARequest] = None,
287
        return_mm_hashes: bool = False,
288
    ) -> MultiModalInputs:
289
290
291
292
        """
        Async version of
        [`_process_multimodal`][vllm.inputs.preprocess.InputPreprocessor._process_multimodal].
        """
293
        tokenizer = await self._get_mm_tokenizer_async(lora_request)
294

295
296
        mm_processor = self.mm_registry.create_processor(self.model_config,
                                                         tokenizer=tokenizer)
297
298
299
        if mm_processor_kwargs is None:
            mm_processor_kwargs = {}

300
301
302
303
304
        return mm_processor.apply(prompt,
                                  mm_data,
                                  hf_processor_mm_kwargs=mm_processor_kwargs,
                                  tokenization_kwargs=tokenization_kwargs,
                                  return_mm_hashes=return_mm_hashes)
305

306
    def _process_embeds(
307
        self,
308
309
        parsed_content: EmbedsPrompt,
    ) -> EmbedsInputs:
310
311
312
        if not self.model_config.enable_prompt_embeds:
            raise ValueError("You must set `--enable-prompt-embeds` to input "
                             "`prompt_embeds`.")
313

314
        prompt_embeds = parsed_content["prompt_embeds"]
315

316
317
318
319
320
321
        # prompt_embeds must be (seq_len, hidden_size), but if the user
        # passes in a batch of size 1, i.e. (1, seq_len, hidden_size),
        # we can unambiguously process the intent by squeezing the batch
        # dimension.
        if prompt_embeds.ndim == 3:
            prompt_embeds = prompt_embeds.squeeze(dim=0)
322

323
324
325
        if prompt_embeds.ndim != 2:
            raise ValueError(
                "prompt_embeds must be of shape (seq_len, hidden_size).")
326

327
328
        return embeds_inputs(prompt_embeds=prompt_embeds,
                             cache_salt=parsed_content.get("cache_salt"))
329

330
331
332
333
334
335
336
337
338
    async def _process_embeds_async(
        self,
        parsed_content: EmbedsPrompt,
    ) -> EmbedsInputs:
        return self._process_embeds(parsed_content)

    def _process_tokens(
        self,
        parsed_content: TokensPrompt,
339
        tokenization_kwargs: Optional[dict[str, Any]] = None,
340
341
342
343
344
345
346
347
348
349
350
351
        lora_request: Optional[LoRARequest] = None,
        return_mm_hashes: bool = False,
    ) -> Union[TokenInputs, MultiModalInputs]:
        prompt_token_ids = parsed_content["prompt_token_ids"]
        token_type_ids = parsed_content.get("token_type_ids")

        inputs: Union[TokenInputs, MultiModalInputs]
        if multi_modal_data := parsed_content.get("multi_modal_data"):
            inputs = self._process_multimodal(
                prompt_token_ids,
                multi_modal_data,
                parsed_content.get("mm_processor_kwargs"),
352
                tokenization_kwargs=tokenization_kwargs,
353
                lora_request=lora_request,
354
                return_mm_hashes=return_mm_hashes,
355
            )
356
        else:
357
            inputs = token_inputs(
358
                prompt_token_ids=prompt_token_ids,
359
                token_type_ids=token_type_ids,
360
361
            )

362
363
364
365
        if cache_salt := parsed_content.get("cache_salt"):
            inputs["cache_salt"] = cache_salt

        return inputs
366

367
368
369
    async def _process_tokens_async(
        self,
        parsed_content: TokensPrompt,
370
        tokenization_kwargs: Optional[dict[str, Any]] = None,
371
372
373
374
375
376
377
378
379
380
381
382
        lora_request: Optional[LoRARequest] = None,
        return_mm_hashes: bool = False,
    ) -> Union[TokenInputs, MultiModalInputs]:
        prompt_token_ids = parsed_content["prompt_token_ids"]
        token_type_ids = parsed_content.get("token_type_ids")

        inputs: Union[TokenInputs, MultiModalInputs]
        if multi_modal_data := parsed_content.get("multi_modal_data"):
            inputs = await self._process_multimodal_async(
                prompt_token_ids,
                multi_modal_data,
                parsed_content.get("mm_processor_kwargs"),
383
                tokenization_kwargs=tokenization_kwargs,
384
385
386
387
388
                lora_request=lora_request,
                return_mm_hashes=return_mm_hashes,
            )
        else:
            inputs = token_inputs(
389
                prompt_token_ids=prompt_token_ids,
390
                token_type_ids=token_type_ids,
391
392
            )

393
394
395
396
397
398
399
400
401
402
403
404
405
        if cache_salt := parsed_content.get("cache_salt"):
            inputs["cache_salt"] = cache_salt

        return inputs

    def _process_text(
        self,
        parsed_content: TextPrompt,
        tokenization_kwargs: Optional[dict[str, Any]] = None,
        lora_request: Optional[LoRARequest] = None,
        return_mm_hashes: bool = False,
    ) -> Union[TokenInputs, MultiModalInputs]:
        prompt_text = parsed_content["prompt"]
406

407
408
409
410
411
412
        inputs: Union[TokenInputs, MultiModalInputs]
        if multi_modal_data := parsed_content.get("multi_modal_data"):
            inputs = self._process_multimodal(
                prompt_text,
                multi_modal_data,
                parsed_content.get("mm_processor_kwargs"),
413
                tokenization_kwargs=tokenization_kwargs,
414
415
416
417
                lora_request=lora_request,
                return_mm_hashes=return_mm_hashes,
            )
        else:
418
            prompt_token_ids = self._tokenize_prompt(
419
                prompt_text,
420
                lora_request=lora_request,
421
                tokenization_kwargs=tokenization_kwargs,
422
            )
423
            inputs = token_inputs(
424
425
426
                prompt=prompt_text,
                prompt_token_ids=prompt_token_ids,
            )
427

428
429
        if cache_salt := parsed_content.get("cache_salt"):
            inputs["cache_salt"] = cache_salt
430

431
        return inputs
432

433
    async def _process_text_async(
434
        self,
435
436
        parsed_content: TextPrompt,
        tokenization_kwargs: Optional[dict[str, Any]] = None,
437
        lora_request: Optional[LoRARequest] = None,
438
        return_mm_hashes: bool = False,
439
440
    ) -> Union[TokenInputs, MultiModalInputs]:
        prompt_text = parsed_content["prompt"]
441

442
443
444
445
446
447
        inputs: Union[TokenInputs, MultiModalInputs]
        if multi_modal_data := parsed_content.get("multi_modal_data"):
            inputs = await self._process_multimodal_async(
                prompt_text,
                multi_modal_data,
                parsed_content.get("mm_processor_kwargs"),
448
                tokenization_kwargs=tokenization_kwargs,
449
450
451
452
                lora_request=lora_request,
                return_mm_hashes=return_mm_hashes,
            )
        else:
453
            prompt_token_ids = await self._tokenize_prompt_async(
454
                prompt_text,
455
                lora_request=lora_request,
456
                tokenization_kwargs=tokenization_kwargs,
457
            )
458
            inputs = token_inputs(
459
460
461
462
                prompt=prompt_text,
                prompt_token_ids=prompt_token_ids,
            )

463
464
        if cache_salt := parsed_content.get("cache_salt"):
            inputs["cache_salt"] = cache_salt
465

466
        return inputs
467
468

    def _prompt_to_llm_inputs(
469
        self,
470
        prompt: SingletonPrompt,
471
        tokenization_kwargs: Optional[dict[str, Any]] = None,
472
        lora_request: Optional[LoRARequest] = None,
473
        return_mm_hashes: bool = False,
474
    ) -> SingletonInputs:
475
476
477
        """
        Extract the singleton inputs from a prompt.

478
479
        Arguments:

480
        * prompt: single encoder or decoder input prompt
481
        * lora_request: this is only valid for decoder prompts
482
        * return_mm_hashes: whether to return multimodal hashes
483
484
485

        Returns:

486
        * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance
487
        """
488
        parsed = parse_singleton_prompt(prompt)
489
490

        if parsed["type"] == "embeds":
491
492
493
494
            return self._process_embeds(parsed["content"])
        if parsed["type"] == "tokens":
            return self._process_tokens(
                parsed["content"],
495
                lora_request=lora_request,
496
                return_mm_hashes=return_mm_hashes,
497
            )
498
499
500
501
        if parsed["type"] == "text":
            return self._process_text(
                parsed["content"],
                tokenization_kwargs=tokenization_kwargs,
502
                lora_request=lora_request,
503
504
505
506
507
                return_mm_hashes=return_mm_hashes,
            )
        if parsed["type"] == "str":
            return self._process_text(
                TextPrompt(prompt=parsed["content"]),
508
                tokenization_kwargs=tokenization_kwargs,
509
                lora_request=lora_request,
510
                return_mm_hashes=return_mm_hashes,
511
            )
512

513
514
        assert_never(parsed)

515
    async def _prompt_to_llm_inputs_async(
516
        self,
517
        prompt: SingletonPrompt,
518
        tokenization_kwargs: Optional[dict[str, Any]] = None,
519
        lora_request: Optional[LoRARequest] = None,
520
        return_mm_hashes: bool = False,
521
    ) -> SingletonInputs:
522
523
524
525
        """
        Async version of
        [`_prompt_to_llm_inputs`][vllm.inputs.preprocess.InputPreprocessor._prompt_to_llm_inputs].
        """
526
        parsed = parse_singleton_prompt(prompt)
527

528
        if parsed["type"] == "embeds":
529
530
531
532
            return await self._process_embeds_async(parsed["content"])
        if parsed["type"] == "tokens":
            return await self._process_tokens_async(
                parsed["content"],
533
534
                lora_request=lora_request,
                return_mm_hashes=return_mm_hashes,
535
            )
536
537
538
539
        if parsed["type"] == "text":
            return await self._process_text_async(
                parsed["content"],
                tokenization_kwargs=tokenization_kwargs,
540
                lora_request=lora_request,
541
542
543
544
545
                return_mm_hashes=return_mm_hashes,
            )
        if parsed["type"] == "str":
            return await self._process_text_async(
                TextPrompt(prompt=parsed["content"]),
546
                tokenization_kwargs=tokenization_kwargs,
547
548
                lora_request=lora_request,
                return_mm_hashes=return_mm_hashes,
549
            )
550

551
        assert_never(parsed)
552
553
554

    def _build_enc_dec_llm_inputs(
        self,
555
556
        encoder_inputs: SingletonInputs,
        decoder_inputs: Optional[SingletonInputs],
557
    ) -> EncoderDecoderInputs:
558
559
560
561
        if (encoder_inputs["type"] == "embeds"
                or decoder_inputs and decoder_inputs["type"] == "embeds"):
            raise ValueError("Embedding inputs are not supported for encoder-"
                             "decoder models")
562

563
564
565
566
567
        # Needed for mypy
        encoder_inputs = cast(Union[TokenInputs, MultiModalInputs],
                              encoder_inputs)
        decoder_inputs = cast(Optional[Union[TokenInputs, MultiModalInputs]],
                              decoder_inputs)
568
569

        if decoder_inputs is None:
570
571
572
573
574
575
576
577
578
            if self.model_config.hf_config.model_type == "whisper":
                # For Whisper models, the text prompt should go to the decoder.
                # If no explicit encoder/decoder inputs, then copy the prompt
                # from the encoder to the decoder. The encoder tokens are later
                # overridden by the audio features.
                dec_token_ids = encoder_inputs["prompt_token_ids"].copy()
            else:
                dec_token_ids = self._prepare_decoder_input_ids_for_generation(
                    None)
579
            decoder_inputs = token_inputs(dec_token_ids)
580
        else:
581
582
583
            if "multi_modal_data" in decoder_inputs:
                raise ValueError("Multi-modal decoder inputs of encoder-"
                                 "decoder models are not supported yet")
584
585
586
587

            dec_token_ids = self._prepare_decoder_input_ids_for_generation(
                decoder_inputs["prompt_token_ids"])
            decoder_inputs["prompt_token_ids"] = dec_token_ids
588

589
        return EncoderDecoderInputs(
590
591
            encoder=encoder_inputs,
            decoder=decoder_inputs,
592
593
        )

594
    def _split_enc_dec_mm_inputs(
595
        self,
596
        inputs: Union[SingletonInputs, MultiModalEncDecInputs],
597
        decoder_inputs_to_override: Optional[SingletonInputs] = None,
598
    ) -> tuple[SingletonInputs, SingletonInputs]:
599
600
601
602
        """
        For encoder/decoder models only:
        Separate Encoder/Decoder inputs from a MultiModalEncDecInputs
        """
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
        if (inputs["type"] == "embeds" or decoder_inputs_to_override
                and decoder_inputs_to_override["type"] == "embeds"):
            raise ValueError("Embedding inputs are not supported for encoder-"
                             "decoder models")

        # Needed for mypy
        inputs = cast(
            Union[TokenInputs, MultiModalInputs, MultiModalEncDecInputs],
            inputs,
        )
        decoder_inputs_to_override = cast(
            Optional[Union[TokenInputs, MultiModalInputs]],
            decoder_inputs_to_override,
        )

618
619
        encoder_inputs: SingletonInputs
        decoder_inputs: SingletonInputs
620
621
622
623
624
625
626

        if inputs["type"] == "multimodal":  # Multimodal data inputs
            if not ("encoder_prompt" in inputs
                    and "encoder_prompt_token_ids" in inputs):
                raise RuntimeError("You should register an encoder-decoder "
                                   "multi-modal processor for encoder-decoder "
                                   "models.")
627
            inputs = cast(MultiModalEncDecInputs, inputs)
628

629
630
631
632
            encoder_inputs = token_inputs(
                prompt=inputs["encoder_prompt"],
                prompt_token_ids=inputs["encoder_prompt_token_ids"],
            )
633

634
635
636
637
638
639
640
641
642
643
            decoder_prompt_inputs = decoder_inputs_to_override or inputs
            decoder_inputs = MultiModalInputs(
                type="multimodal",
                prompt=decoder_prompt_inputs.get("prompt", ""),
                prompt_token_ids=decoder_prompt_inputs["prompt_token_ids"],
                mm_kwargs=inputs["mm_kwargs"],
                mm_hashes=inputs["mm_hashes"],
                mm_placeholders=inputs["mm_placeholders"],
            )
            if cache_salt := inputs.get("cache_salt"):
644
645
                decoder_inputs["cache_salt"] = cache_salt

646
        elif inputs["type"] == "token":  # Text-only inputs
647
648
649
650
            encoder_inputs = token_inputs(prompt="", prompt_token_ids=[])
            decoder_inputs = decoder_inputs_to_override or inputs
        else:
            assert_never(inputs)  # type: ignore[arg-type]
651

652
653
        return encoder_inputs, decoder_inputs

654
655
    def _process_encoder_decoder_prompt(
        self,
656
        prompt: PromptType,
657
        tokenization_kwargs: Optional[dict[str, Any]] = None,
658
    ) -> EncoderDecoderInputs:
659
        """
660
        For encoder/decoder models only:
661
662
663
        Process an input prompt into an
        [`EncoderDecoderInputs`][vllm.inputs.data.EncoderDecoderInputs]
        instance.
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681

        There are two types of input prompts:
        singleton prompts which carry only the
        encoder prompt, and explicit encoder/decoder
        prompts which carry both the encoder and the
        decoder prompts as member variables.

        This function handles the following scenarios:
        * Singleton encoder prompt: extract encoder prompt
          token ids & infer default decoder prompt token ids
        * Explicit encoder/decoder prompt: extract encoder
          and decoder prompt token ids

        Note that for Explicit encoder/decoder prompts,
        each sub-prompt (encoder or decoder prompt) can
        have any possible singleton type; thus this
        method relies on helper functions to obtain
        token ids for the sub-prompts.
682

683
684
        Arguments:

685
        * prompt: an input prompt
686
687
688

        Returns:

689
690
        * [`EncoderDecoderInputs`][vllm.inputs.data.EncoderDecoderInputs]
          instance
691
        """
692
693
        encoder_inputs: SingletonInputs
        decoder_inputs: Optional[SingletonInputs]
694

695
        if is_explicit_encoder_decoder_prompt(prompt):
696
            encoder_inputs = self._prompt_to_llm_inputs(
697
698
699
                prompt["encoder_prompt"],
                tokenization_kwargs=tokenization_kwargs,
            )
700
            if (decoder_input := prompt["decoder_prompt"]) is None:
701
                decoder_inputs = None
702
            else:
703
                decoder_inputs = self._prompt_to_llm_inputs(decoder_input)
704
705
            # For multimodal model, override decoder prompt from processor
            # with explicit decoder prompt.
706
            if self.model_config.is_multimodal_model:
707
                encoder_inputs, decoder_inputs = (
708
709
                    self._split_enc_dec_mm_inputs(encoder_inputs,
                                                  decoder_inputs))
710
        else:
711
712
713
714
            inputs = self._prompt_to_llm_inputs(
                prompt,
                tokenization_kwargs=tokenization_kwargs,
            )
715
            if self.model_config.is_multimodal_model:
716
717
                # Encoder-Decoder Multimodal model
                encoder_inputs, decoder_inputs = (
718
                    self._split_enc_dec_mm_inputs(inputs))
719
720
721
            else:
                encoder_inputs = inputs
                decoder_inputs = None
722
723

        return self._build_enc_dec_llm_inputs(encoder_inputs, decoder_inputs)
724
725
726

    async def _process_encoder_decoder_prompt_async(
        self,
727
        prompt: PromptType,
728
        tokenization_kwargs: Optional[dict[str, Any]] = None,
729
    ) -> EncoderDecoderInputs:
730
731
732
733
        """
        Async version of
        [`_process_encoder_decoder_prompt`][vllm.inputs.preprocess.InputPreprocessor._process_encoder_decoder_prompt].
        """
734
735
        encoder_inputs: SingletonInputs
        decoder_inputs: Optional[SingletonInputs]
736

737
        if is_explicit_encoder_decoder_prompt(prompt):
738
            encoder_task = self._prompt_to_llm_inputs_async(
739
740
741
                prompt["encoder_prompt"],
                tokenization_kwargs=tokenization_kwargs,
            )
742

743
            if (decoder_input := prompt["decoder_prompt"]) is None:
744
745
                encoder_inputs = await encoder_task
                decoder_inputs = None
746
            else:
747
748
749
750
                decoder_task = self._prompt_to_llm_inputs_async(
                    decoder_input,
                    tokenization_kwargs=tokenization_kwargs,
                )
751

752
                encoder_inputs, decoder_inputs = await asyncio.gather(
753
                    encoder_task, decoder_task)
754
755
756

            # For multimodal model, override decoder prompt from processor
            # with explicit decoder prompt.
757
            if self.model_config.is_multimodal_model:
758
                encoder_inputs, decoder_inputs = (
759
760
                    self._split_enc_dec_mm_inputs(encoder_inputs,
                                                  decoder_inputs))
761
        else:
762
763
764
765
            inputs = await self._prompt_to_llm_inputs_async(
                prompt,
                tokenization_kwargs=tokenization_kwargs,
            )
766
            if self.model_config.is_multimodal_model:
767
768
                # Encoder-Decoder Multimodal model
                encoder_inputs, decoder_inputs = (
769
                    self._split_enc_dec_mm_inputs(inputs))
770
771
772
            else:
                encoder_inputs = inputs
                decoder_inputs = None
773
774

        return self._build_enc_dec_llm_inputs(encoder_inputs, decoder_inputs)
775
776
777

    def _build_decoder_only_llm_inputs(
        self,
778
        prompt_inputs: DecoderOnlyInputs,
779
    ) -> DecoderOnlyInputs:
780
781
782
        if "prompt_token_ids" in prompt_inputs:
            prompt_inputs = cast(Union[TokenInputs, MultiModalInputs],
                                 prompt_inputs)  # Needed for mypy
783

784
        return prompt_inputs
785
786
787

    def _process_decoder_only_prompt(
        self,
788
        prompt: SingletonPrompt,
789
        tokenization_kwargs: Optional[dict[str, Any]] = None,
790
        lora_request: Optional[LoRARequest] = None,
791
        return_mm_hashes: bool = False,
792
    ) -> DecoderOnlyInputs:
793
        """
794
        For decoder-only models:
795
796
        Process an input prompt into a
        [`DecoderOnlyInputs`][vllm.inputs.data.DecoderOnlyInputs] instance.
797
798
799

        Arguments:

800
        * prompt: input prompt
801
        * lora_request
802
        * return_mm_hashes
803
804
805

        Returns:

806
        * [`DecoderOnlyInputs`][vllm.inputs.data.DecoderOnlyInputs] instance
807
        """
808

809
        prompt_comps = self._prompt_to_llm_inputs(
810
            prompt,
811
            tokenization_kwargs=tokenization_kwargs,
812
            lora_request=lora_request,
813
            return_mm_hashes=return_mm_hashes,
814
815
        )

816
        return self._build_decoder_only_llm_inputs(prompt_comps)
817
818
819

    async def _process_decoder_only_prompt_async(
        self,
820
        prompt: SingletonPrompt,
821
        tokenization_kwargs: Optional[dict[str, Any]] = None,
822
        lora_request: Optional[LoRARequest] = None,
823
        return_mm_hashes: bool = False,
824
    ) -> DecoderOnlyInputs:
825
826
827
828
        """
        Async version of
        [`_process_decoder_only_prompt`][vllm.inputs.preprocess.InputPreprocessor._process_decoder_only_prompt].
        """
829
        prompt_comps = await self._prompt_to_llm_inputs_async(
830
            prompt,
831
            tokenization_kwargs=tokenization_kwargs,
832
            lora_request=lora_request,
833
            return_mm_hashes=return_mm_hashes,
834
835
        )

836
        return self._build_decoder_only_llm_inputs(prompt_comps)
837
838
839

    def preprocess(
        self,
840
        prompt: PromptType,
841
        tokenization_kwargs: Optional[dict[str, Any]] = None,
842
        lora_request: Optional[LoRARequest] = None,
843
        return_mm_hashes: bool = False,
844
    ) -> ProcessorInputs:
845
        """Preprocess the input prompt."""
846
        if self.model_config.is_encoder_decoder:
847
848
849
            assert not return_mm_hashes, (
                "Multimodal hashes for encoder-decoder models should not be ",
                "returned until they are supported on vLLM V1.")
850
851
            # Encoder-decoder model requires special mapping of
            # input prompts to encoder & decoder
852
853
            return self._process_encoder_decoder_prompt(
                prompt, tokenization_kwargs)
854

855
        if is_explicit_encoder_decoder_prompt(prompt):
856
857
858
859
860
            raise ValueError("Cannot pass encoder-decoder prompt "
                             "to decoder-only models")

        # Decoder-only operation
        return self._process_decoder_only_prompt(
861
            prompt,
862
            tokenization_kwargs=tokenization_kwargs,
863
            lora_request=lora_request,
864
            return_mm_hashes=return_mm_hashes,
865
866
867
868
        )

    async def preprocess_async(
        self,
869
        prompt: PromptType,
870
        tokenization_kwargs: Optional[dict[str, Any]] = None,
871
        lora_request: Optional[LoRARequest] = None,
872
        return_mm_hashes: bool = False,
873
    ) -> ProcessorInputs:
874
875
876
877
        """
        Async version of
        [`preprocess`][vllm.inputs.preprocess.InputPreprocessor.preprocess].
        """
878
        if self.model_config.is_encoder_decoder:
879
880
881
            assert not return_mm_hashes, (
                "Multimodal hashes for encoder-decoder models should not be ",
                "returned until they are supported on vLLM V1.")
882
883
            # Encoder-decoder model requires special mapping of
            # input prompts to encoder & decoder
884
            return await self._process_encoder_decoder_prompt_async(prompt)
885

886
        if is_explicit_encoder_decoder_prompt(prompt):
887
888
889
890
891
            raise ValueError("Cannot pass encoder-decoder prompt "
                             "to decoder-only models")

        # Decoder-only operation
        return await self._process_decoder_only_prompt_async(
892
            prompt,
893
            tokenization_kwargs=tokenization_kwargs,
894
            lora_request=lora_request,
895
            return_mm_hashes=return_mm_hashes,
896
        )