[https://nvbugs/6550127][fix] Added a fourth dispatch branch for plain dim-0 stacks, resolved per modality… - #17209
Conversation
… encoder cache Gemma4's raw multimodal inputs are a plain per-item dim-0 stack (`pixel_values` plus `image_position_ids` / `image_seq_lens`) with no companion `image_sizes` or `*_grid_thw` field, so none of the default `build_multimodal_encoder_input` branches matched and every partial encoder-cache hit raised NotImplementedError. Enabling `supports_encoder_cache` on Gemma4 exposed the gap. Add a fourth branch for that layout. The key is resolved per modality and only accepted when the tensor's first axis matches the request's declared item count, so Gemma4 video -- which flattens (B, frames, patches, C) to (B * frames, patches, C) and therefore counts frames on dim 0 -- still falls through instead of being sliced by the wrong axis. The existing sibling-slice pass then carries `image_position_ids` and `image_seq_lens` along, so no model-specific field handling is needed. Remove the now-obsolete waiver. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
WalkthroughThe multimodal mixin now supports slicing plain per-item dim-0 image, video, and audio stacks. It validates item counts, preserves per-item metadata, and rejects flattened layouts. Tests cover these cases, and one integration waiver was removed. ChangesPlain multimodal slicing
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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
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_multimodal_mixin.py (1)
608-625: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocument the public method arguments.
build_multimodal_encoder_inputis a public method. Its modified docstring does not documentparamoritem_indicesin a Google-styleArgs:section.As per coding guidelines, public functions must use Google-style docstrings and document their 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_multimodal_mixin.py` around lines 608 - 625, Update the docstring for the public build_multimodal_encoder_input method to add a Google-style Args: section documenting both param and item_indices, including their expected roles and types based on the method signature. Preserve the existing layout documentation.Source: Coding guidelines
🧹 Nitpick comments (1)
tests/unittest/_torch/multimodal/test_multimodal_mixin.py (1)
707-746: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover the
pixel_values_videosplain-stack key.This test uses
pixel_valuesfor video. It does not exercise the newpixel_values_videoscandidate in_MM_STACKED_ITEM_KEYS. Add a video case withpixel_values_videosand assert the same raw-input and sibling-field reordering.Proposed test expansion
-@pytest.mark.parametrize("modality", ["image", "video", "audio"]) -def test_build_multimodal_encoder_input_slices_plain_dim0_stack(modality): - pixel_key = "audio_features" if modality == "audio" else "pixel_values" +@pytest.mark.parametrize( + ("modality", "pixel_key"), + [ + ("image", "pixel_values"), + ("video", "pixel_values"), + ("video", "pixel_values_videos"), + ("audio", "audio_features"), + ], +) +def test_build_multimodal_encoder_input_slices_plain_dim0_stack(modality, pixel_key):Test coverage summary: Added
test_build_multimodal_encoder_input_slices_plain_dim0_stackandtest_build_multimodal_encoder_input_frame_flattened_video_raises. Both are covered by the directory-scoped CI selection intests/integration/test_lists/test-db/l0_h100.yml. No suppliedqa/list confirms manual-QA selection. Coverage verdict: insufficient untilpixel_values_videosis exercised.As per path instructions, test-code changes require a coverage verdict and test-list status. Based on learnings,
tests/integration/test_lists/test-db/l0_h100.ymlselectsunittest/_torch/multimodalat directory scope.🤖 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/multimodal/test_multimodal_mixin.py` around lines 707 - 746, Extend test_build_multimodal_encoder_input_slices_plain_dim0_stack to use pixel_values_videos for the video modality instead of only pixel_values, while retaining the existing image and audio key behavior. Assert that the video raw input and its image_position_ids and image_seq_lens siblings are reordered identically, and preserve the unchanged per-request scalar assertion; the existing directory-scoped CI selection already covers this test.Sources: Path instructions, Learnings
🤖 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_multimodal_mixin.py`:
- Line 734: Update the _stacked_item_key signature to use built-in generic and
union annotations: replace Dict[str, Any] with dict[str, object] or Mapping[str,
object], and Optional[str] with str | None. Remove any now-unused typing
imports.
---
Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_multimodal_mixin.py`:
- Around line 608-625: Update the docstring for the public
build_multimodal_encoder_input method to add a Google-style Args: section
documenting both param and item_indices, including their expected roles and
types based on the method signature. Preserve the existing layout documentation.
---
Nitpick comments:
In `@tests/unittest/_torch/multimodal/test_multimodal_mixin.py`:
- Around line 707-746: Extend
test_build_multimodal_encoder_input_slices_plain_dim0_stack to use
pixel_values_videos for the video modality instead of only pixel_values, while
retaining the existing image and audio key behavior. Assert that the video raw
input and its image_position_ids and image_seq_lens siblings are reordered
identically, and preserve the unchanged per-request scalar assertion; the
existing directory-scoped CI selection already covers this test.
🪄 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: 9d244c01-143c-47cf-a7ac-11bbc1ff68f1
📒 Files selected for processing (3)
tensorrt_llm/_torch/models/modeling_multimodal_mixin.pytests/integration/test_lists/waives.txttests/unittest/_torch/multimodal/test_multimodal_mixin.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
| ) | ||
|
|
||
| @staticmethod | ||
| def _stacked_item_key(modality_data: Dict[str, Any], item_count: int) -> Optional[str]: |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Use built-in generic annotations.
Replace Dict[str, Any] with a built-in generic such as dict[str, object] or Mapping[str, object]. Replace Optional[str] with str | None.
As per coding guidelines, Python code must prefer built-in generic types and | unions.
🤖 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_multimodal_mixin.py` at line 734, Update
the _stacked_item_key signature to use built-in generic and union annotations:
replace Dict[str, Any] with dict[str, object] or Mapping[str, object], and
Optional[str] with str | None. Remove any now-unused typing imports.
Source: Coding guidelines
|
Closing in favor of #17231 |
2 participants
Summary
pixel_values+image_position_ids/image_seq_lens, noimage_sizesor*_grid_thw) matched none of the default slicer's three branches, so every partial encoder-cache hit raisedNotImplementedError._MM_STACKED_ITEM_KEYSand accepted only when the tensor's first axis equals the declared item count (so Gemma4's frame-flattened video still falls through); the existing sibling-slice pass then carriesimage_position_ids/image_seq_lensautomatically.Test plan
Links
Dev Engineer Review
image_position_idsorimage_seq_lens.NotImplementedErrorbehavior.TestGemma4_26B_A4B::test_nvfp4waiver for NVBUG 6550127.QA Engineer Review
test_build_multimodal_encoder_input_slices_plain_dim0_stack(modality).test_build_multimodal_encoder_input_frame_flattened_video_raises().tests/integration/test_lists/waives.txt; no matchingtest-db/orqa/coverage entry was identified.