[Usage]: Pytorch Workflow Add Multimodal Embedding into Qwen3

[bnuzhanyu]

System Info

System Information:

• OS: Ubuntu
• Python version: 3.12
• CUDA version: 12.8
• GPU model(s): 4090
• Driver version: 570.124.04
• TensorRT-LLM version: 1.1.0rc1

Detailed output:

Python 3.12.3
Name: tensorrt_llm
Version: 1.1.0rc1
Summary: TensorRT-LLM: A TensorRT Toolbox for Large Language Models
Home-page: https://github.com/NVIDIA/TensorRT-LLM
Author: NVIDIA Corporation
Author-email: 
License: Apache License 2.0
Location: /usr/local/lib/python3.12/dist-packages
Requires: accelerate, aenum, backoff, blake3, blobfile, build, click, click_option_group, colored, cuda-python, datasets, diffusers, einops, etcd3, evaluate, fastapi, flashinfer-python, h5py, jsonschema, lark, llguidance, matplotlib, meson, mpi4py, mpmath, ninja, numpy, nvidia-cuda-nvrtc-cu12, nvidia-ml-py, nvidia-modelopt, nvidia-nccl-cu12, nvtx, omegaconf, onnx, onnx_graphsurgeon, openai, opencv-python-headless, optimum, ordered-set, pandas, peft, pillow, polygraphy, prometheus_client, prometheus_fastapi_instrumentator, protobuf, psutil, pulp, pydantic, pydantic-settings, pynvml, pyzmq, sentencepiece, setuptools, soundfile, StrEnum, tensorrt, tiktoken, torch, torchvision, transformers, triton, uvicorn, wheel, xgrammar

How would you like to use TensorRT-LLM

Question 1

I want to add multimodal embeddings to Qwen3, so I intend to use the data from the multimodal part when the input_id is greater than or equal to the vocab_size. However, when there are ids in my input_ids that are larger than the vocab_size, an error occurs somewhere (I haven't found the cause of the error, and it seems that it hasn't even entered my forward function).

