[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,89 @@ 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,
+            max_length: Optional[int] = None) -> 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")
+        arrays = {}
+
+        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)
+            arrays[name] = array
+        text_config = getattr(self.config, "text_config", self.config)
+        max_positions = getattr(text_config, "max_position_embeddings", None)
+        if type(max_positions) is not int or not 0 < max_positions <= 2**31:
+            raise ValueError("invalid Qwen mRoPE model position bound")
+        if max_length is None:
+            max_length = max_positions
+        if prompt_length <= 0 or max_length < prompt_length:
+            raise ValueError("invalid Qwen mRoPE prompt length")
+        positions = arrays["mrope_position_ids"]
+        if np.any(positions < 0) or np.any(positions >= max_positions):
+            raise ValueError("Qwen mRoPE positions exceed model bounds")
+        delta = int(arrays["mrope_position_deltas"][0, 0])
+        if delta != int(positions.max()) + 1 - prompt_length:
+            raise ValueError("inconsistent Qwen mRoPE position delta")
+        # Include every possible decode position before handing state to the GPU.
+        last_position = max_length - 1 + delta
+        if not 0 <= last_position < max_positions:
+            raise ValueError("Qwen mRoPE decode positions exceed model bounds")
+        arrays["mrope_position_deltas"] = arrays[
+            "mrope_position_deltas"].astype(np.int32)
+        return {
+            "mrope_config": {
+                name: torch.from_numpy(array)
+                for name, array in arrays.items()
+            }
+        }
+
     def __init__(self,
                  model_path: str,
                  config: PretrainedConfig,

tensorrt_llm/grpc/openengine/README.md

@@ -73,7 +73,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/disagg.py

@@ -43,7 +43,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]:
@@ -224,6 +229,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.
 
@@ -245,6 +251,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/preparation.py

