Skip to content

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

Open
2ez4bz wants to merge 2 commits into
NVIDIA:mainfrom
2ez4bz:dev-nvbug-6550127
Open

[https://nvbugs/6550127][fix] Support Gemma4 multimodal cache partial hits#17231
2ez4bz wants to merge 2 commits into
NVIDIA:mainfrom
2ez4bz:dev-nvbug-6550127

Conversation

@2ez4bz

@2ez4bz 2ez4bz commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Dev Engineer Review

  • Image and audio partial cache hits preserve item-major tensors 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.
  • Full-hit video cache reuse is also disabled. Update the related scope and TRTLLM-14981 documentation if this behavior is intentional.
  • The obsolete Gemma4 NVFP4 waiver entry was removed.
  • No configuration files changed.
  • Verdict: needs follow-up.

QA Engineer Review

  • Updated Gemma4MultimodalModelBase cache tests.
  • Added coverage for partial image and audio cache hits, metadata slicing, unsupported item-axis layouts, and video cache behavior.
  • Video tests cover full-hit reuse and partial-hit re-encoding.
  • TestGemma4_26B_A4B::test_nvfp4 is covered by the removed waiver entry.
  • No new test-db/ or qa/ entries are reported for the changed test functions.
  • 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

2ez4bz commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@coderabbitai

coderabbitai Bot commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

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: ce8ac2a4-33ca-4917-8c06-f1638f3fcdf1

📥 Commits

Reviewing files that changed from the base of the PR and between 7608520 and d3b295b.

📒 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 (2)
  • tests/unittest/_torch/modeling/test_gemma4_multimodal.py
  • tensorrt_llm/_torch/models/modeling_gemma4mm.py

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
The cache harness supports audio inputs and keyed multi-item fixtures. Tests cover partial image and audio hits, metadata slicing, unsupported layouts, and video cache behavior.
Gemma4 test waiver cleanup
tests/integration/test_lists/waives.txt
The Gemma4 NVFP4 multimodal test waiver entry is removed.

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

Suggested reviewers: mlefeb01, emmaqiaoch, yechank-nvidia

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 follows the required ticket and type format and clearly summarizes support for Gemma4 multimodal cache partial hits.
Description check ✅ Passed The description clearly explains the issue and solution, but the Test Coverage section does not list the relevant tests.
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 coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

BowenFu commented Aug 4, 2026

Copy link
Copy Markdown

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 xinhe-nv changed the title [https://nvbugs/6550127] Support Gemma4 multimodal cache partial hits [https://nvbugs/6550127][fix] Support Gemma4 multimodal cache partial hits Aug 4, 2026
@xinhe-nv

xinhe-nv commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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 2ez4bz changed the title [https://nvbugs/6550127][fix] Support Gemma4 multimodal cache partial hits [https://nvbugs/6550127][fix] Support Gemma4 multimodal cache partial hits Aug 4, 2026
* 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.

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
@2ez4bz
2ez4bz force-pushed the dev-nvbug-6550127 branch from ceb71dc to 9e3a01f Compare August 4, 2026 16:59
@2ez4bz

2ez4bz commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@2ez4bz

2ez4bz commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator Author

@BowenFu addressed, please take another look.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

2ez4bz commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

2ez4bz commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

2ez4bz commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@BowenFu BowenFu left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The updated cache partition keeps full video hits reusable and re-encodes only partial hits. The new image, audio, and video regressions cover the requested behavior. Approved.

@brnguyen2 brnguyen2 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving — the comments below are optional touch-ups, not blockers.

Approach looks right, and switching _Gemma4EncoderCacheHarness to subclass Gemma4MultimodalModelBase is a real improvement — the tests now go through the actual encode_multimodal_inputs. Two things below on the fallback path; neither blocks.

On the unwaive: TestGemma4_26B_A4B::test_nvfp4 is in test-db/l0_b200.yml and encoder_cache_max_bytes defaults to 128 MiB, so the B200 stage does exercise this. Please link a green run of that stage before merging.

Unrelated but adjacent: the encoder_cache_max_bytes description in llm_args.py:569 still says "reuse is all-or-nothing for each request: every item in the request must hit the cache", which stopped being true when partial hits landed. Worth a one-line fix while you're in this area.

I could not run the unit tests locally (no torch in this checkout), so my read is static only.

Comment thread tensorrt_llm/_torch/models/modeling_gemma4mm.py
Comment thread tensorrt_llm/_torch/models/modeling_gemma4mm.py
Comment thread tests/unittest/_torch/modeling/test_gemma4_multimodal.py
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63886 [ run ] completed with state SUCCESS. Commit: 9e3a01f
/LLM/main/L0_MergeRequest_PR pipeline #51826 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
@2ez4bz
2ez4bz force-pushed the dev-nvbug-6550127 branch from 9e3a01f to d3b295b Compare August 5, 2026 06:35
@coderabbitai

coderabbitai Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

@2ez4bz

2ez4bz commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63966 [ run ] triggered by Bot. Commit: d3b295b Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants