@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8c9da6be22

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Core] Simplify mm processing cache (#22457)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Aug 8, 2025 at 00:47 UTC 8c9da6be229336a769d9c904415daaa250824c89

4 files changed +97 -206

vllm/model_executor/models/qwen2_5_omni_thinker.py

+6 -6

@@ -431,7 +431,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(

431 tokenization_kwargs: Mapping[str, object],

432 *,

433 enable_hf_prompt_update: bool,

434 - ) -> tuple[list[int], MultiModalKwargs, bool]:

434 + ) -> tuple[list[int], BatchFeature, bool]:

435 """

436 Qwen2.5-Omni reimplements this function to handle text only.

437 """

@@ -448,20 +448,20 @@ class Qwen2_5OmniThinkerMultiModalProcessor(

448 else:

449 prompt_ids = self._apply_hf_processor_tokens_only(prompt)

450

451 - mm_kwargs = self._apply_hf_processor_mm_only(

451 + mm_processed_data = self._apply_hf_processor_mm_only(

452 mm_items=mm_items,

453 hf_processor_mm_kwargs=hf_processor_mm_kwargs,

454 tokenization_kwargs=tokenization_kwargs,

455 )

456

457 - return prompt_ids, mm_kwargs, False

457 + return prompt_ids, mm_processed_data, False

458

459 def _apply_hf_processor_mm_only(

460 self,

461 mm_items: MultiModalDataItems,

462 hf_processor_mm_kwargs: Mapping[str, object],

463 tokenization_kwargs: Mapping[str, object],

464 - ) -> MultiModalKwargs:

464 + ) -> BatchFeature:

465 """

466 Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.

467 """

@@ -473,14 +473,14 @@ class Qwen2_5OmniThinkerMultiModalProcessor(

473 assert "audio" in mm_counts

474 mm_counts["audio"] -= mm_counts["video"]

475

476 - _, mm_kwargs, _ = self._apply_hf_processor_text_mm(

476 + _, mm_processed_data, _ = self._apply_hf_processor_text_mm(

477 prompt_text=self.dummy_inputs.get_dummy_text(mm_counts),

478 mm_items=mm_items,

479 hf_processor_mm_kwargs=hf_processor_mm_kwargs,

480 tokenization_kwargs=tokenization_kwargs,

481 )

482

483 - return mm_kwargs

483 + return mm_processed_data

484

485 def _validate_mm_placeholders(

486 self,

vllm/model_executor/models/transformers.py

+3 -2

@@ -22,7 +22,8 @@ from typing import Literal, Optional, Union

22 import regex as re

23 import torch

24 from torch import nn

25 -from transformers import AutoModel, PretrainedConfig, PreTrainedModel

25 +from transformers import (AutoModel, BatchFeature, PretrainedConfig,

26 + PreTrainedModel)

27 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

29 from vllm.attention import Attention

@@ -269,7 +270,7 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):

270 mm_items: MultiModalDataItems,

271 hf_processor_mm_kwargs: Mapping[str, object],

272 tokenization_kwargs: Mapping[str, object],

272 - ):

273 + ) -> tuple[list[int], BatchFeature, bool]:

274 """

275 Apply the HF processor on the prompt text and multi-modal data

276 together.

vllm/multimodal/processing.py

+69 -179

@@ -18,7 +18,7 @@ from vllm.inputs import InputProcessingContext

18 from vllm.logger import init_logger

19 from vllm.transformers_utils.tokenizer import (AnyTokenizer, decode_tokens,

20 encode_tokens)

21 -from vllm.utils import GiB_bytes, flatten_2d_lists, full_groupby

21 +from vllm.utils import flatten_2d_lists, full_groupby

23 from .cache import MultiModalCache

24 from .hasher import MultiModalHasher

@@ -887,120 +887,19 @@ def find_mm_placeholders(

887 return dict(full_groupby_modality(it))

888

889

890 -class ProcessingCacheOptionalItem(NamedTuple):

891 - key: str

892 - value: Optional[MultiModalKwargsItem]

893 -

894 -

895 -class ProcessingCacheItem(NamedTuple):

896 - key: str

897 - value: MultiModalKwargsItem

898 -

899 -

890 class ProcessingCache(MultiModalCache):

891

902 - def __init__(

903 - self,

904 - capacity_gb: float,

905 - *,

906 - debug_cache_hit_ratio_steps: Optional[int] = None,

907 - ) -> None:

892 + def __init__(self, capacity_gb: float) -> None:

893 super().__init__()

894

910 - self.debug_cache_hit_ratio_steps = debug_cache_hit_ratio_steps

911 - self.debug_cache_hits = 0

912 - self.debug_cache_total = 0

913 -

914 - self._cache = self.get_lru_cache(

915 - capacity_gb,

916 - MultiModalKwargsItem,

917 - debug=bool(debug_cache_hit_ratio_steps),

918 - )

919 -

920 - def _maybe_log_cache_stats(self) -> None:

921 - steps = self.debug_cache_hit_ratio_steps

922 - if not steps:

923 - return

924 -

925 - total = self.debug_cache_total

926 - if total > 0 and total % steps == 0:

927 - logger.debug("ProcessingCache: hit_ratio = %.2f",

928 - self.debug_cache_hits / total)

929 - logger.debug("ProcessingCache: size = %.2f / %.2f GiB",

930 - self._cache.currsize / GiB_bytes,

931 - self._cache.maxsize / GiB_bytes)

932 -

933 - def get(

934 - self,

935 - model_id: str,

936 - modality: str,

937 - input_item: object,

938 - input_kwargs: Mapping[str, object],

939 - ) -> Optional[MultiModalKwargsItem]:

940 - """

941 - Get a processed multi-modal item from the cache

942 - according to its dependencies, including:

943 -

944 - - The model ID

945 - - The modality of the item

946 - - The original data item passed to the HF processor

947 - - The configuration options of the HF processor

948 - """

949 - self._maybe_log_cache_stats()

895 + self._cache = self.get_lru_cache(capacity_gb, MultiModalKwargsItem)

896

951 - cache_key = MultiModalHasher.hash_kwargs(model_id=model_id,

952 - **{modality: input_item},

953 - **input_kwargs)

897 + self.get = self._cache.get

898 + self.put = self._cache.put

899 + self.reset = self._cache.clear

900

955 - if self.debug_cache_hit_ratio_steps:

956 - if cache_key in self._cache:

957 - self.debug_cache_hits += 1

901

959 - self.debug_cache_total += 1

960 -

961 - return self._cache.get(cache_key)

962 -

963 - def get_item(

964 - self,

965 - model_id: str,

966 - modality: str,

967 - input_item: object,

968 - input_kwargs: Mapping[str, object],

969 - ) -> ProcessingCacheOptionalItem:

970 - cache_key = MultiModalHasher.hash_kwargs(model_id=model_id,

971 - **{modality: input_item},

972 - **input_kwargs)

973 -

974 - return ProcessingCacheOptionalItem(

975 - key=cache_key,

976 - value=self._cache.get(cache_key),

977 - )

978 -

979 - def put(

980 - self,

981 - model_id: str,

982 - modality: str,

983 - input_item: object,

984 - input_kwargs: Mapping[str, object],

985 - output_kwargs: MultiModalKwargsItem,

986 - ) -> None:

987 - """

988 - Put a processed multi-modal item into the cache

989 - according to its dependencies

990 - (see [`get`][vllm.multimodal.processing.ProcessingCache.get]).

991 - """

992 - cache_key = MultiModalHasher.hash_kwargs(model_id=model_id,

993 - **{modality: input_item},

994 - **input_kwargs)

995 - self._cache[cache_key] = output_kwargs

996 -

997 - def put_item(self, item: ProcessingCacheItem) -> None:

998 - self._cache[item.key] = item.value

999 -

1000 - def reset(self) -> bool:

1001 - self._cache.clear()

1002 -

1003 - return True

902 +_CacheItemOrHash = Union[MultiModalKwargsItem, str]

903

904

905 class BaseProcessingInfo:

@@ -1279,7 +1178,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1178 mm_items: MultiModalDataItems,

1179 hf_processor_mm_kwargs: Mapping[str, object],

1180 tokenization_kwargs: Mapping[str, object],

1282 - ) -> tuple[list[int], MultiModalKwargs, bool]:

1181 + ) -> tuple[list[int], "BatchFeature", bool]:

1182 """

1183 Apply the HF processor on the prompt text and multi-modal data

1184 together.

@@ -1298,11 +1197,6 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1197

1198 prompt_ids, = processed_data.pop("input_ids").tolist()

1199

1301 - mm_kwargs = MultiModalKwargs.from_hf_inputs(

1302 - processed_data,

1303 - self._get_mm_fields_config(processed_data, hf_processor_mm_kwargs),

1304 - )

1305 -

1200 is_update_applied = self._hf_processor_applies_updates(

1201 prompt_text=prompt_text,

1202 mm_items=mm_items,

@@ -1310,11 +1204,13 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1204 tokenization_kwargs=tokenization_kwargs,

1205 )

1206

1313 - return prompt_ids, mm_kwargs, is_update_applied

1207 + return prompt_ids, processed_data, is_update_applied

1208

1209 def _apply_hf_processor_text_only(

1316 - self, prompt_text: str,

1317 - tokenization_kwargs: Mapping[str, object]) -> list[int]:

1210 + self,

1211 + prompt_text: str,

1212 + tokenization_kwargs: Mapping[str, object],

1213 + ) -> list[int]:

