Skip to content

[None][feat] Gemma4 MM: native vision + audio towers - #14300

Merged
Hudayday merged 10 commits into
NVIDIA:mainfrom
Hudayday:feat/gemma4-mm-tower-native
May 21, 2026
Merged

[None][feat] Gemma4 MM: native vision + audio towers#14300
Hudayday merged 10 commits into
NVIDIA:mainfrom
Hudayday:feat/gemma4-mm-tower-native

Conversation

@Hudayday

@Hudayday Hudayday commented May 19, 2026

Copy link
Copy Markdown
Collaborator

Summary

Per multimedia team's PR feedback on the initial Gemma4 MM port, this PR replaces the temporary AutoModel.from_config HF fallback with native TensorRT-LLM implementations of the Gemma4 vision and audio
(Conformer) towers, so that the multimodal path goes through TRT-LLM's own attention/MLP modules rather than re-entering HuggingFace at runtime.

  • Native Gemma4 vision tower (modeling_gemma4_vision.py) — replaces AutoModel.from_config; vision attention workspace now bound to the vision scale rather than inheriting the LM-side bound (avoids the
    5589 GiB resize previously observed on 31B MMMU).
  • Native Gemma4 audio Conformer tower (modeling_gemma4_audio.py) — replaces AutoModel.from_config; chunked local attention + rel-pos bias ported directly from the HF reference.
  • HF Gemma4RMSNorm convention fix (audio + vision) — Gemma4's MM-tower RMSNorm is the plain w * x / rms, not the Gemma3 LLM-side (1 + w) * x / rms. Wiring the correct norm convention is what brings
    BLEU back into parity (E4B CoVoST2 1.90 → ~41.9 in earlier validation runs).
  • modeling_gemma4mm.py updated to import the native towers instead of the HF fallback.
  • Test suite tests/unittest/_torch/multimodal/test_gemma4_multimodal.py registered in l0_b200.yml, covering Tier 0/1/2 vision-tower parity plus the new RMSNorm regression tests.
  • Debug probes used during the port (DEBUG_AUDIO_TOWER env-gated stats) have been scrubbed before this commit; only Signed-off-by trailers in the history.

Alignment with HF tower baseline

The fix-mode (native towers) results have been aligned with the HF-fallback baseline. All completed fix-mode cells so far are within ±2pp of the HF baseline (well inside MMMU/CoVoST seed noise):

variant task HF baseline native fix Δ
E2B mmmu 42.67 42.11 −0.56
E2B covost 37.70 BLEU 37.78 BLEU +0.08
E4B covost 41.91 BLEU 41.91 BLEU exact
26B mmmu 58.33 58.56 +0.23

Summary by CodeRabbit

Release Notes

  • New Features

    • Added native audio and vision tower implementations for Gemma4 multimodal model, replacing previous placeholder approach with optimized TensorRT-LLM components.
  • Tests

    • Expanded test coverage for multimodal audio/vision components, including forward pass validation, weight loading parity, and end-to-end generation testing.

Review Change Stack

Hudayday added 6 commits May 19, 2026 01:39
- modeling_gemma4_vision.py: Gemma4VisionModel inherits Attention base class,
  participates in TRT-LLM backend dispatch (FLASHINFER for LM, TRTLLM for vision).
- modeling_gemma4_audio.py: stub for Phase 2 audio tower migration.
- modeling_gemma4mm.py: construct native TRT-LLM vision tower instead of HF
  AutoModel.from_config().

Pre-nvidia/main merge snapshot.

Signed-off-by: Hudayday <tianruih@nvidia.com>
modeling_gemma4_vision.py: drop ``[:]`` slicing on clamp scalar buffers in
``load_weights`` and ``_remap_clamp_buffers`` (use ``.clone()`` instead).
HF registers clamp buffers as 0-dim ``torch.tensor(-float("inf"))``; ``[:]``
on a 0-dim tensor raises and breaks ``use_clipped_linears=True`` loads.
Also switch ``seq_lens`` pinning to ``maybe_pin_memory`` per repo policy.

