[https://nvbugs/6550127][fix] Support Gemma4 multimodal cache partial hits

[2ez4bz]

Dev Engineer Review

• Gemma4 image and audio partial cache hits preserve item-major tensor and per-item metadata alignment.
• Unsupported layouts and modalities use the generic implementation.
• Partial video cache hits fall back to full misses and emit a one-time warning because frame-level slicing is not supported.
• This behavior disables persistent video cache reuse, including full-hit reuse. Update the related scope and TRTLLM-14981 documentation if intentional.
• The waiver removal enables the Gemma4 NVFP4 multimodal test.
• No configuration files changed.
• Verdict: needs follow-up.

QA Engineer Review

• Added and updated Gemma4MultimodalModelBase cache tests.
• Added coverage for partial image and audio cache hits, multi-item metadata slicing, and video cache behavior.
• Video tests cover full-hit reuse and re-encoding for partial-hit requests.
• TestGemma4_26B_A4B::test_nvfp4 is covered by the removed waiver entry.
• The changed test functions are covered by the test code, but no new test-db/ or qa/ entries are reported.
• Verdict: needs follow-up.

Description

• Why?

The generic multimodal cache path cannot slice Gemma4 image and audio layouts, causing partial encoder-cache hits to fail during input construction.

• What?

Override partial-hit input construction for Gemma4 images and audio, keeping their per-item metadata aligned. Bypass persistent caching for videos with a warning until frame-level slicing is supported, and remove the obsolete accuracy waiver.

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]

/bot run

[coderabbitai]

Walkthrough

Gemma4 now applies modality-specific encoder-cache handling. Image and audio partial hits use sliced inputs with aligned metadata. Partial video hits trigger full re-encoding with a one-time warning. Tests cover these paths, and the related waiver is removed.

Changes

Gemma4 multimodal cache handling

Layer / File(s)	Summary
Cache policy and multimodal input slicing tensorrt_llm/_torch/models/modeling_gemma4mm.py	Gemma4 partitions encoder-cache requests. Image and audio partial misses use item-major tensors and aligned fields. Partial video hits become full re-encodes with a one-time warning.
Cache behavior validation tests/unittest/_torch/modeling/test_gemma4_multimodal.py, tests/integration/test_lists/waives.txt	Tests cover multi-item image, audio, and video cache behavior, metadata alignment, and audio inputs. The Gemma4 waiver entry is removed.

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

Possibly related PRs

• NVIDIA/TensorRT-LLM#16051: Both changes modify Gemma4 encoder-cache partitioning and partial item-level input slicing.
• NVIDIA/TensorRT-LLM#17209: This change builds on dispatch logic for plain dim-0 stacked inputs in build_multimodal_encoder_input.

Suggested labels: api-compatible

Suggested reviewers: brnguyen2, qijune

Sequence Diagram(s)

sequenceDiagram
  participant Request
  participant Gemma4MultimodalModelBase
  participant TensorLRUCache
  participant MultimodalEncoder
  Request->>Gemma4MultimodalModelBase: submit multimodal items
  Gemma4MultimodalModelBase->>TensorLRUCache: partition cache keys
  TensorLRUCache-->>Gemma4MultimodalModelBase: return hits and misses
  Gemma4MultimodalModelBase->>MultimodalEncoder: encode sliced image or audio inputs
  MultimodalEncoder-->>Gemma4MultimodalModelBase: return embeddings
  Gemma4MultimodalModelBase-->>Request: return ordered embeddings

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the Gemma4 multimodal cache partial-hit fix and uses the required ticket and type tags.
Description check	✅ Passed	The description explains the issue and solution and includes the required sections, but the Test Coverage section is not filled in.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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]

🧹 Nitpick comments (1)

tensorrt_llm/_torch/models/modeling_gemma4mm.py (1)

599-612: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document the override contract.

Add Google-style Args and Returns sections for param, item_indices, and the residual MultimodalParams. This method is a public override point for the encoder-cache flow.

As per coding guidelines, “Prefer docstrings for external interfaces, use Google-style docstrings, document public function arguments.”

🤖 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_gemma4mm.py` around lines 599 - 612,
Update the docstring for build_multimodal_encoder_input to add Google-style Args
entries describing param and item_indices, plus a Returns entry describing the
residual MultimodalParams produced for the selected items. Keep the existing
behavior description and document this public encoder-cache override contract
without changing implementation logic.

Source: Coding guidelines

🤖 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.

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_gemma4mm.py`:
- Around line 599-612: Update the docstring for build_multimodal_encoder_input
to add Google-style Args entries describing param and item_indices, plus a
Returns entry describing the residual MultimodalParams produced for the selected
items. Keep the existing behavior description and document this public
encoder-cache override contract without changing implementation logic.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 060df0c7-4077-4117-9301-cacdb3d0864a

📥 Commits

Reviewing files that changed from the base of the PR and between dbe6a41 and 889fab9.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/models/modeling_gemma4mm.py
• tests/integration/test_lists/waives.txt
• tests/unittest/_torch/modeling/test_gemma4_multimodal.py

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

