[None][perf] bound multimodal request CPU memory in trtllm-serve

docs/source/commands/trtllm-serve/trtllm-serve.rst

@@ -214,6 +214,36 @@ model.
          {"type": "audio_url", "audio_url": {"url": "https://example.com/audio.mp3"}}
      ]}
 
+Multimodal CPU Memory Limits
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Decoded media can need far more CPU memory than the request that carried it.
+Two prototype options bound this memory in each serving frontend process; both
+are disabled by default:
+
+* ``--max_multimodal_cpu_bytes_per_request`` caps one request. Request bodies
+  above the cap on ``/v1/chat/completions``, ``/v1/messages``, and
+  ``/v1/responses`` are rejected before they are parsed, including text-only
+  requests. For chat and Messages requests, the raw body plus the decoded media
+  and processed multimodal tensors must also stay under the cap; media decoding
+  stops as soon as it is exceeded. Rejected requests get HTTP 413.
+* ``--max_multimodal_cpu_bytes`` caps the requests with media that are decoded
+  and preprocessed at the same time. Each one reserves the per-request cap, so
+  at most ``max_multimodal_cpu_bytes // max_multimodal_cpu_bytes_per_request``
+  run concurrently and the rest wait. Text-only requests do not wait. If
+  ``--max_multimodal_cpu_bytes_per_request`` is not set, it defaults to
+  ``--max_multimodal_cpu_bytes``, so only one request with media is decoded and
+  preprocessed at a time.
+
+.. code-block:: bash
+
+   trtllm-serve Qwen/Qwen2-VL-7B-Instruct \
+       --max_multimodal_cpu_bytes 8GiB \
+       --max_multimodal_cpu_bytes_per_request 1GiB
+
+The limits count tensor, array, and image storage, not process RSS. Leave
+headroom for decoder temporaries and allocator overhead.
+
 
 
 Visual Generation Serving

tensorrt_llm/commands/serve.py

@@ -58,7 +58,8 @@
                                               parse_disagg_config_file,
                                               parse_metadata_server_config_file,
                                               validate_config_bool)
-from tensorrt_llm.llmapi.llm_args import MultimodalConfig, TorchLlmArgs
+from tensorrt_llm.llmapi.llm_args import (MultimodalConfig, TorchLlmArgs,
+                                          _parse_binary_byte_string)
 from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict
 from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr, split_mpi_env
 from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory,
@@ -1211,6 +1212,29 @@ def launch_visual_gen_server(
                   "\"num_frames\": 16}}' to enable audio extraction from "
                   "video files.",
                   status="prototype")
+@stability_option(
+    "--max_multimodal_cpu_bytes",
+    type=str,
+    default=None,
+    help=("Maximum CPU working set for concurrently preprocessing multimodal "
+          "requests in each serving frontend process. Each request with media "
+          "reserves --max_multimodal_cpu_bytes_per_request, so "
+          "max_multimodal_cpu_bytes // max_multimodal_cpu_bytes_per_request "
+          "of them run at a time and the rest wait. If "
+          "--max_multimodal_cpu_bytes_per_request is not set, it defaults to "
+          "this value and only one runs at a time. Accepts byte values such "
+          "as '8GiB'. Disabled by default."),
+    status="prototype")
+@stability_option(
+    "--max_multimodal_cpu_bytes_per_request",
+    type=str,
+    default=None,
+    help=("Maximum CPU storage for one multimodal request: its raw HTTP body "
+          "plus decoded media and processed multimodal tensors. Chat, "
+          "Messages, and Responses request bodies above this size are "
+          "rejected with HTTP 413. Accepts byte values such as '1GiB'. "
+          "Defaults to the total multimodal CPU limit when that limit is set."),
+    status="prototype")
 @stability_option("--video_pruning_rate",
                   type=float,
                   default=None,
@@ -1332,6 +1356,8 @@ def serve(
     enable_attention_dp: bool,
     disagg_cluster_uri: Optional[str],
     media_io_kwargs: Optional[str],
+    max_multimodal_cpu_bytes: Optional[str],
+    max_multimodal_cpu_bytes_per_request: Optional[str],
     agent_percentage: float,
     agent_types: Optional[str],
     video_pruning_rate: Optional[float],
@@ -1495,7 +1521,11 @@ def _serve_llm():
                 raise ValueError(f"Invalid JSON for media_io_kwargs: {e}")
 
         multimodal_server_config = MultimodalServerConfig(
-            media_io_kwargs=parsed_media_io_kwargs)
+            media_io_kwargs=parsed_media_io_kwargs,
+            max_cpu_bytes=_parse_binary_byte_string(max_multimodal_cpu_bytes),
+            max_cpu_bytes_per_request=_parse_binary_byte_string(
+                max_multimodal_cpu_bytes_per_request),
+        )
 
         if grpc:
             if num_serve_frontends != 1:
@@ -1522,6 +1552,10 @@ def _serve_llm():
                 server_role,
                 "disagg_cluster_config":
                 disagg_cluster_config,
+                "max_multimodal_cpu_bytes":
+                max_multimodal_cpu_bytes,
+                "max_multimodal_cpu_bytes_per_request":
+                max_multimodal_cpu_bytes_per_request,
             }
             for name, value in unsupported_args.items():
                 if value is not None:

