From e67d50841e08ffa6ac23c26ed14dfe8659124b5e Mon Sep 17 00:00:00 2001 From: SWivid Date: Fri, 24 Oct 2025 09:12:08 +0000 Subject: [PATCH] runtime trtllm: fix batch inference skipping last words in shorter sentences #1039 #1179 --- .../triton_trtllm/patch/f5tts/model.py | 35 ++++++++++++-- .../triton_trtllm/patch/f5tts/modules.py | 47 +++++++++---------- src/f5_tts/runtime/triton_trtllm/run.sh | 2 +- 3 files changed, 54 insertions(+), 30 deletions(-) diff --git a/src/f5_tts/runtime/triton_trtllm/patch/f5tts/model.py b/src/f5_tts/runtime/triton_trtllm/patch/f5tts/model.py index baaef95..03132f8 100644 --- a/src/f5_tts/runtime/triton_trtllm/patch/f5tts/model.py +++ b/src/f5_tts/runtime/triton_trtllm/patch/f5tts/model.py @@ -4,11 +4,20 @@ import os 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 @@ class InputEmbedding(Module): 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 @@ class F5TTS(PretrainedModel): 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 diff --git a/src/f5_tts/runtime/triton_trtllm/patch/f5tts/modules.py b/src/f5_tts/runtime/triton_trtllm/patch/f5tts/modules.py index f79af77..312e508 100644 --- a/src/f5_tts/runtime/triton_trtllm/patch/f5tts/modules.py +++ b/src/f5_tts/runtime/triton_trtllm/patch/f5tts/modules.py @@ -16,7 +16,6 @@ from ...functional import ( chunk, concat, constant, - expand, expand_dims, expand_dims_like, expand_mask, @@ -95,15 +94,24 @@ class ConvPositionEmbedding(Module): 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 @@ class Attention(Module): rope_cos, rope_sin, input_lengths, + mask=None, c=None, # context c scale=1.0, rope=None, @@ -283,6 +292,7 @@ class AttnProcessor: 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 @@ class AttnProcessor: 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 @@ class DiTBlock(Module): 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 diff --git a/src/f5_tts/runtime/triton_trtllm/run.sh b/src/f5_tts/runtime/triton_trtllm/run.sh index 4be94ef..3c84d8f 100644 --- a/src/f5_tts/runtime/triton_trtllm/run.sh +++ b/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}