feat: add layer-wise KV cache H2D copy optimization.

Author: Kang-Meng

xllm/core/framework/kv_cache/kv_cache_store.cpp

@@ -128,16 +128,14 @@ uint32_t KVCacheStore::batch_put(
     return block_transfer_info.size();
   }
 
-  uint64_t success_cnt = str_keys.size();
+  uint64_t success_cnt = block_transfer_info.size() - str_keys.size();
   auto results = client_ptr_->BatchPut(str_keys, slices, rep_config_);
 
   for (int i = 0; i < str_keys.size(); i++) {
     if (!results[i].has_value()) {
-      success_cnt = i;
-      // LOG(ERROR) << "success_cnt: " << success_cnt
-      //            << ", failed to BatchPut: " << toString(results[i].error());
       break;
     }
+    success_cnt++;
   }
   return success_cnt;
 }
@@ -179,15 +177,13 @@ uint32_t KVCacheStore::batch_get(
     return 0;
   }
 
-  uint64_t success_cnt = str_keys.size();
+  uint64_t success_cnt = 0;
   auto results = client_ptr_->BatchGet(str_keys, slices);
   for (int i = 0; i < str_keys.size(); i++) {
     if (!results[i].has_value()) {
-      success_cnt = i;
-      // LOG(ERROR) << "success_cnt: " << success_cnt
-      //            << ", failed to BatchGet: " << toString(results[i].error());
       break;
     }
+    success_cnt++;
   }
   return success_cnt;
 }

xllm/core/framework/model/model_input_params.h

@@ -122,6 +122,8 @@ struct ModelInputParams {
     // Copy graph_buffer to device
     params.graph_buffer = safe_to(graph_buffer, device, true);
 
+    params.batch_id = batch_id;
+
     return params;
   }
 
@@ -207,6 +209,8 @@ struct ModelInputParams {
 
 #if defined(USE_NPU)
   std::shared_ptr<NPULayerSynchronizerImpl> layer_synchronizer = nullptr;
+  std::shared_ptr<NPULayerSynchronizerImpl> layer_wise_load_synchronizer =
+      nullptr;
 #endif
 
   DpEpPaddingData dp_ep_padding_data;

xllm/core/platform/device.cpp

@@ -96,8 +96,8 @@ int Device::synchronize_default_stream() {
 #endif
 }
 
-std::unique_ptr<Stream> Device::get_stream_from_pool() {
-  return std::make_unique<Stream>();
+std::unique_ptr<Stream> Device::get_stream_from_pool(const int32_t timeout) {
+  return std::make_unique<Stream>(timeout);
 }
 
 }  // namespace xllm

xllm/core/platform/device.h

@@ -44,7 +44,7 @@ class Device {
   int64_t free_memory();
 
   int synchronize_default_stream();
-  std::unique_ptr<Stream> get_stream_from_pool();
+  std::unique_ptr<Stream> get_stream_from_pool(const int32_t timeout = -1);
 
  private:
   struct DeviceMem {

xllm/core/platform/npu/npu_layer_synchronizer.cpp

@@ -19,8 +19,11 @@ limitations under the License.
 
 namespace xllm {
 
-NPULayerSynchronizerImpl::NPULayerSynchronizerImpl(const int64_t num_layers)
-    : events_(num_layers, nullptr), event_record_flags_(num_layers) {
+NPULayerSynchronizerImpl::NPULayerSynchronizerImpl(const int64_t num_layers,
+                                                   const int32_t timeout)
+    : events_(num_layers, nullptr),
+      event_record_flags_(num_layers),
+      timeout_(timeout) {
   uint32_t flags = ACL_EVENT_SYNC;
   for (int64_t i = 0; i < num_layers; ++i) {
     auto ret = aclrtCreateEventWithFlag(&events_[i], flags);
@@ -45,9 +48,9 @@ std::atomic<bool>* NPULayerSynchronizerImpl::get_event_flag(
 
 bool NPULayerSynchronizerImpl::synchronize_layer(const int64_t layer_index) {
   while (!event_record_flags_[layer_index].load(std::memory_order_acquire));
-  auto ret = aclrtSynchronizeEvent(events_[layer_index]);
+  auto ret = aclrtSynchronizeEventWithTimeout(events_[layer_index], timeout_);
   if (ret != ACL_SUCCESS) {
-    LOG(ERROR) << "Synchronize event failed.";
+    LOG(ERROR) << "Synchronize event failed: " << ret;
     return false;
   }
   return true;

xllm/core/platform/npu/npu_layer_synchronizer.h

@@ -24,7 +24,8 @@ namespace xllm {
 
 class NPULayerSynchronizerImpl {
  public:
-  NPULayerSynchronizerImpl(const int64_t num_layers);
+  NPULayerSynchronizerImpl(const int64_t num_layers,
+                           const int32_t timeout = -1);
   virtual ~NPULayerSynchronizerImpl();
 
   aclrtEvent* get_event(const int64_t layer_index);
@@ -34,6 +35,7 @@ class NPULayerSynchronizerImpl {
  private:
   std::vector<aclrtEvent> events_;
   std::vector<std::atomic<bool>> event_record_flags_;
+  const int32_t timeout_;
 };
 
 }  // namespace xllm

xllm/core/platform/stream.cpp

@@ -18,14 +18,16 @@ limitations under the License.
 namespace xllm {
 
 #if defined(USE_NPU)
-Stream::Stream() : stream_(c10_npu::getNPUStreamFromPool()) {}
+Stream::Stream(const int32_t timeout)
+    : stream_(c10_npu::getNPUStreamFromPool()), timeout_(timeout) {}
 #elif defined(USE_MLU)
-Stream::Stream() : stream_(torch_mlu::getStreamFromPool()) {}
+Stream::Stream(const int32_t timeout)
+    : stream_(torch_mlu::getStreamFromPool()), timeout_(timeout) {}
 #endif
 
 int Stream::synchronize() const {
 #if defined(USE_NPU)
-  return aclrtSynchronizeStream(stream_.stream());
+  return aclrtSynchronizeStreamWithTimeout(stream_.stream(), timeout_);
 #elif defined(USE_MLU)
   stream_.unwrap().synchronize();
   return 0;

xllm/core/platform/stream.h

@@ -34,7 +34,7 @@ namespace xllm {
 
 class Stream {
  public:
-  Stream();
+  Stream(const int32_t timeout = -1);
   ~Stream() = default;
 
   Stream(const Stream&) = delete;
@@ -44,13 +44,19 @@ class Stream {
 
   int synchronize() const;
   c10::StreamGuard set_stream_guard() const;
+#if defined(USE_NPU)
+  c10_npu::NPUStream* get_stream() { return &stream_; }
+#elif defined(USE_MLU)
+  torch_mlu::MLUStream* get_stream() { return &stream_; }
+#endif
 
  private:
 #if defined(USE_NPU)
   c10_npu::NPUStream stream_;
 #elif defined(USE_MLU)
   torch_mlu::MLUStream stream_;
 #endif
+  const int32_t timeout_;
 };
 
 }  // namespace xllm