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

import asyncio
5
import io
6
from abc import ABC, abstractmethod
7
from dataclasses import dataclass
8
from typing import Annotated
9

10
11
import pybase64
import torch
12
13
14
from pydantic import Field

from vllm.config import ModelConfig
15
from vllm.inputs.data import EmbedsPrompt as EngineEmbedsPrompt
16
from vllm.inputs.data import TextPrompt as EngineTextPrompt
17
from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
18
from vllm.inputs.parse import get_prompt_components, parse_raw_prompts
19
from vllm.transformers_utils.tokenizer import AnyTokenizer
20
from vllm.utils.asyncio import AsyncMicrobatchTokenizer
21
22


23
24
25
26
@dataclass(frozen=True)
class RenderConfig:
    """Configuration to control how prompts are prepared."""

27
    max_length: int | None = None
28
29
30
    """Maximum allowable total input token length. If provided,
    token inputs longer than this raise ``ValueError``."""

31
    truncate_prompt_tokens: int | None = None
32
33
34
35
    """Number of tokens to keep. ``None`` means no truncation.
    ``0`` yields an empty list (and skips embeds).
    ``-1`` maps to ``model_config.max_model_len``."""

36
    add_special_tokens: bool | None = True
37
38
    """Whether to add model-specific special tokens during tokenization."""

39
    cache_salt: str | None = None
40
41
    """String to disambiguate prefix cache entries."""

42
    needs_detokenization: bool | None = False
43
44
    """If True, detokenize IDs back to text for inclusion in outputs."""

45
    def verify_truncate_prompt_tokens(self, model_config: ModelConfig) -> int | None:
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
        """Validate and normalize `truncate_prompt_tokens` parameter."""
        truncate_prompt_tokens = self.truncate_prompt_tokens
        if truncate_prompt_tokens is None:
            return None

        if truncate_prompt_tokens == 0:
            return 0

        if truncate_prompt_tokens < 0:
            truncate_prompt_tokens = model_config.max_model_len

        max_length = self.max_length
        if max_length is not None and truncate_prompt_tokens > max_length:  # type: ignore[operator]
            raise ValueError(
                f"{truncate_prompt_tokens=} cannot be greater than "
61
62
                f"{max_length=}. Please select a smaller truncation size."
            )
63
64
65

        return truncate_prompt_tokens

66

67
68
69
class BaseRenderer(ABC):
    """
    Base class for unified input processing and rendering.
70

71
72
73
74
75
    The Renderer serves as a unified input processor that consolidates
    tokenization, chat template formatting, and multimodal input handling
    into a single component.
    It converts high-level API requests (OpenAI-style JSON) into token IDs and
    multimodal features ready for engine consumption.
76

77
78
79
80
81
82
83
84
85
86
87
    Key responsibilities:
    - Convert text prompts to token sequences with proper special tokens
    - Apply chat templates and format conversations
    - Handle multimodal inputs (images, audio, etc.) when applicable
    - Manage prompt truncation and length validation
    - Provide clean separation between API layer and engine core
    """

    def __init__(
        self,
        model_config: ModelConfig,
88
        tokenizer: AnyTokenizer | None = None,
89
90
91
92
93
94
95
96
    ):
        super().__init__()
        self.model_config = model_config
        self.tokenizer = tokenizer

    @abstractmethod
    async def render_prompt(
        self,
97
        *,
98
        prompt_or_prompts: str | list[str] | list[int] | list[list[int]],
99
        config: RenderConfig,
100
101
    ) -> list[EngineTokensPrompt]:
        """
102
103
104
105
106
107
108
109
110
111
112
113
        Convert text or token inputs into engine-ready TokensPrompt objects.

        This method accepts text or token inputs and produces a
        list of [`TokensPrompt`][vllm.inputs.data.TokensPrompt] objects
        for the engine.

        Args:
            prompt_or_prompts: One of:
                - ``str``: Single text prompt.
                - ``list[str]``: Batch of text prompts.
                - ``list[int]``: Single pre-tokenized sequence.
                - ``list[list[int]]``: Batch of pre-tokenized sequences.
114
            config: Render configuration controlling how prompts are prepared
115
                (e.g., tokenization and length handling).
116
117
118
119
120
121
122
123
124
125
126
127

        Returns:
            list[EngineTokensPrompt]: Engine-ready token prompts.

        Raises:
            ValueError: If input formats are invalid or length limits exceeded.
        """
        raise NotImplementedError

    @abstractmethod
    async def render_prompt_and_embeds(
        self,
128
        *,
129
130
        prompt_or_prompts: str | list[str] | list[int] | list[list[int]] | None = None,
        prompt_embeds: bytes | list[bytes] | None = None,
131
        config: RenderConfig,
132
    ) -> list[EngineTokensPrompt | EngineEmbedsPrompt]:
