[Bug]: find_input_mm_embeds is order-dependent for mixed full-prefill and partial VLM requests

[HOYALIM]

Description

I found a CPU-reproducible correctness gap in the current individual-batching contract of find_input_mm_embeds and would like to contribute a narrow fix if maintainers approve the direction.

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 VLM embedding fusion with chunked prefill/KV reuse

On current main (573bd5f0b8414566bd6183f74b06f8c8a99e6a69), each context request independently receives MultimodalRuntimeData only when multimodal_embed_mask_cumsum is available. Full-prefill requests are explicitly allowed to keep multimodal_runtime=None, while chunked-prefill or KV-reused requests require runtime metadata. A scheduled batch can therefore contain both states.

find_input_mm_embeds currently makes a batch-wide no-slicing decision from only multimodal_params[0].multimodal_runtime. In individual mode it also enumerates the filtered list of runtime-bearing slices, losing their original request indices. The result is order-dependent:

homogeneous control:
  observed [[11, 12], [22]]
  expected [[11, 12], [22]]

runtime states [None, partial]:
  observed [[30, 31, 32], [40, 41, 42, 43]]
  expected [[30, 31, 32], [42]]

runtime states [partial, None, partial]:
  observed [[51], [62]]
  expected [[51], [60, 61, 62], [72, 73]]

The exact current-source function was extracted from its AST and exercised with CPU tensors. All three cases repeated identically three times. No CUDA or model weights are involved.

This can fuse embeddings outside the current chunk, apply a later request's slice to an earlier request's tensor, or omit a request's embeddings. It also defeats the intended KV-cache/chunked-prefill efficiency in the early-return ordering.

Before implementation, could a maintainer confirm whether heterogeneous runtime states are intended to be supported in one batch and approve a narrow correction?

If approved, I propose to:

1. make the all-or-none/full-prefill decision over every MultimodalParams, not only index 0;
2. preserve original request indices in individual-batching mode and apply each request's own slice;
3. define or validate the corresponding pre-concatenated contract when a runtime-less request is mixed with runtime-bearing requests;
4. add CPU unit regressions for both request orderings plus homogeneous controls;
5. leave encoder kernels and CUDA paths unchanged, with upstream GPU CI covering integration behavior.

I checked current open PRs #16051 (multimodal encoder runtime scheduling) and #16337 (generic mixed-modality support). They touch adjacent files but their current patches do not change find_input_mm_embeds or its mixed-runtime slicing logic. I will not open an implementation PR until maintainers approve the issue and confirm there is no overlap.