code clean up.

Author: hiworldwzj

lightllm/common/basemodel/basemodel.py

@@ -303,10 +303,9 @@ def _prefill(
         infer_state.req_manager = self.req_manager
 
         infer_state.mem_index = mem_indexes
-        infer_state.kv_buffer = torch.empty(
+        infer_state.kv_buffer_shapedtype = (
             (input_ids.shape[0], self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
-            dtype=self.data_type,
-            device="cuda",
+            self.data_type,
         )
         infer_state.dist_group = dist_group_manager.get_default_group()
 
@@ -351,10 +350,9 @@ def _decode(
         infer_state.req_manager = self.req_manager
 
         infer_state.mem_index = mem_indexes
-        infer_state.kv_buffer = torch.empty(
+        infer_state.kv_buffer_shapedtype = (
             (batch_size, self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
-            dtype=self.data_type,
-            device="cuda",
+            self.data_type,
         )
         infer_state.dist_group = dist_group_manager.get_default_group()
         copy_kv_index_to_req(self.req_manager.req_to_token_indexs, b_req_idx, b_seq_len, infer_state.mem_index)
@@ -395,10 +393,9 @@ def create_inferstate(cur_batch: DecodeMicroBatch, batch_index):
             infer_state.req_manager = self.req_manager
 
             infer_state.mem_index = cur_batch.mem_indexes
-            infer_state.kv_buffer = torch.empty(
+            infer_state.kv_buffer_shapedtype = (
                 (cur_batch.batch_size, self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
-                dtype=self.data_type,
-                device="cuda",
+                self.data_type,
             )
             infer_state.dist_group = dist_group_manager.get_group(batch_index)
             copy_kv_index_to_req(
@@ -469,10 +466,9 @@ def create_inferstate(cur_batch: PrefillMicroBatch, batch_index):
             infer_state.req_manager = self.req_manager
 
             infer_state.mem_index = cur_batch.mem_indexes
-            infer_state.kv_buffer = torch.empty(
+            infer_state.kv_buffer_shapedtype = (
                 (cur_batch.input_ids.shape[0], self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
-                dtype=self.data_type,
-                device="cuda",
+                self.data_type,
             )
             infer_state.dist_group = dist_group_manager.get_group(batch_index)
             init_req_to_token_indexes(

lightllm/common/basemodel/infer_struct.py

@@ -26,7 +26,7 @@ def __init__(self):
         self.req_manager: ReqManager = None
 
         self.mem_index = None
-        self.kv_buffer = None
+        self.kv_buffer_shapedtype = None
 
         self.is_token_healing = False
         self.return_all_prompt_logics = False

lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py

@@ -31,7 +31,13 @@ def _ffn_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.T
         raise Exception("need to impl")
 
     def _pre_cache_kv(self, infer_state: InferStateInfo, layer_weight) -> Tuple[torch.Tensor, torch.Tensor]:
-        cache_kv = infer_state.kv_buffer
+        cache_kv = self.alloc_tensor(
+            shape=infer_state.kv_buffer_shapedtype[0],
+            dtype=infer_state.kv_buffer_shapedtype[1],
+            device="cuda",
+            is_graph_out=False,
+            microbatch_index=infer_state.microbatch_index,
+        )
         return cache_kv
 
     def _get_qkv(