improve cpu kv kernel setting

Author: wangzaijun

lightllm/common/basemodel/triton_kernel/kv_cache_offload.py

@@ -116,6 +116,7 @@ def offload_gpu_kv_to_cpu(
     page_readies: torch.Tensor,
     tp_index: int,
     tp_world_size: int,
+    grid_num: int,
     _cache_data={},
 ):
     """
@@ -231,7 +232,7 @@ def offload_gpu_kv_to_cpu(
     assert token_block_size == triton.next_power_of_2(token_block_size)
     page_num = page_indexes.shape[0]
 
-    grid = (1,)
+    grid = (grid_num,)
     num_warps = 4
 
     _offload_gpu_kv_to_cpu[grid](

lightllm/server/core/objs/start_args_type.py

@@ -122,3 +122,6 @@ class StartArgs:
     # multi_modal
     enable_multimodal: bool = field(default=False)
     enable_multimodal_audio: bool = field(default=False)
+
+    # kernel setting
+    enable_fa3: bool = field(default=False)

lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py

@@ -202,6 +202,7 @@ def _start_kv_cache_offload_task(
                 page_readies=page_readies,
                 tp_index=self.backend.rank_in_dp,
                 tp_world_size=self.backend.dp_world_size,
+                grid_num=1 if self.args.enable_fa3 else 16,  # TODO 更有效的分配策略。
             )
 
             sync_event = torch.cuda.Event()