1214 """

1215 Apply the HF processor on the prompt text only.

1216

@@ -1353,7 +1249,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1249 mm_items: MultiModalDataItems,

1250 hf_processor_mm_kwargs: Mapping[str, object],

1251 tokenization_kwargs: Mapping[str, object],

1356 - ) -> MultiModalKwargs:

1252 + ) -> "BatchFeature":

1253 """

1254 Apply the HF processor on the multi-modal data only.

1255

@@ -1364,14 +1260,14 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1260 """

1261 mm_counts = mm_items.get_all_counts()

1262

1367 - _, mm_kwargs, _ = self._apply_hf_processor_text_mm(

1263 + _, mm_processed_data, _ = self._apply_hf_processor_text_mm(

1264 prompt_text=self.dummy_inputs.get_dummy_text(mm_counts),

1265 mm_items=mm_items,

1266 hf_processor_mm_kwargs=hf_processor_mm_kwargs,

1267 tokenization_kwargs=tokenization_kwargs,

1268 )

1269

1374 - return mm_kwargs

1270 + return mm_processed_data

1271

1272 def _apply_hf_processor_main(

1273 self,

@@ -1381,7 +1277,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1277 tokenization_kwargs: Mapping[str, object],

1278 *,

1279 enable_hf_prompt_update: bool,

1384 - ) -> tuple[list[int], MultiModalKwargs, bool]:

1280 + ) -> tuple[list[int], "BatchFeature", bool]:

1281 """

