[Usage]: Does Medusa support multi batch multimolal model (like intervl3) inference ?

System Info

System Information:

• OS: linux ubuntu22.04
• Python version: 3.12.3
• CUDA version:12.9
• GPU model(s): L4
• Driver version: 575.64.03
• TensorRT-LLM version: 0.16.0

Detailed output:

Paste the output of the above commands here

How would you like to use TensorRT-LLM

I want to run inference of a Intervl3 1B with medusa heads. However, when it comes to batch size more than 1， it wiill occur error like this:
[TRT] [E] IExecutionContext::enqueueV3: Error Code 7: Internal Error (GenericMedusaForCausalLM/transformer/vocab_embedding/__add___L322/elementwise_binary_L2890/ELEMENTWISE_SUM_0: dimensions not compatible for elementwise. Broadcast has incompatible dimensions: 2 != 128 && 2 != 1 && 128 != 1. Instruction: CHECK_BROADCAST 2 128.)
Traceback (most recent call last):
output_dict = self.decoder.decode(
^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/runtime/generation.py", line 1168, in wrapper
ret = func(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/runtime/generation.py", line 4271, in decode
return self.decode_regular(
^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/runtime/generation.py", line 3809, in decode_regular
should_stop, next_step_tensors, tasks, context_lengths, host_context_lengths, attention_mask, context_logits, generation_logits, encoder_input_lengths = self.handle_per_step(
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/runtime/generation.py", line 3388, in handle_per_step
raise RuntimeError(f"Executing TRT engine failed step={step}!")
RuntimeError: Executing TRT engine failed step=1!

My core inference code like this：
self.decoder = ensorrt_llm.runtime.GenerationSession
self.decoder.setup(
batch_size=batch_size,
max_context_length=max_input_length,
max_new_tokens=self.generation_config["max_new_tokens"],
beam_width=self.generation_config["num_beams"],
medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]
)

output_dict = self.decoder.decode(
batch_ids,
input_lengths,
self.sampling_config,
prompt_table,
tasks,
task_vocab_size,
output_sequence_lengths=True,
return_dict=True
)
batch_ids shape 1N
input_lengths: B,
prompt_table: Ndimension_size
tasks: B
taks_vocab_size: (1,)

My prompt_table generate process like this:

def ptuning_setup(self, prompt_table, batch_size, image_patches=None):
"""
prompt_table: [B, num_patch, C]
input_ids: [B, T]
image_patches: Optional[Tensor] # shape: [B], 每个样本 块数（int32），每张图片切成的不一样

      返回：
           - prompt_table_flattened: [max(image_patches)*num_patch, C]  不足的补0
           - tasks: [B]
           - task_vocab_size: [B]
      """
      B, P, C = prompt_table.shape   # B=batch, P=patch数, C=embedding dim 值得注意的是，一张图片，B!=1
      prompt_table = prompt_table.view(B * P, C)  # flatten to [B*P, C]

      if image_patches is None or batch_size == 1:
           # 单图（或假设每个样本 patch 数相同）
           task_vocab_size = torch.tensor([B * P], dtype=torch.int32, device="cuda")
           tasks = torch.zeros([batch_size], dtype=torch.int32).cuda()
      else:
           max_patch = torch.max(image_patches).item()
           min_patch = torch.min(image_patches).item()
           task_vocab_size = torch.tensor([max_patch * P], dtype=torch.int32, device="cuda")  # shape: [B]
           if max_patch != min_patch:
                table = torch.zeros((max_patch * P * batch_size, C), dtype=torch.float16, device="cuda")
                cum = 0
                # 对patch数不同的样本补0
                for i, image_patch in enumerate(image_patches):
                        n_patch = int(image_patch.item())
                        start = i * (max_patch * P)
                        end = start + n_patch * P
                        table[start:end] = prompt_table[cum:cum + image_patch * P]
                        cum = cum + n_patch * P
                prompt_table = table
           
           # 多图：每个样本 patch 数量不同
           tasks = torch.arange(0, batch_size, dtype=torch.int32).cuda()
      return prompt_table, tasks, task_vocab_size

It goes well when without medusa heads。 I don't k
now how to tackle this.

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.

Start in tensorrt_llm/runtime/generation.py at GenerationSession.decode_regular and handle_per_step, using the reported batch-size-greater-than-one Medusa failure as the reproduction case. Compare the multimodal prompt-table, tasks, and task_vocab_size shapes for batch size 1 versus larger batches, then establish whether multi-batch InterVL3 Medusa inference is supported and document or validate the resulting behavior.