| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import asyncio |
| 6 | import copy |
| 7 | import inspect |
| 8 | import itertools |
| 9 | import weakref |
| 10 | from collections import defaultdict, deque |
| 11 | from collections.abc import Sequence |
| 12 | from functools import lru_cache |
| 13 | from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload |
| 14 | |
| 15 | import jinja2 |
| 16 | import jinja2.ext |
| 17 | import jinja2.meta |
| 18 | import jinja2.nodes |
| 19 | import jinja2.parser |
| 20 | import jinja2.sandbox |
| 21 | import torch |
| 22 | from typing_extensions import override |
| 23 | |
| 24 | from vllm import envs |
| 25 | from vllm.entrypoints.chat_utils import ( |
| 26 | PROMPT_EMBEDS_PLACEHOLDER_TOKEN, |
| 27 | ChatTemplateResolutionError, |
| 28 | load_chat_template, |
| 29 | parse_chat_messages, |
| 30 | parse_chat_messages_async, |
| 31 | ) |
| 32 | from vllm.exceptions import VLLMValidationError |
| 33 | from vllm.inputs import EmbedsPrompt |
| 34 | from vllm.inputs.engine import MultiModalInput |
| 35 | from vllm.logger import init_logger |
| 36 | from vllm.multimodal.hasher import MultiModalHasher |
| 37 | from vllm.multimodal.inputs import ( |
| 38 | MultiModalFieldElem, |
| 39 | MultiModalKwargsItem, |
| 40 | MultiModalKwargsItems, |
| 41 | MultiModalSharedField, |
| 42 | PlaceholderRange, |
| 43 | ) |
| 44 | from vllm.multimodal.processing.processor import ( |
| 45 | PromptReplacement, |
| 46 | apply_token_matches, |
| 47 | find_mm_placeholders, |
| 48 | ) |
| 49 | from vllm.tokenizers.hf import HfTokenizer, maybe_make_thread_pool |
| 50 | from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path |
| 51 | from vllm.transformers_utils.processor import cached_get_processor |
| 52 | from vllm.utils.async_utils import make_async |
| 53 | from vllm.utils.func_utils import supports_kw |
| 54 | |
| 55 | from .base import BaseRenderer |
| 56 | from .inputs.preprocess import parse_dec_only_prompt |
| 57 | |
| 58 | if TYPE_CHECKING: |
| 59 | from collections.abc import Set |
| 60 | |
| 61 | from vllm.config import ModelConfig, VllmConfig |
| 62 | from vllm.entrypoints.chat_utils import ( |
| 63 | ChatCompletionMessageParam, |
| 64 | ChatTemplateContentFormat, |
| 65 | ChatTemplateContentFormatOption, |
| 66 | ConversationMessage, |
| 67 | ) |
| 68 | from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict, TokensPrompt |
| 69 | from vllm.inputs.engine import TokensInput |
| 70 | from vllm.multimodal.processing.processor import ( |
| 71 | MultiModalPromptUpdates, |
| 72 | ResolvedPromptUpdate, |
| 73 | ) |
| 74 | |
| 75 | from .inputs import DictPrompt |
| 76 | from .params import ChatParams |
| 77 | |
| 78 | logger = init_logger(__name__) |
| 79 | |
| 80 | |
| 81 | # Cache of `tokenizer -> prompt_embeds placeholder token ID`. Keyed by the |
| 82 | # tokenizer object (not `id(tokenizer)`) so a fresh tokenizer landing at a |
| 83 | # recycled memory address can't pick up a stale tid. Entries evict atomically |
| 84 | # with the tokenizer's garbage-collection. |
| 85 | _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE: Final[ |
| 86 | weakref.WeakKeyDictionary[HfTokenizer, int] |
| 87 | ] = weakref.WeakKeyDictionary() |
| 88 | _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_ERROR: Final[str] = ( |
| 89 | "Expected {token!r} to tokenize to exactly 1 token, got {num_ids} ({ids!r})." |
| 90 | ) |
| 91 | _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR: Final[str] = ( |
| 92 | "Expected {expected} prompt_embeds placeholder spans in the " |
| 93 | "tokenized prompt, found {actual}." |
| 94 | ) |
| 95 | _MISSING_PROMPT_TOKEN_IDS_ERROR: Final[str] = ( |
| 96 | "Expected prompt_token_ids in rendered prompt when prompt_embeds " |
| 97 | "are present. This indicates the chat template was invoked with " |
| 98 | "tokenize=False." |
| 99 | ) |
| 100 | _TOKENIZE_OVERRIDE_WARNING: Final[str] = ( |
| 101 | "Overriding `tokenize=False` to `True` because `prompt_embeds` " |
| 102 | "post-processing requires tokenized IDs." |
| 103 | ) |
| 104 | |
| 105 | |
| 106 | def _ensure_prompt_embeds_placeholder_token(tokenizer: HfTokenizer) -> int: |
| 107 | """Register `PROMPT_EMBEDS_PLACEHOLDER_TOKEN` as a special token and return |
| 108 | its token ID.""" |
| 109 | cached = _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE.get(tokenizer) |
| 110 | if cached is not None: |
| 111 | return cached |
| 112 | |
| 113 | tokenizer.add_special_tokens( |
| 114 | {"additional_special_tokens": [PROMPT_EMBEDS_PLACEHOLDER_TOKEN]} |
| 115 | ) |
| 116 | |
| 117 | ids = tokenizer.encode(PROMPT_EMBEDS_PLACEHOLDER_TOKEN, add_special_tokens=False) |
| 118 | if len(ids) != 1: |
| 119 | raise RuntimeError( |
| 120 | _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_ERROR.format( |
| 121 | token=PROMPT_EMBEDS_PLACEHOLDER_TOKEN, |
| 122 | num_ids=len(ids), |
| 123 | ids=ids, |
| 124 | ) |
| 125 | ) |
| 126 | |
| 127 | token_id = ids[0] |
| 128 | _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE[tokenizer] = token_id |
| 129 | return token_id |
| 130 | |
| 131 | |
| 132 | def _build_prompt_embeds_updates( |
| 133 | prompt_embeds_tensors: Sequence[torch.Tensor], |
| 134 | placeholder_token_id: int, |
| 135 | ) -> MultiModalPromptUpdates: |
| 136 | """Build `MultiModalPromptUpdates` for `prompt_embeds` expansion. |
| 137 | |
| 138 | Each tensor produces a `PromptReplacement` that maps |
| 139 | `[placeholder_token_id]` -> `[placeholder_token_id] x N` |
| 140 | (where `N = tensor.shape[0]`). |
| 141 | """ |
| 142 | updates: list[Sequence[ResolvedPromptUpdate]] = [] |
| 143 | for i, tensor in enumerate(prompt_embeds_tensors): |
| 144 | update = PromptReplacement( |
| 145 | modality="prompt_embeds", |
| 146 | target=[placeholder_token_id], |
| 147 | replacement=[placeholder_token_id] * tensor.shape[0], |
| 148 | ) |
| 149 | updates.append([update.resolve(item_idx=i)]) |
| 150 | return {"prompt_embeds": updates} |
| 151 | |
| 152 | |
| 153 | def _expand_prompt_embeds_placeholders( |
| 154 | token_ids: list[int], |
| 155 | mm_prompt_updates: MultiModalPromptUpdates, |
| 156 | ) -> list[int]: |
| 157 | """Expand each 1-token `prompt_embeds` sentinel into an N-token span. |
| 158 | |
| 159 | Uses `apply_token_matches`. Each single placeholder token in |
| 160 | `token_ids` is replaced with a consecutive span of |
| 161 | `tensor.shape[0]` copies, following tensors in order. |
| 162 | """ |
| 163 | expanded, _ = apply_token_matches(token_ids, mm_prompt_updates) |
| 164 | return expanded |
| 165 | |
| 166 | |
| 167 | def _build_prompt_embeds_positions( |
| 168 | token_ids: list[int], |
| 169 | num_tensors: int, |
| 170 | mm_prompt_updates: MultiModalPromptUpdates, |
| 171 | ) -> list[tuple[int, int]]: |
| 172 | """Locate each prompt_embeds placeholder span in `token_ids`. |
| 173 | |
| 174 | Expects `token_ids` to already contain expanded N-token spans. |
| 175 | Returns `[(start_idx, length), ...]` aligned with the tensors. |
| 176 | """ |
| 177 | placeholders = find_mm_placeholders(token_ids, mm_prompt_updates) |
| 178 | features = placeholders.get("prompt_embeds", []) |
| 179 | |
| 180 | if len(features) != num_tensors: |
| 181 | raise ValueError( |
| 182 | _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR.format( |
| 183 | expected=num_tensors, |
| 184 | actual=len(features), |
| 185 | ) |
| 186 | ) |
| 187 | |
| 188 | return [(f.start_idx, f.length) for f in features] |
| 189 | |
| 190 | |
| 191 | def _build_mixed_prompt_embeds( |
| 192 | token_ids: list[int], |
| 193 | prompt_embeds_tensors: Sequence[torch.Tensor], |
| 194 | positions: list[tuple[int, int]], |
| 195 | ) -> tuple[torch.Tensor, list[bool]]: |
| 196 | """Build the full-length `prompt_embeds` tensor and the `is_token_ids` |
| 197 | mask aligned to `token_ids`.""" |
| 198 | total_len = len(token_ids) |
| 199 | hidden_size = prompt_embeds_tensors[0].shape[1] |
| 200 | dtype = prompt_embeds_tensors[0].dtype |
| 201 | |
| 202 | full_embeds = torch.zeros(total_len, hidden_size, dtype=dtype) |
| 203 | is_token_ids = torch.ones(total_len, dtype=torch.bool) |
| 204 | |
| 205 | for (start, length), tensor in zip(positions, prompt_embeds_tensors, strict=True): |
| 206 | full_embeds[start : start + length] = tensor |
| 207 | is_token_ids[start : start + length] = False |
| 208 | |
| 209 | return full_embeds, is_token_ids.tolist() |
| 210 | |
| 211 | |
| 212 | _PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, str | None, str | None, bool], str | None]() |
| 213 | """ |
| 214 | Used in `_try_get_processor_chat_template` to avoid calling |
| 215 | `cached_get_processor` again if the processor fails to be loaded. |
| 216 | |
| 217 | This is needed because `lru_cache` does not cache when an exception happens. |
| 218 | """ |
| 219 | |
| 220 | |
| 221 | def _try_get_processor_chat_template( |
| 222 | tokenizer: HfTokenizer, |
| 223 | *, |
| 224 | revision: str | None, |
| 225 | code_revision: str | None, |
| 226 | trust_remote_code: bool, |
| 227 | ) -> str | None: |
| 228 | cache_key = ( |
| 229 | tokenizer.name_or_path, |
| 230 | revision, |
| 231 | code_revision, |
| 232 | trust_remote_code, |
| 233 | ) |
| 234 | if cache_key in _PROCESSOR_CHAT_TEMPLATES: |
| 235 | return _PROCESSOR_CHAT_TEMPLATES[cache_key] |
| 236 | |
| 237 | from transformers import ProcessorMixin, PythonBackend, TokenizersBackend |
| 238 | |
| 239 | try: |
| 240 | processor = cached_get_processor( |
| 241 | tokenizer.name_or_path, |
| 242 | processor_cls=(PythonBackend, TokenizersBackend, ProcessorMixin), |
| 243 | revision=revision, |
| 244 | code_revision=code_revision, |
| 245 | trust_remote_code=trust_remote_code, |
| 246 | ) |
| 247 | if ( |
| 248 | isinstance(processor, ProcessorMixin) |
| 249 | and hasattr(processor, "chat_template") |
| 250 | and (chat_template := processor.chat_template) is not None |
| 251 | ): |
| 252 | _PROCESSOR_CHAT_TEMPLATES[cache_key] = chat_template |
| 253 | return chat_template |
| 254 | except Exception: |
| 255 | logger.debug( |
| 256 | "Failed to load AutoProcessor chat template for %s", |
| 257 | tokenizer.name_or_path, |
| 258 | exc_info=True, |
| 259 | ) |
| 260 | |
| 261 | _PROCESSOR_CHAT_TEMPLATES[cache_key] = None |
| 262 | return None |
| 263 | |
| 264 | |
| 265 | def resolve_chat_template( |
| 266 | tokenizer: HfTokenizer, |
| 267 | chat_template: str | None, |
| 268 | tools: list[dict[str, Any]] | None, |
| 269 | *, |
| 270 | model_config: ModelConfig, |
| 271 | ) -> str | None: |
| 272 | # 1st priority: The given chat template |
| 273 | if chat_template is not None: |
| 274 | # Resolve template names (e.g. "tool_use") to actual Jinja content |
| 275 | # so that downstream kwargs detection can parse template variables. |
| 276 | return tokenizer.get_chat_template(chat_template, tools=tools) |
| 277 | |
| 278 | # 2nd priority: AutoProcessor chat template, unless tool calling is enabled |
| 279 | if tools is None: |
| 280 | chat_template = _try_get_processor_chat_template( |
| 281 | tokenizer, |
| 282 | revision=model_config.tokenizer_revision, |
| 283 | code_revision=( |
| 284 | model_config.code_revision |
| 285 | if model_config.tokenizer == model_config.model |
| 286 | else None |
| 287 | ), |
| 288 | trust_remote_code=model_config.trust_remote_code, |
| 289 | ) |
| 290 | if chat_template is not None: |
| 291 | return chat_template |
| 292 | |
| 293 | # 3rd priority: AutoTokenizer chat template |
| 294 | try: |
| 295 | return tokenizer.get_chat_template(chat_template, tools=tools) |
| 296 | except Exception: |
| 297 | logger.debug( |
| 298 | "Failed to load AutoTokenizer chat template for %s", |
| 299 | tokenizer.name_or_path, |
| 300 | exc_info=True, |
| 301 | ) |
| 302 | |
| 303 | # 4th priority: Predefined fallbacks |
| 304 | path = get_chat_template_fallback_path( |
| 305 | model_type=model_config.hf_config.model_type, |
| 306 | tokenizer_name_or_path=tokenizer.name_or_path, |
| 307 | ) |
| 308 | if path is not None: |
| 309 | logger.info_once( |
| 310 | "Loading chat template fallback for %s as there isn't one " |
| 311 | "defined on HF Hub.", |
| 312 | tokenizer.name_or_path, |
| 313 | ) |
| 314 | chat_template = load_chat_template(path) |
| 315 | else: |
| 316 | logger.debug_once( |
| 317 | "There is no chat template fallback for %s", tokenizer.name_or_path |
| 318 | ) |
| 319 | |
| 320 | return chat_template |
| 321 | |
| 322 | |
| 323 | def _is_var_access(node: jinja2.nodes.Node, varname: str) -> bool: |
| 324 | if isinstance(node, jinja2.nodes.Name): |
| 325 | return node.ctx == "load" and node.name == varname |
| 326 | |
| 327 | return False |
| 328 | |
| 329 | |
| 330 | def _is_attr_access(node: jinja2.nodes.Node, varname: str, key: str) -> bool: |
| 331 | if isinstance(node, jinja2.nodes.Getitem): |
| 332 | return ( |
| 333 | _is_var_access(node.node, varname) |
| 334 | and isinstance(node.arg, jinja2.nodes.Const) |
| 335 | and node.arg.value == key |
| 336 | ) |
| 337 | |
| 338 | if isinstance(node, jinja2.nodes.Getattr): |
| 339 | return _is_var_access(node.node, varname) and node.attr == key |
| 340 | |
| 341 | return False |
| 342 | |
| 343 | |
| 344 | def _is_var_or_elems_access( |
| 345 | node: jinja2.nodes.Node, |
| 346 | varname: str, |
| 347 | key: str | None = None, |
| 348 | ) -> bool: |
| 349 | if isinstance(node, jinja2.nodes.Filter): |
| 350 | return node.node is not None and _is_var_or_elems_access( |
| 351 | node.node, varname, key |
| 352 | ) |
| 353 | if isinstance(node, jinja2.nodes.Test): |
| 354 | return _is_var_or_elems_access(node.node, varname, key) |
| 355 | |
| 356 | if isinstance(node, jinja2.nodes.Getitem) and isinstance( |
| 357 | node.arg, jinja2.nodes.Slice |
| 358 | ): |
| 359 | return _is_var_or_elems_access(node.node, varname, key) |
| 360 | |
| 361 | return _is_attr_access(node, varname, key) if key else _is_var_access(node, varname) |
| 362 | |
| 363 | |
| 364 | def _iter_nodes_assign_var_or_elems(root: jinja2.nodes.Node, varname: str): |
| 365 | # Global variable that is implicitly defined at the root |
| 366 | yield root, varname |
| 367 | |
| 368 | # Iterative BFS |
| 369 | related_varnames = deque([varname]) |
| 370 | while related_varnames: |
| 371 | related_varname = related_varnames.popleft() |
| 372 | |
| 373 | for assign_ast in root.find_all(jinja2.nodes.Assign): |
| 374 | lhs = assign_ast.target |
| 375 | rhs = assign_ast.node |
| 376 | |
| 377 | if _is_var_or_elems_access(rhs, related_varname): |
| 378 | assert isinstance(lhs, jinja2.nodes.Name) |
| 379 | yield assign_ast, lhs.name |
| 380 | |
| 381 | # Avoid infinite looping for self-assignment |
| 382 | if lhs.name != related_varname: |
| 383 | related_varnames.append(lhs.name) |
| 384 | |
| 385 | |
| 386 | # NOTE: The proper way to handle this is to build a CFG so that we can handle |
| 387 | # the scope in which each variable is defined, but that is too complicated |
| 388 | def _iter_nodes_assign_messages_item(root: jinja2.nodes.Node): |
| 389 | messages_varnames = [ |
| 390 | varname for _, varname in _iter_nodes_assign_var_or_elems(root, "messages") |
| 391 | ] |
| 392 | |
| 393 | # Search for {%- for message in messages -%} loops |
| 394 | for loop_ast in root.find_all(jinja2.nodes.For): |
| 395 | loop_iter = loop_ast.iter |
| 396 | loop_target = loop_ast.target |
| 397 | |
| 398 | for varname in messages_varnames: |
| 399 | if _is_var_or_elems_access(loop_iter, varname): |
| 400 | assert isinstance(loop_target, jinja2.nodes.Name) |
| 401 | yield loop_ast, loop_target.name |
| 402 | break |
| 403 | |
| 404 | |
| 405 | def _iter_nodes_assign_content_item(root: jinja2.nodes.Node): |
| 406 | """Yield loops that iterate over message content or macro-bound content.""" |
| 407 | message_varnames = [ |
| 408 | varname for _, varname in _iter_nodes_assign_messages_item(root) |
| 409 | ] |
| 410 | |
| 411 | # Track macro parameters that receive message.content as an argument. |
| 412 | # Some templates pass message.content through a macro parameter whose |
| 413 | # name is not literally "content". |
| 414 | macro_content_params_by_loop: dict[int, set[str]] = {} |
| 415 | loops_in_macros: set[int] = set() |
| 416 | for macro_node in root.find_all(jinja2.nodes.Macro): |
| 417 | macro_param_names = {arg.name for arg in macro_node.args} |
| 418 | macro_content_params: set[str] = set() |
| 419 | for call_node in root.find_all(jinja2.nodes.Call): |
| 420 | if ( |
| 421 | isinstance(call_node.node, jinja2.nodes.Name) |
| 422 | and call_node.node.name == macro_node.name |
| 423 | ): |
| 424 | for i, arg in enumerate(call_node.args): |
| 425 | if i < len(macro_node.args) and any( |
| 426 | _is_var_or_elems_access(arg, varname, "content") |
| 427 | for varname in message_varnames |
| 428 | ): |
| 429 | macro_content_params.add(macro_node.args[i].name) |
| 430 | for kwarg in call_node.kwargs: |
| 431 | if ( |
| 432 | isinstance(kwarg, jinja2.nodes.Keyword) |
| 433 | and kwarg.key in macro_param_names |
| 434 | and any( |
| 435 | _is_var_or_elems_access(kwarg.value, varname, "content") |
| 436 | for varname in message_varnames |
| 437 | ) |
| 438 | ): |
| 439 | macro_content_params.add(kwarg.key) |
| 440 | |
| 441 | for loop_ast in macro_node.find_all(jinja2.nodes.For): |
| 442 | loops_in_macros.add(id(loop_ast)) |
| 443 | if macro_content_params: |
| 444 | macro_content_params_by_loop[id(loop_ast)] = macro_content_params |
| 445 | |
| 446 | # Search for {%- for content in message['content'] -%} loops |
| 447 | # or {%- for item in content -%} loops |
| 448 | for loop_ast in root.find_all(jinja2.nodes.For): |
| 449 | loop_iter = loop_ast.iter |
| 450 | loop_target = loop_ast.target |
| 451 | |
| 452 | for varname in message_varnames: |
| 453 | if _is_var_or_elems_access(loop_iter, varname, "content"): |
| 454 | assert isinstance(loop_target, jinja2.nodes.Name) |
| 455 | yield loop_ast, loop_target.name |
| 456 | break |
| 457 | else: |
| 458 | macro_content_params_for_loop = macro_content_params_by_loop.get( |
| 459 | id(loop_ast) |
| 460 | ) |
| 461 | if ( |
| 462 | isinstance(loop_iter, jinja2.nodes.Name) |
| 463 | and macro_content_params_for_loop is not None |
| 464 | and loop_iter.name in macro_content_params_for_loop |
| 465 | ): |
| 466 | assert isinstance(loop_target, jinja2.nodes.Name) |
| 467 | yield loop_ast, loop_target.name |
| 468 | continue |
| 469 | |
| 470 | if ( |
| 471 | id(loop_ast) not in loops_in_macros |
| 472 | and isinstance(loop_iter, jinja2.nodes.Name) |
| 473 | and loop_iter.name == "content" |
| 474 | ): |
| 475 | assert isinstance(loop_target, jinja2.nodes.Name) |
| 476 | yield loop_ast, loop_target.name |
| 477 | |
| 478 | |
| 479 | def _try_extract_ast(chat_template: str) -> jinja2.nodes.Template | None: |
| 480 | import transformers.utils.chat_template_utils as hf_chat_utils |
| 481 | |
| 482 | try: |
| 483 | jinja_compiled = hf_chat_utils._compile_jinja_template(chat_template) |
| 484 | return jinja_compiled.environment.parse(chat_template) |
| 485 | except Exception: |
| 486 | logger.exception("Error when compiling Jinja template") |
| 487 | return None |
| 488 | |
| 489 | |
| 490 | @lru_cache(maxsize=32) |
| 491 | def _detect_content_format( |
| 492 | chat_template: str, |
| 493 | *, |
| 494 | default: ChatTemplateContentFormat, |
| 495 | ) -> ChatTemplateContentFormat: |
| 496 | jinja_ast = _try_extract_ast(chat_template) |
| 497 | if jinja_ast is None: |
| 498 | return default |
| 499 | |
| 500 | try: |
| 501 | next(_iter_nodes_assign_content_item(jinja_ast)) |
| 502 | except StopIteration: |
| 503 | return "string" |
| 504 | except Exception: |
| 505 | logger.exception("Error when parsing AST of Jinja template") |
| 506 | return default |
| 507 | else: |
| 508 | return "openai" |
| 509 | |
| 510 | |
| 511 | @lru_cache(maxsize=32) |
| 512 | def _detect_developer_role_support(chat_template: str) -> bool: |
| 513 | return '"developer"' in chat_template or "'developer'" in chat_template |
| 514 | |
| 515 | |
| 516 | def _convert_developer_to_system( |
| 517 | conversation: list[ConversationMessage], |
| 518 | ) -> list[ConversationMessage]: |
| 519 | converted: list[ConversationMessage] = [] |
| 520 | for msg in conversation: |
| 521 | if msg["role"] == "developer": |
| 522 | new_msg = dict(msg) |
| 523 | new_msg["role"] = "system" |
| 524 | new_msg.pop("tools", None) |
| 525 | converted.append(new_msg) # type: ignore[arg-type] |
| 526 | else: |
| 527 | converted.append(msg) |
| 528 | return converted |
| 529 | |
| 530 | |
| 531 | def _consolidate_system_messages( |
| 532 | conversation: list[ConversationMessage], |
| 533 | ) -> list[ConversationMessage]: |
| 534 | """Merge all system messages into one at position 0. |
| 535 | |
| 536 | Some chat templates (e.g. Qwen 3.6) require the system message to be the |
| 537 | very first message. After developer-to-system conversion, system messages |
| 538 | may appear at non-first positions; this merges them into a single message. |
| 539 | """ |
| 540 | system_contents: list[str] = [] |
| 541 | non_system: list[ConversationMessage] = [] |
| 542 | needs_consolidation = False |
| 543 | for i, msg in enumerate(conversation): |
| 544 | if msg["role"] == "system": |
| 545 | if i > 0 or system_contents: |
| 546 | needs_consolidation = True |
| 547 | content = msg.get("content", "") |
| 548 | if isinstance(content, list): |
| 549 | parts = [] |
| 550 | for part in content: |
| 551 | if isinstance(part, dict) and "text" in part: |
| 552 | parts.append(part["text"]) |
| 553 | elif isinstance(part, str): |
| 554 | parts.append(part) |
| 555 | content = "\n".join(parts) |
| 556 | if content: |
| 557 | system_contents.append(content) |
| 558 | else: |
| 559 | non_system.append(msg) |
| 560 | |
| 561 | if not needs_consolidation: |
| 562 | return conversation |
| 563 | |
| 564 | merged: ConversationMessage = { |
| 565 | "role": "system", |
| 566 | "content": "\n\n".join(system_contents), |
| 567 | } |
| 568 | return [merged, *non_system] |
| 569 | |
| 570 | |
| 571 | _CONTENT_FORMATS_MAXSIZE = 32 |
| 572 | _CONTENT_FORMATS = dict[ |
| 573 | tuple[str | None, bool, str, str | None, str | None, bool], |
| 574 | "ChatTemplateContentFormat", |
| 575 | ]() |
| 576 | """ |
| 577 | Used in `_resolve_chat_template_content_format` to avoid resolving and parsing |
| 578 | the chat template on every request. |
| 579 | |
| 580 | `lru_cache` cannot be used because `tools` is a list and `ModelConfig` defines |
| 581 | `__eq__` without `__hash__`, so the key holds the fields `resolve_chat_template` |
| 582 | selects the template by. Only the presence of `tools` is part of it: it decides |
| 583 | whether the processor template is tried, and its contents never reach the AST. |
| 584 | """ |
| 585 | |
| 586 | |
| 587 | def _resolve_chat_template_content_format( |
| 588 | chat_template: str | None, |
| 589 | tools: list[dict[str, Any]] | None, |
| 590 | tokenizer: HfTokenizer, |
| 591 | *, |
| 592 | model_config: ModelConfig, |
| 593 | ) -> ChatTemplateContentFormat: |
| 594 | cache_key = ( |
| 595 | chat_template, |
| 596 | tools is None, |
| 597 | tokenizer.name_or_path, |
| 598 | model_config.revision, |
| 599 | model_config.code_revision, |
| 600 | model_config.trust_remote_code, |
| 601 | ) |
| 602 | if (cached_format := _CONTENT_FORMATS.get(cache_key)) is not None: |
| 603 | return cached_format |
| 604 | |
| 605 | resolved_chat_template = resolve_chat_template( |
| 606 | tokenizer, |
| 607 | chat_template=chat_template, |
| 608 | tools=tools, |
| 609 | model_config=model_config, |
| 610 | ) |
| 611 | |
| 612 | jinja_text = ( |
| 613 | resolved_chat_template |
| 614 | if isinstance(resolved_chat_template, str) |
| 615 | else load_chat_template(chat_template, is_literal=True) |
| 616 | ) |
| 617 | |
| 618 | detected_format = ( |
| 619 | "string" |
| 620 | if jinja_text is None |
| 621 | else _detect_content_format(jinja_text, default="string") |
| 622 | ) |
| 623 | |
| 624 | # Requests may carry their own chat template, so bound the cache. |
| 625 | if len(_CONTENT_FORMATS) >= _CONTENT_FORMATS_MAXSIZE: |
| 626 | _CONTENT_FORMATS.clear() |
| 627 | _CONTENT_FORMATS[cache_key] = detected_format |
| 628 | |
| 629 | return detected_format |
| 630 | |
| 631 | |
| 632 | def resolve_chat_template_content_format( |
| 633 | chat_template: str | None, |
| 634 | tools: list[dict[str, Any]] | None, |
| 635 | given_format: ChatTemplateContentFormatOption, |
| 636 | tokenizer: HfTokenizer, |
| 637 | *, |
| 638 | model_config: ModelConfig, |
| 639 | ) -> ChatTemplateContentFormat: |
| 640 | detected_format = _resolve_chat_template_content_format( |
| 641 | chat_template, |
| 642 | tools, |
| 643 | tokenizer, |
| 644 | model_config=model_config, |
| 645 | ) |
| 646 | |
| 647 | if given_format == "auto": |
| 648 | logger.info_once( |
| 649 | "Detected the chat template content format to be '%s'. " |
| 650 | "You can set `--chat-template-content-format` to override this.", |
| 651 | detected_format, |
| 652 | ) |
| 653 | elif given_format != detected_format: |
| 654 | logger.warning_once( |
| 655 | "You specified `--chat-template-content-format %s` " |
| 656 | "which is different from the detected format '%s'. " |
| 657 | "If our automatic detection is incorrect, please consider " |
| 658 | "opening a GitHub issue so that we can improve it: " |
| 659 | "https://github.com/vllm-project/vllm/issues/new/choose", |
| 660 | given_format, |
| 661 | detected_format, |
| 662 | ) |
| 663 | |
| 664 | return detected_format if given_format == "auto" else given_format |
| 665 | |
| 666 | |
| 667 | # adapted from https://github.com/huggingface/transformers/blob/v4.56.2/src/transformers/utils/chat_template_utils.py#L398-L412 |
| 668 | # only preserve the parse function used to resolve chat template kwargs |
| 669 | class AssistantTracker(jinja2.ext.Extension): |
| 670 | tags = {"generation"} |
| 671 | |
| 672 | def parse(self, parser: jinja2.parser.Parser) -> jinja2.nodes.Node: |
| 673 | lineno = next(parser.stream).lineno |
| 674 | body = parser.parse_statements(("name:endgeneration",), drop_needle=True) |
| 675 | call = self.call_method("_generation_support") |
| 676 | call_block = jinja2.nodes.CallBlock(call, [], [], body) |
| 677 | return call_block.set_lineno(lineno) |
| 678 | |
| 679 | |
| 680 | def _resolve_chat_template_kwargs(chat_template: str) -> Set[str]: |
| 681 | env = jinja2.sandbox.ImmutableSandboxedEnvironment( |
| 682 | trim_blocks=True, |
| 683 | lstrip_blocks=True, |
| 684 | extensions=[AssistantTracker, jinja2.ext.loopcontrols], |
| 685 | ) |
| 686 | parsed_content = env.parse(chat_template) |
| 687 | template_vars = jinja2.meta.find_undeclared_variables(parsed_content) |
| 688 | return template_vars |
| 689 | |
| 690 | |
| 691 | _cached_resolve_chat_template_kwargs = lru_cache(_resolve_chat_template_kwargs) |
| 692 | |
| 693 | |
| 694 | @lru_cache |
| 695 | def _get_hf_base_chat_template_params() -> frozenset[str]: |
| 696 | from transformers import PythonBackend |
| 697 | |
| 698 | # Get standard parameters from HuggingFace's base tokenizer class. |
| 699 | # This dynamically extracts parameters from PythonBackend's |
| 700 | # apply_chat_template method, ensuring compatibility with tokenizers |
| 701 | # that use **kwargs to receive standard parameters. |
| 702 | |
| 703 | # Read signature from HF's base class - the single source of truth |
| 704 | base_sig = inspect.signature(PythonBackend.apply_chat_template) |
| 705 | |
| 706 | # Exclude VAR_KEYWORD (**kwargs) and VAR_POSITIONAL (*args) placeholders |
| 707 | return frozenset( |
| 708 | p.name |
| 709 | for p in base_sig.parameters.values() |
| 710 | if p.kind |
| 711 | not in (inspect.Parameter.VAR_KEYWORD, inspect.Parameter.VAR_POSITIONAL) |
| 712 | ) |
| 713 | |
| 714 | |
| 715 | def resolve_chat_template_kwargs( |
| 716 | tokenizer: HfTokenizer, |
| 717 | chat_template: str, |
| 718 | chat_template_kwargs: dict[str, Any], |
| 719 | raise_on_unexpected: bool = True, |
| 720 | ) -> dict[str, Any]: |
| 721 | # We exclude chat_template from kwargs here, because |
| 722 | # chat template has been already resolved at this stage |
| 723 | unexpected_vars = {"chat_template", "tokenize"} |
| 724 | if raise_on_unexpected and ( |
| 725 | unexpected_in_kwargs := unexpected_vars & chat_template_kwargs.keys() |
| 726 | ): |
| 727 | raise ValueError( |
| 728 | "Found unexpected chat template kwargs from request: " |
| 729 | f"{unexpected_in_kwargs}" |
| 730 | ) |
| 731 | |
| 732 | fn_kw = { |
| 733 | k |
| 734 | for k in chat_template_kwargs |
| 735 | if supports_kw(tokenizer.apply_chat_template, k, allow_var_kwargs=False) |
| 736 | } |
| 737 | template_vars = _cached_resolve_chat_template_kwargs(chat_template) |
| 738 | |
| 739 | # Allow standard HF parameters even if tokenizer uses **kwargs to receive them |
| 740 | hf_base_params = _get_hf_base_chat_template_params() |
| 741 | |
| 742 | accept_vars = (fn_kw | template_vars | hf_base_params) - unexpected_vars |
| 743 | return {k: v for k, v in chat_template_kwargs.items() if k in accept_vars} |
| 744 | |
| 745 | |
| 746 | def _template_error_reason(exc: BaseException) -> str: |
| 747 | # Extract the most specific reason from a chat template error chain. |
| 748 | seen: set[int] = set() |
| 749 | current: BaseException | None = exc |
| 750 | while current is not None and id(current) not in seen: |
| 751 | seen.add(id(current)) |
| 752 | if isinstance(current, jinja2.TemplateError): |
| 753 | return str(current) |
| 754 | current = current.__cause__ or current.__context__ |
| 755 | return str(exc) |
| 756 | |
| 757 | |
| 758 | @overload |
| 759 | def safe_apply_chat_template( |
| 760 | model_config: ModelConfig, |
| 761 | tokenizer: HfTokenizer, |
| 762 | conversation: list[ConversationMessage], |
| 763 | *, |
| 764 | tools: list[dict[str, Any]] | None = ..., |
| 765 | chat_template: str | None = ..., |
| 766 | tokenize: Literal[True] = ..., |
| 767 | **kwargs, |
| 768 | ) -> list[int]: ... |
| 769 | @overload |
| 770 | def safe_apply_chat_template( |
| 771 | model_config: ModelConfig, |
| 772 | tokenizer: HfTokenizer, |
| 773 | conversation: list[ConversationMessage], |
| 774 | *, |
| 775 | tools: list[dict[str, Any]] | None = ..., |
| 776 | chat_template: str | None = ..., |
| 777 | tokenize: Literal[False] = ..., |
| 778 | **kwargs, |
| 779 | ) -> str: ... |
| 780 | def safe_apply_chat_template( |
| 781 | model_config: ModelConfig, |
| 782 | tokenizer: HfTokenizer, |
| 783 | conversation: list[ConversationMessage], |
| 784 | *, |
| 785 | tools: list[dict[str, Any]] | None = None, |
| 786 | chat_template: str | None = None, |
| 787 | tokenize: bool = True, |
| 788 | **kwargs, |
| 789 | ) -> str | list[int]: |
| 790 | chat_template = resolve_chat_template( |
| 791 | tokenizer, |
| 792 | chat_template=chat_template, |
| 793 | tools=tools, |
| 794 | model_config=model_config, |
| 795 | ) |
| 796 | if chat_template is None: |
| 797 | raise ChatTemplateResolutionError( |
| 798 | "As of transformers v4.44, default chat template is no longer " |
| 799 | "allowed, so you must provide a chat template if the tokenizer " |
| 800 | "does not define one." |
| 801 | ) |
| 802 | if any( |
| 803 | msg["role"] == "developer" for msg in conversation |
| 804 | ) and not _detect_developer_role_support(chat_template): |
| 805 | conversation = _convert_developer_to_system(conversation) |
| 806 | conversation = _consolidate_system_messages(conversation) |
| 807 | logger.info_once( |
| 808 | "Chat template does not support the 'developer' message role. " |
| 809 | "Converting developer messages to 'system' role.", |
| 810 | ) |
| 811 | resolved_kwargs = resolve_chat_template_kwargs( |
| 812 | tokenizer=tokenizer, |
| 813 | chat_template=chat_template, |
| 814 | chat_template_kwargs=kwargs, |
| 815 | ) |
| 816 | |
| 817 | # transformers v5 changed the default of `return_dict` to True, which |
| 818 | # makes `apply_chat_template(tokenize=True)` return a `BatchEncoding` |
| 819 | # instead of `list[int]`. Force `return_dict=False` so downstream code |
| 820 | # that expects a flat token list (e.g. `parse_dec_only_prompt`) works |
| 821 | # consistently across v4 and v5. |
| 822 | if tokenize and "return_dict" not in resolved_kwargs: |
| 823 | resolved_kwargs["return_dict"] = False |
| 824 | |
| 825 | try: |
| 826 | plain = tokenizer.apply_chat_template( |
| 827 | conversation=conversation, # type: ignore[arg-type] |
| 828 | tools=tools, # type: ignore[arg-type] |
| 829 | chat_template=chat_template, |
| 830 | tokenize=tokenize, |
| 831 | **resolved_kwargs, |
| 832 | ) |
| 833 | except Exception as e: |
| 834 | # Chat templates reject invalid user input (e.g. an unsupported |
| 835 | # `reasoning_effort` value) by raising from within the template. |
| 836 | # Surface those as a 400 Bad Request carrying the template's own |
| 837 | # reason (which typically lists the supported values) instead of a |
| 838 | # 500 or any generic upstream wrapper message. |
| 839 | logger.warning("Chat template rejected the request: %s", e) |
| 840 | raise VLLMValidationError(_template_error_reason(e)) from e |
| 841 | |
| 842 | return plain |
| 843 | |
| 844 | |
| 845 | def rebuild_mm_uuids_from_mm_data( |
| 846 | mm_uuids: MultiModalUUIDDict, |
| 847 | mm_data: MultiModalDataDict, |
| 848 | ) -> MultiModalUUIDDict: |
| 849 | """Rebuild mm_uuids after vision_chunk processing. |
| 850 | |
| 851 | When videos are split into chunks, the original UUIDs need to be updated |
| 852 | to reflect the new UUIDs generated for each chunk. |
| 853 | |
| 854 | Args: |
| 855 | mm_uuids: Original UUIDs dictionary |
| 856 | mm_data: Processed multimodal data with vision_chunk items |
| 857 | |
| 858 | Returns: |
| 859 | Updated UUIDs dictionary with chunk UUIDs |
| 860 | |
| 861 | """ |
| 862 | vision_chunks = mm_data.get("vision_chunk") |
| 863 | if vision_chunks is None: |
| 864 | return mm_uuids |
| 865 | |
| 866 | assert all(isinstance(item, dict) for item in vision_chunks), ( |
| 867 | "Expected all vision_chunk items to be dicts" |
| 868 | ) |
| 869 | vision_chunks = cast(list[dict[str, Any]], vision_chunks) |
| 870 | vision_chunk_uuids = [ |
| 871 | uuid_val for item in vision_chunks if (uuid_val := item.get("uuid")) is not None |
| 872 | ] |
| 873 | |
| 874 | if vision_chunk_uuids: |
| 875 | mm_uuids = dict(mm_uuids) |
| 876 | mm_uuids["vision_chunk"] = vision_chunk_uuids |
| 877 | |
| 878 | return mm_uuids |
| 879 | |
| 880 | |
| 881 | def build_video_prompts_from_mm_data( |
| 882 | mm_data: MultiModalDataDict, |
| 883 | ) -> list[str]: |
| 884 | """Build video prompts from vision_chunk data. |
| 885 | |
| 886 | Collects prompts from video chunks and groups them by video_idx. |
| 887 | |
| 888 | Args: |
| 889 | mm_data: Processed multimodal data with vision_chunk items |
| 890 | |
| 891 | Returns: |
| 892 | List of video prompts, one per video. |
| 893 | |
| 894 | """ |
| 895 | vision_chunks = mm_data.get("vision_chunk") |
| 896 | if vision_chunks is None: |
| 897 | return [] |
| 898 | |
| 899 | # Group chunks by video_idx |
| 900 | video_prompts_dict: dict[int, list[str]] = defaultdict(list) |
| 901 | |
| 902 | for item in vision_chunks: |
| 903 | # vision_chunk items are always dicts (VisionChunkImage/VisionChunkVideo) |
| 904 | assert isinstance(item, dict) |
| 905 | if item.get("type") == "video_chunk": |
| 906 | video_idx = item.get("video_idx", 0) |
| 907 | prompt = item.get("prompt", "") |
| 908 | video_prompts_dict[video_idx].append(prompt) |
| 909 | |
| 910 | # Build prompts in video order |
| 911 | video_prompts = [ |
| 912 | "".join(video_prompts_dict[video_idx]) |
| 913 | for video_idx in sorted(video_prompts_dict.keys()) |
| 914 | ] |
| 915 | |
| 916 | return video_prompts |
| 917 | |
| 918 | |
| 919 | def replace_vision_chunk_video_placeholder( |
| 920 | prompt_raw: str | list[int], |
| 921 | mm_data: MultiModalDataDict, |
| 922 | video_placeholder: str | None, |
| 923 | ) -> str | list[int]: |
| 924 | # get video placeholder, replace it with runtime video-chunk prompts |
| 925 | if video_placeholder and isinstance(prompt_raw, str): |
| 926 | video_prompts = build_video_prompts_from_mm_data(mm_data) |
| 927 | |
| 928 | # replace in order |
| 929 | prompt_raw_parts = prompt_raw.split(video_placeholder) |
| 930 | if len(prompt_raw_parts) == len(video_prompts) + 1: |
| 931 | prompt_raw = "".join( |
| 932 | itertools.chain.from_iterable(zip(prompt_raw_parts, video_prompts)) |
| 933 | ) |
| 934 | prompt_raw += prompt_raw_parts[-1] |
| 935 | else: |
| 936 | logger.warning( |
| 937 | "Number of video placeholders (%d) does not match " |
| 938 | "number of videos (%d) in the request.", |
| 939 | len(prompt_raw_parts) - 1, |
| 940 | len(video_prompts), |
| 941 | ) |
| 942 | return prompt_raw |
| 943 | |
| 944 | |
| 945 | class HfRenderer(BaseRenderer[HfTokenizer]): |
| 946 | def __init__( |
| 947 | self, |
| 948 | config: VllmConfig, |
| 949 | tokenizer: HfTokenizer | None, |
| 950 | ) -> None: |
| 951 | # Ensure the og tokenizer is never modified by maybe_make_thread_pool |
| 952 | tokenizer = copy.copy(tokenizer) |
| 953 | if ( |
| 954 | # Skip for mock configs and tokenizers |
| 955 | getattr(config.model_config, "enable_prompt_embeds", False) |
| 956 | and isinstance(tokenizer, HfTokenizer) |
| 957 | ): |
| 958 | _ensure_prompt_embeds_placeholder_token(tokenizer) |
| 959 | super().__init__(config, tokenizer) |
| 960 | |
| 961 | self.use_unified_vision_chunk = getattr( |
| 962 | config.model_config.hf_config, "use_unified_vision_chunk", False |
| 963 | ) |
| 964 | |
| 965 | self._apply_chat_template_async = make_async( |
| 966 | safe_apply_chat_template, executor=self._executor |
| 967 | ) |
| 968 | |
| 969 | if self.tokenizer is not None: |
| 970 | maybe_make_thread_pool( |
| 971 | self.tokenizer, config.model_config.renderer_num_workers + 1 |
| 972 | ) |
| 973 | |
| 974 | def _replace_executor(self) -> None: |
| 975 | """Replace the inner render pool after a stuck worker. |
| 976 | |
| 977 | Called after a render timeout so that the stuck thread (which CPython |
| 978 | cannot interrupt) does not permanently block all subsequent renders. |
| 979 | The executor *object* is kept so ``make_async`` wrappers bound at |
| 980 | init (tokenize, decode, pooling, derender) keep submitting work. |
| 981 | The old inner pool is abandoned with wait=False — its thread will |
| 982 | eventually finish or be reaped on process exit. |
| 983 | """ |
| 984 | self._executor.replace_inner() |
| 985 | logger.warning( |
| 986 | "Chat template render timed out — executor pool replaced " |
| 987 | "(%d worker(s)). The stuck thread may continue consuming " |
| 988 | "CPU until the process exits.", |
| 989 | self.model_config.renderer_num_workers, |
| 990 | ) |
| 991 | |
| 992 | async def _render_with_timeout(self, coro): |
| 993 | """Run a template-rendering coroutine with a timeout. |
| 994 | |
| 995 | On timeout, replaces the executor so subsequent renders are not |
| 996 | permanently blocked by the stuck thread. |
| 997 | """ |
| 998 | timeout = envs.VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT |
| 999 | if timeout <= 0: |
| 1000 | return await coro |
| 1001 | try: |
| 1002 | return await asyncio.wait_for(coro, timeout=timeout) |
| 1003 | except asyncio.TimeoutError: |
| 1004 | self._replace_executor() |
| 1005 | raise TimeoutError( |
| 1006 | f"Chat template rendering timed out after " |
| 1007 | f"{timeout}s. This may indicate a malicious or " |
| 1008 | f"excessively complex template. Adjust " |
| 1009 | f"VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT to override." |
| 1010 | ) from None |
| 1011 | |
| 1012 | def _can_produce_offsets(self) -> bool: |
| 1013 | # Only fast tokenizers expose offset_mapping. |
| 1014 | return self.tokenizer is not None and self.tokenizer.is_fast |
| 1015 | |
| 1016 | def render_messages( |
| 1017 | self, |
| 1018 | messages: list[ChatCompletionMessageParam], |
| 1019 | params: ChatParams, |
| 1020 | ) -> tuple[list[ConversationMessage], DictPrompt]: |
| 1021 | model_config = self.model_config |
| 1022 | tokenizer = self.get_tokenizer() |
| 1023 | |
| 1024 | prompt_embeds_placeholder_token_id: int | None = None |
| 1025 | if model_config.enable_prompt_embeds: |
| 1026 | prompt_embeds_placeholder_token_id = ( |
| 1027 | _ensure_prompt_embeds_placeholder_token(tokenizer) |
| 1028 | ) |
| 1029 | |
| 1030 | conversation, mm_data, mm_uuids = parse_chat_messages( |
| 1031 | messages, |
| 1032 | model_config, |
| 1033 | content_format=resolve_chat_template_content_format( |
| 1034 | chat_template=params.chat_template, |
| 1035 | tools=params.chat_template_kwargs.get("tools"), |
| 1036 | given_format=params.chat_template_content_format, |
| 1037 | tokenizer=tokenizer, |
| 1038 | model_config=model_config, |
| 1039 | ), |
| 1040 | media_io_kwargs=params.media_io_kwargs, |
| 1041 | mm_processor_kwargs=params.mm_processor_kwargs, |
| 1042 | ) |
| 1043 | |
| 1044 | # prompt_embeds tensors are carried by the tracker through mm_data, |
| 1045 | # but they must NOT be fed to the MM processor (which would reject |
| 1046 | # the unknown key). Extract them here. |
| 1047 | prompt_embeds_tensors: list[torch.Tensor] | None = None |
| 1048 | if mm_data is not None and "prompt_embeds" in mm_data: |
| 1049 | prompt_embeds_tensors = list( |
| 1050 | cast(Sequence[torch.Tensor], mm_data["prompt_embeds"]) |
| 1051 | ) |
| 1052 | mm_data = {k: v for k, v in mm_data.items() if k != "prompt_embeds"} |
| 1053 | if not mm_data: |
| 1054 | mm_data = None |
| 1055 | |
| 1056 | chat_template_kwargs = params.get_apply_chat_template_kwargs() |
| 1057 | if prompt_embeds_tensors: |
| 1058 | # prompt_embeds post-processing requires prompt_token_ids. |
| 1059 | if chat_template_kwargs.get("tokenize") is False: |
| 1060 | logger.warning_once(_TOKENIZE_OVERRIDE_WARNING) |
| 1061 | chat_template_kwargs["tokenize"] = True |
| 1062 | |
| 1063 | prompt_raw = safe_apply_chat_template( |
| 1064 | model_config, |
| 1065 | tokenizer, |
| 1066 | conversation, |
| 1067 | **chat_template_kwargs, |
| 1068 | ) |
| 1069 | |
| 1070 | # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5 |
| 1071 | # model which uses unified vision chunks for both images and videos. |
| 1072 | if ( |
| 1073 | self.use_unified_vision_chunk |
| 1074 | and mm_uuids is not None |
| 1075 | and mm_data is not None |
| 1076 | ): |
| 1077 | mm_uuids = rebuild_mm_uuids_from_mm_data(mm_uuids, mm_data) |
| 1078 | |
| 1079 | # get video placeholder, replace it with runtime video-chunk prompts |
| 1080 | video_placeholder = getattr( |
| 1081 | model_config.hf_config, "video_placeholder", None |
| 1082 | ) |
| 1083 | prompt_raw = cast( |
| 1084 | list[int], |
| 1085 | replace_vision_chunk_video_placeholder( |
| 1086 | prompt_raw, |
| 1087 | mm_data, |
| 1088 | video_placeholder, |
| 1089 | ), |
| 1090 | ) |
| 1091 | |
| 1092 | prompt = parse_dec_only_prompt(prompt_raw) |
| 1093 | |
| 1094 | # When `prompt_embeds` is mixed with other modality data, |
| 1095 | # `_process_tokens` runs `_process_multimodal` first (expanding |
| 1096 | # `<|AUDIO|>` / `<|IMAGE|>` placeholders) and then |
| 1097 | # `_apply_prompt_embeds_to_engine_input` augments the result. |
| 1098 | # Stash the tensors and placeholder ID for that override to consume. |
| 1099 | if prompt_embeds_tensors and mm_data: |
| 1100 | assert prompt_embeds_placeholder_token_id is not None |
| 1101 | cast(dict, prompt)["_prompt_embeds"] = ( |
| 1102 | prompt_embeds_tensors, |
| 1103 | prompt_embeds_placeholder_token_id, |
| 1104 | ) |
| 1105 | if params.mm_processor_kwargs: |
| 1106 | cast(dict, prompt)["mm_processor_kwargs"] = params.mm_processor_kwargs |
| 1107 | elif prompt_embeds_tensors: |
| 1108 | # Pure mode: no other MM data, mutate prompt to EmbedsPrompt shape. |
| 1109 | assert prompt_embeds_placeholder_token_id is not None |
| 1110 | self._apply_prompt_embeds_to_prompt( |
| 1111 | prompt, |
| 1112 | prompt_embeds_tensors, |
| 1113 | prompt_embeds_placeholder_token_id, |
| 1114 | ) |
| 1115 | |
| 1116 | if mm_data is not None: |
| 1117 | prompt["multi_modal_data"] = mm_data |
| 1118 | if mm_uuids is not None: |
| 1119 | prompt["multi_modal_uuids"] = mm_uuids |
| 1120 | |
| 1121 | return conversation, prompt |
| 1122 | |
| 1123 | async def render_messages_async( |
| 1124 | self, |
| 1125 | messages: list[ChatCompletionMessageParam], |
| 1126 | params: ChatParams, |
| 1127 | ) -> tuple[list[ConversationMessage], DictPrompt]: |
| 1128 | model_config = self.model_config |
| 1129 | tokenizer = self.get_tokenizer() |
| 1130 | |
| 1131 | prompt_embeds_placeholder_token_id: int | None = None |
| 1132 | if model_config.enable_prompt_embeds: |
| 1133 | prompt_embeds_placeholder_token_id = ( |
| 1134 | _ensure_prompt_embeds_placeholder_token(tokenizer) |
| 1135 | ) |
| 1136 | |
| 1137 | conversation, mm_data, mm_uuids = await parse_chat_messages_async( |
| 1138 | messages, |
| 1139 | model_config, |
| 1140 | content_format=resolve_chat_template_content_format( |
| 1141 | chat_template=params.chat_template, |
| 1142 | tools=params.chat_template_kwargs.get("tools"), |
| 1143 | given_format=params.chat_template_content_format, |
| 1144 | tokenizer=tokenizer, |
| 1145 | model_config=model_config, |
| 1146 | ), |
| 1147 | media_io_kwargs=params.media_io_kwargs, |
| 1148 | mm_processor_kwargs=params.mm_processor_kwargs, |
| 1149 | ) |
| 1150 | |
| 1151 | prompt_embeds_tensors: list[torch.Tensor] | None = None |
| 1152 | if mm_data is not None and "prompt_embeds" in mm_data: |
| 1153 | prompt_embeds_tensors = list( |
| 1154 | cast(Sequence[torch.Tensor], mm_data["prompt_embeds"]) |
| 1155 | ) |
| 1156 | mm_data = {k: v for k, v in mm_data.items() if k != "prompt_embeds"} |
| 1157 | if not mm_data: |
| 1158 | mm_data = None |
| 1159 | |
| 1160 | chat_template_kwargs = params.get_apply_chat_template_kwargs() |
| 1161 | if prompt_embeds_tensors: |
| 1162 | # prompt_embeds post-processing requires prompt_token_ids. |
| 1163 | if chat_template_kwargs.get("tokenize") is False: |
| 1164 | logger.warning_once(_TOKENIZE_OVERRIDE_WARNING) |
| 1165 | chat_template_kwargs["tokenize"] = True |
| 1166 | |
| 1167 | prompt_raw = await self._render_with_timeout( |
| 1168 | self._apply_chat_template_async( |
| 1169 | model_config, |
| 1170 | tokenizer, |
| 1171 | conversation, |
| 1172 | **chat_template_kwargs, |
| 1173 | ) |
| 1174 | ) |
| 1175 | |
| 1176 | # NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5 |
| 1177 | # model which uses unified vision chunks for both images and videos. |
| 1178 | if ( |
| 1179 | self.use_unified_vision_chunk |
| 1180 | and mm_uuids is not None |
| 1181 | and mm_data is not None |
| 1182 | ): |
| 1183 | mm_uuids = rebuild_mm_uuids_from_mm_data(mm_uuids, mm_data) |
| 1184 | |
| 1185 | # get video placeholder, replace it with runtime video-chunk prompts |
| 1186 | video_placeholder = getattr( |
| 1187 | model_config.hf_config, "video_placeholder", None |
| 1188 | ) |
| 1189 | prompt_raw = cast( |
| 1190 | list[int], |
| 1191 | replace_vision_chunk_video_placeholder( |
| 1192 | prompt_raw, |
| 1193 | mm_data, |
| 1194 | video_placeholder, |
| 1195 | ), |
| 1196 | ) |
| 1197 | |
| 1198 | prompt = parse_dec_only_prompt(prompt_raw) |
| 1199 | |
| 1200 | # See `render_messages` for the rationale. |
| 1201 | if prompt_embeds_tensors and mm_data: |
| 1202 | assert prompt_embeds_placeholder_token_id is not None |
| 1203 | cast(dict, prompt)["_prompt_embeds"] = ( |
| 1204 | prompt_embeds_tensors, |
| 1205 | prompt_embeds_placeholder_token_id, |
| 1206 | ) |
| 1207 | if params.mm_processor_kwargs: |
| 1208 | cast(dict, prompt)["mm_processor_kwargs"] = params.mm_processor_kwargs |
| 1209 | elif prompt_embeds_tensors: |
| 1210 | assert prompt_embeds_placeholder_token_id is not None |
| 1211 | self._apply_prompt_embeds_to_prompt( |
| 1212 | prompt, |
| 1213 | prompt_embeds_tensors, |
| 1214 | prompt_embeds_placeholder_token_id, |
| 1215 | ) |
| 1216 | |
| 1217 | if mm_data is not None: |
| 1218 | prompt["multi_modal_data"] = mm_data |
| 1219 | if mm_uuids is not None: |
| 1220 | prompt["multi_modal_uuids"] = mm_uuids |
| 1221 | |
| 1222 | return conversation, prompt |
| 1223 | |
| 1224 | @override |
| 1225 | def _process_tokens( |
| 1226 | self, |
| 1227 | prompt: TokensPrompt, |
| 1228 | *, |
| 1229 | skip_mm_cache: bool = False, |
| 1230 | ) -> TokensInput | MultiModalInput: |
| 1231 | """Pre-expand `prompt_embeds` sentinels before delegating to the MM |
| 1232 | processor, then attach `prompt_embeds` modality data to the result. |
| 1233 | |
| 1234 | Mixed mode only: the `_prompt_embeds` stash is set by |
| 1235 | `render_messages` when `prompt_embeds` co-exist with other MM data |
| 1236 | (images, audio, …). We expand each 1-token sentinel to an N-token |
| 1237 | span *before* calling `super()._process_tokens()` so the MM |
| 1238 | processor records all placeholder offsets in the final (post-expansion) |
| 1239 | coordinate space, no offset shifting needed afterwards. |
| 1240 | """ |
| 1241 | prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) |
| 1242 | if prompt_embeds_info is None: |
| 1243 | return super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache) |
| 1244 | |
| 1245 | tensors, placeholder_token_id = prompt_embeds_info |
| 1246 | mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) |
| 1247 | cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( |
| 1248 | list(prompt["prompt_token_ids"]), mm_updates |
| 1249 | ) |
| 1250 | engine_input = super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache) |
| 1251 | self._apply_prompt_embeds_to_engine_input( |
| 1252 | cast(MultiModalInput, engine_input), |
| 1253 | tensors, |
| 1254 | mm_updates, |
| 1255 | ) |
| 1256 | |
| 1257 | return engine_input |
| 1258 | |
| 1259 | @override |
| 1260 | async def _process_tokens_async( |
| 1261 | self, |
| 1262 | prompt: TokensPrompt, |
| 1263 | *, |
| 1264 | skip_mm_cache: bool = False, |
| 1265 | ) -> TokensInput | MultiModalInput: |
| 1266 | """Async equivalent of `_process_tokens`.""" |
| 1267 | prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None) |
| 1268 | if prompt_embeds_info is None: |
| 1269 | return await super()._process_tokens_async( |
| 1270 | prompt, skip_mm_cache=skip_mm_cache |
| 1271 | ) |
| 1272 | |
| 1273 | tensors, placeholder_token_id = prompt_embeds_info |
| 1274 | mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id) |
| 1275 | cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders( |
| 1276 | list(prompt["prompt_token_ids"]), mm_updates |
| 1277 | ) |
| 1278 | engine_input = await super()._process_tokens_async( |
| 1279 | prompt, skip_mm_cache=skip_mm_cache |
| 1280 | ) |
| 1281 | self._apply_prompt_embeds_to_engine_input( |
| 1282 | cast(MultiModalInput, engine_input), |
| 1283 | tensors, |
| 1284 | mm_updates, |
| 1285 | ) |
| 1286 | |
| 1287 | return engine_input |
| 1288 | |
| 1289 | @staticmethod |
| 1290 | def _apply_prompt_embeds_to_prompt( |
| 1291 | prompt: DictPrompt, |
| 1292 | prompt_embeds_tensors: list[torch.Tensor], |
| 1293 | placeholder_token_id: int, |
| 1294 | ) -> None: |
| 1295 | """Mutate `prompt` from `TokensPrompt` to `EmbedsPrompt` shape. |
| 1296 | |
| 1297 | Pure `prompt_embeds` path only (no other MM modalities). Expands |
| 1298 | each `<prompt_embeds>` sentinel token into an N-token span and builds |
| 1299 | the full-length `prompt_embeds` tensor + `prompt_is_token_ids` mask |
| 1300 | that the engine's `enable_prompt_embeds` worker branch consumes. |
| 1301 | """ |
| 1302 | token_ids = cast(list[int] | None, prompt.get("prompt_token_ids")) |
| 1303 | if token_ids is None: |
| 1304 | raise RuntimeError(_MISSING_PROMPT_TOKEN_IDS_ERROR) |
| 1305 | |
| 1306 | embeds_orig_positions: list[int] = [ |
| 1307 | i for i, tok in enumerate(token_ids) if tok == placeholder_token_id |
| 1308 | ] |
| 1309 | if len(embeds_orig_positions) != len(prompt_embeds_tensors): |
| 1310 | raise ValueError( |
| 1311 | f"Expected {len(prompt_embeds_tensors)} prompt_embeds " |
| 1312 | f"placeholder tokens in the rendered prompt, found " |
| 1313 | f"{len(embeds_orig_positions)}." |
| 1314 | ) |
| 1315 | |
| 1316 | mm_updates = _build_prompt_embeds_updates( |
| 1317 | prompt_embeds_tensors, placeholder_token_id |
| 1318 | ) |
| 1319 | expanded = _expand_prompt_embeds_placeholders(token_ids, mm_updates) |
| 1320 | positions = _build_prompt_embeds_positions( |
| 1321 | expanded, len(prompt_embeds_tensors), mm_updates |
| 1322 | ) |
| 1323 | |
| 1324 | embeds_prompt = cast(EmbedsPrompt, prompt) |
| 1325 | embeds_prompt["prompt_token_ids"] = expanded |
| 1326 | full_embeds, is_token_ids_mask = _build_mixed_prompt_embeds( |
| 1327 | expanded, prompt_embeds_tensors, positions |
| 1328 | ) |
| 1329 | embeds_prompt["prompt_embeds"] = full_embeds |
| 1330 | embeds_prompt["prompt_is_token_ids"] = is_token_ids_mask |
| 1331 | |
| 1332 | def _apply_prompt_embeds_to_engine_input( |
| 1333 | self, |
| 1334 | engine_input: MultiModalInput, |
| 1335 | prompt_embeds_tensors: list[torch.Tensor], |
| 1336 | mm_updates: MultiModalPromptUpdates, |
| 1337 | ) -> None: |
| 1338 | """Augment `engine_input` in-place with a `prompt_embeds` modality. |
| 1339 | |
| 1340 | Mixed mode: called after `_process_multimodal` has already run on the |
| 1341 | pre-expanded token IDs (expansion was done in `_process_tokens` before |
| 1342 | calling `super()`). Locates the already-expanded `prompt_embeds` spans |
| 1343 | and adds `prompt_embeds` entries to `mm_kwargs`, `mm_hashes`, and |
| 1344 | `mm_placeholders`. |
| 1345 | """ |
| 1346 | # token_ids already contain the pre-expanded N-token spans. |
| 1347 | token_ids = list(engine_input["prompt_token_ids"]) |
| 1348 | |
| 1349 | positions = _build_prompt_embeds_positions( |
| 1350 | token_ids, len(prompt_embeds_tensors), mm_updates |
| 1351 | ) |
| 1352 | |
| 1353 | pe_kwargs_items: list[MultiModalKwargsItem] = [] |
| 1354 | pe_hashes: list[str] = [] |
| 1355 | pe_placeholders: list[PlaceholderRange] = [] |
| 1356 | mm_config = self.model_config.get_multimodal_config() |
| 1357 | for tensor, (start, length) in zip( |
| 1358 | prompt_embeds_tensors, positions, strict=True |
| 1359 | ): |
| 1360 | pe_kwargs_items.append( |
| 1361 | MultiModalKwargsItem( |
| 1362 | { |
| 1363 | "embedding": MultiModalFieldElem( |
| 1364 | data=tensor, |
| 1365 | field=MultiModalSharedField(batch_size=1), |
| 1366 | ) |
| 1367 | } |
| 1368 | ) |
| 1369 | ) |
| 1370 | pe_hashes.append( |
| 1371 | MultiModalHasher.hash_kwargs( |
| 1372 | mm_config.mm_hasher_algorithm, prompt_embeds=tensor |
| 1373 | ) |
| 1374 | ) |
| 1375 | # `is_embed=None` matches the existing image_embeds-style |
| 1376 | # "no encoder, just splice the tensor directly" semantics. |
| 1377 | pe_placeholders.append( |
| 1378 | PlaceholderRange(offset=start, length=length, is_embed=None) |
| 1379 | ) |
| 1380 | |
| 1381 | cast( |
| 1382 | MultiModalKwargsItems[MultiModalKwargsItem | None], |
| 1383 | engine_input["mm_kwargs"], |
| 1384 | )["prompt_embeds"] = pe_kwargs_items |
| 1385 | engine_input["mm_hashes"] = { |
| 1386 | **engine_input["mm_hashes"], |
| 1387 | "prompt_embeds": pe_hashes, |
| 1388 | } |
| 1389 | cast(dict, engine_input["mm_placeholders"])["prompt_embeds"] = pe_placeholders |