From b0e104508d18d5bf7c00fa778a78f6013913f443 Mon Sep 17 00:00:00 2001 From: Aswin Visva <31215515+aswinvisva@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:23:32 -0700 Subject: [PATCH 1/2] [None][feat] Multimodal encoder cache: per-item partial hits * Why? Persistent multimodal encoder cache is currently all-or-nothing per `MultimodalParams`: if any item of a request misses the cache, every item is re-encoded. In practice this makes the cache inert for common VLM workloads where consecutive requests share most of their images but differ in one -- e.g. multi-turn conversations that grow context, or comparison tasks that anchor on a shared reference set. * What? Add an opt-in per-item partial-hit path to `MultimodalModelMixin`: - `EncoderCachePartition` records the per-item cache lookup result (hits / miss indices / keys). - `_partition_encoder_cache` replaces `_attach_encoder_cache_hit` and records every item lookup rather than bailing on first miss. - `_encode_with_partial_cache` encodes only the miss items via a residual `MultimodalParams`, then interleaves cached and freshly encoded rows in original per-item order so downstream `get_multimodal_embeddings` treats the param as fully cached. - `_apply_metadata_slice` re-slices `multimodal_embedding_lengths` and `multimodal_hashes` on the residual so every model gets the parallel per-item metadata slice identically. - Models opt in with `supports_granular_encoder_cache = True` and implement `slice_multimodal_data_items`, which slices raw modality tensors (`pixel_values`, `image_grid_thw`, ...) for the given item indices. Default remains `False`; existing models continue on the all-or-nothing path unchanged. Onboard Mistral 3 (Mistral-Small-3.1-24B) as the first opted-in model by slicing `pixel_values` on dim 0 and `image_sizes` list in parallel. * Numbers Real trtllm-serve run against `mistralai/Mistral-Small-3.1-24B-Instruct-2503` on 1x H200 (BF16, TP=1, max_seq_len=65536, encoder_cache_max_bytes=2GiB). Two requests each with 100 items of 512x512 images; request 2 shares 99 items with request 1 and adds 1 novel item: | config | req1 [100 fresh] | req2 [99 shared + 1 novel] | req2 delta | |----------|-----------------:|---------------------------:|-----------:| | baseline | 5966 ms | 5543 ms | -423 ms | | granular | 5962 ms | 5136 ms | -826 ms | Direct req2 comparison: granular is 406 ms faster than the all-or-nothing baseline, matching the encoder cost of 99 Pixtral + projector passes that granular skips. Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com> --- .../_torch/models/modeling_mistral.py | 26 +++ .../models/modeling_multimodal_mixin.py | 200 +++++++++++++++--- 2 files changed, 193 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 2246a73f8fdd..4bc2ebe2ab88 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -670,6 +670,7 @@ class Mistral3VLM(MultimodalModelMixin, PreTrainedModel): """ supports_encoder_cache = True + supports_granular_encoder_cache = True def __init__( self, @@ -825,6 +826,31 @@ def encode_multimodal_inputs( mm_embeds = self._vision_forward(list(multimodal_params)) return mm_embeds[0] + def slice_multimodal_data_items( + self, + param: MultimodalParams, + item_indices: Sequence[int], + ) -> MultimodalParams: + # One item == one image. `pixel_values` is `[B, C, H, W]` with B parallel to + # `image_sizes` (a Python list of `[h, w]`), so both are sliced along the item axis. + src_image = param.multimodal_data["image"] + indices = list(item_indices) + sliced_image = { + **src_image, + "pixel_values": src_image["pixel_values"][indices], + "image_sizes": [src_image["image_sizes"][i] for i in indices], + } + new_multimodal_data = {**param.multimodal_data, "image": sliced_image} + + # Shallow-copy `multimodal_input` so the mixin's metadata slice can rewrite + # `multimodal_hashes` on the residual without mutating the source. + residual_input = (copy.copy(param.multimodal_input) + if param.multimodal_input is not None else None) + return MultimodalParams( + multimodal_data=new_multimodal_data, + multimodal_input=residual_input, + ) + def get_language_model_forward_kwargs( self, *, diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index b1a2c5df6845..ef0061a0af71 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -115,6 +115,28 @@ class PreparedLlmInputs: extra_embeds: Sequence[torch.Tensor] = () +@dataclass(frozen=True) +class EncoderCachePartition: + """Per-item cache partition for a single `MultimodalParams`. + + `hits` maps item index to its cached embedding row-block; `miss_indices` lists item + indices that still require encoder work; `keys` is aligned to item order so miss + embeddings can be written back after they are computed. + """ + + hits: Dict[int, torch.Tensor] + miss_indices: list[int] + keys: list[Hashable] + + @property + def is_full_hit(self) -> bool: + return bool(self.keys) and not self.miss_indices + + @property + def is_full_miss(self) -> bool: + return bool(self.keys) and not self.hits + + class MultimodalModelMixin: """Template-method mixin for PyTorch multimodal causal LM models. @@ -124,15 +146,24 @@ class MultimodalModelMixin: Current limitations: * For the time being, the persistent multimodal encoder cache stores per-item embeddings for - single-modality `MultimodalParams` objects. - * Cache reuse is all-or-nothing for each `MultimodalParams` object: every item in that object - hit the cache before cached embeddings are reused. Mixed-modality `MultimodalParams` objects - bypass the persistent cache. + single-modality `MultimodalParams` objects. Mixed-modality objects bypass the cache. + * Per-item cache reuse is all-or-nothing by default. Models can opt into partial reuse by + implementing `slice_multimodal_data_items` and setting + `supports_granular_encoder_cache = True`; miss items are re-encoded and interleaved with + cached items in original per-item order. """ supports_encoder_cache: ClassVar[bool] = False """Whether the model's production forward path uses the persistent encoder cache.""" + supports_granular_encoder_cache: ClassVar[bool] = False + """Whether the mixin may reuse a partial cache hit for this model. + + When True, a partially cached `MultimodalParams` is handled by encoding only its miss + items and interleaving cached items back in per-item order. Requires the model to + implement `slice_multimodal_data_items`. + """ + model_config: ModelConfig _multimodal_encoder_cache: Optional[TensorLRUCache] = None @@ -240,6 +271,21 @@ def after_active_multimodal_embeddings( # them as extra embeds without changing the base flow. return active_embeddings, () + def slice_multimodal_data_items( + self, + param: MultimodalParams, + item_indices: Sequence[int], + ) -> MultimodalParams: + """Return a `MultimodalParams` whose raw modality inputs contain only + `item_indices` from `param`, in that order. + + Only the model knows how to slice its own modality tensors (`pixel_values`, + `image_grid_thw`, packed audio features, ...). The parallel per-item metadata + (`multimodal_embedding_lengths`, `multimodal_hashes`) is re-sliced by the mixin + after this returns, so implementations should focus on modality data only. + """ + raise NotImplementedError + # A future optional mixin-owned forward can build on the same template method. def prepare_multimodal_inputs( self, @@ -312,16 +358,33 @@ def _get_or_encode_multimodal_embeddings( the single tensor contract for both encoded and cached-only paths. """ encoder_cache = self._get_multimodal_encoder_cache() - cache_misses = [] + cache_misses: list[MultimodalParams] = [] + partial_hits: list[tuple[MultimodalParams, EncoderCachePartition]] = [] 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. continue - if not self._attach_encoder_cache_hit(param, encoder_cache): + partition = self._partition_encoder_cache(param, encoder_cache) + if partition is None or partition.is_full_miss: + cache_misses.append(param) + continue + if partition.is_full_hit: + param.multimodal_data["multimodal_embedding"] = self._assemble_full_embedding( + partition.hits, len(partition.keys) + ) + continue + if self.supports_granular_encoder_cache: + partial_hits.append((param, partition)) + else: cache_misses.append(param) + if partial_hits: + # `encoder_cache` is non-None here because partitions are only produced when the cache + # exists. + self._encode_with_partial_cache(partial_hits, encoder_cache) + embeddings = get_multimodal_embeddings( encoder_forward_fn=self.encode_multimodal_inputs, multimodal_params=list(multimodal_params), @@ -476,47 +539,118 @@ def _encoder_cache_keys( ] @classmethod - def _attach_encoder_cache_hit( + def _partition_encoder_cache( cls, param: MultimodalParams, encoder_cache: TensorLRUCache, - ) -> bool: - """Attach a full persistent-cache hit and report whether one was found.""" + ) -> Optional[EncoderCachePartition]: + """Look up every item of `param` and return the per-item partition. + + Returns `None` when the param is not cacheable (mixed modality, missing metadata, + request-local embedding already attached); the caller treats that as a full miss. + """ if param.multimodal_data.get("multimodal_embedding") is not None: logger.debug( f"{_MM_ENCODER_CACHE_LOG_NAME}: request-local multimodal embedding present; " "skipping persistent cache lookup" ) - return False + return None keys = cls._encoder_cache_keys(param) if not keys: - return False - - cached_embeddings = [] - for key in keys: - cached_embedding = encoder_cache.get(key) - if cached_embedding is None: - # TODO(TRTLLM-13996): allow re-computing only the uncached items. - # `get_multimodal_embeddings` treats a param as either fully cached or uncached. - # Attaching partial hits would make the later concatenated tensor ambiguous because - # there is no placeholder for missing item rows inside `multimodal_embedding`. - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: cache miss; hit_items={len(cached_embeddings)}," - f" total_items={len(keys)}." - ) - return False - cached_embeddings.append(cached_embedding) + return None + + hits: Dict[int, torch.Tensor] = {} + miss_indices: list[int] = [] + for i, key in enumerate(keys): + cached = encoder_cache.get(key) + if cached is None: + miss_indices.append(i) + else: + hits[i] = cached - if len(cached_embeddings) == 1: - param.multimodal_data["multimodal_embedding"] = cached_embeddings[0] - else: - param.multimodal_data["multimodal_embedding"] = torch.cat(cached_embeddings, dim=0) logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: full cache hit for {len(keys)} item entries, " - f"rows={param.multimodal_data['multimodal_embedding'].shape[0]}." + f"{_MM_ENCODER_CACHE_LOG_NAME}: partition hit_items={len(hits)}, " + f"miss_items={len(miss_indices)}, total_items={len(keys)}." ) - return True + return EncoderCachePartition(hits=hits, miss_indices=miss_indices, keys=keys) + + @staticmethod + def _assemble_full_embedding( + item_tensors: Dict[int, torch.Tensor], + total_items: int, + ) -> torch.Tensor: + """Concatenate per-item embedding tensors in item-index order. + + `item_tensors` must contain every index in `[0, total_items)`. + """ + if total_items == 1: + return item_tensors[0] + return torch.cat([item_tensors[i] for i in range(total_items)], dim=0) + + @staticmethod + def _apply_metadata_slice( + residual: MultimodalParams, + source: MultimodalParams, + item_indices: Sequence[int], + ) -> None: + """Overwrite `residual`'s per-item metadata to match the sliced items. + + Models slice raw modality tensors in `slice_multimodal_data_items`; the mixin owns + the parallel per-item metadata slice so every model gets it identically. + """ + source_lengths = source.multimodal_data["multimodal_embedding_lengths"] + residual.multimodal_data["multimodal_embedding_lengths"] = [ + source_lengths[i] for i in item_indices + ] + if residual.multimodal_input is not None and source.multimodal_input is not None: + source_hashes = source.multimodal_input.multimodal_hashes + residual.multimodal_input.multimodal_hashes = [source_hashes[i] for i in item_indices] + + def _encode_with_partial_cache( + self, + partials: Sequence[tuple[MultimodalParams, EncoderCachePartition]], + encoder_cache: TensorLRUCache, + ) -> None: + """Encode only the miss items of each partial-hit param and stitch results. + + After this returns, each param's `multimodal_embedding` has the same shape as a + full encoder run so downstream `get_multimodal_embeddings` treats it as fully + cached. + """ + for param, partition in partials: + residual = self.slice_multimodal_data_items(param, partition.miss_indices) + self._apply_metadata_slice(residual, param, partition.miss_indices) + + residual_output = self.encode_multimodal_inputs([residual]) + miss_lengths = [ + param.multimodal_data["multimodal_embedding_lengths"][i] + for i in partition.miss_indices + ] + miss_tensors = torch.split(residual_output, miss_lengths, dim=0) + + by_item: Dict[int, torch.Tensor] = dict(partition.hits) + for miss_idx, tensor in zip(partition.miss_indices, miss_tensors, strict=True): + by_item[miss_idx] = tensor + param.multimodal_data["multimodal_embedding"] = self._assemble_full_embedding( + by_item, len(partition.keys) + ) + + inserted = 0 + rejected = 0 + for miss_idx, tensor in zip(partition.miss_indices, miss_tensors, strict=True): + if encoder_cache.put(partition.keys[miss_idx], tensor): + inserted += 1 + else: + rejected += 1 + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: partial-hit encode " + f"total_items={len(partition.keys)} " + f"hit_items={len(partition.hits)} " + f"encoded_items={len(partition.miss_indices)} " + f"cache_writes_inserted={inserted} " + f"cache_writes_rejected={rejected}" + ) @classmethod def _write_encoder_cache_entries( From 89f669466c08ce0624045739bb0bf7dd252a2ba3 Mon Sep 17 00:00:00 2001 From: Aswin Visva <31215515+aswinvisva@users.noreply.github.com> Date: Fri, 24 Jul 2026 13:43:49 -0700 Subject: [PATCH 2/2] [None][refactor] Multimodal encoder cache: public API + mixin default slicer * Why? To land ahead of the item-level MM encoder scheduling work on yechank-nvidia/multimodal-encoder-runtime-scheduling without forcing a churny rebase, the partial-cache primitives need a stable public surface his executor can call into. Two smaller shifts fall out naturally: - the per-model slicing hook can be a mixin default for the two common layouts, removing the only Mistral 3 code delta; - the second opt-in flag is redundant with `supports_encoder_cache`. * What? - Rename the partial-cache primitives to drop the underscore prefix so callers outside the mixin can consume them: `_partition_encoder_cache` -> `partition_encoder_cache` `_assemble_full_embedding` -> `assemble_full_embedding` `slice_multimodal_data_items` -> `build_multimodal_encoder_input` The last rename also aligns naming with the item-scheduling branch, which already uses `build_multimodal_encoder_input`. - Provide a default `build_multimodal_encoder_input` in the mixin covering the two common single-modality layouts: * Stacked on dim 0 (Pixtral / LLaVA-family): slice `pixel_values` dim 0 and the parallel `image_sizes` list. * Packed with `*_grid_thw` offsets (Qwen2-VL family): prefix-sum patch counts to locate each item's slab, concat requested items in item-index order, slice `image_grid_thw` in parallel. Models with a different layout override. - Drop `supports_granular_encoder_cache`. `supports_encoder_cache = True` already carries the opt-in signal; a model that opts into the cache and has an unhandled layout gets a loud `NotImplementedError` from the default slicer -- the correct signal to override. - Delete Mistral 3's `build_multimodal_encoder_input` override: the default now covers it. Mistral 3's net diff shrinks to zero. * Numbers Same 1x H200 / Mistral-Small-3.1-24B / N=100 / 99-shared A/B run as the previous commit, now with the mixin default doing the slicing: req1 [100 fresh] 5928 ms req2 [99 shared + 1 novel] 5114 ms (delta -814 ms / -13.7%) Within measurement noise of the override-based numbers, confirming the default reproduces the override end-to-end. Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com> --- .../_torch/models/modeling_mistral.py | 26 ----- .../models/modeling_multimodal_mixin.py | 106 +++++++++++++----- 2 files changed, 78 insertions(+), 54 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 4bc2ebe2ab88..2246a73f8fdd 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -670,7 +670,6 @@ class Mistral3VLM(MultimodalModelMixin, PreTrainedModel): """ supports_encoder_cache = True - supports_granular_encoder_cache = True def __init__( self, @@ -826,31 +825,6 @@ def encode_multimodal_inputs( mm_embeds = self._vision_forward(list(multimodal_params)) return mm_embeds[0] - def slice_multimodal_data_items( - self, - param: MultimodalParams, - item_indices: Sequence[int], - ) -> MultimodalParams: - # One item == one image. `pixel_values` is `[B, C, H, W]` with B parallel to - # `image_sizes` (a Python list of `[h, w]`), so both are sliced along the item axis. - src_image = param.multimodal_data["image"] - indices = list(item_indices) - sliced_image = { - **src_image, - "pixel_values": src_image["pixel_values"][indices], - "image_sizes": [src_image["image_sizes"][i] for i in indices], - } - new_multimodal_data = {**param.multimodal_data, "image": sliced_image} - - # Shallow-copy `multimodal_input` so the mixin's metadata slice can rewrite - # `multimodal_hashes` on the residual without mutating the source. - residual_input = (copy.copy(param.multimodal_input) - if param.multimodal_input is not None else None) - return MultimodalParams( - multimodal_data=new_multimodal_data, - multimodal_input=residual_input, - ) - def get_language_model_forward_kwargs( self, *, diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index ef0061a0af71..c5f428936fb5 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -15,6 +15,7 @@ # SPDX-License-Identifier: Apache-2.0 import contextlib +import copy from dataclasses import dataclass from typing import ( TYPE_CHECKING, @@ -147,23 +148,15 @@ class MultimodalModelMixin: * For the time being, the persistent multimodal encoder cache stores per-item embeddings for single-modality `MultimodalParams` objects. Mixed-modality objects bypass the cache. - * Per-item cache reuse is all-or-nothing by default. Models can opt into partial reuse by - implementing `slice_multimodal_data_items` and setting - `supports_granular_encoder_cache = True`; miss items are re-encoded and interleaved with - cached items in original per-item order. + * A partially cached `MultimodalParams` is handled by encoding only its miss items + and interleaving cached items back in original per-item order. The default + `build_multimodal_encoder_input` handles stacked-on-dim-0 and packed-with-grid-thw + layouts; models with other layouts override that method. """ supports_encoder_cache: ClassVar[bool] = False """Whether the model's production forward path uses the persistent encoder cache.""" - supports_granular_encoder_cache: ClassVar[bool] = False - """Whether the mixin may reuse a partial cache hit for this model. - - When True, a partially cached `MultimodalParams` is handled by encoding only its miss - items and interleaving cached items back in per-item order. Requires the model to - implement `slice_multimodal_data_items`. - """ - model_config: ModelConfig _multimodal_encoder_cache: Optional[TensorLRUCache] = None @@ -271,7 +264,7 @@ def after_active_multimodal_embeddings( # them as extra embeds without changing the base flow. return active_embeddings, () - def slice_multimodal_data_items( + def build_multimodal_encoder_input( self, param: MultimodalParams, item_indices: Sequence[int], @@ -279,12 +272,72 @@ def slice_multimodal_data_items( """Return a `MultimodalParams` whose raw modality inputs contain only `item_indices` from `param`, in that order. - Only the model knows how to slice its own modality tensors (`pixel_values`, - `image_grid_thw`, packed audio features, ...). The parallel per-item metadata + Default handles two common single-modality layouts: + + - Stacked on dim 0 (Mistral 3 / Pixtral / LLaVA-family): `pixel_values` + `[B, C, H, W]` with a parallel `image_sizes` list; both sliced by item. + - Packed with `*_grid_thw` offsets (Qwen2-VL family): `pixel_values` + `[total_patches, feat]` + `image_grid_thw` `[B, 3]`; prefix-summed patch + counts locate each item's slice, and `image_grid_thw` is sliced in parallel. + + Models with a different layout (e.g. mixed-modality per param, custom packed + formats) override this method. The parallel per-item metadata (`multimodal_embedding_lengths`, `multimodal_hashes`) is re-sliced by the mixin - after this returns, so implementations should focus on modality data only. + after this returns, so overrides need only handle modality data. """ - raise NotImplementedError + modality = self._encoder_cache_modality(param) + if modality is None: + raise NotImplementedError( + "Default `build_multimodal_encoder_input` only supports single-modality " + "params. Override for other layouts." + ) + modality_data = param.multimodal_data[modality] + if not isinstance(modality_data, dict): + raise TypeError( + f"multimodal_data[{modality!r}] must be a dict, got {type(modality_data).__name__}" + ) + + indices = list(item_indices) + grid_key = {"image": "image_grid_thw", "video": "video_grid_thw"}.get(modality) + pixel_key = {"image": "pixel_values", "video": "pixel_values_videos"}.get(modality) + + if grid_key and pixel_key and grid_key in modality_data and pixel_key in modality_data: + # Packed layout: prefix-sum patch counts to locate each item's slab, then + # concat the requested subset in item-index order. + grids = modality_data[grid_key] + patch_counts = [int(c) for c in torch.prod(grids, dim=1).tolist()] + per_item = torch.split(modality_data[pixel_key], patch_counts, dim=0) + sliced = { + **modality_data, + pixel_key: torch.cat([per_item[i] for i in indices], dim=0), + grid_key: grids[indices], + } + elif ( + modality == "image" + and "pixel_values" in modality_data + and "image_sizes" in modality_data + ): + # Stacked layout: dim-0 select from `pixel_values` and list-index `image_sizes`. + sliced = { + **modality_data, + "pixel_values": modality_data["pixel_values"][indices], + "image_sizes": [modality_data["image_sizes"][i] for i in indices], + } + else: + raise NotImplementedError( + f"Default `build_multimodal_encoder_input` cannot slice {modality} layout " + f"with fields {sorted(modality_data)}; override this method." + ) + + # Shallow-copy `multimodal_input` so `_apply_metadata_slice` can rewrite + # `multimodal_hashes` on the residual without mutating the source. + residual_input = ( + copy.copy(param.multimodal_input) if param.multimodal_input is not None else None + ) + return MultimodalParams( + multimodal_data={**param.multimodal_data, modality: sliced}, + multimodal_input=residual_input, + ) # A future optional mixin-owned forward can build on the same template method. def prepare_multimodal_inputs( @@ -366,19 +419,16 @@ def _get_or_encode_multimodal_embeddings( # The forward that attached this request-local embedding already populated the # persistent cache. continue - partition = self._partition_encoder_cache(param, encoder_cache) + partition = self.partition_encoder_cache(param, encoder_cache) if partition is None or partition.is_full_miss: cache_misses.append(param) continue if partition.is_full_hit: - param.multimodal_data["multimodal_embedding"] = self._assemble_full_embedding( + param.multimodal_data["multimodal_embedding"] = self.assemble_full_embedding( partition.hits, len(partition.keys) ) continue - if self.supports_granular_encoder_cache: - partial_hits.append((param, partition)) - else: - cache_misses.append(param) + partial_hits.append((param, partition)) if partial_hits: # `encoder_cache` is non-None here because partitions are only produced when the cache @@ -539,7 +589,7 @@ def _encoder_cache_keys( ] @classmethod - def _partition_encoder_cache( + def partition_encoder_cache( cls, param: MultimodalParams, encoder_cache: TensorLRUCache, @@ -576,7 +626,7 @@ def _partition_encoder_cache( return EncoderCachePartition(hits=hits, miss_indices=miss_indices, keys=keys) @staticmethod - def _assemble_full_embedding( + def assemble_full_embedding( item_tensors: Dict[int, torch.Tensor], total_items: int, ) -> torch.Tensor: @@ -596,7 +646,7 @@ def _apply_metadata_slice( ) -> None: """Overwrite `residual`'s per-item metadata to match the sliced items. - Models slice raw modality tensors in `slice_multimodal_data_items`; the mixin owns + Models slice raw modality tensors in `build_multimodal_encoder_input`; the mixin owns the parallel per-item metadata slice so every model gets it identically. """ source_lengths = source.multimodal_data["multimodal_embedding_lengths"] @@ -619,7 +669,7 @@ def _encode_with_partial_cache( cached. """ for param, partition in partials: - residual = self.slice_multimodal_data_items(param, partition.miss_indices) + residual = self.build_multimodal_encoder_input(param, partition.miss_indices) self._apply_metadata_slice(residual, param, partition.miss_indices) residual_output = self.encode_multimodal_inputs([residual]) @@ -632,7 +682,7 @@ def _encode_with_partial_cache( by_item: Dict[int, torch.Tensor] = dict(partition.hits) for miss_idx, tensor in zip(partition.miss_indices, miss_tensors, strict=True): by_item[miss_idx] = tensor - param.multimodal_data["multimodal_embedding"] = self._assemble_full_embedding( + param.multimodal_data["multimodal_embedding"] = self.assemble_full_embedding( by_item, len(partition.keys) )