Code clean up (#860)

Author: hiworldwzj

lightllm/common/basemodel/basemodel.py

@@ -302,12 +302,10 @@ 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(
+        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,14 +349,10 @@ 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(
+        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)
@@ -398,14 +392,10 @@ 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(
+            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(
@@ -475,14 +465,10 @@ 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(
+            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

@@ -25,11 +25,8 @@ 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.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,12 +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]:
-        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 = 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(
@@ -51,9 +52,8 @@ def _tpsp_get_qkv(
 
     def _post_cache_kv(self, cache_kv, infer_state: InferStateInfo, layer_weight):
         mem_manager = infer_state.mem_manager
-        if not infer_state.mem_is_contiguous:
-            self._copy_kv_to_mem_cache(cache_kv, infer_state.mem_index, mem_manager)
-            return
+        self._copy_kv_to_mem_cache(cache_kv, infer_state.mem_index, mem_manager)
+        return
 
     def _copy_kv_to_mem_cache(self, buffer, mem_index, mem_manager):
         destindex_copy_kv(buffer, mem_index, mem_manager.kv_buffer[self.layer_num_])

lightllm/models/deepseek2/infer_struct.py

@@ -14,7 +14,6 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
         super().init_some_extra_state(model, input_ids)
         if not self.is_prefill:
             self.kv_starts = torch.cat([self.b_start_loc, self.b_start_loc[-1:] + self.b_seq_len[-1:]], dim=0)
-            self.total_token_num_tensor = torch.sum(self.b_seq_len)
 
         if self.is_prefill:
             self.b_kv_start_loc = self.b_seq_len.cumsum(dim=0) - self.b_seq_len

test/kernel/tuning/deepseekv2_gqa_decode_tuning.py

@@ -50,7 +50,6 @@ def test_decode_attentions(
     ).cuda()
     infer_state.b_req_idx = torch.arange(0, infer_state.batch_size, step=1, dtype=torch.int32).cuda()
     infer_state.b_seq_len = torch.full((infer_state.batch_size,), fill_value=test_seq_len, dtype=torch.int32).cuda()
-    infer_state.total_token_num_tensor = torch.sum(infer_state.b_seq_len)
 
     input_tuples = []
     for _ in range(test_count):

test/kernel/tuning/llama_gqa_decode_vsm_tuning.py

@@ -51,8 +51,6 @@ def test_decode_attentions(
     ).cuda()
     state.b_req_idx = torch.arange(0, state.batch_size, step=1, dtype=torch.int32).cuda()
     state.b_seq_len = torch.full((state.batch_size,), fill_value=test_seq_len, dtype=torch.int32).cuda()
-    total_token_num_tensor = torch.tensor([state.batch_size * test_seq_len], dtype=torch.int32, device="cuda")
-    state.total_token_num = total_token_num_tensor
 
     args = []
     q_head_dim = q_shape[2]
@@ -63,7 +61,7 @@ def test_decode_attentions(
     state.q_head_dim = q_head_dim
     state.kv_head_num = kv_head_num
     state.softmax_scale = 1 / (q_head_dim ** 0.5)
-    state.total_token_num = total_token_num_tensor
+    state.total_token_num = state.batch_size * test_seq_len
 
     infer_state = state
     for _ in range(test_count):