feat:添加耗时统计

fish_speech/models/text2semantic/inference.py

@@ -7,7 +7,7 @@ import traceback
 from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Callable, Literal, Optional, Tuple, Union
+from typing import Callable, Literal, Optional, Tuple, Union, Any
 
 import click
 import numpy as np
@@ -243,6 +243,96 @@ def decode_n_tokens(
     return torch.cat(new_tokens, dim=1)
 
 
+def decode_n_tokens_optimized(
+        model: DualARTransformer,
+        cur_token: torch.Tensor,
+        input_pos: torch.Tensor,
+        num_new_tokens: int,
+        temperature: torch.Tensor,
+        top_p: torch.Tensor,
+        top_k: int,
+        semantic_logit_bias: torch.Tensor,
+        audio_masks: torch.Tensor,
+        audio_parts: torch.Tensor,
+        previous_tokens: torch.Tensor,
+        im_end_id: Any,
+        decode_one_token=decode_one_token_ar,
+):
+    """
+    Optimized version:
+    - no roll (ring buffer)
+    - flash attention
+    - reduced view/reshape
+    """
+
+    device = cur_token.device
+    num_streams = model.config.num_codebooks + 1
+
+    # =========================
+    # 1. ring buffer index (替代 roll)
+    # =========================
+    history_len = previous_tokens.size(1)
+    write_idx = history_len - 1
+
+    new_tokens = []
+
+    # =========================
+    # 2. precompute reshape shape
+    # =========================
+    batch = 1
+
+    # =========================
+    # 3. main loop
+    # =========================
+    for i in range(num_new_tokens):
+
+        # ⚡ use flash attention (重要优化)
+        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+            next_token = decode_one_token(
+                model=model,
+                x=cur_token,
+                input_pos=input_pos,
+                previous_tokens=previous_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                semantic_logit_bias=semantic_logit_bias,
+                audio_masks=audio_masks,
+                audio_parts=audio_parts,
+            ).clone()
+
+        # =========================
+        # 4. update position
+        # =========================
+        input_pos += 1
+
+        # =========================
+        # 5. reshape once (reuse view logic)
+        # =========================
+        next_token_2d = next_token.view(num_streams, -1)
+
+        cur_token = next_token_2d.unsqueeze(0)
+
+        # =========================
+        # 6. ring buffer update (NO roll)
+        # =========================
+        previous_tokens[:, write_idx] = next_token_2d[:, 0]
+        write_idx = (write_idx + 1) % history_len
+
+        # =========================
+        # 7. store output
+        # =========================
+        new_tokens.append(next_token)
+
+        # =========================
+        # 8. EOS check
+        # =========================
+        if cur_token[0, 0, -1] == im_end_id:
+            break
+
+    return new_tokens
+
+
 @torch.no_grad()
 @torch.inference_mode()
 def generate(
@@ -252,6 +342,7 @@ def generate(
         max_new_tokens: int,
         audio_masks: torch.Tensor,
         audio_parts: torch.Tensor,
+        prompt_tokens = None,
         decode_one_token=decode_one_token_ar,
         num_samples: int = 1,
         **sampling_kwargs,
@@ -355,7 +446,33 @@ def generate(
     input_pos = torch.tensor([T], device=device, dtype=torch.int)
     step9 = time.perf_counter()
 
-    x = decode_n_tokens(
+    im_end_id = model.tokenizer.get_token_id(IM_END_TOKEN)
+    codebook_dim = 1 + model.config.num_codebooks
+    window_size = 64
+
+    previous_tokens = torch.zeros(
+        (1, window_size, codebook_dim),
+        device=device,
+        dtype=first_token.dtype,
+    )
+
+    # =========================
+    # 1. warm start prompt
+    # =========================
+    if prompt_tokens is not None:
+        # 确保 shape = [B, T, C]
+        if prompt_tokens.dim() == 2:
+            prompt_tokens = prompt_tokens.unsqueeze(0)
+
+        T = min(prompt_tokens.size(1), window_size)
+        previous_tokens[:, -T:] = prompt_tokens[:, -T:]
+
+    # =========================
+    # 2. insert first token
+    # =========================
+    previous_tokens[:, -1, :] = first_token.view(codebook_dim)
+
+    x = decode_n_tokens_optimized(
         model,
         first_token.view(1, codebook_dim, -1),
         input_pos,
@@ -366,6 +483,8 @@ def generate(
         semantic_logit_bias=semantic_logit_bias,
         audio_masks=audio_masks,
         audio_parts=audio_parts,
+        im_end_id=im_end_id,
+        previous_tokens=previous_tokens,
         decode_one_token=decode_one_token,
     )
     seq = seq[:, : T + 1 + x.size(1)]
@@ -725,6 +844,7 @@ def generate_long(
                 temperature=temperature,
                 top_p=top_p,
                 top_k=top_k,
+                prompt_tokens=all_codes,
             )
 
             if sample_idx == 0 and batch_idx == 0 and compile: