feat: optimize layer wise copy.

Kang-Meng · Kang-Meng · commit f52f4dd1d11e · 2025-12-08T10:57:18.000+08:00
diff --git a/xllm/core/common/global_flags.cpp b/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.");
diff --git a/xllm/core/common/global_flags.h b/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);
diff --git a/xllm/core/common/options.cpp b/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()
diff --git a/xllm/core/common/options.h b/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
diff --git a/xllm/core/distributed_runtime/worker_service.cpp b/xllm/core/distributed_runtime/worker_service.cpp
@@ -489,7 +489,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;
           }
diff --git a/xllm/core/framework/request/sequence.cpp b/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);
diff --git a/xllm/core/framework/request/sequence.h b/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
diff --git a/xllm/core/platform/npu/npu_layer_synchronizer.h b/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_;
diff --git a/xllm/core/runtime/llm_master.cpp b/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()) {
diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp
@@ -614,6 +614,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));
 
@@ -874,9 +875,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) {
@@ -885,47 +891,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,
@@ -935,11 +948,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;
   }
diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h
@@ -266,6 +266,7 @@ class WorkerImpl {
 
   uint64_t key_cache_size_per_layer_;
   uint64_t value_cache_size_per_layer_;
+  uint32_t layers_per_copy_;
 
   bool is_spec_draft_ = false;
 
diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp
@@ -214,7 +214,8 @@ void ContinuousScheduler::handle_prefill_requests(
 
     bool prefetch_result = true;
     for (auto& prefill_sequence : request->sequences()) {
-      prefetch_result &= prefill_sequence->update_prefetch_result();
+      prefetch_result &=
+          prefill_sequence->update_prefetch_result(options_.prefetch_timeout());
     }
 
     if (!prefetch_result) {
diff --git a/xllm/core/scheduler/continuous_scheduler.h b/xllm/core/scheduler/continuous_scheduler.h
@@ -117,6 +117,8 @@ class ContinuousScheduler : public Scheduler {
     PROPERTY(int32_t, max_global_ttft_ms) = std::numeric_limits<int32_t>::max();
     // all requests use single global tpot
     PROPERTY(int32_t, max_global_tpot_ms) = std::numeric_limits<int32_t>::max();
+    // Prefetch timeout for prefetch from kv cache store
+    PROPERTY(uint32_t, prefetch_timeout) = 0;
   };
 
   ContinuousScheduler(Engine* engine, const Options& options);
diff --git a/xllm/core/scheduler/prefill_only_scheduler.cpp b/xllm/core/scheduler/prefill_only_scheduler.cpp
@@ -87,7 +87,8 @@ void PrefillOnlyScheduler::handle_prefill_requests(
 
     bool prefetch_result = true;
     for (auto& prefill_sequence : request->sequences()) {
-      prefetch_result &= prefill_sequence->update_prefetch_result();
+      prefetch_result &=
+          prefill_sequence->update_prefetch_result(options_.prefetch_timeout());
     }
 
     if (!prefetch_result) {
diff --git a/xllm/models/llm/qwen3.h b/xllm/models/llm/qwen3.h
@@ -165,6 +165,13 @@ class QWen3ModelImpl : public LlmModelImplBase<QWen3DecoderLayer> {
     }
 
 #if defined(USE_NPU)
+    uint32_t layers_per_copy = layers_.size();
+    if (input_params.layer_wise_load_synchronizer != nullptr) {
+      uint32_t event_cnt =
+          input_params.layer_wise_load_synchronizer->get_event_size();
+      layers_per_copy = layers_per_copy / event_cnt +
+                        uint32_t(layers_per_copy % event_cnt == 0);
+    }
     for (size_t i = 0; i < layers_.size(); i++) {
       aclrtEvent* event{nullptr};
       std::atomic<bool>* event_flag{nullptr};
@@ -173,8 +180,9 @@ class QWen3ModelImpl : public LlmModelImplBase<QWen3DecoderLayer> {
         event = input_params.layer_synchronizer->get_event(i);
         event_flag = input_params.layer_synchronizer->get_event_flag(i);
       }
-      if (input_params.layer_wise_load_synchronizer != nullptr) {
-        if (!input_params.layer_wise_load_synchronizer->synchronize_layer(i)) {
+      if (layers_per_copy != layers_.size() && i % layers_per_copy == 0) {
+        if (!input_params.layer_wise_load_synchronizer->synchronize_layer(
+                i / layers_per_copy)) {
           return torch::Tensor();
         }
       }
diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp
@@ -181,6 +181,7 @@ int run() {
       .enable_kvcache_store(FLAGS_enable_kvcache_store &&
                             FLAGS_enable_prefix_cache &&
                             (FLAGS_host_blocks_factor > 0.0))
+      .prefetch_timeout(FLAGS_prefetch_timeout)
       .store_protocol(FLAGS_store_protocol)
       .store_master_server_address(FLAGS_store_master_server_address)
       .store_metadata_server(FLAGS_store_metadata_server)

Original file line number	Diff line number	Diff line change
`@@ -489,7 +489,6 @@ void WorkerService::PrefetchFromStorage(`
`489`	`489`	`butil::IOBuf buf;`
`490`	`490`	`buf.append(std::to_string(success_cnt));`
`491`	`491`	`if (brpc::StreamWrite(*stream_id.get(), buf) != 0) {`
`492`		`- brpc::StreamClose(*stream_id.get());`
`493`	`492`	`is_completed = false;`
`494`	`493`	`break;`
`495`	`494`	`}`