fix

Author: hiworldwzj

lightllm/common/basemodel/basemodel.py

@@ -302,7 +302,6 @@ def _prefill(
         infer_state.mem_manager = self.mem_manager
         infer_state.req_manager = self.req_manager
 
-        infer_state.mem_is_contiguous = False
         infer_state.mem_index = mem_indexes
         infer_state.kv_buffer = torch.empty(
             (input_ids.shape[0], self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
@@ -351,9 +350,6 @@ def _decode(
         infer_state.mem_manager = self.mem_manager
         infer_state.req_manager = self.req_manager
 
-        # 在使用 cuda graph 特性的时候，必须保证每次推理的流程一致
-        # 所以不再使用分配连续的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_),
@@ -398,9 +394,6 @@ def create_inferstate(cur_batch: DecodeMicroBatch, batch_index):
             infer_state.mem_manager = self.mem_manager
             infer_state.req_manager = self.req_manager
 
-            # 在使用 cuda graph 特性的时候，必须保证每次推理的流程一致
-            # 所以不再使用分配连续的mem带来的优化，保证推理流程的一致
-            infer_state.mem_is_contiguous = False
             infer_state.mem_index = cur_batch.mem_indexes
             infer_state.kv_buffer = torch.empty(
                 (cur_batch.batch_size, self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
@@ -475,9 +468,6 @@ def create_inferstate(cur_batch: PrefillMicroBatch, batch_index):
             infer_state.mem_manager = self.mem_manager
             infer_state.req_manager = self.req_manager
 
-            # 在使用 cuda graph 特性的时候，必须保证每次推理的流程一致
-            # 所以不再使用分配连续的mem带来的优化，保证推理流程的一致
-            infer_state.mem_is_contiguous = False
             infer_state.mem_index = cur_batch.mem_indexes
             infer_state.kv_buffer = torch.empty(
                 (cur_batch.input_ids.shape[0], self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),

lightllm/common/basemodel/infer_struct.py

@@ -25,10 +25,7 @@ def __init__(self):
         self.mem_manager: MemoryManager = None
         self.req_manager: ReqManager = None
 
-        self.mem_is_contiguous = None
         self.mem_index = None
-        self.mem_start = None
-        self.mem_end = None
         self.kv_buffer = None
 
         self.is_token_healing = False

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

@@ -31,12 +31,7 @@ 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]:
-        if infer_state.mem_is_contiguous:
-            cache_kv = infer_state.mem_manager.kv_buffer[self.layer_num_][
-                infer_state.mem_start : infer_state.mem_end, :, :
-            ]
-        else:
-            cache_kv = infer_state.kv_buffer
+        cache_kv = infer_state.kv_buffer
         return cache_kv
 
     def _get_qkv(