[bugfix] index out of range (#223)

* 解决数组越界

* 解决数组越界

Author: zbb200819

ucm/ucm_sparse/gsa.py

@@ -152,6 +152,7 @@ def updata_req_state(
                 self.calc_block_table = []
                 self.calc_repre_slot_mapping = []
         if len(self.repre_slot_mapping) > len(self.blocks):
+            self.topk_buf_tmp = None
             self.repre_slot_mapping = self.repre_slot_mapping[: len(self.blocks)]
 
     def _get_sparse_and_free_block(self):
@@ -265,6 +266,8 @@ def get_model_input(
                 input_batch.req_id_to_index[req_id],
             )
         for new_req in scheduler_output.scheduled_new_reqs:
+            if new_req.req_id in self.gsa_stats:
+                del self.gsa_stats[new_req.req_id]
             self.gsa_stats[new_req.req_id] = GSAReqStat(new_req.req_id)
             self.gsa_stats[new_req.req_id].add_req_new(
                 scheduler_output.num_scheduled_tokens[new_req.req_id],

ucm/ucm_sparse/prefetch_engine.py

@@ -517,6 +517,16 @@ def _no_gsa_input_deal(
                 self.is_gsa_req_id[req_id]
                 and gsa_metadata.gsa_stats[req_id].topk_buf_tmp != None
             ):
+                if (torch.max(gsa_metadata.gsa_stats[req_id].topk_buf_tmp) > 
+                    (len(self.block_table_list_bs[index]) - 1)
+                ):
+                    self.gsa_seq_len[:, bs_index] = gsa_metadata.gsa_stats[
+                        req_id
+                    ].get_seq_len()
+                    self.use_block_table[
+                        :, bs_index, : len(gsa_metadata.gsa_stats[req_id].blocks)
+                    ] = one_block_table
+                    continue
                 remain_slot = (
                     gsa_metadata.gsa_stats[req_id].get_seq_len() % self.block_size
                 )