[Bug]: find_input_mm_embeds is order-dependent for mixed full-prefill and partial VLM requests
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
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
mainat573bd5f0b8414566bd6183f74b06f8c8a99e6a69 - 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:
- make the all-or-none/full-prefill decision over every
MultimodalParams, not only index 0; - preserve original request indices in individual-batching mode and apply each request's own slice;
- define or validate the corresponding pre-concatenated contract when a runtime-less request is mixed with runtime-bearing requests;
- add CPU unit regressions for both request orderings plus homogeneous controls;
- 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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at the current find_input_mm_embeds entry point and reproduce the three CPU cases described in the issue, including both request orderings and the homogeneous controls. Confirm the intended mixed-runtime contract with maintainers, then add CPU unit regressions preserving original request indices; done means each request receives only its own embeddings without changing encoder or CUDA paths.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100