[bugfix] fix gsa coredump (#265)

* [bugfix] fix gsa coredump

* [bugfix] fix build_sparse_meta param lost

* clean code

Author: HaoLi980405

ucm/sparse/gsa/gsa.py

@@ -497,9 +497,7 @@ def init_topk_cal(
         self.gsa_offload_ops = gsa_offload_ops.CalKpreAndTopk(
             self.layer_num, block_size, MAX_BS, att_num_heads, head_size
         )
-        self.gsa_offload_ops.set_kpre_method_param(
-            int(max_model_len / block_size) * MAX_BS, kv_num_heads, 1
-        )
+        self.gsa_offload_ops.set_kpre_method_param(kv_num_heads, 1)
         self.gsa_offload_ops.set_kpre_cache(prefetch_engine.kpre_caches)
         self.is_cal_kpre = [False] * self.layer_num
         self.gsa_q_cache = torch.zeros(
@@ -868,10 +866,7 @@ def execute_finished(self):
                 )
 
     def build_sparse_meta(
-        self,
-        scheduler_output: SchedulerOutput,
-        requests,
-        input_batch,
+        self, scheduler_output: SchedulerOutput, requests, input_batch, attn_metadata
     ) -> None:
         self.gsa_metadata = self.build_gsa_metadata(
             scheduler_output, requests, input_batch

ucm/sparse/gsa/offload_ops/include/cal_kpre_and_topk.h

@@ -46,7 +46,7 @@ class __attribute__((visibility("hidden"))) CalKpreAndTopk
 public:
     CalKpreAndTopk(uint32_t layerNum, uint32_t blockSize, uint32_t maxBs, uint32_t numHeads, uint32_t headSize);
     ~CalKpreAndTopk();
-    void SetKpreMethodParam(uint32_t maxBlockNum, uint32_t numHeads, uint32_t numKpre);
+    void SetKpreMethodParam(uint32_t numHeads, uint32_t numKpre);
     void SetKpreCache(std::vector<torch::Tensor>& kpreCache);
     void SetTopkCache(std::vector<torch::Tensor>& topkCache, std::vector<uint32_t>& topkLens);
     void SetCommonParam(std::vector<uint32_t>& calTopkIdx, std::vector<bool>& isDecode);

ucm/sparse/gsa/offload_ops/src/cal_kpre_and_topk.cpp

@@ -27,15 +27,10 @@ CalKpreAndTopk::CalKpreAndTopk(uint32_t layerNum, uint32_t blockSize, uint32_t m
     m_count = 0;
 }
 
-void CalKpreAndTopk::SetKpreMethodParam(uint32_t maxBlockNum, uint32_t numHeads, uint32_t numKpre)
+void CalKpreAndTopk::SetKpreMethodParam(uint32_t numHeads, uint32_t numKpre)
 {
-    // m_kNumHeads = numHeads;
-    // m_numKpre = numKpre;
-    // auto optionsForKCache = torch::TensorOptions().device("cpu").dtype(torch::kFloat32);
-    // for (uint32_t i = 0; i < m_layerNum; i++) {
-    //     torch::Tensor layerKCache = torch::zeros({maxBlockNum, m_kNumHeads, m_blockSize, m_headSize}, optionsForKCache);
-    //     m_kCache.push_back(layerKCache);
-    // }
+    m_kNumHeads = numHeads;
+    m_numKpre = numKpre;
 }
 
 void CalKpreAndTopk::SetKpreCache(std::vector<torch::Tensor>& kpreCache)