fix: fix a bug in flashinfer

Author: niushengxiao

lightllm/models/deepseek2/flashinfer_struct.py

@@ -38,7 +38,7 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                     self.b_req_idx,
                     self.b_seq_len,
                     self.b_start_loc,
-                    self.max_len_in_batch,
+                    self.max_kv_seq_len,
                     self.kv_indices,
                 )
                 if self.decode_wrapper is None:

lightllm/models/llama/flashinfer_struct.py

@@ -41,7 +41,7 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                     self.b_req_idx,
                     self.b_seq_len,
                     self.b_start_loc,
-                    self.max_len_in_batch,
+                    self.max_kv_seq_len,
                     self.kv_indices,
                 )
                 self.kv_starts = self.b1_cu_kv_seq_len.int()
@@ -81,8 +81,8 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                     self.req_manager.req_to_token_indexs,
                     self.b_req_idx,
                     self.b_seq_len,
-                    kv_starts,
-                    self.max_len_in_batch,
+                    kv_starts[:-1],
+                    self.max_kv_seq_len,
                     kv_indices,
                 )
                 self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper(