[bugfix] fix whl install gsa error and gsa kpre reslotmapping out of range (#204)

* md max seq len bug

* clean code

---------

Co-authored-by: zbb200819 <1130072360@qq.com>

Author: HaoLi980405

setup.py

@@ -130,7 +130,7 @@ def _copy_so_files(self, ext: CMakeExtension):
             build_install_dir = "ucm/store"
         else:
             install_dir = GSA_INSTALL_DIR
-            build_install_dir = "ucm_sparse"
+            build_install_dir = "ucm/ucm_sparse"
 
         for so_file in so_files:
             src_path = os.path.join(so_search_dir, so_file)

ucm/ucm_sparse/gsa.py

@@ -457,9 +457,7 @@ def __init__(self, vllm_config: VllmConfig, role: UcmSparseRole):
             self.prefetch_engine = GSAPrefetchBase(
                 vllm_config, 16, True, True, False, 1
             )
-        self.topk_kpre_manger = TopKAndKpreManger(
-            vllm_config.scheduler_config.max_num_seqs
-        )
+        self.topk_kpre_manger = TopKAndKpreManger(MAX_BS)
         self.k_cache = {}
         self.v_cache = {}
         self.tasks_dump = {}
@@ -505,7 +503,7 @@ def init_topk_cal(
         self.gsa_q_cache = torch.zeros(
             (
                 self.layer_num,
-                vllm_config.scheduler_config.max_num_seqs,
+                MAX_BS,
                 att_num_heads,
                 head_size,
             ),

ucm/ucm_sparse/prefetch_engine.py

@@ -169,9 +169,14 @@ def _topk_tmp_deal(self, gsa_metadata, topk_buf_tmp):
         for index, topk_info in enumerate(self.topk_bs):
             if topk_info[1]:
                 if topk_info[0] in gsa_metadata.gsa_stats:
-                    gsa_metadata.gsa_stats[topk_info[0]].topk_buf_tmp = (
-                        self.topk_buf_tmp[:, index, : topk_info[2]].clone()
-                    )
+                    if not self.is_cpu_topk:
+                        gsa_metadata.gsa_stats[topk_info[0]].topk_buf_tmp = (
+                            self.topk_buf_tmp[:, index, : topk_info[2]].cpu()
+                        )
+                    else:
+                        gsa_metadata.gsa_stats[topk_info[0]].topk_buf_tmp = (
+                            self.topk_buf_tmp[:, index, : topk_info[2]].clone()
+                        )
         self.topk_bs = []
         for index, req_id in enumerate(self.req_ids_bs):
             one_block_len = len(gsa_metadata.gsa_stats[req_id].blocks)