tensorrt_llm/inputs/multimodal.py

@@ -823,6 +823,23 @@ def has_content(self) -> bool:
 @dataclass
 class MultimodalServerConfig():
     media_io_kwargs: Optional[dict] = None
+    max_cpu_bytes: Optional[int] = None
+    max_cpu_bytes_per_request: Optional[int] = None
+
+    def __post_init__(self) -> None:
+        if self.max_cpu_bytes is not None and self.max_cpu_bytes <= 0:
+            raise ValueError("max_cpu_bytes must be positive")
+        if (self.max_cpu_bytes_per_request is not None
+                and self.max_cpu_bytes_per_request <= 0):
+            raise ValueError("max_cpu_bytes_per_request must be positive")
+        if (self.max_cpu_bytes is not None
+                and self.max_cpu_bytes_per_request is None):
+            self.max_cpu_bytes_per_request = self.max_cpu_bytes
+        if (self.max_cpu_bytes is not None
+                and self.max_cpu_bytes_per_request is not None
+                and self.max_cpu_bytes_per_request > self.max_cpu_bytes):
+            raise ValueError(
+                "max_cpu_bytes_per_request cannot exceed max_cpu_bytes")
 
 
 def _update_hash(hasher, item: object) -> None:

tensorrt_llm/inputs/utils.py

@@ -5,6 +5,7 @@
 import base64
 import tempfile
 from collections import defaultdict
