[TRTLLM-14304][feat] Integrate embeddings cache with encoder side-stream

docs/source/models/supported-models.md

@@ -137,12 +137,15 @@ The following optimizations are available to models that implement
   CUDA stream, allowing it to overlap with work on the main stream. Set
   `multimodal_config.encoder_side_stream_max_ahead` to a positive value to enable it; the value
   limits the number of prefetched requests that can be ahead of admission. This option is mutually
-  exclusive with `multimodal_config.encoder_cuda_graph` and can increase peak GPU memory use.
+  exclusive with `multimodal_config.encoder_cuda_graph` and can increase peak GPU memory use. It
+  can be combined with the multimodal embeddings cache so side-stream cache hits skip encoder work
+  and misses populate the cache.
 - **Multimodal embeddings cache** is a per-model, cross-request LRU cache of encoder embeddings.
   Set `multimodal_config.encoder_cache_max_bytes` to its capacity (for example, `"512MiB"`), or
   `0` to disable it. Entries are cached per multimodal item, but a request reuses cached embeddings
   only when all of its items hit the cache. At present, only single-modality requests are cacheable;
-  mixed-modality requests bypass the cache.
+  mixed-modality requests bypass the cache. When combined with side-stream prefetch, peak memory is
+  the cache capacity plus any in-flight prefetched encoder inputs and outputs.
 
 # Visual Generation Models
 

tensorrt_llm/_torch/models/modeling_multimodal_mixin.py

@@ -33,11 +33,11 @@
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache
 from tensorrt_llm._utils import prefer_pinned
-from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData
+from tensorrt_llm.inputs.multimodal import MultimodalInput, MultimodalParams, MultimodalRuntimeData
 from tensorrt_llm.logger import logger
 
 from .modeling_multimodal_utils import (
-    _cache_multimodal_embeddings,
+    _store_chunked_prefill_embeddings,
     find_input_mm_embeds,
     fuse_input_embeds,
     get_multimodal_embeddings,
@@ -52,6 +52,26 @@
 _MM_ENCODER_CACHE_LOG_NAME = "mm_encoder_cache"
 
 
+def _build_request_multimodal_input(
+    request: "LlmRequest", cache_enabled: bool
+) -> Optional[MultimodalInput]:
+    """Build the encoder-cache key metadata carried by one request."""
+    # Skip construction (and `from_components` validation) when no persistent cache consumes it.
+    if not cache_enabled or request.multimodal_hashes is None:
+        return None
+    # `MultimodalModelMixin._encoder_cache_keys` uses UUID-aware multimodal hashes internally.
+    # Although the UUIDs are not exposed as an attribute, they remain in the backing C++ request
+    # for KV-cache block keys and cache events.
+    return MultimodalInput.from_components(
+        request.multimodal_hashes,
+        request.multimodal_positions,
+        request.multimodal_lengths,
+        mm_item_run_cu_offsets=request.multimodal_item_run_cu_offsets,
+        mm_run_positions=request.multimodal_run_positions,
+        mm_run_lengths=request.multimodal_run_lengths,
+    )
+
+
 def _get_mm_aux_stream(max_prefetch_ahead: int = 0) -> Optional[torch.cuda.Stream]:
     """Return the side CUDA stream used for multimodal encoder prefetch.
 
@@ -190,6 +210,25 @@ def embedding_dtype(self) -> torch.dtype:
         """Return the dtype of each cached multimodal embedding row."""
         raise NotImplementedError
 
+    @property
+    def encoder_cache_active(self) -> bool:
+        """Whether the persistent encoder cache is active for this model.
+
+        Single source of truth shared by:
+
+        * the in-iter consume path
+        * the side-stream prefetch dispatch
+        * the engine's cache-related gate
+        * the KV-cache memory reservation.
+
+        The cache is only active when the model opts in via `supports_encoder_cache` and configures
+        a positive capacity.
+        """
+        if not self.supports_encoder_cache:
+            return False
+        multimodal_config = self.model_config.multimodal_config
+        return multimodal_config is not None and multimodal_config.encoder_cache_max_bytes > 0
+
     def select_multimodal_params(
         self,
         multimodal_params: Sequence[MultimodalParams],
@@ -310,14 +349,22 @@ def _get_or_encode_multimodal_embeddings(
 
         Delegates cache lookup and gather behavior to `get_multimodal_embeddings`, then validates
         the single tensor contract for both encoded and cached-only paths.
+
+        During side-stream prefetch, this runs with the auxiliary stream current, so the H2D copies,
+        the encoder, and every persistent-cache `put()` are issued on that stream. `TensorLRUCache`
+        records each entry's producer event on the issuing (aux) stream; the next iteration's
+        main-stream consumer waits on the request-level `encoder_event` for ordering.
         """
         encoder_cache = self._get_multimodal_encoder_cache()
         cache_misses = []
         if encoder_cache is not None:
             for param in multimodal_params:
                 if param.multimodal_data.get("multimodal_embedding") is not None:
-                    # The forward that attached this request-local embedding already populated the
-                    # persistent cache.
+                    # A present embedding means either an earlier forward already wrote the
+                    # persistent cache, or a prefetch hit attached a cache-owned tensor. Either
+                    # way, skip re-lookup and re-write. `get_multimodal_embeddings` waits on the
+                    # request event and records the attached tensor on the consuming stream before
+                    # gathering it (see `_attach_encoder_cache_hit`).
                     continue
                 if not self._attach_encoder_cache_hit(param, encoder_cache):
                     cache_misses.append(param)
@@ -341,19 +388,16 @@ def _get_multimodal_encoder_cache(self) -> Optional[TensorLRUCache]:
         The cache stores per-item embeddings for params that can be represented by one modality.
         See `_encoder_cache_keys` for the mixed-modality skip path and its technical limitation.
         """
-        multimodal_config = self.model_config.multimodal_config
-        if multimodal_config is None:
-            return None
-
-        max_bytes = multimodal_config.encoder_cache_max_bytes
-        if max_bytes <= 0:
+        if not self.encoder_cache_active:
             logger.debug_once(
-                f"{_MM_ENCODER_CACHE_LOG_NAME}: disabled because "
-                "multimodal_config.encoder_cache_max_bytes=0.",
+                f"{_MM_ENCODER_CACHE_LOG_NAME}: disabled because the model does not opt in via "
+                "supports_encoder_cache or multimodal_config.encoder_cache_max_bytes=0.",
                 key="mm_encoder_cache_disabled",
             )
             return None
 
+        multimodal_config = self.model_config.multimodal_config
+        max_bytes = multimodal_config.encoder_cache_max_bytes
         if self._multimodal_encoder_cache is None:
             # Per-item embeddings are views produced by splitting a request-level encoder output.
             # Clone them so a cached item neither aliases mutable caller output nor retains the
@@ -363,6 +407,7 @@ def _get_multimodal_encoder_cache(self) -> Optional[TensorLRUCache]:
             self._multimodal_encoder_cache = TensorLRUCache(
                 max_bytes,
                 name=_MM_ENCODER_CACHE_LOG_NAME,
+                cuda_stream_aware=multimodal_config.encoder_side_stream_max_ahead > 0,
             )
             try:
                 embedding_dim = self.embedding_dim
@@ -508,6 +553,10 @@ def _attach_encoder_cache_hit(
                 return False
             cached_embeddings.append(cached_embedding)
 
+        # On a side-stream prefetch, this attaches a cache-owned (aliased) tensor while the aux
+        # stream is current. `TensorLRUCache.get` protects the aux gather below. The next
+        # iteration's `get_multimodal_embeddings` call separately waits on the request event and
+        # records the attached tensor on the main consuming stream before gathering it.
         if len(cached_embeddings) == 1:
             param.multimodal_data["multimodal_embedding"] = cached_embeddings[0]
         else:
@@ -711,26 +760,29 @@ def _dispatch_cross_iter_prefetch(
     The event covers all work queued in the aux-stream block, so the same
     event object is shared across all candidates.
     """
+    encoder_cache_enabled = model.encoder_cache_active
     params_list = [
         MultimodalParams(
+            multimodal_input=_build_request_multimodal_input(req, encoder_cache_enabled),
             multimodal_data=mm_data,
             multimodal_runtime=MultimodalRuntimeData(
                 past_seen_token_num=0,
                 chunk_end_pos=cumsum.numel(),
                 embed_mask_cumsum=cumsum,
             ),
         )
-        for _, mm_data, cumsum in candidates
+        for req, mm_data, cumsum in candidates
     ]
 
     # Prefetch targets requests outside the current iteration, so their
     # multimodal tensors are not touched by the main stream. The caller queues
     # this after the iteration's LLM kernels so aux-stream H2D copies and
     # encoder work can overlap them.
     #
-    # Ordering is handled by `encoder_event`, and tensor lifetime is anchored
-    # by `req.py_multimodal_data` until the request is consumed or terminated.
-    # Keeping `record_stream` out also keeps this path modality-neutral.
+    # Request-local ordering is handled by `encoder_event`; the consume path also `record_stream`s
+    # attached embeddings before gathering them so post-prefill request cleanup cannot release
+    # storage while main-stream work is pending. Persistent-cache clones use their own producer
+    # events and consumer `record_stream` calls inside `TensorLRUCache`.
     encoder_event = None
     try:
         with _run_on_aux_stream(aux_stream) as encoder_event:
@@ -748,8 +800,11 @@ def _dispatch_cross_iter_prefetch(
             # model_engine._prepare_inputs.
             for (req, _, _), p in zip(candidates, params_list):
                 req.py_multimodal_data = p.multimodal_data
-            encoder_output = model.encode_multimodal_inputs(params_list)
-            _cache_multimodal_embeddings(params_list, [encoder_output])
+            if encoder_cache_enabled:
+                model._get_or_encode_multimodal_embeddings(params_list)
+            else:
+                encoder_output = model.encode_multimodal_inputs(params_list)
+                _store_chunked_prefill_embeddings(params_list, [encoder_output])
     finally:
         # Stash the event on every candidate's durable LlmRequest (not the
         # per-iter `MultimodalParams`), since `_prepare_inputs` rebuilds the

tensorrt_llm/_torch/models/modeling_multimodal_utils.py

@@ -178,13 +178,14 @@ def _get_uncached_multimodal_params(
     return params_to_run
 
 
-def _cache_multimodal_embeddings(
+def _store_chunked_prefill_embeddings(
     multimodal_params: List[MultimodalParams],
     embeddings: List[torch.Tensor],
 ) -> None:
     """
-    Cache computed multimodal embeddings back to multimodal_data to avoid recomputation.
-    Note this function only caches multimodal embeddings within the current request context,
+    Store computed multimodal embeddings back to multimodal_data to avoid recomputation.
+
+    NOTE: this function only stores multimodal embeddings within the current request context,
     mostly for chunked prefill. It does not persist embeddings across different requests or sessions.
     """
     # TODO: support multiple multimodal modalities per request
@@ -267,12 +268,23 @@ def get_multimodal_embeddings(
     if not multimodal_params:
         return []
 
-    # Wait before touching tensors produced on the MM side stream. Do not
-    # clear the event here; repeated stream-side waits are cheap, and leaving
-    # the event field untouched avoids races if a caller accidentally reuses it.
+    # Wait before touching tensors produced on the MM side stream. Do not clear the event here;
+    # repeated stream-side waits are cheap, and leaving the event field untouched avoids races if a
+    # caller accidentally reuses it.
+    # Register attached embeddings with the consumer stream as well: the request drops its Python
+    # references after prefill, potentially before the asynchronous gather below has finished
+    # reading them.
     for param in multimodal_params:
         if param.encoder_event is not None:
-            torch.cuda.current_stream().wait_event(param.encoder_event)
+            consumer_stream = torch.cuda.current_stream()
+            consumer_stream.wait_event(param.encoder_event)
+            embeds = param.multimodal_data.get("multimodal_embedding")
+            if isinstance(embeds, torch.Tensor):
+                embeds = [embeds]
+            if isinstance(embeds, list):
+                for embed in embeds:
+                    if isinstance(embed, torch.Tensor) and embed.is_cuda:
+                        embed.record_stream(consumer_stream)
 
     # Step 1: Find uncached multimodal params that need encoder processing
     uncached_multimodal_params = _get_uncached_multimodal_params(
@@ -307,8 +319,8 @@ def get_multimodal_embeddings(
             return encoder_embeddings
 
         # Step 3: Cache the computed embeddings to multimodal_data["multimodal_embedding"]
-        _cache_multimodal_embeddings(uncached_multimodal_params,
-                                     encoder_embeddings)
+        _store_chunked_prefill_embeddings(uncached_multimodal_params,
+                                          encoder_embeddings)
 
     # Step 4: Gather all embeddings for the batch
     for param in multimodal_params:
@@ -333,9 +345,11 @@ def get_attached_multimodal_embeddings(
         multimodal_params: List[MultimodalParams]) -> List[torch.Tensor]:
     """Gather embeddings already stored on MultimodalParams.
 
-    Use this on E/P prefill workers and cached-only paths. The encoder already
-    ran somewhere else. This only makes the tensor list that
-    find_input_mm_embeds slices.
+    Use this on E/P prefill workers and cached-only paths. The encoder already ran somewhere else.
+    This only makes the tensor list that `find_input_mm_embeds` slices.
+
+    Side-stream-prefetched requests must use `get_multimodal_embeddings`, which waits on
+    `encoder_event` and registers attached tensors with the consuming stream before gathering them.
     """
     attached_embeddings = []
     for param in multimodal_params:

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -640,13 +640,10 @@ def _reserve_multimodal_encoder_cache_memory(
         """Reserve encoder-cache capacity when the model implements that cache."""
         model = self._model_engine.model
         if (not isinstance(model, MultimodalModelMixin)
-                or not model.supports_encoder_cache):
+                or not model.encoder_cache_active):
             return peak_memory
 
-        multimodal_config = model.model_config.multimodal_config
-        if multimodal_config is None:
-            return peak_memory
-        return peak_memory + multimodal_config.encoder_cache_max_bytes
+        return peak_memory + model.model_config.multimodal_config.encoder_cache_max_bytes
 
     def _get_token_num_for_estimation(self) -> int:
         """Compute KV cache capacity required for estimate_max_kv_cache_tokens to succeed."""

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -22,7 +22,7 @@
                                  prefer_pinned, release_gc, torch_dtype_to_str,
                                  trace_func)
 from tensorrt_llm.bindings.internal.runtime import TaskLayerModuleConfig
-from tensorrt_llm.inputs.multimodal import (MultimodalInput, MultimodalParams,
+from tensorrt_llm.inputs.multimodal import (MultimodalParams,
                                             MultimodalRuntimeData,
                                             _has_mm_payload_keys,
                                             check_mm_embed_cumsum_if_needed,
@@ -54,7 +54,8 @@
 from ..metadata import KVCacheParams
 from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader
 from ..models.modeling_multimodal_encoder import MultimodalEncoderMixin
-from ..models.modeling_multimodal_mixin import MultimodalModelMixin
+from ..models.modeling_multimodal_mixin import (MultimodalModelMixin,
+                                                _build_request_multimodal_input)
 from ..models.modeling_multimodal_utils import filter_mm_token_from_input_ids
 from ..models.modeling_utils import DecoderModelForCausalLM
 from ..modules.fused_moe.moe_load_balancer import (MoeLoadBalancer,
@@ -163,25 +164,6 @@ def _filter_piecewise_capture_num_tokens(
     return kept, unrecordable
 
 
-def _build_request_multimodal_input(
-        request: LlmRequest, cache_enabled: bool) -> Optional[MultimodalInput]:
-    # Skip building this input (and its `from_components` validation) when the cache is disabled.
-    if not cache_enabled or request.multimodal_hashes is None:
-        return None
-    # `multimodal_input` is consumed only by the encoder-cache key path
-    # (`MultimodalModelMixin._encoder_cache_keys`), which uses UUID-aware multimodal hashes
-    # internally. Although the `multimodal_uuids` are not exposed as an attribute, they remain in
-    # the backing C++ request for KV-cache block keys and cache events.
-    return MultimodalInput.from_components(
-        request.multimodal_hashes,
-        request.multimodal_positions,
-        request.multimodal_lengths,
-        mm_item_run_cu_offsets=request.multimodal_item_run_cu_offsets,
-        mm_run_positions=request.multimodal_run_positions,
-        mm_run_lengths=request.multimodal_run_lengths,
-    )
-
-
 def _filter_cuda_graph_batch_sizes(cuda_graph_batch_sizes: list[int],
                                    max_batch_size: int, max_num_tokens: int,
                                    max_total_draft_tokens: int,
@@ -894,9 +876,9 @@ def use_mrope(self):
     @functools.cached_property
     def _mm_encoder_cache_enabled(self) -> bool:
         """Whether the multimodal encoder cache is active for this model."""
-        multimodal_config = self.model.model_config.multimodal_config
-        return (multimodal_config is not None
-                and multimodal_config.encoder_cache_max_bytes > 0)
+        model = self.model
+        return (isinstance(model, MultimodalModelMixin)
+                and model.encoder_cache_active)
 
     @property
     def is_warmup(self):

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -6254,10 +6254,10 @@ def _maybe_prefetch_next_iter_mm_encoders(
             )
         except Exception:
             # Speculative prefetch is best-effort and must never crash the
-            # executor loop. On failure, `py_mm_encoder_event` is not stamped,
-            # so the next iteration's `_prepare_inputs` falls back to the
-            # standard in-iter encode path (which re-runs `to_device` and the
-            # encoder unconditionally when no cached embedding is present).
+            # executor loop. The dispatch helper stamps an event when it queued
+            # any auxiliary-stream work, even on a partial failure, so the next
+            # iteration can safely inspect the request-local data and fall back
+            # to the standard in-iter encode path for any missing embedding.
             logger.warning(
                 f"Cross-iter MM encoder prefetch failed; falling back to "
                 f"in-iter encode.\n{traceback.format_exc()}")