1282 Apply the HF processor on the prompt text and multi-modal data.

1283

@@ -1407,52 +1303,46 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1303 else:

1304 prompt_ids = self._apply_hf_processor_tokens_only(prompt)

1305

1410 - mm_kwargs = self._apply_hf_processor_mm_only(

1306 + mm_processed_data = self._apply_hf_processor_mm_only(

1307 mm_items=mm_items,

1308 hf_processor_mm_kwargs=hf_processor_mm_kwargs,

1309 tokenization_kwargs=tokenization_kwargs,

1310 )

1311

1416 - return prompt_ids, mm_kwargs, False

1312 + return prompt_ids, mm_processed_data, False

1313

1314 def _get_cache_missing_items(

1315 self,

1316 cache: ProcessingCache,

1317 mm_data_items: MultiModalDataItems,

1422 - hf_processor_mm_kwargs: Mapping[str, object],

1423 - tokenization_kwargs: Mapping[str, object],

1424 - ) -> tuple[dict[str, list[ProcessingCacheOptionalItem]], dict[

1425 - str, list[object]]]:

1426 - model_id = self.info.model_id

1427 -

1428 - mm_cache_items = {

1429 - modality: [

1430 - cache.get_item(

1431 - model_id, modality, item,

1432 - dict(**hf_processor_mm_kwargs, **tokenization_kwargs))

1433 - for item in items

1434 - ]

1435 - for modality, items in mm_data_items.items()

1318 + mm_hashes: MultiModalHashes,

1319 + ) -> tuple[dict[str, list[_CacheItemOrHash]], MultiModalDataItems]:

1320 + mm_cache_items_or_hashes: dict[str, list[_CacheItemOrHash]] = {

1321 + modality: [(h if (v := cache.get(h)) is None else v)

1322 + for h in hashes]

1323 + for modality, hashes in mm_hashes.items()

1324 }

1325

1326 mm_missing_idxs = {

1327 modality: [

1440 - idx for idx, item in enumerate(cache_items)

1441 - if item.value is None

1328 + idx for idx, item_or_hash in enumerate(items_or_hashes)

1329 + if isinstance(item_or_hash, str)

1330 ]

1443 - for modality, cache_items in mm_cache_items.items()

1331 + for modality, items_or_hashes in mm_cache_items_or_hashes.items()

1332 }

1333 mm_missing_data = {

1334 modality: [mm_data_items[modality][idx] for idx in idxs]

1335 for modality, idxs in mm_missing_idxs.items()

1336 }

1337

1450 - return mm_cache_items, mm_missing_data

1338 + return mm_cache_items_or_hashes, self._to_mm_items(mm_missing_data)

1339

1340 def _hash_mm_items(

1453 - self, mm_items: MultiModalDataItems,

1454 - hf_processor_mm_kwargs: Mapping[str, object],

1455 - tokenization_kwargs: Mapping[str, object]) -> MultiModalHashes:

1341 + self,

1342 + mm_items: MultiModalDataItems,

1343 + hf_processor_mm_kwargs: Mapping[str, object],

1344 + tokenization_kwargs: Mapping[str, object],

1345 + ) -> MultiModalHashes:

