bugfix: set up kv cache compleition guard to prevent kv cache transfer error.

tests/core/framework/kv_cache_transfer/CMakeLists.txt

@@ -20,6 +20,16 @@ cc_test(
     GTest::gtest_main
 )
 
+cc_test(
+  NAME
+    kv_transfer_completion_test
+  SRCS
+    kv_transfer_completion_test.cpp
+  DEPS
+    :kv_transfer_completion
+    GTest::gtest_main
+)
+
 if(USE_NPU OR USE_MLU)
 cc_test(
   NAME
@@ -35,4 +45,4 @@ cc_test(
 # Resolve static link order between xtensor and xllm_server for this test target.
 target_link_libraries(mooncake_transfer_engine_test PRIVATE
   "$<LINK_GROUP:RESCAN,xtensor,xllm_server>")
-endif()
+endif()

tests/core/framework/kv_cache_transfer/kv_transfer_completion_test.cpp

@@ -0,0 +1,86 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "core/framework/kv_cache_transfer/kv_transfer_completion.h"
+
+#include <gtest/gtest.h>
+
+#include <chrono>
+#include <future>
+
+namespace xllm {
+namespace {
+
+using namespace std::chrono_literals;
+
+TEST(KVTransferCompletionTest, WaitsForEveryTransfer) {
+  folly::Promise<bool> first_promise;
+  folly::Promise<bool> second_promise;
+  KVTransferCompletion completion;
+  completion.add(first_promise.getSemiFuture());
+  completion.add(second_promise.getSemiFuture());
+
+  std::promise<void> waiter_started;
+  std::future<void> started = waiter_started.get_future();
+  std::future<bool> result = std::async(std::launch::async, [&]() {
+    waiter_started.set_value();
+    return completion.wait();
+  });
+
+  started.wait();
+  first_promise.setValue(true);
+  EXPECT_EQ(result.wait_for(50ms), std::future_status::timeout);
+  second_promise.setValue(true);
+  EXPECT_TRUE(result.get());
+}
+
+TEST(KVTransferCompletionTest, ReportsTransferFailure) {
+  folly::Promise<bool> success_promise;
+  folly::Promise<bool> failure_promise;
+  KVTransferCompletion completion;
+  completion.add(success_promise.getSemiFuture());
+  completion.add(failure_promise.getSemiFuture());
+  success_promise.setValue(true);
+  failure_promise.setValue(false);
+
+  EXPECT_FALSE(completion.wait());
+}
+
+TEST(KVTransferCompletionTest, RejectsPendingTransferAfterTimeout) {
+  EXPECT_DEATH(
+      {
+        folly::Promise<bool> promise;
+        KVTransferCompletion completion(1ms);
+        completion.add(promise.getSemiFuture());
+        try {
+          completion.wait();
+        } catch (const folly::FutureTimeout&) {
+        }
+      },
+      "pending KV transfers");
+}
+
+TEST(KVTransferCompletionTest, RejectsPendingTransferAtDestruction) {
+  EXPECT_DEATH(
+      {
+        folly::Promise<bool> promise;
+        KVTransferCompletion completion;
+        completion.add(promise.getSemiFuture());
+      },
+      "pending KV transfers");
+}
+
+}  // namespace
+}  // namespace xllm

xllm/core/framework/kv_cache_transfer/CMakeLists.txt

@@ -21,6 +21,18 @@ cc_library(
     glog::glog
 )
 
+cc_library(
+  NAME
+    kv_transfer_completion
+  HDRS
+    kv_transfer_completion.h
+  SRCS
+    kv_transfer_completion.cpp
+  DEPS
+    glog::glog
+    Folly::folly
+)
+
 cc_library(
   NAME
     kv_cache_transfer

xllm/core/framework/kv_cache_transfer/kv_transfer_completion.cpp

@@ -0,0 +1,63 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "core/framework/kv_cache_transfer/kv_transfer_completion.h"
+
+#include <glog/logging.h>
+
+#include <algorithm>
+#include <chrono>
+#include <utility>
+
+namespace xllm {
+namespace {
+
+constexpr std::chrono::seconds kKVTransferWaitTimeout{60};
+
+}  // namespace
+
+KVTransferCompletion::KVTransferCompletion()
+    : KVTransferCompletion(kKVTransferWaitTimeout) {}
+
+KVTransferCompletion::KVTransferCompletion(
+    std::chrono::milliseconds wait_timeout)
+    : wait_timeout_(wait_timeout) {
+  CHECK_GT(wait_timeout_.count(), 0) << "wait timeout must be positive";
+}
+
+KVTransferCompletion::~KVTransferCompletion() {
+  CHECK(futures_.empty())
+      << "pending KV transfers must finish before source blocks are released";
+}
+
+void KVTransferCompletion::add(folly::SemiFuture<bool> future) {
+  futures_.emplace_back(std::move(future));
+}
+
+bool KVTransferCompletion::wait() {
+  if (futures_.empty()) {
+    return true;
+  }
+
+  std::vector<folly::Try<bool>> results =
+      folly::collectAll(futures_).get(wait_timeout_);
+  futures_.clear();
+  return std::all_of(
+      results.begin(), results.end(), [](const folly::Try<bool>& result) {
+        return result.hasValue() && result.value();
+      });
+}
+
+}  // namespace xllm

xllm/core/framework/kv_cache_transfer/kv_transfer_completion.h

@@ -0,0 +1,49 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <folly/futures/Future.h>
+
+#include <chrono>
+#include <vector>
+
+namespace xllm {
+
+// Owns asynchronous KV transfers until every transfer reaches a terminal
+// state. Source KV blocks must not be released while this object is pending.
+class KVTransferCompletion final {
+ public:
+  KVTransferCompletion();
+  explicit KVTransferCompletion(std::chrono::milliseconds wait_timeout);
+  ~KVTransferCompletion();
+
+  KVTransferCompletion(const KVTransferCompletion&) = delete;
+  KVTransferCompletion& operator=(const KVTransferCompletion&) = delete;
+  KVTransferCompletion(KVTransferCompletion&&) = delete;
+  KVTransferCompletion& operator=(KVTransferCompletion&&) = delete;
+
+  void add(folly::SemiFuture<bool> future);
+
+  // Waits until all owned transfers finish. Returns false when any transfer
+  // reports failure or completes with an exception.
+  bool wait();
+
+ private:
+  std::chrono::milliseconds wait_timeout_;
+  std::vector<folly::SemiFuture<bool>> futures_;
+};
+
+}  // namespace xllm

xllm/core/framework/kv_cache_transfer/mooncake_transfer_engine.cpp

@@ -72,6 +72,28 @@ bool check_buf_range(uint64_t buf_len,
   return true;
 }
 
+bool wait_batch(TransferEngine* engine, BatchID batch_id) {
+  while (true) {
+    TransferStatus status;
+    mooncake::Status result = engine->getBatchTransferStatus(batch_id, status);
+    if (!result.ok()) {
+      LOG(ERROR) << "getBatchTransferStatus not ok";
+      return false;
+    }
+    if (status.s == TransferStatusEnum::COMPLETED) {
+      return true;
+    }
+    if (status.s == TransferStatusEnum::FAILED) {
+      LOG(ERROR) << "getBatchTransferStatus failed";
+      return false;
+    }
+    if (status.s == TransferStatusEnum::TIMEOUT) {
+      LOG(ERROR) << "Sync data transfer timeout";
+      return false;
+    }
+  }
+}
+
 }  // namespace
 
 // ============================================================================
@@ -544,49 +566,15 @@ bool MooncakeTransferEngine::move_memory_blocks(
     return false;
   }
 
-  TransferStatus status;
-  bool completed = false;
-#if defined(USE_DCU)
-  bool transfer_success = true;
-#endif
-  while (!completed) {
-    s = engine->getBatchTransferStatus(batch_id, status);
-    if (!s.ok()) {
-      LOG(ERROR) << "getBatchTransferStatus not ok";
-#if defined(USE_DCU)
-      transfer_success = false;
-#endif
-      completed = true;
-    }
-
-    if (status.s == TransferStatusEnum::COMPLETED) {
-      completed = true;
-    } else if (status.s == TransferStatusEnum::FAILED) {
-      LOG(ERROR) << "getBatchTransferStatus failed";
-#if defined(USE_DCU)
-      transfer_success = false;
-#endif
-      completed = true;
-    } else if (status.s == TransferStatusEnum::TIMEOUT) {
-      LOG(ERROR) << "Sync data transfer timeout";
-#if defined(USE_DCU)
-      transfer_success = false;
-#endif
-      completed = true;
-    }
-  }
+  const bool transfer_success = wait_batch(engine, batch_id);
 
   s = engine->freeBatchID(batch_id);
   if (!s.ok()) {
     LOG(ERROR) << "freeBatchID failed";
     return false;
   }
 
-#if defined(USE_DCU)
   return transfer_success;
-#else
-  return true;
-#endif
 }
 
 bool MooncakeTransferEngine::move_memory_by_global_offsets(
@@ -656,33 +644,15 @@ bool MooncakeTransferEngine::move_memory_by_global_offsets(
     return false;
   }
 
-  TransferStatus status;
-  bool completed = false;
-  while (!completed) {
-    s = engine->getBatchTransferStatus(batch_id, status);
-    if (!s.ok()) {
-      LOG(ERROR) << "getBatchTransferStatus not ok";
-      completed = true;
-    }
-
-    if (status.s == TransferStatusEnum::COMPLETED) {
-      completed = true;
-    } else if (status.s == TransferStatusEnum::FAILED) {
-      LOG(ERROR) << "getBatchTransferStatus failed";
-      completed = true;
-    } else if (status.s == TransferStatusEnum::TIMEOUT) {
-      LOG(ERROR) << "Sync data transfer timeout";
-      completed = true;
-    }
-  }
+  const bool transfer_success = wait_batch(engine, batch_id);
 
   s = engine->freeBatchID(batch_id);
   if (!s.ok()) {
     LOG(ERROR) << "freeBatchID failed";
     return false;
   }
 
-  return true;
+  return transfer_success;
 }
 
 bool MooncakeTransferEngine::pull_memory_blocks(

xllm/core/runtime/CMakeLists.txt

@@ -85,6 +85,7 @@ cc_library(
     :parallel_state
     :kv_cache
     :kv_cache_transfer
+    :kv_transfer_completion
     :models
     :sampler
     :tokenizer