Draft: [None][refactor] Mixed Modality Support for Nemotron Nano Omni V3#16764
Draft
aswinvisva wants to merge 9 commits into
Draft
Draft: [None][refactor] Mixed Modality Support for Nemotron Nano Omni V3#16764aswinvisva wants to merge 9 commits into
aswinvisva wants to merge 9 commits into
Conversation
…der groups Introduces a small framework abstraction so any multimodal model can opt into mixed-modality encoding by declaring which modalities share an encoder, without reimplementing the interleave logic per model. Framework additions: - ``EncoderGroup`` (dataclass): a set of modalities that share one encoder call, plus a model-provided ``build_batched_input`` that cats items from those modalities into one encoder input. - ``_lengths_by_modality``: inverts the prompt-ordered ``multimodal_embedding_lengths`` (already computed by ``create_input_processor_with_hash`` and plumbed via ``py_multimodal_data``) into per-modality per-item lists via each request's ``mm_item_order``. Enforces the "mixed request must carry manifest" invariant once, centrally, instead of per-model tripwires. - ``_reorder_embeds_by_manifest``: slices per-modality tensors item-by-item and concatenates in each request's prompt-order manifest. - ``MultimodalModelMixin.encode_multimodal_by_groups``: iterates the model's declared groups, issues one encoder call per group, splits the output back per-modality, and reorders into prompt order. Qwen3-VL migration: - Deletes ``_parse_and_batch_multimodal_data``, ``_interleave_multimodal_data``, and the two-branch ``Qwen3VisionModelBase.forward`` — all subsumed by the generic path. - Adds ``Qwen3VisionModelBase.encode_batched(pixel_values, grid_thw)``: one ViT call, folds deepstack streams into hidden_dim. - Adds ``_qwen3vl_build_batched_input`` free function: cats image then video items across requests into one ViT input. - Registers ``mm_encoder_groups`` on ``Qwen3VLModelBase`` at ``__init__``. - ``encode_multimodal_inputs`` routes through ``encode_multimodal_by_groups`` via the existing ``get_multimodal_embeddings`` cache-aware orchestrator. Net Qwen3-VL diff: -114 lines. Zero changes to input processor, MM hashing, IPC, or ``MultimodalParams`` — all lengths ride the existing ``multimodal_embedding_lengths`` field on ``multimodal_data``. Tests: - ``test_encoder_group.py`` (new): pure-helper coverage — manifest inversion, single-modality fallback, mixed-request rejection, per-modality cursor advancement across requests. - ``test_qwen3vl_mixed_modality_encoder.py`` (updated): reworked to route through the mixin path via a minimal ``_StubModel`` with a marker encoder. Cases cover single-modality, mixed I-V-I, missing-manifest rejection, cross-request batching, and heterogeneous single-modality batches. Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
- Critical: ``Qwen3VLModelBase.forward``'s raw-MM branch was still passing ``self.mm_encoder.forward`` to ``get_multimodal_embeddings``, but that method no longer exists after the refactor (replaced by ``encode_batched``). Route through ``self.encode_multimodal_by_groups`` instead so raw image/video requests go through the grouped path. - Add ``test_raw_inputs_route_through_encode_multimodal_by_groups`` covering the callable used by that branch. - Annotate ``_synthesize_single_modality_manifest``'s ``modalities`` parameter as ``Iterable[str]``. - Switch ``_qwen3vl_build_batched_input`` to Python 3.10+ built-in generics (``list[MultimodalParams]`` / ``dict[str, Any]``). - Add ``-> None`` and parameter annotations to newly added test callables in ``test_qwen3vl_mixed_modality_encoder.py`` and ``test_chat_utils.py::test_item_order``. Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
…trings - Move `EncoderGroup`, `_lengths_by_modality`, `_reorder_embeds_by_manifest`, and `_synthesize_single_modality_manifest` from `modeling_multimodal_utils.py` to `modeling_multimodal_mixin.py`; the framework is the only consumer. - Freeze `EncoderGroup` (`@dataclass(frozen=True)`) and add per-field docstrings covering the `build_batched_input` / `encoder_fn` contract. - Switch `:class:` / `:attr:` RST roles to single backticks in comments. - Drop the `if total > 0` special case in `encode_multimodal_by_groups`; reuse `lengths` directly via `per_modality_lengths.update`. - Replace the manual prefix-sum with `itertools.accumulate`. - Add a comment in `_lengths_by_modality` explaining that raw-prompt entrypoints are the single enforcement point for the mm_item_order requirement on >1-modality requests. - Update import sites in `modeling_qwen3vl.py` and the two multimodal tests. Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
…cale Qwen3-VL preferred `format="np"` for videos so the HF processor could rescale/permute uint8 HWC frames in one pass, while `ImageMediaIO` defaulted to `format="pt"` (pre-rescaled float32 Tensor). The HF processor applies a single `do_rescale` flag to both modalities, so a mixed image+video request set `do_rescale=False` (because the image was pre-rescaled) and left the video at raw 0-255. The ViT then produced garbage for the video item and the model failed to describe any video color in mixed prompts. Fix: - Add "np" format to `ImageMediaIO` so images can arrive as uint8 HWC too. - Teach `BaseMultimodalInputProcessor.get_num_tokens_per_image` / `get_num_tokens_per_video` (and the Qwen2-VL overrides) to read H/W from `np.ndarray` shape. - Update `Qwen3VLInputProcessorBase.get_preferred_media_io_kwargs` to request "np" for both modalities so the HF processor rescales+ normalizes uniformly. - Replace the two `do_rescale=False` shortcuts in `Qwen3VLInputProcessorBase._preprocess` with a fail-fast guard that raises if image and video arrive in mismatched pre-rescale states. Verified end-to-end: the color-interleave correctness harness goes from 0/8 to 6/8 passing on cosmos3-nano-reasoner, and a 10-case IV/VI sanity harness goes from 0/10 to 7/10 (remaining failures are model-quality issues on shared-RGB-channel color pairs that reproduce on main). Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
The generic grouped-encoder path (encode_multimodal_by_groups -> _reorder_embeds_by_manifest) unconditionally called torch.cat(slices), which raises on an empty slice list. This happens on the executor's KV-cache profiling pass: _encode_dummy_inputs runs the encoder on a worst-case dummy batch that carries the encoder tensors but no multimodal_embedding_lengths, so per-modality lengths (and thus the sliced embeds) come back empty and no slices are collected. It surfaced as "Executor worker returned error" during executor initialization. Return a correctly-typed empty embedding tensor when no slices resolve. The encoder forward still runs, so its activation memory is captured by peak-memory profiling. Add unit tests for the empty-bookkeeping and no-embeds cases. Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
The generic mixed-modality refactor replaced Qwen3VisionModelBase.forward with encode_batched, but the standalone/disaggregated MM-encoder executor path (_forward_step_mm_encoder_only) still calls self.model.forward(params) and expects a list of embeddings. Qwen3-VL fell through to PyTorch's _forward_unimplemented and raised NotImplementedError, while the sibling Qwen2VisionModelBase still exposes forward for this path. Restore forward on Qwen3VisionModelBase, delegating to the same batching (_qwen3vl_build_batched_input) and encode_batched the encoder group uses, and return a single-element list to match the mm-encoder-only contract. Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
The mm-encoder-only executor (_forward_step_mm_encoder_only) splits the vision model's output into per-request slices using request-ordered split_lengths. Qwen3VisionModelBase.forward was returning rows in modality-grouped order (all images across requests, then all videos) via _qwen3vl_build_batched_input, so mixed image+video batches got silently misassigned across requests. Extract the group-encode-then-reorder algorithm out of MultimodalModelMixin.encode_multimodal_by_groups into a module-level free helper of the same name. Both the aggregated path (Qwen3VLModelBase via functools.partial) and the mm-encoder-only path (Qwen3VisionModelBase via its own mm_encoder_groups property) now call the same helper, so row ordering is applied once. Deletes the one-line instance-method wrapper and consolidates the Qwen3-VL EncoderGroup definition on the vision class so Qwen3VLModelBase.__init__ reads it via self.mm_encoder.mm_encoder_groups instead of redeclaring it. Adds a regression test covering the reviewer's exact [video-only req0, image-only req1] scenario. Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
…ty requests `not images or ...` conflates "modality absent" with "modality is pre-rescaled": if only one modality is present, the absent side collapsed to image_is_pre=True and a uint8 numpy payload on the present side would spuriously trip the mismatch check. Masked in practice by get_preferred_media_io_kwargs shipping `np` for both, but caller media_io_kwargs overrides would surface it. Compute each side's pre-rescale flag only when its modality is present, gate the mismatch check on both being present, and derive do_rescale from whichever modality actually is. Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
Contributor
WalkthroughMultimodal encoding now supports configurable encoder groups, prompt-order reconstruction, raw modality detection, and typed empty outputs. NanoV2VL and Qwen3-VL use grouped routing, while media loading and token sizing support NumPy inputs. ChangesMultimodal encoder grouping
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Request
participant GroupedEncoder
participant VisionEncoder
participant PromptAssembler
Request->>GroupedEncoder: submit modality payloads and ordering metadata
GroupedEncoder->>VisionEncoder: build and encode grouped inputs
VisionEncoder-->>GroupedEncoder: return concatenated embeddings
GroupedEncoder->>PromptAssembler: slice and reorder embeddings
PromptAssembler-->>Request: return prompt-ordered embeddings
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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_nemotron_nano.py`:
- Around line 2269-2279: Ensure the payload-building logic near the modalities
computation always provides the modality_type expected by merge_evs_mm_embeds(),
including mixed image/video/audio requests, or bypass EVS merging before that
method reads the key. Preserve the existing single-modality values and ensure
mixed requests do not raise KeyError.
- Around line 954-965: Update the local _vision function to split mixed
image/video multimodal parameters into virtual single-modality params before
calling NanoV2VLMultimodalEncoder.forward, mirroring _encode_vision. Preserve
the returned embedding concatenation while ensuring each forward call receives
exactly one modality and remains aligned with encode_multimodal_by_groups.
🪄 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: f9bd18db-b985-49c6-9d6d-a5df81fc1a4d
📒 Files selected for processing (10)
tensorrt_llm/_torch/models/modeling_multimodal_mixin.pytensorrt_llm/_torch/models/modeling_multimodal_utils.pytensorrt_llm/_torch/models/modeling_nemotron_nano.pytensorrt_llm/_torch/models/modeling_qwen2vl.pytensorrt_llm/_torch/models/modeling_qwen3vl.pytensorrt_llm/inputs/media_io.pytensorrt_llm/inputs/registry.pytests/unittest/_torch/multimodal/test_encoder_group.pytests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.pytests/unittest/llmapi/apps/test_chat_utils.py
Comment on lines
+954
to
+965
| def _vision(multimodal_params): | ||
| embeds, _ = super(NanoV2VLMultimodalEncoder, self).forward(multimodal_params) | ||
| return torch.cat(embeds, dim=0) if embeds else torch.empty(0) | ||
|
|
||
| return [encode_multimodal_by_groups( | ||
| (EncoderGroup( | ||
| modalities=("image", "video"), | ||
| encoder_fn=_vision, | ||
| build_batched_input=lambda params: {"multimodal_params": params}, | ||
| ),), | ||
| multimodal_params, | ||
| )] |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -nP -C4 'NanoV2VLMultimodalEncoder|_forward_step_mm_encoder_only' --type=py
ast-grep run --pattern '_lengths_by_modality($$$)' --lang python tensorrt_llm/_torch/models/modeling_multimodal_mixin.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map relevant symbols and inspect the target section plus supporting helpers.
ast-grep outline tensorrt_llm/_torch/models/modeling_nemotron_nano.py --view expanded | sed -n '1,260p'
printf '\n===== TARGET SECTION =====\n'
sed -n '930,990p' tensorrt_llm/_torch/models/modeling_nemotron_nano.py
printf '\n===== RELATED HELPERS =====\n'
rg -n "def _encode_vision|def merge_evs_mm_embeds|_lengths_by_modality|encode_multimodal_by_groups|modality_type|_forward_step_mm_encoder_only|mm_encoder_only" tensorrt_llm/_torch/models/modeling_nemotron_nano.py tensorrt_llm/_torch/models/modeling_multimodal_mixin.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 11050
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tensorrt_llm/_torch/models/modeling_nemotron_nano.py --view expanded | sed -n '1,260p'
printf '\n===== TARGET SECTION =====\n'
sed -n '930,990p' tensorrt_llm/_torch/models/modeling_nemotron_nano.py
printf '\n===== RELATED HELPERS =====\n'
rg -n "def _encode_vision|def merge_evs_mm_embeds|_lengths_by_modality|encode_multimodal_by_groups|modality_type|_forward_step_mm_encoder_only|mm_encoder_only" tensorrt_llm/_torch/models/modeling_nemotron_nano.py tensorrt_llm/_torch/models/modeling_multimodal_mixin.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 11050
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path('tensorrt_llm/_torch/models/modeling_nemotron_nano.py')
text = path.read_text()
for needle in [
'class NanoV2VLMultimodalEncoder',
'def _vision',
'def _encode_vision',
'def merge_evs_mm_embeds',
]:
idx = text.find(needle)
print(f'== {needle} ==')
if idx == -1:
print('NOT FOUND')
continue
start = max(0, text.rfind('\n', 0, idx - 400))
end = min(len(text), text.find('\n', idx + 1200))
print(text[start:end])
print()
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 6896
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether this code path is even reachable from a mixed-modality batch builder.
rg -n "encode_multimodal_by_groups\(|EncoderGroup\(|modalities=\(\"image\", \"video\"\)" tensorrt_llm/_torch/models/modeling_nemotron_nano.py tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
printf '\n===== call sites for multimodal encoder =====\n'
rg -n "NanoV2VLMultimodalEncoder|MultimodalEncoder|multimodal_params" tensorrt_llm/_torch/models/modeling_nemotron_nano.py
printf '\n===== mixed-modality ordering helpers =====\n'
rg -n "_lengths_by_modality|modality_type|mm_item_order|video_pruning_rate" tensorrt_llm/_torch/models/modeling_nemotron_nano.py tensorrt_llm/_torch/models/modeling_multimodal_mixin.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 9182
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '===== multimodal mixin ordering contract =====\n'
sed -n '55,220p' tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
printf '\n===== vision encoder forward =====\n'
sed -n '734,790p' tensorrt_llm/_torch/models/modeling_nemotron_nano.py
printf '\n===== mm-encoder-only and full-model group setup =====\n'
sed -n '3077,3148p' tensorrt_llm/_torch/models/modeling_nemotron_nano.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 14268
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '===== NanoV2VLVisionEncoder.forward =====\n'
sed -n '734,832p' tensorrt_llm/_torch/models/modeling_nemotron_nano.py
printf '\n===== any internal bucket ordering =====\n'
sed -n '832,930p' tensorrt_llm/_torch/models/modeling_nemotron_nano.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 10557
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '===== call_with_text_prompt / prompt assembly =====\n'
sed -n '2202,2295p' tensorrt_llm/_torch/models/modeling_nemotron_nano.py
printf '\n===== encode / forward caller context =====\n'
sed -n '3086,3168p' tensorrt_llm/_torch/models/modeling_nemotron_nano.py
printf '\n===== multimodal embedding getters =====\n'
rg -n "def get_multimodal_embeddings|def get_attached_multimodal_embeddings|has_raw_multimodal_payload|mm_item_order" tensorrt_llm/_torch/models/modeling_multimodal_mixin.py tensorrt_llm/_torch/models/modeling_nemotron_nano.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 10158
_vision needs to split image/video params before calling super().forward.
encode_multimodal_by_groups can pass mixed image/video requests here, but NanoV2VLVisionEncoder.forward only accepts exactly one modality per param. Mirror _encode_vision by wrapping each modality into virtual single-modality params; otherwise mixed batches will raise and the group split won’t line up.
🤖 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_nemotron_nano.py` around lines 954 - 965,
Update the local _vision function to split mixed image/video multimodal
parameters into virtual single-modality params before calling
NanoV2VLMultimodalEncoder.forward, mirroring _encode_vision. Preserve the
returned embedding concatenation while ensuring each forward call receives
exactly one modality and remains aligned with encode_multimodal_by_groups.
Comment on lines
+2269
to
+2279
| # Legacy `modality_type` disambiguates the single-modality bucket for | ||
| # code paths (EVS merge, encoder-only EPD) that predate the mixin's | ||
| # per-modality dispatch. Mixed requests carry `mm_item_order` from | ||
| # chat parsing, which the mixin uses instead. | ||
| modalities = [m for m in ("image", "video", "audio") if m in multimodal_data] | ||
| if len(modalities) == 1: | ||
| m = modalities[0] | ||
| multimodal_data["modality_type"] = m | ||
| if self.video_pruning_rate > 0 and m != "video": | ||
| multimodal_data[m]["evs_ids"] = input_ids[0].to(torch.int32) | ||
| return input_ids[0].to(torch.int32).tolist(), {"multimodal_data": multimodal_data} |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm merge_evs path and whether mixed modality can coexist with EVS.
rg -nP -C3 'modality_type|video_pruning_rate' tensorrt_llm/_torch/models/modeling_nemotron_nano.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 9446
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE=tensorrt_llm/_torch/models/modeling_nemotron_nano.py
# Inspect the relevant methods around the suspected call sites.
sed -n '2810,2865p' "$FILE"
printf '\n----\n'
sed -n '3090,3165p' "$FILE"
printf '\n----\n'
sed -n '2238,2285p' "$FILE"Repository: NVIDIA/TensorRT-LLM
Length of output: 8722
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE=tensorrt_llm/_torch/models/modeling_nemotron_nano.py
# Find the multimodal dispatch path and any virtual-param wrapping.
rg -n -C4 'modality_type|mm_item_order|MultimodalParams\(|build_disagg_prefill_multimodal_inputs|encode_multimodal_by_groups|get_attached_multimodal_embeddings|merge_evs_mm_embeds' "$FILE"Repository: NVIDIA/TensorRT-LLM
Length of output: 8646
Mixed requests still need a modality_type guard
merge_evs_mm_embeds() unconditionally reads multimodal_data["modality_type"], but the payload builder only sets that key for single-modality requests. A mixed image/video/audio request with EVS enabled reaches this path and raises KeyError before the "video" not in modalities guard. Keep modality_type on mixed payloads or skip the EVS merge earlier.
🤖 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_nemotron_nano.py` around lines 2269 -
2279, Ensure the payload-building logic near the modalities computation always
provides the modality_type expected by merge_evs_mm_embeds(), including mixed
image/video/audio requests, or bypass EVS merging before that method reads the
key. Preserve the existing single-modality values and ensure mixed requests do
not raise KeyError.
aswinvisva
force-pushed
the
avisva/nemotron-mm-mixin-clean
branch
from
July 22, 2026 23:50
f31981e to
92e5a98
Compare
Migrate NemotronH_Nano_VL_V2 to the generic mixed-modality mixin so
image + video + audio can coexist in a single request.
Model
-----
* Inherits `MultimodalModelMixin` and registers two `EncoderGroup`s:
`("image", "video")` sharing the ViT and `("audio",)` on the sound
encoder. `encode_multimodal_inputs` delegates through
`encode_multimodal_by_groups`, which runs one encoder call per
group and reorders rows per each request's `mm_item_order`
manifest.
* `_encode_vision` replaces `_encode_multimodal`: image and video
items are dispatched to the shared ViT via per-modality virtual
params (since `NanoV2VLVisionEncoder.forward` still keys on the
legacy `modality_type`). Video-with-embedded-audio is encoded and
interleaved per-video here because it shares the video item's
prompt-token run. `num_tokens_in_video` is stashed on each video
param for the EVS post-step.
* `_encode_audio_group` handles standalone audio items.
* Bucket-key detection replaces `modality_type` reads in
`NanoV2VLVisionEncoder.forward`, `apply_evs`,
`_check_encoders_exist`, and the attached-embeddings EVS gate.
`modality_type` is still emitted by the input processor for
single-modality requests so the legacy EVS merge path keeps
working unchanged.
* `NanoV2VLMultimodalEncoder` (EPD encoder-only wrapper) routes
through `encode_multimodal_by_groups`. Its `_vision` encoder_fn
partitions the incoming batch by modality bucket and invokes the
parent `NanoV2VLVisionEncoder.forward` once per modality (image
first, then video). Parent reassembles output in input-param
order; without the pre-partition, the resulting `torch.cat` would
be request-order, which the group split would then slice as if it
were modality-major and hand each request the wrong rows.
* Guard against mixed-modality + EVS in `forward` before reaching
`merge_evs_mm_embeds`, which reads the legacy `modality_type` tag
on every context param. Mixed requests drop that tag, so without
the guard the merge would `KeyError` deep in its loop.
* `has_raw_multimodal_payload` in `modeling_multimodal_utils.py` is
now bucket-key based so mixed requests are not misclassified as
attached-embed E/P prefills.
Input processor
---------------
* Removes the single-modality-per-request gate. `_process_images`,
`_process_images_dynamic`, `_process_video_prompts`, and
`_process_audio` now return the expanded prompt string, and
`call_with_text_prompt` chains them and tokenizes once at the end.
Emits per-modality `multimodal_data["image"|"video"|"audio"]`
buckets in parallel; `mm_item_order` from chat parsing carries
send order.
* Pins `content_format=ContentFormat.STRING` on Nano's placeholder
metadata. The OPENAI auto-detection would otherwise let Nano's
Jinja regroup placeholders by hardcoded modality order, breaking
the `mm_item_order == prompt order` invariant the mixin reorder
relies on — mixed audio+image requests would lose one row when the
mixin sliced with mislabeled lengths. STRING routes serve through
`add_multimodal_placeholders`, which lays placeholders in
`mm_item_order` order into a plain string that Nano's template
consumes verbatim via its `message.content is string` branch (same
pattern Qwen3-VL uses).
Smoke
-----
Verified end-to-end on `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-FP8`
with a 7-case mixed-modality suite (image / video / audio /
image+audio / audio+image / image+video / image+video+audio). Each
case asserts per-item cues appear in the reply in prompt-item order:
image/video color words and an audio-descriptor keyword for audio
items. All 7 pass with matching in-order cues in the model output.
Net effect on `modeling_nemotron_nano.py`: -20 lines. Deletion of
`_encode_multimodal` and the modality_type single-modality machinery
outweighs the mixin wiring.
Signed-off-by: Aswin Visva <avisva@nvidia.com>
Signed-off-by: Aswin Visva <31215515+aswinvisva@users.noreply.github.com>
aswinvisva
force-pushed
the
avisva/nemotron-mm-mixin-clean
branch
from
July 23, 2026 18:14
92e5a98 to
a490685
Compare
2ez4bz
reviewed
Jul 23, 2026
| # One `EncoderGroup` per modality — matches how the parent's | ||
| # forward emits one row block per param and lets the framework | ||
| # do the per-modality reorder in one place. | ||
| def _encoder_fn(modality): |
Collaborator
There was a problem hiding this comment.
Nit:
_make_encoder_fn?- You could consider having
_encoder_fn(multimodal_params, modality)without the inner_fn, and passpartial(_encoder_fn, modality="image")in e.g. line 972.
| def _process_images( | ||
| self, images: List[Image.Image | torch.Tensor], text_prompt: str | ||
| ) -> Tuple[Dict[str, Any], torch.Tensor]: | ||
| ) -> Tuple[Dict[str, Any], str]: |
Collaborator
There was a problem hiding this comment.
Can you explain why it was necessary to switch from input ids (int tensor) to a text prompt? Does this change affect any consumer of this class's public methods?
| }, | ||
| placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, | ||
| placeholders_separator="\n", | ||
| # STRING pre-inserts placeholders in `mm_item_order` order; the OPENAI |
Collaborator
There was a problem hiding this comment.
Hm... I remember I specifically had to make this change to make the nemotron v3 outputs match that of vLLM. If you run the E2E accuracy tests, do they show the same scores (within noise) both before and after the change?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Work in progress - make Nemotron Nano Omni V3 a child class of
MultimodalModelMixinand leverage the generic mixed modality support (not yet merged, so folded into this PR as well)Dev Engineer Review
EncoderGroupdefinitions and shared prompt-order reconstruction viamm_item_order.modality_type, and consistency with project coding guidelines.QA Engineer Review
Test-code changes include:
tests/unittest/_torch/multimodal/test_encoder_group.pycovering modality-length mapping, manifest synthesis, prompt-order reassembly, cursor advancement, and typed empty outputs.tests/integration/test_lists/ortest-db//qa/was not identified for these new unit tests.Verdict: needs follow-up — confirm CI/manual test-list coverage and run the seven mixed-modality smoke cases.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.