modeling_gemma4mm.py: add no-op ``_check_and_adjust_experts_implementation``
on ``Gemma4ForConditionalGeneration``. HF transformers 5.5.3
``PreTrainedModel.__init__`` dispatches this on VL wrappers and raises when
the wrapper has no direct MoE layers (mirrors ``Qwen3VLModelBase`` override).
TRT-LLM manages expert implementations independently.

test_gemma4_multimodal.py: full Tier 0/1/2 rewrite against the post-Option B
vision tower.
 - Tier 0 ``TestGemma4VisionTower``: build / forward random-weights (fp32 —
   bf16 + head_dim=32 + trtllm-gen unfused MHA fallback at sm_100 is NaN-
   prone) / load-weights remap / clamp buffer parity.
 - Tier 1 ``TestGemma4VisionPipeline``: HF vs TRT vision tower parity.
 - Tier 1 ``TestGemma4ForConditionalGeneration``: vision_tower attach/skip.
 - Tier 2 ``TestGemma4MultimodalE2E`` via shared ``TestModelingMultimodal``
   skeleton; gated by ``_gemma4_e2e_model_available()`` which also requires
   ``{LLM_MODELS_ROOT}/multimodals/test_data/seashore.png`` (skeleton image
   fixture used by ``modality="image"``).

B200 (sm_100, c163c72 wheel + transformers 5.5.3) result: Tier 0/1 12/12
PASS, Tier 2 13 PASS + 1 skip (E2E gracefully gated on missing fixture).

Signed-off-by: Hudayday <tianruih@nvidia.com>
HF Gemma4RMSNorm computes `w * x / rms(x)` (plain w), NOT the
`(1 + w) * x / rms(x)` convention used by Gemma3 LLM. The previous
audio+vision ports re-implemented RMSNorm by analogy with Gemma3, which
silently corrupted activations because the Gemma4 checkpoint weights are
not zero-centered (mean=+2.7, absmax=9.2). Compounded across ~36 norm
applications in the conformer, this produced grammatical-but-wrong
German translations: E4B CoVoST BLEU=1.90 vs HF baseline 24.54.

Fix: replace inline norm classes with thin adapters inheriting
`transformers.models.gemma4.modeling_gemma4.Gemma4RMSNorm`, preserving
the existing kwarg + dtype + has_weights call signature. Vision uses
`with_scale=has_weights` for the weightless v_norm slot. Drop the 2D
reshape helper in vision _head_norm (HF norm is pure PyTorch, no
flashinfer dispatcher constraint on shape). Also wire mm wrapper to use
the native `Gemma4AudioModel` (with `load_weights`) instead of
`AutoModel.from_config`.

Verification (B200): E4B CoVoST BLEU = 24.69 (job 2206264), slightly
above HF baseline 24.54. 31B vision (job 2206265) past the original
`flashinfer_rmsnorm` dispatcher crash on bf16 1152-wide SigLIP tensors;
remaining 31B issue is a separate downstream attention resize bug
unrelated to RMSNorm.

Tests: add `TestGemma4MMTowerRMSNormConvention` (numerical regression
locking in plain-w convention, inheritance + dtype + weightless guards)
and `TestGemma4AudioTowerStructure` (module exports + mm wiring) in
tests/unittest/_torch/modeling/test_modeling_gemma4.py.

