fix config search type error. (#685)

Co-authored-by: wangzaijun <wangzaijun@sensetime.com>
ModelTC · Dec 25, 2024 · 06c31fd · 06c31fd
1 parent d9305a6
commit 06c31fd
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/lightllm/common/fused_moe/moe_kernel_configs.py b/lightllm/common/fused_moe/moe_kernel_configs.py
@@ -37,7 +37,7 @@ def try_to_get_best_config(
         finded_config = cls.get_the_config(key_params)
 
         if finded_config:
-            config = finded_config[min(finded_config.keys(), key=lambda x: abs(x - M))]
+            config = finded_config[min(finded_config.keys(), key=lambda x: abs(int(x) - M))]
             return config
         else:
             if M <= expert_num:

diff --git a/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_config.py b/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_config.py
@@ -31,9 +31,9 @@ def try_to_get_best_config(
         if finded_config:
             # two search dim, first is avg_seq_len_in_batch, second is batch_size
             batch_size_config: dict = finded_config[
-                min(finded_config.keys(), key=lambda x: abs(x - avg_seq_len_in_batch))
+                min(finded_config.keys(), key=lambda x: abs(int(x) - avg_seq_len_in_batch))
             ]
-            config = batch_size_config[min(batch_size_config.keys(), key=lambda x: abs(x - batch_size))]
+            config = batch_size_config[min(batch_size_config.keys(), key=lambda x: abs(int(x) - batch_size))]
 
             return config
         else: