fix

Author: sufubao

lightllm/common/fused_moe/grouped_fused_moe.py

@@ -492,7 +492,7 @@ def _get_grouped_matmul_configs():
     name="grouped_matmul:v1",
     configs_gen_func=_get_grouped_matmul_configs,
     static_key_func=_get_grouped_matmul_static_key,
-    run_key_func=lambda token_num_mul_topk_num: token_num_mul_topk_num,
+    run_key_func=lambda token_inputs: token_inputs.shape[0],
 )
 def grouped_matmul(
     token_num_mul_topk_num: int,

lightllm/common/quantization/triton_quant/fp8/fp8w8a8_block_gemm_kernel.py

@@ -7,7 +7,7 @@
 from functools import lru_cache
 from typing import Any, Dict, List, Optional, Tuple
 from triton import Config
-from lightllm.common.triton_utils.autotuner import autotune, nearest_power_of_2
+from lightllm.common.triton_utils.autotuner import autotune
 
 
 class Fp8BlockMMKernelConfig(KernelConfigs):
@@ -180,7 +180,7 @@ def _get_static_key(A, B, block_size, dtype):
     name="w8a8_block_fp8_matmul:v1",
     configs_gen_func=get_test_configs,
     static_key_func=_get_static_key,
-    run_key_func=lambda M: M,
+    run_key_func=lambda A: A.shape[0],
 )
 def w8a8_block_fp8_matmul(
     A: torch.Tensor,

lightllm/common/triton_utils/autotuner.py

@@ -225,12 +225,15 @@ def _autotune(self, args, kwargs, static_key, run_key):
 
         # save configs to file
         if not dist.is_initialized() or get_global_rank() == 0:
-            cache_file = os.path.join(self.cache_dir, f"{KernelConfigs.get_config_file_name(static_key)}.json")
+            cache_file = os.path.join(self.cache_dir, KernelConfigs.get_config_file_name(static_key))
             with open(cache_file, "wb") as f:
                 fcntl.flock(f, fcntl.LOCK_EX)
                 try:
                     f.write(
-                        orjson.dumps(self.cached_configs[static_key], option=orjson.OPT_INDENT_2 | orjson.OPT_SORT_KEYS)
+                        orjson.dumps(
+                            self.cached_configs[static_key],
+                            option=orjson.OPT_INDENT_2 | orjson.OPT_SORT_KEYS | orjson.OPT_NON_STR_KEYS,
+                        )
                     )
                 finally:
                     fcntl.flock(f, fcntl.LOCK_UN)