From ef4b7d4afb0eac9caf4aace36f7ccbfe8c86e675 Mon Sep 17 00:00:00 2001
From: rubik <huajianfeng@zhejianglab.com>
Date: Mon, 6 Jul 2026 17:01:36 +0800
Subject: [PATCH] [performance] InferEngine: drastically improve initialization
 speed

- Replace sync `set_zeros` with async `set_zeros_device_async` to eliminate CPU buffer allocation and H2D overhead.
- Decouple RankWorker creation from initialization wait to allow parallel startup.
- Add explicit synchronization after KV cache allocation to ensure readiness.
---
 csrc/cache/kv_cache.cpp        | 4 +---
 csrc/engine/infer_engine.cpp   | 4 ++++
 csrc/engine/rank_worker.cpp    | 9 +++++++--
 csrc/engine/rank_worker.hpp    | 2 ++
 csrc/models/infinilm_model.cpp | 1 +
 5 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/csrc/cache/kv_cache.cpp b/csrc/cache/kv_cache.cpp
index 92734c05f..a57db544c 100644
--- a/csrc/cache/kv_cache.cpp
+++ b/csrc/cache/kv_cache.cpp
@@ -136,9 +136,7 @@ infinicore::Tensor create_layer_kv_cache(
         kv_shape,
         dtype,
         rank_info.device);
-    set_zeros(kv_cache);
-
-    infinicore::context::syncStream();
+    set_zeros_device_async(kv_cache);
 
     return kv_cache;
 }
diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp
index b48a10900..bb1136d33 100644
--- a/csrc/engine/infer_engine.cpp
+++ b/csrc/engine/infer_engine.cpp
@@ -57,6 +57,10 @@ InferEngine::InferEngine(
             enable_graph_compiling,
             attention_backend_));
     }
+
+    for (auto &worker : workers_) {
+        worker->wait_for_init();
+    }
     // Graphs must be compiled after weights are loaded and post-processed.
     // Quantized models may replace their linear implementations during
     // process_weights_after_loading(), so compiling here would capture stale
diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp
index 57b5b23b4..5cebd351b 100644
--- a/csrc/engine/rank_worker.cpp
+++ b/csrc/engine/rank_worker.cpp
@@ -30,9 +30,14 @@ RankWorker::RankWorker(
     }
     // start the thread
     thread_ = std::thread(&RankWorker::thread_loop, this);
-    // Wait until the worker thread finishes initialization (model created)
+}
+
+void RankWorker::wait_for_init() {
     std::unique_lock<std::mutex> lk(mutex_);
-    cv_.wait(lk, [&] { return init_done_; });
+    cv_.wait(lk, [&] { return init_done_ || should_exit_; });
+    if (should_exit_) {
+        throw std::runtime_error("RankWorker failed to initialize");
+    }
 }
 
 std::string RankWorker::info() const {
diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp
index 69516a5b8..e262b5d87 100644
--- a/csrc/engine/rank_worker.hpp
+++ b/csrc/engine/rank_worker.hpp
@@ -83,6 +83,8 @@ class RankWorker {
                bool enable_graph_compiling,
                backends::AttentionBackend attention_backend);
 
+    void wait_for_init();
+
     // Submit a parameter load job and wait until the load completes on the worker thread.
     void load_param(const std::string &name,
                     const infinicore::Tensor &param);
diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp
index 1706d2aad..c2cc68762 100644
--- a/csrc/models/infinilm_model.cpp
+++ b/csrc/models/infinilm_model.cpp
@@ -78,6 +78,7 @@ std::vector<infinicore::Tensor> InfinilmModel::default_allocate_kv_cache_tensors
                 *paged_kv_cache_config);
             kv_cache_vec.push_back(kv_cache);
         }
+        infinicore::context::syncStream();
         break;
     }
     default: