[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.
• Video inputs bypass persistent encoder caching and emit a one-time warning until frame-level slicing is supported.
• Unsupported layouts and modalities use the generic implementation.
• The waiver removal enables the Gemma4 NVFP4 multimodal test.
• No configuration or public API changes were introduced.
• Verdict: sufficient.

QA Engineer Review

• Added and updated Gemma4 multimodal cache tests for image, audio, and video inputs.
• Added coverage for multi-item partial cache hits and metadata slicing.
• Video tests verify re-encoding and exclusion from persistent caching.
• The TestGemma4_26B_A4B::test_nvfp4 coverage is represented by the removed waiver entry.
• Verdict: sufficient.

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 multimodal encoder caching. Image and audio cache misses use sliced inputs, while video inputs bypass persistent caching. Tests cover multi-item cache behavior, 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	Image and audio inputs retain cache support. Video inputs are excluded from persistent caching. Partial cache misses use item-major tensors and aligned metadata.
Cache behavior validation tests/unittest/_torch/modeling/test_gemma4_multimodal.py, tests/integration/test_lists/waives.txt	The test harness supports audio inputs and multi-item image, audio, and video cases. Tests verify partial cache hits, metadata alignment, video re-encoding, and removal of the Gemma4 waiver.

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

Possibly related PRs

• NVIDIA/TensorRT-LLM#17209: Both changes address Gemma4 multimodal encoder-cache slicing and related waiver removal.

Suggested reviewers: bowenfu, stanleysun639

Sequence Diagram(s)

sequenceDiagram
  participant Request
  participant Gemma4
  participant EncoderCache
  participant MultimodalEncoder
  Request->>Gemma4: submit multimodal items
  Gemma4->>EncoderCache: check item hashes
  EncoderCache-->>Gemma4: return cached and missing items
  Gemma4->>MultimodalEncoder: encode sliced image or audio inputs
  MultimodalEncoder-->>Gemma4: return embeddings
  Gemma4->>EncoderCache: store image and audio embeddings
  Gemma4-->>Request: return ordered multimodal embeddings

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description explains the motivation and solution, but the required Test Coverage section contains no test names or results.	List the relevant tests, including the new Gemma4 multimodal cache tests and QA verification results.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title uses the required NVBugs and fix tags and clearly describes Gemma4 multimodal cache partial-hit support.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 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