Draft: [None][feat] Multimodal encoder cache: per-item partial hits#16817
Draft
aswinvisva wants to merge 2 commits into
Draft
Draft: [None][feat] Multimodal encoder cache: per-item partial hits#16817aswinvisva wants to merge 2 commits into
aswinvisva wants to merge 2 commits into
Conversation
* 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>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
WalkthroughAdds granular multimodal encoder-cache reuse to ChangesGranular encoder cache
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant MultimodalModelMixin
participant EncoderCache
participant Mistral3VLM
MultimodalModelMixin->>EncoderCache: Partition per-item cache keys
EncoderCache-->>MultimodalModelMixin: Return hits and miss indices
MultimodalModelMixin->>Mistral3VLM: Slice multimodal data for misses
Mistral3VLM-->>MultimodalModelMixin: Return miss-only parameters
MultimodalModelMixin->>EncoderCache: Store newly encoded miss items
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
1 task
… 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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
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.Add an opt-in per-item partial-hit path to
MultimodalModelMixin:EncoderCachePartitionrecords the per-item cache lookup result (hits / miss indices / keys)._partition_encoder_cachereplaces_attach_encoder_cache_hitand records every item lookup rather than bailing on first miss._encode_with_partial_cacheencodes only the miss items via a residualMultimodalParams, then interleaves cached and freshly encoded rows in original per-item order so downstreamget_multimodal_embeddingstreats the param as fully cached._apply_metadata_slicere-slicesmultimodal_embedding_lengthsandmultimodal_hasheson the residual so every model gets the parallel per-item metadata slice identically.supports_granular_encoder_cache = Trueand implementslice_multimodal_data_items, which slices raw modality tensors (pixel_values,image_grid_thw, ...) for the given item indices. Default remainsFalse; 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_valueson dim 0 andimage_sizeslist in parallel.Perf
Real trtllm-serve run against
mistralai/Mistral-Small-3.1-24B-Instruct-2503on 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: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.
Dev Engineer Review
MultimodalModelMixin.pixel_valuesandimage_sizes.QA Engineer Review
No test changes.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.