feat: optimize layer wise copy.

Author: Kang-Meng

xllm/core/common/global_flags.cpp

@@ -164,6 +164,10 @@ DEFINE_int32(
     256,
     "Max decode token per sequence which used for ZeroEvictionScheduler.");
 
+DEFINE_uint32(prefetch_timeout,
+              0,
+              "Prefetch timeout for prefetch from kv cache store.");
+
 // --- parallel config ---
 
 DEFINE_int32(dp_size, 1, "Data parallel size for MLA attention.");

xllm/core/common/global_flags.h

@@ -153,6 +153,8 @@ DECLARE_bool(use_zero_evict);
 
 DECLARE_int32(max_decode_token_per_sequence);
 
+DECLARE_uint32(prefetch_timeout);
+
 DECLARE_string(priority_strategy);
 
 DECLARE_bool(enable_online_preempt_offline);

xllm/core/common/options.cpp

@@ -53,6 +53,7 @@ std::string Options::to_string() const {
      << ", enable_service_routing: " << enable_service_routing()
      << ", enable_cache_upload: " << enable_cache_upload()
      << ", enable_kvcache_store: " << enable_kvcache_store()
+     << ", prefetch_timeout: " << prefetch_timeout()
      << ", store_protocol: " << store_protocol()
      << ", store_master_server_address: " << store_master_server_address()
      << ", store_metadata_server: " << store_metadata_server()

xllm/core/common/options.h

@@ -189,6 +189,9 @@ class Options {
 
   // whether the worker and master are on the same machine.
   PROPERTY(bool, is_local) = false;
+
+  // Prefetch timeout for prefetch from kv cache store
+  PROPERTY(uint32_t, prefetch_timeout) = 0;
 };
 
 }  // namespace xllm

xllm/core/distributed_runtime/worker_service.cpp

