[Bug]: Preserve per-item processed metadata when computing multimodal token lengths

[HOYALIM]

System Info

• TensorRT-LLM: current main at 573bd5f0b8414566bd6183f74b06f8c8a99e6a69
• Reproduction hardware: Apple M1 Pro CPU, macOS arm64
• Python 3.11.15, PyTorch 2.12.1
• Backend surface: PyTorch multimodal input preprocessing/hashing contract

Information

find_mm_token_lengths accepts flattened processor-produced video_grid_thw, but it only preserves that metadata when there is one source video or exactly one grid row per video. When several source videos each expand to a variable number of processed rows, the function cannot recover item boundaries and drops video_grid_thw entirely.

That fallback is not generally equivalent: processed metadata may reflect frame sampling, segmentation, resizing, or cropping decisions that cannot be reconstructed from the raw video item. The returned lengths feed multimodal hashing, positions, masks, and layout validation.

The current source contains TODO(TRTLLM-11951) at this branch suggesting a generic per-item processed-metadata route.

Reproduction

I executed the exact current-source find_mm_token_lengths definition with a deterministic synthetic processor whose processed grid metadata is authoritative.

layout	expected	current result
one video, two grid rows	[14]	[14]
two videos, one grid row each	[2, 3]	[2, 3]
two videos, two grid rows each	[5, 9]	[-1, -1]

For the failing case, current main emits:

find_mm_token_lengths: video_grid_thw row count (4) does not match number of videos in mm_data (2); falling back to per-item recompute without video_grid_thw.

The processor receives video_grid_thw=None for both items. The result reproduced identically in three repetitions.

Expected behavior

Processor-produced multimodal metadata should retain source-item boundaries so each get_num_tokens_per_<modality> invocation receives only the metadata belonging to that item. The contract should be modality-generic rather than adding another Qwen-specific argument.

Actual behavior

Flattened multi-row video metadata loses item ownership. Current code drops it and recomputes from raw inputs, which may produce token lengths inconsistent with the already processed tensors.

Active adjacent work

I found three active PRs touching this area and do not want to overlap them:

• [None][feat] Enforce multimodal encoder runtime budgets with budgeted output storage #16051 introduces typed multimodal encoder item metadata for runtime scheduling;
• [None][feat] Generic Mixed Modality Support #16337 introduces a prompt-order mixed-modality manifest;
• [Model]Support MiniCPM-V 4.5 #15766 adds model-specific temporal_ids routing in find_mm_token_lengths.

Their current patches do not partition flattened multi-row video_grid_thw by source item or remove this fallback, but one of their metadata contracts may be the right foundation. I will not implement until maintainers confirm whether this should be handled by one of those PRs or as an independent follow-up.

Proposed contribution scope

If maintainers approve, I would like to contribute a narrow, API-compatible change that:

1. defines a generic per-item processed-metadata representation or adapter at the input-processor boundary;
2. preserves the existing single-video and one-row-per-video behavior;
3. covers variable rows per item, malformed boundary metadata, image/audio controls, and deterministic item ordering;
4. adds CPU-only synthetic regression tests and leaves full-model GPU validation clearly separated.

Could a TensorRT-LLM engineer confirm the preferred metadata shape and whether this work is available for an external contribution?