Repository navigation
Description
System Info
- TensorRT-LLM
mainatfec3d9191, PyTorch backend, default KV cache settings (enable_block_reuse=True, KV cache manager V2). - Image embeddings passed in by the caller:
multi_modal_embeddingsin the LLM API, orimage_embedscontent parts in trtllm-serve chat completions. Dynamo's TRT-LLM encode/prefill flow uses the same path. - Found while reading the input processors and reproduced on CPU (details below). No GPU run.
Who can help?
No response
Information
- The official example scripts
- My own modified scripts
Tasks
- An officially supported task in the
examplesfolder (such as GLUE/SQuAD, ...) - My own task or dataset (give details below)
Reproduction
Every implementation of _attach_multimodal_embeddings_impl (Qwen2VLInputProcessorBase, used by Qwen2.5-VL, Qwen3-VL and the Cosmos3 processors; LLaVA-NeXT; Llama 4) returns (prompt_token_ids, {"multimodal_data": ...}) without a multimodal_input. These requests therefore carry no multimodal_hashes, multimodal_positions or multimodal_lengths, and KVCacheManagerV2._augment_tokens_for_block_reuse uses the prompt tokens unchanged as reuse keys. Every image token is the same image_token_id, so two requests with the same text and the same number of embedding rows have the same block keys, whatever the embedding values.
When the second request is scheduled, block reuse matches the first request's blocks across the image span. Since #12944, llm.py computes multimodal_embed_mask_cumsum for this path, so _pack_context_multimodal builds MultimodalRuntimeData and the model fuses only the embeddings of the uncached tail. Nothing fails, and the second request is answered from the KV that was computed from the first request's embeddings.
With the Qwen2.5-VL input processor (config, tokenizer and processor of Qwen/Qwen2.5-VL-3B-Instruct, no weights):
import torch
from transformers import AutoConfig, AutoTokenizer
from tensorrt_llm._torch.models.modeling_qwen2vl import Qwen2_5VLInputProcessorBase
from tensorrt_llm.sampling_params import SamplingParams
model = "Qwen/Qwen2.5-VL-3B-Instruct" # config, tokenizer and processor only
proc = Qwen2_5VLInputProcessorBase(model, AutoConfig.from_pretrained(model),
AutoTokenizer.from_pretrained(model))
prompt = proc.tokenizer.apply_chat_template(
[{"role": "user", "content": [{"type": "image"},
{"type": "text", "text": "Describe the image."}]}],
tokenize=False, add_generation_prompt=True)
hidden = proc.config.text_config.hidden_size
def attach(embedding):
return proc.attach_multimodal_embeddings(
{"prompt": prompt}, {"image": [embedding]}, SamplingParams())
ids_a, extra_a = attach(torch.randn(391, hidden, dtype=torch.bfloat16))
ids_b, extra_b = attach(torch.randn(391, hidden, dtype=torch.bfloat16))
print(len(ids_a), ids_a == ids_b) # 416 True
print(extra_a.get("multimodal_input")) # None: no hashes, positions or lengthsI passed the request fields that llm.py builds from these results to KVCacheManagerV2._augment_tokens_for_block_reuse: the two requests get identical keys for all 13 full 32-token blocks of the 416-token prompt (391 of the tokens are image tokens). With Qwen3VLInputProcessorBase (Qwen/Qwen3-VL-2B-Instruct) all 12 full blocks of the 405-token prompt match. This machine has no TensorRT-LLM build, so I ran it on CPU with stubs for the native bindings.
Expected behavior
Attached embeddings are keyed by their content, as media items are keyed by apply_mm_hashes: requests with identical embeddings and prompt can reuse each other's blocks, requests with different embeddings cannot.
actual behavior
Different embeddings with the same shape and the same prompt share KV blocks. The second request is computed from the first request's image KV, which may come from another user, and no error is reported.
additional notes
- feat: support KV reuse in llava raw embedding case #11967 reported an earlier form of this for LLaVA-NeXT raw embeddings: a second identical request crashed in
fuse_input_embeds. Since [TRTLLM-11951][feat] Chunked prefill for non-contiguous multimodal data + preproc fixes #12944 the cached embeddings are sliced off, so these requests no longer fail, including when the cached blocks come from different embeddings. Dynamo's TRT-LLM multimodal engine configs setenable_block_reuse: false(qwen3-vl-2b-instruct prefill/agg, llava-v1.6 agg, llama4 Scout prefill); fix(trtllm): enable KV block reuse for LLaVA raw embeddings ai-dynamo/dynamo#10640 proposed enabling it for LLaVA after [TRTLLM-11951][feat] Chunked prefill for non-contiguous multimodal data + preproc fixes #12944 and tested it by sending the same image twice. MultimodalParams.from_embeddingssays that without hashes only the processed data is kept, "as for external embeddings without KV-cache reuse". The KV cache lender (keyed_by_placeholders) and the Mooncake store connector (_multimodal_item_digests) also treat multimodal requests keyed by placeholder tokens as unsafe to share. The Mooncake check looks for positions or lengths, which this path does not set either.- Proposed fix: hash each attached embedding (dtype, shape and raw bytes, since bfloat16 cannot go through NumPy) and pass the digests to
build_disagg_prefill_multimodal_inputs, which already builds the positions, lengths and runs for the encoder/prefill handoff, then return itsmultimodal_input. I will open a PR for the Qwen VL processors. LLaVA-NeXT (_postprocess) and Llama 4 build their prompts differently and would need their own layout code.
Activity
[NVIDIA#22117][test] Cover hash order and dtype for attached embeddings
[NVIDIA#22117][test] Check each attached image's hash and non-contigu…