runtime trtllm: fix batch inference skipping last words in shorter sentences #1039 #1179

Author: SWivid

src/f5_tts/runtime/triton_trtllm/patch/f5tts/model.py

@@ -4,11 +4,20 @@
 import sys
 from collections import OrderedDict
 
+import numpy as np
 import tensorrt as trt
 from tensorrt_llm._common import default_net
 
 from ..._utils import str_dtype_to_trt
-from ...functional import Tensor, concat
+from ...functional import (
+    Tensor,
+    concat,
+    constant,
+    expand,
+    shape,
+    slice,
+    unsqueeze,
+)
 from ...layers import Linear
 from ...module import Module, ModuleList
 from ...plugin import current_all_reduce_helper
@@ -27,9 +36,9 @@ def __init__(self, mel_dim, text_dim, out_dim):
         self.proj = Linear(mel_dim * 2 + text_dim, out_dim)
         self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
 
-    def forward(self, x, cond):
+    def forward(self, x, cond, mask=None):
         x = self.proj(concat([x, cond], dim=-1))
-        return self.conv_pos_embed(x) + x
+        return self.conv_pos_embed(x, mask=mask) + x
 
 
 class F5TTS(PretrainedModel):
@@ -69,10 +78,26 @@ def forward(
         input_lengths,
         scale=1.0,
     ):
+        if default_net().plugin_config.remove_input_padding:
+            mask = None
+        else:
+            N = shape(noise, 1)
+            B = shape(noise, 0)
+            seq_len_2d = concat([1, N])
+            max_position_embeddings = 4096
+            # create position ids
+            position_ids_buffer = constant(np.expand_dims(np.arange(max_position_embeddings).astype(np.int32), 0))
+            tmp_position_ids = slice(position_ids_buffer, starts=[0, 0], sizes=seq_len_2d)
+            tmp_position_ids = expand(tmp_position_ids, concat([B, N]))  # [B, N]
+            tmp_input_lengths = unsqueeze(input_lengths, 1)  # [B, 1]
+            tmp_input_lengths = expand(tmp_input_lengths, concat([B, N]))  # [B, N]
+            mask = tmp_position_ids < tmp_input_lengths  # [B, N]
+            mask = mask.cast("int32")
+
         t = self.time_embed(time)
-        x = self.input_embed(noise, cond)
+        x = self.input_embed(noise, cond, mask=mask)
         for block in self.transformer_blocks:
-            x = block(x, t, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale)
+            x = block(x, t, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale, mask=mask)
         denoise = self.proj_out(self.norm_out(x, t))
         denoise.mark_output("denoised", self.dtype)
         return denoise

src/f5_tts/runtime/triton_trtllm/patch/f5tts/modules.py

@@ -16,7 +16,6 @@
     chunk,
     concat,
     constant,
-    expand,
     expand_dims,
     expand_dims_like,
     expand_mask,
@@ -95,15 +94,24 @@ def __init__(self, dim, kernel_size=31, groups=16):
         self.conv1d2 = Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2)
         self.mish = Mish()
 
-    def forward(self, x, mask=None):  # noqa: F722
+    def forward(self, x, mask=None):
         if default_net().plugin_config.remove_input_padding:
             x = unsqueeze(x, 0)
-        x = permute(x, [0, 2, 1])
-        x = self.mish(self.conv1d2(self.mish(self.conv1d1(x))))
-        out = permute(x, [0, 2, 1])
+        if mask is not None:
+            mask = mask.view(concat([shape(mask, 0), 1, shape(mask, 1)]))  # [B 1 N]
+            mask = expand_dims_like(mask, x)  # [B D N]
+            mask = cast(mask, x.dtype)
+        x = permute(x, [0, 2, 1])  # [B D N]
+
+        if mask is not None:
+            x = self.mish(self.conv1d2(self.mish(self.conv1d1(x * mask) * mask)) * mask)
+        else:
+            x = self.mish(self.conv1d2(self.mish(self.conv1d1(x))))
+
+        x = permute(x, [0, 2, 1])  # [B N D]
         if default_net().plugin_config.remove_input_padding:
-            out = squeeze(out, 0)
-        return out
+            x = squeeze(x, 0)
+        return x
 
 
 class Attention(Module):
@@ -185,6 +193,7 @@ def forward(
         rope_cos,
         rope_sin,
         input_lengths,
+        mask=None,
         c=None,  # context c
         scale=1.0,
         rope=None,
@@ -283,6 +292,7 @@ def __call__(
         input_lengths,
         scale=1.0,
         rope=None,
+        mask=None,
     ) -> torch.FloatTensor:
         query = attn.to_q(x)
         key = attn.to_k(x)
@@ -295,20 +305,8 @@ def __call__(
         inner_dim = key.shape[-1]
         norm_factor = math.sqrt(attn.attention_head_size)
         q_scaling = 1.0 / norm_factor
-        mask = None
-        if not default_net().plugin_config.remove_input_padding:
-            N = shape(x, 1)
-            B = shape(x, 0)
-            seq_len_2d = concat([1, N])
-            max_position_embeddings = 4096
-            # create position ids
-            position_ids_buffer = constant(np.expand_dims(np.arange(max_position_embeddings).astype(np.int32), 0))
-            tmp_position_ids = slice(position_ids_buffer, starts=[0, 0], sizes=seq_len_2d)
-            tmp_position_ids = expand(tmp_position_ids, concat([B, N]))  # BxL
-            tmp_input_lengths = unsqueeze(input_lengths, 1)  # Bx1
-            tmp_input_lengths = expand(tmp_input_lengths, concat([B, N]))  # BxL
-            mask = tmp_position_ids < tmp_input_lengths  # BxL
-            mask = mask.cast("int32")
+        if default_net().plugin_config.remove_input_padding:
+            mask = None
 
         if default_net().plugin_config.bert_attention_plugin:
             qkv = concat([query, key, value], dim=-1)
@@ -393,14 +391,15 @@ def __init__(self, dim, heads, dim_head, ff_mult=2, dropout=0.1, pe_attn_head=No
         self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout)
 
     def forward(
-        self, x, t, rope_cos, rope_sin, input_lengths, scale=1.0, rope=ModuleNotFoundError
+        self, x, t, rope_cos, rope_sin, input_lengths, scale=1.0, rope=ModuleNotFoundError, mask=None
     ):  # x: noised input, t: time embedding
         # pre-norm & modulation for attention input
         norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
         # attention
         # norm ----> (2,1226,1024)
-        attn_output = self.attn(x=norm, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale)
-
+        attn_output = self.attn(
+            x=norm, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale, mask=mask
+        )
         # process attention output for input x
         if default_net().plugin_config.remove_input_padding:
             x = x + gate_msa * attn_output

src/f5_tts/runtime/triton_trtllm/run.sh

@@ -73,7 +73,7 @@ fi
 
 if [ $stage -le 7 ] && [ $stop_stage -ge 7 ]; then
     echo "TRT-LLM: offline decoding benchmark test"
-    batch_size=1
+    batch_size=2
     split_name=wenetspeech4tts
     backend_type=trt
     log_dir=./tests/benchmark_${model}_batch_size_${batch_size}_${split_name}_${backend_type}