remove kv buffer for decode

Author: shihaobai

lightllm/common/basemodel/basemodel.py

@@ -334,11 +334,6 @@ def _decode(
         # 所以不再使用分配连续的mem带来的优化，保证推理流程的一致
         infer_state.mem_is_contiguous = False
         infer_state.mem_index = mem_indexes
-        infer_state.kv_buffer = torch.empty(
-            (batch_size, self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
-            dtype=self.data_type,
-            device="cuda",
-        )
         copy_kv_index_to_req(self.req_manager.req_to_token_indexs, b_req_idx, b_seq_len, infer_state.mem_index)
 
         infer_state.init_some_extra_state(self, input_ids)

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

@@ -35,7 +35,12 @@ def _pre_cache_kv(self, infer_state: InferStateInfo, layer_weight) -> Tuple[torc
                 infer_state.mem_start : infer_state.mem_end, :, :
             ]
         else:
-            cache_kv = infer_state.kv_buffer
+            dtype = infer_state.mem_manager.kv_buffer.dtype
+            cache_kv = self.alloc_tensor(
+                [infer_state.batch_size, self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_],
+                dtype=dtype,
+                device="cuda",
+            )
         return cache_kv
 
     def _get_qkv(