@@ -0,0 +1,172 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Versioned, portable CPU metadata for generation-only preparation reuse."""
+
+import base64
+import hashlib
+import json
+from dataclasses import asdict
+from itertools import chain
+from typing import Any
+
+import numpy as np
+from google.protobuf import struct_pb2
+
+from tensorrt_llm.inputs.model_config import resolve_top_level_model_type
+from tensorrt_llm.inputs.multimodal import MultimodalInput, MultimodalParams
+from tensorrt_llm.llmapi.llm import PreprocessedInputs
+
+_DECODE_PREPARATION_KEY = "trtllm_decode_preparation"
+_MAX_DECODE_PREPARATION_BYTES = 2 * 1024 * 1024
+
+
+def _input_digest(request: Any) -> str:
+    digest = hashlib.sha256()
+    kind = request.WhichOneof("input")
+    data = request.prompt.encode() if kind == "prompt" else request.token_ids.SerializeToString()
+    for part in chain((kind.encode(), data), (item.SerializeToString() for item in request.media)):
+        digest.update(len(part).to_bytes(8, "big"))
+        digest.update(part)
+    return digest.hexdigest()
+
+
+def _config_digest(llm: Any) -> str:
+    processor = llm.input_processor
+    config = processor.config.to_dict()
+    config.pop("_name_or_path", None)
+    image_processor = getattr(processor.processor, "image_processor", None)
+    identity = {
+        "model": config,
+        "processor_class": f"{type(processor).__module__}.{type(processor).__qualname__}",
+        "media_io": processor.get_preferred_media_io_kwargs(),
+        "image_processor": image_processor.to_dict() if image_processor is not None else None,
+    }
+    return hashlib.sha256(json.dumps(identity, sort_keys=True, default=str).encode()).hexdigest()
+
+
+def _decode_preparation(
+    request: Any, llm: Any, capture: dict, preprocessed: PreprocessedInputs | None = None
+) -> str | None:
+    if not capture:
+        return None
+    payload = {
+        "version": 1,
+        "model_type": resolve_top_level_model_type(llm.input_processor.config),
+        "input_digest": _input_digest(request),
+        "config_digest": _config_digest(llm),
+        "state": capture["state"],
+        "prompt_token_ids_u32le": base64.b64encode(
+            np.asarray(capture["prompt_token_ids"], dtype="<u4").tobytes()
+        ).decode("ascii"),
+    }
+    params = getattr(preprocessed, "multimodal_params", None)
+    if params is not None and params.multimodal_input is not None:
+        payload["multimodal_input"] = asdict(params.multimodal_input)
+    wire_value = struct_pb2.Value(
+        string_value=json.dumps(payload, separators=(",", ":"), allow_nan=False)
+    )
+    if wire_value.ByteSize() > _MAX_DECODE_PREPARATION_BYTES:
+        return None
+    return wire_value.string_value
+
+
+def _prepared_decode_from_request(request: Any, llm: Any) -> PreprocessedInputs | None:
+    if not request.HasField("kv") or not request.kv.HasField("session"):
+        return None
+    attrs = request.kv.session.attributes_struct
+    if _DECODE_PREPARATION_KEY not in attrs.fields:
+        return None
+    value = attrs.fields[_DECODE_PREPARATION_KEY]
+    if value.ByteSize() > _MAX_DECODE_PREPARATION_BYTES:
+        raise ValueError("prepared decode metadata exceeds the size limit")
+    if value.WhichOneof("kind") != "string_value":
+        raise ValueError("prepared decode metadata must be a JSON string")
+    try:
+        payload = json.loads(value.string_value)
+    except (ValueError, RecursionError) as error:
+        raise ValueError("invalid prepared decode JSON") from error
+    if not isinstance(payload, dict):
+        raise ValueError("prepared decode metadata must contain a JSON object")
+    version = payload.get("version")
+    if type(version) not in (int, float):
+        raise ValueError("prepared decode metadata version must be a number")
+    if version != 1:
+        return None  # An unknown version retains the ordinary media preparation path.
+    if payload.get("model_type") != resolve_top_level_model_type(llm.input_processor.config):
+        raise ValueError("prepared decode metadata belongs to another model type")
+    if payload.get("input_digest") != _input_digest(request):
+        raise ValueError("prepared decode metadata does not match the request inputs")
+    if payload.get("config_digest") != _config_digest(llm):
+        raise ValueError("prepared decode metadata belongs to another model configuration")
+    encoded_tokens = payload.get("prompt_token_ids_u32le")
+    limit = getattr(llm.args, "max_seq_len", None) or 131072
+    if (
+        not isinstance(encoded_tokens, str)
+        or not encoded_tokens
+        or len(encoded_tokens) > ((limit * 4 + 2) // 3) * 4
+    ):
+        raise ValueError("invalid prepared prompt length")
+    try:
+        raw_tokens = base64.b64decode(encoded_tokens, validate=True)
+    except (ValueError, UnicodeError) as error:
+        raise ValueError("invalid prepared prompt encoding") from error
+    if not raw_tokens or len(raw_tokens) % 4 or len(raw_tokens) > limit * 4:
+        raise ValueError("invalid prepared prompt byte length")
+    tokens = np.frombuffer(raw_tokens, dtype="<u4").tolist()
+    state = payload.get("state")
+    if not isinstance(state, dict):
+        raise ValueError("prepared decode model state must be an object")
+    max_length = getattr(llm.args, "max_seq_len", None)
+    if request.HasField("stopping") and request.stopping.HasField("max_tokens"):
+        requested_length = len(tokens) + request.stopping.max_tokens
+        max_length = min(max_length, requested_length) if max_length else requested_length
+    restored = llm.input_processor.restore_decode_state(state, len(tokens), max_length=max_length)
+    mm_input = payload.get("multimodal_input")
+    if mm_input is not None:
+
+        def normalize(value: Any) -> Any:
+            if isinstance(value, dict):
+                return {key: normalize(item) for key, item in value.items()}
+            if isinstance(value, list):
+                return [normalize(item) for item in value]
+            if value is None or isinstance(value, str):
+                return value
+            if (
+                type(value) not in (int, float)
+                or not -(2**31) <= value < 2**31
+                or not float(value).is_integer()
+            ):
+                raise ValueError("invalid prepared multimodal identity")
+            return int(value)
+
+        if not isinstance(mm_input, dict):
+            raise ValueError("prepared multimodal identity must be an object")
+        try:
+            mm_input = MultimodalInput(**normalize(mm_input))
+        except (TypeError, RecursionError) as error:
+            raise ValueError("invalid prepared multimodal identity fields") from error
+        if any(
+            len(digest) != 8 or any(type(value) is not int for value in digest)
+            for digest in mm_input.multimodal_hashes
+        ):
+            raise ValueError("prepared multimodal hashes must contain eight int32 values")
+        if any(
+            position + length > len(tokens)
+            for position, length in zip(mm_input.multimodal_positions, mm_input.multimodal_lengths)
+        ):
+            raise ValueError("prepared multimodal identity exceeds the prompt")
+        if mm_input.multimodal_run_positions is not None and any(
+            position + length > len(tokens)
+            for position, length in zip(
+                mm_input.multimodal_run_positions, mm_input.multimodal_run_lengths
+            )
+        ):
+            raise ValueError("prepared multimodal runs exceed the prompt")
+    params = MultimodalParams(multimodal_input=mm_input, multimodal_data=restored)
+    try:
+        params.to_handle("multimodal_data")
+    except Exception:
+        params.release_handles()
+        raise
+    return PreprocessedInputs(prompt_token_ids=tokens, multimodal_params=params)

tensorrt_llm/grpc/openengine/request_mapping.py

@@ -625,7 +625,9 @@ def _count_token_sequences(
     return counts
 
 
-def _validate_media_placeholders(prepared: _PreparedInput, llm: Any) -> None:
+def _validate_media_placeholders(
+    prepared: _PreparedInput, llm: Any, *, add_special_tokens: bool = False
+) -> list[int] | None:
     """Reject placeholder mismatches before probing the stateful MM processor."""
     if not prepared.media or prepared.capabilities is None:
         return
@@ -634,7 +636,7 @@ def _validate_media_placeholders(prepared: _PreparedInput, llm: Any) -> None:
         return
     tokenizer = llm.tokenizer
     prompt_ids = (
-        tokenizer.encode(prepared.inputs, add_special_tokens=False)
+        tokenizer.encode(prepared.inputs, add_special_tokens=add_special_tokens)
         if isinstance(prepared.inputs, str)
         else prepared.inputs["prompt_token_ids"]
     )
@@ -671,6 +673,7 @@ def _validate_media_placeholders(prepared: _PreparedInput, llm: Any) -> None:
         actual_tokens = Counter(token_id for token_id in prompt_ids if token_id in content_ids)
         if actual_tokens != expected_tokens:
             raise ValueError("multimodal placeholder token counts do not match media items")
+    return prompt_ids
 
 
 async def _load_prepared_input(prepared: _PreparedInput, llm: Any) -> str | dict[str, Any]:

tensorrt_llm/grpc/openengine/server.py

@@ -99,6 +99,8 @@ class OpenEngineServer:
         llm: Initialized TensorRT-LLM LLM instance.
         model: Model name accepted by Generate requests.
         subagent_affinity_auth_key: Shared key for authenticated parent affinity hints.
+        max_media_preparations: Concurrent media preparation requests. Defaults to
+            TRTLLM_OPENENGINE_MEDIA_PREPARATION_LIMIT, or four when unset.
     """
 
     def __init__(
@@ -111,7 +113,14 @@ def __init__(
         frontend: FrontendClient | None = None,
         instance_id: str | None = None,
         subagent_affinity_auth_key: str | None = None,
+        max_media_preparations: int | None = None,
     ) -> None:
+        if max_media_preparations is None:
+            max_media_preparations = int(
+                os.getenv("TRTLLM_OPENENGINE_MEDIA_PREPARATION_LIMIT", "4")
+            )
+        if type(max_media_preparations) is not int or max_media_preparations < 1:
+            raise ValueError("max_media_preparations must be a positive integer")
         self.host = host
         self.port = port
         self._llm = llm
@@ -143,6 +152,7 @@ def __init__(
             subagent_affinity_auth_key=subagent_affinity_auth_key,
             input_executor=self._input_executor,
             media_executor=self._media_executor,
+            max_media_preparations=max_media_preparations,
         )
         openengine_pb2_grpc.add_InferenceServicer_to_server(inference, self._server)
         # Control shares the inference servicer's in-flight request table so