Skip to content

[Bug]: Requests with attached image embeddings reuse KV blocks computed from different embeddings #22117

Open

Description

@Hang-get

System Info

  • TensorRT-LLM main at fec3d9191, PyTorch backend, default KV cache settings (enable_block_reuse=True, KV cache manager V2).
  • Image embeddings passed in by the caller: multi_modal_embeddings in the LLM API, or image_embeds content 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 examples folder (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 lengths

I 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 set enable_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_embeddings says 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 its multimodal_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

  1. added
    MultimodalLabel for issues & PRs regarding Multimodal related objects
    on Nov 11, 2026
  2. added 2 commits that reference this issue on Nov 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    MultimodalLabel for issues & PRs regarding Multimodal related objects

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions