KV Cache quantization is not working with Whisper

System Info

• TensorRT-LLM 0.17.0.post1
• ModelOpt: https://github.com/MahmoudAshraf97/TensorRT-Model-Optimizer
  A fork with added support for quantizing and exporting TRT checkpoints for whisper

Who can help?

@Tracin

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

from modelopt.torch.export import export_tensorrt_llm_checkpoint
import torch
import modelopt.torch.quantization as mtq
from transformers import AutoModelForSpeechSeq2Seq
from accelerate.hooks import remove_hook_from_module

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    "openai/whisper-tiny",
    use_safetensors=True,
)

EMPTY_CFG = {
    "quant_cfg": {
        "*weight_quantizer": {"enable": False},
        "*input_quantizer": {"enable": False},
        "*lm_head*": {"enable": False},
        "*output_layer*": {"enable": False},
        "default": {"enable": False},
    },
    "algorithm": "max",
}
QUANT_CFG_CHOICES = {
    "int8_sq": mtq.INT8_SMOOTHQUANT_CFG,
    "fp8": mtq.FP8_DEFAULT_CFG,
    "fp8_pc": mtq.FP8_PER_CHANNEL_CFG,
    None: EMPTY_CFG,
}
calib_dataloader = None  # implement any dataloader here


def calibrate_loop():
    if calib_dataloader is None:
        return
    with torch.no_grad():
        for idx, data in enumerate(calib_dataloader):
            model(**data)


quant_dtype = "fp8"
kv_cache_dtype = "fp8"
quant_cfg = QUANT_CFG_CHOICES[quant_dtype]
quant_cfg["quant_cfg"]["*proj_out*"] = {"enable": False}
if kv_cache_dtype is not None:
    if kv_cache_dtype == "fp8":
        for value in KV_CACHE_CFG.values():
            value.update({"num_bits": (4, 3)})  # type: ignore
    quant_cfg["quant_cfg"].update(KV_CACHE_CFG)  # type: ignore


mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
remove_hook_from_module(model, recurse=True)

with torch.inference_mode():
    export_tensorrt_llm_checkpoint(
        model,
        "whisper",
        torch.float16,
        export_dir="checkpoint_dir",
        inference_tensor_parallel=1,
        inference_pipeline_parallel=1,
    )

and then build the engine using

trtllm-build  --checkpoint_dir checkpoint_dir/encoder \
              --output_dir engine_dir/encoder \
              --kv_cache_type paged \
              --moe_plugin disable \
              --max_batch_size 512 \
              --gemm_plugin disable \
              --bert_attention_plugin auto \
              --max_input_len 3000 \
              --max_seq_len 3000 \
              --gather_all_token_logits

trtllm-build  --checkpoint_dir checkpoint_dir/decoder \
              --output_dir engine_dir/decoder \
              --kv_cache_type paged \
              --moe_plugin disable \
              --max_beam_width 5 \
              --max_batch_size 512 \
              --max_seq_len 214 \
              --max_input_len 50 \
              --max_encoder_input_len 3000 \
              --gemm_plugin disable \
              --bert_attention_plugin auto \
              --gpt_attention_plugin auto \
              --gather_all_token_logits

Expected behavior

To infer successfully

actual behavior

The engine can be loaded using the executor API but the first inference request throws this error:

[TensorRT-LLM][ERROR] Encountered an error in forwardSync function: [TensorRT-LLM][ERROR] CUDA runtime error in ::cudaEventSynchronize(get()): an illegal memory access was encountered (/home/jenkins/agent/workspace/LLM/release-0.17/L0_Test-x86_64/tensorrt_llm/cpp/include/tensorrt_llm/runtime/cudaEvent.h:66)
1       0x7f341909fc4c void tensorrt_llm::common::check<cudaError>(cudaError, char const*, char const*, int) + 92
2       0x7f3419281d8b tensorrt_llm::runtime::GptDecoderBatched::forwardSync(tensorrt_llm::runtime::decoder_batch::DecoderFinishedEvent const&) + 59
3       0x7f3419bc8749 tensorrt_llm::batch_manager::TrtGptModelInflightBatching::decoderSync(tensorrt_llm::batch_manager::ScheduledRequests const&, std::unique_ptr<tensorrt_llm::runtime::decoder_batch::DecoderFinishedEvent const, std::default_delete<tensorrt_llm::runtime::decoder_batch::DecoderFinishedEvent const> > const&) + 617
4       0x7f3419bc8ac4 tensorrt_llm::batch_manager::TrtGptModelInflightBatching::forwardSync() + 596
5       0x7f3419c59d26 tensorrt_llm::executor::Executor::Impl::forwardSync(std::list<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest>, std::allocator<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest> > >&) + 54
6       0x7f3419c659b7 tensorrt_llm::executor::Executor::Impl::executionLoop() + 439
7       0x7f36c4f0d5c0 /root/mahmoud/mahmoud/lib/python3.10/site-packages/torch/lib/libtorch.so(+0x145c0) [0x7f36c4f0d5c0]
8       0x7f36c50a81c4 /lib/x86_64-linux-gnu/libc.so.6(+0x891c4) [0x7f36c50a81c4]
9       0x7f36c512885c /lib/x86_64-linux-gnu/libc.so.6(+0x10985c) [0x7f36c512885c]
[TensorRT-LLM][ERROR] Encountered an error in forwardAsync function: [TensorRT-LLM][ERROR] CUDA runtime error in ::cudaMallocAsync(ptr, n, mMemPool->getPool(), mCudaStream->get()): an illegal memory access was encountered (/home/jenkins/agent/workspace/LLM/release-0.17/L0_Test-x86_64/tensorrt_llm/cpp/tensorrt_llm/runtime/tllmBuffers.h:125)
1       0x7f341909fc4c void tensorrt_llm::common::check<cudaError>(cudaError, char const*, char const*, int) + 92
2       0x7f341922c406 virtual thunk to tensorrt_llm::runtime::GenericTensor<tensorrt_llm::runtime::CudaAllocatorAsync>::reshape(nvinfer1::Dims64 const&) + 422
3       0x7f3419b38675 tensorrt_llm::batch_manager::EncoderBuffers::updateBufferSizes(std::vector<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest>, std::allocator<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest> > > const&, tensorrt_llm::runtime::ModelConfig const&, tensorrt_llm::runtime::WorldConfig const&, tensorrt_llm::runtime::TllmRuntime const&) + 885
4       0x7f3419b3c707 tensorrt_llm::batch_manager::EncoderBuffers::prepareIO(std::vector<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest>, std::allocator<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest> > > const&, tensorrt_llm::runtime::ModelConfig const&, tensorrt_llm::runtime::WorldConfig const&, tensorrt_llm::runtime::TllmRuntime const&) + 55
5       0x7f3419baf3a5 tensorrt_llm::batch_manager::TrtEncoderModel::executeBatch(tensorrt_llm::batch_manager::ScheduledRequests const&) + 85
6       0x7f3419bb3a12 tensorrt_llm::batch_manager::TrtEncoderModel::forwardAsync(std::list<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest>, std::allocator<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest> > > const&) + 1074
7       0x7f3419c59e58 tensorrt_llm::executor::Executor::Impl::forwardAsync(std::list<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest>, std::allocator<std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest> > >&) + 136
8       0x7f3419c65cb6 tensorrt_llm::executor::Executor::Impl::executionLoop() + 1206
9       0x7f36c4f0d5c0 /root/mahmoud/mahmoud/lib/python3.10/site-packages/torch/lib/libtorch.so(+0x145c0) [0x7f36c4f0d5c0]
10      0x7f36c50a81c4 /lib/x86_64-linux-gnu/libc.so.6(+0x891c4) [0x7f36c50a81c4]
11      0x7f36c512885c /lib/x86_64-linux-gnu/libc.so.6(+0x10985c) [0x7f36c512885c]

additional notes

FP8 quantization without KV cache quantization works fine, quantizing KV Cache to fp8 or int8 throws this error in both FP8 and FP16 inference

Start with the provided Whisper quantization script, especially KV_CACHE_CFG and export_tensorrt_llm_checkpoint, then build the encoder and decoder with the shown trtllm-build commands. Reproduce through the executor API with FP8 or INT8 KV-cache quantization and compare against FP8 quantization without KV-cache quantization. Done means the first inference request completes without the reported illegal memory access in both FP8 and FP16 inference.