Update TensorRT backend tensorrt_llm==0.9.0

Signed-off-by: makaveli10 <vineet.suryan@collabora.com>

README.md

@@ -115,7 +115,18 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b
   docker run -it --gpus all -p 9090:9090 ghcr.io/collabora/whisperlive-gpu:latest
   ```
 
-  - TensorRT. Follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) in order to setup docker and use TensorRT backend. We provide a pre-built docker image which has TensorRT-LLM built and ready to use.
+  - TensorRT. 
+  ```bash
+  docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt
+
+  # Build tiny.en engine
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
+
+  # Run server with tiny.en
+  python3 run_server.py --port 9090 \
+                        --backend tensorrt \
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
+  ```
 
 - CPU
 ```bash

TensorRT_whisper.md

@@ -1,6 +1,6 @@
-# Whisper-TensorRT
+# WhisperLive-TensorRT
 We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup.
-**Note**: We use [our fork to setup TensorRT](https://github.com/makaveli10/TensorRT-LLM)
+**Note**: We use `tensorrt_llm==0.9.0`
 
 ## Installation
 - Install [docker](https://docs.docker.com/engine/install/)
@@ -12,56 +12,31 @@ git clone https://github.com/collabora/WhisperLive.git
 cd WhisperLive
 ```
 
-- Pull the TensorRT-LLM docker image which we prebuilt for WhisperLive TensorRT backend.
+- Run WhisperLive TensorRT in docker
 ```bash
-docker pull ghcr.io/collabora/whisperbot-base:latest
+docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt:latest
 ```
 
-- Next, we run the docker image and mount WhisperLive repo to the containers `/home` directory.
-```bash
-docker run -it --gpus all --shm-size=8g \
-       --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
-       -p 9090:9090 -v /path/to/WhisperLive:/home/WhisperLive \
-       ghcr.io/collabora/whisperbot-base:latest
-```
-
-- Make sure to test the installation. 
-```bash
-# export ENV=${ENV:-/etc/shinit_v2} 
-# source $ENV
-python -c "import torch; import tensorrt; import tensorrt_llm"
-```
-**NOTE**: Uncomment and update library paths if imports fail.
-
 ## Whisper TensorRT Engine
-- We build `small.en` and `small` multilingual TensorRT engine. The script logs the path of the directory with Whisper TensorRT engine. We need the model_path to run the server.
+- We build `small.en` and `small` multilingual TensorRT engine as examples below. The script logs the path of the directory with Whisper TensorRT engine. We need that model_path to run the server.
 ```bash
 # convert small.en
-bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small.en
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
 
 # convert small multilingual model
-bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
 ```
 
 ## Run WhisperLive Server with TensorRT Backend
 ```bash
-cd /home/WhisperLive
-
-# Install requirements
-apt update && bash scripts/setup.sh
-pip install -r requirements/server.txt
-
-# Required to create mel spectogram
-wget --directory-prefix=assets assets/mel_filters.npz https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
-
 # Run English only model
 python3 run_server.py --port 9090 \
                       --backend tensorrt \
-                      --trt_model_path "path/to/whisper_trt/from/build/step"
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
 
 # Run Multilingual model
 python3 run_server.py --port 9090 \
                       --backend tensorrt \
-                      --trt_model_path "path/to/whisper_trt/from/build/step" \
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small" \
                       --trt_multilingual
 ```

docker/Dockerfile.tensorrt

@@ -0,0 +1,28 @@
+FROM nvidia/cuda:12.1.0-devel-ubuntu22.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && apt-get install -y \
+    python3.10 python3-pip openmpi-bin libopenmpi-dev git wget \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN pip3 install --no-cache-dir -U tensorrt_llm==0.9.0 --extra-index-url https://pypi.nvidia.com
+
+WORKDIR /app
+
+RUN git clone -b v0.9.0 --depth 1 https://github.com/NVIDIA/TensorRT-LLM.git && \
+    mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \
+    rm -rf TensorRT-LLM
+
+COPY assets/ ./assets
+RUN wget -nc -P assets/ https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
+
+COPY scripts/setup.sh ./
+RUN apt update && bash setup.sh && rm setup.sh
+
+COPY requirements/server.txt .
+RUN pip install --no-cache-dir -r server.txt && rm server.txt
+
+COPY whisper_live ./whisper_live
+COPY scripts/build_whisper_tensorrt.sh .
+COPY run_server.py .

scripts/build_whisper_tensorrt.sh

@@ -57,7 +57,7 @@ download_and_build_model() {
     local output_dir="whisper_${model_name//./_}"
     echo "$output_dir"
     echo "Running build script for $model_name with output directory $output_dir"
-    python3 build.py --output_dir "$output_dir" --use_gpt_attention_plugin --use_gemm_plugin --use_bert_attention_plugin --model_name "$model_name"
+    python3 build.py --output_dir "$output_dir" --use_gpt_attention_plugin --use_gemm_plugin  --use_bert_attention_plugin --enable_context_fmha --model_name "$model_name"
     echo "Whisper $model_name TensorRT engine built."
     echo "========================================="
     echo "Model is located at: $(pwd)/$output_dir"

whisper_live/transcriber_tensorrt.py

@@ -35,6 +35,8 @@ class WhisperEncoding:
         with open(config_path, 'r') as f:
             config = json.load(f)
 
+        use_gpt_attention_plugin = config['plugin_config'][
+            'gpt_attention_plugin']
         dtype = config['builder_config']['precision']
         n_mels = config['builder_config']['n_mels']
         num_languages = config['builder_config']['num_languages']
@@ -51,12 +53,21 @@ class WhisperEncoding:
         return session
 
     def get_audio_features(self, mel):
-        inputs = OrderedDict()
-        output_list = []
 
-        inputs.update({'x': mel})
+        input_lengths = torch.tensor(
-        output_list.append(
+            [mel.shape[2] // 2 for _ in range(mel.shape[0])],
-            TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape))
+            dtype=torch.int32,
+            device=mel.device)
+
+        inputs = OrderedDict()
+        inputs['x'] = mel
+        inputs['input_lengths'] = input_lengths
+
+        output_list = [
+            TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape),
+            TensorInfo('input_lengths', str_dtype_to_trt('int32'),
+                       input_lengths.shape)
+        ]
 
         output_info = (self.session).infer_shapes(output_list)
 
@@ -101,6 +112,8 @@ class WhisperDecoding:
             decoder_engine_buffer = f.read()
 
         decoder_model_config = ModelConfig(
+            max_batch_size=self.decoder_config['max_batch_size'],
+            max_beam_width=self.decoder_config['max_beam_width'],
             num_heads=self.decoder_config['num_heads'],
             num_kv_heads=self.decoder_config['num_heads'],
             hidden_size=self.decoder_config['hidden_size'],
@@ -141,6 +154,10 @@ class WhisperDecoding:
                                              device='cuda')
         decoder_max_input_length = torch.max(decoder_input_lengths).item()
 
+        cross_attention_mask = torch.ones(
+            [encoder_outputs.shape[0], 1,
+             encoder_outputs.shape[1]]).int().cuda()
+
         # generation config
         sampling_config = SamplingConfig(end_id=eot_id,
                                          pad_id=eot_id,
@@ -161,6 +178,7 @@ class WhisperDecoding:
             sampling_config,
             encoder_output=encoder_outputs,
             encoder_input_lengths=encoder_input_lengths,
+            cross_attention_mask=cross_attention_mask,
         )
         torch.cuda.synchronize()