tune config to speedup

Author: WANDY666

cosyvoice/cosyvoice/cli/frontend.py

@@ -64,8 +64,16 @@ def __init__(self,
         if self.use_ttsfrd:
             self.frd = ttsfrd.TtsFrontendEngine()
             ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
-            assert self.frd.initialize('{}/../../pretrained_models/CosyVoice-ttsfrd/resource'.format(ROOT_DIR)) is True, \
-                'failed to initialize ttsfrd resource'
+            resource_paths = [
+                '{}/../../../pretrained_models/CosyVoice-ttsfrd/resource'.format(ROOT_DIR),
+                '{}/../../pretrained_models/CosyVoice-ttsfrd/resource'.format(ROOT_DIR)
+            ]
+            initialized = False
+            for path in resource_paths:
+                if self.frd.initialize(path):
+                    initialized = True
+                    break
+            assert initialized, 'failed to initialize ttsfrd resource'
             self.frd.set_lang_type('pinyinvg')
         else:
             self.zh_tn_model = ZhNormalizer(remove_erhua=False)

light_tts/common/all_kernel_configs/triton_flashdecoding/{head_dim=64,kv_head_num=2,out_dtype=torch.float16,q_head_num=14}_NVIDIA_GeForce_RTX_5090.json

@@ -0,0 +1 @@
+{"32": {"1": {"BLOCK_SEQ": 16, "BLOCK_N": 16, "stage1_num_warps": 1, "stage1_num_stages": 4, "stage2_num_warps": 1, "stage2_num_stages": 1}, "2": {"BLOCK_SEQ": 256, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 6, "stage2_num_warps": 1, "stage2_num_stages": 1}, "4": {"BLOCK_SEQ": 256, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 6, "stage2_num_warps": 1, "stage2_num_stages": 1}, "8": {"BLOCK_SEQ": 256, "BLOCK_N": 32, "stage1_num_warps": 4, "stage1_num_stages": 1, "stage2_num_warps": 1, "stage2_num_stages": 1}}, "64": {"1": {"BLOCK_SEQ": 16, "BLOCK_N": 16, "stage1_num_warps": 1, "stage1_num_stages": 6, "stage2_num_warps": 1, "stage2_num_stages": 1}, "2": {"BLOCK_SEQ": 32, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 8, "stage2_num_warps": 1, "stage2_num_stages": 1}, "4": {"BLOCK_SEQ": 256, "BLOCK_N": 64, "stage1_num_warps": 4, "stage1_num_stages": 8, "stage2_num_warps": 1, "stage2_num_stages": 1}, "8": {"BLOCK_SEQ": 256, "BLOCK_N": 64, "stage1_num_warps": 4, "stage1_num_stages": 3, "stage2_num_warps": 1, "stage2_num_stages": 1}}, "128": {"1": {"BLOCK_SEQ": 16, "BLOCK_N": 16, "stage1_num_warps": 1, "stage1_num_stages": 8, "stage2_num_warps": 1, "stage2_num_stages": 1}, "2": {"BLOCK_SEQ": 32, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 1, "stage2_num_warps": 1, "stage2_num_stages": 1}, "4": {"BLOCK_SEQ": 32, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 8, "stage2_num_warps": 1, "stage2_num_stages": 1}, "8": {"BLOCK_SEQ": 64, "BLOCK_N": 64, "stage1_num_warps": 4, "stage1_num_stages": 4, "stage2_num_warps": 1, "stage2_num_stages": 1}}, "256": {"1": {"BLOCK_SEQ": 16, "BLOCK_N": 16, "stage1_num_warps": 1, "stage1_num_stages": 6, "stage2_num_warps": 1, "stage2_num_stages": 1}, "2": {"BLOCK_SEQ": 32, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 8, "stage2_num_warps": 1, "stage2_num_stages": 1}, "4": {"BLOCK_SEQ": 32, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 8, "stage2_num_warps": 1, "stage2_num_stages": 1}, "8": {"BLOCK_SEQ": 64, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 7, "stage2_num_warps": 1, "stage2_num_stages": 1}}, "512": {"1": {"BLOCK_SEQ": 16, "BLOCK_N": 16, "stage1_num_warps": 1, "stage1_num_stages": 7, "stage2_num_warps": 1, "stage2_num_stages": 1}, "2": {"BLOCK_SEQ": 16, "BLOCK_N": 16, "stage1_num_warps": 1, "stage1_num_stages": 6, "stage2_num_warps": 1, "stage2_num_stages": 1}, "4": {"BLOCK_SEQ": 32, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 8, "stage2_num_warps": 1, "stage2_num_stages": 1}, "8": {"BLOCK_SEQ": 32, "BLOCK_N": 32, "stage1_num_warps": 1, "stage1_num_stages": 8, "stage2_num_warps": 1, "stage2_num_stages": 1}}}

light_tts/models/llama/triton_kernel/flash_decoding.py

@@ -4,6 +4,7 @@
 from frozendict import frozendict
 from typing import Dict
 
+
 class LlamaFlashDecodingStage1KernelConfig(KernelConfigs):
     kernel_name: str = "triton_flashdecoding"
 
@@ -44,7 +45,7 @@ def try_to_get_best_config(
                 "stage2_num_stages": 2,
             }
             return config
-    
+
     @classmethod
     def save_config(cls, *args, **kwargs) -> None:
         key_params = {
@@ -57,13 +58,13 @@ def save_config(cls, *args, **kwargs) -> None:
 
         cls.store_config(key_params, kwargs["store_json_ans"])
 
+
 def token_decode_attention_flash_decoding(
     q, infer_state, q_head_num, head_dim, cache_k, cache_v, out=None, alloc_tensor_func=torch.empty
 ):
     batch_size = infer_state.batch_size
-    avg_seq_len_in_batch = infer_state.total_token_num // batch_size
     run_config = LlamaFlashDecodingStage1KernelConfig.try_to_get_best_config(
-        batch_size, avg_seq_len_in_batch, head_dim, q_head_num, cache_k.shape[1], torch.float16
+        batch_size, infer_state.max_len_in_batch, head_dim, q_head_num, cache_k.shape[1], torch.float16
     )
     BLOCK_SEQ = run_config["BLOCK_SEQ"]
 

light_tts/server/api_cli.py

@@ -116,7 +116,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--graph_max_len_in_batch",
         type=int,
-        default=32768,
+        default=2048,
         help="""Maximum sequence length that can be captured by the cuda graph for decodign stage.
                 The default value is 8192. It will turn into eagar mode if encounters a larger value. """,
     )