Whisper Encoder issues with Executor API

Hello, 0.15.0.dev2024101500 added a new issue when using the executor API with whisper

[TensorRT-LLM][ERROR] IExecutionContext::inferShapes: Error Code 7: Internal Error (WhisperEncoder/__add___L322/elementwise_binary_L2856/ELEMENTWISE_SUM_0: dimensions not compatible for elementwise. Broadcast has incompatible dimensions: 3000 != 1500 && 3000 != 1 && 1500 != 1. Instruction: CHECK_BROADCAST 3000 1500.)
[TensorRT-LLM][ERROR] Encountered an error in forwardAsync function: Invalid input shape (/home/jenkins/agent/workspace/LLM/main/L0_PostMerge/tensorrt_llm/cpp/tensorrt_llm/runtime/tllmRuntime.cpp:313)

This error was introduced after the change in positional embeddings in the last update, when commenting the problematic addition here:
https://github.com/NVIDIA/TensorRT-LLM/blob/75057cd036af25e288c004d8ac9e52fd2d6224aa/tensorrt_llm/models/enc_dec/model.py#L1928

the error is solved, but new error comes up which is also reproducible in 0.14.0.dev2024100800

[TensorRT-LLM][ERROR] Encountered an error in forwardAsync function: slice 1152000 exceeds buffer size 576000

this error occurs in all whisper models, the slice is always double the buffer size no matter what the model size is

Who can help?

@byshiue

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

INFERENCE_PRECISION=float16
WEIGHT_ONLY_PRECISION=int8
MAX_BEAM_WIDTH=2
MAX_BATCH_SIZE=8

trtllm-build  --checkpoint_dir ${checkpoint_dir}/encoder \
              --output_dir ${output_dir}/encoder \
              --kv_cache_type paged \
              --moe_plugin disable \
              --enable_xqa disable \
              --max_batch_size ${MAX_BATCH_SIZE} \
              --gemm_plugin disable \
              --bert_attention_plugin ${INFERENCE_PRECISION} \
              --max_input_len 3000 \
              --max_seq_len 3000 

trtllm-build  --checkpoint_dir ${checkpoint_dir}/decoder \
              --output_dir ${output_dir}/decoder \
              --kv_cache_type paged \
              --moe_plugin disable \
              --enable_xqa enable \
              --max_beam_width ${MAX_BEAM_WIDTH} \
              --max_batch_size ${MAX_BATCH_SIZE} \
              --max_seq_len 114 \
              --max_input_len 14 \
              --max_encoder_input_len 1500 \
              --gemm_plugin ${INFERENCE_PRECISION} \
              --bert_attention_plugin ${INFERENCE_PRECISION} \
              --gpt_attention_plugin ${INFERENCE_PRECISION} 

import tensorrt_llm.bindings.executor as trtllm
import torch

# Create the executor.
executor = trtllm.Executor(
    encoder_model_path=engine_dir + "/encoder/",
    decoder_model_path=engine_dir + "/decoder/",
    model_type=trtllm.ModelType.ENCODER_DECODER,
    executor_config=trtllm.ExecutorConfig(1, gpu_weights_percent=0.5),
)

request = trtllm.Request(
    input_token_ids=[50257, 50258, 50358, 50362],
    max_tokens=100,
    encoder_input_features=torch.randn(3000,80).half().cuda(), # mel features
)

# Enqueue the request.
request_id = executor.enqueue_request(request)

Expected behavior

It should complete the request successfuly

additional notes

The decoder works perfectly with executor API but loading the encoder alone or the encoder and the decoder together is not working