@@ -491,7 +491,6 @@ void WorkerService::PrefetchFromStorage(
           butil::IOBuf buf;
           buf.append(std::to_string(success_cnt));
           if (brpc::StreamWrite(*stream_id.get(), buf) != 0) {
-            brpc::StreamClose(*stream_id.get());
             is_completed = false;
             break;
           }

xllm/core/framework/request/sequence.cpp

@@ -381,7 +381,7 @@ void Sequence::add_host_kv_blocks(const std::vector<Block>& blocks) {
 void Sequence::reset() {
   kv_state_.reset();
   host_kv_state_.reset();
-  timeout_checker_.reset();
+  timer_.reset();
   volatile_num_prompt_tokens_ = num_tokens_;
 }
 
@@ -456,14 +456,20 @@ Slice<int32_t> Sequence::get_generated_tokens() const {
   return {tokens_.data(), 0};
 }
 
-bool Sequence::update_prefetch_result() {
+bool Sequence::update_prefetch_result(uint32_t timeout) {
   if (prefetch_results_.empty()) {
     return true;
   }
 
-  if (!termination_flag_.load(std::memory_order_acquire) &&
-      timeout_checker_.check_timeout()) {
-    return false;
+  if (timeout != 0 && !termination_flag_.load(std::memory_order_acquire)) {
+    if (timer_ != nullptr) {
+      timer_ = std::make_shared<Timer>();
+      return false;
+    }
+
+    if (timer_->elapsed_milliseconds() < timeout) {
+      return false;
+    }
   }
 
   termination_flag_.store(true, std::memory_order_release);

xllm/core/framework/request/sequence.h

@@ -35,6 +35,7 @@ limitations under the License.
 #include "sequence_kv_state.h"
 #include "sequence_logprob_state.h"
 #include "stopping_checker.h"
+#include "util/timer.h"
 
 namespace xllm {
 
@@ -82,44 +83,6 @@ struct SequenceParams {
   StoppingChecker* stopping_checker;  // not owned
 };
 
-static uint32_t timeout_ms = 0;
-class TimeoutChecker {
- private:
-  std::chrono::steady_clock::time_point timeout_start_;
-  bool is_timeout_set_ = false;
-
- public:
-  TimeoutChecker() { init(); }
-
-  bool check_timeout() {
-    if (!is_timeout_set_) {
-      timeout_start_ = std::chrono::steady_clock::now();
-      is_timeout_set_ = true;
-
-      return false;
-    } else {
-      auto now = std::chrono::steady_clock::now();
-      auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
-          now - timeout_start_);
-
-      return elapsed.count() >= timeout_ms;
-    }
-  }
-
-  void reset() { is_timeout_set_ = false; }
-
- private:
-  static void init_timeout() {
-    const char* env_str = std::getenv("PREFETCH_TIMEOUT_MS");
-    timeout_ms = env_str ? std::strtoul(env_str, nullptr, 10) : 0;
-    LOG(INFO) << "Prefetch timeout set as: " << timeout_ms;
-  }
-  static void init() {
-    static std::once_flag flag_;
-    std::call_once(flag_, init_timeout);
-  }
-};
-
 class Sequence final {
  public:
   Sequence(size_t index,
@@ -286,7 +249,7 @@ class Sequence final {
     return &prefetch_results_;
   }
 
-  bool update_prefetch_result();
+  bool update_prefetch_result(uint32_t timeout = 30);
 
   void reset();
 
@@ -401,7 +364,7 @@ class Sequence final {
   std::atomic<bool> termination_flag_{false};
   std::vector<std::shared_ptr<std::atomic<uint32_t>>> prefetch_results_;
 
-  TimeoutChecker timeout_checker_;
+  std::shared_ptr<Timer> timer_ = nullptr;
 };
 
 }  // namespace xllm

xllm/core/platform/npu/npu_layer_synchronizer.h

@@ -31,6 +31,7 @@ class NPULayerSynchronizerImpl {
   aclrtEvent* get_event(const int64_t layer_index);
   std::atomic<bool>* get_event_flag(const int64_t layer_index);
   bool synchronize_layer(const int64_t layer_index);
+  uint32_t get_event_size() { return events_.size(); };
 
  private:
   std::vector<aclrtEvent> events_;

xllm/core/runtime/llm_master.cpp

@@ -95,7 +95,8 @@ LLMMaster::LLMMaster(const Options& options)
       .disable_ttft_profiling(options_.disable_ttft_profiling())
       .enable_forward_interruption(options_.enable_forward_interruption())
       .max_global_ttft_ms(options_.max_global_ttft_ms())
-      .max_global_tpot_ms(options_.max_global_tpot_ms());
+      .max_global_tpot_ms(options_.max_global_tpot_ms())
+      .prefetch_timeout(options_.prefetch_timeout());
   scheduler_ = create_continuous_scheduler(engine_.get(), scheduler_options);
 
   if (options_.enable_service_routing()) {

xllm/core/runtime/worker_impl.cpp

@@ -638,6 +638,7 @@ bool WorkerImpl::init_model(const std::string& model_weights_path,
   if (!status) {
     return false;
   }
+  layers_per_copy_ = context_.get_model_args().n_layers() / 4;
 
   this->load_model(std::move(model_loader));
 
@@ -898,9 +899,14 @@ bool WorkerImpl::h2d_batch_copy(const uint64_t batch_id,
   }
 
   const int64_t num_layers = context_.get_model_args().n_layers();
+  uint32_t layers_per_copy = layers_per_copy_;
   uint32_t num_batches = block_transfer_info.size() * 2;
+  while (num_batches * layers_per_copy > BATCH_COPY_MAX_SIZE) {
+    layers_per_copy--;
+  }
 
-  auto synchronizer = std::make_shared<NPULayerSynchronizerImpl>(num_layers);
+  uint32_t copy_cnt = (num_layers + layers_per_copy - 1) / layers_per_copy;
+  auto synchronizer = std::make_shared<NPULayerSynchronizerImpl>(copy_cnt);
   {
     std::lock_guard<std::mutex> lock(mutex_);
     if (layer_wise_load_synchronizer_.count(batch_id) != 0) {
@@ -909,47 +915,54 @@ bool WorkerImpl::h2d_batch_copy(const uint64_t batch_id,
     layer_wise_load_synchronizer_[batch_id] = synchronizer;
   }
 
-  void** srcs = new void*[num_batches];
-  void** dsts = new void*[num_batches];
-  size_t* copy_size = new size_t[num_batches];
   aclrtMemcpyBatchAttr attrs[1] = {h2d_attrs_};
   size_t attrs_indexes[1] = {0};
 
   std::unique_ptr<Stream> stream;
   copy_stream_.wait_dequeue(stream);
   c10::StreamGuard streamGuard = stream->set_stream_guard();
-
   aclError ret = 0;
 
-  for (int layer_id = 0; layer_id < num_layers; layer_id++) {
-    auto dst_k_cache = kv_caches_.at(layer_id).get_k_cache();
-    auto dst_v_cache = kv_caches_.at(layer_id).get_v_cache();
+  void** srcs = new void*[num_batches * layers_per_copy];
+  void** dsts = new void*[num_batches * layers_per_copy];
+  size_t* copy_size = new size_t[num_batches * layers_per_copy];
+
+  for (int index = 0; index < copy_cnt; index++) {
+    int layer_id = index * layers_per_copy;
     size_t fail_index = 0;
     uint32_t curr_index = 0;
-    auto* event = synchronizer->get_event(layer_id);
-    auto* event_flag = synchronizer->get_event_flag(layer_id);
+    uint32_t layer_cnt = 0;
 
-    for (const auto& info : block_transfer_info) {
-      auto src_k_cache = host_kv_caches_.at(info.src_block_id).get_k_cache();
-      auto src_v_cache = host_kv_caches_.at(info.src_block_id).get_v_cache();
+    while (layer_id < (index + 1) * layers_per_copy && layer_id < num_layers) {
+      auto dst_k_cache = kv_caches_.at(layer_id).get_k_cache();
+      auto dst_v_cache = kv_caches_.at(layer_id).get_v_cache();
 
-      srcs[curr_index] = src_k_cache[layer_id].data_ptr();
-      dsts[curr_index] = dst_k_cache[info.dst_block_id].data_ptr();
-      copy_size[curr_index] = key_cache_size_per_layer_;
-      curr_index++;
+      for (const auto& info : block_transfer_info) {
+        auto src_k_cache = host_kv_caches_.at(info.src_block_id).get_k_cache();
+        auto src_v_cache = host_kv_caches_.at(info.src_block_id).get_v_cache();
 
-      srcs[curr_index] = src_v_cache[layer_id].data_ptr();
-      dsts[curr_index] = dst_v_cache[info.dst_block_id].data_ptr();
-      copy_size[curr_index] = value_cache_size_per_layer_;
-      curr_index++;
+        srcs[curr_index] = src_k_cache[layer_id].data_ptr();
+        dsts[curr_index] = dst_k_cache[info.dst_block_id].data_ptr();
+        copy_size[curr_index] = key_cache_size_per_layer_;
+        curr_index++;
+
+        srcs[curr_index] = src_v_cache[layer_id].data_ptr();
+        dsts[curr_index] = dst_v_cache[info.dst_block_id].data_ptr();
+        copy_size[curr_index] = value_cache_size_per_layer_;
+        curr_index++;
+      }
+      layer_id++;
+      layer_cnt++;
     }
 
     // TODO(kangmeng): change to async API
+    CHECK(layer_cnt <= layers_per_copy)
+        << "layer_cnt should less equal to layers_per_copy.";
     ret = aclrtMemcpyBatch(dsts,
                            copy_size,
                            srcs,
                            copy_size,
-                           num_batches,
+                           num_batches * layer_cnt,
                            attrs,
                            attrs_indexes,
                            1,
@@ -959,11 +972,13 @@ bool WorkerImpl::h2d_batch_copy(const uint64_t batch_id,
       LOG(ERROR) << "aclrtMemcpyBatch error: " << ret
                  << ", fail_index:" << fail_index;
     } else {
+      auto* event = synchronizer->get_event(index);
       ret = aclrtRecordEvent(*event, stream->get_stream()->stream());
       if (ret != 0) {
         LOG(ERROR) << "aclrtRecordEvent error: " << ret;
       }
     }
+    auto* event_flag = synchronizer->get_event_flag(index);
     event_flag->store(true, std::memory_order_release);
     if (ret != 0) break;
   }