[TRTLLM-14304][feat] Integrate embeddings cache with encoder side-stream by 2ez4bz · Pull Request #16554 · NVIDIA/TensorRT-LLM

2ez4bz · 2026-07-17T22:05:33Z

Summary by CodeRabbit

New Features
- Multimodal encoder side-stream prefetching can now be combined with persistent embeddings caching.
- Cross-iteration prefetch becomes encoder-cache-aware: cache hits reuse previously stored embeddings; misses trigger encoding and populate the persistent cache.
- Stream-safe cache consumption is improved so prefetched embeddings remain valid on the consumer stream.
Bug Fixes / Correctness
- Request-local (“uncacheable”) multimodal embeddings continue to follow the correct fallback path (no persistent-cache insert/update).
- TensorLRUCache now supports CUDA stream-aware entries via producer events, with consumers waiting/recording streams to avoid premature reuse.
Documentation
- Updated supported-models docs to clarify how encoder_side_stream_max_ahead interacts with multimodal embeddings caching and the combined peak-memory implications.

Dev Engineer Review

Config / API consistency
- MultimodalConfig validation was updated to allow encoder_side_stream_max_ahead + encoder_cache_max_bytes together (memory limits treated as additive), while keeping the existing incompatibility with encoder_cuda_graph.
- Updated unit tests to remove mutual-exclusion expectations for the cache + side-stream lookahead knobs.
Core implementation
- Introduced MultimodalModelMixin.encoder_cache_active to centralize gating for “cache enabled” based on model opt-in (supports_encoder_cache) and configured cache capacity.
- Persistent encoder cache creation and stream behavior are now aligned with side-stream prefetch via TensorLRUCache(..., cuda_stream_aware=encoder_side_stream_max_ahead>0).
- Cross-iteration prefetch dispatcher now builds/propagates multimodal metadata (MultimodalInput and MultimodalRuntimeData) when caching is active and routes through persistent-cache lookup/write (_get_or_encode_multimodal_embeddings); otherwise it uses the chunked prefill embedding storage path (_store_chunked_prefill_embeddings).
- Correctness/lifetime handling:
  - get_multimodal_embeddings now waits on param.encoder_event (when present) and then record_stream(consumer_stream) for CUDA embedding tensors to keep storage valid on the consuming stream.
  - Doc/comments updated to ensure side-stream prefetched requests must go through get_multimodal_embeddings so consumer-side stream attachment occurs before reading.
Memory reservation
- PyExecutor memory reservation now adds multimodal encoder-cache bytes only when the model is a MultimodalModelMixin instance and model.encoder_cache_active is true, preventing incorrect reservation for models that don’t opt in / aren’t compatible.
Supporting plumbing
- TensorLRUCache stream-aware mode adds optional producer events on put() and forces cross-stream waiting + record_stream() in get()/pop() when enabled, with associated unit tests.
Testing artifacts / CI health
- Author reported multiple CI failures for specific commits followed by later runs showing bot SUCCESS but L0 pipeline FAILURE; needs follow-up confirmation that the failures are fully resolved and green before merge.

QA Engineer Review

Test code changes (under tests/): yes

`tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py`

Added
- _CacheStubModel stub with supports_encoder_cache = True
- _make_cacheable_request
- CUDA tests:
  - test_cross_iter_prefetch_populates_and_reuses_persistent_cache
  - test_cross_iter_prefetch_mixed_cache_hit_and_miss_encodes_only_miss
  - test_cross_iter_prefetch_cache_model_preserves_uncacheable_fallbacks
  - test_cross_iter_prefetch_does_not_synchronize_main_stream
  - test_cross_iter_prefetch_does_not_rewrite_request_local_embedding
  - test_prefetched_embedding_records_main_consumer_stream (parametrized)

`tests/unittest/_torch/multimodal/test_multimodal_mixin.py`

Modified
- Dummy models opt-in via supports_encoder_cache = True
Added
- test_encoder_cache_requires_model_opt_in

`tests/unittest/_torch/test_tensor_lru_cache.py`

Added
- test_stream_aware_mode_leaves_cpu_cache_behavior_unchanged
- CUDA-only cross-stream ordering/lifetime tests:
  - test_cuda_stream_aware_cache_orders_cross_stream_producer_and_consumer
  - test_cuda_stream_aware_cache_pop_orders_cross_stream_consumer
  - test_cuda_stream_aware_cache_preserves_evicted_consumer_storage

