refactor

Author: sufubao

lightllm/common/basemodel/basemodel.py

@@ -81,6 +81,10 @@ def __init__(self, kvargs):
 
         self._init_datatype()
         self._init_config()
+
+        if os.environ.get("LIGHTLLM_TRITON_AUTOTUNE", "0") == "1":
+            self.layers_num = self.autotune_layers()
+
         self._verify_must()
         self._verify_params()
         self._init_quant()
@@ -744,8 +748,6 @@ def _autotune_warmup(self):
 
         warmup_lengths.sort(reverse=True)
 
-        layer_num_bak = self.layers_num
-        self.layers_num = self.autotune_layers()
         for input_len in warmup_lengths:
             try:
                 logger.info(f"autotune warmup for length {input_len}")
@@ -777,14 +779,16 @@ def _autotune_warmup(self):
                 del model_output
                 self.req_manager.free_all()
                 self.mem_manager.free_all()
+                torch.cuda.empty_cache()
                 logger.info(f"autotune warmup for length {input_len} ok")
             except Exception as e:
                 logger.warning(f"autotune warmup for length {input_len} failed: {str(e)}")
                 self.req_manager.free_all()
                 self.mem_manager.free_all()
-        self.layers_num = layer_num_bak
+                torch.cuda.empty_cache()
         torch.distributed.barrier()
-        os.environ["LIGHTLLM_TRITON_AUTOTUNE"] = "0"
+        logger.info("autotune warmup done, exit!")
+        exit(0)
 
     @final
     @torch.no_grad()

lightllm/common/fused_moe/grouped_fused_moe.py

@@ -478,7 +478,7 @@ def get_grouped_matmul_static_key(
             "BLOCK_SIZE_N": bn,
             "BLOCK_SIZE_K": bk,
             "GROUP_SIZE_M": gm,
-            "NUM_WARPS": nw,
+            "num_warps": nw,
             "NUM_STAGE": ns,
         }
         for ns in [1, 2, 3, 4, 5]
@@ -493,7 +493,7 @@ def get_grouped_matmul_static_key(
         "BLOCK_SIZE_N": 64,
         "BLOCK_SIZE_K": 32,
         "GROUP_SIZE_M": 8,
-        "NUM_WARPS": 4,
+        "num_warps": 4,
         "NUM_STAGE": 1,
     },
     static_key_func=get_grouped_matmul_static_key,
@@ -550,7 +550,7 @@ def grouped_matmul(
     BLOCK_SIZE_N = run_config["BLOCK_SIZE_N"]
     BLOCK_SIZE_K = run_config["BLOCK_SIZE_K"]
     GROUP_SIZE_M = run_config["GROUP_SIZE_M"]
-    num_warps = run_config["NUM_WARPS"]
+    num_warps = run_config["num_warps"]
     num_stages = run_config["NUM_STAGE"]
 
     if block_size_k != 0: