Skip to content

[None][feat] Multimodal encoder cache unification #17876

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yechank-nvidia wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[None][feat] Multimodal encoder cache unification #17876

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
164 changes: 103 additions & 61 deletions .claude/skills/trtllm-model-onboard-multimodal/SKILL.md
View file
Open in desktop

Large diffs are not rendered by default.

81 changes: 36 additions & 45 deletions tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1066,7 +1066,7 @@ def _get_or_encode_multimodal_embeddings(
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()
encoder_cache = self._multimodal_encoder_cache
cache_misses: list[MultimodalParams] = []
partial_hits: list[tuple[MultimodalParams, EncoderCachePartition]] = []
if encoder_cache is not None:
Expand Down Expand Up @@ -1107,60 +1107,51 @@ def _get_or_encode_multimodal_embeddings(
self._validate_embeddings(embeddings, multimodal_params)
return embeddings[0]

def _get_multimodal_encoder_cache(self) -> Optional[TensorLRUCache]:
"""Return the per-model full-request-path encoder clone cache, if enabled.
def _initialize_multimodal_encoder_cache(self, max_bytes: int) -> Optional[TensorLRUCache]:
"""Initialize the model-owned multimodal encoder-output cache once.

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.

Scope: the single encoder cache instance for a cache-enabled model
(`supports_encoder_cache`). The full-request (legacy inline-encode) consumers — side-stream
prefetch, `mm_encoder_only`/disagg encoding — populate and read it inline; the
item-scheduling path consumes the same instance read-through at encode time
(`ModelEngine.forward_multimodal_encoder_items`). The key format is shared
(`_encoder_cache_item_key`) so hits cross between paths. The item path's recorded outputs
are cloned, so cache eviction never invalidates an in-flight request.
`ModelEngine` resolves `max_bytes` as the larger of the item-scheduling
output budget and persistent-reuse capacity before runtime access.
Zero leaves the cache disabled.
"""
if not self.encoder_cache_active:
if self._multimodal_encoder_cache is not None:
return self._multimodal_encoder_cache

if max_bytes == 0:
logger.debug_once(
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.",
f"{_MM_ENCODER_CACHE_LOG_NAME}: disabled because neither item scheduling nor "
"persistent reuse requires storage.",
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
# entire batch allocation while cache accounting charges only its logical size. This
# briefly needs source and clone memory during insertion, but preserves existing cache
# entries when the copy cannot be allocated.
self._multimodal_encoder_cache = TensorLRUCache(
max_bytes,
name=_MM_ENCODER_CACHE_LOG_NAME,
cuda_stream_aware=multimodal_config.encoder_side_stream_max_ahead > 0,
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
embedding_dtype = self.embedding_dtype
except NotImplementedError:
logger.info(
f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, "
"embedding row capacity unavailable because the model does not implement "
"embedding_dim and embedding_dtype."
)
else:
bytes_per_embedding_row = (
embedding_dim * torch.empty((), dtype=embedding_dtype).element_size()
)
max_embedding_rows = max_bytes // bytes_per_embedding_row
logger.info(
f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, "
f"max_embedding_rows={max_embedding_rows}, embedding_dim={embedding_dim}, "
f"embedding_dtype={embedding_dtype}"
)
try:
embedding_dim = self.embedding_dim
embedding_dtype = self.embedding_dtype
except NotImplementedError:
logger.info(
f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, "
"embedding row capacity unavailable because the model does not implement "
"embedding_dim and embedding_dtype."
)
else:
bytes_per_embedding_row = (
embedding_dim * torch.empty((), dtype=embedding_dtype).element_size()
)
max_embedding_rows = max_bytes // bytes_per_embedding_row
logger.info(
f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, "
f"max_embedding_rows={max_embedding_rows}, embedding_dim={embedding_dim}, "
f"embedding_dtype={embedding_dtype}"
)
return self._multimodal_encoder_cache

@staticmethod
Expand Down Expand Up @@ -1678,7 +1669,7 @@ def _dispatch_cross_iter_prefetch(
encoder_event = None
try:
with _run_on_aux_stream(aux_stream) as encoder_event:
encoder_cache = model._get_multimodal_encoder_cache() if encoder_cache_enabled else None
encoder_cache = model._multimodal_encoder_cache if encoder_cache_enabled else None
cache_misses: list[MultimodalParams] = []
partial_hits: list[tuple[MultimodalParams, EncoderCachePartition]] = []
if encoder_cache is None:
Expand Down
36 changes: 17 additions & 19 deletions tensorrt_llm/_torch/pyexecutor/_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -955,7 +955,7 @@ def _create_dummy_encoder_inputs(self) -> List[MultimodalParams]:
return [MultimodalParams(multimodal_data=mm_data)]

def _encode_dummy_inputs(self) -> Optional[torch.Tensor]:
"""Run the full-budget MM encoder and retain request-owned output storage."""
"""Run the full-budget MM encoder and retain equivalent store output."""
if not self._dummy_encoder_inputs:
return None

Expand All @@ -975,7 +975,7 @@ def _encode_dummy_inputs(self) -> Optional[torch.Tensor]:
)
output = self._model_engine.model.encode_multimodal_inputs(
encoder_inputs)
# Runtime item state owns detached copies rather than views of
# The runtime store owns detached copies rather than views of
# an encoder batch. Reproduce that allocation boundary here.
return output.detach().clone()
finally:
Expand All @@ -984,20 +984,11 @@ def _encode_dummy_inputs(self) -> Optional[torch.Tensor]:
def _get_multimodal_encoder_memory_reserve(self,
profiled_output_bytes: int = 0
) -> int:
"""Return output and cache capacity absent from the measured peak."""
output_budget = getattr(self._model_engine,
"mm_encoder_output_budget_bytes", None)
unprofiled_output_bytes = max(0, (output_budget or 0) -
profiled_output_bytes)

model = self._model_engine.model
cache_bytes = 0
if (isinstance(model, MultimodalModelMixin)
and model.encoder_cache_active
and model.model_config.multimodal_config is not None):
cache_bytes = (
model.model_config.multimodal_config.encoder_cache_max_bytes)
return unprofiled_output_bytes + cache_bytes
"""Return unified encoder-store capacity absent from the measured peak."""
encoder_cache = self._model_engine.mm_encoder_cache
if encoder_cache is None:
return 0
return max(0, encoder_cache.max_bytes - profiled_output_bytes)

def _get_token_num_for_estimation(self) -> int:
"""Compute KV cache capacity required for estimate_max_kv_cache_tokens to succeed."""
Expand Down Expand Up @@ -1221,7 +1212,7 @@ def configure_kv_cache_capacity(self,

if py_executor is not None and not self._skip_est:
# Run the MM encoder at its independent token budget, then keep the
# resulting request-owned embeddings resident while the text-only
# equivalent cache-owned embeddings resident while the text-only
# LLM dummy fills max_num_tokens.
encoder_profile_output = self._encode_dummy_inputs()
if encoder_profile_output is not None:
Expand Down Expand Up @@ -3384,7 +3375,8 @@ def create_py_executor_instance(
reorder_policy_config.policy_args.agent_inflight_seq_num)
scheduler = SimpleScheduler(capacity_scheduler, mb_scheduler)

if getattr(model_engine, "mm_encoder_item_scheduling_enabled", False):
if (getattr(model_engine, "mm_encoder_item_scheduling_enabled", False)
and model_engine.mapping.is_first_pp_rank()):
# Wrap the LLM scheduler with atomic MM item budgeting. ModelEngine
# already validated model-capability-dependent feature combinations.
multimodal_config = llm_args.multimodal_config
Expand All @@ -3397,13 +3389,19 @@ def create_py_executor_instance(
logger.info("Eager multimodal encoder scheduling is enabled for "
"capacity-rejected active requests.")
scheduler_cls = MultimodalEagerEncoderScheduler
encoder_cache = model_engine.mm_encoder_cache
if encoder_cache is None:
raise RuntimeError(
"MM encoder item scheduling requires its unified output cache")
scheduler = scheduler_cls(
scheduler,
max_batch_size=model_engine.encoder_batch_size,
max_num_tokens=model_engine.encoder_max_num_tokens,
output_budget_bytes=model_engine.mm_encoder_output_budget_bytes,
encoder_cache=encoder_cache,
get_item_cache_keys=model_engine.get_mm_encoder_item_cache_keys,
bytes_per_encoder_embedding=(
model_engine.bytes_per_mm_encoder_embedding),
retain_cache_entries=model_engine.model.encoder_cache_active,
)

config = model_engine.model.model_config.pretrained_config
Expand Down
Loading
Loading