我想在qwen3上加上多模emb，所以想沿用之前的input_id>=vocab_size就取multimodal里的数据。
但当我的input_ids里有大于vocab_size的id后，就会在某个地方报错(我没有找到错误的原因，好像还没有进入我的forward函数）。

[08/26/2025-14:10:33] [TRT-LLM] [E] Error in thread proxy_dispatch_result_thread: 2
Traceback (most recent call last):
File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/llmapi/utils.py", line 268, in run
if not task(**self.kwargs):
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/executor/proxy.py", line 191, in dispatch_result_task
process_res(i)
File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/executor/proxy.py", line 172, in process_res
queue = self._results[client_id].queue
~~~~~~~~~~~~~^^^^^^^^^^^
KeyError: 2

It seems that I could not do any operation to the input_ids(even copy).

class MyQwen3Model(Qwen3Model):

    def __init__(self, model_config: ModelConfig[Qwen3Config]):
        super().__init__(model_config)
        logger.info(f"MyQwen3Model init..........................")

    def forward(
        self,
        attn_metadata: AttentionMetadata,
        input_ids: Optional[torch.IntTensor] = None,
        position_ids: Optional[torch.IntTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        mrope_config: Optional[Tuple[torch.Tensor, int]] = None,
        spec_metadata: Optional[SpecMetadata] = None,
        **kwargs,
    ) -> torch.Tensor:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError(
                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
            )
        logger.info("forward start")
        #TODO: replace input_embeds to mmodal embedding

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        
        hidden_states = inputs_embeds
        if 'multimodal_params' in kwargs:
            multimodal_params: List[MultimodalParams] = kwargs['multimodal_params']
            logger.info(f"multimodal_params: {len(multimodal_params)}")
            for mm_param in multimodal_params:
                logger.info(f"mm_param: {mm_param.multimodal_data}")
                mm_param.multimodal_data = None


        residual = None
        for decoder_layer in self.layers:
            hidden_states, residual = decoder_layer(
                position_ids=position_ids,
                hidden_states=hidden_states,
                attn_metadata=attn_metadata,
                residual=residual,
                mrope_config=mrope_config,
                spec_metadata=spec_metadata,
            )

        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


@register_auto_model("Qwen3ForCausalLM")
class MyQwen3ForCausalLM(SpecDecOneEngineForCausalLM[MyQwen3Model, Qwen3Config]):

    def __init__(
        self,
        model_config: ModelConfig[Qwen3Config],
    ):
        super().__init__(
            MyQwen3Model(model_config),
            model_config,
        )

Question 2

One problem I want to solve is that after using multimodality, the same token_id may have different embeddings. However, when enable_block_reuse is used, these ids may be cached, resulting in different embeddings being unable to be passed to this token_id again. But in our scenario, multimodal tokens are few in number and not in very early positions.

Is there a solution that can both use enable_block_reuse and prevent certain token_ids(and following input_ids) from being cached?

我想解决的一个问题是，使用多模后，相同的token_id可能拥有不同的embedding，但使用enable_block_reuse后，这些id可能会被cache，导致不同的embedding无法被再次传给这个token_id。但我们场景多模token是少数，而且也不是在很靠前的位置。

有没有既能使用 enable_block_reuse，又能让某些token_id之后不被cache的方案？

Question 3

It seems that in the _prepare_tp_inputs method of PyTorchModelEngine, there might be an issue where the multimodal_params_list array does not correctly correspond to the requests. This makes it difficult to distinguish which multimodal parameter (mmodal) belongs to which request in the model's forward function.

PyTorchModelEngine里_prepare_tp_inputs里，感觉会导致multimodal_params_list数组和request对应不上，在模型的forward里，不太好区分某个mmodal是对应哪个request的。（虽然可以通过input_ids长度是否为1判断是prefill还是decode，但也不是所有的prefill都会有多模信息）。

 for request in scheduled_requests.context_requests:
            request_ids.append(request.py_request_id)
            all_prompt_tokens = request.get_tokens(0)
            draft_lens.append(0)
            begin_compute = request.context_current_position
            end_compute = begin_compute + request.context_chunk_size
            prompt_tokens = all_prompt_tokens[begin_compute:end_compute]
            position_ids.extend(
                range(begin_compute, begin_compute + len(prompt_tokens)))
            input_ids.extend(prompt_tokens)
            gather_ids.append(len(input_ids) - 1)
            sequence_lengths.append(len(prompt_tokens))
            prompt_lengths.append(len(prompt_tokens))
            past_seen_token_num = begin_compute
            num_cached_tokens_per_seq.append(past_seen_token_num)

            # Multimodal
            # TODO: enable chunk prefill for multimodal (maybe need to pass prompt_tokens to MultimodalRuntimeData)
            py_multimodal_runtime = MultimodalRuntimeData(
                mm_token_lengths=request.multimodal_lengths,
                mm_token_positions=request.multimodal_positions,
                num_cached_tokens=past_seen_token_num
            ) if request.multimodal_hashes is not None else None

            multimodal_params = MultimodalParams(
                multimodal_data=request.py_multimodal_data,
                multimodal_runtime=py_multimodal_runtime)

            if multimodal_params.has_content():
                multimodal_params.to_device("multimodal_data",
                                            "cuda",
                                            pin_memory=True)
                #re-assign the multimodal_data to the request after to_device for generation requests
                request.py_multimodal_data = multimodal_params.multimodal_data
                multimodal_params_list.append(multimodal_params)

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[onestardao]

sorry to jump in, maintainers.
from the logs and code, this looks less like a TRT-LLM bug and more like a schema / param mismatch around input_ids vs inputs_embeds with block reuse turned on. the KV cache is keyed off token ids from the first bind, then you swap to embeds or change vocab space and the cache handle keeps pointing to the old key. in our map this is ProblemMap No.2 (schema-param mixup / key mismatch), and sometimes pairs with No.14 (first forward binds cache before multimodal params are registered).

quick triage to confirm: run one path only (ids or embeds) with block reuse on, then switch and clear KV before the next run; verify tokenizer vocab size matches the embed matrix; log that the cache key includes a modality flag. if helpful, i can share the short checklist we use for No.2/No.14.

[nv-guomingz]

Hi @yechank-nvidia ,could you please take a look this question?

[yechank-nvidia]

Hi, thanks for the question! I will answer them in order.

A1. We suggest you to do register_input_processor on the modeling side and adjust the input_ids there before getting into forward() call. In the given token_ids, insert vision_token_ids(which normally done by AutoProcessor), replace the vision_ids with ids that are larger than the vocab_size and hand over to the LLM. You can refer to this.

A2. In order to enable block reuse, we need some extra metadata from the inputprocessor by calling get_num_tokens_per_image(Ref). If the inputprocessor have this function, it will try to enable block_reuse for multimodal. You can see them more in detail in here.

A3. In _prepare_tp_inputs , we don't tie the request_id per multimodal_params, rather we just append them in order. This design is introduced because there might be a case of doing mixture of text-only prompts and multimodal prompts. On the text-only prompts, there will be no multimodal_data found on the request, so that len(requests) != len(multimodal_params) in all cases. Also, AFAIK, the ordering is correct. If you find any counter-cases, can you share them, please?

If you have any further questions, feel free to reply.

[bnuzhanyu]

Hi, thanks for the question! I will answer them in order.

A1. We suggest you to do register_input_processor on the modeling side and adjust the input_ids there before getting into forward() call. In the given token_ids, insert vision_token_ids(which normally done by AutoProcessor), replace the vision_ids with ids that are larger than the vocab_size and hand over to the LLM. You can refer to this.

A2. In order to enable block reuse, we need some extra metadata from the inputprocessor by calling get_num_tokens_per_image(Ref). If the inputprocessor have this function, it will try to enable block_reuse for multimodal. You can see them more in detail in here.

A3. In _prepare_tp_inputs , we don't tie the request_id per multimodal_params, rather we just append them in order. This design is introduced because there might be a case of doing mixture of text-only prompts and multimodal prompts. On the text-only prompts, there will be no multimodal_data found on the request, so that len(requests) != len(multimodal_params) in all cases. Also, AFAIK, the ordering is correct. If you find any counter-cases, can you share them, please?

If you have any further questions, feel free to reply.

Thanks for your answer, but I still have some questions.

For A1

I found the _postprocess also make self.tllm_multimodal_token_id (== vocab_size+1) into the input_ids, does it mean you allow vocab_size+1 but not allow arbitrary larger number within int32?
Which I really need is fuse_input_embeds in forward function, so I have to pass some id larger than vocab_size which fails before forward.

For A2

I'm not sure how get_num_tokens_per_image related to the question.
Let me simplify this:
if the input_ids is [1, 2, 3, 4, img_token_id, 5, 6], I hope it can only cache [1,2,3,4] (where img_token_id is fixed).

For A3

If the ordering is correct, then it will not be a question.

[bnuzhanyu]

@yechank-nvidia

I think there is no relation with the processor.
I fixed the mm_token_id and use fuse_input_embeds to combine the mmodal and input_id's embedding.

In my case, Model.num_embeddings = 160128.
When mm_token_id < 160128, it will run ok, when mm_token_id >= 160128, it will crash.
So we can not pass a value larger than model.num_embeddings, and this may be a bug?

The input is built by this code:

    token_ids = tokenizer.encode(prompt) 
    token_ids[100] = 160120  # error if  >=160128
    prompt_inputs = {
        'prompt_token_ids': token_ids,
        'multi_modal_data': {
            'image': torch.ones(1, 2048)
        }
    }

class MyQwen3Model(Qwen3Model):

    def __init__(self, model_config: ModelConfig[Qwen3Config]):
        super().__init__(model_config)
        logger.info(f"MyQwen3Model init, {model_config.pretrained_config.vocab_size}..........................")
        self.mm_token_id = torch.tensor(160120, dtype=torch.int32, device=torch.cuda.current_device())

    def forward(
        self,
        attn_metadata: AttentionMetadata,
        input_ids: Optional[torch.IntTensor] = None,
        position_ids: Optional[torch.IntTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        mrope_config: Optional[Tuple[torch.Tensor, int]] = None,
        spec_metadata: Optional[SpecMetadata] = None,
        **kwargs,
    ) -> torch.Tensor:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError(
                "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
            )

        if 'multimodal_params' in kwargs:
            multimodal_params: List[MultimodalParams] = kwargs['multimodal_params']
            logger.info(f"multimodal_params: {len(multimodal_params)}, num_embeddings, {self.embed_tokens.num_embeddings}")
            logger.info(f"input_ids large than num_embeddings: {sum(input_ids >= self.mm_token_id)}")
            mm_embeds = [mm_param.multimodal_data['image'] for mm_param in multimodal_params]
            _, inputs_embeds = fuse_input_embeds(self.embed_tokens,
                                                    input_ids, mm_embeds, self.mm_token_id)
                                                
         ...

(Of course llm.generate_async has a little fix to pass the mmodal_data.)

[chang-l]

When mm_token_id < 160128, it will run ok, when mm_token_id >= 160128, it will crash.
So we can not pass a value larger than model.num_embeddings, and this may be a bug?

@bnuzhanyu were you seeing a crash around here? If so, you may need to adjust the validation to bypass qwen3

[bnuzhanyu]

When mm_token_id < 160128, it will run ok, when mm_token_id >= 160128, it will crash.
So we can not pass a value larger than model.num_embeddings, and this may be a bug?

@bnuzhanyu were you seeing a crash around here? If so, you may need to adjust the validation to bypass qwen3

@chang-l
I did not know the exact position of crash, it was just some cuda assertion. And we decided to use an exist token as the mm-token.