`tests/unittest/llmapi/test_llm_args.py`

Modified
- Updated/expanded multimodal config test expectations to allow combining encoder_side_stream_max_ahead and encoder_cache_max_bytes
Removed
- test_encoder_cache_and_side_stream_max_ahead_are_exclusive
Added
- test_encoder_cache_and_side_stream_max_ahead_can_be_combined
Coverage in tests/integration/test_lists/
- Not provided / not verifiable from the supplied context which integration/CBTS entries cover these new unit tests.
- Verdict: needs follow-up (confirm whether CI test-db/qa runs exercise the new cache + cross-stream behaviors beyond unit coverage).

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

2ez4bz · 2026-07-17T22:11:51Z

 _MM_ENCODER_CACHE_LOG_NAME = "mm_encoder_cache"


+def _build_request_multimodal_input(


This was moved from tensorrt_llm/_torch/pyexecutor/model_engine.py to avoid circular dependencies.

2ez4bz · 2026-07-20T21:14:29Z

/bot run

coderabbitai · 2026-07-20T21:22:35Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8119e924-ebff-4c59-a8e9-3629fc495734

📥 Commits

Reviewing files that changed from the base of the PR and between bd94ed0 and 76324bd.

📒 Files selected for processing (9)

docs/source/models/supported-models.md
tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
tensorrt_llm/_torch/models/modeling_multimodal_utils.py
tensorrt_llm/_torch/pyexecutor/_util.py
tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/_torch/pyexecutor/py_executor.py
tensorrt_llm/_torch/tensor_lru_cache.py
tensorrt_llm/inputs/multimodal.py
tensorrt_llm/llmapi/llm_args.py

🚧 Files skipped from review as they are similar to previous changes (9)

tensorrt_llm/inputs/multimodal.py
tensorrt_llm/_torch/pyexecutor/_util.py
docs/source/models/supported-models.md
tensorrt_llm/_torch/pyexecutor/py_executor.py
tensorrt_llm/_torch/models/modeling_multimodal_utils.py
tensorrt_llm/_torch/tensor_lru_cache.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/_torch/models/modeling_multimodal_mixin.py

Walkthrough

Multimodal encoder caching now works with cross-iteration side-stream prefetch. Cache activation is centralized, CUDA stream ordering is tracked in TensorLRUCache, prefetch handles cache hits and misses, and configuration, documentation, and CUDA tests cover the combined behavior.

Changes

Multimodal encoder cache and prefetch

Layer / File(s)	Summary
Cache activation and configuration `tensorrt_llm/_torch/models/modeling_multimodal_mixin.py`, `tensorrt_llm/_torch/pyexecutor/...`, `tensorrt_llm/llmapi/llm_args.py`, `tests/unittest/_torch/multimodal/test_multimodal_mixin.py`	Encoder cache enablement is centralized through `encoder_cache_active`, shared request metadata construction is used, and side-stream prefetch can be combined with encoder caching.
Stream-aware tensor cache `tensorrt_llm/_torch/tensor_lru_cache.py`, `tests/unittest/_torch/test_tensor_lru_cache.py`	`TensorLRUCache` records producer events, waits on consuming streams, records tensor usage, and tests cross-stream `get`, `pop`, eviction, and CPU behavior.
Cross-iteration prefetch integration `tensorrt_llm/_torch/models/modeling_multimodal_mixin.py`, `tensorrt_llm/_torch/models/modeling_multimodal_utils.py`, `tensorrt_llm/inputs/multimodal.py`, `tensorrt_llm/_torch/pyexecutor/py_executor.py`	Prefetch candidates carry cache metadata, cache-enabled requests use persistent lookup and encoding, and uncached requests use request-local embedding storage with explicit consumer-stream attachment.
Validation and documentation `tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py`, `tests/unittest/llmapi/test_llm_args.py`, `docs/source/models/supported-models.md`	Tests cover cache reuse, mixed hits and misses, fallbacks, synchronization, and stream recording; configuration and memory behavior documentation are updated.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PyExecutor
  participant MultimodalModelMixin
  participant TensorLRUCache
  participant ConsumerStream
  PyExecutor->>MultimodalModelMixin: Dispatch cross-iteration prefetch
  MultimodalModelMixin->>TensorLRUCache: Lookup or store encoder embeddings
  TensorLRUCache-->>MultimodalModelMixin: Return embeddings and producer event
  MultimodalModelMixin->>ConsumerStream: Attach encoder event and embeddings
  ConsumerStream->>ConsumerStream: Wait and record tensor lifetime

Possibly related PRs

NVIDIA/TensorRT-LLM#14322: Adds the auxiliary side-stream cross-iteration multimodal encoder prefetch flow extended by this change.

Suggested reviewers: chienchunhung, qijune, brnguyen2, yihuilu512, dongxuy04

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is mostly the template comments and checkbox, with no actual issue, solution, or test coverage details filled in.	Fill in the Description and Test Coverage sections with the problem, the fix, and tests covering the new cache and side-stream paths.
Docstring Coverage	⚠️ Warning	Docstring coverage is 18.60% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly matches the main change: integrating the embeddings cache with encoder side-stream prefetching.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/multimodal/test_multimodal_mixin.py`:
- Around line 262-269: Update test_encoder_cache_requires_model_opt_in to
reference the public encoder_cache_active property on DummyMultimodalModel
instead of the nonexistent _encoder_cache_active attribute, while preserving the
assertion that the cache is inactive and _get_multimodal_encoder_cache() returns
None.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 643c2135-56e9-45e6-92f6-84de30895f71

📥 Commits

Reviewing files that changed from the base of the PR and between b520c6f and fa22df4.

📒 Files selected for processing (13)

docs/source/models/supported-models.md
tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
tensorrt_llm/_torch/models/modeling_multimodal_utils.py
tensorrt_llm/_torch/pyexecutor/_util.py
tensorrt_llm/_torch/pyexecutor/model_engine.py
tensorrt_llm/_torch/pyexecutor/py_executor.py
tensorrt_llm/_torch/tensor_lru_cache.py
tensorrt_llm/inputs/multimodal.py
tensorrt_llm/llmapi/llm_args.py
tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py
tests/unittest/_torch/multimodal/test_multimodal_mixin.py
tests/unittest/_torch/test_tensor_lru_cache.py
tests/unittest/llmapi/test_llm_args.py

tensorrt-cicd · 2026-07-20T21:23:53Z

PR_Github #60462 [ run ] triggered by Bot. Commit: fa22df4 Link to invocation

tensorrt-cicd · 2026-07-20T21:45:37Z

PR_Github #60462 [ run ] completed with state FAILURE. Commit: fa22df4
/LLM/main/L0_MergeRequest_PR pipeline #48795 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-20T23:17:53Z

/bot run

tensorrt-cicd · 2026-07-20T23:23:46Z

PR_Github #60491 [ run ] triggered by Bot. Commit: fa22df4 Link to invocation

tensorrt-cicd · 2026-07-21T01:41:47Z

PR_Github #60491 [ run ] completed with state FAILURE. Commit: fa22df4
/LLM/main/L0_MergeRequest_PR pipeline #48817 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-21T03:52:20Z

/bot run

tensorrt-cicd · 2026-07-21T03:57:42Z

PR_Github #60565 [ run ] triggered by Bot. Commit: e084f4d Link to invocation

tensorrt-cicd · 2026-07-21T06:13:08Z

PR_Github #60565 [ run ] completed with state FAILURE. Commit: e084f4d
/LLM/main/L0_MergeRequest_PR pipeline #48877 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-21T06:33:46Z

/bot run

tensorrt-cicd · 2026-07-21T06:39:58Z

PR_Github #60613 [ run ] triggered by Bot. Commit: e084f4d Link to invocation

tensorrt-cicd · 2026-07-21T08:12:21Z

PR_Github #60613 [ run ] completed with state FAILURE. Commit: e084f4d
/LLM/main/L0_MergeRequest_PR pipeline #48920 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-21T15:53:04Z

/bot run

tensorrt-cicd · 2026-07-21T15:59:14Z

PR_Github #60720 [ run ] triggered by Bot. Commit: e084f4d Link to invocation

tensorrt-cicd · 2026-07-21T17:37:25Z

PR_Github #60720 [ run ] completed with state FAILURE. Commit: e084f4d
/LLM/main/L0_MergeRequest_PR pipeline #49007 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-21T17:49:19Z

/bot run

tensorrt-cicd · 2026-07-21T17:55:36Z

PR_Github #60745 [ ] completed with state FAILURE. Commit: ``

Link to invocation

2ez4bz · 2026-07-21T23:29:44Z

/bot run

tensorrt-cicd · 2026-07-21T23:35:10Z

PR_Github #60816 [ run ] triggered by Bot. Commit: e084f4d Link to invocation

tensorrt-cicd · 2026-07-22T04:20:07Z

PR_Github #60816 [ run ] completed with state SUCCESS. Commit: e084f4d
/LLM/main/L0_MergeRequest_PR pipeline #49092 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

2ez4bz · 2026-07-24T20:35:33Z

/bot run

tensorrt-cicd · 2026-07-24T20:41:21Z

PR_Github #61638 [ run ] triggered by Bot. Commit: bd94ed0 Link to invocation

tensorrt-cicd · 2026-07-24T22:35:15Z

PR_Github #61638 [ run ] completed with state SUCCESS. Commit: bd94ed0
/LLM/main/L0_MergeRequest_PR pipeline #49845 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

* 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>

2ez4bz · 2026-07-25T04:51:56Z

/bot run

tensorrt-cicd · 2026-07-25T04:57:21Z

PR_Github #61681 [ run ] triggered by Bot. Commit: 76324bd Link to invocation

github-actions Bot assigned 2ez4bz Jul 17, 2026

2ez4bz commented Jul 17, 2026

View reviewed changes

2ez4bz force-pushed the dev-mm-encoder-cache branch 2 times, most recently from 220dd1d to fa22df4 Compare July 20, 2026 20:08

2ez4bz added the api-compatible Accepted LLM API contract change that is backwards-compatible label Jul 20, 2026

2ez4bz marked this pull request as ready for review July 20, 2026 21:14

2ez4bz requested review from a team as code owners July 20, 2026 21:14

2ez4bz requested review from QiJune, YihuiLu512, brnguyen2, chienchunhung, laikhtewari, nv-xtf and zhaoyangwang-nvidia July 20, 2026 21:14

coderabbitai Bot reviewed Jul 20, 2026

View reviewed changes

Comment thread tests/unittest/_torch/multimodal/test_multimodal_mixin.py Outdated

2ez4bz force-pushed the dev-mm-encoder-cache branch from fa22df4 to e084f4d Compare July 21, 2026 03:51

coderabbitai Bot mentioned this pull request Jul 21, 2026

[None][feat] Enable MM encoder cache on Qwen3.x and Gemma4 VLMs #16662

Draft

1 task

cascade812 approved these changes Jul 22, 2026

View reviewed changes

Comment thread tests/unittest/_torch/multimodal/test_multimodal_mixin.py Outdated

2ez4bz force-pushed the dev-mm-encoder-cache branch from e084f4d to bd94ed0 Compare July 23, 2026 23:42

arysef approved these changes Jul 24, 2026

View reviewed changes

2ez4bz force-pushed the dev-mm-encoder-cache branch from bd94ed0 to 76324bd Compare July 25, 2026 04:51

		_MM_ENCODER_CACHE_LOG_NAME = "mm_encoder_cache"


		def _build_request_multimodal_input(

Uh oh!

Conversation

2ez4bz commented Jul 17, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Dev Engineer Review

QA Engineer Review

tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py

tests/unittest/_torch/multimodal/test_multimodal_mixin.py

tests/unittest/_torch/test_tensor_lru_cache.py

tests/unittest/llmapi/test_llm_args.py

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

2ez4bz Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

2ez4bz commented Jul 20, 2026

Uh oh!

coderabbitai Bot commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (2 warnings)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

2ez4bz commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

2ez4bz commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

2ez4bz commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

2ez4bz commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

2ez4bz commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

2ez4bz commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

Uh oh!

2ez4bz commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

tensorrt-cicd commented Jul 24, 2026

Uh oh!

2ez4bz commented Jul 25, 2026

Uh oh!

tensorrt-cicd commented Jul 25, 2026

2ez4bz commented Jul 17, 2026 •

edited by coderabbitai Bot

Loading

`tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py`

`tests/unittest/_torch/multimodal/test_multimodal_mixin.py`

`tests/unittest/_torch/test_tensor_lru_cache.py`

`tests/unittest/llmapi/test_llm_args.py`

coderabbitai Bot commented Jul 20, 2026 •

edited

Loading