[None][feat] Enable MM encoder cache on Qwen3.x and Gemma4 VLMs - #16662
Conversation
7da572a to
f6fabcc
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe PR centralizes multimodal forwarding hooks, migrates Gemma 4 and Qwen3-VL models to shared contracts, enables encoder caching for additional models, updates token-ID compatibility, and adds CUDA coverage for cross-iteration prefetch behavior. ChangesMultimodal model integration
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/modeling/test_gemma4_multimodal.py (1)
765-776: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd production-path cache coverage for Qwen and Gemma.
Coverage is insufficient: this harness calls
_get_or_encode_multimodal_embeddings()directly, so it does not exercise Gemma4forward()or Qwen’s new_get_qwen_multimodal_embeddings()routing. Add cache-hit/miss tests totests/unittest/_torch/modeling/test_gemma4_multimodal.pyand the Qwen3-VL model test module, asserting two identical raw requests invoke the encoder once; parameterize the Qwen test across dense, MoE, and Qwen3.5 wrappers. Run the targeted tests underpytest tests/unittest/.As per path instructions, “Keep feedback actionable: suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR.”
🤖 Prompt for 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. In `@tests/unittest/_torch/modeling/test_gemma4_multimodal.py` around lines 765 - 776, The cache test only exercises the helper directly and does not cover production routing. Add cache miss/hit tests through Gemma4ForConditionalGeneration.forward and Qwen’s _get_qwen_multimodal_embeddings, asserting identical raw requests invoke the encoder once; parameterize Qwen coverage across dense, MoE, and Qwen3.5 wrappers, and run the targeted tests under pytest tests/unittest/.Source: Path instructions
🤖 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 `@docs/source/models/supported-models.md`:
- Around line 129-138: Add Gemma4ForConditionalGeneration to the
MultimodalModelMixin optimization matrix, mark Multimodal Embeddings Cache as
Yes, and set Multimodal Encoder Side Stream according to its verified support
status.
---
Nitpick comments:
In `@tests/unittest/_torch/modeling/test_gemma4_multimodal.py`:
- Around line 765-776: The cache test only exercises the helper directly and
does not cover production routing. Add cache miss/hit tests through
Gemma4ForConditionalGeneration.forward and Qwen’s
_get_qwen_multimodal_embeddings, asserting identical raw requests invoke the
encoder once; parameterize Qwen coverage across dense, MoE, and Qwen3.5
wrappers, and run the targeted tests under pytest tests/unittest/.
🪄 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: 2af008bf-08c0-455b-8e97-c6273f5a8e3d
📒 Files selected for processing (6)
docs/source/models/supported-models.mdtensorrt_llm/_torch/models/modeling_gemma4mm.pytensorrt_llm/_torch/models/modeling_qwen3_5.pytensorrt_llm/_torch/models/modeling_qwen3vl.pytensorrt_llm/_torch/models/modeling_qwen3vl_moe.pytests/unittest/_torch/modeling/test_gemma4_multimodal.py
389bb36 to
a426936
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/models/modeling_mistral.py (2)
619-634: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
load_tokenizerhardcodestrust_remote_code=True, overriding the caller's setting.
MistralCommonInputProcessor.__init__acceptstrust_remote_codeand forwards it tosuper().__init__, but this fallback ignores it and always enables remote code execution (also hardcodesuse_fast=True). Thread the flag through instead.Separately,
MistralTokenizer.from_pretrainedcan fail withOSError/FileNotFoundErrorfor a checkpoint that simply has no mistral-common artifacts; catching onlyValueErrormeans those cases abort instead of falling back.🔒️ Proposed fix
`@staticmethod` def load_tokenizer(model_path: str, config: PretrainedConfig, - tokenizer: AutoTokenizer | None = None): + tokenizer: AutoTokenizer | None = None, + trust_remote_code: bool = False, + use_fast: bool = True): if getattr(config, "input_processor_type", None) == "mistral_large_3": try: return MistralTokenizer.from_pretrained(model_path) - except ValueError: + except (ValueError, OSError): logger.info( f"Could not load mistral-common tokenizer from {model_path}, falling back to HuggingFace" ) tokenizer = tokenizer if tokenizer is not None else AutoTokenizer.from_pretrained( - model_path, config=config, use_fast=True, trust_remote_code=True) + model_path, + config=config, + use_fast=use_fast, + trust_remote_code=trust_remote_code) return tokenizerAnd at the call site:
- tokenizer = self.load_tokenizer(model_path, - config=config, - tokenizer=tokenizer) + tokenizer = self.load_tokenizer(model_path, + config=config, + tokenizer=tokenizer, + trust_remote_code=trust_remote_code)🤖 Prompt for 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. In `@tensorrt_llm/_torch/models/modeling_mistral.py` around lines 619 - 634, Update load_tokenizer to accept the caller’s trust_remote_code and use_fast settings instead of hardcoding them when calling AutoTokenizer.from_pretrained; update MistralCommonInputProcessor.__init__ to pass those flags into load_tokenizer. Expand the MistralTokenizer.from_pretrained exception handling to fall back for missing tokenizer artifacts by catching OSError, including FileNotFoundError, alongside ValueError.Source: Linters/SAST tools
64-67: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse a mapping-aware RoPE lookup here
config.rope_scaling/config.rope_parameterscan be dicts, sogetattr(..., "rope_type")returnsNoneand"yarn"falls through torope_gpt_neox. Use a dict-aware lookup likeget("rope_type", get("type"))before the fallback.🤖 Prompt for 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. In `@tensorrt_llm/_torch/models/modeling_mistral.py` around lines 64 - 67, Update the RoPE type lookup near rope_params_section so it supports mapping-based rope_scaling or rope_parameters values, retrieving rope_type and falling back to type before using attribute access for object configurations. Ensure dictionary configurations with a "yarn" type enter the existing yarn branch instead of falling through to rope_gpt_neox.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/models/modeling_mistral.py (1)
428-438: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDrop the image-only
do_rescalekwarg on the text-only branch.
do_rescaleis an image-processor argument; forwarding it totext_processorfor a prompt with no images is meaningless and can raise for processors that validate kwargs.♻️ Proposed change
else: - processed = self.text_processor( - text=inputs["prompt"], - do_rescale=do_rescale, - ) + processed = self.text_processor(text=inputs["prompt"])🤖 Prompt for 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. In `@tensorrt_llm/_torch/models/modeling_mistral.py` around lines 428 - 438, Update the text-only branch in the processing flow to remove the do_rescale keyword argument from the self.text_processor call. Keep do_rescale passed to self.processor in the images branch, and preserve the existing prompt handling.tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
2998-3002: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value
getattr(..., None)won't catch the mixin'sNotImplementedError.
MultimodalModelMixin.multimodal_token_idsis a property whose default body raisesNotImplementedError, andgetattrwith a default only suppressesAttributeError. Every mixin subclass overrides it today, so this is latent rather than broken — but the first model that inherits without overriding will crash here instead of falling back tomm_token_ids. Consider returningNonefrom the mixin's default property (it already documentsNoneas the out-of-vocabulary sentinel behavior) rather than raising.🤖 Prompt for 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. In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 2998 - 3002, Update the default multimodal_token_ids property in MultimodalModelMixin to return None instead of raising NotImplementedError, preserving the documented out-of-vocabulary sentinel behavior so model_engine.py can fall back to mm_token_ids.
🤖 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 `@tensorrt_llm/_torch/models/modeling_gemma4mm.py`:
- Around line 745-752: Guard the image token assignment in the mm_token_type_ids
initialization block so self.image_token_ids is checked for None before
accessing [0]. Preserve the existing image classification when image_token_ids
is present and the current independent video and audio guards.
- Around line 595-699: The multimodal embedding assembly in the mixed-modality
path currently concatenates modality buckets and loses prompt order. Update the
logic around multimodal_params and multimodal_embeddings to preserve each
parameter’s original content order when combining image, audio, and video
embeddings; alternatively reject mixed-modality parameters explicitly until
ordering metadata is available.
In `@tensorrt_llm/_torch/models/modeling_mistral.py`:
- Around line 817-831: The model-specific
get_language_model_extra_forward_kwargs overrides must preserve all engine
kwargs required by their inner language models. In
tensorrt_llm/_torch/models/modeling_mistral.py lines 817-831, update
get_language_model_extra_forward_kwargs to include lora_params in the returned
kwargs, or verify and document that MistralForCausalLM.forward ignores it. In
tensorrt_llm/_torch/models/modeling_gemma4mm.py lines 734-744, accept and return
spec_metadata and resource_manager alongside lora_params, matching the Mistral
and Qwen3-VL overrides.
In `@tensorrt_llm/_torch/models/modeling_qwen_image_bench.py`:
- Around line 100-109: Update the model_config cloning logic around replace() to
preserve the existing extra_attrs metadata when constructing the new config.
Carry model_config.extra_attrs into the replacement so runtime entries such as
nvfp4_gemm_allowed_backends and allreduce_* remain available, while keeping the
multimodal_config encoder_cache_max_bytes override unchanged.
In `@tensorrt_llm/_torch/models/modeling_qwen3vl.py`:
- Around line 1478-1493: Update the validation in the Qwen multimodal request
flow around `_get_requests_with_mm_data` so batches carrying only attached
`multimodal_embedding` handles are accepted without requiring
`support_mm_disagg`; retain the `NotImplementedError` for genuinely
disaggregated inputs without embeddings. Replace the hardcoded
environment-variable name in the error message with the imported
`_MULTIMODAL_ENV_NAME`.
---
Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_mistral.py`:
- Around line 619-634: Update load_tokenizer to accept the caller’s
trust_remote_code and use_fast settings instead of hardcoding them when calling
AutoTokenizer.from_pretrained; update MistralCommonInputProcessor.__init__ to
pass those flags into load_tokenizer. Expand the
MistralTokenizer.from_pretrained exception handling to fall back for missing
tokenizer artifacts by catching OSError, including FileNotFoundError, alongside
ValueError.
- Around line 64-67: Update the RoPE type lookup near rope_params_section so it
supports mapping-based rope_scaling or rope_parameters values, retrieving
rope_type and falling back to type before using attribute access for object
configurations. Ensure dictionary configurations with a "yarn" type enter the
existing yarn branch instead of falling through to rope_gpt_neox.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_mistral.py`:
- Around line 428-438: Update the text-only branch in the processing flow to
remove the do_rescale keyword argument from the self.text_processor call. Keep
do_rescale passed to self.processor in the images branch, and preserve the
existing prompt handling.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2998-3002: Update the default multimodal_token_ids property in
MultimodalModelMixin to return None instead of raising NotImplementedError,
preserving the documented out-of-vocabulary sentinel behavior so model_engine.py
can fall back to mm_token_ids.
🪄 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: 56028f2a-821b-433d-ada8-05c37c9e8b56
📒 Files selected for processing (10)
docs/source/models/supported-models.mdtensorrt_llm/_torch/models/modeling_gemma4_unified.pytensorrt_llm/_torch/models/modeling_gemma4mm.pytensorrt_llm/_torch/models/modeling_mistral.pytensorrt_llm/_torch/models/modeling_multimodal_mixin.pytensorrt_llm/_torch/models/modeling_qwen3_5.pytensorrt_llm/_torch/models/modeling_qwen3vl.pytensorrt_llm/_torch/models/modeling_qwen3vl_moe.pytensorrt_llm/_torch/models/modeling_qwen_image_bench.pytensorrt_llm/_torch/pyexecutor/model_engine.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py
- tensorrt_llm/_torch/models/modeling_qwen3_5.py
- docs/source/models/supported-models.md
a426936 to
6efe549
Compare
|
PR_Github #62408 [ run ] triggered by Bot. Commit: |
|
PR_Github #62408 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62438 [ run ] triggered by Bot. Commit: |
|
PR_Github #62438 [ run ] completed with state
|
* Why? Repeated or chunked multimodal requests re-ran encoder work on supported Qwen3.x and Gemma4 models, increasing latency and GPU utilization. Duplicated model-specific multimodal plumbing also made this behavior harder to maintain consistently. * What? Refactor Gemma4 and Qwen3.x VLMs to reuse the multimodal base mixin and shared encoder flow, removing duplicated multimodal boilerplate. Enable cached encoder embeddings for supported models, preserve model-specific behavior, exclude Qwen Image Bench from caching, and document support. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
8050707 to
c104620
Compare
|
/bot run |
|
PR_Github #62568 [ run ] triggered by Bot. Commit: |
|
PR_Github #62568 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62626 [ run ] triggered by Bot. Commit: |
|
PR_Github #62626 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62685 [ run ] triggered by Bot. Commit: |
|
PR_Github #62685 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62723 [ run ] triggered by Bot. Commit: |
|
PR_Github #62723 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62798 [ run ] triggered by Bot. Commit: |
|
PR_Github #62798 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62872 [ run ] triggered by Bot. Commit: |
|
PR_Github #62872 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62945 [ run ] triggered by Bot. Commit: |
|
PR_Github #62945 [ run ] completed with state |
6 participants
Description
Repeated or chunked multimodal requests re-ran encoder work on supported
Qwen3.x and Gemma4 models, increasing latency and GPU utilization.
Duplicated model-specific multimodal plumbing also made this behavior
harder to maintain consistently.
Refactor Gemma4 and Qwen3.x VLMs to reuse the multimodal base mixin and
shared encoder flow, removing duplicated multimodal boilerplate.
Enable cached encoder embeddings for supported models, preserve
model-specific behavior, exclude Qwen Image Bench from caching, and
document support.
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.Dev Engineer Review
MultimodalModelMixin/model base flow:Gemma4MultimodalModelBaseintroduced andGemma4ForConditionalGenerationupdated to inherit it (centralizing multimodal encoding + LM integration hooks, and renaming token-id exposure tomultimodal_token_ids).Gemma4UnifiedForConditionalGenerationnow inheritsGemma4MultimodalModelBase, removing the prior custom multimodal forward/fusion path and switching to encoder-free vision/audio feature extraction + shared multimodal fusion.Mistral3VLMhook renamed/reshaped:get_language_model_forward_kwargs→get_language_model_extra_forward_kwargsto align with the shared extra-forward-kwargs pattern.Qwen3VLModelBaseforward rewritten to use shared multimodal lifecycle methods (select_multimodal_params,after_active_multimodal_embeddings,_fuse_multimodal_embeddings,get_language_model_extra_forward_kwargs) and to enforce the new multimodal token-id property name (multimodal_token_ids).MultimodalModelMixin.multimodal_token_idsnow defaults toNone(instead of throwing), enabling the sentinel/OOV path infuse_input_embeds.MultimodalModelMixin.select_multimodal_paramsnow filters params only whenparam.has_content()and the chunk has multimodal tokens (viamultimodal_runtime/num_mm_tokens_in_chunkrules).mm_item_order) intoMultimodalParams.PyTorchModelEngine._prepare_multimodal_indicesnow derivesmm_token_idsvia the mixin-stylemultimodal_token_idsfirst, falling back to legacymm_token_ids.supports_encoder_cache = Trueadded for_Qwen3_5VLModelandQwen3MoeVLModel.encoder_cache_max_bytes = 0to exclude it from caching.config.image_token_id(raisesValueErrorif missing).mm_encoderfor encoder-backed multimodal embedding; validates exactly one packed embedding tensor is returned.encode_multimodal_inputsand shared multimodal selection helpers.docs/source/models/supported-models.mdupdated to enumerate all architectures implementingMultimodalModelMixinthat support multimodal encoder side-stream + embeddings cache (replacing the previous single-model statement).QA Engineer Review
Test changes (files under
tests/touched)tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py_CacheStubModel(cache-capable stub),_make_cacheable_request(...).test_cross_iter_prefetch_populates_and_reuses_persistent_cachetest_cross_iter_prefetch_mixed_cache_hit_and_miss_encodes_only_misstest_cross_iter_prefetch_cache_model_preserves_uncacheable_fallbackstest_cross_iter_prefetch_does_not_synchronize_main_streamtest_cross_iter_prefetch_does_not_rewrite_request_local_embeddingtest_prefetched_embedding_records_main_consumer_stream@requires_cudafortest_cross_iter_prefetch_materializes_on_side_stream.tests/unittest/_torch/executor/test_pytorch_model_engine.py_prepare_multimodal_indicesbehavior for both:multimodal_token_idsmm_token_idstests/unittest/_torch/modeling/test_gemma4_multimodal.pyencode_multimodal_inputs(instead of_forward_multimodal_encoder)._Gemma4EncoderCacheHarnessandtest_encoder_cache_reuses_image_embedding_across_requests.test_single_request_with_multiple_modalities_raises(behavior update from prior “allowed” test).tests/unittest/_torch/modeling/test_modeling_gemma4_unified.pytest_wrapper_rejects_missing_image_token_id.Coverage / registration
tests/integration/test_lists/(test-db/qa entries) could not be confirmed from the provided change summary.