[feat] Add pcstore for enhanced PrefixCache performance (#393)

# Purpose
Add `pcstore` for enhanced PrefixCache performance
# Modifications 
- Data is paged out to Host and asynchronously written to SSD, freeing HBM space earlier.
- Reads and writes are aggregated at Block-level granularity to increase SSD I/O size and improve performance.
- For MLA models, data is loaded once from SSD and shared across Devices in DRAM, reducing SSD bandwidth pressure.
# Test
- `ucm/store/test/e2e/pcstore_embed.py`
- `ucm/store/test/e2e/pcstore_fetch.py`

Author: FangRun2

ucm/shared/CMakeLists.txt

@@ -1,4 +1,3 @@
-include_directories(${CMAKE_CURRENT_SOURCE_DIR})
 add_subdirectory(vendor)
 add_subdirectory(trans)
 add_subdirectory(test)

ucm/shared/trans/CMakeLists.txt

@@ -7,6 +7,7 @@ endif()
 if(RUNTIME_ENVIRONMENT STREQUAL "simu")
     add_subdirectory(simu)
 endif()
+target_include_directories(trans PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/..)
 
 file(GLOB_RECURSE UCMTRANS_CPY_SOURCE_FILES "./cpy/*.cc")
 pybind11_add_module(ucmtrans ${UCMTRANS_CPY_SOURCE_FILES})

ucm/shared/trans/ascend/ascend_buffer.cc

@@ -42,35 +42,15 @@ std::shared_ptr<void> Trans::AscendBuffer::MakeHostBuffer(size_t size)
     return nullptr;
 }
 
-Status Trans::AscendBuffer::RegisterHostBuffer(void* ptr, size_t size)
+Status Buffer::RegisterHostBuffer(void* host, size_t size, void** pDevice)
 {
-    if (registerOnHost_) {
-        aclrtHostUnregister(registerOnHost_);
-        registerOnHost_ = nullptr;
-        registerOnDevice_ = nullptr;
-    }
-    auto ret = aclrtHostRegister(ptr, size, ACL_HOST_REGISTER_MAPPED, &registerOnDevice_);
-    if (ret == ACL_SUCCESS) {
-        registerOnHost_ = ptr;
-        return Status::OK();
-    }
-    registerOnDevice_ = nullptr;
-    return Status{ret, std::to_string(ret)};
-}
-
-void Trans::AscendBuffer::UnregisterHostBuffer(void* ptr)
-{
-    aclrtHostUnregister(ptr);
-    if (registerOnHost_ == ptr) {
-        registerOnHost_ = nullptr;
-        registerOnDevice_ = nullptr;
-    }
+    void* device = nullptr;
+    auto ret = aclrtHostRegister(host, size, ACL_HOST_REGISTER_MAPPED, &device);
+    if (ret != ACL_SUCCESS) [[unlikely]] { return Status{ret, std::to_string(ret)}; }
+    if (pDevice) { *pDevice = device; }
+    return Status::OK();
 }
 
-void* Trans::AscendBuffer::GetHostPtrOnDevice(void* ptr)
-{
-    if (registerOnHost_ == ptr) { return registerOnDevice_; }
-    return nullptr;
-}
+void Buffer::UnregisterHostBuffer(void* host) { aclrtHostUnregister(host); }
 
 } // namespace UC::Trans

ucm/shared/trans/ascend/ascend_buffer.h

@@ -29,16 +29,9 @@
 namespace UC::Trans {
 
 class AscendBuffer : public ReservedBuffer {
-    void* registerOnHost_{nullptr};
-    void* registerOnDevice_{nullptr};
-
 public:
     std::shared_ptr<void> MakeDeviceBuffer(size_t size) override;
     std::shared_ptr<void> MakeHostBuffer(size_t size) override;
-
-    Status RegisterHostBuffer(void* ptr, size_t size) override;
-    void UnregisterHostBuffer(void* ptr) override;
-    void* GetHostPtrOnDevice(void* ptr) override;
 };
 
 } // namespace UC::Trans

ucm/shared/trans/buffer.h

