back the infer_struct

Author: shihaobai

lightllm/common/basemodel/basemodel.py

@@ -341,7 +341,7 @@ def _prefill(
             infer_state.mem_index,
         )
 
-        infer_state.init_some_extra_state(self, model_input)
+        infer_state.init_some_extra_state(self, model_input.input_ids)
         return self._context_forward(model_input.input_ids, infer_state)
 
     def _decode(
@@ -365,7 +365,7 @@ def _decode(
                 infer_state.b_seq_len,
                 infer_state.mem_index,
             )
-            infer_state.init_some_extra_state(self, padded_model_input)
+            infer_state.init_some_extra_state(self, padded_model_input.input_ids)
 
             if self.graph.need_capture(find_graph_batch_size):
                 infer_state.is_cuda_graph = True
@@ -386,7 +386,7 @@ def _decode(
                 infer_state.b_seq_len,
                 infer_state.mem_index,
             )
-            infer_state.init_some_extra_state(self, model_input)
+            infer_state.init_some_extra_state(self, model_input.input_ids)
             model_output = self._token_forward(model_input.input_ids, infer_state)
 
         return model_output

lightllm/common/basemodel/infer_struct.py

@@ -65,7 +65,7 @@ def __init__(self):
         # 的输入会用到，其他模型和场景都不会用到
         self.deepseekv3_mtp_draft_input_hiddens: Optional[torch.Tensor] = None
 
-    def init_some_extra_state(self, model, model_input: ModelInput):
+    def init_some_extra_state(self, model, input_ids: torch.Tensor):
         if self.is_prefill:
             (
                 self.b_q_seq_len,
@@ -76,7 +76,7 @@ def init_some_extra_state(self, model, model_input: ModelInput):
                 self.max_q_seq_len,
                 self.max_kv_seq_len,
             ) = gen_prefill_params(
-                input_token_num=model_input.input_ids.shape[0],
+                input_token_num=input_ids.shape[0],
                 b_ready_cache_len=self.b_ready_cache_len,
                 b_seq_len=self.b_seq_len,
             )
@@ -88,10 +88,10 @@ def init_some_extra_state(self, model, model_input: ModelInput):
                 self.b_kv_seq_len,
                 self.b1_cu_kv_seq_len,
                 self.position_ids,
+                self.max_q_seq_len,
+                self.max_kv_seq_len,
             ) = gen_decode_params(self.b_seq_len)
             self.b_start_loc = self.b1_cu_kv_seq_len[0:-1]
-            self.max_q_seq_len = 1
-            self.max_kv_seq_len = model_input.max_len_in_batch
 
     def copy_for_cuda_graph(self, new_infer_state: "InferStateInfo"):
         for attr_name, attr_value in vars(new_infer_state).items():

lightllm/common/basemodel/triton_kernel/gen_decode_params.py

@@ -10,5 +10,7 @@ def gen_decode_params(b_seq_len: torch.Tensor):
     position_ids = b_seq_len - 1
     b_q_seq_len = torch.ones_like(b_seq_len)
     b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len)
+    max_q_seq_len = b_q_seq_len.max().item()
+    max_kv_seq_len = b_kv_seq_len.max().item()
 
-    return b_q_seq_len, b1_cu_q_seq_len, b_kv_seq_len, b1_cu_kv_seq_len, position_ids
+    return b_q_seq_len, b1_cu_q_seq_len, b_kv_seq_len, b1_cu_kv_seq_len, position_ids, max_q_seq_len, max_kv_seq_len

lightllm/models/llama/flashattention_infer_struct.py

@@ -24,13 +24,13 @@ def get_page_table_buffer(cls, graph_max_batch_size: int, max_seq_len: int):
             ]
         return cls._shared_page_table_buffer
 
-    def init_some_extra_state(self, model, model_input: ModelInput):
-        super().init_some_extra_state(model, model_input)
+    def init_some_extra_state(self, model, input_ids: torch.Tensor):
+        super().init_some_extra_state(model, input_ids)
         if self.is_prefill:
             self.cu_seqlens_q = self.b1_cu_q_seq_len.int()
             self.cu_seqlens_k = self.b1_cu_kv_seq_len.int()
             self.page_table = torch.empty(
-                (self.batch_size, self.max_seq_len), dtype=torch.int32, device=model_input.input_ids.device
+                (self.batch_size, self.max_seq_len), dtype=torch.int32, device=input_ids.device
             )
             self.page_table.copy_(model.req_manager.req_to_token_indexs[self.b_req_idx, : self.max_seq_len])
         else:
@@ -47,7 +47,7 @@ def init_some_extra_state(self, model, model_input: ModelInput):
                 ].reshape(self.batch_size, model.graph_max_len_in_batch)
             else:
                 self.page_table = torch.empty(
-                    (self.batch_size, self.max_len_in_batch), dtype=torch.int32, device=model_input.input_ids.device
+                    (self.batch_size, self.max_len_in_batch), dtype=torch.int32, device=input_ids.device
                 )
 
             self.page_table[:, :max_seq_len_k].copy_(
@@ -58,7 +58,7 @@ def init_some_extra_state(self, model, model_input: ModelInput):
 
         if "offline_calibration_fp8kv" in model.mode:
             if self.is_prefill:
-                device = model_input.input_ids.device
+                device = input_ids.device
                 # q_scale和token_batch_ids在对q做per head量化使用，为了节省资源在推理外部初始化
                 self.q_scale = torch.empty(
                     (self.batch_size, self.mem_manager.head_num), dtype=torch.float32, device=device
@@ -78,7 +78,7 @@ def init_some_extra_state(self, model, model_input: ModelInput):
                 else torch.ones(
                     (self.mem_manager.layer_num, self.batch_size, head_num),
                     dtype=torch.float32,
-                    device=model_input.input_ids.device,
+                    device=input_ids.device,
                 )
             )
             self.v_descale = (
@@ -89,7 +89,7 @@ def init_some_extra_state(self, model, model_input: ModelInput):
                 else torch.ones(
                     (self.mem_manager.layer_num, self.batch_size, head_num),
                     dtype=torch.float32,
-                    device=model_input.input_ids.device,
+                    device=input_ids.device,
                 )
             )
         return