133
134
        """
        Convert text/token and/or base64-encoded embeddings inputs into
135
        engine-ready prompt objects using a unified RenderConfig.
136
137
138
139
140

        At least one of ``prompt_or_prompts`` or ``prompt_embeds`` must be
        provided and non-empty. If both are omitted or empty (e.g., empty
        string and empty list), a ``ValueError`` is raised.

141
        Args:
142
143
144
            prompt_or_prompts: Text or token inputs to include.
            prompt_embeds: Base64-encoded bytes (or list thereof) containing a
                torch-saved tensor to be used as prompt embeddings.
145
            config: Render configuration controlling how prompts are prepared
146
                (e.g., tokenization and length handling).
147

148
        Returns:
149
150
151
            list[Union[EngineTokensPrompt, EngineEmbedsPrompt]]:
                Engine-ready prompt objects.

152
        Raises:
153
154
155
            ValueError: If both ``prompt_or_prompts`` and ``prompt_embeds``
                are omitted or empty (decoder prompt cannot be empty), or if
                length limits are exceeded.
156
157
158
        """
        raise NotImplementedError

159
160
161
    @classmethod
    def load_prompt_embeds(
        cls,
162
163
164
        prompt_embeds: bytes | list[bytes],
        truncate_prompt_tokens: Annotated[int, Field(ge=0)] | None = None,
        cache_salt: str | None = None,
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
    ) -> list[EngineEmbedsPrompt]:
        """Load and validate base64-encoded embeddings into prompt objects."""

        def _load_and_validate_embed(embed: bytes) -> EngineEmbedsPrompt:
            tensor = torch.load(
                io.BytesIO(pybase64.b64decode(embed, validate=True)),
                weights_only=True,
                map_location=torch.device("cpu"),
            )
            assert isinstance(tensor, torch.Tensor) and tensor.dtype in (
                torch.float32,
                torch.bfloat16,
                torch.float16,
            )
            tensor = tensor.to_dense()
            if tensor.dim() > 2:
                tensor = tensor.squeeze(0)
                assert tensor.dim() == 2
            if truncate_prompt_tokens is not None:
                tensor = tensor[-truncate_prompt_tokens:]
            embeds_prompt = EngineEmbedsPrompt(prompt_embeds=tensor)
            if cache_salt is not None:
                embeds_prompt["cache_salt"] = cache_salt
            return embeds_prompt

        if isinstance(prompt_embeds, list):
            return [_load_and_validate_embed(embed) for embed in prompt_embeds]
192
193

        return [_load_and_validate_embed(prompt_embeds)]
194

195
196
197
198
199

class CompletionRenderer(BaseRenderer):
    def __init__(
        self,
        model_config: ModelConfig,
200
201
202
        tokenizer: AnyTokenizer | None = None,
        async_tokenizer_pool: dict[AnyTokenizer, AsyncMicrobatchTokenizer]
        | None = None,
203
204
    ):
        super().__init__(model_config, tokenizer)
205
        self.async_tokenizer_pool = async_tokenizer_pool
206
        self.async_tokenizer: AsyncMicrobatchTokenizer | None = None
207
208
209

    async def render_prompt(
        self,
210
        *,
211
        prompt_or_prompts: str | list[str] | list[int] | list[list[int]],
212
        config: RenderConfig,
213
214
    ) -> list[EngineTokensPrompt]:
        """Implementation of prompt rendering for completion-style requests.
215

216
217
218
        Uses async tokenizer pooling for improved performance. See base class
        for detailed parameter documentation.
        """
219
        truncate_prompt_tokens = config.verify_truncate_prompt_tokens(self.model_config)
220
221
        if truncate_prompt_tokens == 0:
            return []
222

223
224
225
226
227
228
229
230
        tasks = (
            self._create_prompt(
                prompt_input,
                config=config,
                truncate_prompt_tokens=truncate_prompt_tokens,
            )
            for prompt_input in parse_raw_prompts(prompt_or_prompts)
        )
231
232

        return await asyncio.gather(*tasks)
233
234
235

    async def render_prompt_and_embeds(
        self,
236
        *,
237
238
        prompt_or_prompts: str | list[str] | list[int] | list[list[int]] | None = None,
        prompt_embeds: bytes | list[bytes] | None = None,
239
        config: RenderConfig,
240
    ) -> list[EngineTokensPrompt | EngineEmbedsPrompt]:
241
242
243
244
        """
        Render text/token prompts and/or precomputed embedding prompts. At
        least one of `prompt_or_prompts` or `prompt_embeds` must be provided.
        """
245
        truncate_prompt_tokens = config.verify_truncate_prompt_tokens(self.model_config)
246
247
248
        if truncate_prompt_tokens == 0:
            return []

249
        rendered: list[EngineTokensPrompt | EngineEmbedsPrompt] = []
250
251
252

        if prompt_embeds is not None:
            rendered.extend(
253
254
255
256
                self.load_prompt_embeds(
                    prompt_embeds, truncate_prompt_tokens, config.cache_salt
                )
            )
257
258
259
260
261
        if prompt_or_prompts is None or prompt_or_prompts == "":
            return rendered

        token_prompts = await self.render_prompt(
            prompt_or_prompts=prompt_or_prompts,
262
            config=config,
263
264
265
266
267
        )
        rendered.extend(token_prompts)

        return rendered

268
    def _maybe_apply_truncation(
269
        self, token_ids: list[int], truncate_prompt_tokens: int | None
270
    ) -> list[int]:
271
272
273
274
275
276
277
278
        """Apply truncation to token sequence."""
        if truncate_prompt_tokens is None:
            return token_ids
        if truncate_prompt_tokens >= len(token_ids):
            return token_ids

        return token_ids[-truncate_prompt_tokens:]

279
280
    async def _create_prompt(
        self,
281
        prompt_input: EngineTextPrompt | EngineTokensPrompt,
282
        config: RenderConfig,
283
        truncate_prompt_tokens: int | None,
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
    ) -> EngineTokensPrompt:
        prompt, prompt_token_ids, _ = get_prompt_components(prompt_input)

        if prompt_token_ids is not None:
            # NOTE: detokenization is needed when echo is enabled,
            # where the input token IDs are decoded back to text.
            return await self._create_prompt_from_token_ids(
                prompt_token_ids,
                config.max_length,
                truncate_prompt_tokens,
                config.cache_salt,
                config.needs_detokenization,
            )

        if prompt is not None:
            return await self._create_prompt_from_text(
                prompt,
                config.max_length,
                truncate_prompt_tokens,
                config.add_special_tokens,
                config.cache_salt,
            )

        # TODO: Also handle embeds prompt using this method
        raise NotImplementedError

    async def _create_prompt_from_text(
311
312
        self,
        text: str,
313
314
315
316
        max_length: int | None,
        truncate_prompt_tokens: int | None,
        add_special_tokens: bool | None,
        cache_salt: str | None,
317
318
319
320
321
    ) -> EngineTokensPrompt:
        """Tokenize text input asynchronously."""
        async_tokenizer = self._get_async_tokenizer()

        # Handle encoder-specific preprocessing
322
323
324
325
        if (
            self.model_config.encoder_config is not None
            and self.model_config.encoder_config.get("do_lower_case", False)
        ):
326
327
328
329
            text = text.lower()

        # Tokenize texts
        if truncate_prompt_tokens is None:
330
            encoded = await async_tokenizer(text, add_special_tokens=add_special_tokens)
331
332
333
334
335
        else:
            encoded = await async_tokenizer(
                text,
                add_special_tokens=add_special_tokens,
                truncation=True,
336
337
                max_length=truncate_prompt_tokens,
            )
338

339
340
341
        return self._create_tokens_prompt(
            encoded.input_ids, max_length, cache_salt, text
        )
342

343
    async def _create_prompt_from_token_ids(
344
345
        self,
        token_ids: list[int],
346
347
348
349
        max_length: int | None,
        truncate_prompt_tokens: int | None,
        cache_salt: str | None,
        needs_detokenization: bool | None = False,
350
351
    ) -> EngineTokensPrompt:
        """Optionally detokenize token IDs and build a tokens prompt."""
352
        token_ids = self._maybe_apply_truncation(token_ids, truncate_prompt_tokens)
353
354

        prompt = None
355
        if needs_detokenization:
356
357
358
            async_tokenizer = self._get_async_tokenizer()
            prompt = await async_tokenizer.decode(token_ids)

359
360
361
362
363
364
        return self._create_tokens_prompt(
            token_ids=token_ids,
            max_length=max_length,
            cache_salt=cache_salt,
            prompt=prompt,
        )
365
366
367

    def _get_async_tokenizer(self) -> AsyncMicrobatchTokenizer:
        """Get or create async tokenizer using shared pool."""
368
369
370
371
372
        async_tokenizer = self.async_tokenizer
        if async_tokenizer is not None:
            return async_tokenizer

        tokenizer = self.tokenizer
373
        if self.tokenizer is None:
374
            raise ValueError("No tokenizer available for text input processing")
375

376
377
378
379
380
381
382
383
384
        if self.async_tokenizer_pool is None:
            async_tokenizer = AsyncMicrobatchTokenizer(tokenizer)
        else:
            async_tokenizer = self.async_tokenizer_pool.get(tokenizer)
            if async_tokenizer is None:
                async_tokenizer = AsyncMicrobatchTokenizer(tokenizer)
                self.async_tokenizer_pool[tokenizer] = async_tokenizer
        self.async_tokenizer = async_tokenizer
        return async_tokenizer
385
386
387
388

    def _create_tokens_prompt(
        self,
        token_ids: list[int],
389
390
391
        max_length: int | None = None,
        cache_salt: str | None = None,
        prompt: str | None = None,
392
393
394
395
    ) -> EngineTokensPrompt:
        """Create validated EngineTokensPrompt."""
        if max_length is not None and len(token_ids) > max_length:
            raise ValueError(
396
                f"This model's maximum context length is {max_length} tokens. "
397
                f"However, your request has {len(token_ids)} input tokens. "
398
399
                "Please reduce the length of the input messages."
            )
400
401
402
403

        tokens_prompt = EngineTokensPrompt(prompt_token_ids=token_ids)
        if cache_salt is not None:
            tokens_prompt["cache_salt"] = cache_salt
404
405
        if prompt is not None:
            tokens_prompt["prompt"] = prompt
406
        return tokens_prompt