feat: support async layer wise batch copy.

Author: Kang-Meng

xllm/core/common/global_flags.cpp

@@ -168,6 +168,12 @@ DEFINE_uint32(prefetch_timeout,
               0,
               "Prefetch timeout for prefetch from kv cache store.");
 
+DEFINE_uint32(prefetch_bacth_size,
+              2,
+              "Prefetch from kvcache store copy batch size.");
+
+DEFINE_uint32(layers_wise_copy_batchs, 4, "Layer wise H2D copy batchs.");
+
 // --- parallel config ---
 
 DEFINE_int32(dp_size, 1, "Data parallel size for MLA attention.");

xllm/core/common/global_flags.h

@@ -155,6 +155,10 @@ DECLARE_int32(max_decode_token_per_sequence);
 
 DECLARE_uint32(prefetch_timeout);
 
+DECLARE_uint32(prefetch_bacth_size);
+
+DECLARE_uint32(layers_wise_copy_batchs);
+
 DECLARE_string(priority_strategy);
 
 DECLARE_bool(enable_online_preempt_offline);

xllm/core/common/options.cpp

@@ -54,6 +54,8 @@ std::string Options::to_string() const {
      << ", enable_cache_upload: " << enable_cache_upload()
      << ", enable_kvcache_store: " << enable_kvcache_store()
      << ", prefetch_timeout: " << prefetch_timeout()
+     << ", prefetch_bacth_size: " << prefetch_bacth_size()
+     << ", layers_wise_copy_batchs: " << layers_wise_copy_batchs()
      << ", store_protocol: " << store_protocol()
      << ", store_master_server_address: " << store_master_server_address()
      << ", store_metadata_server: " << store_metadata_server()

xllm/core/common/options.h

@@ -192,6 +192,12 @@ class Options {
 
   // Prefetch timeout for prefetch from kv cache store
   PROPERTY(uint32_t, prefetch_timeout) = 0;
+
+  // Prefetch from kvcache store copy batch size
+  PROPERTY(uint32_t, prefetch_bacth_size) = 2;
+
+  // Layer wise H2D copy batchs
+  PROPERTY(uint32_t, layers_wise_copy_batchs) = 4;
 };
 
 }  // namespace xllm

xllm/core/distributed_runtime/worker_service.cpp

@@ -33,8 +33,6 @@ limitations under the License.
 
 namespace xllm {
 
-constexpr uint32_t COPY_BATCH_SIZE = 1;
-
 WorkerService::WorkerService(runtime::Options options,
                              const torch::Device& device)
     : options_(options), device_(device), initialized_(false) {
@@ -477,21 +475,23 @@ void WorkerService::PrefetchFromStorage(
         auto close_future = stream_handler->get_close_future();
         bool is_completed = false;
 
-        for (size_t i = 0; i < transfer_slice.size(); i += COPY_BATCH_SIZE) {
-          auto current_slice = transfer_slice.slice(
-              i, std::min(i + COPY_BATCH_SIZE, transfer_slice.size()));
+        for (size_t i = 0; i < transfer_slice.size();
+             i += options_.prefetch_bacth_size()) {
+          auto current_slice =
+              transfer_slice.slice(i,
+                                   std::min(i + options_.prefetch_bacth_size(),
+                                            transfer_slice.size()));
 
           auto success_cnt = worker_->prefetch_from_storage(current_slice);
 
           if (success_cnt != current_slice.size() ||
-              i + COPY_BATCH_SIZE >= transfer_slice.size()) {
+              i + options_.prefetch_bacth_size() >= transfer_slice.size()) {
             is_completed = true;
           }
 
           butil::IOBuf buf;
           buf.append(std::to_string(success_cnt));
           if (brpc::StreamWrite(*stream_id.get(), buf) != 0) {
-            is_completed = false;
             break;
           }
 

xllm/core/framework/request/sequence.cpp

@@ -382,6 +382,7 @@ void Sequence::reset() {
   kv_state_.reset();
   host_kv_state_.reset();
   timer_.reset();
+  is_timeout_set_ = false;
   volatile_num_prompt_tokens_ = num_tokens_;
 }
 
@@ -462,12 +463,13 @@ bool Sequence::update_prefetch_result(uint32_t timeout) {
   }
 
   if (timeout != 0 && !termination_flag_.load(std::memory_order_acquire)) {
-    if (timer_ != nullptr) {
-      timer_ = std::make_shared<Timer>();
+    if (!is_timeout_set_) {
+      timer_.reset();
+      is_timeout_set_ = true;
       return false;
     }
 
-    if (timer_->elapsed_milliseconds() < timeout) {
+    if (timer_.elapsed_milliseconds() < timeout) {
       return false;
     }
   }

xllm/core/framework/request/sequence.h

@@ -364,7 +364,8 @@ class Sequence final {
   std::atomic<bool> termination_flag_{false};
   std::vector<std::shared_ptr<std::atomic<uint32_t>>> prefetch_results_;
 
-  std::shared_ptr<Timer> timer_ = nullptr;
+  Timer timer_;
+  bool is_timeout_set_ = false;
 };
 
 }  // namespace xllm

xllm/core/kernels/npu/xllm_ops/top_k_top_p.h

@@ -20,7 +20,7 @@ limitations under the License.
 #include <vector>
 
 #include "acl/acl.h"
-#include "aclnn_apply_top_k_top_p.h"
+#include "aclnnop/aclnn_apply_top_k_top_p.h"
 #include "acltensor_utils.h"
 #include "util/tensor_helper.h"
 

xllm/core/runtime/master.cpp

@@ -218,6 +218,8 @@ Master::Master(const Options& options, EngineType type) : options_(options) {
         .store_master_server_address(options_.store_master_server_address())
         .store_metadata_server(options_.store_metadata_server())
         .store_local_hostname(options_.store_local_hostname())
+        .prefetch_bacth_size(options_.prefetch_bacth_size())
+        .layers_wise_copy_batchs(options_.layers_wise_copy_batchs())
         .enable_continuous_kvcache(options_.enable_continuous_kvcache())
         .enable_offline_inference(options_.enable_offline_inference())
         .spawn_worker_path(options_.spawn_worker_path())

xllm/core/runtime/options.h

@@ -158,6 +158,12 @@ struct Options {
   //  value used if port is not included)
   PROPERTY(std::string, store_local_hostname) = "";
 
+  // Prefetch from kvcache store copy batch size
+  PROPERTY(uint32_t, prefetch_bacth_size) = 2;
+
+  // Layer wise H2D copy batchs
+  PROPERTY(uint32_t, layers_wise_copy_batchs) = 4;
+
   // dit
   // max requests per batch
   PROPERTY(int, max_requests_per_batch) = 0;