[Model Runner V2] Fix mm input embeddings lookup (#36588)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Nick Hill
2026-03-10 00:24:58 -07:00
committed by GitHub
parent 9efc3bdcd6
commit ddbb0d230a

View File

@@ -98,8 +98,11 @@ class DefaultModelState(ModelState):
req_states.prefill_len.np[input_batch.idx_mapping_np],
req_states.num_computed_prefill_tokens[input_batch.idx_mapping_np],
)
# Use unpadded input_ids to match is_mm_embed size (num_tokens).
# input_batch.input_ids may be padded for CUDA graphs.
input_ids_unpadded = input_batch.input_ids[: input_batch.num_tokens]
inputs_embeds = self.encoder_runner.get_inputs_embeds(
input_batch.input_ids, mm_embeds, is_mm_embed
input_ids_unpadded, mm_embeds, is_mm_embed
)
return inputs_embeds[: input_batch.num_tokens_after_padding]