From 5d8174f82f956f4353a4f776ade43f86fbbf4f53 Mon Sep 17 00:00:00 2001 From: phantomlei3 Date: Sat, 11 Jul 2026 15:20:59 +0800 Subject: [PATCH] bugfix: set up kv cache compleition guard to prevent kv cache transfer error. --- .../kv_cache_transfer/CMakeLists.txt | 12 ++- .../kv_transfer_completion_test.cpp | 86 +++++++++++++++++++ .../kv_cache_transfer/CMakeLists.txt | 12 +++ .../kv_transfer_completion.cpp | 63 ++++++++++++++ .../kv_transfer_completion.h | 49 +++++++++++ .../mooncake_transfer_engine.cpp | 80 ++++++----------- xllm/core/runtime/CMakeLists.txt | 1 + xllm/core/runtime/llm_worker_impl.cpp | 40 +++------ 8 files changed, 259 insertions(+), 84 deletions(-) create mode 100644 tests/core/framework/kv_cache_transfer/kv_transfer_completion_test.cpp create mode 100644 xllm/core/framework/kv_cache_transfer/kv_transfer_completion.cpp create mode 100644 xllm/core/framework/kv_cache_transfer/kv_transfer_completion.h diff --git a/tests/core/framework/kv_cache_transfer/CMakeLists.txt b/tests/core/framework/kv_cache_transfer/CMakeLists.txt index b51ccaf743..6b16c034dd 100644 --- a/tests/core/framework/kv_cache_transfer/CMakeLists.txt +++ b/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 "$") -endif() \ No newline at end of file +endif() diff --git a/tests/core/framework/kv_cache_transfer/kv_transfer_completion_test.cpp b/tests/core/framework/kv_cache_transfer/kv_transfer_completion_test.cpp new file mode 100644 index 0000000000..1e89eb065a --- /dev/null +++ b/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 + +#include +#include + +namespace xllm { +namespace { + +using namespace std::chrono_literals; + +TEST(KVTransferCompletionTest, WaitsForEveryTransfer) { + folly::Promise first_promise; + folly::Promise second_promise; + KVTransferCompletion completion; + completion.add(first_promise.getSemiFuture()); + completion.add(second_promise.getSemiFuture()); + + std::promise waiter_started; + std::future started = waiter_started.get_future(); + std::future 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 success_promise; + folly::Promise 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 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 promise; + KVTransferCompletion completion; + completion.add(promise.getSemiFuture()); + }, + "pending KV transfers"); +} + +} // namespace +} // namespace xllm diff --git a/xllm/core/framework/kv_cache_transfer/CMakeLists.txt b/xllm/core/framework/kv_cache_transfer/CMakeLists.txt index 73f473319b..834709a232 100644 --- a/xllm/core/framework/kv_cache_transfer/CMakeLists.txt +++ b/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 diff --git a/xllm/core/framework/kv_cache_transfer/kv_transfer_completion.cpp b/xllm/core/framework/kv_cache_transfer/kv_transfer_completion.cpp new file mode 100644 index 0000000000..cbc95adf37 --- /dev/null +++ b/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 + +#include +#include +#include + +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 future) { + futures_.emplace_back(std::move(future)); +} + +bool KVTransferCompletion::wait() { + if (futures_.empty()) { + return true; + } + + std::vector> results = + folly::collectAll(futures_).get(wait_timeout_); + futures_.clear(); + return std::all_of( + results.begin(), results.end(), [](const folly::Try& result) { + return result.hasValue() && result.value(); + }); +} + +} // namespace xllm diff --git a/xllm/core/framework/kv_cache_transfer/kv_transfer_completion.h b/xllm/core/framework/kv_cache_transfer/kv_transfer_completion.h new file mode 100644 index 0000000000..b7e2bd53bf --- /dev/null +++ b/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 + +#include +#include + +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 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> futures_; +}; + +} // namespace xllm diff --git a/xllm/core/framework/kv_cache_transfer/mooncake_transfer_engine.cpp b/xllm/core/framework/kv_cache_transfer/mooncake_transfer_engine.cpp index cc826e367a..7fc51f10be 100644 --- a/xllm/core/framework/kv_cache_transfer/mooncake_transfer_engine.cpp +++ b/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,37 +566,7 @@ 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()) { @@ -582,11 +574,7 @@ bool MooncakeTransferEngine::move_memory_blocks( return false; } -#if defined(USE_DCU) return transfer_success; -#else - return true; -#endif } bool MooncakeTransferEngine::move_memory_by_global_offsets( @@ -656,25 +644,7 @@ 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()) { @@ -682,7 +652,7 @@ bool MooncakeTransferEngine::move_memory_by_global_offsets( return false; } - return true; + return transfer_success; } bool MooncakeTransferEngine::pull_memory_blocks( diff --git a/xllm/core/runtime/CMakeLists.txt b/xllm/core/runtime/CMakeLists.txt index b571aa13fc..1c2e3efcfd 100644 --- a/xllm/core/runtime/CMakeLists.txt +++ b/xllm/core/runtime/CMakeLists.txt @@ -85,6 +85,7 @@ cc_library( :parallel_state :kv_cache :kv_cache_transfer + :kv_transfer_completion :models :sampler :tokenizer diff --git a/xllm/core/runtime/llm_worker_impl.cpp b/xllm/core/runtime/llm_worker_impl.cpp index 96c6b4b2d0..9114f2700c 100644 --- a/xllm/core/runtime/llm_worker_impl.cpp +++ b/xllm/core/runtime/llm_worker_impl.cpp @@ -34,6 +34,7 @@ limitations under the License. #include "core/framework/config/kv_cache_config.h" #include "core/framework/config/load_config.h" #include "framework/kv_cache/kv_cache.h" +#include "framework/kv_cache_transfer/kv_transfer_completion.h" #include "framework/model/model_input_params.h" #include "framework/state_dict/state_dict.h" #if defined(USE_CUDA) || defined(USE_ILU) || defined(USE_MUSA) @@ -218,7 +219,7 @@ std::optional LLMWorkerImpl::step_internal( Timer timer; auto& sampling_params = input.sampling_params; - std::vector> futures; + KVTransferCompletion kv_transfers; if (options_.kv_cache_transfer_mode() == "PUSH" && !input.transfer_kv_infos.empty()) { @@ -239,13 +240,16 @@ std::optional LLMWorkerImpl::step_internal( const_cast(&(input.input_params)) ->parallel.layer_synchronizer = layer_synchronizer; - futures.emplace_back( + kv_transfers.add( kv_cache_transfer_->push_kv_blocks_async(input.transfer_kv_infos, context_.get_parallel_args(), layer_synchronizer, is_spec_draft_)); #endif } + auto wait_kv_push = [&kv_transfers]() { + CHECK(kv_transfers.wait()) << "KV cache push failed"; + }; if (::xllm::EPLBConfig::get_instance().enable_eplb()) { eplb_executor_->eplb_execute(input.input_params.expert.eplb_info); } @@ -254,6 +258,7 @@ std::optional LLMWorkerImpl::step_internal( auto model_output = model_executor_->forward( input.token_ids, input.positions, kv_caches_, input.input_params); if (!model_output.hidden_states.defined()) { + wait_kv_push(); return std::nullopt; } @@ -291,23 +296,12 @@ std::optional LLMWorkerImpl::step_internal( !options_.enable_speculative_decode()) { MULTI_MODEL_STEP_UNLOCK(); if (sync_policy == ForwardSyncPolicy::NO_SYNC) { + wait_kv_push(); return std::nullopt; } int ret = device_.synchronize_default_stream(); - // in p-d disaggregation scene, all micro batches should be in same - // prefill/decode stage, so, to judge transfer_kv_infos.empty, - if (options_.kv_cache_transfer_mode() == "PUSH" && - !input.transfer_kv_infos.empty()) { - auto results = - folly::collectAll(futures).within(std::chrono::seconds(60)).get(); - for (const auto& result : results) { - // TODO: Add error handling - if (!result.value()) { - LOG(ERROR) << "kv_cache_transfer_ failed"; - break; - } - } - } + CHECK_EQ(ret, 0) << "synchronize_default_stream failed"; + wait_kv_push(); if (::xllm::EPLBConfig::get_instance().enable_eplb()) { return output; } @@ -380,6 +374,7 @@ std::optional LLMWorkerImpl::step_internal( !can_skip_npu_graph_decode_sync(input.input_params); #endif if (sync_policy == ForwardSyncPolicy::NO_SYNC) { + wait_kv_push(); output.retained_input = std::make_shared(input); if (enable_schedule_overlap()) { output.ready_event = record_current_stream_event(device_); @@ -391,18 +386,7 @@ std::optional LLMWorkerImpl::step_internal( CHECK_EQ(ret, 0) << "synchronize_default_stream failed"; } - if (options_.kv_cache_transfer_mode() == "PUSH" && - !input.transfer_kv_infos.empty()) { - auto results = - folly::collectAll(futures).within(std::chrono::seconds(60)).get(); - for (const auto& result : results) { - // TODO: Add error handling - if (!result.value()) { - LOG(ERROR) << "kv_cache_transfer_ failed"; - break; - } - } - } + wait_kv_push(); COUNTER_ADD(execution_latency_seconds_model, timer.elapsed_seconds()); if (should_sync_default_stream) {