feat: change mode name to offline_calibration_fp8kv

Author: niushengxiao

lightllm/common/mem_utils.py

@@ -22,7 +22,7 @@ def select_mem_manager_class(mode):
         logger.info("Model kv cache using mode triton int8kv")
     elif "triton_fp8kv" in mode:
         raise Exception("currently only for deepseek")
-    elif "calibration_fp8kv" in mode:
+    elif "offline_calibration_fp8kv" in mode:
         memory_manager_class = CalibrationFP8KVMemoryManager
         logger.info("Model kv cache using mode calibration fp8kv")
     elif "export_fp8kv_calibration" in mode:

lightllm/models/llama/flashattention_infer_struct.py

@@ -32,7 +32,7 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                 (self.batch_size, self.max_seq_len), dtype=torch.int32, device=input_ids.device
             )
             self.page_table.copy_(model.req_manager.req_to_token_indexs[self.b_req_idx, : self.max_seq_len])
-            if "calibration_fp8kv" in model.mode:
+            if "offline_calibration_fp8kv" in model.mode:
                 device = input_ids.device
                 self.q_scale = torch.empty(
                     (self.batch_size, self.mem_manager.head_num), dtype=torch.float32, device=device
@@ -60,7 +60,7 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
             )
             self.page_table[:, max_seq_len_k:].fill_(0)
 
-        if "calibration_fp8kv" in model.mode:
+        if "offline_calibration_fp8kv" in model.mode:
             offline_scales = self.mem_manager.scales
             head_num = self.mem_manager.head_num
             self.k_descale = (

lightllm/models/llama/layer_infer/transformer_layer_infer.py

@@ -69,15 +69,15 @@ def _bind_norm(self):
 
     def _bind_attention(self):
         if get_env_start_args().enable_fa3:
-            if "calibration_fp8kv" in self.mode:
+            if "offline_calibration_fp8kv" in self.mode:
                 self._context_attention_kernel = partial(
                     LlamaTransformerLayerInfer._context_attention_flashattention_fp8, self
                 )
                 self._token_attention_kernel = partial(
                     LlamaTransformerLayerInfer._token_decode_attention_flashattention_fp8, self
                 )
                 self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_fp8kv, self)
-            else:
+            elif not self.mode:
                 self._context_attention_kernel = partial(
                     LlamaTransformerLayerInfer._context_attention_flashattention, self
                 )
@@ -90,6 +90,8 @@ def _bind_attention(self):
                     )
                 else:
                     self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
+            else:
+                raise Exception(f"Unsupported mode for fa3 backend: {self.mode}")
             return
         elif get_env_start_args().enable_flashinfer_prefill:
             self._context_attention_kernel = partial(
@@ -127,7 +129,7 @@ def _bind_attention(self):
         elif "triton_int8kv" in self.mode:
             self._token_attention_kernel = partial(LlamaTransformerLayerInfer._token_decode_attention_int8kv, self)
             self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_int8kv, self)
-        elif "calibration_fp8kv" in self.mode:
+        elif "offline_calibration_fp8kv" in self.mode:
             raise Exception("calibration fp8 kvcache only support fa3 backend")
         elif "triton_flashdecoding" in self.mode:
             self._token_attention_kernel = partial(
@@ -147,14 +149,16 @@ def _bind_attention(self):
                 LlamaTransformerLayerInfer._token_decode_attention_gqa_flashdecoding_vsm, self
             )
             self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
-        else:
+        elif not self.mode:
             if get_env_start_args().enable_flashinfer_decode:
                 self._token_attention_kernel = partial(
                     LlamaTransformerLayerInfer._token_decode_attention_flashinfer, self
                 )
             else:
                 self._token_attention_kernel = partial(LlamaTransformerLayerInfer._token_decode_attention_normal, self)
             self._copy_kv_to_mem_cache = partial(LlamaTransformerLayerInfer._copy_kv_to_mem_cache_normal, self)
+        else:
+            raise Exception(f"Unsupported mode: {self.mode}")
 
         return
 

lightllm/server/api_cli.py

@@ -164,13 +164,13 @@ def make_argument_parser() -> argparse.ArgumentParser:
         default=[],
         nargs="+",
         help="""Model mode: [triton_int8kv | ppl_int8kv | ppl_fp16 | triton_flashdecoding
-                        | triton_gqa_attention | triton_gqa_flashdecoding | triton_fp8kv | calibration_fp8kv
+                        | triton_gqa_attention | triton_gqa_flashdecoding | triton_fp8kv | offline_calibration_fp8kv
                         | export_fp8kv_calibration
                         triton_flashdecoding mode is for long context, current support llama llama2 qwen;
                         triton_gqa_attention and triton_gqa_flashdecoding is fast kernel for model which use GQA;
                         triton_int8kv mode use int8 to store kv cache, can increase token capacity, use triton kernel;
                         triton_fp8kv mode use float8 to store kv cache, currently only for deepseek2;
-                        calibration_fp8kv mode use float8 to store kv cache, need fa3 backend,
+                        offline_calibration_fp8kv mode use float8 to store kv cache, need fa3 backend,
                         currently only for llama and qwen model;
                         export_fp8kv_calibration record and export kv cache quant calibration results to a json file.
                         It can be used for llama and qwen model. Calibration need to disable cudagraph and fa3 backend.