[None][perf] Bound OpenEngine image preparation and reuse prefill state

tensorrt_llm/_torch/models/modeling_qwen2vl.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import base64
 import copy
 import math
 import re
@@ -145,6 +146,61 @@ def _prepare_qwen_vl_mrope_config(
 class Qwen2VLInputProcessorBase(BaseMultimodalInputProcessor,
                                 BaseMultimodalDummyInputsBuilder):
 
+    def export_decode_state(
+            self, multimodal_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        config = multimodal_data.get("mrope_config")
+        if config is None:
+            return None
+        return {
+            "mrope_config": {
+                name: {
+                    "shape":
+                    list(config[name].shape),
+                    "data":
+                    base64.b64encode(config[name].detach().cpu().numpy().astype(
+                        "<i8").tobytes()).decode("ascii"),
+                }
+                for name in ("mrope_position_ids", "mrope_position_deltas")
+            }
+        }
+
+    def restore_decode_state(self, state: Dict[str, Any],
+                             prompt_length: int) -> Dict[str, Any]:
+        if set(state) != {"mrope_config"}:
+            raise ValueError("invalid Qwen decode state")
+        config = state["mrope_config"]
+        if not isinstance(config, dict) or set(config) != {
+                "mrope_position_ids", "mrope_position_deltas"
+        }:
+            raise ValueError("invalid Qwen mRoPE fields")
+        tensors = {}
+
+        for name, value in config.items():
+            expected_shape = (
+                3, 1, prompt_length) if name == "mrope_position_ids" else (1, 1)
+            if not isinstance(value, dict) or set(value) != {"shape", "data"}:
+                raise ValueError(f"invalid Qwen {name} descriptor")
+            shape = value["shape"]
+            if (not isinstance(shape, list)
+                    or any(type(size) not in (int, float) for size in shape)
+                    or shape != list(expected_shape)):
+                raise ValueError(f"invalid Qwen {name} shape")
+            expected_bytes = math.prod(expected_shape) * 8
+            encoded = value["data"]
+            if (not isinstance(encoded, str)
+                    or len(encoded) > ((expected_bytes + 2) // 3) * 4):
+                raise ValueError(f"invalid Qwen {name} encoded length")
+            try:
+                raw = base64.b64decode(encoded, validate=True)
+            except (ValueError, UnicodeError) as error:
+                raise ValueError(f"invalid Qwen {name} encoding") from error
+            if len(raw) != expected_bytes:
+                raise ValueError(f"invalid Qwen {name} byte length")
+            array = np.frombuffer(raw, dtype="<i8").astype(
+                np.int64).reshape(expected_shape)
+            tensors[name] = torch.from_numpy(array)
+        return {"mrope_config": tensors}
+
     def __init__(self,
                  model_path: str,
                  config: PretrainedConfig,

tensorrt_llm/_torch/shared_tensor/shared_tensor.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import base64
 import itertools
 import logging
@@ -6,12 +9,7 @@
 from typing import Any, Callable, Dict, Optional, Tuple
 
 import torch
-from torch.multiprocessing import get_sharing_strategy, set_sharing_strategy
-from torch.multiprocessing.reductions import (rebuild_cuda_tensor,
-                                              rebuild_storage_filename,
-                                              rebuild_tensor,
-                                              rebuild_typed_storage,
-                                              reduce_storage, reduce_tensor)
+from torch.multiprocessing import reductions
 
 logger = logging.getLogger(__name__)
 
@@ -88,8 +86,8 @@ class _SharedTensorRebuildMethodRegistry:
     def initialize(cls):
         """Initialize the registry with common rebuild methods."""
         # Register common methods with fixed keys
-        cls._registry[cls.REBUILD_CUDA] = rebuild_cuda_tensor
-        cls._registry[cls.REBUILD_CPU] = rebuild_tensor
+        cls._registry[cls.REBUILD_CUDA] = reductions.rebuild_cuda_tensor
+        cls._registry[cls.REBUILD_CPU] = reductions.rebuild_tensor
 
     @classmethod
     def register(cls, method: Callable) -> int:
@@ -101,9 +99,9 @@ def register(cls, method: Callable) -> int:
         Returns:
             The numeric key assigned to the method
         """
-        if method == rebuild_cuda_tensor:
+        if method == reductions.rebuild_cuda_tensor:
             return cls.REBUILD_CUDA
-        if method == rebuild_tensor:
+        if method == reductions.rebuild_tensor:
             return cls.REBUILD_CPU
         raise NotImplementedError("Other rebuild methods are not supported yet")
 
@@ -326,9 +324,9 @@ def dict_to_cpu_handle(tensor_info: Dict[str, Any]) -> Tuple:
                                 storage_handle, tensor_info['storage_size'],
                                 str_to_torch_dtype(
                                     tensor_info['storage_dtype']))
-            storage = rebuild_storage_filename(*storage_metadata)
+            storage = reductions.rebuild_storage_filename(*storage_metadata)
             if not isinstance(storage, torch.storage.TypedStorage):
-                storage = rebuild_typed_storage(
+                storage = reductions.rebuild_typed_storage(
                     storage, str_to_torch_dtype(tensor_info['storage_dtype']))
 
             meta_data = (tensor_info['tensor_storage_offset'],
@@ -371,7 +369,7 @@ def from_tensor(cls,
                 "local_token": token,
                 "producer_pid": os.getpid()
             })
-        rebuild_method, tensor_handle = reduce_tensor(tensor)
+        rebuild_method, tensor_handle = reductions.reduce_tensor(tensor)
         method_key = _SharedTensorRebuildMethodRegistry.register(rebuild_method)
         return cls(method_key, tensor_handle)
 
@@ -433,6 +431,18 @@ def get_local_view(self) -> torch.Tensor:
             self.method_key)
         return rebuild_method(*self.tensor_handle)
 
+    @classmethod
+    def release_dict(cls, tensor_info: Dict[str, Any]) -> None:
+        """Consume an exported handle that will not be delivered to a consumer."""
+        if tensor_info[
+                'method_key'] == _SharedTensorRebuildMethodRegistry.REBUILD_CUDA:
+            torch.UntypedStorage._release_ipc_counter(
+                base64.b64decode(tensor_info['ref_counter_handle']),
+                tensor_info['ref_counter_offset'],
+                device=tensor_info['storage_device'])
+        else:
+            cls.from_dict(tensor_info).get_local_view()
+
     def dump_to_dict(self) -> Dict[str, Any]:
         """Convert this container to a dictionary for direct IPC (Producer side).
 
@@ -452,18 +462,16 @@ def dump_to_dict(self) -> Dict[str, Any]:
             tensor_dict = SharedTensorContainer.cuda_handle_to_dict(
                 self.tensor_handle)
         elif self.method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CPU:
-            sharing_strategy = get_sharing_strategy()
-            # Here we use file_system sharing strategy to make it serializable between two non-python independent processes
-            set_sharing_strategy("file_system")
             storage = self.tensor_handle[1]
             meta_data = self.tensor_handle[2]
-            storage_handle = reduce_storage(storage)
-            # restore the original sharing strategy
-            set_sharing_strategy(sharing_strategy)
-            # exclude the first element which is the type of the storage
-            storage_metadata = storage_handle[-1][1:]
+            # Use filename sharing directly without changing PyTorch's process-wide strategy.
+            storage_metadata = storage._share_filename_cpu_() + (
+                storage.dtype, )
             tensor_dict = SharedTensorContainer.cpu_handle_to_dict(
                 meta_data, storage_metadata)
+            reductions.shared_cache[
+                storage_metadata[1]] = reductions.StorageWeakRef(storage)
+            storage._shared_incref()
         else:
             raise ValueError(f"Unsupported tensor device: {self.method_key}")
 

tensorrt_llm/grpc/openengine/README.md

@@ -30,7 +30,30 @@ The `Inference.Generate` RPC loads the selected model through TensorRT-LLM's PyT
 
 Clients must continuously consume the response stream. If response delivery remains stalled for 30 seconds, the server aborts the engine request and terminates the stream with a retryable overload error.
 
-Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, multimodal media, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`.
+Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, unsupported media modalities, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`.
+
+### Media preparation
+
+Image inputs are accepted when the loaded processor advertises image support.
+The prompt must follow that processor's placeholder policy. Media decoding and
+preprocessing use separate worker pools. Set
+`TRTLLM_OPENENGINE_MEDIA_PREPARATION_LIMIT` to a positive integer to bound the
+number of requests admitted before decoding (default: 4, per frontend process).
+Cancellation retains admission until outstanding worker threads finish. This
+does not bound encoded requests waiting for admission or engine KV-cache memory.
+
+Compatible token-ID processors reuse prompt-validation tokens. Other processors
+retain their text-prompt path. For Qwen vision models, context workers also attach
+versioned prepared prompt, position, and multimodal cache-identity metadata to
+`KvSessionRef.attributes_struct` as a bounded JSON string. A matching decode
+worker can reuse it without fetching or preparing the media again. Missing or
+unknown-version metadata uses the ordinary preparation path. The handoff contains
+bounded CPU metadata, not cross-process tensor handles; encoder disaggregation
+is a separate feature.
+
+Both workers must use matching model/tokenizer revisions and processor settings.
+Input, model, and processor-setting fingerprints detect mismatches; they do not
+authenticate caller-provided metadata or verify model weights.
 
 ### Disaggregated serving
 

tensorrt_llm/grpc/openengine/capabilities.py

@@ -1,11 +1,17 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""What this engine can actually do, defined once for Generate and Control.
+"""Shared capability resolution for OpenEngine Control and Generate."""
 
-Control advertises from this table and Generate enforces it, so a client can
-never discover a capability that Generate then rejects.
-"""
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Any, Collection, Mapping
+
+from tensorrt_llm.inputs.model_config import resolve_top_level_model_type
+from tensorrt_llm.inputs.registry import (
+    MULTIMODAL_PLACEHOLDER_REGISTRY,
+    BaseMultimodalInputProcessor,
+)
 
 # Guides each grammar backend can build, keyed by the `GuidedDecoding` oneof
 # field name. Control advertises from this table and Generate enforces it, so
@@ -18,6 +24,30 @@
 }
 
 
+@dataclass(frozen=True)
+class MediaCapabilities:
+    """Media support shared by Control and Generate."""
+
+    placeholder_templates: Mapping[str, str]
+
+    def __post_init__(self) -> None:
+        object.__setattr__(
+            self, "placeholder_templates", MappingProxyType(dict(self.placeholder_templates))
+        )
+
+    def __hash__(self) -> int:
+        return hash(frozenset(self.placeholder_templates.items()))
+
+    @property
+    def modalities(self) -> frozenset[str]:
+        """Modalities supported by both the model and request mapper."""
+        return frozenset(self.placeholder_templates)
+
+    def placeholder(self, modality: str, item_number: int) -> str:
+        """Format the model placeholder for a one-based media item."""
+        return self.placeholder_templates[modality].format(item_number)
+
+
 def supported_guides(backend: str | None) -> frozenset[str]:
     """Guides `backend` can build; empty when none can be.
 
@@ -32,7 +62,47 @@ def supported_guides(backend: str | None) -> frozenset[str]:
     return GUIDE_SUPPORT_BY_BACKEND.get(backend.lower(), frozenset())
 
 
+def _loaded_model_type(llm: Any) -> str:
+    """Return the top-level model type for the loaded input processor."""
+    return resolve_top_level_model_type(llm.input_processor.config)
+
+
+def resolve_media_capabilities(
+    llm: Any, implemented_modalities: Collection[str]
+) -> MediaCapabilities:
+    """Intersect model support with implemented request handlers.
+
+    A live multimodal processor prevents registry-only false positives.
+    """
+    input_processor = getattr(llm, "input_processor", None)
+    if not isinstance(input_processor, BaseMultimodalInputProcessor):
+        return MediaCapabilities({})
+    if getattr(getattr(llm, "args", None), "disable_mm_encoder", False):
+        return MediaCapabilities({})
+    model_type = _loaded_model_type(llm)
+    if not model_type:
+        return MediaCapabilities({})
+
+    placeholders = {
+        modality: MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality)
+        for modality in implemented_modalities
+        if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality)
+    }
+    # Fail startup before Control advertises a malformed template.
+    for modality, template in placeholders.items():
+        try:
+            template.format(1)
+        except (IndexError, KeyError, ValueError) as error:
+            raise RuntimeError(
+                f"invalid multimodal placeholder for model type '{model_type}' "
+                f"and modality '{modality}': {error}"
+            ) from error
+    return MediaCapabilities(placeholders)
+
+
 __all__ = [
     "GUIDE_SUPPORT_BY_BACKEND",
+    "MediaCapabilities",
+    "resolve_media_capabilities",
     "supported_guides",
 ]

tensorrt_llm/grpc/openengine/control.py

@@ -284,6 +284,7 @@ async def GetModelInfo(
     ) -> model_pb2.ModelInfo:
         # Single-model server: like Generate, any non-empty name resolves to the
         # loaded model.
+        media_modalities = self._inference.supported_media_modalities
         info = model_pb2.ModelInfo(
             model_id=self._model,
             served_model_name=self._model,
@@ -292,9 +293,10 @@ async def GetModelInfo(
             supports_token_ids_input=True,
             # Generate rejects `lora_name` even on an enable_lora build.
             supports_lora=False,
-            supports_multimodal=False,
+            supports_multimodal=bool(media_modalities),
             reasoning_parser=str(self._llm.args.reasoning_parser or ""),
         )
+        info.extra.update({"trtllm_supported_modalities": sorted(media_modalities)})
 
         max_seq_len = _positive_int(self._llm.args.max_seq_len)
         if max_seq_len is not None:

tensorrt_llm/grpc/openengine/disagg.py

@@ -41,7 +41,12 @@
 
 
 def _struct_to_dict(struct: struct_pb2.Struct) -> dict[str, Any]:
-    return MessageToDict(struct)
+    # Prepared media metadata is decoded separately, not as KV rendezvous state.
+    return {
+        key: MessageToDict(value)
+        for key, value in struct.fields.items()
+        if key != "trtllm_decode_preparation"
+    }
 
 
 def _endpoint_to_str(endpoints: Sequence[Any]) -> Optional[str]:
@@ -218,6 +223,7 @@ def prefill_ready_response(
     disaggregated_params: Any,
     transfer_backend: str,
     usage: Optional[generation_pb2.Usage] = None,
+    decode_preparation: Optional[str] = None,
 ) -> generation_pb2.GenerateResponse:
     """Build the PrefillReady event a context worker returns after prefill.
 
@@ -239,6 +245,8 @@ def prefill_ready_response(
         dp_rank=disaggregated_params.ctx_dp_rank or 0,
     )
     attributes: dict[str, Any] = {}
+    if decode_preparation is not None:
+        attributes["trtllm_decode_preparation"] = decode_preparation
     endpoint = disaggregated_params.ctx_info_endpoint
     if endpoint:
         # Carried verbatim as well as split: the generation worker connects a

tensorrt_llm/grpc/openengine/errors.py

@@ -8,6 +8,10 @@ class UnsupportedFeatureError(ValueError):
     """Raised when OpenEngine requests a feature this adapter cannot map."""
 
 
+class MediaUnavailableError(RuntimeError):
+    """Raised when fetching remote request media fails transiently."""
+
+
 class AbortFailedError(RuntimeError):
     """Raised when an active request could not be aborted on the engine."""
 
@@ -18,5 +22,6 @@ def __init__(self, request_id: str, cause: BaseException) -> None:
 
 __all__ = [
     "AbortFailedError",
+    "MediaUnavailableError",
     "UnsupportedFeatureError",
 ]