[performance] InferEngine: drastically improve initialization speed

csrc/cache/kv_cache.cpp

@@ -132,13 +132,10 @@ infinicore::Tensor create_layer_kv_cache(
     }
 
     // [1+1, num_blocks, num_rank_k_heads, block_size, k_dim]
-    infinicore::Tensor kv_cache = infinicore::Tensor::empty(
+    infinicore::Tensor kv_cache = infinicore::Tensor::zeros(
         kv_shape,
         dtype,
         rank_info.device);
-    set_zeros(kv_cache);
-
-    infinicore::context::syncStream();
 
     return kv_cache;
 }

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

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");
+    }
 }
 
 RankWorker::~RankWorker() {

csrc/engine/rank_worker.hpp

@@ -87,6 +87,8 @@ class RankWorker {
                bool enable_graph_compiling,
                backends::AttentionBackend attention_backend);
 
+    void wait_for_init();
+
     ~RankWorker();
 
     // Submit a parameter load job and wait until the load completes on the worker thread.

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: