From f94e2d9489d1d0c85e6d04206a169dafdd3eca03 Mon Sep 17 00:00:00 2001 From: kangmeng3 Date: Wed, 19 Aug 2026 10:25:41 +0800 Subject: [PATCH] feat: add composite host cache support for MTP. --- .../block/composite_block_manager_test.cpp | 22 + .../hierarchy_block_manager_pool_test.cpp | 81 ++ .../kv_cache_transfer/CMakeLists.txt | 1 + .../hierarchy_kv_cache_transfer_test.cpp | 13 +- .../kv_transfer_completion_test.cpp | 123 ++ tests/core/runtime/CMakeLists.txt | 6 +- tests/core/runtime/mtp_host_offload_test.cpp | 54 +- .../core/scheduler/scheduler_policy_test.cpp | 125 ++ .../block/composite_block_manager.cpp | 40 +- .../block/hierarchy_block_manager_pool.cpp | 27 +- .../block/sliding_window_block_manager.cpp | 101 +- .../block/sliding_window_block_manager.h | 14 +- .../hierarchy_kv_cache_transfer.cpp | 1081 +++++++++++------ .../hierarchy_kv_cache_transfer.h | 195 ++- .../kv_cache_transfer/kv_cache_store.cpp | 429 +++++-- .../kv_cache_transfer/kv_cache_store.h | 30 +- .../framework/prefix_cache/prefix_cache.cpp | 19 +- .../framework/prefix_cache/prefix_cache.h | 4 + xllm/core/runtime/mtp_worker_impl.cpp | 322 +++-- xllm/core/runtime/mtp_worker_impl.h | 25 +- xllm/core/runtime/worker_impl.cpp | 61 +- xllm/core/runtime/worker_impl.h | 24 +- 22 files changed, 2166 insertions(+), 631 deletions(-) diff --git a/tests/core/framework/block/composite_block_manager_test.cpp b/tests/core/framework/block/composite_block_manager_test.cpp index f7b9c9d561..6bdc19e5cc 100644 --- a/tests/core/framework/block/composite_block_manager_test.cpp +++ b/tests/core/framework/block/composite_block_manager_test.cpp @@ -321,6 +321,28 @@ TEST(CompositeBlockManagerTest, FailedGrowthRollsBackNewBlocks) { manager.deallocate_for_sequence(&seq); } +TEST(CompositeBlockManagerTest, FreshSequenceEvictsSwaPrefixForCapacity) { + BlockManager::Options opts = MakeCompositeOptions(/*base_num_blocks=*/256, + kBaseBlockSize, + /*window_size=*/12, + /*max_seqs_per_batch=*/4); + CompositeBlockManager manager(build_composite_leaves(opts)); + const std::vector first_prompt(kMaxTokensPerBatch, 1); + Sequence first = MakeTestSequence(0, first_prompt); + ASSERT_TRUE(manager.allocate_sequence(&first, first_prompt.size())); + first.kv_state().incr_kv_cache_tokens_num(first_prompt.size()); + manager.deallocate_for_sequence(&first); + first.reset(); + + const std::vector second_prompt(kMaxTokensPerBatch, 2); + Sequence second = MakeTestSequence(1, second_prompt); + EXPECT_TRUE(manager.allocate_sequence(&second, second_prompt.size())); + EXPECT_EQ(SwaBlocks(second).size(), + ExpectedSwaLogicalBlocks(second_prompt.size())); + + manager.deallocate_for_sequence(&second); +} + TEST(CompositeBlockManagerTest, DeallocateToleratesRolledBackEmptySequence) { BlockManager::Options opts = MakeCompositeOptions(/*base_num_blocks=*/128, kBaseBlockSize, diff --git a/tests/core/framework/block/hierarchy_block_manager_pool_test.cpp b/tests/core/framework/block/hierarchy_block_manager_pool_test.cpp index 559dc3f2eb..cf1b24857d 100644 --- a/tests/core/framework/block/hierarchy_block_manager_pool_test.cpp +++ b/tests/core/framework/block/hierarchy_block_manager_pool_test.cpp @@ -314,6 +314,87 @@ TEST(HierarchyBlockManagerPoolTest, TypedLayoutHasPerDpRankLeaves) { } } +TEST(HierarchyBlockManagerPoolTest, + RuntimeSizedTypedLayoutAllocatesColdPromptChunks) { + BlockManagerPool::Options options = make_typed_cache_options(); + options.num_blocks(6400) + .max_tokens_per_batch(25000) + .max_seqs_per_batch(8) + .host_num_blocks_by_type({{BlockType::SWA, 856}, + {BlockType::C4, 6400}, + {BlockType::C128, 200}}); + HierarchyBlockManagerPool pool(options, + /*engine=*/nullptr, + /*dp_size=*/1); + + std::vector tokens(36025, 97); + Sequence sequence = make_test_sequence(/*index=*/0, tokens); + + pool.allocate_shared(&sequence); + const HostCacheRestorePoint selected = pool.select_host_cache_restore( + &sequence, std::numeric_limits::max()); + pool.trim_host_cache(&sequence, selected); + EXPECT_TRUE(pool.allocate(&sequence, /*num_tokens=*/16384)); + EXPECT_GE(sequence.kv_state().current_max_tokens_capacity(), 16384u); + + sequence.kv_state().set_kv_cache_tokens_num(16384); + EXPECT_TRUE(pool.allocate(&sequence, /*num_tokens=*/32768)); + EXPECT_GE(sequence.kv_state().current_max_tokens_capacity(), 32768u); + const Slice swa_blocks = sequence.kv_state().blocks(BlockType::SWA); + ASSERT_EQ(swa_blocks.size(), 256u); + EXPECT_FALSE(swa_blocks[0].is_valid()); + EXPECT_TRUE(swa_blocks[127].is_valid()); + EXPECT_TRUE(swa_blocks[128].is_valid()); + + sequence.kv_state().set_kv_cache_tokens_num(32768); + EXPECT_TRUE(pool.allocate(&sequence, /*num_tokens=*/36025)); + EXPECT_GE(sequence.kv_state().current_max_tokens_capacity(), 36025u); + + pool.deallocate(&sequence); +} + +TEST(HierarchyBlockManagerPoolTest, + RuntimeSizedTypedLayoutRestoresColdPromptTail) { + BlockManagerPool::Options options = make_typed_cache_options(); + options.num_blocks(6400) + .max_tokens_per_batch(25000) + .max_seqs_per_batch(8) + .host_num_blocks_by_type({{BlockType::SWA, 856}, + {BlockType::C4, 6400}, + {BlockType::C128, 200}}); + HierarchyBlockManagerPool pool(options, + /*engine=*/nullptr, + /*dp_size=*/1); + + std::vector tokens(36025, 101); + auto& host_leaves = + HierarchyPoolTestPeer::mutable_host_block_managers(pool).front(); + seed_host_prefix(host_leaves.at(BlockType::SWA).leaf.get(), tokens); + seed_host_prefix(host_leaves.at(BlockType::C4).leaf.get(), tokens); + seed_host_prefix(host_leaves.at(BlockType::C128).leaf.get(), tokens); + + Sequence restored = make_test_sequence(/*index=*/0, tokens); + ASSERT_TRUE(allocate_with_host_cache_budget( + &pool, + &restored, + /*num_tokens=*/tokens.size(), + /*max_copy_units=*/std::numeric_limits::max())); + EXPECT_EQ(restored.kv_state().kv_cache_tokens_num(), 32768u); + + const Slice swa_blocks = restored.kv_state().blocks(BlockType::SWA); + ASSERT_EQ(swa_blocks.size(), 282u); + EXPECT_FALSE(swa_blocks[254].is_valid()); + EXPECT_TRUE(swa_blocks[255].is_valid()); + EXPECT_TRUE(swa_blocks.back().is_valid()); + const size_t valid_swa_blocks = static_cast(std::count_if( + swa_blocks.begin(), swa_blocks.end(), [](const Block& block) { + return block.is_valid(); + })); + EXPECT_EQ(valid_swa_blocks, 27u); + + pool.deallocate(&restored); +} + TEST(HierarchyBlockManagerPoolTest, DecodeTypedLayoutKeepsSwaHostLeafForOffloadOnly) { BlockManagerPool::Options options = make_typed_cache_options(); diff --git a/tests/core/framework/kv_cache_transfer/CMakeLists.txt b/tests/core/framework/kv_cache_transfer/CMakeLists.txt index 3950c9dd2c..1439df98ad 100644 --- a/tests/core/framework/kv_cache_transfer/CMakeLists.txt +++ b/tests/core/framework/kv_cache_transfer/CMakeLists.txt @@ -27,6 +27,7 @@ cc_test( kv_transfer_completion_test.cpp DEPS :kv_transfer_completion + :kv_cache_transfer GTest::gtest_main ) diff --git a/tests/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer_test.cpp b/tests/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer_test.cpp index a93704b0a2..6cba05f3f8 100644 --- a/tests/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer_test.cpp +++ b/tests/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer_test.cpp @@ -46,15 +46,22 @@ class HierarchyKVCacheTransferTestPeer final { static void set_layer_batch_ranges( HierarchyKVCacheTransfer* transfer, std::vector ranges) { - transfer->layer_batch_ranges_ = std::move(ranges); + transfer->participant_states_.at(CacheParticipant::TARGET) + .layer_batch_ranges = std::move(ranges); } static bool load_from_host( HierarchyKVCacheTransfer* transfer, std::shared_ptr synchronizer, const std::vector& block_transfer_info) { - return transfer->load_from_host(std::move(synchronizer), - block_transfer_info); + auto transaction = + std::make_shared(); + transaction->synchronizers[CacheParticipant::TARGET] = + std::move(synchronizer); + transaction->required_participant_mask = + HierarchyKVCacheTransfer::participant_mask(CacheParticipant::TARGET); + return transfer->load_from_host( + CacheParticipant::TARGET, transaction, block_transfer_info); } }; 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 index 4fcd015db7..adc76afdae 100644 --- a/tests/core/framework/kv_cache_transfer/kv_transfer_completion_test.cpp +++ b/tests/core/framework/kv_cache_transfer/kv_transfer_completion_test.cpp @@ -17,15 +17,138 @@ limitations under the License. #include +#include #include #include #include +#include +#include +#include + +#include "core/framework/kv_cache_transfer/kv_cache_store.h" namespace xllm { + +class KVCacheStoreTestPeer final { + public: + static void set_config(KVCacheStore* store, + const KVCacheStoreInitConfig& config) { + store->config_ = config; + } + + static void set_components(KVCacheStore* store, + std::vector components) { + store->components_ = std::move(components); + } + + static std::string build_component_key( + const KVCacheStore& store, + const HostCacheComponentSchema& component, + const BlockTransferInfo& block_info) { + return store.build_component_key(component, block_info); + } + + static size_t required_component_count(const KVCacheStore& store, + BlockType block_type) { + return store.required_components(block_type).size(); + } +}; + namespace { using namespace std::chrono_literals; +BlockTransferInfo make_block_info(BlockType block_type) { + std::array hash_key{}; + for (size_t index = 0; index < hash_key.size(); ++index) { + hash_key[index] = static_cast(index + 1); + } + return BlockTransferInfo(/*src_id=*/3, + /*dst_id=*/7, + hash_key.data(), + TransferType::D2H2G, + block_type); +} + +HostCacheComponentSchema make_component(CacheParticipant participant, + const std::string& model_identity, + BlockType block_type = BlockType::KV) { + HostCacheComponentSchema component; + component.participant = participant; + component.block_type = block_type; + component.model_identity = model_identity; + component.schema_fingerprint = "schema-fingerprint"; + component.tp_rank = 0; + component.tp_size = 8; + return component; +} + +TEST(KVCacheStoreKeyTest, SeparatesTargetAndDraftComponents) { + KVCacheStore store; + KVCacheStoreInitConfig config; + config.model_id = "deepseek-v4"; + KVCacheStoreTestPeer::set_config(&store, config); + const BlockTransferInfo block_info = make_block_info(BlockType::C128); + HostCacheComponentSchema target = make_component( + CacheParticipant::TARGET, "deepseek-v4-target", BlockType::C128); + HostCacheComponentSchema draft = target; + draft.participant = CacheParticipant::DRAFT; + + const std::string target_key = + KVCacheStoreTestPeer::build_component_key(store, target, block_info); + const std::string draft_key = + KVCacheStoreTestPeer::build_component_key(store, draft, block_info); + + EXPECT_EQ(target_key.rfind("xllm-kv-v3:", 0), 0u); + EXPECT_EQ(draft_key.rfind("xllm-kv-v3:", 0), 0u); + EXPECT_NE(target_key, draft_key); +} + +TEST(KVCacheStoreKeyTest, SeparatesParticipantModelAndSchemaIdentity) { + KVCacheStore store; + KVCacheStoreInitConfig config; + config.model_id = "deepseek-v4"; + KVCacheStoreTestPeer::set_config(&store, config); + const BlockTransferInfo block_info = make_block_info(BlockType::SWA); + HostCacheComponentSchema first = make_component( + CacheParticipant::DRAFT, "draft-revision-a", BlockType::SWA); + HostCacheComponentSchema second = first; + second.model_identity = "draft-revision-b"; + HostCacheComponentSchema third = first; + third.schema_fingerprint = "different-schema"; + + const std::string first_key = + KVCacheStoreTestPeer::build_component_key(store, first, block_info); + const std::string second_key = + KVCacheStoreTestPeer::build_component_key(store, second, block_info); + const std::string third_key = + KVCacheStoreTestPeer::build_component_key(store, third, block_info); + + EXPECT_NE(first_key, second_key); + EXPECT_NE(first_key, third_key); +} + +TEST(KVCacheStoreKeyTest, ExpandsLogicalTypeToEveryRequiredParticipant) { + KVCacheStore store; + std::vector components; + components.emplace_back( + make_component(CacheParticipant::TARGET, "target", BlockType::C4)); + components.emplace_back( + make_component(CacheParticipant::DRAFT, "draft", BlockType::C4)); + components.emplace_back( + make_component(CacheParticipant::TARGET, "target", BlockType::SWA)); + KVCacheStoreTestPeer::set_components(&store, std::move(components)); + + EXPECT_EQ( + KVCacheStoreTestPeer::required_component_count(store, BlockType::C4), 2u); + EXPECT_EQ( + KVCacheStoreTestPeer::required_component_count(store, BlockType::SWA), + 1u); + EXPECT_EQ( + KVCacheStoreTestPeer::required_component_count(store, BlockType::C128), + 0u); +} + TEST(KVTransferCompletionTest, WaitsForEveryTransfer) { folly::Promise first_promise; folly::Promise second_promise; diff --git a/tests/core/runtime/CMakeLists.txt b/tests/core/runtime/CMakeLists.txt index b25434b8ca..012c8e1612 100644 --- a/tests/core/runtime/CMakeLists.txt +++ b/tests/core/runtime/CMakeLists.txt @@ -116,7 +116,9 @@ if(USE_MLU) GTest::gtest_main torch_mlu ) +endif() +if(USE_NPU OR USE_MLU) cc_test( NAME mtp_host_offload_test @@ -129,8 +131,10 @@ if(USE_MLU) :kv_cache :platform GTest::gtest_main - torch_mlu ) + if(USE_MLU) + target_link_libraries(mtp_host_offload_test PRIVATE torch_mlu) + endif() endif() if(USE_CUDA) diff --git a/tests/core/runtime/mtp_host_offload_test.cpp b/tests/core/runtime/mtp_host_offload_test.cpp index 6d6a911813..25fb6c0975 100644 --- a/tests/core/runtime/mtp_host_offload_test.cpp +++ b/tests/core/runtime/mtp_host_offload_test.cpp @@ -36,9 +36,11 @@ class RecordingTransferWorker final : public LLMWorkerImpl { RecordingTransferWorker(const ParallelArgs& parallel_args, const torch::Device& device, const runtime::Options& options, - uint32_t transfer_result) + uint32_t transfer_result, + uint8_t prefetch_hit = 1) : LLMWorkerImpl(parallel_args, device, options), - transfer_result_(transfer_result) {} + transfer_result_(transfer_result), + prefetch_hit_(prefetch_hit) {} uint32_t transfer_kv_blocks( uint64_t batch_id, @@ -49,6 +51,12 @@ class RecordingTransferWorker final : public LLMWorkerImpl { return transfer_result_; } + std::vector prefetch_kv_blocks( + Slice& block_transfer_info) override { + ++prefetch_count_; + return std::vector(block_transfer_info.size(), prefetch_hit_); + } + uint32_t transfer_kv_blocks( uint64_t batch_id, Slice& block_transfer_info) override { @@ -60,6 +68,7 @@ class RecordingTransferWorker final : public LLMWorkerImpl { uint32_t vector_transfer_count() const { return vector_transfer_count_; } uint32_t slice_transfer_count() const { return slice_transfer_count_; } + uint32_t prefetch_count() const { return prefetch_count_; } uint64_t last_batch_id() const { return last_batch_id_; } size_t last_transfer_size() const { return last_transfer_size_; } @@ -67,6 +76,8 @@ class RecordingTransferWorker final : public LLMWorkerImpl { uint32_t transfer_result_ = 0; uint32_t vector_transfer_count_ = 0; uint32_t slice_transfer_count_ = 0; + uint32_t prefetch_count_ = 0; + uint8_t prefetch_hit_ = 1; uint64_t last_batch_id_ = 0; size_t last_transfer_size_ = 0; }; @@ -89,14 +100,14 @@ class MTPHostOffloadTest : public ::testing::Test { protected: void SetUp() override { if (Platform::device_count() < 1) { - GTEST_SKIP() << "MLU device is required for MTP host offload tests."; + GTEST_SKIP() << "An accelerator is required for MTP host offload tests."; } } }; TEST_F(MTPHostOffloadTest, TransfersEveryBlockToTargetAndDraft) { constexpr uint64_t kBatchId = 42; - const torch::Device device("mlu:0"); + const torch::Device device(Platform::type_torch(), /*index=*/0); ParallelArgs parallel_args( /*rank=*/0, /*world_size=*/1, /*process_group=*/nullptr); runtime::Options options; @@ -134,7 +145,7 @@ TEST_F(MTPHostOffloadTest, TransfersEveryBlockToTargetAndDraft) { TEST_F(MTPHostOffloadTest, RejectsMismatchedTargetAndDraftTransferCounts) { constexpr uint64_t kBatchId = 73; - const torch::Device device("mlu:0"); + const torch::Device device(Platform::type_torch(), /*index=*/0); ParallelArgs parallel_args( /*rank=*/0, /*world_size=*/1, /*process_group=*/nullptr); runtime::Options options; @@ -160,5 +171,38 @@ TEST_F(MTPHostOffloadTest, RejectsMismatchedTargetAndDraftTransferCounts) { EXPECT_EQ(draft_ptr->slice_transfer_count(), 1); } +TEST_F(MTPHostOffloadTest, PrefetchRequiresEveryParticipantHit) { + const torch::Device device(Platform::type_torch(), /*index=*/0); + ParallelArgs parallel_args( + /*rank=*/0, /*world_size=*/1, /*process_group=*/nullptr); + runtime::Options options; + options.block_size(16).num_speculative_tokens(1); + TestMTPWorker worker(parallel_args, device, options); + + const std::vector transfer_info = { + BlockTransferInfo(/*src_block_id=*/5, /*dst_block_id=*/6)}; + auto target = std::make_unique(parallel_args, + device, + options, + /*transfer_result=*/1, + /*prefetch_hit=*/1); + auto draft = std::make_unique(parallel_args, + device, + options, + /*transfer_result=*/1, + /*prefetch_hit=*/0); + RecordingTransferWorker* target_ptr = target.get(); + RecordingTransferWorker* draft_ptr = draft.get(); + worker.replace_transfer_workers(std::move(target), std::move(draft)); + Slice transfer_slice(transfer_info); + + const std::vector hits = worker.prefetch_kv_blocks(transfer_slice); + + ASSERT_EQ(hits.size(), 1u); + EXPECT_EQ(hits.front(), 0); + EXPECT_EQ(target_ptr->prefetch_count(), 1u); + EXPECT_EQ(draft_ptr->prefetch_count(), 1u); +} + } // namespace } // namespace xllm diff --git a/tests/core/scheduler/scheduler_policy_test.cpp b/tests/core/scheduler/scheduler_policy_test.cpp index b0662efd25..4c85d64bf8 100644 --- a/tests/core/scheduler/scheduler_policy_test.cpp +++ b/tests/core/scheduler/scheduler_policy_test.cpp @@ -30,6 +30,7 @@ limitations under the License. #include "core/framework/config/scheduler_config.h" #include "distributed_runtime/engine.h" #include "framework/block/block_manager_pool.h" +#include "framework/block/hierarchy_block_manager_pool.h" #include "framework/model/model_args.h" #include "util/utils.h" @@ -521,6 +522,130 @@ TEST(SchedulerPolicyTest, UnifiedRetryRefreshesHostRestoreBeforeChunkSizing) { EXPECT_TRUE(finished.empty()); } +TEST(SchedulerPolicyTest, PrefillFirstHostColdRequestAllocatesAllDsv4Chunks) { + ContinuousScheduler::Options options = create_scheduler_options( + /*max_tokens_per_batch=*/25000, + /*max_seqs_per_batch=*/8, + /*num_speculative_tokens=*/1, + /*max_tokens_per_chunk_for_prefill=*/16384, + /*dp_size=*/1, + /*priority_strategy=*/"fcfs"); + options.enable_schedule_overlap() = true; + BatchMode mode{ + .enable_mix_batch = false, + .enable_chunked_prefill = true, + .priority_strategy = "fcfs", + }; + PrefillFirstPolicy policy(mode, options); + + BlockManagerPool::Options block_options; + block_options.num_blocks(6400) + .block_size(128) + .enable_prefix_cache(true) + .enable_host_offload(true) + .sliding_window_size(128) + .swa_blocks_per_seq(1) + .max_tokens_per_batch(25000) + .max_seqs_per_batch(8) + .num_speculative_tokens(1) + .num_embedding_blocks(214) + .manager_types({1u, 0u, 0u}) + .compress_ratios({0u, 4u, 128u}) + .host_num_blocks_by_type({{BlockType::SWA, 856}, + {BlockType::C4, 6400}, + {BlockType::C128, 200}}); + HierarchyBlockManagerPool block_manager_pool(block_options, + /*engine=*/nullptr, + /*dp_size=*/1); + auto profile_engine = std::make_unique(6400, 128); + ProfileManager::Options profile_options; + profile_options.max_tokens_per_batch(25000).max_seqs_per_batch(8); + ProfileManager profile_manager(profile_engine.get(), profile_options); + + std::vector> requests = generate_request( + {36025}, {32}, std::nullopt, std::nullopt, /*max_context_len=*/40000); + DequeQueue prefill_queue; + prefill_queue.push(requests.front()); + DequeQueue chunk_queue; + DequeQueue decode_queue; + std::list> unified_queue; + std::vector> running_requests; + std::vector running_sequences; + std::vector running_sequence_budgets; + bool last_step_prefill = false; + SchedulerState state{ + .prefill_queue = prefill_queue, + .chunk_queue = chunk_queue, + .decode_queue = decode_queue, + .unified_queue = unified_queue, + .running_requests = running_requests, + .running_sequences = running_sequences, + .running_sequences_budgets = running_sequence_budgets, + .kv_cache_manager = &block_manager_pool, + .profile_manager = &profile_manager, + .response_processor = nullptr, + .last_step_prefill = last_step_prefill, + .options = options, + .min_speculative_tokens_required = 2, + .enable_prefix_cache = true, + .has_linear_attention_layers = false, + }; + ScheduleBudget budget{ + .remaining_token_budget = 25000, + .remaining_seq_budget = 8, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + std::vector> finished; + + policy.schedule(state, budget, finished); + + ASSERT_EQ(running_sequences.size(), 1u); + EXPECT_EQ(running_sequence_budgets, (std::vector{16384})); + EXPECT_GE(running_sequences.front()->kv_state().current_max_tokens_capacity(), + 16384u); + EXPECT_TRUE(prefill_queue.empty()); + EXPECT_TRUE(unified_queue.empty()); + EXPECT_TRUE(finished.empty()); + + running_sequences.front()->kv_state().set_kv_cache_tokens_num(16384); + ScheduleBudget second_budget{ + .remaining_token_budget = 25000, + .remaining_seq_budget = 8, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + policy.schedule(state, second_budget, finished); + + ASSERT_EQ(running_sequences.size(), 1u); + EXPECT_EQ(running_sequence_budgets, (std::vector{16384})); + EXPECT_GE(running_sequences.front()->kv_state().current_max_tokens_capacity(), + 32768u); + EXPECT_TRUE(chunk_queue.empty()); + EXPECT_TRUE(finished.empty()); + + running_sequences.front()->kv_state().set_kv_cache_tokens_num(32768); + ScheduleBudget final_budget{ + .remaining_token_budget = 25000, + .remaining_seq_budget = 8, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + policy.schedule(state, final_budget, finished); + + ASSERT_EQ(running_sequences.size(), 1u); + EXPECT_EQ(running_sequence_budgets, (std::vector{3257})); + EXPECT_GE(running_sequences.front()->kv_state().current_max_tokens_capacity(), + 36025u); + EXPECT_TRUE(chunk_queue.empty()); + EXPECT_TRUE(finished.empty()); + + block_manager_pool.deallocate(running_sequences.front()); +} + TEST(SchedulerPolicyTest, DefersWhileAsyncBlockReleaseIsPending) { ContinuousScheduler::Options options = create_scheduler_options( /*max_tokens_per_batch=*/16384, diff --git a/xllm/core/framework/block/composite_block_manager.cpp b/xllm/core/framework/block/composite_block_manager.cpp index 1223aa2ce6..7fd273529a 100644 --- a/xllm/core/framework/block/composite_block_manager.cpp +++ b/xllm/core/framework/block/composite_block_manager.cpp @@ -368,6 +368,12 @@ bool CompositeBlockManager::allocate_sequence(Sequence* seq, } KVCacheState& kv_state = seq->kv_state(); + // Publish completed blocks before any growth. SWA retirement is deferred + // until every independent leaf has allocated successfully: releasing SWA + // mutates the sequence in place and therefore must be the final operation + // that can precede an allocation failure. + cache_full_blocks_for_sequence(seq); + // Fan out growth. Each leaf returns its newly allocated blocks (or nullopt // on failure). Stage keyed by BlockType; commit only after every leaf // succeeds so a failure rolls back cleanly. @@ -380,16 +386,41 @@ bool CompositeBlockManager::allocate_sequence(Sequence* seq, staged.clear(); }; - for (auto& [type, entry] : leaves_) { + auto allocate_leaf = [&](BlockType type, LeafEntry& entry) { std::optional> blocks = entry.leaf->allocate_for_sequence(seq, num_tokens); if (!blocks.has_value()) { + VLOG(1) << "Composite block allocation failed: type=" + << static_cast(type) << ", num_tokens=" << num_tokens + << ", held=" << kv_state.num_blocks(type) + << ", free=" << entry.leaf->num_free_blocks() + << ", total=" << entry.leaf->num_total_blocks(); release_staged(); return false; } if (!blocks->empty()) { staged.emplace(type, std::move(*blocks)); } + return true; + }; + + for (auto& [type, entry] : leaves_) { + if (type == BlockType::SWA) { + continue; + } + if (!allocate_leaf(type, entry)) { + return false; + } + } + + const auto swa_it = leaves_.find(BlockType::SWA); + if (swa_it != leaves_.end()) { + // SlidingWindowBlockManager performs its capacity preflight and prefix + // retirement atomically inside allocate_for_sequence(). Keep it last so + // no independent leaf failure can follow that in-place mutation. + if (!allocate_leaf(swa_it->first, swa_it->second)) { + return false; + } } // Grow-or-fail: every cache-bearing leaf must cover num_tokens now. @@ -412,6 +443,13 @@ bool CompositeBlockManager::allocate_sequence(Sequence* seq, staged_it == staged.end() ? 0 : staged_it->second.size(); const size_t total = seq->kv_state().num_blocks(type) + staged_for_type; if (total < needed) { + VLOG(1) << "Composite block coverage failed: type=" + << static_cast(type) << ", num_tokens=" << num_tokens + << ", needed=" << needed + << ", held=" << seq->kv_state().num_blocks(type) + << ", staged=" << staged_for_type + << ", free=" << entry.leaf->num_free_blocks() + << ", total=" << entry.leaf->num_total_blocks(); release_staged(); return false; } diff --git a/xllm/core/framework/block/hierarchy_block_manager_pool.cpp b/xllm/core/framework/block/hierarchy_block_manager_pool.cpp index ffd16140e4..f094d9fbd6 100644 --- a/xllm/core/framework/block/hierarchy_block_manager_pool.cpp +++ b/xllm/core/framework/block/hierarchy_block_manager_pool.cpp @@ -468,6 +468,11 @@ bool HierarchyBlockManagerPool::allocate(Sequence* sequence, auto* composite = static_cast(block_managers_[dp_rank].get()); if (!composite->allocate_sequence(sequence, num_tokens)) { + VLOG(1) << "[HostCache][AdmissionFailed] sequence_id=" << sequence->seq_id() + << " num_tokens=" << num_tokens + << " device_tokens=" << hbm_state.kv_cache_tokens_num() + << " host_tokens=" + << sequence->host_kv_state().kv_cache_tokens_num(); release_host_match(sequence, dp_rank); return false; } @@ -486,12 +491,12 @@ bool HierarchyBlockManagerPool::allocate(Sequence* sequence, staged.clear(); }; - for (auto& [type, entry] : host_block_managers_[dp_rank]) { + auto allocate_host_leaf = [&](BlockType type, auto& entry) { std::optional> blocks = entry.leaf->allocate_for_sequence(sequence, host_state, num_tokens); if (!blocks.has_value()) { release_staged(); - break; + return false; } if (!blocks->empty()) { staged.emplace(type, std::move(*blocks)); @@ -504,9 +509,27 @@ bool HierarchyBlockManagerPool::allocate(Sequence* sequence, const size_t total = host_state.num_blocks(type) + staged_for_type; if (total < needed) { release_staged(); + return false; + } + return true; + }; + + bool host_growth_succeeded = true; + for (auto& [type, entry] : host_block_managers_[dp_rank]) { + if (type == BlockType::SWA) { + continue; + } + if (!allocate_host_leaf(type, entry)) { + host_growth_succeeded = false; break; } } + const auto host_swa_it = host_block_managers_[dp_rank].find(BlockType::SWA); + if (host_growth_succeeded && + host_swa_it != host_block_managers_[dp_rank].end()) { + host_growth_succeeded = + allocate_host_leaf(host_swa_it->first, host_swa_it->second); + } for (auto& [type, blocks] : staged) { host_state.add_blocks(type, blocks); diff --git a/xllm/core/framework/block/sliding_window_block_manager.cpp b/xllm/core/framework/block/sliding_window_block_manager.cpp index 7f001fc071..f4ec32abc7 100644 --- a/xllm/core/framework/block/sliding_window_block_manager.cpp +++ b/xllm/core/framework/block/sliding_window_block_manager.cpp @@ -34,6 +34,96 @@ SlidingWindowBlockManager::SlidingWindowBlockManager(const Options& options) } } +std::optional> +SlidingWindowBlockManager::allocate_for_sequence(Sequence* seq, + size_t num_tokens) { + if (seq == nullptr) { + return std::nullopt; + } + return allocate_for_sequence(seq, seq->kv_state(), num_tokens); +} + +std::optional> +SlidingWindowBlockManager::allocate_for_sequence(Sequence* seq, + KVCacheState& kv_state, + size_t num_tokens) { + if (seq == nullptr) { + return std::nullopt; + } + const size_t block_size = options_.block_size(); + if (block_size == 0) { + return std::vector{}; + } + + const size_t held = kv_state.num_blocks(block_type()); + const size_t needed = (num_tokens + block_size - 1) / block_size; + if (needed <= held) { + return std::vector{}; + } + + const bool device_state = &kv_state == &seq->kv_state(); + const size_t completed_tokens = device_state ? seq->kv_cache_tokens_num() + : kv_state.kv_cache_tokens_num(); + const size_t skipped = + std::min(num_slid_out_blocks(completed_tokens), needed); + const size_t first_physical_position = std::max(held, skipped); + const size_t physical_count = needed - first_physical_position; + + // Releasing the completed prefix mutates the sequence in place. Prove that + // its uniquely owned blocks plus the current free list can satisfy the new + // tail before making that mutation. The SWA pool is sized from the global + // burst budget, so it must not depend on evicting unrelated cached blocks. + const Slice current_blocks = kv_state.blocks(block_type()); + const size_t release_positions = std::min(skipped, held); + size_t releasable_blocks = 0; + for (size_t position = 0; position < release_positions; ++position) { + const Block& block = current_blocks[position]; + if (block.is_valid() && block.ref_count() <= 2u) { + ++releasable_blocks; + } + } + const size_t available_without_eviction = + num_free_blocks() + releasable_blocks; + if (physical_count > available_without_eviction) { + const size_t eviction_deficit = physical_count - available_without_eviction; + if (prefix_cache_ == nullptr || + !prefix_cache_->can_evict(eviction_deficit)) { + return std::nullopt; + } + } + + release_out_of_window(seq, kv_state, completed_tokens); + std::vector physical_blocks = allocate(physical_count); + CHECK_EQ(physical_blocks.size(), physical_count) + << "SWA capacity preflight succeeded but allocation failed"; + + std::vector blocks(needed - held); + for (size_t position = first_physical_position; position < needed; + ++position) { + blocks[position - held] = + std::move(physical_blocks[position - first_physical_position]); + } + return blocks; +} + +size_t SlidingWindowBlockManager::num_slid_out_blocks( + size_t cached_tokens) const { + const size_t block_size = options_.block_size(); + if (block_size == 0) { + return 0; + } + const size_t num_spec_tokens = + static_cast(options_.num_speculative_tokens()); + const size_t sliding_window_tokens = + std::max(options_.sliding_window_size(), 1); + if (cached_tokens < sliding_window_tokens + num_spec_tokens) { + return 0; + } + const size_t skipped_tokens = + cached_tokens - sliding_window_tokens - num_spec_tokens + 1; + return skipped_tokens / block_size; +} + void SlidingWindowBlockManager::release_out_of_window(Sequence* seq) { if (seq == nullptr) { return; @@ -57,16 +147,7 @@ void SlidingWindowBlockManager::release_out_of_window(Sequence* seq, if (block_size == 0 || swa_blocks.empty()) { return; } - const size_t num_spec_tokens = - static_cast(options_.num_speculative_tokens()); - const size_t sliding_window_tokens = - std::max(options_.sliding_window_size(), 1); - if (cached_tokens < (sliding_window_tokens + num_spec_tokens)) { - return; - } - const size_t skipped_tokens = - cached_tokens - sliding_window_tokens - num_spec_tokens + 1; - const size_t skipped_blocks = skipped_tokens / block_size; + const size_t skipped_blocks = num_slid_out_blocks(cached_tokens); const size_t release_blocks = std::min(skipped_blocks, swa_blocks.size()); if (release_blocks == 0) { return; diff --git a/xllm/core/framework/block/sliding_window_block_manager.h b/xllm/core/framework/block/sliding_window_block_manager.h index 21f11e3334..1c9056e0ff 100644 --- a/xllm/core/framework/block/sliding_window_block_manager.h +++ b/xllm/core/framework/block/sliding_window_block_manager.h @@ -20,7 +20,10 @@ limitations under the License. namespace xllm { // Sliding-window leaf of CompositeBlockManager. Reuses BlockManagerImpl's -// physical pool and flat-append growth. SWA-specific behavior: +// physical pool and position-preserving growth. SWA-specific behavior: +// - allocate_for_sequence() appends invalid placeholders for positions that +// have already slid out, and allocates physical blocks only for the active +// window plus the current chunk. // - release_out_of_window() drops leading slid-out blocks; released // positions stay as invalid placeholders so DSA modulo indexing // (`(pos/block_size) % semantic_cols`) remains stable. @@ -31,6 +34,14 @@ class SlidingWindowBlockManager : public BlockManagerImpl { explicit SlidingWindowBlockManager(const Options& options); ~SlidingWindowBlockManager() override = default; + std::optional> allocate_for_sequence( + Sequence* seq, + size_t num_tokens) override; + std::optional> allocate_for_sequence( + Sequence* seq, + KVCacheState& kv_state, + size_t num_tokens) override; + // Deallocate leading blocks that have slid out of the window; leaves // invalid placeholders in their slots. Called by the composite after a // successful allocate commit. @@ -49,6 +60,7 @@ class SlidingWindowBlockManager : public BlockManagerImpl { uint32_t swa_blocks_per_seq() const { return options_.swa_blocks_per_seq(); } private: + size_t num_slid_out_blocks(size_t cached_tokens) const; void release_out_of_window(Sequence* seq, KVCacheState& kv_state, size_t cached_tokens); diff --git a/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.cpp b/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.cpp index 13aa8fdc2e..c3f00207bd 100644 --- a/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.cpp +++ b/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.cpp @@ -16,18 +16,40 @@ limitations under the License. #include "framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h" #include +#include #include -#include #include +#include #include +#include #include #include "framework/kv_cache_transfer/kv_cache_store.h" +#include "util/hash_util.h" namespace xllm { namespace { -constexpr uint32_t TIMEOUT_MS = 60000; +constexpr uint32_t kTimeoutMs = 60000; +constexpr size_t kOffloadStreamCount = 4; + +const std::vector kBlockTypes = {BlockType::KV, + BlockType::LINEAR, + BlockType::SWA, + BlockType::C4, + BlockType::C128}; + +std::string participant_name(CacheParticipant participant) { + switch (participant) { + case CacheParticipant::TARGET: + return "TARGET"; + case CacheParticipant::DRAFT: + return "DRAFT"; + } + LOG(FATAL) << "Unsupported cache participant: " + << static_cast(participant); + return "UNKNOWN"; +} std::string make_store_local_hostname(const std::string& configured, uint32_t worker_id) { @@ -57,13 +79,10 @@ std::string make_store_local_hostname(const std::string& configured, } if (port_begin != std::string::npos) { const std::string port_text = configured.substr(port_begin); - bool numeric = true; - for (const char character : port_text) { - if (character < '0' || character > '9') { - numeric = false; - break; - } - } + const bool numeric = + std::all_of(port_text.begin(), port_text.end(), [](char character) { + return character >= '0' && character <= '9'; + }); if (numeric) { port = static_cast(std::stoul(port_text)); host = configured.substr(0, host_end); @@ -75,12 +94,6 @@ std::string make_store_local_hostname(const std::string& configured, return host + ":" + std::to_string(port + worker_id); } -// Streams reserved for concurrent D2H offload callers. D2H runs synchronously -// on the RemoteWorker copy threadpool (4 threads, see remote_worker.h); reserve -// one stream per such thread so concurrent offloads never block on the stream -// queue. -constexpr size_t kOffloadStreamCount = 4; - using CopyStreamQueue = moodycamel::BlockingConcurrentQueue>; @@ -102,7 +115,7 @@ class CopyStreamLease final { void drain_or_die(const char* reason) const { try { - const int synchronize_result = stream_->synchronize(); + const int32_t synchronize_result = stream_->synchronize(); if (synchronize_result != 0) { LOG(FATAL) << "Failed to drain KV Cache copy stream: reason=" << reason << ", result=" << synchronize_result; @@ -128,13 +141,11 @@ std::vector build_layer_batch_ranges( if (num_layers <= 0) { return ranges; } - uint32_t layers_per_batch = requested_batches == 0 ? static_cast(num_layers) : static_cast(num_layers) / requested_batches; layers_per_batch = std::max(layers_per_batch, 1); - for (int64_t begin = 0; begin < num_layers; begin += layers_per_batch) { ranges.push_back( {begin, std::min(begin + layers_per_batch, num_layers)}); @@ -142,164 +153,82 @@ std::vector build_layer_batch_ranges( return ranges; } -bool has_tensor(const torch::Tensor& tensor) { - return tensor.defined() && tensor.numel() > 0; +std::vector device_block_shape(const torch::Tensor& tensor) { + CHECK(tensor.defined() && tensor.dim() > 0); + std::vector shape; + shape.reserve(static_cast(tensor.dim() - 1)); + for (int64_t dim = 1; dim < tensor.dim(); ++dim) { + shape.emplace_back(tensor.size(dim)); + } + return shape; } -BlockTypeTensorMap build_block_type_tensor_map(const KVCache& kv_cache, - BlockType type) { - BlockTypeTensorMap map; +std::vector host_block_layer_shape(const torch::Tensor& tensor) { + CHECK(tensor.defined() && tensor.dim() > 1); + std::vector shape; + shape.reserve(static_cast(tensor.dim() - 2)); + for (int64_t dim = 2; dim < tensor.dim(); ++dim) { + shape.emplace_back(tensor.size(dim)); + } + return shape; +} - const torch::Tensor key_cache = kv_cache.get_k_cache(); - const torch::Tensor value_cache = kv_cache.get_v_cache(); - const torch::Tensor index_cache = kv_cache.get_index_cache(); - const torch::Tensor conv_cache = kv_cache.get_conv_cache(); - const torch::Tensor ssm_cache = kv_cache.get_ssm_cache(); - const torch::Tensor swa_cache = kv_cache.get_swa_cache(); - const std::optional index_cache_scale = - kv_cache.get_indexer_cache_scale(); +std::string hash_schema_string(const std::string& schema) { + const XXH3Key hash = hash_string(schema); + return std::string(reinterpret_cast(hash.data), + sizeof(hash.data)); +} - switch (type) { - case BlockType::KV: - if (has_tensor(conv_cache) || has_tensor(ssm_cache) || - has_tensor(swa_cache)) { - return {}; - } - if (has_tensor(key_cache)) { - map.emplace(KVCacheTensorRole::KEY, key_cache); - } - if (has_tensor(value_cache)) { - map.emplace(KVCacheTensorRole::VALUE, value_cache); - } - if (has_tensor(index_cache)) { - map.emplace(KVCacheTensorRole::INDEX, index_cache); - } - // INT8 indexer cache carries a per-token fp32 scale that must travel with - // the int8 index values during offload/reload. - if (index_cache_scale.has_value() && - has_tensor(index_cache_scale.value())) { - map.emplace(KVCacheTensorRole::INDEX_SCALE, index_cache_scale.value()); - } - return map; - case BlockType::LINEAR: - if (has_tensor(conv_cache)) { - map.emplace(KVCacheTensorRole::CONV, conv_cache); - } - if (has_tensor(ssm_cache)) { - map.emplace(KVCacheTensorRole::SSM, ssm_cache); - } - return map; - case BlockType::SWA: - // Every DSV4 layer reads the SWA cache, including ratio-4/128 layers. - // Host restore is restricted to C128-aligned boundaries by the hierarchy - // pool; compressor/index state is partial-block scratch at those - // boundaries and is regenerated by the resumed forward. The persistent - // SWA window itself must be restored for every layer. - if (!has_tensor(swa_cache)) { - return {}; - } - map.emplace(KVCacheTensorRole::SWA, swa_cache); - return map; - case BlockType::C4: - // DSV4 compress-ratio-4 layer: has swa + key + index (no value). - if (!has_tensor(swa_cache) || has_tensor(value_cache) || - !has_tensor(key_cache) || !has_tensor(index_cache)) { - return {}; - } - map.emplace(KVCacheTensorRole::KEY, key_cache); - map.emplace(KVCacheTensorRole::INDEX, index_cache); - // INT8 indexer cache carries a per-token fp16 scale that must travel - // with the int8 index values during offload/reload. - if (index_cache_scale.has_value() && - has_tensor(index_cache_scale.value())) { - map.emplace(KVCacheTensorRole::INDEX_SCALE, index_cache_scale.value()); - } - return map; - case BlockType::C128: - // DSV4 compress-ratio-128 layer: has swa + key, but no index/value. - if (!has_tensor(swa_cache) || has_tensor(value_cache) || - !has_tensor(key_cache) || has_tensor(index_cache)) { - return {}; - } - map.emplace(KVCacheTensorRole::KEY, key_cache); - return map; - default: - return {}; +} // namespace + +void HierarchyKVCacheTransfer::LoadTransaction::abort() { + std::vector> synchronizers_to_abort; + { + std::lock_guard lock(mutex); + if (aborted) { + return; + } + aborted = true; + for (const auto& [participant, synchronizer] : synchronizers) { + (void)participant; + synchronizers_to_abort.emplace_back(synchronizer); + } + } + for (const std::shared_ptr& synchronizer : + synchronizers_to_abort) { + if (synchronizer != nullptr) { + synchronizer->abort(); + } } } -} // namespace +HierarchyKVCacheTransfer::HierarchyKVCacheTransfer(const Options& options, + const torch::Device& device) + : options_(options), device_(device) {} HierarchyKVCacheTransfer::HierarchyKVCacheTransfer( const Options& options, const torch::Device& device, const Stream* compute_stream, - std::vector* kv_caches_ptr, + std::vector* kv_caches_ptr, const KVCacheShape& kv_cache_shape, const KVCacheCreateOptions& create_options) - : options_(options), - device_(device), - compute_stream_(compute_stream), - kv_caches_ptr_(kv_caches_ptr), - kv_cache_shape_(kv_cache_shape), - create_options_(create_options) { - CHECK(kv_caches_ptr_ != nullptr) << "kv_caches_ptr must not be null."; - CHECK(compute_stream_ != nullptr) << "compute stream must not be null."; - - device_.set_device(); - device_.init_device_context(); - load_threadpool_ = std::make_unique( - /*num_threads=*/2, - /*init_func=*/[this]() mutable { device_.set_device(); }, - /*cpu_binding=*/false, - /*pool_name=*/"HierarchyKVCacheTransfer.load"); - // D2H offload runs synchronously on the caller (RemoteWorker copy thread) so - // its copied-block count can be returned to the scheduler; it is not posted - // to a local pool. Size the shared stream pool to cover the H2D load threads - // plus the concurrent D2H callers. - const size_t num_streams = load_threadpool_->size() + kOffloadStreamCount; - for (size_t i = 0; i < num_streams; ++i) { - copy_stream_.enqueue(device_.get_stream_from_pool(TIMEOUT_MS)); - } - - build_device_block_type_map(); - layer_batch_ranges_ = build_layer_batch_ranges( - options_.layers(), options_.layers_wise_copy_batchs()); - - if (options_.host_blocks_factor() > 1.0) { - batch_memcpy_ = create_batch_memcpy(device_); - create_host_cache(); - } - - if (options_.enable_kvcache_store()) { - CHECK(options_.host_blocks_factor() > 1.0) - << "Mooncake Store requires Host cache capacity."; - KVCacheStoreInitConfig store_config; - const std::string store_local_hostname = make_store_local_hostname( - options_.store_local_hostname(), options_.store_worker_id()); - store_config.localhost_name = store_local_hostname; - store_config.protocol = options_.store_protocol(); - store_config.metadata_server = options_.store_metadata_server(); - store_config.master_server_address = options_.store_master_server_address(); - store_config.model_id = options_.store_namespace(); - store_config.tp_rank = options_.tp_rank(); - store_config.tp_size = options_.tp_size(); - LOG(INFO) << "[Mooncake][StoreEngine] initialize, endpoint=" - << store_local_hostname << ", protocol=" << store_config.protocol - << ", tp_rank=" << store_config.tp_rank - << ", tp_size=" << store_config.tp_size; - kv_cache_store_ = std::make_unique(); - CHECK(kv_cache_store_->init(store_config, &host_kv_caches_)) - << "Failed to initialize Mooncake Store."; - LOG(INFO) << "[Mooncake][StoreEngine] ready, endpoint=" - << store_local_hostname << ", protocol=" << store_config.protocol - << ", tp_rank=" << store_config.tp_rank; - } + : HierarchyKVCacheTransfer(options, device) { + ParticipantRegistration registration; + registration.participant = CacheParticipant::TARGET; + registration.actual_compute_stream = compute_stream; + registration.device_caches = kv_caches_ptr; + registration.cache_shape = kv_cache_shape; + registration.create_options = create_options; + registration.model_identity = + create_options.model_id() + "|" + create_options.model_type(); + registration.tp_rank = options.tp_rank(); + registration.tp_size = options.tp_size(); + register_participant(std::move(registration)); + CHECK(finalize_registration()); } HierarchyKVCacheTransfer::~HierarchyKVCacheTransfer() { - // Joining the load pool first guarantees no producer can enqueue more work - // while the copy streams and host cache storage are being released. load_threadpool_.reset(); device_.set_device(); @@ -313,104 +242,319 @@ HierarchyKVCacheTransfer::~HierarchyKVCacheTransfer() { } std::lock_guard lock(mutex_); - layer_wise_load_synchronizer_.clear(); + for (auto& [batch_id, transaction] : load_transactions_) { + (void)batch_id; + transaction->abort(); + } + load_transactions_.clear(); } -void HierarchyKVCacheTransfer::build_device_block_type_map() { - device_kv_caches_.clear(); - device_block_type_layer_ids_.clear(); +void HierarchyKVCacheTransfer::register_participant( + ParticipantRegistration registration) { + CHECK(registration_state_ != RegistrationState::READY) + << "Cannot register a participant after finalization."; + CHECK(registration.actual_compute_stream != nullptr) + << "actual compute stream must not be null."; + CHECK(registration.device_caches != nullptr) + << "device caches must not be null."; + CHECK(!registration.device_caches->empty()) + << "device caches must not be empty."; + CHECK(!registration.model_identity.empty()) + << "participant model identity must not be empty."; + CHECK_GT(registration.tp_size, 0u); + CHECK_LT(registration.tp_rank, registration.tp_size); + CHECK(participant_states_.find(registration.participant) == + participant_states_.end()) + << "Duplicate cache participant: " + << participant_name(registration.participant); + + ParticipantState state; + state.registration = std::move(registration); + participant_states_.emplace(state.registration.participant, std::move(state)); + registration_state_ = RegistrationState::REGISTERING; +} - const std::vector kBlockTypes = {BlockType::KV, - BlockType::LINEAR, - BlockType::SWA, - BlockType::C4, - BlockType::C128}; +bool HierarchyKVCacheTransfer::finalize_registration() { + CHECK(registration_state_ == RegistrationState::REGISTERING) + << "Hierarchy KV cache registration is not active."; + CHECK(!participant_states_.empty()); + device_.set_device(); + device_.init_device_context(); + for (auto& [participant, state] : participant_states_) { + (void)participant; + build_participant_state(state); + } + validate_composite_schema(); + initialize_resources(); + if (options_.enable_kvcache_store()) { + initialize_store(); + } + registration_state_ = RegistrationState::READY; + return true; +} + +void HierarchyKVCacheTransfer::initialize_resources() { + CHECK_GT(options_.host_blocks_factor(), 1.0) + << "Hierarchy KV cache transfer requires Host cache capacity."; + const size_t load_threads = std::max(2, participant_states_.size()); + load_threadpool_ = std::make_unique( + load_threads, + /*init_func=*/[this]() mutable { device_.set_device(); }, + /*cpu_binding=*/false, + /*pool_name=*/"HierarchyKVCacheTransfer.load"); + const size_t num_streams = load_threadpool_->size() + kOffloadStreamCount; + for (size_t i = 0; i < num_streams; ++i) { + copy_stream_.enqueue(device_.get_stream_from_pool(kTimeoutMs)); + } + batch_memcpy_ = create_batch_memcpy(device_); + CHECK(batch_memcpy_ != nullptr); +} + +void HierarchyKVCacheTransfer::build_participant_state( + ParticipantState& state) { + CHECK_EQ(state.registration.device_caches->size(), + static_cast(state.registration.create_options.num_layers())) + << "Participant cache layer count does not match allocation metadata: " + << participant_name(state.registration.participant); + build_device_block_type_map(state); + CHECK(!state.device_grouped_caches.empty()) + << "Participant has no Host-cache-compatible tensors: " + << participant_name(state.registration.participant); + state.layer_batch_ranges = build_layer_batch_ranges( + static_cast(state.registration.device_caches->size()), + options_.layers_wise_copy_batchs()); + create_host_cache(state); + build_and_validate_schema(state); +} + +void HierarchyKVCacheTransfer::build_device_block_type_map( + ParticipantState& state) { + state.device_grouped_caches.clear(); + state.absolute_layer_ids.clear(); for (int64_t layer_id = 0; - layer_id < static_cast(kv_caches_ptr_->size()); + layer_id < + static_cast(state.registration.device_caches->size()); ++layer_id) { - KVCache& kv_cache = kv_caches_ptr_->at(static_cast(layer_id)); - for (BlockType type : kBlockTypes) { - BlockTypeTensorMap tensor_map = - build_block_type_tensor_map(kv_cache, type); - if (!tensor_map.empty()) { - device_kv_caches_[type].push_back(&kv_cache); - device_block_type_layer_ids_[type].push_back(layer_id); + KVCache& kv_cache = + state.registration.device_caches->at(static_cast(layer_id)); + for (BlockType block_type : kBlockTypes) { + const BlockTypeTensorMap tensors = + kv_cache.get_block_type_tensors(block_type); + if (tensors.empty()) { + continue; } + state.device_grouped_caches[block_type].push_back(&kv_cache); + state.absolute_layer_ids[block_type].push_back(layer_id); } } } -void HierarchyKVCacheTransfer::create_host_cache() { - CHECK(!device_kv_caches_.empty()) - << "device block type caches must not be empty."; - - for (const auto& [block_type, group_caches] : device_kv_caches_) { - if (group_caches.empty()) { - continue; - } - - const int64_t layer_count = static_cast(group_caches.size()); - - KVCacheCreateOptions host_opts = create_options_; - host_opts.device(torch::Device(torch::kCPU)) +void HierarchyKVCacheTransfer::create_host_cache(ParticipantState& state) { + for (const auto& [block_type, group_caches] : state.device_grouped_caches) { + CHECK(!group_caches.empty()); + KVCacheCreateOptions host_options = state.registration.create_options; + host_options.device(torch::Device(torch::kCPU)) .enable_xtensor(false) .tensor_allocator(nullptr) .host_blocks_factor(options_.host_blocks_factor()); #if defined(USE_NPU) - host_opts.enable_kv_cache_huge_page_allocator(false); + host_options.enable_kv_cache_huge_page_allocator(false); #endif + state.host_grouped_caches[block_type] = + std::make_unique(state.registration.cache_shape, + host_options, + block_type, + static_cast(group_caches.size())); + } +} - host_kv_caches_[block_type] = std::make_unique( - kv_cache_shape_, host_opts, block_type, layer_count); +void HierarchyKVCacheTransfer::build_and_validate_schema( + ParticipantState& state) { + state.schema.tensor_specs.clear(); + state.schema.component_fingerprints.clear(); + for (const auto& [block_type, group_caches] : state.device_grouped_caches) { + const auto layer_ids_it = state.absolute_layer_ids.find(block_type); + const auto host_it = state.host_grouped_caches.find(block_type); + CHECK(layer_ids_it != state.absolute_layer_ids.end()); + CHECK(host_it != state.host_grouped_caches.end() && + host_it->second != nullptr); + const std::vector& layer_ids = layer_ids_it->second; + CHECK_EQ(layer_ids.size(), group_caches.size()); + const BlockTypeTensorMap host_tensors = + host_it->second->get_block_type_tensors(block_type); + CHECK(!host_tensors.empty()); + + std::string component_schema = + participant_name(state.registration.participant) + "|" + + state.registration.model_identity + + "|tp=" + std::to_string(state.registration.tp_size) + ":" + + std::to_string(state.registration.tp_rank) + + "|type=" + std::to_string(static_cast(block_type)); + for (size_t layer_slot = 0; layer_slot < group_caches.size(); + ++layer_slot) { + const BlockTypeTensorMap device_tensors = + group_caches[layer_slot]->get_block_type_tensors(block_type); + CHECK(!device_tensors.empty()); + for (const auto& [role, device_tensor] : device_tensors) { + const auto host_tensor_it = host_tensors.find(role); + CHECK(host_tensor_it != host_tensors.end()) + << "Missing required Host tensor: participant=" + << participant_name(state.registration.participant) + << ", type=" << static_cast(block_type) + << ", layer=" << layer_ids[layer_slot] + << ", role=" << static_cast(role); + const torch::Tensor& host_tensor = host_tensor_it->second; + CHECK_EQ(device_tensor.scalar_type(), host_tensor.scalar_type()); + CHECK_EQ(host_tensor.size(1), + static_cast(group_caches.size())); + const std::vector block_shape = + device_block_shape(device_tensor); + CHECK(block_shape == host_block_layer_shape(host_tensor)); + + HostCacheTensorSpec spec; + spec.participant = state.registration.participant; + spec.block_type = block_type; + spec.absolute_layer_id = layer_ids[layer_slot]; + spec.host_layer_slot = static_cast(layer_slot); + spec.role = role; + spec.coverage = HostPayloadCoverage::REQUIRED; + spec.dtype = device_tensor.scalar_type(); + spec.block_shape = block_shape; + state.schema.tensor_specs.emplace_back(std::move(spec)); + + component_schema.append("|layer="); + component_schema.append(std::to_string(layer_ids[layer_slot])); + component_schema.append(",slot="); + component_schema.append(std::to_string(layer_slot)); + component_schema.append(",role="); + component_schema.append(std::to_string(static_cast(role))); + component_schema.append(",dtype="); + component_schema.append( + std::to_string(static_cast(device_tensor.scalar_type()))); + component_schema.append(",shape="); + for (int64_t dimension : block_shape) { + component_schema.append(std::to_string(dimension)); + component_schema.push_back('x'); + } + } + } + state.schema.component_fingerprints[block_type] = + hash_schema_string(component_schema); + } +} + +void HierarchyKVCacheTransfer::validate_composite_schema() const { + std::map device_block_counts; + std::map host_block_counts; + int64_t block_size = -1; + for (const auto& [participant, state] : participant_states_) { + (void)participant; + const int64_t participant_block_size = + state.registration.create_options.block_size(); + if (block_size < 0) { + block_size = participant_block_size; + } else { + CHECK_EQ(block_size, participant_block_size) + << "Composite participants use different logical block sizes."; + } + for (const auto& [block_type, group_caches] : state.device_grouped_caches) { + CHECK(!group_caches.empty()); + const BlockTypeTensorMap device_tensors = + group_caches.front()->get_block_type_tensors(block_type); + CHECK(!device_tensors.empty()); + const int64_t device_blocks = device_tensors.begin()->second.size(0); + auto device_count_it = device_block_counts.find(block_type); + if (device_count_it == device_block_counts.end()) { + device_block_counts[block_type] = device_blocks; + } else { + CHECK_EQ(device_count_it->second, device_blocks) + << "Composite participants use different device block counts for " + << static_cast(block_type); + } + + const auto host_it = state.host_grouped_caches.find(block_type); + CHECK(host_it != state.host_grouped_caches.end() && + host_it->second != nullptr); + const BlockTypeTensorMap host_tensors = + host_it->second->get_block_type_tensors(block_type); + CHECK(!host_tensors.empty()); + const int64_t host_blocks = host_tensors.begin()->second.size(0); + auto host_count_it = host_block_counts.find(block_type); + if (host_count_it == host_block_counts.end()) { + host_block_counts[block_type] = host_blocks; + } else { + CHECK_EQ(host_count_it->second, host_blocks) + << "Composite participants use different Host block counts for " + << static_cast(block_type); + } + CHECK_GE(host_blocks, device_blocks); + } } } +void HierarchyKVCacheTransfer::initialize_store() { + CHECK_GT(options_.host_blocks_factor(), 1.0) + << "Mooncake Store requires Host cache capacity."; + KVCacheStoreInitConfig store_config; + const std::string store_local_hostname = make_store_local_hostname( + options_.store_local_hostname(), options_.store_worker_id()); + store_config.localhost_name = store_local_hostname; + store_config.protocol = options_.store_protocol(); + store_config.metadata_server = options_.store_metadata_server(); + store_config.master_server_address = options_.store_master_server_address(); + store_config.model_id = options_.store_namespace(); + store_config.tp_rank = options_.tp_rank(); + store_config.tp_size = options_.tp_size(); + LOG(INFO) << "[Mooncake][StoreEngine] initialize, endpoint=" + << store_local_hostname << ", protocol=" << store_config.protocol + << ", tp_rank=" << store_config.tp_rank + << ", tp_size=" << store_config.tp_size; + kv_cache_store_ = std::make_unique(); + CHECK(kv_cache_store_->init(store_config, this)) + << "Failed to initialize Mooncake Store."; + LOG(INFO) << "[Mooncake][StoreEngine] ready, endpoint=" + << store_local_hostname << ", protocol=" << store_config.protocol + << ", tp_rank=" << store_config.tp_rank; +} + HierarchyKVCacheTransfer::CopyPlan HierarchyKVCacheTransfer::build_copy_plan( + const ParticipantState& state, const std::vector& block_transfer_info, const LayerBatchRange& layer_batch_range) const { CopyPlan plan; if (block_transfer_info.empty()) { return plan; } - const TransferType transfer_type = block_transfer_info.front().transfer_type; - - for (const auto& info : block_transfer_info) { - BlockType type = info.block_type; - auto device_it = device_kv_caches_.find(type); - auto layer_ids_it = device_block_type_layer_ids_.find(type); - auto host_it = host_kv_caches_.find(type); - if (device_it == device_kv_caches_.end() || - layer_ids_it == device_block_type_layer_ids_.end() || - host_it == host_kv_caches_.end()) { + for (const BlockTransferInfo& info : block_transfer_info) { + const auto device_it = state.device_grouped_caches.find(info.block_type); + const auto layer_ids_it = state.absolute_layer_ids.find(info.block_type); + const auto host_it = state.host_grouped_caches.find(info.block_type); + if (device_it == state.device_grouped_caches.end() || + layer_ids_it == state.absolute_layer_ids.end() || + host_it == state.host_grouped_caches.end()) { continue; } - - const auto& group_caches = device_it->second; - const auto& layer_ids = layer_ids_it->second; - const KVCache* host_cache = host_it->second.get(); - CHECK(host_cache != nullptr) << "host cache instance must not be null."; + const std::vector& group_caches = device_it->second; + const std::vector& layer_ids = layer_ids_it->second; const BlockTypeTensorMap host_tensors = - host_cache->get_block_type_tensors(type); + host_it->second->get_block_type_tensors(info.block_type); int32_t host_block_id = -1; int32_t device_block_id = -1; - switch (transfer_type) { - case TransferType::H2D: - host_block_id = info.src_block_id; - device_block_id = info.dst_block_id; - break; - case TransferType::D2H2G: - host_block_id = info.dst_block_id; - device_block_id = info.src_block_id; - break; - default: - LOG(FATAL) << "Unsupported transfer type for copy plan: " - << static_cast(transfer_type); + if (transfer_type == TransferType::H2D) { + host_block_id = info.src_block_id; + device_block_id = info.dst_block_id; + } else if (transfer_type == TransferType::D2H2G) { + host_block_id = info.dst_block_id; + device_block_id = info.src_block_id; + } else { + LOG(FATAL) << "Unsupported transfer type for copy plan: " + << static_cast(transfer_type); } - - CHECK_GE(host_block_id, 0) << "host block id must be non-negative."; + CHECK_GE(host_block_id, 0); + CHECK_GE(device_block_id, 0); for (size_t layer_slot = 0; layer_slot < group_caches.size(); ++layer_slot) { @@ -419,24 +563,18 @@ HierarchyKVCacheTransfer::CopyPlan HierarchyKVCacheTransfer::build_copy_plan( absolute_layer_id >= layer_batch_range.end_layer) { continue; } - - BlockTypeTensorMap device_tensors = - build_block_type_tensor_map(*group_caches[layer_slot], type); + const BlockTypeTensorMap device_tensors = + group_caches[layer_slot]->get_block_type_tensors(info.block_type); for (const auto& [role, device_tensor] : device_tensors) { - auto host_tensor_it = host_tensors.find(role); - if (host_tensor_it == host_tensors.end()) { - continue; - } - - // device_tensor shape: [num_blocks, ...per_block_dims] - // host_tensor shape: [num_host_blocks, num_layers, ...per_block_dims] + const auto host_tensor_it = host_tensors.find(role); + CHECK(host_tensor_it != host_tensors.end()) + << "Required Host tensor disappeared after schema finalization."; const torch::Tensor& host_tensor = host_tensor_it->second; - CHECK_LT(host_block_id, host_tensor.size(0)) - << "host block id out of range."; + CHECK_LT(host_block_id, host_tensor.size(0)); + CHECK_LT(device_block_id, device_tensor.size(0)); torch::Tensor device_block = device_tensor[device_block_id]; torch::Tensor host_block_layer = host_tensor[host_block_id][static_cast(layer_slot)]; - if (transfer_type == TransferType::H2D) { plan.src_tensors.emplace_back(host_block_layer); plan.dst_tensors.emplace_back(device_block); @@ -447,64 +585,70 @@ HierarchyKVCacheTransfer::CopyPlan HierarchyKVCacheTransfer::build_copy_plan( } } } - return plan; } uint32_t HierarchyKVCacheTransfer::transfer_kv_blocks( uint64_t batch_id, const std::vector& block_transfer_info) { + CHECK(registration_state_ == RegistrationState::READY); CHECK(!block_transfer_info.empty()); - - // This runs synchronously on the caller's thread (a brpc RPC worker thread - // for remote workers), which has no ACL context of its own. Both branches - // below touch device resources on this thread — D2H offload issues the copy - // inline, and H2D creates the layer synchronizer's events - // (aclrtCreateEventWithFlag). Establish the context here so those calls do - // not fail with ACL_ERROR_RT_CONTEXT_NULL (107002). Idempotent; the async H2D - // copy posted to load_threadpool_ already has context via that pool's - // init_func. device_.set_device(); - switch (block_transfer_info[0].transfer_type) { + switch (block_transfer_info.front().transfer_type) { case TransferType::D2H2G: return offload(block_transfer_info); case TransferType::H2D: { - const uint32_t scheduled_blocks = - static_cast(block_transfer_info.size()); - // Register before scheduling the load. RemoteWorker serializes this RPC - // ahead of the matching forward RPC, so the engine does not wait for H2D - // and the forward still observes this worker-local synchronizer. - auto synchronizer = create_layer_synchronizer( - static_cast(layer_batch_ranges_.size())); - CHECK(synchronizer != nullptr) - << "Failed to create layer synchronizer for H2D batch_id=" << batch_id - << "; H2D copy cannot be backed and the pool has already advanced " - "kv_cache_tokens_num_."; + auto transaction = std::make_shared(); + if (!snapshot_ready_entries(block_transfer_info, transaction.get())) { + LOG(ERROR) << "Composite Host cache entry is not READY for batch_id=" + << batch_id; + return 0; + } + for (const auto& [participant, state] : participant_states_) { + const bool required = std::any_of( + block_transfer_info.begin(), + block_transfer_info.end(), + [&state](const BlockTransferInfo& info) { + return state.device_grouped_caches.find(info.block_type) != + state.device_grouped_caches.end(); + }); + if (!required) { + continue; + } + std::shared_ptr synchronizer = + create_layer_synchronizer( + static_cast(state.layer_batch_ranges.size())); + CHECK(synchronizer != nullptr) + << "Failed to create participant layer synchronizer."; + transaction->synchronizers[participant] = synchronizer; + transaction->required_participant_mask |= participant_mask(participant); + } + CHECK_NE(transaction->required_participant_mask, 0u); { std::lock_guard lock(mutex_); - auto existing = layer_wise_load_synchronizer_.find(batch_id); - if (existing != layer_wise_load_synchronizer_.end()) { - LOG(ERROR) - << "layer_wise_load_synchronizer collision at batch_id=" - << batch_id - << ", previous entry was never consumed (batch cancelled or " - "batch_id reused). Overwriting; stale entry's events will " - "release when its refcount drops."; + auto existing = load_transactions_.find(batch_id); + if (existing != load_transactions_.end()) { + existing->second->abort(); + LOG(ERROR) << "Composite load transaction collision at batch_id=" + << batch_id << "; replacing stale transaction."; } - layer_wise_load_synchronizer_[batch_id] = synchronizer; + load_transactions_[batch_id] = transaction; + } + for (const auto& [participant, synchronizer] : + transaction->synchronizers) { + (void)synchronizer; + load_threadpool_->schedule( + [this, participant, transaction, block_transfer_info]() { + load_from_host(participant, transaction, block_transfer_info); + }); } - load_threadpool_->schedule( - [this, - synchronizer, - block_transfer_info = std::move(block_transfer_info)]() mutable { - load_from_host(synchronizer, block_transfer_info); - }); - return scheduled_blocks; + return static_cast(block_transfer_info.size()); } default: LOG(ERROR) << "Unsupported transfer type: " - << static_cast(block_transfer_info[0].transfer_type); + << static_cast( + block_transfer_info.front().transfer_type); return 0; } } @@ -513,13 +657,14 @@ uint32_t HierarchyKVCacheTransfer::transfer_kv_blocks( uint64_t /*batch_id*/, Slice& block_transfer_info) { CHECK(!block_transfer_info.empty()); - CHECK(kv_cache_store_ != nullptr); - if (block_transfer_info[0].transfer_type == TransferType::G2H) { - return kv_cache_store_->batch_get(block_transfer_info); + if (block_transfer_info[0].transfer_type != TransferType::G2H) { + LOG(ERROR) << "Unsupported slice transfer type: " + << static_cast(block_transfer_info[0].transfer_type); + return 0; } - LOG(ERROR) << "Unsupported slice transfer type: " - << static_cast(block_transfer_info[0].transfer_type); - return 0; + const std::vector statuses = prefetch_kv_blocks(block_transfer_info); + return static_cast( + std::count(statuses.begin(), statuses.end(), static_cast(1))); } std::vector HierarchyKVCacheTransfer::prefetch_kv_blocks( @@ -533,6 +678,7 @@ std::vector HierarchyKVCacheTransfer::prefetch_kv_blocks( } std::vector hits = kv_cache_store_->batch_get_with_status(block_transfer_info); + update_prefetched_entries(block_transfer_info, hits); const size_t hit_count = std::count(hits.begin(), hits.end(), static_cast(1)); VLOG(1) << "[Mooncake][PrefetchGet] type=" @@ -546,117 +692,376 @@ uint32_t HierarchyKVCacheTransfer::offload( if (block_transfer_info.empty()) { return 0; } - - if (batch_memcpy_ == nullptr) { - return block_transfer_info.size(); + if (!begin_entry_write(block_transfer_info)) { + return 0; } - Slice slice(block_transfer_info); if (!offload_to_host(slice)) { - LOG(ERROR) << "Offload to host failed."; + abort_entry_write(block_transfer_info); + LOG(ERROR) << "Composite offload to Host failed."; return 0; } + commit_entry_write(block_transfer_info); if (options_.enable_kvcache_store()) { CHECK(kv_cache_store_ != nullptr); const uint32_t put_count = kv_cache_store_->batch_put(block_transfer_info); if (put_count != block_transfer_info.size()) { - LOG(WARNING) << "Mooncake BatchPut partially failed: " << put_count << "/" - << block_transfer_info.size(); + LOG(WARNING) << "Mooncake composite BatchPut partially failed: " + << put_count << "/" << block_transfer_info.size(); } VLOG(1) << "[Mooncake][OffloadPut] blocks=" << block_transfer_info.size() << ", success=" << put_count; } - return block_transfer_info.size(); + return static_cast(block_transfer_info.size()); } bool HierarchyKVCacheTransfer::offload_to_host( Slice& block_transfer_info) { - if (block_transfer_info.empty()) { - return true; + CHECK(batch_memcpy_ != nullptr); + const std::vector transfer_info = + static_cast>(block_transfer_info); + CopyStreamLease stream(©_stream_); + std::set waited_streams; + for (const auto& [participant, state] : participant_states_) { + (void)participant; + if (waited_streams.emplace(state.registration.actual_compute_stream) + .second) { + stream.get()->wait_stream(*state.registration.actual_compute_stream); + } } - CHECK(batch_memcpy_ != nullptr) << "batch memcpy must be initialized."; - CopyStreamLease stream(©_stream_); - // D2H is issued from an RPC worker, while the preceding forward is enqueued - // on WorkerImpl's compute stream. Establish a device-side dependency before - // reading KV; the RPC thread's current/default stream is unrelated when - // schedule overlap is enabled. - stream.get()->wait_stream(*compute_stream_); - bool success = true; - for (const auto& range : layer_batch_ranges_) { - CopyPlan plan = build_copy_plan( - static_cast>(block_transfer_info), - range); - if (plan.src_tensors.empty()) { - continue; + for (const auto& [participant, state] : participant_states_) { + bool participant_copied = true; + for (const LayerBatchRange& range : state.layer_batch_ranges) { + CopyPlan plan = build_copy_plan(state, transfer_info, range); + if (plan.src_tensors.empty()) { + continue; + } + if (!batch_memcpy_->copy_d2h( + plan.src_tensors, plan.dst_tensors, stream.get())) { + participant_copied = false; + break; + } } - if (!batch_memcpy_->copy_d2h( - plan.src_tensors, plan.dst_tensors, stream.get())) { - success = false; - break; + if (!participant_copied) { + return false; } + complete_participant_write(participant, transfer_info); } - return success; + return true; } bool HierarchyKVCacheTransfer::load_from_host( - std::shared_ptr synchronizer, + CacheParticipant participant, + const std::shared_ptr& transaction, const std::vector& block_transfer_info) { - if (block_transfer_info.empty()) { - return true; - } - - CHECK(synchronizer != nullptr) << "layer synchronizer must not be null."; - CHECK(batch_memcpy_ != nullptr) << "batch memcpy must be initialized."; + const auto state_it = participant_states_.find(participant); + CHECK(state_it != participant_states_.end()); + const ParticipantState& state = state_it->second; + const auto synchronizer_it = transaction->synchronizers.find(participant); + CHECK(synchronizer_it != transaction->synchronizers.end()); + const std::shared_ptr& synchronizer = + synchronizer_it->second; CopyStreamLease stream(©_stream_); bool success = true; bool stream_has_async_h2d = false; - for (size_t range_idx = 0; range_idx < layer_batch_ranges_.size(); - ++range_idx) { - CopyPlan plan = - build_copy_plan(block_transfer_info, layer_batch_ranges_[range_idx]); + for (size_t range_index = 0; range_index < state.layer_batch_ranges.size(); + ++range_index) { + if (!entry_snapshot_matches(block_transfer_info, *transaction)) { + success = false; + break; + } + CopyPlan plan = build_copy_plan( + state, block_transfer_info, state.layer_batch_ranges[range_index]); if (!plan.src_tensors.empty()) { if (!batch_memcpy_->submit_h2d( plan.src_tensors, plan.dst_tensors, stream.get())) { - stream_has_async_h2d = false; success = false; break; } stream_has_async_h2d = true; } - if (!synchronizer->record_stream(static_cast(range_idx), + if (!synchronizer->record_stream(static_cast(range_index), stream.get())) { if (stream_has_async_h2d) { - stream.drain_or_die("layer-ready event recording failed"); - stream_has_async_h2d = false; + stream.drain_or_die("participant layer-ready event recording failed"); } success = false; break; } } - - // On failure some ranges were never recorded; abort the synchronizer so a - // forward thread spinning on those layers unblocks and reports failure - // (aborting the forward) instead of hanging or reading uncopied KV cache. + if (success && !entry_snapshot_matches(block_transfer_info, *transaction)) { + success = false; + } if (!success) { - synchronizer->abort(); + transaction->abort(); } return success; } void HierarchyKVCacheTransfer::set_layer_synchronizer( ModelInputParams& params) { + CHECK_EQ(participant_states_.size(), 1u) + << "Single-participant synchronizer API used by a composite transfer."; + set_layer_synchronizer(participant_states_.begin()->first, params); +} + +void HierarchyKVCacheTransfer::set_layer_synchronizer( + CacheParticipant participant, + ModelInputParams& params) { + std::lock_guard lock(mutex_); + const auto transaction_it = load_transactions_.find(params.meta.batch_id); + if (transaction_it == load_transactions_.end()) { + return; + } + const std::shared_ptr& transaction = transaction_it->second; + const auto synchronizer_it = transaction->synchronizers.find(participant); + if (synchronizer_it == transaction->synchronizers.end()) { + return; + } + params.parallel.layer_wise_load_synchronizer = synchronizer_it->second; + const ParticipantState& state = participant_states_.at(participant); + params.parallel.layers_per_bacth_copy = + state.layer_batch_ranges.empty() + ? static_cast(state.registration.device_caches->size()) + : static_cast(state.layer_batch_ranges.front().end_layer - + state.layer_batch_ranges.front().begin_layer); + transaction->consumed_participant_mask |= participant_mask(participant); + if (transaction->consumed_participant_mask == + transaction->required_participant_mask) { + load_transactions_.erase(transaction_it); + } +} + +std::vector +HierarchyKVCacheTransfer::store_components() const { + std::vector components; + for (const auto& [participant, state] : participant_states_) { + for (const auto& [block_type, fingerprint] : + state.schema.component_fingerprints) { + HostCacheComponentSchema component; + component.participant = participant; + component.block_type = block_type; + component.model_identity = state.registration.model_identity; + component.schema_fingerprint = fingerprint; + component.tp_rank = state.registration.tp_rank; + component.tp_size = state.registration.tp_size; + components.emplace_back(std::move(component)); + } + } + return components; +} + +std::vector HierarchyKVCacheTransfer::component_storage_tensors( + CacheParticipant participant, + BlockType block_type) const { + const ParticipantState& state = participant_states_.at(participant); + const auto cache_it = state.host_grouped_caches.find(block_type); + CHECK(cache_it != state.host_grouped_caches.end() && + cache_it->second != nullptr); + const BlockTypeTensorMap tensors = + cache_it->second->get_block_type_tensors(block_type); + std::vector storage; + storage.reserve(tensors.size()); + for (const auto& [role, tensor] : tensors) { + (void)role; + storage.emplace_back(tensor); + } + return storage; +} + +std::vector HierarchyKVCacheTransfer::component_tensors( + CacheParticipant participant, + BlockType block_type, + int32_t host_block_id) const { + const std::vector storage = + component_storage_tensors(participant, block_type); + std::vector blocks; + blocks.reserve(storage.size()); + for (const torch::Tensor& tensor : storage) { + CHECK_GE(host_block_id, 0); + CHECK_LT(host_block_id, tensor.size(0)); + torch::Tensor block = tensor[host_block_id]; + CHECK(block.is_contiguous()); + blocks.emplace_back(std::move(block)); + } + return blocks; +} + +uint32_t HierarchyKVCacheTransfer::participant_mask( + CacheParticipant participant) { + return 1u << static_cast(participant); +} + +uint64_t HierarchyKVCacheTransfer::schema_fingerprint_value( + const std::string& fingerprint) { + CHECK_GE(fingerprint.size(), sizeof(uint64_t)); + uint64_t value = 0; + std::memcpy(&value, fingerprint.data(), sizeof(value)); + return value; +} + +uint32_t HierarchyKVCacheTransfer::required_participant_mask( + BlockType block_type) const { + uint32_t mask = 0; + for (const auto& [participant, state] : participant_states_) { + if (participant_requires_type(state, block_type)) { + mask |= participant_mask(participant); + } + } + return mask; +} + +uint64_t HierarchyKVCacheTransfer::composite_schema_fingerprint( + BlockType block_type) const { + std::string composite_schema; + for (const auto& [participant, state] : participant_states_) { + const auto fingerprint_it = + state.schema.component_fingerprints.find(block_type); + if (fingerprint_it == state.schema.component_fingerprints.end()) { + continue; + } + composite_schema.append(participant_name(participant)); + composite_schema.append(fingerprint_it->second); + } + CHECK(!composite_schema.empty()); + return schema_fingerprint_value(hash_schema_string(composite_schema)); +} + +bool HierarchyKVCacheTransfer::participant_requires_type( + const ParticipantState& state, + BlockType block_type) const { + return state.device_grouped_caches.find(block_type) != + state.device_grouped_caches.end(); +} + +bool HierarchyKVCacheTransfer::begin_entry_write( + const std::vector& block_transfer_info) { + std::lock_guard lock(mutex_); + for (const BlockTransferInfo& info : block_transfer_info) { + const uint32_t required_mask = required_participant_mask(info.block_type); + if (required_mask == 0) { + LOG(ERROR) << "No participant owns requested BlockType " + << static_cast(info.block_type); + return false; + } + HostCacheEntryMetadata& metadata = + entry_metadata_[{info.block_type, info.dst_block_id}]; + ++metadata.generation; + metadata.schema_fingerprint = composite_schema_fingerprint(info.block_type); + metadata.required_participant_mask = required_mask; + metadata.completed_participant_mask = 0; + metadata.state = HostEntryState::WRITING; + } + return true; +} + +void HierarchyKVCacheTransfer::complete_participant_write( + CacheParticipant participant, + const std::vector& block_transfer_info) { + std::lock_guard lock(mutex_); + const uint32_t mask = participant_mask(participant); + for (const BlockTransferInfo& info : block_transfer_info) { + HostCacheEntryMetadata& metadata = + entry_metadata_.at({info.block_type, info.dst_block_id}); + if ((metadata.required_participant_mask & mask) != 0) { + metadata.completed_participant_mask |= mask; + } + } +} + +void HierarchyKVCacheTransfer::commit_entry_write( + const std::vector& block_transfer_info) { + std::lock_guard lock(mutex_); + for (const BlockTransferInfo& info : block_transfer_info) { + HostCacheEntryMetadata& metadata = + entry_metadata_.at({info.block_type, info.dst_block_id}); + CHECK_EQ(metadata.completed_participant_mask, + metadata.required_participant_mask); + metadata.state = HostEntryState::READY; + } +} + +void HierarchyKVCacheTransfer::abort_entry_write( + const std::vector& block_transfer_info) { + std::lock_guard lock(mutex_); + for (const BlockTransferInfo& info : block_transfer_info) { + HostCacheEntryMetadata& metadata = + entry_metadata_[{info.block_type, info.dst_block_id}]; + metadata.state = HostEntryState::INVALID; + } +} + +bool HierarchyKVCacheTransfer::snapshot_ready_entries( + const std::vector& block_transfer_info, + LoadTransaction* transaction) const { + CHECK(transaction != nullptr); + std::lock_guard lock(mutex_); + transaction->entry_generations.clear(); + for (const BlockTransferInfo& info : block_transfer_info) { + const std::pair entry_key = {info.block_type, + info.src_block_id}; + const auto metadata_it = entry_metadata_.find(entry_key); + if (metadata_it == entry_metadata_.end()) { + return false; + } + const HostCacheEntryMetadata& metadata = metadata_it->second; + if (metadata.state != HostEntryState::READY || + metadata.required_participant_mask != + required_participant_mask(info.block_type) || + metadata.schema_fingerprint != + composite_schema_fingerprint(info.block_type)) { + return false; + } + transaction->entry_generations[entry_key] = metadata.generation; + } + return true; +} + +bool HierarchyKVCacheTransfer::entry_snapshot_matches( + const std::vector& block_transfer_info, + const LoadTransaction& transaction) const { + std::lock_guard lock(mutex_); + for (const BlockTransferInfo& info : block_transfer_info) { + const std::pair entry_key = {info.block_type, + info.src_block_id}; + const auto generation_it = transaction.entry_generations.find(entry_key); + const auto metadata_it = entry_metadata_.find(entry_key); + if (generation_it == transaction.entry_generations.end() || + metadata_it == entry_metadata_.end()) { + return false; + } + const HostCacheEntryMetadata& metadata = metadata_it->second; + if (metadata.state != HostEntryState::READY || + metadata.generation != generation_it->second || + metadata.required_participant_mask != + required_participant_mask(info.block_type) || + metadata.schema_fingerprint != + composite_schema_fingerprint(info.block_type)) { + return false; + } + } + return true; +} + +void HierarchyKVCacheTransfer::update_prefetched_entries( + Slice& block_transfer_info, + const std::vector& statuses) { + CHECK_EQ(block_transfer_info.size(), statuses.size()); std::lock_guard lock(mutex_); - auto it = layer_wise_load_synchronizer_.find(params.meta.batch_id); - if (it != layer_wise_load_synchronizer_.end()) { - params.parallel.layer_wise_load_synchronizer = it->second; - params.parallel.layers_per_bacth_copy = - layer_batch_ranges_.empty() - ? options_.layers() - : static_cast(layer_batch_ranges_[0].end_layer - - layer_batch_ranges_[0].begin_layer); - layer_wise_load_synchronizer_.erase(it); + for (size_t index = 0; index < block_transfer_info.size(); ++index) { + const BlockTransferInfo& info = block_transfer_info[index]; + HostCacheEntryMetadata& metadata = + entry_metadata_[{info.block_type, info.dst_block_id}]; + ++metadata.generation; + metadata.schema_fingerprint = composite_schema_fingerprint(info.block_type); + metadata.required_participant_mask = + required_participant_mask(info.block_type); + metadata.completed_participant_mask = + statuses[index] == 1 ? metadata.required_participant_mask : 0; + metadata.state = + statuses[index] == 1 ? HostEntryState::READY : HostEntryState::INVALID; } } diff --git a/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h b/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h index 581f58aa2b..53b7743582 100644 --- a/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h +++ b/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h @@ -17,9 +17,11 @@ limitations under the License. #include +#include #include #include #include +#include #include #include @@ -36,9 +38,53 @@ limitations under the License. #include "util/threadpool.h" namespace xllm { + class KVCacheStore; -class HierarchyKVCacheTransfer { +using HostGroupedCaches = std::map>; + +enum class CacheParticipant : int8_t { + TARGET = 0, + DRAFT = 1, +}; + +struct HostCacheComponentSchema { + CacheParticipant participant = CacheParticipant::TARGET; + BlockType block_type = BlockType::KV; + std::string model_identity; + std::string schema_fingerprint; + uint32_t tp_rank = 0; + uint32_t tp_size = 1; +}; + +class HostCacheSliceProvider { + public: + virtual ~HostCacheSliceProvider() = default; + + virtual std::vector store_components() const = 0; + virtual std::vector component_storage_tensors( + CacheParticipant participant, + BlockType block_type) const = 0; + virtual std::vector component_tensors( + CacheParticipant participant, + BlockType block_type, + int32_t host_block_id) const = 0; +}; + +enum class HostPayloadCoverage : int8_t { + REQUIRED = 0, + NOT_APPLICABLE = 1, + REGENERATED = 2, +}; + +enum class HostEntryState : int8_t { + EMPTY = 0, + WRITING = 1, + READY = 2, + INVALID = 3, +}; + +class HierarchyKVCacheTransfer final : public HostCacheSliceProvider { public: struct LayerBatchRange { int64_t begin_layer = 0; @@ -52,11 +98,6 @@ class HierarchyKVCacheTransfer { using GroupedCaches = std::map>; - // Host prefix caches: one real KVCache per block type, allocated over - // page-aligned + mlock'd + NPU-registered host memory. Shape per tensor is - // [host_blocks, layer_count, ...per_block_dims]. - using HostGroupedCaches = std::map>; - struct Options { PROPERTY(uint32_t, tp_rank); PROPERTY(uint32_t, tp_size); @@ -73,63 +114,167 @@ class HierarchyKVCacheTransfer { PROPERTY(uint32_t, store_worker_id) = 0; }; + struct HostCacheTensorSpec { + CacheParticipant participant = CacheParticipant::TARGET; + BlockType block_type = BlockType::KV; + int64_t absolute_layer_id = 0; + int64_t host_layer_slot = 0; + KVCacheTensorRole::Value role = KVCacheTensorRole::KEY; + HostPayloadCoverage coverage = HostPayloadCoverage::REQUIRED; + torch::ScalarType dtype = torch::kFloat32; + std::vector block_shape; + }; + + struct ParticipantHostCacheSchema { + std::vector tensor_specs; + std::map component_fingerprints; + }; + + struct ParticipantRegistration { + CacheParticipant participant = CacheParticipant::TARGET; + const Stream* actual_compute_stream = nullptr; + std::vector* device_caches = nullptr; + KVCacheShape cache_shape; + KVCacheCreateOptions create_options; + std::string model_identity; + uint32_t tp_rank = 0; + uint32_t tp_size = 1; + }; + + struct HostCacheEntryMetadata { + uint64_t schema_fingerprint = 0; + uint64_t generation = 0; + uint32_t required_participant_mask = 0; + uint32_t completed_participant_mask = 0; + HostEntryState state = HostEntryState::EMPTY; + }; + + HierarchyKVCacheTransfer(const Options& options, const torch::Device& device); HierarchyKVCacheTransfer(const Options& options, const torch::Device& device, const Stream* compute_stream, - std::vector* kv_caches_ptr, + std::vector* kv_caches_ptr, const KVCacheShape& kv_cache_shape, const KVCacheCreateOptions& create_options); - ~HierarchyKVCacheTransfer(); + ~HierarchyKVCacheTransfer() override; + + void register_participant(ParticipantRegistration registration); + bool finalize_registration(); uint32_t transfer_kv_blocks( - const uint64_t batch_id, + uint64_t batch_id, const std::vector& block_transfer_info); - uint32_t transfer_kv_blocks(const uint64_t batch_id, + uint32_t transfer_kv_blocks(uint64_t batch_id, Slice& block_transfer_info); std::vector prefetch_kv_blocks( Slice& block_transfer_info); void set_layer_synchronizer(ModelInputParams& params); + void set_layer_synchronizer(CacheParticipant participant, + ModelInputParams& params); + + std::vector store_components() const override; + std::vector component_storage_tensors( + CacheParticipant participant, + BlockType block_type) const override; + std::vector component_tensors( + CacheParticipant participant, + BlockType block_type, + int32_t host_block_id) const override; private: friend class HierarchyKVCacheTransferTestPeer; - void build_device_block_type_map(); - void create_host_cache(); + enum class RegistrationState : int8_t { + CREATED = 0, + REGISTERING = 1, + READY = 2, + }; + + struct ParticipantState { + ParticipantRegistration registration; + GroupedCaches device_grouped_caches; + std::map> absolute_layer_ids; + HostGroupedCaches host_grouped_caches; + std::vector layer_batch_ranges; + ParticipantHostCacheSchema schema; + }; + + struct LoadTransaction { + std::mutex mutex; + std::map> + synchronizers; + std::map, uint64_t> entry_generations; + uint32_t required_participant_mask = 0; + uint32_t consumed_participant_mask = 0; + bool aborted = false; + + void abort(); + }; + + static uint32_t participant_mask(CacheParticipant participant); + static uint64_t schema_fingerprint_value(const std::string& fingerprint); + + void initialize_resources(); + void build_participant_state(ParticipantState& state); + void build_device_block_type_map(ParticipantState& state); + void create_host_cache(ParticipantState& state); + void build_and_validate_schema(ParticipantState& state); + void validate_composite_schema() const; + void initialize_store(); + CopyPlan build_copy_plan( + const ParticipantState& state, const std::vector& block_transfer_info, const LayerBatchRange& layer_batch_range) const; uint32_t offload(const std::vector& block_transfer_info); bool offload_to_host(Slice& block_transfer_info); bool load_from_host( - std::shared_ptr synchronizer, + CacheParticipant participant, + const std::shared_ptr& transaction, + const std::vector& block_transfer_info); + + uint32_t required_participant_mask(BlockType block_type) const; + uint64_t composite_schema_fingerprint(BlockType block_type) const; + bool participant_requires_type(const ParticipantState& state, + BlockType block_type) const; + bool begin_entry_write( + const std::vector& block_transfer_info); + void complete_participant_write( + CacheParticipant participant, + const std::vector& block_transfer_info); + void commit_entry_write( const std::vector& block_transfer_info); + void abort_entry_write( + const std::vector& block_transfer_info); + bool snapshot_ready_entries( + const std::vector& block_transfer_info, + LoadTransaction* transaction) const; + bool entry_snapshot_matches( + const std::vector& block_transfer_info, + const LoadTransaction& transaction) const; + void update_prefetched_entries(Slice& block_transfer_info, + const std::vector& statuses); private: Options options_; Device device_; - const Stream* compute_stream_ = nullptr; + RegistrationState registration_state_ = RegistrationState::CREATED; + std::map participant_states_; std::unique_ptr load_threadpool_; moodycamel::BlockingConcurrentQueue> copy_stream_; - - std::vector* kv_caches_ptr_ = nullptr; - KVCacheShape kv_cache_shape_; - KVCacheCreateOptions create_options_; - GroupedCaches device_kv_caches_; - std::map> device_block_type_layer_ids_; - HostGroupedCaches host_kv_caches_; - std::vector layer_batch_ranges_; - std::unique_ptr batch_memcpy_; std::unique_ptr kv_cache_store_; mutable std::mutex mutex_; - std::unordered_map> - layer_wise_load_synchronizer_; + std::map, HostCacheEntryMetadata> + entry_metadata_; + std::unordered_map> + load_transactions_; }; } // namespace xllm diff --git a/xllm/core/framework/kv_cache_transfer/kv_cache_store.cpp b/xllm/core/framework/kv_cache_transfer/kv_cache_store.cpp index ebf2e8393a..90d35e7fbc 100644 --- a/xllm/core/framework/kv_cache_transfer/kv_cache_store.cpp +++ b/xllm/core/framework/kv_cache_transfer/kv_cache_store.cpp @@ -15,26 +15,120 @@ limitations under the License. #include "framework/kv_cache_transfer/kv_cache_store.h" +#include #include #include #include #include +#include #include #include +#include #include #include "util/hash_util.h" namespace xllm { +namespace { + +std::string participant_name(CacheParticipant participant) { + switch (participant) { + case CacheParticipant::TARGET: + return "TARGET"; + case CacheParticipant::DRAFT: + return "DRAFT"; + } + LOG(FATAL) << "Unsupported cache participant: " + << static_cast(participant); + return "UNKNOWN"; +} + +void append_key_field(std::string& key, const std::string& value) { + key.append(std::to_string(value.size())); + key.push_back(':'); + key.append(value); + key.push_back(':'); +} + +struct ComponentRequest final { + size_t logical_index = 0; + const HostCacheComponentSchema* component = nullptr; + std::string key; +}; + +std::vector generate_mooncake_slices( + const HostCacheSliceProvider* slice_provider, + const HostCacheComponentSchema& component, + int32_t block_id) { + CHECK(slice_provider != nullptr); + const std::vector tensors = slice_provider->component_tensors( + component.participant, component.block_type, block_id); + CHECK(!tensors.empty()) << "Missing Host cache slices for participant=" + << participant_name(component.participant) + << ", type=" + << static_cast(component.block_type); + + std::vector slices; + slices.reserve(tensors.size()); + for (const torch::Tensor& tensor : tensors) { + CHECK(tensor.defined() && tensor.is_contiguous()); + slices.emplace_back( + mooncake::Slice{tensor.data_ptr(), + static_cast(tensor.numel()) * + static_cast(tensor.element_size())}); + } + return slices; +} + +bool copy_slices(const std::vector& source, + const std::vector& destination) { + if (source.size() != destination.size()) { + return false; + } + for (size_t index = 0; index < source.size(); ++index) { + if (source[index].size != destination[index].size || + source[index].ptr == nullptr || destination[index].ptr == nullptr) { + return false; + } + if (source[index].ptr != destination[index].ptr) { + std::memcpy( + destination[index].ptr, source[index].ptr, source[index].size); + } + } + return true; +} + +} // namespace + +struct KVCacheStore::Impl final { + mooncake::ReplicateConfig rep_config; + std::shared_ptr client; +}; + +KVCacheStore::KVCacheStore() : impl_(std::make_unique()) {} bool KVCacheStore::init(const KVCacheStoreInitConfig& config, - HostGroupedCaches* host_kv_caches) { + const HostCacheSliceProvider* slice_provider) { CHECK(!is_initialized_) << "KVCacheStore is already initialized."; - CHECK(host_kv_caches != nullptr && !host_kv_caches->empty()) - << "KVCacheStore requires typed Host caches."; + CHECK(slice_provider != nullptr) + << "KVCacheStore requires a Host cache slice provider."; config_ = config; - host_kv_caches_ = host_kv_caches; + slice_provider_ = slice_provider; + components_ = slice_provider_->store_components(); + CHECK(!components_.empty()) + << "KVCacheStore requires at least one Host cache component."; + std::sort(components_.begin(), + components_.end(), + [](const HostCacheComponentSchema& lhs, + const HostCacheComponentSchema& rhs) { + if (lhs.participant != rhs.participant) { + return static_cast(lhs.participant) < + static_cast(rhs.participant); + } + return static_cast(lhs.block_type) < + static_cast(rhs.block_type); + }); std::optional device_names = std::nullopt; if (config_.protocol == "rdma") { @@ -58,21 +152,29 @@ bool KVCacheStore::init(const KVCacheStoreInitConfig& config, << config_.localhost_name; return false; } - client_ptr_ = client.value(); - rep_config_.replica_num = config_.replica_num; + impl_->client = client.value(); + impl_->rep_config.replica_num = config_.replica_num; - std::string cache_schema = "tp=" + std::to_string(config_.tp_size); - for (const auto& [type, cache] : *host_kv_caches_) { - CHECK(cache != nullptr); - const BlockTypeTensorMap tensors = cache->get_block_type_tensors(type); - CHECK(!tensors.empty()) << "Host cache has no tensors for BlockType " - << static_cast(type); + std::unordered_set component_ids; + for (const HostCacheComponentSchema& component : components_) { + CHECK_GT(component.tp_size, 0u); + CHECK_LT(component.tp_rank, component.tp_size); + CHECK(!component.model_identity.empty()); + CHECK(!component.schema_fingerprint.empty()); + const std::string component_id = + participant_name(component.participant) + ":" + + std::to_string(static_cast(component.block_type)); + CHECK(component_ids.emplace(component_id).second) + << "Duplicate Host cache Store component: " << component_id; - size_t slot_bytes = 0; + const std::vector tensors = + slice_provider_->component_storage_tensors(component.participant, + component.block_type); + CHECK(!tensors.empty()) + << "Host cache component has no tensors: " << component_id; int64_t host_blocks = -1; - cache_schema.append("|type="); - cache_schema.append(std::to_string(static_cast(type))); - for (const auto& [role, tensor] : tensors) { + size_t slot_bytes = 0; + for (const torch::Tensor& tensor : tensors) { CHECK(tensor.defined() && tensor.dim() > 0 && tensor.is_contiguous()); if (host_blocks < 0) { host_blocks = tensor.size(0); @@ -81,117 +183,162 @@ bool KVCacheStore::init(const KVCacheStoreInitConfig& config, } slot_bytes += static_cast(tensor[0].numel()) * static_cast(tensor.element_size()); - cache_schema.append(",role="); - cache_schema.append(std::to_string(static_cast(role))); - cache_schema.append(",dtype="); - cache_schema.append( - std::to_string(static_cast(tensor.scalar_type()))); - cache_schema.append(",shape="); - for (int64_t dim = 1; dim < tensor.dim(); ++dim) { - cache_schema.append(std::to_string(tensor.size(dim))); - cache_schema.push_back('x'); + if (config_.protocol != "rdma") { + continue; } - - if (config_.protocol == "rdma") { - void* address = tensor.data_ptr(); - const size_t bytes = static_cast(tensor.numel()) * - static_cast(tensor.element_size()); - auto result = - client_ptr_->RegisterLocalMemory(address, + void* address = tensor.data_ptr(); + const size_t bytes = static_cast(tensor.numel()) * + static_cast(tensor.element_size()); + auto result = + impl_->client->RegisterLocalMemory(address, bytes, /*location=*/"cpu:0", /*remote_accessible=*/false, /*update_metadata=*/false); - if (!result.has_value()) { - LOG(ERROR) << "Failed to register Mooncake Host tensor: " - << toString(result.error()); - return false; - } - registered_addresses_.emplace_back(address); + if (!result.has_value()) { + LOG(ERROR) << "Failed to register Mooncake Host tensor: " + << toString(result.error()); + return false; } + registered_addresses_.emplace_back(address); } - LOG(INFO) << "KVCacheStore init OK: type=" << static_cast(type) + LOG(INFO) << "KVCacheStore init OK: type=" + << static_cast(component.block_type) + << ", participant=" << participant_name(component.participant) << ", host_blocks=" << host_blocks << ", slot_bytes=" << slot_bytes << ", protocol=" << config_.protocol; } - const XXH3Key schema_hash = hash_string(cache_schema); - cache_schema_hash_.assign(reinterpret_cast(schema_hash.data), - sizeof(schema_hash.data)); is_initialized_ = true; return true; } KVCacheStore::~KVCacheStore() { - if (client_ptr_ != nullptr) { + if (impl_->client != nullptr) { for (void* address : registered_addresses_) { - auto result = client_ptr_->unregisterLocalMemory( + auto result = impl_->client->unregisterLocalMemory( address, /*update_metadata=*/false); if (!result.has_value()) { LOG(WARNING) << "Failed to unregister Mooncake Host tensor: " << toString(result.error()); } } - client_ptr_.reset(); + impl_->client.reset(); } } -std::string KVCacheStore::build_key(const BlockTransferInfo& block_info) const { - std::string key = "xllm-kv-v2:"; - key.append(std::to_string(config_.model_id.size())); - key.push_back(':'); - key.append(config_.model_id); - key.push_back(':'); - key.append(std::to_string(config_.tp_size)); +std::string KVCacheStore::build_component_key( + const HostCacheComponentSchema& component, + const BlockTransferInfo& block_info) const { + std::string key = "xllm-kv-v3:"; + append_key_field(key, config_.model_id); + append_key_field(key, participant_name(component.participant)); + append_key_field(key, component.model_identity); + append_key_field(key, "composite-host-v1"); + key.append(std::to_string(component.tp_size)); key.push_back(':'); - key.append(std::to_string(static_cast(block_info.block_type))); + key.append(std::to_string(component.tp_rank)); key.push_back(':'); - key.append(std::to_string(config_.tp_rank)); + key.append(std::to_string(static_cast(component.block_type))); key.push_back(':'); - key.append(cache_schema_hash_); + append_key_field(key, component.schema_fingerprint); key.append(reinterpret_cast(block_info.hash_key), XXH3_128BITS_HASH_VALUE_LEN); return key; } +std::vector KVCacheStore::required_components( + BlockType block_type) const { + std::vector components; + for (const HostCacheComponentSchema& component : components_) { + if (component.block_type == block_type) { + components.emplace_back(&component); + } + } + return components; +} + uint32_t KVCacheStore::batch_put( Slice& block_transfer_info) { if (!is_initialized_ || block_transfer_info.empty()) { return 0; } + std::vector requests; + for (size_t logical_index = 0; logical_index < block_transfer_info.size(); + ++logical_index) { + const BlockTransferInfo& block_info = block_transfer_info[logical_index]; + const std::vector components = + required_components(block_info.block_type); + for (const HostCacheComponentSchema* component : components) { + requests.push_back({logical_index, + component, + build_component_key(*component, block_info)}); + } + } + if (requests.empty()) { + return 0; + } + std::vector all_keys; - all_keys.reserve(block_transfer_info.size()); - for (const BlockTransferInfo& block_info : block_transfer_info) { - all_keys.emplace_back(build_key(block_info)); + all_keys.reserve(requests.size()); + for (const ComponentRequest& request : requests) { + all_keys.emplace_back(request.key); } - const auto exists = client_ptr_->BatchIsExist(all_keys); + const auto exists = impl_->client->BatchIsExist(all_keys); std::vector put_keys; std::vector> put_slices; - put_keys.reserve(block_transfer_info.size()); - put_slices.reserve(block_transfer_info.size()); - uint32_t success_count = 0; - for (size_t i = 0; i < block_transfer_info.size(); ++i) { - const bool already_exists = - i < exists.size() && exists[i].has_value() && exists[i].value(); + std::vector> put_request_groups; + std::unordered_map put_key_indices; + std::vector completed_components(block_transfer_info.size(), 0); + std::vector required_component_counts(block_transfer_info.size(), + 0); + for (size_t request_index = 0; request_index < requests.size(); + ++request_index) { + const ComponentRequest& request = requests[request_index]; + ++required_component_counts[request.logical_index]; + const bool already_exists = request_index < exists.size() && + exists[request_index].has_value() && + exists[request_index].value(); if (already_exists) { - ++success_count; + ++completed_components[request.logical_index]; continue; } - put_keys.emplace_back(all_keys[i]); - put_slices.emplace_back( - generate_mooncake_slices(block_transfer_info[i].block_type, - block_transfer_info[i].dst_block_id)); + const auto [put_it, inserted] = + put_key_indices.emplace(request.key, put_keys.size()); + if (inserted) { + put_keys.emplace_back(request.key); + put_slices.emplace_back(generate_mooncake_slices( + slice_provider_, + *request.component, + block_transfer_info[request.logical_index].dst_block_id)); + put_request_groups.emplace_back(); + } + put_request_groups[put_it->second].emplace_back(request_index); } - if (put_keys.empty()) { - return success_count; + if (!put_keys.empty()) { + const auto results = + impl_->client->BatchPut(put_keys, put_slices, impl_->rep_config); + for (size_t i = 0; i < put_request_groups.size() && i < results.size(); + ++i) { + if (!results[i].has_value()) { + continue; + } + for (size_t request_index : put_request_groups[i]) { + ++completed_components[requests[request_index].logical_index]; + } + } } - const auto results = client_ptr_->BatchPut(put_keys, put_slices, rep_config_); - for (size_t i = 0; i < put_keys.size() && i < results.size(); ++i) { - if (results[i].has_value()) { + + uint32_t success_count = 0; + for (size_t logical_index = 0; logical_index < block_transfer_info.size(); + ++logical_index) { + if (required_component_counts[logical_index] > 0 && + completed_components[logical_index] == + required_component_counts[logical_index]) { ++success_count; } } @@ -213,40 +360,98 @@ std::vector KVCacheStore::batch_get_with_status( return statuses; } + std::vector requests; + std::vector required_component_counts(block_transfer_info.size(), + 0); + for (size_t logical_index = 0; logical_index < block_transfer_info.size(); + ++logical_index) { + const BlockTransferInfo& block_info = block_transfer_info[logical_index]; + const std::vector components = + required_components(block_info.block_type); + for (const HostCacheComponentSchema* component : components) { + requests.push_back({logical_index, + component, + build_component_key(*component, block_info)}); + ++required_component_counts[logical_index]; + } + } + if (requests.empty()) { + return statuses; + } + std::vector all_keys; - all_keys.reserve(block_transfer_info.size()); - for (const BlockTransferInfo& block_info : block_transfer_info) { - all_keys.emplace_back(build_key(block_info)); + all_keys.reserve(requests.size()); + for (const ComponentRequest& request : requests) { + all_keys.emplace_back(request.key); + } + const auto exists = impl_->client->BatchIsExist(all_keys); + + std::vector existing_component_counts(block_transfer_info.size(), + 0); + for (size_t request_index = 0; request_index < requests.size(); + ++request_index) { + if (request_index < exists.size() && exists[request_index].has_value() && + exists[request_index].value()) { + ++existing_component_counts[requests[request_index].logical_index]; + } } - const auto exists = client_ptr_->BatchIsExist(all_keys); std::vector get_keys; std::unordered_map> get_slices; - std::vector get_positions; - get_keys.reserve(block_transfer_info.size()); - get_positions.reserve(block_transfer_info.size()); - get_slices.reserve(block_transfer_info.size()); - for (size_t i = 0; i < block_transfer_info.size(); ++i) { - const bool exists_in_store = - i < exists.size() && exists[i].has_value() && exists[i].value(); - if (!exists_in_store) { + std::vector> get_request_groups; + std::unordered_map get_key_indices; + for (size_t request_index = 0; request_index < requests.size(); + ++request_index) { + const ComponentRequest& request = requests[request_index]; + const size_t logical_index = request.logical_index; + if (required_component_counts[logical_index] == 0 || + existing_component_counts[logical_index] != + required_component_counts[logical_index]) { continue; } - get_positions.emplace_back(i); - get_keys.emplace_back(all_keys[i]); - get_slices.emplace( - all_keys[i], - generate_mooncake_slices(block_transfer_info[i].block_type, - block_transfer_info[i].dst_block_id)); + const auto [get_it, inserted] = + get_key_indices.emplace(request.key, get_keys.size()); + if (inserted) { + get_keys.emplace_back(request.key); + get_slices.emplace(request.key, + generate_mooncake_slices( + slice_provider_, + *request.component, + block_transfer_info[logical_index].dst_block_id)); + get_request_groups.emplace_back(); + } + get_request_groups[get_it->second].emplace_back(request_index); } - if (get_keys.empty()) { return statuses; } - const auto results = client_ptr_->BatchGet(get_keys, get_slices); - for (size_t i = 0; i < get_keys.size() && i < results.size(); ++i) { - if (results[i].has_value()) { - statuses[get_positions[i]] = 1; + + const auto results = impl_->client->BatchGet(get_keys, get_slices); + std::vector fetched_component_counts(block_transfer_info.size(), 0); + for (size_t i = 0; i < get_request_groups.size() && i < results.size(); ++i) { + if (!results[i].has_value()) { + continue; + } + const std::vector& source_slices = + get_slices.at(get_keys[i]); + for (size_t request_index : get_request_groups[i]) { + const ComponentRequest& request = requests[request_index]; + const std::vector destination_slices = + generate_mooncake_slices( + slice_provider_, + *request.component, + block_transfer_info[request.logical_index].dst_block_id); + if (copy_slices(source_slices, destination_slices)) { + ++fetched_component_counts[request.logical_index]; + } + } + } + for (size_t logical_index = 0; logical_index < block_transfer_info.size(); + ++logical_index) { + if (required_component_counts[logical_index] > 0 && + fetched_component_counts[logical_index] == + required_component_counts[logical_index]) { + statuses[logical_index] = 1; } } return statuses; @@ -256,37 +461,11 @@ uint32_t KVCacheStore::batch_exist(std::vector&& keys) { if (!is_initialized_) { return 0; } - const auto exists = client_ptr_->BatchIsExist(keys); + const auto exists = impl_->client->BatchIsExist(keys); return static_cast( std::count_if(exists.begin(), exists.end(), [](const auto& result) { return result.has_value() && result.value(); })); } -std::vector KVCacheStore::generate_mooncake_slices( - BlockType type, - int32_t block_id) const { - CHECK(host_kv_caches_ != nullptr); - const auto cache_it = host_kv_caches_->find(type); - CHECK(cache_it != host_kv_caches_->end() && cache_it->second != nullptr) - << "Missing Host cache for BlockType " << static_cast(type); - const BlockTypeTensorMap tensors = - cache_it->second->get_block_type_tensors(type); - - std::vector slices; - slices.reserve(tensors.size()); - for (const auto& tensor_entry : tensors) { - const torch::Tensor& tensor = tensor_entry.second; - CHECK_GE(block_id, 0); - CHECK_LT(block_id, tensor.size(0)); - torch::Tensor block = tensor[block_id]; - CHECK(block.is_contiguous()); - slices.emplace_back( - mooncake::Slice{block.data_ptr(), - static_cast(block.numel()) * - static_cast(block.element_size())}); - } - return slices; -} - } // namespace xllm diff --git a/xllm/core/framework/kv_cache_transfer/kv_cache_store.h b/xllm/core/framework/kv_cache_transfer/kv_cache_store.h index 7e606c3de5..78941d452c 100644 --- a/xllm/core/framework/kv_cache_transfer/kv_cache_store.h +++ b/xllm/core/framework/kv_cache_transfer/kv_cache_store.h @@ -15,22 +15,17 @@ limitations under the License. #pragma once -#include - #include -#include #include #include +#include #include -#include "framework/kv_cache/kv_cache.h" -#include "framework/model/model_input_params.h" +#include "framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h" #include "util/slice.h" namespace xllm { -using HostGroupedCaches = std::map>; - struct KVCacheStoreInitConfig { std::string localhost_name = "127.0.0.1"; std::string protocol = "tcp"; @@ -44,11 +39,11 @@ struct KVCacheStoreInitConfig { class KVCacheStore final { public: - KVCacheStore() = default; + KVCacheStore(); ~KVCacheStore(); bool init(const KVCacheStoreInitConfig& config, - HostGroupedCaches* host_kv_caches); + const HostCacheSliceProvider* slice_provider); uint32_t batch_put( const std::vector& block_transfer_info) { @@ -70,21 +65,24 @@ class KVCacheStore final { uint32_t batch_exist(std::vector&& keys); private: + friend class KVCacheStoreTestPeer; + struct Impl; + KVCacheStore(const KVCacheStore&) = delete; KVCacheStore& operator=(const KVCacheStore&) = delete; - std::string build_key(const BlockTransferInfo& block_info) const; - std::vector generate_mooncake_slices(BlockType type, - int32_t block_id) const; + std::string build_component_key(const HostCacheComponentSchema& component, + const BlockTransferInfo& block_info) const; + std::vector required_components( + BlockType block_type) const; private: bool is_initialized_ = false; KVCacheStoreInitConfig config_; - std::string cache_schema_hash_; - mooncake::ReplicateConfig rep_config_; - HostGroupedCaches* host_kv_caches_ = nullptr; + const HostCacheSliceProvider* slice_provider_ = nullptr; + std::vector components_; std::vector registered_addresses_; - std::shared_ptr client_ptr_; + std::unique_ptr impl_; }; } // namespace xllm diff --git a/xllm/core/framework/prefix_cache/prefix_cache.cpp b/xllm/core/framework/prefix_cache/prefix_cache.cpp index 983bacf44d..2fab5abc5c 100644 --- a/xllm/core/framework/prefix_cache/prefix_cache.cpp +++ b/xllm/core/framework/prefix_cache/prefix_cache.cpp @@ -30,7 +30,7 @@ std::vector PrefixCache::match(const Slice& token_ids, const Slice& existed_shared_blocks, const MMData& mm_data, const Slice& block_hashes) { - // allign tokens to block boundary + // align tokens to block boundary const size_t n_tokens = round_down(token_ids.size(), block_size_); if (n_tokens == 0) { return std::vector(); @@ -109,7 +109,7 @@ size_t PrefixCache::insert(const Slice& token_ids, const MMData& mm_data, const Slice& block_hashes) { const int64_t now = absl::ToUnixMicros(absl::Now()); - // allign tokens to block boundary + // align tokens to block boundary const size_t n_blocks = std::min(token_ids.size() / block_size_, blocks.size()); @@ -274,6 +274,21 @@ size_t PrefixCache::evict(size_t n_blocks) { return evict_count; } +bool PrefixCache::can_evict(size_t n_blocks) const { + if (n_blocks == 0) { + return true; + } + size_t evictable = 0; + for (const auto& [hash, node] : cached_blocks_) { + (void)hash; + CHECK(node != nullptr); + if (!node->block.is_shared() && ++evictable >= n_blocks) { + return true; + } + } + return false; +} + uint32_t PrefixCache::compute_hash_keys(const Slice& token_ids, std::vector& blocks, const size_t cached_blocks) { diff --git a/xllm/core/framework/prefix_cache/prefix_cache.h b/xllm/core/framework/prefix_cache/prefix_cache.h index e95bef219e..6160c414fd 100644 --- a/xllm/core/framework/prefix_cache/prefix_cache.h +++ b/xllm/core/framework/prefix_cache/prefix_cache.h @@ -106,6 +106,10 @@ class PrefixCache { // Evict up to `n_blocks` LRU-oldest entries. Returns the number evicted. virtual size_t evict(size_t n_blocks); + // Return whether at least `n_blocks` entries are currently evictable without + // mutating cache state. Shared entries are pinned by active sequence owners. + bool can_evict(size_t n_blocks) const; + virtual size_t num_blocks() const { CHECK(num_blocks_ == cached_blocks_.size()) << "check block num failed"; return num_blocks_; diff --git a/xllm/core/runtime/mtp_worker_impl.cpp b/xllm/core/runtime/mtp_worker_impl.cpp index f3607b28db..0192d528b4 100644 --- a/xllm/core/runtime/mtp_worker_impl.cpp +++ b/xllm/core/runtime/mtp_worker_impl.cpp @@ -276,22 +276,6 @@ void clear_ready_events(ForwardInput& input) { input.metadata_ready_event.reset(); } -std::optional run_llm_no_sync_impl( - LLMWorkerImpl& worker, - const ForwardInput& input, - Stream& prepare_stream, - Stream& compute_stream, - ForwardInput& processed_input) { - worker.prepare_work_before_execute_on_stream( - input, - processed_input, - prepare_stream, - /*record_ready_event=*/&prepare_stream != &compute_stream); - worker.set_hierarchy_layer_synchronizer(processed_input.input_params); - return worker.execute_no_sync_on_stream( - processed_input, compute_stream, /*record_ready_event=*/false); -} - torch::Tensor clone_host_tensor(const torch::Tensor& tensor) { if (!tensor.defined()) { return tensor; @@ -443,6 +427,17 @@ runtime::Options mtp_draft_options(const runtime::Options& options) { return draft_options; } +runtime::Options hierarchy_child_options( + const runtime::Options& options, + MTPWorkerImpl::HierarchyCacheOwnershipMode ownership_mode) { + runtime::Options child_options = options; + if (ownership_mode == + MTPWorkerImpl::HierarchyCacheOwnershipMode::PARENT_COMPOSITE) { + child_options.host_blocks_factor(0.0).enable_kvcache_store(false); + } + return child_options; +} + ParallelArgs MTPDraftParallelArgs(const ParallelArgs& parallel_args, const runtime::Options& options) { if (!options.enable_mtp_draft_body_tp1()) { @@ -781,21 +776,31 @@ MTPWorkerImpl::MTPWorkerImpl(const ParallelArgs& parallel_args, MTPTargetOptions(options), mtp_draft_options(options), ::xllm::SpeculativeConfig::get_instance().enable_opt_validate_probs(), - /*enable_adaptive_speculative_decode=*/true) {} - -MTPWorkerImpl::MTPWorkerImpl(const ParallelArgs& parallel_args, - const torch::Device& device, - const runtime::Options& options, - const runtime::Options& target_options, - const runtime::Options& draft_options, - bool enable_opt_validate_probs, - bool enable_adaptive_speculative_decode) - : SpeculativeWorkerImpl(parallel_args, device, options, target_options), - enable_opt_validate_probs_(enable_opt_validate_probs) { + /*enable_adaptive_speculative_decode=*/true, + HierarchyCacheOwnershipMode::PARENT_COMPOSITE) {} + +MTPWorkerImpl::MTPWorkerImpl( + const ParallelArgs& parallel_args, + const torch::Device& device, + const runtime::Options& options, + const runtime::Options& target_options, + const runtime::Options& draft_options, + bool enable_opt_validate_probs, + bool enable_adaptive_speculative_decode, + HierarchyCacheOwnershipMode hierarchy_cache_ownership_mode) + : SpeculativeWorkerImpl( + parallel_args, + device, + options, + hierarchy_child_options(target_options, + hierarchy_cache_ownership_mode)), + enable_opt_validate_probs_(enable_opt_validate_probs), + hierarchy_cache_ownership_mode_(hierarchy_cache_ownership_mode) { draft_impl_ = std::make_unique( MTPDraftParallelArgs(parallel_args, options), device, - mtp_draft_options(draft_options)); + hierarchy_child_options(mtp_draft_options(draft_options), + hierarchy_cache_ownership_mode)); const bool enable_parallel_adaptive_sl = parallel_args.dp_size() <= 1 && parallel_args.ep_size() <= 1; if (enable_adaptive_speculative_decode && enable_parallel_adaptive_sl) { @@ -812,6 +817,34 @@ MTPWorkerImpl::MTPWorkerImpl(const ParallelArgs& parallel_args, MTPWorkerImpl::~MTPWorkerImpl() = default; +std::optional MTPWorkerImpl::run_child_llm_no_sync( + CacheParticipant participant, + LLMWorkerImpl& worker, + const ForwardInput& input, + Stream& prepare_stream, + Stream& compute_stream, + ForwardInput& processed_input) { + if (participant == CacheParticipant::TARGET) { + CHECK_EQ(&worker, impl_.get()); + } else { + CHECK(participant == CacheParticipant::DRAFT); + CHECK_EQ(&worker, draft_impl_.get()); + } + worker.prepare_work_before_execute_on_stream( + input, + processed_input, + prepare_stream, + /*record_ready_event=*/&prepare_stream != &compute_stream); + if (composite_hierarchy_kv_cache_transfer_ != nullptr) { + composite_hierarchy_kv_cache_transfer_->set_layer_synchronizer( + participant, processed_input.input_params); + } else { + worker.set_hierarchy_layer_synchronizer(processed_input.input_params); + } + return worker.execute_no_sync_on_stream( + processed_input, compute_stream, /*record_ready_event=*/false); +} + bool MTPWorkerImpl::init_model(const std::string& model_weights_path, int32_t random_seed, MasterStatus master_status) { @@ -1015,7 +1048,90 @@ bool MTPWorkerImpl::allocate_kv_cache(const KVCacheShape& kv_cache_shape) { CHECK_EQ(draft_status, WorkerImpl::Status::READY); } - return target_allocated && draft_allocated; + const bool allocated = target_allocated && draft_allocated; + if (allocated) { + initialize_composite_hierarchy_kv_cache_transfer(); + } + return allocated; +} + +void MTPWorkerImpl::initialize_composite_hierarchy_kv_cache_transfer() { + if (hierarchy_cache_ownership_mode_ != + HierarchyCacheOwnershipMode::PARENT_COMPOSITE || + composite_hierarchy_kv_cache_transfer_ != nullptr || + options_.host_blocks_factor() <= 1.0) { + return; + } + CHECK(impl_ != nullptr); + CHECK(draft_impl_ != nullptr); + CHECK(compute_stream_ != nullptr); + if (options_.enable_kvcache_store()) { + CHECK_GT(options_.host_blocks_factor(), 1.0) + << "KV cache Store requires Host cache blocks."; + } + + CHECK_GT(options_.dp_size(), 0u); + CHECK_EQ(options_.world_size() % options_.dp_size(), 0u); + CHECK_GT(options_.cp_size(), 0u); + const uint32_t dp_local_size = + static_cast(options_.world_size() / options_.dp_size()); + CHECK_EQ(dp_local_size % options_.cp_size(), 0u); + const bool mlu_overlap = options_.cp_size() > 1 && + Platform::uses_model_cp_sharding() && + Platform::is_mlu(); + const uint32_t target_tp_size = + mlu_overlap ? dp_local_size : dp_local_size / options_.cp_size(); + CHECK_GT(target_tp_size, 0u); + const int32_t worker_rank = context_.get_parallel_args().rank(); + CHECK_GE(worker_rank, 0); + const uint32_t worker_id = static_cast(worker_rank); + + HierarchyKVCacheTransfer::Options transfer_options; + transfer_options.tp_rank(worker_id % target_tp_size) + .tp_size(target_tp_size) + .layers(context_.get_model_args().n_layers()) + .host_blocks_factor(options_.host_blocks_factor()) + .layers_wise_copy_batchs(options_.layers_wise_copy_batchs()) + .enable_mla(options_.enable_mla()) + .enable_kvcache_store(options_.enable_kvcache_store()) + .store_protocol(options_.store_protocol()) + .store_master_server_address(options_.store_master_server_address()) + .store_metadata_server(options_.store_metadata_server()) + .store_local_hostname(options_.store_local_hostname()) + .store_namespace(options_.model_id()) + .store_worker_id(worker_id); + auto transfer = + std::make_unique(transfer_options, device_); + + const HierarchyKVCacheAllocationDescriptor& target_descriptor = + impl_->hierarchy_kv_cache_allocation_descriptor(); + HierarchyKVCacheTransfer::ParticipantRegistration target_registration; + target_registration.participant = CacheParticipant::TARGET; + target_registration.actual_compute_stream = compute_stream_.get(); + target_registration.device_caches = target_descriptor.device_caches; + target_registration.cache_shape = target_descriptor.cache_shape; + target_registration.create_options = target_descriptor.create_options; + target_registration.model_identity = target_descriptor.model_identity; + target_registration.tp_rank = worker_id % target_tp_size; + target_registration.tp_size = target_tp_size; + transfer->register_participant(std::move(target_registration)); + + const HierarchyKVCacheAllocationDescriptor& draft_descriptor = + draft_impl_->hierarchy_kv_cache_allocation_descriptor(); + HierarchyKVCacheTransfer::ParticipantRegistration draft_registration; + draft_registration.participant = CacheParticipant::DRAFT; + draft_registration.actual_compute_stream = compute_stream_.get(); + draft_registration.device_caches = draft_descriptor.device_caches; + draft_registration.cache_shape = draft_descriptor.cache_shape; + draft_registration.create_options = draft_descriptor.create_options; + draft_registration.model_identity = draft_descriptor.model_identity; + draft_registration.tp_rank = + options_.enable_mtp_draft_body_tp1() ? 0 : worker_id % target_tp_size; + draft_registration.tp_size = + options_.enable_mtp_draft_body_tp1() ? 1 : target_tp_size; + transfer->register_participant(std::move(draft_registration)); + CHECK(transfer->finalize_registration()); + composite_hierarchy_kv_cache_transfer_ = std::move(transfer); } uint32_t MTPWorkerImpl::transfer_kv_blocks( @@ -1024,6 +1140,13 @@ uint32_t MTPWorkerImpl::transfer_kv_blocks( CHECK(impl_ != nullptr); CHECK(draft_impl_ != nullptr); + if (composite_hierarchy_kv_cache_transfer_ != nullptr) { + return schedule_hierarchy_kv_cache_transfer( + composite_hierarchy_kv_cache_transfer_.get(), + batch_id, + block_transfer_info); + } + const uint32_t target_transferred = impl_->transfer_kv_blocks(batch_id, block_transfer_info); const uint32_t draft_transferred = @@ -1037,6 +1160,11 @@ uint32_t MTPWorkerImpl::transfer_kv_blocks( CHECK(impl_ != nullptr); CHECK(draft_impl_ != nullptr); + if (composite_hierarchy_kv_cache_transfer_ != nullptr) { + return composite_hierarchy_kv_cache_transfer_->transfer_kv_blocks( + batch_id, block_transfer_info); + } + const uint32_t target_transferred = impl_->transfer_kv_blocks(batch_id, block_transfer_info); const uint32_t draft_transferred = @@ -1044,6 +1172,27 @@ uint32_t MTPWorkerImpl::transfer_kv_blocks( return validate_paired_transfer_counts(target_transferred, draft_transferred); } +std::vector MTPWorkerImpl::prefetch_kv_blocks( + Slice& block_transfer_info) { + CHECK(impl_ != nullptr); + CHECK(draft_impl_ != nullptr); + if (composite_hierarchy_kv_cache_transfer_ != nullptr) { + return composite_hierarchy_kv_cache_transfer_->prefetch_kv_blocks( + block_transfer_info); + } + + std::vector target_hits = + impl_->prefetch_kv_blocks(block_transfer_info); + const std::vector draft_hits = + draft_impl_->prefetch_kv_blocks(block_transfer_info); + CHECK_EQ(target_hits.size(), draft_hits.size()); + for (size_t index = 0; index < target_hits.size(); ++index) { + target_hits[index] = + target_hits[index] == 1 && draft_hits[index] == 1 ? 1 : 0; + } + return target_hits; +} + #if defined(USE_NPU) || defined(USE_MLU) bool MTPWorkerImpl::allocate_kv_cache_with_transfer( const KVCacheShape& kv_cache_shape) { @@ -1121,7 +1270,11 @@ bool MTPWorkerImpl::allocate_kv_cache_with_transfer( torch::zeros({size}, torch::dtype(dtype_).device(device_))); } } - return target_allocated && draft_allocated; + const bool allocated = target_allocated && draft_allocated; + if (allocated) { + initialize_composite_hierarchy_kv_cache_transfer(); + } + return allocated; } #endif @@ -1140,7 +1293,7 @@ MTPWorkerImpl::update_input_by_last_step_output_for_schedule_overlap( void MTPWorkerImpl::prepare_work_before_execute(const ForwardInput& input, ForwardInput& processed_input) { - // Composite skips CP prepare; leaves run it in run_llm_no_sync_impl. + // Composite skips CP prepare; leaves run it in run_child_llm_no_sync. SpeculativeWorkerImpl::prepare_work_before_execute(input, processed_input); } @@ -1167,13 +1320,18 @@ std::optional MTPWorkerImpl::step_empty( if (!input.input_params.meta.batch_forward_type.is_decode()) { ForwardInput target_prepared; ForwardInput draft_prepared; - auto output = run_llm_no_sync_impl( - *impl_, input, *prepare_stream_, *compute_stream_, target_prepared); - auto draft_output = run_llm_no_sync_impl(*draft_impl_, - input, - *prepare_stream_, - *compute_stream_, - draft_prepared); + auto output = run_child_llm_no_sync(CacheParticipant::TARGET, + *impl_, + input, + *prepare_stream_, + *compute_stream_, + target_prepared); + auto draft_output = run_child_llm_no_sync(CacheParticipant::DRAFT, + *draft_impl_, + input, + *prepare_stream_, + *compute_stream_, + draft_prepared); if (draft_output.has_value()) { transfer_retained_inputs(*output, draft_output.value()); } @@ -1204,20 +1362,22 @@ std::optional MTPWorkerImpl::step_empty( draft_extend_prepared = std::move(pending_draft_context_.prepared_input); pending_draft_context_ = PendingDraftContext(); } else { - draft_outputs.emplace_back(run_llm_no_sync_impl(*draft_impl_, - new_input, - *prepare_stream_, - *compute_stream_, - draft_extend_prepared) + draft_outputs.emplace_back(run_child_llm_no_sync(CacheParticipant::DRAFT, + *draft_impl_, + new_input, + *prepare_stream_, + *compute_stream_, + draft_extend_prepared) .value()); } for (int32_t i = 1; i < options_.num_speculative_tokens(); ++i) { - draft_outputs.emplace_back(run_llm_no_sync_impl(*draft_impl_, - input, - *prepare_stream_, - *compute_stream_, - draft_step_prepared[i]) + draft_outputs.emplace_back(run_child_llm_no_sync(CacheParticipant::DRAFT, + *draft_impl_, + input, + *prepare_stream_, + *compute_stream_, + draft_step_prepared[i]) .value()); } @@ -1230,11 +1390,12 @@ std::optional MTPWorkerImpl::step_empty( new_input.input_params.parallel.raw_dp_global_token_nums) { token_num *= options_.num_speculative_tokens() + 1; } - ForwardOutput output = run_llm_no_sync_impl(*impl_, - new_input, - *prepare_stream_, - *compute_stream_, - target_prepared) + ForwardOutput output = run_child_llm_no_sync(CacheParticipant::TARGET, + *impl_, + new_input, + *prepare_stream_, + *compute_stream_, + target_prepared) .value(); for (ForwardOutput& draft_output : draft_outputs) { transfer_retained_inputs(output, draft_output); @@ -1267,10 +1428,13 @@ std::optional MTPWorkerImpl::step_prefill( ForwardInput draft_prepared; // run the target model to get first token and hidden states - ForwardOutput output = - run_llm_no_sync_impl( - *impl_, input, *prepare_stream_, *compute_stream_, target_prepared) - .value(); + ForwardOutput output = run_child_llm_no_sync(CacheParticipant::TARGET, + *impl_, + input, + *prepare_stream_, + *compute_stream_, + target_prepared) + .value(); COUNTER_ADD(speculative_execution_latency_seconds_target, timer.elapsed_seconds()); @@ -1308,11 +1472,12 @@ std::optional MTPWorkerImpl::step_prefill( } // generate kv cache for draft model timer.reset(); - ForwardOutput draft_output = run_llm_no_sync_impl(*draft_impl_, - prefill_input, - *prepare_stream_, - *compute_stream_, - draft_prepared) + ForwardOutput draft_output = run_child_llm_no_sync(CacheParticipant::DRAFT, + *draft_impl_, + prefill_input, + *prepare_stream_, + *compute_stream_, + draft_prepared) .value(); { c10::StreamGuard stream_guard = compute_stream_->set_stream_guard(); @@ -1672,11 +1837,12 @@ std::optional MTPWorkerImpl::step_decode( std::move(pending_draft_context_.prepared_input); pending_draft_context_ = PendingDraftContext(); } else { - draft_output_opt = run_llm_no_sync_impl(*draft_impl_, - current_draft_input, - *compute_stream_, - *compute_stream_, - draft_prepared[draft_idx]); + draft_output_opt = run_child_llm_no_sync(CacheParticipant::DRAFT, + *draft_impl_, + current_draft_input, + *compute_stream_, + *compute_stream_, + draft_prepared[draft_idx]); } if ((use_device_target_context || use_prelaunched_first_draft) && @@ -2170,11 +2336,12 @@ std::optional MTPWorkerImpl::run_validate( per_seq_val_tokens, json_scratch, *compute_stream_); - ForwardOutput target_output = run_llm_no_sync_impl(*impl_, - validate_input, - *compute_stream_, - *compute_stream_, - target_prepared) + ForwardOutput target_output = run_child_llm_no_sync(CacheParticipant::TARGET, + *impl_, + validate_input, + *compute_stream_, + *compute_stream_, + target_prepared) .value(); const double target_latency_ms = timer.elapsed_milliseconds(); COUNTER_ADD(speculative_execution_latency_seconds_target, @@ -2777,11 +2944,12 @@ void MTPWorkerImpl::submit_pending_first_draft( pending_draft_context_.dp_global_batch_generations = batch_identity_input.input_params.parallel.dp_global_batch_generations; pending_draft_context_.output = - run_llm_no_sync_impl(*draft_impl_, - draft_input, - *compute_stream_, - *compute_stream_, - pending_draft_context_.prepared_input); + run_child_llm_no_sync(CacheParticipant::DRAFT, + *draft_impl_, + draft_input, + *compute_stream_, + *compute_stream_, + pending_draft_context_.prepared_input); CHECK(pending_draft_context_.output.has_value()) << "failed to prelaunch next MTP first draft"; } diff --git a/xllm/core/runtime/mtp_worker_impl.h b/xllm/core/runtime/mtp_worker_impl.h index d7f2c7a637..498ff08d1e 100644 --- a/xllm/core/runtime/mtp_worker_impl.h +++ b/xllm/core/runtime/mtp_worker_impl.h @@ -47,6 +47,11 @@ class NpuJsonDraftTokenHandoff; // Eagle3WorkerImpl inherits from this class. class MTPWorkerImpl : public SpeculativeWorkerImpl { public: + enum class HierarchyCacheOwnershipMode : int8_t { + CHILD_OWNED = 0, + PARENT_COMPOSITE = 1, + }; + MTPWorkerImpl(const ParallelArgs& parallel_args, const torch::Device& device, const runtime::Options& options); @@ -64,7 +69,9 @@ class MTPWorkerImpl : public SpeculativeWorkerImpl { const runtime::Options& target_options, const runtime::Options& draft_options, bool enable_opt_validate_probs = false, - bool enable_adaptive_speculative_decode = false); + bool enable_adaptive_speculative_decode = false, + HierarchyCacheOwnershipMode hierarchy_cache_ownership_mode = + HierarchyCacheOwnershipMode::CHILD_OWNED); public: bool init_model(const std::string& model_weights_path, @@ -83,6 +90,9 @@ class MTPWorkerImpl : public SpeculativeWorkerImpl { uint64_t batch_id, Slice& block_transfer_info) override; + std::vector prefetch_kv_blocks( + Slice& block_transfer_info) override; + #if defined(USE_NPU) || defined(USE_MLU) bool allocate_kv_cache_with_transfer( const KVCacheShape& kv_cache_shape) override; @@ -272,6 +282,14 @@ class MTPWorkerImpl : public SpeculativeWorkerImpl { int32_t num_speculative_tokens, const std::vector* pruned_prefix_lengths = nullptr) const; bool adaptive_enabled() const; + void initialize_composite_hierarchy_kv_cache_transfer(); + std::optional run_child_llm_no_sync( + CacheParticipant participant, + LLMWorkerImpl& worker, + const ForwardInput& input, + Stream& prepare_stream, + Stream& compute_stream, + ForwardInput& processed_input); protected: // Draft model worker @@ -299,6 +317,8 @@ class MTPWorkerImpl : public SpeculativeWorkerImpl { // Whether validation directly uses selected-only draft_probs [B, S]. // If false, selected-only cache values are restored to dense [B, S, V]. bool enable_opt_validate_probs_ = false; + HierarchyCacheOwnershipMode hierarchy_cache_ownership_mode_ = + HierarchyCacheOwnershipMode::CHILD_OWNED; // adaptive_spec_controller_ now lives on SpeculativeWorkerImpl (base class). // Classified once when the corresponding models are loaded. Decode-path @@ -329,5 +349,8 @@ class MTPWorkerImpl : public SpeculativeWorkerImpl { #if defined(USE_NPU) || defined(USE_MLU) std::shared_ptr kv_cache_transfer_; #endif + + std::unique_ptr + composite_hierarchy_kv_cache_transfer_; }; } // namespace xllm diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 1638d54955..f115718dda 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -530,6 +530,14 @@ bool WorkerImpl::allocate_kv_cache_storage( // KV cache over a VMM-backed SleepableAllocator region (see kv_cache.cpp), so // sleep()/wake_up() can release / re-acquire it. allocate_kv_caches(kv_caches_, kv_cache_shape, create_options); + HierarchyKVCacheAllocationDescriptor allocation_descriptor; + allocation_descriptor.device_caches = &kv_caches_; + allocation_descriptor.cache_shape = kv_cache_shape; + allocation_descriptor.create_options = create_options; + allocation_descriptor.create_options.tensor_allocator(nullptr); + allocation_descriptor.model_identity = + options_.model_id() + "|" + args.model_type(); + hierarchy_kv_cache_allocation_descriptor_ = std::move(allocation_descriptor); init_hierarchy_kv_cache_transfer(kv_cache_shape, create_options); #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_DCU) @@ -2217,33 +2225,40 @@ uint32_t WorkerImpl::transfer_kv_blocks( const uint64_t batch_id, const std::vector& block_transfer_info) { if (hierarchy_kv_cache_transfer_ != nullptr) { - if (!block_transfer_info.empty() && - block_transfer_info.front().transfer_type == TransferType::D2H2G) { - // Schedule-overlap can deliver the D2H RPC before the preceding forward - // task has enqueued its kernels. Queue D2H on the same single-threaded - // worker executor so it runs after that forward; the hierarchy transfer - // then makes its copy stream wait on compute_stream_. H2D must remain on - // the RPC thread because it is registered before the matching forward. - auto result = std::make_shared>(); - std::future future = result->get_future(); - threadpool_.schedule( - [this, batch_id, block_transfer_info, result]() mutable { - try { - result->set_value( - hierarchy_kv_cache_transfer_->transfer_kv_blocks( - batch_id, block_transfer_info)); - } catch (...) { - result->set_exception(std::current_exception()); - } - }); - return future.get(); - } - return hierarchy_kv_cache_transfer_->transfer_kv_blocks( - batch_id, block_transfer_info); + return schedule_hierarchy_kv_cache_transfer( + hierarchy_kv_cache_transfer_.get(), batch_id, block_transfer_info); } return 0; } +uint32_t WorkerImpl::schedule_hierarchy_kv_cache_transfer( + HierarchyKVCacheTransfer* transfer, + uint64_t batch_id, + const std::vector& block_transfer_info) { + CHECK(transfer != nullptr); + if (!block_transfer_info.empty() && + block_transfer_info.front().transfer_type == TransferType::D2H2G) { + // Schedule-overlap can deliver the D2H RPC before the preceding forward + // task has enqueued its kernels. Queue D2H on the same single-threaded + // worker executor so it runs after that forward; the hierarchy transfer + // then makes its copy stream wait on the actual compute stream. H2D must + // remain on the RPC thread because it is registered before the forward. + auto result = std::make_shared>(); + std::future future = result->get_future(); + threadpool_.schedule( + [transfer, batch_id, block_transfer_info, result]() mutable { + try { + result->set_value( + transfer->transfer_kv_blocks(batch_id, block_transfer_info)); + } catch (...) { + result->set_exception(std::current_exception()); + } + }); + return future.get(); + } + return transfer->transfer_kv_blocks(batch_id, block_transfer_info); +} + uint32_t WorkerImpl::transfer_kv_blocks( const uint64_t batch_id, Slice& block_transfer_info) { diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h index a4e2443515..d73d89b3df 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include #include "common/types.h" #include "executor.h" @@ -55,6 +56,13 @@ class WorkerRendezvous; class WorkerImpl { public: + struct HierarchyKVCacheAllocationDescriptor final { + std::vector* device_caches = nullptr; + KVCacheShape cache_shape; + KVCacheCreateOptions create_options; + std::string model_identity; + }; + enum Status : int8_t { UNINITIALIZED = 0, LOADED, @@ -205,7 +213,7 @@ class WorkerImpl { Slice& block_transfer_info); // Run the model on the given input. async call - // the future returns a successfull status with no meaningful value + // the future returns a successful status with no meaningful value virtual folly::SemiFuture> step_async( const ForwardInput& inputs); @@ -231,6 +239,13 @@ class WorkerImpl { Status get_status() const { return status_; } + const HierarchyKVCacheAllocationDescriptor& + hierarchy_kv_cache_allocation_descriptor() const { + CHECK(hierarchy_kv_cache_allocation_descriptor_.has_value()) + << "KV cache allocation descriptor is not available."; + return hierarchy_kv_cache_allocation_descriptor_.value(); + } + // model context, includes model args, parallel args and date type etc. mutable ModelContext context_; @@ -256,6 +271,11 @@ class WorkerImpl { const KVCacheShape& kv_cache_shape, const KVCacheCreateOptions& kv_cache_create_options); + uint32_t schedule_hierarchy_kv_cache_transfer( + HierarchyKVCacheTransfer* transfer, + uint64_t batch_id, + const std::vector& block_transfer_info); + bool can_prepare_npu_graph_decode_input( const ModelInputParams& input_params) const; bool can_prepare_without_compute_stream_wait( @@ -363,6 +383,8 @@ class WorkerImpl { // kv caches std::vector kv_caches_; + std::optional + hierarchy_kv_cache_allocation_descriptor_; // causal LM model std::unique_ptr model_;