+from dataclasses import fields, is_dataclass
 from io import BytesIO
 from pathlib import Path
 from typing import (Any, Collection, Coroutine, Dict, List, Optional, Tuple,
@@ -14,7 +15,7 @@
 import numpy as np
 import soundfile
 import torch
-from PIL import Image
+from PIL import Image, ImageMode
 from torchvision.transforms import ToTensor
 from transformers import AutoProcessor, PreTrainedTokenizerBase, ProcessorMixin
 from transformers.utils import logging
@@ -47,6 +48,107 @@
 logger = logging.get_logger(__name__)
 
 
+class MultimodalDataTooLargeError(ValueError):
+    """A request's materialized multimodal inputs exceed its CPU limit."""
+
+
+def _is_shared_cpu_tensor_handle(value: dict) -> bool:
+    """Whether value is a CPU handle from `SharedTensorContainer.dump_to_dict`."""
+    return ("method_key" in value and "storage_size" in value
+            and "storage_dtype" in value)
+
+
+def _cpu_storage_bytes(
+        value: Any,
+        seen_storages: Optional[set[tuple[Any, ...]]] = None) -> int:
+    """Count unique CPU tensor, array, and image storage reachable from value."""
+    if seen_storages is None:
+        seen_storages = set()
+
+    if isinstance(value, torch.Tensor):
+        if value.device.type != "cpu":
+            return 0
+        storage = value.untyped_storage()
+        key = ("cpu", storage.data_ptr(), storage.nbytes())
+        if key in seen_storages:
+            return 0
+        seen_storages.add(key)
+        return storage.nbytes()
+
+    if isinstance(value, np.ndarray):
+        owner = value
+        while isinstance(owner.base, np.ndarray):
+            owner = owner.base
+        data_ptr = owner.__array_interface__["data"][0]
+        key = ("cpu", data_ptr, owner.nbytes)
+        if key in seen_storages:
+            return 0
+        seen_storages.add(key)
+        return owner.nbytes
+
+    if isinstance(value, Image.Image):
+        mode = ImageMode.getmode(value.mode)
+        size_bytes = (value.width * value.height * len(mode.bands) *
+                      np.dtype(mode.typestr).itemsize)
+        key = ("pil", id(value), size_bytes)
+        if key in seen_storages:
+            return 0
+        seen_storages.add(key)
+        return size_bytes
+
+    if isinstance(value, BaseModalityData):
+        return _cpu_storage_bytes(vars(value), seen_storages)
+    if isinstance(value, dict):
+        if _is_shared_cpu_tensor_handle(value):
+            dtype = getattr(
+                torch,
+                str(value["storage_dtype"]).removeprefix("torch."),
+                None,
+            )
+            if isinstance(dtype, torch.dtype):
+                size_bytes = int(value["storage_size"]) * dtype.itemsize
+                key = ("shared_cpu", value.get("storage_handle"), size_bytes)
+                if key in seen_storages:
+                    return 0
+                seen_storages.add(key)
+                return size_bytes
+        return sum(
+            _cpu_storage_bytes(item, seen_storages) for item in value.values())
+    if isinstance(value, (list, tuple)):
+        return sum(_cpu_storage_bytes(item, seen_storages) for item in value)
+    if is_dataclass(value) and not isinstance(value, type):
+        return sum(
+            _cpu_storage_bytes(getattr(value, field.name), seen_storages)
+            for field in fields(value))
+    return 0
+
+
+def _release_shared_cpu_tensors(value: Any) -> None:
+    """Free the shared memory behind CPU tensor handles reachable from value.
+
+    Exporting a CPU tensor as a handle (`MultimodalParams.to_handle`) keeps its
+    shared-memory segment alive until a consumer rebuilds the tensor. A request
+    dropped before any worker consumes its handles must rebuild them here, or
+    the segments stay in /dev/shm until the process exits. CUDA handles are
+    left alone.
+    """
+    if isinstance(value, dict):
+        if not _is_shared_cpu_tensor_handle(value):
+            for item in value.values():
+                _release_shared_cpu_tensors(item)
+            return
+        from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
+        try:
+            # The rebuild takes over the reference added by the export, so
+            # dropping the rebuilt view frees the segment.
+            SharedTensorContainer.from_dict(value).get_local_view()
+        except (KeyError, RuntimeError, ValueError) as e:
+            logger.warning(f"Failed to release a shared CPU tensor: {e}")
+    elif isinstance(value, (list, tuple)):
+        for item in value:
+            _release_shared_cpu_tensors(item)
+
+
 def load_base64_image(parsed_url: str) -> Image.Image:
     data_spec, data = parsed_url.path.split(",", 1)
     media_type, data_type = data_spec.split(";", 1)
@@ -375,6 +477,7 @@ def __init__(
         model_type: str,
         multimodal_server_config: Optional[MultimodalServerConfig] = None,
         request_media_io_kwargs: Optional[Dict[str, Dict[str, Any]]] = None,
+        initial_cpu_bytes: int = 0,
     ):
         self._model_type = model_type
         self._data = defaultdict[str, list](list)
@@ -394,6 +497,9 @@ def __init__(
         # Per-request override merged with the server default at media-load
         # time; see `BaseMediaIO.merge_kwargs` in `inputs/media_io.py`.
         self._request_media_io_kwargs = request_media_io_kwargs
+        if initial_cpu_bytes < 0:
+            raise ValueError("initial_cpu_bytes must be non-negative")
+        self._initial_cpu_bytes = initial_cpu_bytes
 
     @property
     def request_media_io_kwargs(self) -> Optional[Dict[str, Dict[str, Any]]]:
@@ -426,7 +532,12 @@ async def _retrieve(
                 out[modality].append(result)
             return dict(out)
 
-        # _data and _embeddings also gathered concurrently
+        max_bytes = self._multimodal_server_config.max_cpu_bytes_per_request
+        if max_bytes is not None:
+            return await self._retrieve_with_byte_limit(max_bytes)
+
+        # _data and _embeddings also gathered concurrently. Preserve this
+        # zero-overhead path when frontend byte admission is not configured.
         data_result, embed_result = await asyncio.gather(
             _retrieve(self._data), _retrieve(self._embeddings))
         return data_result, embed_result
@@ -452,6 +563,52 @@ async def discard_all_async(
             pending.clear()
         return None, None
 
+    async def _retrieve_with_byte_limit(
+        self, max_bytes: int
+    ) -> tuple[Optional[Dict[str, List[Any]]], Optional[Dict[str, List[Any]]]]:
+        """Collect completed items until their unique CPU storage hits the limit."""
+        pairs = [(False, modality, item)
+                 for modality, items in self._data.items() for item in items]
+        pairs.extend((True, modality, item)
+                     for modality, items in self._embeddings.items()
+                     for item in items)
+        if not pairs:
+            return None, None
+
+        async def _indexed_result(index: int,
+                                  item: Coroutine) -> tuple[int, Any]:
+            return index, await item
+
+        tasks = [
+            asyncio.create_task(_indexed_result(index, item))
+            for index, (_, _, item) in enumerate(pairs)
+        ]
+        results: list[Any] = [None] * len(tasks)
+        seen_storages: set[tuple[Any, ...]] = set()
+        resident_bytes = self._initial_cpu_bytes
+        try:
+            for completed in asyncio.as_completed(tasks):
+                index, result = await completed
+                resident_bytes += _cpu_storage_bytes(result, seen_storages)
+                if resident_bytes > max_bytes:
+                    raise MultimodalDataTooLargeError(
+                        "Multimodal inputs require at least "
+                        f"{resident_bytes} CPU bytes, exceeding the per-request "
+                        f"limit of {max_bytes} bytes")
+                results[index] = result
+        except BaseException:
+            for task in tasks:
+                task.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+
+        data_result: dict[str, list] = defaultdict(list)
+        embed_result: dict[str, list] = defaultdict(list)
+        for (is_embedding, modality, _), result in zip(pairs, results):
+            (embed_result
+             if is_embedding else data_result)[modality].append(result)
+        return (dict(data_result) or None, dict(embed_result) or None)
+
     def retrieve_all_sync(
         self
     ) -> tuple[Optional[Dict[str, List[Any]]], Optional[Dict[str, List[Any]]]]:

tensorrt_llm/serve/chat_utils.py

@@ -448,6 +448,7 @@ def parse_chat_messages_coroutines(
     request_media_io_kwargs: Optional[Dict[str, Dict[str, Any]]] = None,
     model_type_override: Optional[str] = None,
     resolve_media: bool = True,
+    initial_cpu_bytes: int = 0,
 ) -> Tuple[List[ConversationMessage], Coroutine[Any, Any, tuple[Optional[Dict[
         str, List[Any]]], Optional[Dict[str, List[Any]]]]], list[dict[str,
                                                                       int]]]:
@@ -468,6 +469,9 @@ def parse_chat_messages_coroutines(
         resolve_media: When False the returned coroutine closes the
             pending loads instead of running them and yields
             `(None, None)`.
+        initial_cpu_bytes: CPU bytes already held by the request (e.g.
+            its raw body), charged against
+            `multimodal_server_config.max_cpu_bytes_per_request`.
 
     Returns:
         `(conversation, mm_coroutine, mm_placeholder_counts)` where
@@ -481,7 +485,8 @@ def parse_chat_messages_coroutines(
     mm_data_tracker = MultimodalDataTracker(
         model_type,
         multimodal_server_config,
-        request_media_io_kwargs=request_media_io_kwargs)
+        request_media_io_kwargs=request_media_io_kwargs,
+        initial_cpu_bytes=initial_cpu_bytes)
 
     # Determine content format to decide placeholder strategy.
     #