From 76324bd4648c40a9f77eaa49ee3d9d2d4845a494 Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Fri, 17 Jul 2026 11:35:36 -0700 Subject: [PATCH] [TRTLLM-14304][feat] Combine embeddings cache with side-stream prefetch * Why? Side-stream multimodal encoder prefetch and the persistent embeddings cache were mutually exclusive, preventing requests from benefiting from both overlapped encoder work and cross-request embedding reuse. * What? Allow supported multimodal models to use both features together. Prefetch now reads and populates the persistent cache while preserving CUDA stream ordering and tensor lifetimes. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- docs/source/models/supported-models.md | 7 +- .../models/modeling_multimodal_mixin.py | 91 ++++-- .../models/modeling_multimodal_utils.py | 38 ++- tensorrt_llm/_torch/pyexecutor/_util.py | 7 +- .../_torch/pyexecutor/model_engine.py | 30 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 8 +- tensorrt_llm/_torch/tensor_lru_cache.py | 44 ++- tensorrt_llm/inputs/multimodal.py | 3 +- tensorrt_llm/llmapi/llm_args.py | 12 +- .../test_mm_encoder_cross_iter_prefetch.py | 271 ++++++++++++++++++ .../multimodal/test_multimodal_mixin.py | 13 + .../unittest/_torch/test_tensor_lru_cache.py | 79 +++++ tests/unittest/llmapi/test_llm_args.py | 19 +- 13 files changed, 536 insertions(+), 86 deletions(-) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index b03f9961f7cc..bf3ee2ec2235 100644 --- a/docs/source/models/supported-models.md +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index b1a2c5df6845..fc07dcd144f5 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/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,8 +760,10 @@ 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, @@ -720,7 +771,7 @@ def _dispatch_cross_iter_prefetch( 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 @@ -728,9 +779,10 @@ def _dispatch_cross_iter_prefetch( # 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 diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py index 3dd866015dd3..63dedd274294 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py +++ b/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: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0149ba7ae03a..71893d2965a6 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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.""" diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 50aa317e9e43..5ee7c4d0572b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 34ebd340a5ee..0060dbcc5b51 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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()}") diff --git a/tensorrt_llm/_torch/tensor_lru_cache.py b/tensorrt_llm/_torch/tensor_lru_cache.py index 2d94e75a1868..58e789e86309 100644 --- a/tensorrt_llm/_torch/tensor_lru_cache.py +++ b/tensorrt_llm/_torch/tensor_lru_cache.py @@ -32,6 +32,10 @@ class _Entry(NamedTuple): value: torch.Tensor size_bytes: int + # CUDA event recorded on the producing stream right after the clone in `put`. Consumers on a + # different stream wait on it before reading `value`. `None` for CPU tensors or when the cache + # is not stream-aware. + producer_event: torch.cuda.Event | None class TensorLRUCacheStats(NamedTuple): @@ -77,9 +81,16 @@ class TensorLRUCache(Generic[K]): temporarily needs both the source tensor and its copy and may exceed the cache limit until eviction completes. + In CUDA-stream-aware mode, each entry owns the event recorded after its clone. Replacement, + eviction, and clear drop that event with the entry; events are not reused because an evicted + tensor may still have outstanding consumers on another stream. + Args: max_bytes: Maximum logical tensor bytes held by this cache. name: Short label used in debug log messages. + cuda_stream_aware: When enabled, synchronize CUDA tensor producers and consumers across + streams and extend allocation lifetime through every consuming stream. CPU tensors are + unaffected. """ def __init__( @@ -87,12 +98,14 @@ def __init__( max_bytes: int, *, name: str = "tensor_lru_cache", + cuda_stream_aware: bool = False, ) -> None: if max_bytes <= 0: raise ValueError("max_bytes must be positive") self._max_bytes = max_bytes self._name = name + self._cuda_stream_aware = cuda_stream_aware self._current_bytes = 0 self._items: OrderedDict[K, _Entry] = OrderedDict() self._lock = RLock() @@ -124,6 +137,7 @@ def get(self, key: K) -> torch.Tensor | None: self._counters.hits += 1 self._items.move_to_end(key) + self._prepare_for_current_stream(entry) return entry.value def put(self, key: K, value: torch.Tensor) -> bool: @@ -144,6 +158,10 @@ def put(self, key: K, value: torch.Tensor) -> bool: return False stored_value = value.detach().clone() + producer_event = None + if self._cuda_stream_aware and stored_value.is_cuda: + producer_event = torch.cuda.Event() + producer_event.record(torch.cuda.current_stream(stored_value.device)) with self._lock: old_entry = self._items.pop(key, None) @@ -153,7 +171,11 @@ def put(self, key: K, value: torch.Tensor) -> bool: else: self._counters.insertions += 1 - self._items[key] = _Entry(value=stored_value, size_bytes=size_bytes) + self._items[key] = _Entry( + value=stored_value, + size_bytes=size_bytes, + producer_event=producer_event, + ) self._current_bytes += size_bytes evicted_count, evicted_bytes = self._evict_until_within_limit() @@ -174,6 +196,7 @@ def pop(self, key: K) -> torch.Tensor | None: return None self._current_bytes -= entry.size_bytes + self._prepare_for_current_stream(entry) return entry.value def clear(self) -> None: @@ -210,6 +233,25 @@ def log_stats(self, reason: str) -> None: def _tensor_size_bytes(tensor: torch.Tensor) -> int: return tensor.numel() * tensor.element_size() + def _prepare_for_current_stream(self, entry: _Entry) -> None: + """Order and anchor a cached tensor for consumption on the current stream. + + Called from `get` / `pop` before returning an entry it: + * makes the current (consuming) stream wait on the entry's producer event, so a cross-stream + read observes fully-written data + * calls `record_stream` on the consuming stream so the caching allocator will not reuse + the storage while consumer-stream work is still pending, even if a later replacement or + eviction drops the cache's own reference. + """ + if not self._cuda_stream_aware or not entry.value.is_cuda: + return + + consumer_stream = torch.cuda.current_stream(entry.value.device) + # The producer event orders the data dependency; `record_stream` separately guards lifetime. + if entry.producer_event is not None: + consumer_stream.wait_event(entry.producer_event) + entry.value.record_stream(consumer_stream) + def _evict_until_within_limit(self) -> tuple[int, int]: evicted_count = 0 evicted_bytes = 0 diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index fb4b33779d32..811553b99672 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -511,7 +511,8 @@ class MultimodalParams: mm_item_order: Optional[List[Dict[str, Union[str, int]]]] = None # CUDA event recorded on a side stream by the MM encoder prefetch path. # When set, the consume site in `get_multimodal_embeddings` issues a - # `wait_event` on the current stream before reading cached embeddings. + # `wait_event` and records attached embeddings on the current stream before + # reading them. # Always `None` unless side-stream prefetch is enabled and a prefetch ran. encoder_event: Optional[torch.cuda.Event] = field(default=None, repr=False, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d09e0b576b5b..5d7a79778bad 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -563,7 +563,7 @@ class MultimodalConfig(StrictBaseModel): ("Maximum number of pending multimodal requests whose encoder work can be prefetched " "on a side CUDA stream ahead of admission. 0 disables side-stream prefetch. " "Incompatible with encoder_cuda_graph because graph replay uses static buffers. " - "For the time being, this is also incompatible with encoder_cache_max_bytes > 0." + "Can be combined with encoder_cache_max_bytes; the two memory limits are additive." ), status="prototype", ) @@ -576,7 +576,7 @@ class MultimodalConfig(StrictBaseModel): "Cache entries are per multimodal item, but reuse is all-or-nothing for each request: " "every item in the request must hit the cache before cached embeddings are reused. " "Only single-modality requests are cacheable for the time being. " - "For the time being, this is incompatible with encoder_side_stream_max_ahead > 0. " + "Can be combined with encoder_side_stream_max_ahead. " "NOTE: This is only valid for child implementations of the `MultimodalModelMixin`." ), status="prototype", @@ -607,14 +607,6 @@ def validate_encoder_optimization_compatibility(self) -> 'MultimodalConfig': "multimodal_config.encoder_side_stream_max_ahead > 0 are " "mutually exclusive. Disable side-stream MM prefetch or " "disable MM encoder CUDA graphs.") - # TODO(TRTLLM-14034): Make encoder side-stream read and write from the cache. - if (self.encoder_cache_max_bytes > 0 - and self.encoder_side_stream_max_ahead > 0): - raise ValueError( - "multimodal_config.encoder_cache_max_bytes > 0 and " - "multimodal_config.encoder_side_stream_max_ahead > 0 are " - "mutually exclusive. Disable side-stream MM prefetch or set " - "the MM encoder cache capacity to 0.") return self diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py index 01e538d16092..935d7a166e75 100644 --- a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py +++ b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py @@ -23,13 +23,19 @@ import pytest import torch +from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_multimodal_mixin as mm_mixin from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( MultimodalModelMixin, _get_mm_aux_stream, maybe_prefetch_mm_encoder_for_next_iter, ) +from tensorrt_llm._torch.models.modeling_multimodal_utils import get_multimodal_embeddings from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig +from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData +from tensorrt_llm.llmapi.llm_args import MultimodalConfig + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") class _StubModel(MultimodalModelMixin): @@ -54,6 +60,38 @@ def encode_multimodal_inputs(self, multimodal_params, **kwargs) -> torch.Tensor: return embeddings +class _CacheStubModel(_StubModel): + supports_encoder_cache = True + + def __init__(self, hidden_size: int, tokens_per_image: int): + super().__init__(hidden_size, tokens_per_image) + self.model_config = ModelConfig( + multimodal_config=MultimodalConfig( + encoder_cache_max_bytes=4096, + encoder_side_stream_max_ahead=2, + ) + ) + self.last_encoder_batch_size = 0 + + @property + def embedding_dim(self) -> int: + return self._hidden_size + + @property + def embedding_dtype(self) -> torch.dtype: + return torch.float32 + + def encode_multimodal_inputs(self, multimodal_params, **kwargs) -> torch.Tensor: + self.encoder_call_count += 1 + self.last_encoder_batch_size = len(multimodal_params) + self.encoder_call_stream_id = torch.cuda.current_stream().cuda_stream + return torch.full( + (len(multimodal_params) * self._tokens_per_image, self._hidden_size), + float(self.encoder_call_count), + device="cuda", + ) + + def _make_request(request_id: int, num_tokens: int) -> LlmRequest: pixel_values = torch.randn(3, 32, 32) # CPU tensor, simulating unscheduled cumsum = torch.arange(1, num_tokens + 1, dtype=torch.int64) @@ -72,6 +110,33 @@ def _make_request(request_id: int, num_tokens: int) -> LlmRequest: ) +def _make_cacheable_request( + request_id: int, + num_tokens: int, + *, + item_hash: list[int] | None = None, +) -> LlmRequest: + if item_hash is None: + item_hash = [1, 2, 3, 4, 5, 6, 7, 8] + return LlmRequest( + request_id=request_id, + max_new_tokens=1, + input_tokens=[0] * num_tokens, + sampling_config=SamplingConfig(beam_width=1), + is_streaming=False, + py_multimodal_data={ + "image": {"pixel_values": torch.randn(3, 32, 32)}, + "multimodal_embed_mask_cumsum": torch.arange(1, num_tokens + 1, dtype=torch.int64), + "multimodal_embedding_lengths": [num_tokens], + "mm_processor_kwargs_hash": "kwargs-a", + }, + multimodal_hashes=[item_hash], + multimodal_positions=[0], + multimodal_lengths=[num_tokens], + multimodal_uuids=[None], + ) + + def _make_metadata_only_request(request_id: int, num_tokens: int) -> LlmRequest: cumsum = torch.arange(1, num_tokens + 1, dtype=torch.int64) return LlmRequest( @@ -222,6 +287,7 @@ def test_cross_iter_prefetch_no_op_when_config_off(): assert req.py_mm_encoder_event is None +@requires_cuda def test_cross_iter_prefetch_materializes_on_side_stream(monkeypatch): max_prefetch_ahead = 1 aux_stream = _get_mm_aux_stream(max_prefetch_ahead) @@ -242,6 +308,102 @@ def test_cross_iter_prefetch_materializes_on_side_stream(monkeypatch): req.py_mm_encoder_event.synchronize() +@requires_cuda +def test_cross_iter_prefetch_populates_and_reuses_persistent_cache(): + max_prefetch_ahead = 2 + aux_stream = _get_mm_aux_stream(max_prefetch_ahead) + assert aux_stream is not None + model = _CacheStubModel(hidden_size=8, tokens_per_image=4) + + first = _make_cacheable_request(request_id=0, num_tokens=4) + assert ( + maybe_prefetch_mm_encoder_for_next_iter( + model, [first], max_prefetch_ahead=max_prefetch_ahead + ) + == 1 + ) + first.py_mm_encoder_event.synchronize() + first_embedding = first.py_multimodal_data["multimodal_embedding"] + + cache = model._multimodal_encoder_cache + assert cache is not None + assert len(cache) == 1 + assert model.encoder_call_count == 1 + assert model.encoder_call_stream_id == aux_stream.cuda_stream + + second = _make_cacheable_request(request_id=1, num_tokens=4) + assert ( + maybe_prefetch_mm_encoder_for_next_iter( + model, [second], max_prefetch_ahead=max_prefetch_ahead + ) + == 1 + ) + second.py_mm_encoder_event.synchronize() + + assert model.encoder_call_count == 1 + torch.testing.assert_close(second.py_multimodal_data["multimodal_embedding"], first_embedding) + + +@requires_cuda +def test_cross_iter_prefetch_mixed_cache_hit_and_miss_encodes_only_miss(): + max_prefetch_ahead = 2 + model = _CacheStubModel(hidden_size=8, tokens_per_image=4) + first = _make_cacheable_request(request_id=0, num_tokens=4) + maybe_prefetch_mm_encoder_for_next_iter(model, [first], max_prefetch_ahead=max_prefetch_ahead) + first.py_mm_encoder_event.synchronize() + + hit = _make_cacheable_request(request_id=1, num_tokens=4) + miss = _make_cacheable_request(request_id=2, num_tokens=4, item_hash=[9] * 8) + assert ( + maybe_prefetch_mm_encoder_for_next_iter( + model, + [hit, miss], + max_prefetch=2, + max_prefetch_ahead=max_prefetch_ahead, + ) + == 2 + ) + hit.py_mm_encoder_event.synchronize() + + assert model.encoder_call_count == 2 + assert model.last_encoder_batch_size == 1 + torch.testing.assert_close( + hit.py_multimodal_data["multimodal_embedding"], + torch.ones((4, 8), device="cuda"), + ) + torch.testing.assert_close( + miss.py_multimodal_data["multimodal_embedding"], + torch.full((4, 8), 2.0, device="cuda"), + ) + + +@requires_cuda +def test_cross_iter_prefetch_cache_model_preserves_uncacheable_fallbacks(): + max_prefetch_ahead = 2 + model = _CacheStubModel(hidden_size=8, tokens_per_image=4) + unkeyable = _make_request(request_id=0, num_tokens=4) + mixed_modality = _make_cacheable_request(request_id=1, num_tokens=4) + mixed_modality.py_multimodal_data["audio"] = {"input_features": torch.empty(1)} + + assert ( + maybe_prefetch_mm_encoder_for_next_iter( + model, + [unkeyable, mixed_modality], + max_prefetch=2, + max_prefetch_ahead=max_prefetch_ahead, + ) + == 2 + ) + unkeyable.py_mm_encoder_event.synchronize() + + assert model.encoder_call_count == 1 + assert model.last_encoder_batch_size == 2 + cache = model._multimodal_encoder_cache + assert cache is not None + assert len(cache) == 0 + + +@requires_cuda def test_cross_iter_prefetch_skips_in_flight_and_cached(monkeypatch): model = _StubModel(hidden_size=8, tokens_per_image=4) in_flight_req = _make_request(request_id=0, num_tokens=4) @@ -259,3 +421,112 @@ def test_cross_iter_prefetch_skips_in_flight_and_cached(monkeypatch): assert n == 1 assert "multimodal_embedding" not in in_flight_req.py_multimodal_data assert fresh_req.py_multimodal_data.get("multimodal_embedding") is not None + + +@requires_cuda +def test_cross_iter_prefetch_does_not_synchronize_main_stream(monkeypatch): + """Routing through the cache must not block the calling (main) stream. + + The prefetch path may only enqueue aux-stream work and cross-stream waits; a + blocking host/stream synchronization would defeat the overlap the side stream + exists for. Spy on the synchronization entry points and assert none fire while + the routed dispatch runs. + """ + model = _CacheStubModel(hidden_size=8, tokens_per_image=4) + req = _make_cacheable_request(request_id=0, num_tokens=4) + original_encoder = model.encode_multimodal_inputs + + def delayed_encoder(multimodal_params, **kwargs): + torch.cuda._sleep(100_000_000) + return original_encoder(multimodal_params, **kwargs) + + monkeypatch.setattr(model, "encode_multimodal_inputs", delayed_encoder) + + sync_calls = [] + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: sync_calls.append("device")) + monkeypatch.setattr( + torch.cuda.Stream, "synchronize", lambda self, *a, **k: sync_calls.append("stream") + ) + + n = maybe_prefetch_mm_encoder_for_next_iter(model, [req], max_prefetch_ahead=2) + + assert n == 1 + assert sync_calls == [] + assert not req.py_mm_encoder_event.query() + # Drain the queued aux-stream work via the (unpatched) event before teardown. + req.py_mm_encoder_event.synchronize() + + +@requires_cuda +def test_cross_iter_prefetch_does_not_rewrite_request_local_embedding(): + """A present request-local embedding skips both the cache lookup and the write. + + Reproduces the next-iteration in-iter consume of a request whose embedding was + already produced by a prefetch miss: the encoder must not run again and the + persistent cache must not be rewritten. + """ + model = _CacheStubModel(hidden_size=8, tokens_per_image=4) + req = _make_cacheable_request(request_id=0, num_tokens=4) + assert maybe_prefetch_mm_encoder_for_next_iter(model, [req], max_prefetch_ahead=2) == 1 + req.py_mm_encoder_event.synchronize() + + cache = model._multimodal_encoder_cache + stats_before = cache.stats() + assert stats_before.insertions == 1 + + cumsum = req.py_multimodal_data["multimodal_embed_mask_cumsum"] + param = MultimodalParams( + multimodal_input=mm_mixin._build_request_multimodal_input(req, cache_enabled=True), + multimodal_data=req.py_multimodal_data, + multimodal_runtime=MultimodalRuntimeData( + past_seen_token_num=0, + chunk_end_pos=cumsum.numel(), + embed_mask_cumsum=cumsum, + ), + ) + model._get_or_encode_multimodal_embeddings([param]) + + assert model.encoder_call_count == 1 # no re-encode + stats_after = cache.stats() + assert stats_after.insertions == stats_before.insertions + assert stats_after.replacements == stats_before.replacements + + +@requires_cuda +@pytest.mark.parametrize("embedding_count", [1, 2]) +def test_prefetched_embedding_records_main_consumer_stream(monkeypatch, embedding_count): + """The request may release an aux-produced embedding before its main-stream gather completes.""" + aux_stream = torch.cuda.Stream() + consumer_stream = torch.cuda.Stream() + producer_event = torch.cuda.Event() + with torch.cuda.stream(aux_stream): + embeddings = [torch.ones(4, device="cuda") for _ in range(embedding_count)] + producer_event.record(aux_stream) + + record_stream_calls = [] + original_record_stream = torch.Tensor.record_stream + + def record_stream(tensor, stream): + record_stream_calls.append((tensor.data_ptr(), stream.cuda_stream)) + return original_record_stream(tensor, stream) + + monkeypatch.setattr(torch.Tensor, "record_stream", record_stream) + param = MultimodalParams( + multimodal_data={ + "multimodal_embedding": embeddings[0] if embedding_count == 1 else embeddings + }, + encoder_event=producer_event, + ) + + with torch.cuda.stream(consumer_stream): + gathered = get_multimodal_embeddings( + encoder_forward_fn=lambda _: pytest.fail("attached embeddings must skip the encoder"), + multimodal_params=[param], + ) + + consumer_stream.synchronize() + expected_calls = { + (embedding.data_ptr(), consumer_stream.cuda_stream) for embedding in embeddings + } + assert expected_calls.issubset(record_stream_calls) + torch.testing.assert_close(gathered[0], torch.ones(4 * embedding_count, device="cuda")) diff --git a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py index bc525e20d805..b43f5233acde 100644 --- a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py +++ b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py @@ -76,6 +76,8 @@ def encode_multimodal_inputs(self, multimodal_params, **encoder_kwargs) -> torch class NoEmbeddingMetadataMultimodalModel(DummyMultimodalModel): + supports_encoder_cache = True + @property def embedding_dim(self) -> int: raise NotImplementedError @@ -86,6 +88,8 @@ def embedding_dtype(self) -> torch.dtype: class CountingEncoderMultimodalModel(DummyMultimodalModel): + supports_encoder_cache = True + def __init__( self, embedding: Embedding, @@ -255,6 +259,15 @@ def test_encoder_cache_first_request_writes_per_item_entries(): assert len(model._multimodal_encoder_cache) == 2 +def test_encoder_cache_requires_model_opt_in(): + model = DummyMultimodalModel(make_embedding(hidden_size=4), torch.tensor([7])) + model.model_config = ModelConfig( + multimodal_config=MultimodalConfig(encoder_cache_max_bytes=4096) + ) + + assert not model.encoder_cache_active + + def test_encoder_cache_creation_logs_embedding_row_capacity(): model = CountingEncoderMultimodalModel( make_embedding(hidden_size=4), diff --git a/tests/unittest/_torch/test_tensor_lru_cache.py b/tests/unittest/_torch/test_tensor_lru_cache.py index 389991e7ab6e..a91a6ff7620d 100644 --- a/tests/unittest/_torch/test_tensor_lru_cache.py +++ b/tests/unittest/_torch/test_tensor_lru_cache.py @@ -180,3 +180,82 @@ def write_and_read(index: int) -> None: hit = cache.get(index) if hit is not None: assert hit.numel() * hit.element_size() == 8 + + +def test_stream_aware_mode_leaves_cpu_cache_behavior_unchanged() -> None: + cache = TensorLRUCache[str](max_bytes=16, cuda_stream_aware=True) + source = torch.arange(4, dtype=torch.float32) + + assert cache.put("key", source) + cached = cache.get("key") + + assert cached is not None + assert cached.device.type == "cpu" + torch.testing.assert_close(cached, source) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cuda_stream_aware_cache_orders_cross_stream_producer_and_consumer() -> None: + cache = TensorLRUCache[str](max_bytes=16, cuda_stream_aware=True) + producer_stream = torch.cuda.Stream() + consumer_stream = torch.cuda.Stream() + source = torch.zeros(4, device="cuda") + + with torch.cuda.stream(producer_stream): + torch.cuda._sleep(10_000_000) + source.fill_(7) + assert cache.put("key", source) + + with torch.cuda.stream(consumer_stream): + cached = cache.get("key") + assert cached is not None + observed = cached.clone() + + consumer_stream.synchronize() + torch.testing.assert_close(observed, torch.full_like(observed, 7)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cuda_stream_aware_cache_pop_orders_cross_stream_consumer() -> None: + cache = TensorLRUCache[str](max_bytes=16, cuda_stream_aware=True) + producer_stream = torch.cuda.Stream() + consumer_stream = torch.cuda.Stream() + source = torch.zeros(4, device="cuda") + + with torch.cuda.stream(producer_stream): + torch.cuda._sleep(10_000_000) + source.fill_(5) + assert cache.put("key", source) + + with torch.cuda.stream(consumer_stream): + popped = cache.pop("key") + assert popped is not None + observed = popped.clone() + + consumer_stream.synchronize() + assert len(cache) == 0 + torch.testing.assert_close(observed, torch.full_like(observed, 5)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cuda_stream_aware_cache_preserves_evicted_consumer_storage() -> None: + cache = TensorLRUCache[str](max_bytes=4 * 1024, cuda_stream_aware=True) + producer_stream = torch.cuda.Stream() + consumer_stream = torch.cuda.Stream() + source = torch.full((1024,), 3.0, device="cuda") + + with torch.cuda.stream(producer_stream): + assert cache.put("key", source) + + with torch.cuda.stream(consumer_stream): + cached = cache.get("key") + assert cached is not None + torch.cuda._sleep(10_000_000) + observed = cached.clone() + del cached + + with torch.cuda.stream(producer_stream): + assert cache.put("key", torch.full_like(source, 9.0)) + + consumer_stream.synchronize() + torch.testing.assert_close(observed, torch.full_like(observed, 3.0)) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 558a57c70208..17512efff3f7 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -749,10 +749,9 @@ def test_encoder_cache_max_bytes_rejects_invalid_values(self, value): def test_torch_llm_args_with_encoder_side_stream_max_ahead(self): args = TorchLlmArgs(model=llama_model_path, multimodal_config=MultimodalConfig( - encoder_side_stream_max_ahead=2, - encoder_cache_max_bytes=0, - )) + encoder_side_stream_max_ahead=2, )) assert args.multimodal_config.encoder_side_stream_max_ahead == 2 + assert args.multimodal_config.encoder_cache_max_bytes == 128 * 1024**2 def test_torch_llm_args_with_multimodal_video_pruning_rate(self): args = TorchLlmArgs( @@ -813,12 +812,14 @@ def test_encoder_cuda_graph_and_side_stream_max_ahead_are_exclusive(self): }, ) - def test_encoder_cache_and_side_stream_max_ahead_are_exclusive(self): - with pytest.raises(ValidationError, match="mutually exclusive"): - MultimodalConfig( - encoder_cache_max_bytes="1MiB", - encoder_side_stream_max_ahead=1, - ) + def test_encoder_cache_and_side_stream_max_ahead_can_be_combined(self): + config = MultimodalConfig( + encoder_cache_max_bytes="1MiB", + encoder_side_stream_max_ahead=1, + ) + + assert config.encoder_cache_max_bytes == 1024**2 + assert config.encoder_side_stream_max_ahead == 1 @pytest.mark.parametrize("kwargs", [