[tensorrt-cicd]

PR_Github #63584 [ run ] triggered by Bot. Commit: 889fab9 Link to invocation

[tensorrt-cicd]

PR_Github #63584 [ run ] completed with state FAILURE. Commit: 889fab9
/LLM/main/L0_MergeRequest_PR pipeline #51548 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

[BowenFu]

Two things before this lands.

The video decision is wider than the comment says. The comment reads "retain the pre-cache behavior", but on main Gemma4MultimodalModelBase.supports_encoder_cache = True with no _encoder_cache_modality override, so Gemma4 video params go through the persistent cache today and all-hit reuse works. What is broken is only the partial hit: the mixin's build_multimodal_encoder_input wants the Qwen2-VL packed layout (video_grid_thw + pixel_values_videos), Gemma4 supplies plain pixel_values under "video", so it falls to the else: and raises NotImplementedError. Returning None from _encoder_cache_modality fixes that crash but also gives up the working all-hit path, i.e. every repeated video is re-encoded. Can you decline only the partial-hit partition and keep video cacheable for full hits? If the wider scope is deliberate, please say so in the comment and on TRTLLM-14981 — as written it reads like an edge case.

Title is failing CI. Check PR Title Format is red because the title has no [type] tag (and a double space): [https://nvbugs/6550127][fix] Support Gemma4 multimodal cache partial hits.

The rest looks right to me. The residual MultimodalParams dropping multimodal_runtime matches the mixin's own construction at modeling_multimodal_mixin.py:705-710, so that's the existing contract, not a new hole; the item_count/dim() guard falling through to super() is a good call. L0 on 889fab9 is red (63584) — worth a look at whether any of it is yours.

[xinhe-nv]

/bot run --only-qa-verify test accuracy/test_llm_api_pytorch_multimodal.py::TestGemma4_26B_A4B::test_nvfp4

[tensorrt-cicd]

PR_Github #63615 [ run ] triggered by Bot. Commit: 889fab9 Link to invocation

[tensorrt-cicd]

PR_Github #63615 [ run ] completed with state SUCCESS. Commit: 889fab9
LLM_FUNCTION_AUTO_V2C #348 completed with status: 'SUCCESS'
QA verify test: accuracy/test_llm_api_pytorch_multimodal.py::TestGemma4_26B_A4B::test_nvfp4 (NVBug 6550127, branch dev-nvbug-6550127, fork 2ez4bz, dry_run_close=true)

Link to invocation

[2ez4bz]

/bot run

[2ez4bz]

@BowenFu addressed, please take another look.

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/models/modeling_gemma4mm.py (1)

828-838: 📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift

Split the unrelated draft-model changes.

The declared encoder-cache objective does not establish a dependency on draft-model delegation or speculative/resource-manager forwarding. Move these changes to a dedicated PR unless a direct dependency on encoder-cache behavior is documented.

• tensorrt_llm/_torch/models/modeling_gemma4mm.py#L828-L838: Move the draft configuration, draft model, and draft-weight delegation.
• tensorrt_llm/_torch/models/modeling_gemma4mm.py#L850-L881: Move the speculative-decoding, resource-manager, and original-input forwarding.

As per coding guidelines: “Keep each pull request focused on one concern and avoid unrelated scope; split unrelated changes into separate PRs.”

🤖 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_gemma4mm.py` around lines 828 - 838,
Remove the unrelated draft-model delegation from modeling_gemma4mm.py lines
828-838, including draft_config, draft_model, and load_draft_weights. Also
remove the speculative-decoding, resource-manager, and original-input forwarding
changes from modeling_gemma4mm.py lines 850-881; retain only
encoder-cache-related changes in this PR and move these features to a separate
PR unless a direct dependency is documented.

Source: Coding guidelines

🤖 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.

Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_gemma4mm.py`:
- Around line 828-838: Remove the unrelated draft-model delegation from
modeling_gemma4mm.py lines 828-838, including draft_config, draft_model, and
load_draft_weights. Also remove the speculative-decoding, resource-manager, and
original-input forwarding changes from modeling_gemma4mm.py lines 850-881;
retain only encoder-cache-related changes in this PR and move these features to
a separate PR unless a direct dependency is documented.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 009402b2-e5e6-440f-b569-f81d5abc7de6

📥 Commits

Reviewing files that changed from the base of the PR and between ceb71dc and 9e3a01f.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/models/modeling_gemma4mm.py
• tests/integration/test_lists/waives.txt
• tests/unittest/_torch/modeling/test_gemma4_multimodal.py

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

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

• tests/unittest/_torch/modeling/test_gemma4_multimodal.py

[tensorrt-cicd]

PR_Github #63802 [ run ] triggered by Bot. Commit: 9e3a01f Link to invocation

[tensorrt-cicd]

PR_Github #63802 [ run ] completed with state SUCCESS. Commit: 9e3a01f
/LLM/main/L0_MergeRequest_PR pipeline #51748 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]

/bot run