fix

Author: randgun

csrc/lightning_indexer/op_host/lightning_indexer.cpp

@@ -15,6 +15,7 @@ namespace sglang::LIHost {
 
 using namespace ge_helper;
 constexpr uint32_t MAX_CAPTURE_NUM = 1024;
+constexpr uint32_t MAX_DECODE_BS = 512;
 // npu tensor max size
 constexpr int SIZE = 8;
 constexpr int DIM_0 = 0;
@@ -117,26 +118,33 @@ HOST_API at::Tensor lightning_indexer(const at::Tensor &query, const at::Tensor
     uint32_t tilingSize = sizeof(LITilingData);
     auto blockDim = tilingData.usedCoreNum;
     auto bs = tilingData.bSize;
-    uint64_t mapKey = tilingData.tilingKey;
-    mapKey = (mapKey << 32) | bs;
-    //std::cout << "mapKey is " << mapKey << std::endl;
-
-    static auto globalTilingData = at::empty({tilingSize * MAX_CAPTURE_NUM},
-                                             at::TensorOptions().dtype(at::kByte).device(query.options().device()));
-    if (captureMap.find(mapKey) == captureMap.end()) {
-        // std::cout << "step in, mapKey is " << mapKey << std::endl;
-        TORCH_CHECK(actualCaptureNum < MAX_CAPTURE_NUM, "lightning_indexer captureNum overflow")
-        captureMap[mapKey] = actualCaptureNum;
-        aclrtMemcpy(globalTilingData.data_ptr<uint8_t>() + actualCaptureNum * tilingSize, tilingSize, &tilingData,
+    at::Tensor tilingTensor;
+
+    if (bs > MAX_DECODE_BS) {
+        static auto preillTilingBuffer = at::empty({tilingSize}, at::TensorOptions().dtype(at::kByte).device(query.options().device()));
+        aclrtMemcpy(preillTilingBuffer.data_ptr<uint8_t>(), tilingSize, &tilingData,
                     tilingSize, ACL_MEMCPY_HOST_TO_DEVICE);
-        actualCaptureNum++;
+        tilingTensor = at::from_blob(preillTilingBuffer.data_ptr<uint8_t>(), tilingSize, at::kByte);       
+    } else {
+        uint64_t mapKey = tilingData.tilingKey;
+        mapKey = (mapKey << 32) | bs;
+
+        static auto globalTilingBuffer = at::empty({tilingSize * MAX_CAPTURE_NUM},
+                                                at::TensorOptions().dtype(at::kByte).device(query.options().device()));
+        if (captureMap.find(mapKey) == captureMap.end()) {
+            // std::cout << "step in, mapKey is " << mapKey << std::endl;
+            TORCH_CHECK(actualCaptureNum < MAX_CAPTURE_NUM, "lightning_indexer captureNum overflow")
+            captureMap[mapKey] = actualCaptureNum;
+            aclrtMemcpy(globalTilingBuffer.data_ptr<uint8_t>() + actualCaptureNum * tilingSize, tilingSize, &tilingData,
+                        tilingSize, ACL_MEMCPY_HOST_TO_DEVICE);
+            actualCaptureNum++;
+        }
+        tilingTensor =
+            at::from_blob(globalTilingBuffer.data_ptr<uint8_t>() + (tilingSize * captureMap[mapKey]), tilingSize, at::kByte);
     }
-    at::Tensor tilingTensor =
-        at::from_blob(globalTilingData.data_ptr<uint8_t>() + (tilingSize * captureMap[mapKey]), tilingSize, at::kByte);
 
     size_t userWorkspaceSize = *context->GetWorkspaceSizes(1);
     workspace = at::empty({userWorkspaceSize}, at::TensorOptions().dtype(at::kByte).device(query.options().device()));
-    //std::cout << "6" << std::endl;
     EXEC_KERNEL_CMD(lightning_indexer, blockDim, query, key, weights, actual_seq_lengths_query, actual_seq_lengths_key,
                     block_table, sparse_indices, workspace, tilingTensor);
     return sparse_indices;