1346 """Create MM hashes to be returned (only used in V1)."""

1347 model_id = self.info.model_id

1348

@@ -1470,34 +1360,25 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1360 def _merge_mm_kwargs(

1361 self,

1362 cache: ProcessingCache,

1473 - mm_cache_items: dict[str, list[ProcessingCacheOptionalItem]],

1474 - mm_missing_data: dict[str, list[object]],

1363 + mm_cache_items_or_hashes: dict[str, list[_CacheItemOrHash]],

1364 mm_missing_kwargs: MultiModalKwargs,

1476 - ) -> dict[str, list[ProcessingCacheItem]]:

1477 - mm_missing_next_idx = {modality: 0 for modality in mm_missing_data}

1365 + ) -> dict[str, list[MultiModalKwargsItem]]:

1366 + mm_missing_next_idx = defaultdict[str, int](lambda: 0)

1367

1479 - merged_items = defaultdict[str, list[ProcessingCacheItem]](list)

1480 - for modality, cache_items in mm_cache_items.items():

1481 - for cache_item in cache_items:

1482 - if cache_item.value is None:

1368 + merged_items = defaultdict[str, list[MultiModalKwargsItem]](list)

1369 + for modality, items_or_hashes in mm_cache_items_or_hashes.items():

1370 + for item_or_hash in items_or_hashes:

1371 + if isinstance(item_or_hash, str):

1372 kw_item = mm_missing_kwargs.get_item(

1373 modality,

1374 mm_missing_next_idx[modality],

1375 )

1487 - cache_item_new = ProcessingCacheItem(

1488 - key=cache_item.key,

1489 - value=kw_item,

1490 - )

1491 -

1492 - cache.put_item(cache_item_new)

1376 + cache.put(item_or_hash, kw_item)

1377 mm_missing_next_idx[modality] += 1

1378 else:

1495 - cache_item_new = ProcessingCacheItem(

1496 - key=cache_item.key,

1497 - value=cache_item.value,

1498 - )

1379 + kw_item = item_or_hash

1380

1500 - merged_items[modality].append(cache_item_new)

1381 + merged_items[modality].append(kw_item)

1382

1383 return dict(merged_items)

1384

@@ -1512,7 +1393,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1393 ) -> tuple[list[int], MultiModalKwargs, Optional[MultiModalHashes], bool]:

1394 (

1395 prompt_ids,

1515 - mm_kwargs,

1396 + mm_processed_data,

1397 is_update_applied,

1398 ) = self._apply_hf_processor_main(

1399 prompt=prompt,

@@ -1522,6 +1403,12 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1403 enable_hf_prompt_update=True,

1404 )

1405

1406 + mm_kwargs = MultiModalKwargs.from_hf_inputs(

1407 + mm_processed_data,

1408 + self._get_mm_fields_config(mm_processed_data,

1409 + hf_processor_mm_kwargs),

1410 + )

1411 +

1412 mm_hashes = (self._hash_mm_items(mm_data_items, hf_processor_mm_kwargs,

1413 tokenization_kwargs)

1414 if return_mm_hashes else None)

@@ -1553,49 +1440,52 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1440 return_mm_hashes=return_mm_hashes,

1441 )

1442

1443 + mm_hashes = self._hash_mm_items(mm_data_items, hf_processor_mm_kwargs,

1444 + tokenization_kwargs)

1445 (

1557 - mm_cache_items,

1558 - mm_missing_data,

1446 + mm_cache_items_or_hashes,

1447 + mm_missing_data_items,

1448 ) = self._get_cache_missing_items(

1449 cache=cache,

1450 mm_data_items=mm_data_items,

1562 - hf_processor_mm_kwargs=hf_processor_mm_kwargs,

1563 - tokenization_kwargs=tokenization_kwargs,

1451 + mm_hashes=mm_hashes,

1452 )

1453

1454 + mm_hashes_to_return = mm_hashes if return_mm_hashes else None

1455 +

1456 # NOTE: `prompt` does not correspond to `mm_missing_data_items`,

1457 # so we can't apply prompt updates until the new multimodal

1458 # items are combined with the cached multimodal items

1459 (

1460 prompt_ids,

1571 - mm_missing_kwargs,

1461 + mm_missing_processed_data,

1462 is_update_applied,

1463 ) = self._apply_hf_processor_main(

1464 prompt=prompt,

1575 - mm_items=self._to_mm_items(mm_missing_data),

1465 + mm_items=mm_missing_data_items,

1466 hf_processor_mm_kwargs=hf_processor_mm_kwargs,

1467 tokenization_kwargs=tokenization_kwargs,

1468 enable_hf_prompt_update=False,

1469 )

1470

1471 + mm_missing_kwargs = MultiModalKwargs.from_hf_inputs(

1472 + mm_missing_processed_data,

1473 + self._get_mm_fields_config(mm_missing_processed_data,

1474 + hf_processor_mm_kwargs),

1475 + )

1476 +

1477 mm_cache_items_merged = self._merge_mm_kwargs(

1478 cache,

1583 - mm_cache_items=mm_cache_items,

1584 - mm_missing_data=mm_missing_data,

1479 + mm_cache_items_or_hashes=mm_cache_items_or_hashes,

1480 mm_missing_kwargs=mm_missing_kwargs,

1481 )

1482

1483 mm_kwargs = MultiModalKwargs.from_items([

1589 - item.value for cache_items in mm_cache_items_merged.values()

1484 + item for cache_items in mm_cache_items_merged.values()

1485 for item in cache_items

1486 ])

1487

1593 - mm_hashes = {

1594 - modality: [item.key for item in cache_items]

1595 - for modality, cache_items in mm_cache_items_merged.items()

1596 - } if return_mm_hashes else None

1597 -

1598 - return prompt_ids, mm_kwargs, mm_hashes, is_update_applied

1488 + return prompt_ids, mm_kwargs, mm_hashes_to_return, is_update_applied

1489

1490 def _bind_and_group_updates(

1491 self,

vllm/v1/serial_utils.py

+19 -19

@@ -312,25 +312,25 @@ class MsgpackDecoder:

312 return arr.view(torch_dtype).view(shape)

313

314 def _decode_mm_items(self, obj: list) -> list[MultiModalKwargsItem]:

315 - decoded_items = []

316 - for item in obj:

317 - elems = []

318 - for v in item:

319 - v["data"] = self._decode_nested_tensors(v["data"])

320 - # Reconstruct the field processor using MultiModalFieldConfig

321 - factory_meth_name, *field_args = v["field"]

322 - factory_meth = getattr(MultiModalFieldConfig,

323 - factory_meth_name)

324 -

325 - # Special case: decode the union "slices" field of

326 - # MultiModalFlatField

327 - if factory_meth_name == "flat":

328 - field_args[0] = self._decode_nested_slices(field_args[0])

329 -

330 - v["field"] = factory_meth(None, *field_args).field

331 - elems.append(MultiModalFieldElem(**v))

332 - decoded_items.append(MultiModalKwargsItem.from_elems(elems))

333 - return decoded_items

315 + return [self._decode_mm_item(v) for v in obj]

316 +

317 + def _decode_mm_item(self, obj: list) -> MultiModalKwargsItem:

318 + return MultiModalKwargsItem.from_elems(

319 + [self._decode_mm_field_elem(v) for v in obj])

320 +

321 + def _decode_mm_field_elem(self, obj: dict) -> MultiModalFieldElem:

322 + obj["data"] = self._decode_nested_tensors(obj["data"])

323 + # Reconstruct the field processor using MultiModalFieldConfig

324 + factory_meth_name, *field_args = obj["field"]

325 + factory_meth = getattr(MultiModalFieldConfig, factory_meth_name)

326 +

327 + # Special case: decode the union "slices" field of

328 + # MultiModalFlatField

329 + if factory_meth_name == "flat":

330 + field_args[0] = self._decode_nested_slices(field_args[0])

331 +

332 + obj["field"] = factory_meth(None, *field_args).field

333 + return MultiModalFieldElem(**obj)

334

335 def _decode_nested_tensors(self, obj: Any) -> NestedTensors:

336 if isinstance(obj, (int, float)):