Signed-off-by: Hudayday <tianruih@nvidia.com>
Signed-off-by: Hudayday <tianruih@nvidia.com>
Vision tower's tower-local TRTLLM attn metadata was pulling
max_num_tokens from the LLM model_config (defaults to LLM's 16384) and
keeping max_num_requests=8192. The C++ attention op sizes its workspace
as ~max_num_requests * heads * max_num_tokens^2 * bytes, which scaled
to TB-range and crashed with "Tried to allocate 1397.40 GiB".

Vision is called per-image (modeling_gemma4mm._get_image_features loops
and feeds one image at a time), so batch and seq are vision-scale, not
LLM-scale. Bind to (max_num_requests=1, max_num_tokens=8192) — matches
the actual call shape and keeps the workspace at ~tens of MB.

Signed-off-by: Hudayday <tianruih@nvidia.com>
Removes the eight DEBUG_PROBE_BEGIN/END blocks (and the _DEBUG_AUDIO flag,
_audio_stats helper, unused os import) from Gemma4AudioModel.forward /
load_weights. The probes were gated by an env var and never fired under
normal runs; carrying them in source is noise.

Signed-off-by: Hudayday <tianruih@nvidia.com>
@Hudayday
Hudayday requested a review from a team as a code owner May 19, 2026 09:37
@Hudayday
Hudayday requested a review from symphonylyh May 19, 2026 09:37
@Hudayday
Hudayday enabled auto-merge (squash) May 19, 2026 09:38
@Hudayday

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@coderabbitai

coderabbitai Bot commented May 19, 2026

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

This PR replaces HuggingFace AutoModel fallbacks with native TRT-LLM implementations of Gemma4's audio and vision towers. The audio tower is a Conformer encoder with relative positional encoding and chunked local attention. The vision tower is a ViT-style encoder with 2D RoPE and spatial pooling. Both are integrated into the multimodal wrapper and validated through comprehensive unit, integration, and E2E tests with HF weight mapping and clamping semantics.

Changes

Gemma4 Audio/Vision Towers and Integration

Layer / File(s) Summary
Audio tower (Conformer encoder)
tensorrt_llm/_torch/models/modeling_gemma4_audio.py
Native Conformer encoder with Gemma4AudioClippableLinear, relative positional encoding, chunked local attention (with relative position bias and softcap), subsampling convolution (stride-4), Macaron feedforward, causal depthwise convolution, and Conformer layer wiring. Full Gemma4AudioModel orchestrates subsampling, layers, final projection, 4D→5D blocked mask conversion, and load_weights() with buffer copying.
Vision tower (ViT with RoPE)
tensorrt_llm/_torch/models/modeling_gemma4_vision.py
Native vision encoder with HF-compatible RMSNorm adapter, 2D RoPE computation, Gemma4VisionClippableLinear, MLP, TRT-LLM Attention subclass (Gemma4 clamping + RoPE), sandwich-norm encoder layer/encoder, patch embedder with learned 2-axis position embeddings, spatial pooler. Gemma4VisionModel orchestrates flattening/scattering, attention metadata prep, optional standardization, and load_weights() with HF clamp-buffer fusion and regex-based weight remapping.
Multimodal wrapper wiring
tensorrt_llm/_torch/models/modeling_gemma4mm.py
Module docstring clarifies native tower implementations. Adds _check_and_adjust_experts_implementation() bypass. Vision/audio tower initialization routes to Gemma4VisionModel/Gemma4AudioModel (replacing HF AutoModel.from_config()). Weight loading calls each tower's load_weights() method instead of load_state_dict().
Vision tower unit and pipeline tests
tests/unittest/_torch/modeling/test_gemma4_multimodal.py (ranges 127–254, 281–509)
Test infrastructure updates with TRTLLM attention backend config. Helpers _make_dummy_pixel_input() and _build_trt_vision_tower(). Vision unit tests: forward shape/NaN checks, HF→TRT QKV fusion validation, forward parity (unclipped). Clamp-specific tests: buffer remapping parity, finite clamp injection, parity under clipped-linears, ValueError on q/k/v clamp mismatch. Vision+embedder pipeline parity mirroring production shape/batching.
Multimodal wrapper and processor tests
tests/unittest/_torch/modeling/test_gemma4_multimodal.py (ranges 516–935)
Wrapper instantiation asserts vision tower is TRT-LLM Gemma4VisionModel class. Image-processing tests add pixel_values dimensionality assertion. Multi-image batch validation. Audio-processing tests assert 3D feature shape. Video-processing constructs frames directly, validates per-frame pixel_values dimensionality (32-frame batch).
E2E validation and regression guards
tests/unittest/_torch/modeling/test_gemma4_multimodal.py (ranges 948–1121), tests/unittest/_torch/modeling/test_modeling_gemma4.py, tests/integration/test_lists/test-db/l0_b200.yml
Tier-2 TestGemma4MultimodalE2E gates full LLM+vision generation parity on LLM_MODELS_ROOT and multimodal fixture availability. TestGemma4MMTowerRMSNormConvention validates audio/vision RMSNorm adapters (inherit from HF, plain w * x / rms(x) semantic, dtype casting, vision attention norm wiring via source checks). TestGemma4AudioTowerStructure confirms native audio tower exports and HF RMSNorm usage (no TRT-LLM flashinfer import). Test list adds test_gemma4_multimodal.py to l0_b200.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Suggested reviewers

  • moraxu
  • QiJune
  • PerkzZheng
  • byshiue
  • brb-nv
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 36.36% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately summarizes the main change: implementing native Gemma4 vision and audio towers to replace HuggingFace fallback.
Description check ✅ Passed The description addresses core PR details including motivation (replacing AutoModel fallback), key changes (native towers, RMSNorm fix), test coverage (test suite in l0_b200.yml), and alignment with HF baseline. However, it lacks explicit PR checklist completion markers.
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.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands and usage tips.

@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_gemma4_audio.py (1)

247-256: 💤 Low value

Unused tuple unpacking in _extract_block_context.

The unpacked variables batch_size, seq_len, num_heads, head_dim are never used. Replace with underscores to silence the static analysis warning.

♻️ Suggested fix
     def _extract_block_context(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Extract overlapping context windows of ``context_size`` per block."""
-        batch_size, seq_len, num_heads, head_dim = hidden_states.shape
+        _, _, _, _ = hidden_states.shape  # Shape validated by caller
         hidden_states = F.pad(
🤖 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_gemma4_audio.py` around lines 247 - 256,
In _extract_block_context replace the unused tuple unpacking "batch_size,
seq_len, num_heads, head_dim = hidden_states.shape" with underscore placeholders
(e.g. _, _, _, _ = hidden_states.shape) or simply omit the assignment to silence
the static-analysis warning; update the line in the method
_extract_block_context in modeling_gemma4_audio.py so only necessary values are
captured or ignored to avoid unused variable warnings while preserving the rest
of the tensor padding, unfold, movedim, and contiguous operations.
🤖 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_gemma4_audio.py`:
- Around line 247-256: In _extract_block_context replace the unused tuple
unpacking "batch_size, seq_len, num_heads, head_dim = hidden_states.shape" with
underscore placeholders (e.g. _, _, _, _ = hidden_states.shape) or simply omit
the assignment to silence the static-analysis warning; update the line in the
method _extract_block_context in modeling_gemma4_audio.py so only necessary
values are captured or ignored to avoid unused variable warnings while
preserving the rest of the tensor padding, unfold, movedim, and contiguous
operations.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2eaf5b81-46cf-4cec-a63b-30ab01df1413

📥 Commits

Reviewing files that changed from the base of the PR and between 989671b and 23a7bee.

📒 Files selected for processing (6)
  • tensorrt_llm/_torch/models/modeling_gemma4_audio.py
  • tensorrt_llm/_torch/models/modeling_gemma4_vision.py
  • tensorrt_llm/_torch/models/modeling_gemma4mm.py
  • tests/integration/test_lists/test-db/l0_b200.yml
  • tests/unittest/_torch/modeling/test_gemma4_multimodal.py
  • tests/unittest/_torch/modeling/test_modeling_gemma4.py

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49169 [ run ] triggered by Bot. Commit: 23a7bee Link to invocation

Two test fixes for failures in L0_Test build 2316:

1. TestGemma4MMTowerRMSNormConvention::test_vision_attention_uses_native_rmsnorm_adapter

   Failure: AssertionError on '_Gemma4VisionRMSNorm(hidden_size=' not
   found in source. The guard read the vision module source as a string
   and did a literal substring match. yapf wraps long ctor calls so the
   actual source has '_Gemma4VisionRMSNorm(\n            hidden_size=...'
   — the assertion passed in CI before the call sites were reformatted
   and broke once the formatter ran. Switch the search to a regex that
   tolerates whitespace + newline between '(' and 'hidden_size=', and
   keep a negative lookbehind so a future direct TRT-LLM RMSNorm
   construction still trips the guard.

2. TestGemma4VisionTower::test_forward_random_weights

   Failure: AssertionError 'Vision tower output contains NaN'. The test
   built the TRT-LLM vision tower with PyTorch-default init on fused
   qkv_proj and pushed activations through the TRTLLM attention backend.
   head_dim=32 + uninitialised fused weights saturate the attention
   kernel and produce NaN even in fp32 (the dtype switch from bf16 to
   fp32 in an earlier commit did not help). Production code always
   calls load_weights(hf_state_dict) before forward, so the appropriate
   sanity-check init is HF defaults. Seed the TRT-LLM tower via an HF
   Gemma4VisionModel state_dict (same path test_forward_parity_no_clamp
   uses) and keep shape + no-NaN assertions.

Local sbatch verification on B200 (job 2233456):
  TestGemma4VisionTower::test_forward_random_weights             PASSED
  TestGemma4MMTowerRMSNormConvention
    ::test_vision_attention_uses_native_rmsnorm_adapter           PASSED
  test_modeling_gemma4.py (full file)                            74/74 PASSED

Pre-existing main-side flake / CI infra failures from build 2316 that
also showed up in build 2317 without this PR's diff and are NOT
addressed here:
  - tests/unittest/auto_deploy/singlegpu/models/test_qwen3_5_moe.py
      ::test_vision_attention_matches_reference
      ::test_vision_block_matches_reference
      ::test_vlm_wrapper_delta_is_request_scoped_no_cross_call_leakage
  - tests/perf/test_perf_sanity.py
      ::test_e2e[aggr_upload-llama3_1_8b_fp8_ad_hopper-llama3_1_8b_ad_ws1_1k1k]

Signed-off-by: Hudayday <tianruih@nvidia.com>
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49169 [ run ] completed with state SUCCESS. Commit: 23a7bee
/LLM/main/L0_MergeRequest_PR pipeline #38847 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

@Hudayday

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49302 [ run ] triggered by Bot. Commit: a1205d8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49302 [ run ] completed with state SUCCESS. Commit: a1205d8
/LLM/main/L0_MergeRequest_PR pipeline #38964 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

@Hudayday

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49343 [ run ] triggered by Bot. Commit: a1205d8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49343 [ run ] completed with state SUCCESS. Commit: a1205d8
/LLM/main/L0_MergeRequest_PR pipeline #39000 completed with status: 'SUCCESS'

CI Report

Link to invocation

Comment thread tensorrt_llm/_torch/models/modeling_gemma4_audio.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_gemma4_audio.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_gemma4_vision.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_gemma4_vision.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_gemma4_vision.py Outdated
Addresses 5 reviewer comments on PR 14300:

C1+C2: Drop leading underscore from private classes per reviewer style
preference. Renames `_Gemma4{Audio,Vision}RMSNorm` and
`_{Audio,Vision}Output` to their public counterparts. Test imports
updated.

C4+C5: Lift the per-image Python loop in `_get_image_features` to a
single batched `vision_tower(pixel_values, image_position_ids)` call —
matches the pattern in Gemma3vl / Qwen2VL / Qwen3VL / LlavaNext /
Nemotron-Nano-dyn-res. The inner `Gemma4VisionModel.forward` already
supported B>1 via varlen `cu_seqlens` (per-image `seq_lens_cpu` from
`valid.sum(dim=-1)`); only the caller-side loop had to go.
`max_num_requests` on the tower-local `attn_metadata` is bumped from a
hardcoded 1 to the LLM-side value (typically 8192) — see below for the
workspace story.

FMHA enablement: vision attention previously fell into the unfused MHA
path (`Fall back to unfused MHA because of unsupported data type`) on
all 4 Gemma4 variants. qkv reaching `thop.attention` was not bf16/half
due to HF `Gemma4RMSNorm` / RoPE intermediate fp32 promotion, and
26B/31B vision head_dim=72 is not in the trtllm-gen sm100a cubin set
(H64/H80/H128/H256/H512). Two changes restore the FMHA path:

  - Explicit dtype cast on q/k/v to the qkv_proj weight dtype right
    before `forward_impl` keeps the FMHA dispatcher's mType check happy.

  - For variants whose HF head_dim is not FMHA-supported (currently
    26B/31B vision: 72), `Gemma4VisionAttention.__init__` overrides
    self.head_dim to the next supported size (72 to 80), and
    `Gemma4VisionModel._pad_attention_head_dim` zero-pads q/k/v/o_proj
    weights along the head dim at load time. RMSNorm + RoPE operate on
    the unpadded HF head_dim — zeros in K/V don't contribute to QK^T,
    zeros in V produce zero output channels, zeros in o_proj's last
    columns ignore those channels. Net: byte-equivalent to unpadded HF
    math, plus ~11% extra compute (80/72) on 26B/31B.

Empirical (B200, samples=30, gemma4_main_c163c72867 wheel + this
overlay):

  variant  head_dim  FMHA fallback  workspace  MMMU acc  CoVoST BLEU
  E2B      64                  0    ~32 MB     42.22     21.87
  E4B      64                  0    ~32 MB     48.00     24.69
  26B-A4B  72 -> 80            0    ~32 MB     56.67     -
  31B      72 -> 80            0    ~32 MB     65.22     -

Workspace at max_num_requests=1 pre-fix (unfused) was 273 MB on 31B at
3135 patches/image. FMHA on path is linear in tokens (no qk_buf, no
attention_mask buffer) so the bump to LLM-side max_num_requests keeps
the workspace ~tens of MB regardless of B. E4B CoVoST BLEU 24.69
matches the HF tower baseline 24.54 — audio path quality preserved.

C3: `valid.sum(dim=-1).to(torch.int).cpu()` sync is API-mandated.
AttentionMetadata.prompt_lens is typed Optional[List[int]] (host);
same pattern in modeling_siglip.py, modeling_qwen2vl.py,
modeling_nemotron_nano.extract_feature_dynamic. No code change; reply
documents the contract.

Tests added:

  - `TestGemma4VisionCrossImageBatching` in test_modeling_gemma4.py:
    structural guards that `_get_image_features` contains exactly one
    vision_tower call (no per-image for loop) and that the vision
    attn_metadata does not hardcode max_num_requests=1.

  - `TestGemma4VisionHeadDimPadding` in test_gemma4_multimodal.py:
    verifies attention overrides head_dim to 80 when HF head_dim is 72,
    skips override when head_dim is already FMHA-supported, zero-pads
    qkv_proj/o_proj weights at load, and matches HF forward output at
    atol/rtol=5e-2.

  - `test_forward_batched_matches_per_image_loop` in
    test_gemma4_multimodal.py: B=3 batched call output matches B=1
    looped x 3 concat to atol=1e-4 (guards against future cross-image
    leakage in shared attn_metadata state).

Signed-off-by: Hudayday <tianruih@nvidia.com>
@Hudayday

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49441 [ run ] triggered by Bot. Commit: c157778 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49441 [ run ] completed with state SUCCESS. Commit: c157778
/LLM/main/L0_MergeRequest_PR pipeline #39086 completed with status: 'SUCCESS'

CI Report

Link to invocation

@Hudayday
Hudayday requested a review from yechank-nvidia May 21, 2026 02:06
…CPU-side

Follow-up to PR 14300 review (yechank-nvidia):
> But it seems Gemma4 derives those values from GPU tensor to CPU which
> occurs CPU-GPU sync. and hurt performance very badly. Is there anyway
> to derive those numbers from host tensors?

Earlier reply conflated two different sync patterns:
  - SigLip / Qwen2VL / Nemotron-Nano derive `prompt_lens` from CPU-side
    processor metadata (e.g. Qwen2VL's `grid_thw.tolist()` on a CPU tensor)
    so the subsequent pinned-memory H2D in `attn_metadata.prepare()` is
    non-blocking.
  - Gemma4 vision was reducing `pixel_position_ids` (GPU-resident for RoPE)
    with `valid.sum(-1).cpu()` — a real GPU→CPU blocking sync that stalled
    the CUDA stream per forward.

Fix mirrors the Qwen2VL pattern:

  Gemma4InputProcessor (CPU)
    pixel_position_ids ───── [executor moves to GPU later] ────► RoPE
    image_seq_lens     ─── List[int], never tensorised to GPU ──► attn_metadata

`Gemma4InputProcessor.__call__` now computes the per-image valid-patch
count on CPU at the moment it inserts the `-1` padding sentinels into
`image_position_ids` (both tensors are still host-resident at that point;
the executor pipeline ships them to GPU later). The count is stored as a
Python `List[int]` under `multimodal_data["image"]["image_seq_lens"]`
(and the analogous `image_data["video"]` field for video frames) — Python
lists aren't touched by the device-migration step, so the value stays CPU
through the entire pipeline.

`Gemma4ForConditionalGeneration.forward` collects per-request lists and
flattens with `list.extend(...)` (no torch ops), then passes through to
`_get_image_features` as a new optional kwarg.

`Gemma4VisionModel.forward` now picks `seq_lens_cpu` from one of two CPU
branches (both `device="cpu"`, neither touches GPU):
  - caller-supplied list → `torch.tensor(list, dtype=torch.int, device="cpu")`
  - no list (synthetic / unit-test input with no -1 padding) →
    `torch.full((B,), N, dtype=torch.int, device="cpu")`

The previous `valid.sum(dim=-1).to(torch.int).cpu()` GPU reduction +
blocking D2H is gone entirely (no fallback path — the assert on
`len(image_seq_lens) == B` catches plumbing bugs, and the all-valid
default covers synthetic inputs). RoPE on GPU is untouched —
`pixel_position_ids` still flows GPU-side into `rotary_emb`.

`attn_metadata.prepare()` does the subsequent pinned-memory H2D async
(non-blocking) — same as SigLip / Qwen2VL / Nemotron-Nano.

Empirical (B200, 30-sample MMMU + 20-sample CoVoST, gemma4_main_c163c72867
wheel + this overlay):

  variant  task         FMHA fall-back  workspace  accuracy
  E2B      MMMU                      0  ~32 MB     42.22
  E4B      MMMU                      0  ~32 MB     48.00
  26B-A4B  MMMU                      0  ~32 MB     58.89
  31B      MMMU                      0  ~32 MB     65.22
  E2B      CoVoST en_de              0  -          BLEU 21.87
  E4B      CoVoST en_de              0  -          BLEU 24.69

No regression vs the previous commit (c157778). Sync count per
MMMU 30-sample run drops by ~1000 D2H syncs (one per `_get_image_features`
call); per-sync cost is ~50-100 us on B200, so ~50-100 ms total wallclock
saved per 30 samples — modest in throughput-bound benchmarks, more visible
in latency-sensitive paths.

New unit test added: `TestGemma4VisionTower.test_forward_image_seq_lens_kwarg_matches_default`
verifies byte-equivalent output between the explicit-list path and the
default all-valid path (atol=1e-4).

Signed-off-by: Hudayday <tianruih@nvidia.com>
@Hudayday

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49565 [ run ] triggered by Bot. Commit: 001a6a0 Link to invocation

…kill gaps

Two fixes flagged by the multimodal-model-onboarding skill (NVIDIA#13842) audit
against this PR:

1. Drop `transformers.models.gemma4` direct imports.

   `Gemma4VisionRMSNorm` and `Gemma4AudioRMSNorm` previously inherited from
   `transformers.models.gemma4.modeling_gemma4.Gemma4RMSNorm`. The skill's
   "Multimodal encoder — port to TRT-LLM modules" rule:

   > No imports of HF modules or computations from
   > transformers.models.<family> — they couple the file to a specific
   > transformers release and break silently on upstream refactors.

   Both classes are now ~12-line native `nn.Module` subclasses implementing
   the plain `w * x / rms` Gemma4 convention (not the Gemma3 LLM
   `(1 + w) * x / rms` variant) with fp32 accumulator and cast-back to
   input dtype — mathematically identical to the HF class.

2. Declare `multimodal_data_device_paths`.

   The skill's Contract 2 / Phase 2 checklist requires every multimodal
   wrapper to list the dotted paths in `multimodal_data` that the engine
   should ship to GPU. Anything not listed stays CPU-resident.

   Added on `Gemma4ForConditionalGeneration`: lists `pixel_values`,
   `image_position_ids`, `pixel_values_videos`, `audio_features`, and
   `audio_features_mask` as GPU paths. Critically, `image_seq_lens` and
   `video.image_seq_lens` (the CPU `List[int]` from the previous commit's
   D2H-sync elimination) are NOT listed — the prior commit's invariant
   ("Python list isn't device-migrated") is now an explicit architectural
   declaration instead of an implicit assumption.

Empirical (B200, gemma4_main_c163c72867 wheel + this overlay):

  variant  task           FMHA fall-back  accuracy
  E2B      MMMU 30s                    0  42.22
  E4B      MMMU 30s                    0  48.00
  26B-A4B  MMMU 30s                    0  57.44 (= prior 58.89 within seed noise)
  31B      MMMU 30s                    0  65.22
  E2B      CoVoST en_de 20s            0  BLEU 21.87
  E4B      CoVoST en_de 20s            0  BLEU 24.69 (= HF baseline 24.54)

All within seed noise of the prior commit (`001a6a01bd`); no perf or
accuracy regression introduced by porting the RMSNorm.

Tests updated:

  - `test_audio_rmsnorm_inherits_from_hf_class` /
    `test_vision_rmsnorm_inherits_from_hf_class` rewritten as
    `test_audio_rmsnorm_does_not_import_hf_class` /
    `test_vision_rmsnorm_does_not_import_hf_class`: guard the new
    invariant (no `transformers.models.gemma4` imports in either modeling
    file; class MRO has no HF module).
  - `test_audio_rmsnorm_matches_hf_class_with_same_weights` /
    `test_vision_rmsnorm_*` continue to verify byte-identical output vs
    HF reference at fp32 with shared weights — confirms math equivalence.

Signed-off-by: Hudayday <tianruih@nvidia.com>
@Hudayday

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@yechank-nvidia yechank-nvidia 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.

Thanks all for the hard work!

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49594 [ run ] triggered by Bot. Commit: 287a601 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49565 [ run ] completed with state ABORTED. Commit: 001a6a0

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #49594 [ run ] completed with state SUCCESS. Commit: 287a601
/LLM/main/L0_MergeRequest_PR pipeline #39218 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

@Hudayday
Hudayday merged commit 4c5500e into NVIDIA:main May 21, 2026
7 checks passed
xxi-nv pushed a commit to xxi-nv/TensorRT-LLM that referenced this pull request May 22, 2026
Signed-off-by: Hudayday <tianruih@nvidia.com>
KleinBlueC pushed a commit to KleinBlueC/TensorRT-LLM that referenced this pull request May 26, 2026
Signed-off-by: Hudayday <tianruih@nvidia.com>
bmarimuthu-nv pushed a commit to nv-auto-deploy/TensorRT-LLM that referenced this pull request May 28, 2026
Signed-off-by: Hudayday <tianruih@nvidia.com>
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.

3 participants