@@ -41,9 +41,8 @@ class Buffer {
     virtual Status MakeHostBuffers(size_t size, size_t number) = 0;
     virtual std::shared_ptr<void> GetHostBuffer(size_t size) = 0;
 
-    virtual Status RegisterHostBuffer(void* ptr, size_t size) = 0;
-    virtual void UnregisterHostBuffer(void* ptr) = 0;
-    virtual void* GetHostPtrOnDevice(void* ptr) = 0;
+    static Status RegisterHostBuffer(void* host, size_t size, void** pDevice = nullptr);
+    static void UnregisterHostBuffer(void* host);
 };
 
 } // namespace UC::Trans

ucm/shared/trans/cuda/cuda_buffer.cc

@@ -42,21 +42,17 @@ std::shared_ptr<void> CudaBuffer::MakeHostBuffer(size_t size)
     return nullptr;
 }
 
-Status CudaBuffer::RegisterHostBuffer(void* ptr, size_t size)
+Status Buffer::RegisterHostBuffer(void* host, size_t size, void** pDevice)
 {
-    auto ret = cudaHostRegister(ptr, size, cudaHostRegisterDefault);
-    if (ret == cudaSuccess) { return Status::OK(); }
-    return Status{ret, cudaGetErrorString(ret)};
+    auto ret = cudaHostRegister(host, size, cudaHostRegisterDefault);
+    if (ret != cudaSuccess) [[unlikely]] { return Status{ret, cudaGetErrorString(ret)}; }
+    if (pDevice) {
+        ret = cudaHostGetDevicePointer(pDevice, host, 0);
+        if (ret != cudaSuccess) [[unlikely]] { return Status{ret, cudaGetErrorString(ret)}; }
+    }
+    return Status::OK();
 }
 
-void CudaBuffer::UnregisterHostBuffer(void* ptr) { cudaHostUnregister(ptr); }
-
-void* CudaBuffer::GetHostPtrOnDevice(void* ptr)
-{
-    void* device = nullptr;
-    auto ret = cudaHostGetDevicePointer(&device, ptr, 0);
-    if (ret == cudaSuccess) { return nullptr; }
-    return device;
-}
+void Buffer::UnregisterHostBuffer(void* host) { cudaHostUnregister(host); }
 
 } // namespace UC::Trans

ucm/shared/trans/cuda/cuda_buffer.h

@@ -32,10 +32,6 @@ class CudaBuffer : public ReservedBuffer {
 public:
     std::shared_ptr<void> MakeDeviceBuffer(size_t size) override;
     std::shared_ptr<void> MakeHostBuffer(size_t size) override;
-
-    Status RegisterHostBuffer(void* ptr, size_t size) override;
-    void UnregisterHostBuffer(void* ptr) override;
-    void* GetHostPtrOnDevice(void* ptr) override;
 };
 
 } // namespace UC::Trans

ucm/shared/trans/simu/simu_buffer.cc

@@ -67,10 +67,12 @@ std::shared_ptr<void> SimuBuffer::MakeHostBuffer(size_t size)
     return std::shared_ptr<void>(device, FreeMemory);
 }
 
-Status SimuBuffer::RegisterHostBuffer(void* ptr, size_t size) { return Status::OK(); }
-
-void SimuBuffer::UnregisterHostBuffer(void* ptr) {}
+Status Buffer::RegisterHostBuffer(void* host, size_t size, void** pDevice)
+{
+    if (pDevice) { *pDevice = host; }
+    return Status::OK();
+}
 
-void* SimuBuffer::GetHostPtrOnDevice(void* ptr) { return ptr; }
+void Buffer::UnregisterHostBuffer(void* host) {}
 
 } // namespace UC::Trans

ucm/shared/trans/simu/simu_buffer.h

@@ -32,10 +32,6 @@ class SimuBuffer : public ReservedBuffer {
 public:
     std::shared_ptr<void> MakeDeviceBuffer(size_t size) override;
     std::shared_ptr<void> MakeHostBuffer(size_t size) override;
-
-    Status RegisterHostBuffer(void* ptr, size_t size) override;
-    void UnregisterHostBuffer(void* ptr) override;
-    void* GetHostPtrOnDevice(void* ptr) override;
 };
 
 } // namespace UC::Trans

ucm/sparse/gsa/prefetch/include/kvcache_pre.h

@@ -37,10 +37,10 @@ namespace ucmprefetch
         int bsIndex;
     } PrefetchReqInfo;
 
-    class ThreadPool 
+    class ThreadPool
     {
     public:
-        static ThreadPool *GetInst() 
+        static ThreadPool *GetInst()
         {
             static ThreadPool pool(1);
             return &pool;
@@ -52,7 +52,7 @@ namespace ucmprefetch
         auto enqueue(F&& f, Args&&... args) -> std::future<typename std::result_of<F(Args...)>::type>;
 
         size_t GetActiveThreads() const;
-    
+
     private:
         explicit ThreadPool(size_t threadCount);
         std::vector<std::thread> workers;
@@ -66,7 +66,7 @@ namespace ucmprefetch
 
     void MutliBSThreadFun(void *args);
 
-    class __attribute__((visibility("hidden"))) GSAPrefetchEngineC 
+    class __attribute__((visibility("hidden"))) GSAPrefetchEngineC
     {
     private:
         std::map<std::string, std::vector<std::map<int, int>>> mDocsTables;
@@ -95,7 +95,7 @@ namespace ucmprefetch
         std::map<std::string, std::vector<std::vector<int>>> allNeedLoadBlock;
         std::map<std::string, std::vector<std::vector<int>>> allMissIdxs;
         std::map<std::string, int> mPromptLen;
-        UC::CCStore *mStore = nullptr;
+        UC::CCStore<> *mStore = nullptr;
         std::vector<torch::Tensor> mKvCaches;
         uint32_t mBlockSize = 128;
         uint32_t mTensorElemSize = 2; // fp16
@@ -108,21 +108,21 @@ namespace ucmprefetch
     public:
         std::mutex mMutex;
         bool mStopPrefetch = false;
-    
+
     private:
         void LoadKVToHBM(std::vector<int> loadNPUBlockIDs,
             std::vector<int> missIdxs, int layerID, std::string reqID);
-        
+
         void GetHitAndMissBlock(PrefetchReqInfo oneBsInfo,
             std::unordered_set<int> &hitBlocks,
             std::map<int, int> &hitBlocksIdx,
             std::vector<int> &missIdxs);
-        
+
         void RunPrefetchH2D(PrefetchReqInfo oneBsInfo,
             std::unordered_set<int> &hitBlocks,
             std::map<int, int> &hitBlocksIdx,
             std::vector<int> &missIdxs);
-        
+
         void RunOneBsPrefetch(std::string reqID, int topkLen,
             int bsIndex, int topkIndex);
 
@@ -144,7 +144,7 @@ namespace ucmprefetch
         void SetBlocksMap(std::string reqID, std::vector<int> &blockTableList,
             std::vector<int> &selectIndex, std::vector<std::string> &blocksHash,
             int maxIdx);
-        
+
         void SetBlocksMapMultiLayer(std::string reqID,
             std::vector<std::map<int, int>> &remainMap,
             std::vector<std::map<int, int>> &prefetchMap,
@@ -168,7 +168,7 @@ namespace ucmprefetch
             std::vector<int> &bsIndexInput,
             std::vector<torch::Tensor> &kvCaches,
             void *storePtr);
-        
+
         int CallPrefetchProcessFun();
 
         void PrintMap(std::string reqID, int i);
@@ -178,7 +178,7 @@ namespace ucmprefetch
         void SetPrefetchStatus(bool flag);
 
         void SetModelRunningStatus(bool flag);
-        
+
         size_t GetOffset(uint32_t layerID, bool isV);
 
         std::map<std::string, std::vector<std::vector<int>>> ObtainLoadBlocks();
@@ -189,7 +189,7 @@ namespace ucmprefetch
 
         std::map<std::string, std::vector<std::map<int, int>>> ObtainDocsMap();
     };
-    
+
 } // namespace uc
 
 #endif