diff --git a/tests/core/framework/batch/batch_test.cpp b/tests/core/framework/batch/batch_test.cpp index 3fc5ea8a8d..551506211e 100644 --- a/tests/core/framework/batch/batch_test.cpp +++ b/tests/core/framework/batch/batch_test.cpp @@ -2056,7 +2056,7 @@ TEST(BatchTest, OverlapMTPReplacementKeepsCompositeKvBlocks) { .max_seqs_per_batch(max_seqs_per_batch) .manager_types({1, 0, 0}) .compress_ratios({0, 4, 128}); - CompositeBlockManager manager(options); + CompositeBlockManager manager(build_composite_leaves(options, /*dp_rank=*/0)); RequestSamplingParam sampling_param; StoppingChecker stopping_checker; @@ -2064,8 +2064,8 @@ TEST(BatchTest, OverlapMTPReplacementKeepsCompositeKvBlocks) { Sequence seq = make_overlap_sequence( {1, 10, 11}, /*seq_capacity=*/128, &sampling_param, &stopping_checker); - ASSERT_TRUE(manager.allocate_for_sequence(&seq, seq.num_prompt_tokens())); - ASSERT_EQ(seq.kv_state().num_kv_blocks(), 0u); + ASSERT_TRUE(manager.allocate_sequence(&seq, seq.num_prompt_tokens())); + ASSERT_EQ(seq.kv_state().num_blocks(BlockType::KV), 0u); ASSERT_GT(seq.kv_state().current_max_tokens_capacity(), 0u); seq.kv_state().incr_kv_cache_tokens_num(seq.num_prompt_tokens() - 1); diff --git a/tests/core/framework/kv_cache/kv_cache_test.cpp b/tests/core/framework/kv_cache/kv_cache_test.cpp index 900bcfa69b..2a31cb6817 100644 --- a/tests/core/framework/kv_cache/kv_cache_test.cpp +++ b/tests/core/framework/kv_cache/kv_cache_test.cpp @@ -22,9 +22,13 @@ limitations under the License. #include #include "core/framework/config/kv_cache_config.h" +#include "framework/block/block.h" +#include "framework/kv_cache/deepseek_v4_cache_policy.h" #include "framework/kv_cache/deepseek_v4_kv_cache_impl.h" +#include "framework/kv_cache/kv_cache_utils.h" #include "kv_cache_estimation.h" #include "kv_cache_shape.h" +#include "platform/device.h" #include "worker.pb.h" namespace xllm { @@ -129,6 +133,17 @@ TEST(Dsv4StateCacheTest, MissingPackedFallsBackWithoutSwappingSplit) { EXPECT_TRUE(torch::equal(state.score(), score)); } +// Host prefix-cache allocation registers page-aligned host memory with the NPU +// via aclrtHostRegister, which requires a live device context. Set one up once. +class HostKVCacheTest : public ::testing::Test { + protected: + void SetUp() override { + Device device(/*device_index=*/0); + device.set_device(); + device.init_device_context(); + } +}; + TEST(KVCacheTest, DeepSeekV4FourDimCachesUseDeviceLayout) { constexpr int64_t kSwaCount = 10; constexpr int64_t kC4Count = 32; @@ -564,4 +579,111 @@ TEST(KVCacheTest, MluIndexerAutoUsesDefaultCacheShapeWithoutScale) { } #endif +TEST_F(HostKVCacheTest, HostKVCacheNormalLayoutAddsLayerDim) { + constexpr int64_t kNumBlocks = 16; + constexpr int64_t kBlockSize = 128; + constexpr int64_t kHeadDim = 64; + constexpr int64_t kNumHeads = 4; + constexpr int64_t kLayerCount = 5; + constexpr double kHostFactor = 2.0; + + KVCacheCapacity capacity; + capacity.n_blocks(kNumBlocks).block_size(kBlockSize); + + ModelArgs model_args; + model_args.model_type("qwen"); + model_args.n_kv_heads(kNumHeads); + model_args.head_dim(kHeadDim); + KVCacheShape shape(capacity, model_args, /*world_size=*/1); + ASSERT_TRUE(shape.has_key_cache_shape()); + + KVCacheCreateOptions options; + options.device(torch::Device(torch::kCPU)) + .dtype(torch::kFloat32) + .num_layers(kLayerCount) + .model_type("qwen") + .host_blocks_factor(kHostFactor); + + KVCache host_cache(shape, options, BlockType::KV, kLayerCount); + + const BlockTypeTensorMap tensors = + host_cache.get_block_type_tensors(BlockType::KV); + ASSERT_TRUE(tensors.count(KVCacheTensorRole::KEY) > 0); + + const std::vector base_key_shape = shape.key_cache_shape(); + const torch::Tensor& host_key = tensors.at(KVCacheTensorRole::KEY); + EXPECT_TRUE(host_key.is_contiguous()); + EXPECT_EQ(host_key.device().type(), torch::kCPU); + // host shape == [scaled_blocks, layer_count, ...per_block_dims] + ASSERT_EQ(host_key.dim(), static_cast(base_key_shape.size()) + 1); + EXPECT_EQ(host_key.size(0), + scale_host_block_count(base_key_shape[0], kHostFactor)); + EXPECT_EQ(host_key.size(1), kLayerCount); + for (size_t i = 1; i < base_key_shape.size(); ++i) { + EXPECT_EQ(host_key.size(static_cast(i) + 1), base_key_shape[i]); + } +} + +TEST_F(HostKVCacheTest, HostKVCacheDeepSeekV4PerBlockType) { + constexpr int64_t kSwaCount = 10; + constexpr int64_t kC4Count = 32; + constexpr int64_t kC128Count = 4; + constexpr int64_t kBlockSize = 128; + constexpr int64_t kHeadDim = 16; + constexpr int64_t kIndexHeadDim = 8; + constexpr double kHostFactor = 3.0; + + KVCacheCapacity capacity; + capacity.block_size(kBlockSize) + .swa_count(kSwaCount) + .c4_count(kC4Count) + .c128_count(kC128Count); + + ModelArgs model_args; + model_args.model_type("deepseek_v4"); + KVCacheShape shape(capacity, model_args, /*world_size=*/1); + + KVCacheCreateOptions options; + options.device(torch::Device(torch::kCPU)) + .dtype(torch::kFloat32) + .num_layers(3) + .model_type("deepseek_v4") + .block_size(kBlockSize) + .head_dim(kHeadDim) + .index_head_dim(kIndexHeadDim) + .window_size(/*window_size=*/512) + .compress_ratios({1, 4, 128}) + .host_blocks_factor(kHostFactor); + + // SWA host cache: 1 layer in this 3-layer config (compress_ratio == 1). + KVCache swa_host(shape, options, BlockType::SWA, /*layer_count=*/1); + const BlockTypeTensorMap swa_tensors = + swa_host.get_block_type_tensors(BlockType::SWA); + ASSERT_TRUE(swa_tensors.count(KVCacheTensorRole::SWA) > 0); + const torch::Tensor& swa = swa_tensors.at(KVCacheTensorRole::SWA); + EXPECT_TRUE(swa.is_contiguous()); + EXPECT_EQ(swa.size(0), scale_host_block_count(kSwaCount, kHostFactor)); + EXPECT_EQ(swa.size(1), 1); + + // C4 host cache: key + index, index uses the DSV4 index dtype. + KVCache c4_host(shape, options, BlockType::C4, /*layer_count=*/1); + const BlockTypeTensorMap c4_tensors = + c4_host.get_block_type_tensors(BlockType::C4); + ASSERT_TRUE(c4_tensors.count(KVCacheTensorRole::KEY) > 0); + ASSERT_TRUE(c4_tensors.count(KVCacheTensorRole::INDEX) > 0); + EXPECT_EQ(c4_tensors.at(KVCacheTensorRole::KEY).size(0), + scale_host_block_count(kC4Count, kHostFactor)); + EXPECT_EQ(c4_tensors.at(KVCacheTensorRole::INDEX).scalar_type(), + get_dsv4_cache_policy(options.dtype()).index_dtype); + + // C128 host cache: key only (no index). + KVCache c128_host(shape, options, BlockType::C128, /*layer_count=*/1); + const BlockTypeTensorMap c128_tensors = + c128_host.get_block_type_tensors(BlockType::C128); + ASSERT_TRUE(c128_tensors.count(KVCacheTensorRole::KEY) > 0); + EXPECT_TRUE(c128_tensors.count(KVCacheTensorRole::INDEX) == 0); + EXPECT_EQ(c128_tensors.at(KVCacheTensorRole::KEY).size(0), + scale_host_block_count(kC128Count, kHostFactor)); +} + } // namespace xllm diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index 8f27c77404..2649601cc6 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -41,8 +41,7 @@ limitations under the License. #include "core/framework/config/service_config.h" #include "core/platform/platform.h" #include "framework/block/block_utils.h" -// hierarchy temporarily disabled during the block-manager refactor -// #include "framework/block/hierarchy_block_manager_pool.h" +#include "framework/block/hierarchy_block_manager_pool.h" #include "framework/kv_cache/kv_cache_estimation.h" #include "framework/kv_cache/kv_cache_shape.h" #include "framework/model/model_args.h" @@ -588,19 +587,23 @@ bool LLMEngine::allocate_kv_cache(const KVCacheCapacity& kv_cache_cap) { kv_cache_cap.swa_count(), std::numeric_limits::max()))); } - if (options_.host_blocks_factor() > 1.0 || options_.enable_kvcache_store()) { - // hierarchy temporarily disabled during the block-manager refactor. - // host-offload / kvcache-store routes the device + host dual - // KVCacheState through HierarchyBlockManagerPool, which is parked while - // the composite block-manager refactor lands in smaller pieces. Until - // then this path fails loudly rather than silently degrading to a - // device-only pool. - LOG(FATAL) << "host-offload / kvcache-store is temporarily disabled during " - "the block-manager refactor (hierarchy rebuild in progress). " - "Please disable --host_blocks_factor and " - "--enable_kvcache_store for now."; + if (options_.host_blocks_factor() > 1.0) { + // Host prefix-cache offload routes device/host blocks through a single flat + // BlockType::KV host pool. DeepSeek-V4 has no KV block group (its device + // caches are SWA/C4/C128), so collect_offload_pairs would find no KV blocks + // and silently offload nothing while still allocating the pinned host + // cache. Fail loud until per-block-type host offload is implemented. + CHECK(!util::is_deepseek_v4_model_type(args_.model_type())) + << "host_blocks_factor > 1 (host prefix-cache offload) does not " + "support " + "DeepSeek-V4 yet: its SWA/C4/C128 block groups have no KV pool to " + "offload. Disable --host_blocks_factor for DeepSeek-V4 models."; + options.enable_host_offload(true); + kv_cache_manager_ = + std::make_unique(options, this, dp_size_); + } else { + kv_cache_manager_ = std::make_unique(options, dp_size_); } - kv_cache_manager_ = std::make_unique(options, dp_size_); // init kv cache for each worker in parallel std::vector> futures; diff --git a/xllm/core/framework/block/CMakeLists.txt b/xllm/core/framework/block/CMakeLists.txt index a9a55aa183..0ae7cf3dd5 100644 --- a/xllm/core/framework/block/CMakeLists.txt +++ b/xllm/core/framework/block/CMakeLists.txt @@ -23,8 +23,7 @@ cc_library( linear_state_block_manager.h block_manager_impl.h concurrent_block_manager_impl.h - # hierarchy temporarily disabled during the block-manager refactor - # hierarchy_block_manager_pool.h + hierarchy_block_manager_pool.h sliding_window_block_manager.h composite_block_manager.h SRCS @@ -34,8 +33,7 @@ cc_library( linear_state_block_manager.cpp concurrent_block_manager_impl.cpp block_manager_impl.cpp - # hierarchy temporarily disabled during the block-manager refactor - # hierarchy_block_manager_pool.cpp + hierarchy_block_manager_pool.cpp sliding_window_block_manager.cpp composite_block_manager.cpp DEPS diff --git a/xllm/core/framework/block/block_manager.h b/xllm/core/framework/block/block_manager.h index e21449842f..e73deaaa8a 100644 --- a/xllm/core/framework/block/block_manager.h +++ b/xllm/core/framework/block/block_manager.h @@ -67,6 +67,11 @@ class BlockManager { // leaves are wrapped in ConcurrentBlockManagerImpl when this (or // enable_disagg_pd) is set. PROPERTY(bool, enable_kvcache_store) = false; + // Whether host prefix-cache offload (host_blocks_factor > 1) is enabled. + // The D2H offload-completion callback frees device/host blocks from a folly + // executor thread, so leaves are wrapped in ConcurrentBlockManagerImpl when + // this is set to make those mutations thread-safe against the scheduler. + PROPERTY(bool, enable_host_offload) = false; // xtensor (VMM) KV leaf parameters. When enable_xtensor is set, the KV leaf // is an XTensorBlockManagerImpl instead of a flat BlockManagerImpl; these // carry the construction args the spec builder needs. diff --git a/xllm/core/framework/block/block_manager_pool.cpp b/xllm/core/framework/block/block_manager_pool.cpp index 411991de81..717ee8e228 100644 --- a/xllm/core/framework/block/block_manager_pool.cpp +++ b/xllm/core/framework/block/block_manager_pool.cpp @@ -49,6 +49,7 @@ BlockManagerPool::BlockManagerPool(const Options& options, int32_t dp_size) .enable_prefix_cache(options_.enable_prefix_cache()) .enable_disagg_pd(options_.enable_disagg_pd()) .enable_kvcache_store(options_.enable_kvcache_store()) + .enable_host_offload(options_.enable_host_offload()) .sliding_window_size(options_.sliding_window_size()) .swa_blocks_per_seq(options_.swa_blocks_per_seq()) .max_tokens_per_batch(options_.max_tokens_per_batch()) @@ -78,13 +79,15 @@ BlockManagerPool::BlockManagerPool(const Options& options, int32_t dp_size) auto leaves = build_composite_leaves(block_options, /*dp_rank=*/i); // SINGLE leaf needs the same concurrency wrapper as the other leaves when // sequence-level entry points run off the scheduler thread (disagg PD / - // kvcache store prefill threadpools call try_allocate concurrently). + // kvcache store prefill threadpools call try_allocate concurrently, and the + // host-offload D2H callback frees blocks off-thread). std::unique_ptr single_leaf = std::make_unique( /*num_blocks=*/num_single_blocks, /*resource_name=*/"single block", /*exhaustion_message=*/"No more single-block ids available"); - if (options_.enable_disagg_pd() || options_.enable_kvcache_store()) { + if (options_.enable_disagg_pd() || options_.enable_kvcache_store() || + options_.enable_host_offload()) { single_leaf = std::make_unique(std::move(single_leaf)); } diff --git a/xllm/core/framework/block/block_manager_pool.h b/xllm/core/framework/block/block_manager_pool.h index 6b18376f85..b7352ef49e 100644 --- a/xllm/core/framework/block/block_manager_pool.h +++ b/xllm/core/framework/block/block_manager_pool.h @@ -46,6 +46,10 @@ class BlockManagerPool : public KVCacheManager { PROPERTY(bool, enable_prefix_cache) = true; PROPERTY(bool, enable_disagg_pd) = false; PROPERTY(bool, enable_kvcache_store) = false; + // Host prefix-cache offload (host_blocks_factor > 1). Wraps composite + // leaves in ConcurrentBlockManagerImpl so the async D2H offload callback + // can free blocks off-thread safely. + PROPERTY(bool, enable_host_offload) = false; PROPERTY(bool, enable_xtensor) = false; PROPERTY(int64_t, num_layers) = 0; // Required when enable_xtensor is true PROPERTY(int64_t, slot_size) = 0; // Memory size per slot (for xtensor) diff --git a/xllm/core/framework/block/composite_block_manager.cpp b/xllm/core/framework/block/composite_block_manager.cpp index a1481f2e06..384f9b717c 100644 --- a/xllm/core/framework/block/composite_block_manager.cpp +++ b/xllm/core/framework/block/composite_block_manager.cpp @@ -40,11 +40,13 @@ uint32_t ceil_div(uint32_t numerator, uint32_t denominator) { } // Wrap a leaf in a concurrency adapter when sequence-level entry points may run -// off the scheduler thread (disagg PD / kvcache store prefill threadpools). +// off the scheduler thread (disagg PD / kvcache store prefill threadpools, or +// the host-offload D2H completion callback). std::unique_ptr maybe_concurrent( std::unique_ptr leaf, const BlockManager::Options& options) { - if (options.enable_disagg_pd() || options.enable_kvcache_store()) { + if (options.enable_disagg_pd() || options.enable_kvcache_store() || + options.enable_host_offload()) { return std::make_unique(std::move(leaf)); } return leaf; @@ -221,6 +223,42 @@ bool CompositeBlockManager::allocate_sequence(Sequence* seq, } } + // Post-condition (grow-or-fail): every cache-bearing leaf must now hold + // enough blocks to cover num_tokens. A leaf that returns an empty vector + // means "no growth needed", so we must verify that was true. SINGLE / LINEAR + // leaves are per-sequence resource slots (one block per sequence), not + // token-cache, and are exempt. Without this check, a leaf that mistakenly + // returned an empty vector under pool pressure (instead of nullopt) would + // let the pool report admission success while the device is under-provisioned + // for num_tokens -- batch_input_builder.cpp:589 CHECK then fires downstream + // with `current_max_tokens_capacity() >= seq_len`. Rolling back here lets the + // scheduler defer the sequence to the next tick, at which point pool state + // may have recovered. + for (const auto& [type, entry] : leaves_) { + if (type == BlockType::SINGLE || type == BlockType::LINEAR) { + continue; + } + const size_t leaf_block_size = entry.leaf->block_size(); + if (leaf_block_size == 0) { + continue; + } + const size_t needed = (num_tokens + leaf_block_size - 1) / leaf_block_size; + size_t staged_for_type = 0; + for (const auto& [staged_type, staged_blocks] : staged) { + if (staged_type == type) { + staged_for_type = staged_blocks.size(); + break; + } + } + const size_t total = seq->kv_state().num_blocks(type) + staged_for_type; + if (total < needed) { + for (auto& [staged_type, staged_blocks] : staged) { + leaf_of(staged_type)->deallocate(staged_blocks); + } + return false; + } + } + // Commit: append staged blocks to the sequence under each leaf's block type. for (auto& [type, blocks] : staged) { kv_state.add_blocks(type, blocks); diff --git a/xllm/core/framework/block/hierarchy_block_manager_pool.cpp b/xllm/core/framework/block/hierarchy_block_manager_pool.cpp index 2855dd1626..7ea3fcdfe6 100644 --- a/xllm/core/framework/block/hierarchy_block_manager_pool.cpp +++ b/xllm/core/framework/block/hierarchy_block_manager_pool.cpp @@ -18,6 +18,7 @@ limitations under the License. #include #include "block_manager_impl.h" +#include "composite_block_manager.h" #include "concurrent_block_manager_impl.h" namespace xllm { @@ -31,21 +32,23 @@ HierarchyBlockManagerPool::HierarchyBlockManagerPool( host_block_managers_.reserve(dp_size); BlockManager::Options host_options; - host_options.num_blocks(options_.num_blocks()) + host_options.num_blocks(options_.host_num_blocks()) .block_size(options_.block_size()) .enable_prefix_cache(options_.enable_prefix_cache()) .enable_disagg_pd(options_.enable_disagg_pd()) - .num_blocks(options_.host_num_blocks()) .hasher_type(options_.hasher_type()); for (int32_t i = 0; i < dp_size; ++i) { - if (options.enable_disagg_pd() || options_.enable_kvcache_store()) { - host_block_managers_.emplace_back( - std::make_unique(host_options)); - } else { - host_block_managers_.emplace_back( - std::make_unique(host_options)); + std::unique_ptr leaf = + std::make_unique(host_options); + // The D2H offload callback (transfer_blocks) frees host blocks from a folly + // executor thread while the scheduler thread also allocates/deallocates on + // this leaf, so it must be lock-guarded. Host offload is always on here. + if (options.enable_disagg_pd() || options_.enable_kvcache_store() || + options_.enable_host_offload()) { + leaf = std::make_unique(std::move(leaf)); } + host_block_managers_.emplace_back(std::move(leaf)); } load_block_transfer_infos_.resize(host_block_managers_.size()); @@ -54,71 +57,110 @@ HierarchyBlockManagerPool::HierarchyBlockManagerPool( void HierarchyBlockManagerPool::deallocate(Sequence* sequence) { DCHECK(sequence != nullptr); - // add blocks to the prefix cache int32_t dp_rank = BlockManagerPool::get_dp_rank(sequence); + // Publish device KV blocks into the device prefix cache first, so that + // offload-eligible blocks reach ref_count == 2 (held only by the sequence + // state and the prefix-cache node). deallocate_for_sequence below re-runs + // this cache step; it is idempotent (hashes recomputed from token ids find + // the nodes inserted here), so moved-out blocks do not corrupt the cache. BlockManagerPool::cache(sequence); - auto* blocks = sequence->kv_state().mutable_kv_blocks(); - auto* host_blocks = sequence->host_kv_state().mutable_kv_blocks(); + collect_offload_pairs(sequence, dp_rank); - if (host_blocks->size() >= blocks->size()) { - host_block_managers_[dp_rank]->deallocate( - sequence->host_kv_state().kv_blocks()); - block_managers_[dp_rank]->deallocate(sequence->kv_state().kv_blocks()); - deallocate_single_block(sequence, dp_rank); - sequence->reset(); + // Release the host blocks still held by the sequence. Blocks moved into the + // offload queue are now invalid in this vector and are skipped by deallocate; + // their host ids stay reserved (held by the queue) until the D2H copy + // completes and the offload callback caches + frees them. + auto host_blocks = sequence->host_kv_state().blocks(BlockType::KV); + if (!host_blocks.empty()) { + host_block_managers_[dp_rank]->deallocate(host_blocks); + } + + // Release device blocks via the composite (includes prefix cache flush). + // Offloaded device blocks were moved out above, so the KV leaf skips them; + // the offload callback releases them once the copy is done. + auto* composite = + static_cast(block_managers_[dp_rank].get()); + composite->deallocate_for_sequence(sequence); + sequence->reset(); +} + +void HierarchyBlockManagerPool::collect_offload_pairs(Sequence* sequence, + int32_t dp_rank) { + if (!options_.enable_prefix_cache()) { return; } - size_t cached_host_block_num = - sequence->host_kv_state().kv_cache_tokens_num() / options_.block_size(); - size_t cached_device_block_num = - sequence->kv_state().kv_cache_tokens_num() / options_.block_size(); + std::vector* device_blocks = + sequence->kv_state().mutable_blocks(BlockType::KV); + std::vector* host_blocks = + sequence->host_kv_state().mutable_blocks(BlockType::KV); + if (device_blocks == nullptr || device_blocks->empty()) { + return; + } - size_t needed_block_num = cached_device_block_num > host_blocks->size() - ? cached_device_block_num - host_blocks->size() - : 0; + const size_t block_size = options_.block_size(); + const size_t cached_host_block_num = + sequence->host_kv_state().kv_cache_tokens_num() / block_size; + const size_t cached_device_block_num = + sequence->kv_state().kv_cache_tokens_num() / block_size; + + const size_t host_block_num = + host_blocks == nullptr ? 0 : host_blocks->size(); + // Host already holds at least as many blocks as the device computed: nothing + // new to offload. + if (host_block_num >= device_blocks->size()) { + return; + } - // allocate additional host blocks for copy + // Allocate the host blocks needed to receive the device blocks that have no + // host counterpart yet. + const size_t needed_block_num = cached_device_block_num > host_block_num + ? cached_device_block_num - host_block_num + : 0; if (needed_block_num != 0) { - sequence->host_kv_state().add_kv_blocks( - host_block_managers_[dp_rank]->allocate(needed_block_num)); + std::vector new_host_blocks = + host_block_managers_[dp_rank]->allocate(needed_block_num); + if (new_host_blocks.size() != needed_block_num) { + // Host pool exhausted; skip offload this round rather than partially + // copy. + return; + } + sequence->add_host_blocks(BlockType::KV, new_host_blocks); + host_blocks = sequence->host_kv_state().mutable_blocks(BlockType::KV); + } + if (host_blocks == nullptr) { + return; } // Only offload blocks that are fully computed on device. In-batch prefix // cache insertion may register blocks before they are computed, so bound the - // offload range by cached_device_block_num to avoid copying uncomputed data - // to host/store. - const size_t offload_end_block_num = - std::min({cached_device_block_num, host_blocks->size(), blocks->size()}); + // offload range by cached_device_block_num to avoid copying uncomputed data. + const size_t offload_end_block_num = std::min( + {cached_device_block_num, host_blocks->size(), device_blocks->size()}); for (size_t i = cached_host_block_num; i < offload_end_block_num; i++) { - if (blocks->at(i).ref_count() != 2) { + // ref_count == 2 means the block is held only by this sequence and the + // prefix-cache node, i.e. it is uniquely owned and safe to offload. Beam + // forks (shared blocks) have ref_count > 2 and are skipped. + if (device_blocks->at(i).ref_count() != 2) { continue; } - - host_blocks->at(i).set_hash_value(blocks->at(i).get_immutable_hash_value()); + host_blocks->at(i).set_hash_value( + device_blocks->at(i).get_immutable_hash_value()); auto block_pair = std::make_shared( - std::move(blocks->at(i)), std::move(host_blocks->at(i))); + std::move(device_blocks->at(i)), std::move(host_blocks->at(i))); offload_block_pair_queues_[dp_rank].enqueue(std::move(block_pair)); } - - host_block_managers_[dp_rank]->deallocate( - sequence->host_kv_state().kv_blocks()); - - block_managers_[dp_rank]->deallocate(sequence->kv_state().kv_blocks()); - deallocate_single_block(sequence, dp_rank); - sequence->reset(); } bool HierarchyBlockManagerPool::allocate(Sequence* sequence, size_t num_tokens, size_t max_copy_in_blocks_num) { - // set needed_kv_cache_tokens_num to overlap computation and data transfer if (!BlockManagerPool::allocate(sequence, num_tokens)) { return false; } - if (sequence->host_kv_state().num_kv_blocks() == 0 && + if (sequence->host_kv_state().num_blocks(BlockType::KV) == 0 && sequence->stage() != SequenceStage::DECODE) { allocate_host_shared(sequence); } @@ -132,34 +174,71 @@ bool HierarchyBlockManagerPool::allocate(Sequence* sequence, hbm_cache_token_num / options_.block_size() : 0; if (max_copy_in_blocks_num > max_can_copy_blocks_num) { - // not enough blocks to copy, return false LOG(ERROR) << "Not enough host blocks to copy, max_copy_in_blocks_num: " << max_copy_in_blocks_num - << ", max_copy_blocks_num: " << max_copy_in_blocks_num; + << ", max_copy_blocks_num: " << max_can_copy_blocks_num; max_copy_in_blocks_num = max_can_copy_blocks_num; } - auto hbm_blocks = sequence->kv_state().kv_blocks(); - auto host_blocks = sequence->host_kv_state().kv_blocks(); - for (int i = hbm_cache_token_num / options_.block_size(); - i < - max_copy_in_blocks_num + (hbm_cache_token_num / options_.block_size()); + auto hbm_blocks = sequence->kv_state().blocks(BlockType::KV); + auto host_blocks = sequence->host_kv_state().blocks(BlockType::KV); + // H2D copies host block i -> device block i, so i must index both vectors. + // The host prefix match (host_cache_token_num) is computed over the full + // prompt and can exceed the device blocks allocated for this (possibly + // chunked) num_tokens, so clamp the copy range to the blocks that actually + // exist on both sides to avoid out-of-bounds reads. + const size_t hbm_block_begin = hbm_cache_token_num / options_.block_size(); + const size_t copy_block_limit = + std::min(hbm_blocks.size(), host_blocks.size()); + if (hbm_block_begin + max_copy_in_blocks_num > copy_block_limit) { + max_copy_in_blocks_num = copy_block_limit > hbm_block_begin + ? copy_block_limit - hbm_block_begin + : 0; + } + for (size_t i = hbm_block_begin; i < max_copy_in_blocks_num + hbm_block_begin; i++) { - load_block_transfer_infos_[dp_rank].emplace_back( - BlockTransferInfo(host_blocks[i].id(), - hbm_blocks[i].id(), - host_blocks[i].get_immutable_hash_value(), - TransferType::H2D)); + const Block& hb = host_blocks[i]; + const Block& db = hbm_blocks[i]; + load_block_transfer_infos_[dp_rank].emplace_back(BlockTransferInfo( + hb.id(), db.id(), hb.get_immutable_hash_value(), TransferType::H2D)); } size_t target_hbm_cache_token_num = max_copy_in_blocks_num == 0 ? hbm_cache_token_num - : (max_copy_in_blocks_num + - (hbm_cache_token_num / options_.block_size())) * - options_.block_size(); - - sequence->kv_state().incr_kv_cache_tokens_num(target_hbm_cache_token_num - - hbm_cache_token_num); + : (max_copy_in_blocks_num + hbm_block_begin) * options_.block_size(); + + // Clamp to num_tokens (this tick's scheduler-committed budget): scheduler + // computed allowed_max_tokens = num_tokens - pre_kv_cache_tokens_num_ before + // this call, and batch_input_builder later computes + // seq_len = min(n_tokens_full - n_kv_cache, allowed_max_tokens) + + // n_kv_cache + // If we bump n_kv_cache past num_tokens via H2D restore, seq_len exceeds + // num_tokens = capacity (post grow-or-fail), tripping the batch CHECK. Cap + // the H2D advance at num_tokens; host prefix cache blocks past this remain + // held and can be restored on a later tick after the scheduler admits a + // larger max_handle_num_tokens. Also clamp to declared capacity as a + // defensive backstop for the DECODE overload where allocate_host_shared + // isn't invoked; incr_kv_cache_tokens_num_up_to CHECK-fails if the counter + // drifted past capacity on an earlier tick. + target_hbm_cache_token_num = + std::min({target_hbm_cache_token_num, + num_tokens, + sequence->kv_state().current_max_tokens_capacity()}); + sequence->kv_state().incr_kv_cache_tokens_num_up_to( + target_hbm_cache_token_num); + + // Symmetric clamp on host_kv_state: Sequence::kv_cache_tokens_num() at + // sequence.h:175 returns max(kv_state, host_kv_state). Schedulers read that + // max view (chunked_prefill_scheduler.cpp:879, continuous_scheduler.cpp:497, + // etc.) to decide q_seq_len for the NEXT tick. Cap host to num_tokens too + // so next-tick admission math sees a value the device can back at that + // point. The host prefix cache blocks themselves stay held so a later tick + // can still restore them once the device grows. + const size_t host_cap = + std::min(num_tokens, sequence->kv_state().current_max_tokens_capacity()); + if (sequence->host_kv_state().kv_cache_tokens_num() > host_cap) { + sequence->host_kv_state().set_kv_cache_tokens_num(host_cap); + } return true; } @@ -170,7 +249,7 @@ bool HierarchyBlockManagerPool::allocate(Sequence* sequence, return false; } - if (sequence->host_kv_state().num_kv_blocks() == 0 && + if (sequence->host_kv_state().num_blocks(BlockType::KV) == 0 && sequence->stage() != SequenceStage::DECODE) { allocate_host_shared(sequence); } @@ -179,27 +258,61 @@ bool HierarchyBlockManagerPool::allocate(Sequence* sequence, size_t hbm_cache_token_num = sequence->kv_state().kv_cache_tokens_num(); size_t host_cache_token_num = sequence->host_kv_state().kv_cache_tokens_num(); if (hbm_cache_token_num < host_cache_token_num) { - auto hbm_blocks = sequence->kv_state().kv_blocks(); - auto host_blocks = sequence->host_kv_state().kv_blocks(); - - for (int i = hbm_cache_token_num / options_.block_size(); - i < host_cache_token_num / options_.block_size(); - i++) { - load_block_transfer_infos_[dp_rank].emplace_back( - BlockTransferInfo(host_blocks[i].id(), - hbm_blocks[i].id(), - host_blocks[i].get_immutable_hash_value(), - TransferType::H2D)); + auto hbm_blocks = sequence->kv_state().blocks(BlockType::KV); + auto host_blocks = sequence->host_kv_state().blocks(BlockType::KV); + + // H2D copies host block i -> device block i. host_cache_token_num is the + // host prefix match over the full prompt and can exceed the device blocks + // allocated for this (chunked) num_tokens, so clamp to the blocks present + // on both sides to avoid out-of-bounds reads on hbm_blocks. + const size_t copy_block_end = + std::min({host_cache_token_num / options_.block_size(), + hbm_blocks.size(), + host_blocks.size()}); + const size_t hbm_block_begin = hbm_cache_token_num / options_.block_size(); + for (size_t i = hbm_block_begin; i < copy_block_end; i++) { + const Block& hb = host_blocks[i]; + const Block& db = hbm_blocks[i]; + load_block_transfer_infos_[dp_rank].emplace_back(BlockTransferInfo( + hb.id(), db.id(), hb.get_immutable_hash_value(), TransferType::H2D)); } - sequence->kv_state().incr_kv_cache_tokens_num(host_cache_token_num - - hbm_cache_token_num); + size_t target_hbm_cache_token_num = + copy_block_end > hbm_block_begin + ? copy_block_end * options_.block_size() + : hbm_cache_token_num; + // Same clamp shape as the 3-arg overload: cap at num_tokens (this tick's + // scheduler-committed budget) so H2D restore never pushes kv_state past + // what batch_input_builder can back given the pre-allocate allowed_max_ + // tokens. Also clamp to declared capacity as a defensive backstop; the + // grow-or-fail post-condition in CompositeBlockManager makes capacity >= + // num_tokens the normal case, but the min() keeps this correct even if a + // future refactor changes that. + target_hbm_cache_token_num = + std::min({target_hbm_cache_token_num, + num_tokens, + sequence->kv_state().current_max_tokens_capacity()}); + sequence->kv_state().incr_kv_cache_tokens_num_up_to( + target_hbm_cache_token_num); + } + // Symmetric clamp on host_kv_state (always, regardless of the H2D branch + // above): allocate_host_shared() at line 247 could have advanced + // host_kv_state.kv_cache_tokens_num_ during this call, past what the + // scheduler admitted for this tick. Sequence::kv_cache_tokens_num() + // (sequence.h:175) returns max(kv_state, host_kv_state), so a stale host + // value leaks a chunk budget that batch_input_builder can't back. Cap host + // to min(num_tokens, capacity); the underlying host prefix cache blocks stay + // held, ready to be restored once the device grows on a later tick. + const size_t host_cap = + std::min(num_tokens, sequence->kv_state().current_max_tokens_capacity()); + if (sequence->host_kv_state().kv_cache_tokens_num() > host_cap) { + sequence->host_kv_state().set_kv_cache_tokens_num(host_cap); } return true; } void HierarchyBlockManagerPool::allocate_shared(Sequence* sequence) { BlockManagerPool::allocate_shared(sequence); - if (sequence->host_kv_state().num_kv_blocks() == 0 && + if (sequence->host_kv_state().num_blocks(BlockType::KV) == 0 && sequence->stage() != SequenceStage::DECODE) { allocate_host_shared(sequence); } @@ -210,7 +323,7 @@ void HierarchyBlockManagerPool::allocate_host_shared(Sequence* sequence) { int32_t dp_rank = BlockManagerPool::get_dp_rank(sequence); std::vector shared_blocks = host_block_managers_[dp_rank]->allocate_shared(sequence->tokens()); - sequence->add_shared_host_kv_blocks(std::move(shared_blocks)); + sequence->add_shared_host_blocks(BlockType::KV, std::move(shared_blocks)); } } @@ -227,11 +340,11 @@ void HierarchyBlockManagerPool::prefetch_from_storage( std::vector shared_blocks = host_block_managers_[dp_rank]->allocate_shared( prefill_sequence->tokens()); - prefill_sequence->add_shared_host_kv_blocks(std::move(shared_blocks)); + prefill_sequence->add_shared_host_blocks(BlockType::KV, + std::move(shared_blocks)); - // round down to the nearest block number size_t shared_blocks_num = - prefill_sequence->host_kv_state().shared_kv_blocks_num(); + prefill_sequence->host_kv_state().shared_blocks_num(BlockType::KV); const size_t num_additional_blocks = (prefill_sequence->num_tokens() + options_.block_size() - 1) / options_.block_size() - @@ -245,21 +358,22 @@ void HierarchyBlockManagerPool::prefetch_from_storage( if (host_blocks.size() != num_additional_blocks) { continue; } - prefill_sequence->host_kv_state().add_kv_blocks(host_blocks); + prefill_sequence->add_host_blocks(BlockType::KV, host_blocks); PrefixCache::compute_hash_keys( prefill_sequence->tokens(), - *prefill_sequence->host_kv_state().mutable_kv_blocks(), + *prefill_sequence->host_kv_state().mutable_blocks(BlockType::KV), shared_blocks_num); if (num_additional_blocks > 1) { - const auto host_blocks = prefill_sequence->host_kv_state().kv_blocks(); + const auto host_blks = + prefill_sequence->host_kv_state().blocks(BlockType::KV); std::vector block_transfer_infos; block_transfer_infos.reserve(num_additional_blocks); - for (int i = 0; i < num_additional_blocks - 1; i++) { + for (size_t i = 0; i < num_additional_blocks - 1; i++) { block_transfer_infos.emplace_back(BlockTransferInfo( -1, - host_blocks[shared_blocks_num + i].id(), - host_blocks[shared_blocks_num + i].get_immutable_hash_value(), + host_blks[shared_blocks_num + i].id(), + host_blks[shared_blocks_num + i].get_immutable_hash_value(), TransferType::G2H)); } @@ -286,9 +400,10 @@ bool HierarchyBlockManagerPool::update_prefetch_result( if (prefetch_result && success_cnt > 0) { int32_t dp_rank = BlockManagerPool::get_dp_rank(prefill_sequence.get()); - auto host_blocks = prefill_sequence->host_kv_state().kv_blocks(); + auto host_blocks = + prefill_sequence->host_kv_state().blocks(BlockType::KV); auto cached_blocks = - prefill_sequence->host_kv_state().shared_kv_blocks_num(); + prefill_sequence->host_kv_state().shared_blocks_num(BlockType::KV); host_block_managers_[dp_rank]->cache( host_blocks.slice(cached_blocks - success_cnt, cached_blocks)); @@ -299,7 +414,6 @@ bool HierarchyBlockManagerPool::update_prefetch_result( } void HierarchyBlockManagerPool::transfer_blocks(std::vector& batches) { - // load blocks from host to device for (size_t i = 0; i < batches.size(); i++) { if (!load_block_transfer_infos_[i].empty()) { batches[i].set_batch_id(); @@ -315,8 +429,7 @@ void HierarchyBlockManagerPool::transfer_blocks(std::vector& batches) { } void HierarchyBlockManagerPool::transfer_blocks() { - // offload blocks from device to host and kvcache store - for (int i = 0; i < offload_block_pair_queues_.size(); i++) { + for (size_t i = 0; i < offload_block_pair_queues_.size(); i++) { std::vector transfer_infos; std::vector src_blocks; std::vector dst_blocks; @@ -329,7 +442,7 @@ void HierarchyBlockManagerPool::transfer_blocks() { BlockTransferInfo(src_blocks.back().id(), dst_blocks.back().id(), dst_blocks.back().get_immutable_hash_value(), - TransferType::D2G)); + TransferType::D2H2G)); block_pair.reset(); } @@ -342,17 +455,31 @@ void HierarchyBlockManagerPool::transfer_blocks() { device_block_mgr_ptr = block_managers_[i].get(), host_block_mgr_ptr = host_block_managers_[i].get()]( std::vector>&& results) mutable { + bool copy_ok = true; for (auto&& result : results) { + if (result.hasException()) { + LOG(ERROR) << "Offload RPC failed: " + << result.exception().what(); + copy_ok = false; + continue; + } if (result.value() != host_blocks.size()) { - LOG(FATAL) << "Offload copy fail, expected " + LOG(ERROR) << "Offload copy fail, expected " << host_blocks.size() << ", got " << result.value(); + copy_ok = false; } } + // Always release the reserved ids so the block pools do not leak. + // On failure, the host copy is incomplete (or its contents are + // undefined on the ranks that raised), so skip host prefix-cache + // publish -- a subsequent H2D would hand back garbage KV. device_block_mgr_ptr->deallocate({device_blocks}); device_blocks.clear(); - host_block_mgr_ptr->cache(host_blocks); + if (copy_ok) { + host_block_mgr_ptr->cache(host_blocks); + } host_block_mgr_ptr->deallocate({host_blocks}); host_blocks.clear(); diff --git a/xllm/core/framework/block/hierarchy_block_manager_pool.h b/xllm/core/framework/block/hierarchy_block_manager_pool.h index 2075f81881..ea06b88b81 100644 --- a/xllm/core/framework/block/hierarchy_block_manager_pool.h +++ b/xllm/core/framework/block/hierarchy_block_manager_pool.h @@ -68,6 +68,9 @@ class HierarchyBlockManagerPool : public BlockManagerPool { private: void allocate_host_shared(Sequence* sequence); + // Move offload-eligible device/host KV block pairs out of the sequence into + // offload_block_pair_queues_[dp_rank] for the next D2H transfer. + void collect_offload_pairs(Sequence* sequence, int32_t dp_rank); private: Engine* engine_; diff --git a/xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.cpp b/xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.cpp index b1392bebdc..194cd9913d 100644 --- a/xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.cpp +++ b/xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.cpp @@ -162,6 +162,79 @@ DeepSeekV4KVCacheImpl::DeepSeekV4KVCacheImpl( compressed_block_type_(tensors.compressed_block_type) { } +DeepSeekV4KVCacheImpl::DeepSeekV4KVCacheImpl( + const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count) { + CHECK(kv_cache_shape.has_key_cache_shape()) + << "DeepSeek V4 host kv cache shape must contain pool counts."; + const std::vector& pool_counts = kv_cache_shape.key_cache_shape(); + CHECK_GE(pool_counts.size(), 3) << "DeepSeek V4 host kv cache shape must be " + << "[swa_count, c4_count, c128_count]."; + CHECK_GT(create_options.block_size(), 0) + << "DeepSeek V4 block_size must be positive."; + CHECK_GT(create_options.head_dim(), 0) + << "DeepSeek V4 head_dim must be positive."; + CHECK_GT(layer_count, 0) + << "DeepSeek V4 host prefix cache layer count must be positive."; + + const double factor = create_options.host_blocks_factor(); + const int64_t host_swa_count = scale_host_block_count(pool_counts[0], factor); + const int64_t host_c4_count = scale_host_block_count(pool_counts[1], factor); + const int64_t host_c128_count = + scale_host_block_count(pool_counts[2], factor); + const int64_t block_size = create_options.block_size(); + const int64_t head_dim = create_options.head_dim(); + const int64_t index_head_dim = + std::max(create_options.index_head_dim(), 1); + const int64_t n_heads = 1; + const int64_t index_n_heads = 1; + const DeepSeekV4CachePolicy cache_policy = + get_dsv4_cache_policy(create_options.dtype()); + + // Host tensor shape: device per-block dims with a layer dimension inserted at + // index 1, i.e. [host_block_count, layer_count, ...per_block_dims]. + auto host_group_shape = [&](int64_t block_count, int64_t heads, int64_t dim) { + std::vector shape = + dsv4_block_shape(block_count, block_size, heads, dim); + shape.insert(shape.begin() + 1, layer_count); + return shape; + }; + + switch (type) { + case BlockType::SWA: + host_page_aligned_regions_.reserve(1); + create_host_tensor(host_group_shape(host_swa_count, n_heads, head_dim), + create_options.dtype(), + &swa_cache_, + nullptr); + break; + case BlockType::C4: + host_page_aligned_regions_.reserve(2); + create_host_tensor(host_group_shape(host_c4_count, n_heads, head_dim), + create_options.dtype(), + &key_cache_, + nullptr); + create_host_tensor( + host_group_shape(host_c4_count, index_n_heads, index_head_dim), + cache_policy.index_dtype, + &index_cache_, + nullptr); + break; + case BlockType::C128: + host_page_aligned_regions_.reserve(1); + create_host_tensor(host_group_shape(host_c128_count, n_heads, head_dim), + create_options.dtype(), + &key_cache_, + nullptr); + break; + default: + LOG(FATAL) << "Unsupported DeepSeek V4 host block type: " + << static_cast(type); + } +} + torch::Tensor DeepSeekV4KVCacheImpl::get_k_cache() const { return key_cache_; } torch::Tensor DeepSeekV4KVCacheImpl::get_index_cache() const { @@ -233,6 +306,34 @@ std::vector DeepSeekV4KVCacheImpl::get_cache_tensors() const { return tensors; } +BlockTypeTensorMap DeepSeekV4KVCacheImpl::get_block_type_tensors( + BlockType type) const { + BlockTypeTensorMap tensor_map; + switch (type) { + case BlockType::SWA: + if (swa_cache_.defined() && swa_cache_.numel() > 0) { + tensor_map.emplace(KVCacheTensorRole::SWA, swa_cache_); + } + break; + case BlockType::C4: + if (key_cache_.defined() && key_cache_.numel() > 0 && + index_cache_.defined() && index_cache_.numel() > 0) { + tensor_map.emplace(KVCacheTensorRole::KEY, key_cache_); + tensor_map.emplace(KVCacheTensorRole::INDEX, index_cache_); + } + break; + case BlockType::C128: + if (key_cache_.defined() && key_cache_.numel() > 0 && + (!index_cache_.defined() || index_cache_.numel() == 0)) { + tensor_map.emplace(KVCacheTensorRole::KEY, key_cache_); + } + break; + default: + break; + } + return tensor_map; +} + bool DeepSeekV4KVCacheImpl::empty() const { return !swa_cache_.defined(); } std::vector> DeepSeekV4KVCacheImpl::get_shapes() const { diff --git a/xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.h b/xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.h index e160692b56..117196485a 100644 --- a/xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.h +++ b/xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.h @@ -50,6 +50,10 @@ class Dsv4StateCache final { class DeepSeekV4KVCacheImpl final : public KVCacheImpl { public: explicit DeepSeekV4KVCacheImpl(const DeepSeekV4KVCacheTensors& tensors); + DeepSeekV4KVCacheImpl(const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count); torch::Tensor get_k_cache() const override; torch::Tensor get_index_cache() const override; @@ -63,6 +67,8 @@ class DeepSeekV4KVCacheImpl final : public KVCacheImpl { torch::Tensor get_compress_index_state() const override; std::vector get_cache_tensors() const override; + BlockTypeTensorMap get_block_type_tensors(BlockType type) const override; + bool empty() const override; std::vector> get_shapes() const override; diff --git a/xllm/core/framework/kv_cache/indexed_kv_cache_impl.cpp b/xllm/core/framework/kv_cache/indexed_kv_cache_impl.cpp index f375f895bb..93fcec921f 100644 --- a/xllm/core/framework/kv_cache/indexed_kv_cache_impl.cpp +++ b/xllm/core/framework/kv_cache/indexed_kv_cache_impl.cpp @@ -65,6 +65,64 @@ IndexedKVCacheImpl::IndexedKVCacheImpl( } } +IndexedKVCacheImpl::IndexedKVCacheImpl( + const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count) + : KVCacheImpl() { + CHECK(type == BlockType::KV) + << "IndexedKVCacheImpl host cache only supports BlockType::KV."; + host_page_aligned_regions_.reserve(4); + if (kv_cache_shape.has_key_cache_shape()) { + create_host_tensor( + build_host_group_tensor_shape(kv_cache_shape.key_cache_shape(), + create_options.host_blocks_factor(), + layer_count), + create_options.dtype(), + &key_cache_, + &key_cache_shape_); + } + if (kv_cache_shape.has_value_cache_shape()) { + create_host_tensor( + build_host_group_tensor_shape(kv_cache_shape.value_cache_shape(), + create_options.host_blocks_factor(), + layer_count), + create_options.dtype(), + &value_cache_, + &value_cache_shape_); + } + if (kv_cache_shape.has_index_cache_shape()) { + // Mirror the device index dtype: INT8 when indexer cache is quantized (see + // create_indexed_kv_cache_tensors), otherwise the model dtype. + const torch::ScalarType index_dtype = + create_options.enable_indexer_cache_quant() ? torch::kChar + : create_options.dtype(); + create_host_tensor( + build_host_group_tensor_shape(kv_cache_shape.index_cache_shape(), + create_options.host_blocks_factor(), + layer_count), + index_dtype, + &index_cache_, + &index_cache_shape_); + } + // The INT8 indexer cache keeps a per-token fp32 scale that must move with the + // int8 values during offload/reload, otherwise dequantization reads + // mismatched coefficients. Allocate it on host alongside the index cache. + if (create_options.enable_indexer_cache_quant() && + kv_cache_shape.has_index_cache_scale_shape()) { + torch::Tensor index_scale; + create_host_tensor( + build_host_group_tensor_shape(kv_cache_shape.index_cache_scale_shape(), + create_options.host_blocks_factor(), + layer_count), + torch::kFloat32, + &index_scale, + &index_cache_scale_shape_); + index_cache_scale_ = index_scale; + } +} + torch::Tensor IndexedKVCacheImpl::get_index_cache() const { return index_cache_; } @@ -94,6 +152,22 @@ std::optional IndexedKVCacheImpl::get_indexer_cache_scale() return std::nullopt; } +BlockTypeTensorMap IndexedKVCacheImpl::get_block_type_tensors( + BlockType type) const { + BlockTypeTensorMap tensor_map = KVCacheImpl::get_block_type_tensors(type); + if (type != BlockType::KV) { + return tensor_map; + } + if (index_cache_.defined() && index_cache_.numel() > 0) { + tensor_map.emplace(KVCacheTensorRole::INDEX, index_cache_); + } + if (index_cache_scale_.has_value() && index_cache_scale_->defined() && + index_cache_scale_->numel() > 0) { + tensor_map.emplace(KVCacheTensorRole::INDEX_SCALE, *index_cache_scale_); + } + return tensor_map; +} + bool IndexedKVCacheImpl::empty() const { return !key_cache_.defined() || !value_cache_.defined() || !index_cache_.defined(); diff --git a/xllm/core/framework/kv_cache/indexed_kv_cache_impl.h b/xllm/core/framework/kv_cache/indexed_kv_cache_impl.h index 66d592aef2..2659c59cc9 100644 --- a/xllm/core/framework/kv_cache/indexed_kv_cache_impl.h +++ b/xllm/core/framework/kv_cache/indexed_kv_cache_impl.h @@ -24,12 +24,18 @@ class IndexedKVCacheImpl final : public KVCacheImpl { explicit IndexedKVCacheImpl(const IndexedKVCacheTensors& tensors); IndexedKVCacheImpl(const KVCacheShape& kv_cache_shape, const KVCacheCreateOptions& create_options); + IndexedKVCacheImpl(const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count); torch::Tensor get_index_cache() const override; std::optional get_k_cache_scale() const override; std::optional get_v_cache_scale() const override; std::optional get_indexer_cache_scale() const override; + BlockTypeTensorMap get_block_type_tensors(BlockType type) const override; + bool empty() const override; std::vector> get_shapes() const override; diff --git a/xllm/core/framework/kv_cache/kv_cache.cpp b/xllm/core/framework/kv_cache/kv_cache.cpp index f12e67d952..f22b5c7f55 100644 --- a/xllm/core/framework/kv_cache/kv_cache.cpp +++ b/xllm/core/framework/kv_cache/kv_cache.cpp @@ -80,6 +80,33 @@ std::unique_ptr create_kv_cache_impl( return std::make_unique(kv_cache_shape, create_options); } +std::unique_ptr create_host_kv_cache_impl( + const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count) { + if (util::is_deepseek_v4_model_type(create_options.model_type())) { + return std::make_unique( + kv_cache_shape, create_options, type, layer_count); + } + + switch (type) { + case BlockType::SINGLE: + return std::make_unique( + kv_cache_shape, create_options, type, layer_count); + case BlockType::KV: + if (create_options.enable_lighting_indexer()) { + return std::make_unique( + kv_cache_shape, create_options, type, layer_count); + } + return std::make_unique( + kv_cache_shape, create_options, type, layer_count); + default: + LOG(FATAL) << "Unsupported non-DSV4 host block type: " + << static_cast(type); + } +} + std::string int32_vector_string(const std::vector& values) { std::ostringstream oss; oss << "["; @@ -180,6 +207,15 @@ KVCache::KVCache(const KVCacheShape& kv_cache_shape, int64_t layer_id) : impl_(create_kv_cache_impl(kv_cache_shape, create_options, layer_id)) {} +KVCache::KVCache(const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count) + : impl_(create_host_kv_cache_impl(kv_cache_shape, + create_options, + type, + layer_count)) {} + torch::Tensor KVCache::get_k_cache() const { return impl_->get_k_cache(); } torch::Tensor KVCache::get_v_cache() const { return impl_->get_v_cache(); } @@ -212,6 +248,10 @@ torch::Tensor KVCache::get_ssm_cache() const { return impl_->get_ssm_cache(); } torch::Tensor KVCache::get_swa_cache() const { return impl_->get_swa_cache(); } +BlockTypeTensorMap KVCache::get_block_type_tensors(BlockType type) const { + return impl_->get_block_type_tensors(type); +} + torch::Tensor KVCache::get_compress_kv_state() const { return impl_->get_compress_kv_state(); } diff --git a/xllm/core/framework/kv_cache/kv_cache.h b/xllm/core/framework/kv_cache/kv_cache.h index 5b714bec3e..493c4280d5 100644 --- a/xllm/core/framework/kv_cache/kv_cache.h +++ b/xllm/core/framework/kv_cache/kv_cache.h @@ -37,6 +37,10 @@ class KVCache final { KVCache(const KVCacheShape& kv_cache_shape, const KVCacheCreateOptions& create_options, int64_t layer_id); + KVCache(const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count); KVCache(const KVCache&) = delete; KVCache& operator=(const KVCache&) = delete; KVCache(KVCache&&) noexcept = default; @@ -61,6 +65,7 @@ class KVCache final { torch::Tensor get_compress_index_score_state() const; torch::Tensor get_compress_state() const; torch::Tensor get_compress_index_state() const; + BlockTypeTensorMap get_block_type_tensors(BlockType type) const; std::vector> get_shapes(); bool empty() const; diff --git a/xllm/core/framework/kv_cache/kv_cache_estimation.cpp b/xllm/core/framework/kv_cache/kv_cache_estimation.cpp index f7b317ed34..fdb0d7458c 100644 --- a/xllm/core/framework/kv_cache/kv_cache_estimation.cpp +++ b/xllm/core/framework/kv_cache/kv_cache_estimation.cpp @@ -253,6 +253,9 @@ int64_t calculate_linear_state_blocks(int64_t cache_size_in_bytes, linear_slot_size, num_full_attention_layers, full_attention_block_size); + if (linear_slot_size <= 0 || num_linear_attention_layers <= 0) { + return kPaddingLinearStateBlocks; + } if (max_linear_state_cache_slots > 0) { const int64_t requested_blocks = max_linear_state_cache_slots + kPaddingLinearStateBlocks; diff --git a/xllm/core/framework/kv_cache/kv_cache_impl.cpp b/xllm/core/framework/kv_cache/kv_cache_impl.cpp index 34281043c0..0eae09e8fb 100644 --- a/xllm/core/framework/kv_cache/kv_cache_impl.cpp +++ b/xllm/core/framework/kv_cache/kv_cache_impl.cpp @@ -15,12 +15,27 @@ limitations under the License. #include "framework/kv_cache/kv_cache_impl.h" +#include + #include "framework/kv_cache/kv_cache_shape.h" #include "framework/kv_cache/kv_cache_utils.h" #include "util/tensor_helper.h" namespace xllm { +void KVCacheImpl::create_host_tensor(const std::vector& dims, + torch::ScalarType dtype, + torch::Tensor* tensor, + std::vector* shape) { + CHECK(tensor != nullptr) << "tensor must not be null."; + HostPageAlignedRegion region; + create_host_page_aligned_tensor(dims, dtype, tensor, ®ion); + host_page_aligned_regions_.emplace_back(std::move(region)); + if (shape != nullptr) { + *shape = dims; + } +} + KVCacheImpl::KVCacheImpl(const KVCacheTensors& tensors) : key_cache_(tensors.key_cache), value_cache_(tensors.value_cache), @@ -36,6 +51,33 @@ KVCacheImpl::KVCacheImpl(const KVCacheShape& kv_cache_shape, } } +KVCacheImpl::KVCacheImpl(const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count) { + CHECK(type == BlockType::KV) + << "Base KVCacheImpl host cache only supports BlockType::KV."; + host_page_aligned_regions_.reserve(2); + if (kv_cache_shape.has_key_cache_shape()) { + create_host_tensor( + build_host_group_tensor_shape(kv_cache_shape.key_cache_shape(), + create_options.host_blocks_factor(), + layer_count), + create_options.dtype(), + &key_cache_, + &key_cache_shape_); + } + if (kv_cache_shape.has_value_cache_shape()) { + create_host_tensor( + build_host_group_tensor_shape(kv_cache_shape.value_cache_shape(), + create_options.host_blocks_factor(), + layer_count), + create_options.dtype(), + &value_cache_, + &value_cache_shape_); + } +} + torch::Tensor KVCacheImpl::get_k_cache() const { return key_cache_; } torch::Tensor KVCacheImpl::get_v_cache() const { return value_cache_; } @@ -120,6 +162,20 @@ std::vector KVCacheImpl::get_cache_tensors() const { return tensors; } +BlockTypeTensorMap KVCacheImpl::get_block_type_tensors(BlockType type) const { + BlockTypeTensorMap tensor_map; + if (type != BlockType::KV) { + return tensor_map; + } + if (key_cache_.defined() && key_cache_.numel() > 0) { + tensor_map.emplace(KVCacheTensorRole::KEY, key_cache_); + } + if (value_cache_.defined() && value_cache_.numel() > 0) { + tensor_map.emplace(KVCacheTensorRole::VALUE, value_cache_); + } + return tensor_map; +} + bool KVCacheImpl::empty() const { return !key_cache_.defined() || !value_cache_.defined(); } diff --git a/xllm/core/framework/kv_cache/kv_cache_impl.h b/xllm/core/framework/kv_cache/kv_cache_impl.h index 315ae488e6..93040918d8 100644 --- a/xllm/core/framework/kv_cache/kv_cache_impl.h +++ b/xllm/core/framework/kv_cache/kv_cache_impl.h @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include "framework/block/block.h" #include "framework/kv_cache/kv_cache_utils.h" namespace xllm { @@ -32,6 +33,10 @@ class KVCacheImpl { explicit KVCacheImpl(const KVCacheTensors& tensors); KVCacheImpl(const KVCacheShape& kv_cache_shape, const KVCacheCreateOptions& create_options); + KVCacheImpl(const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count); virtual ~KVCacheImpl() = default; @@ -52,6 +57,12 @@ class KVCacheImpl { virtual torch::Tensor get_compress_index_state() const; virtual std::vector get_cache_tensors() const; + // Host prefix cache extraction by block type. Unlike the device-side + // build_block_type_tensor_map in the transfer layer, this carries no + // device-pool discriminator guards: a per-type host cache holds exactly the + // tensors for that block type. + virtual BlockTypeTensorMap get_block_type_tensors(BlockType type) const; + virtual bool empty() const; virtual std::vector> get_shapes() const; @@ -60,6 +71,12 @@ class KVCacheImpl { torch::Tensor& dst_tensor); protected: + void create_host_tensor(const std::vector& dims, + torch::ScalarType dtype, + torch::Tensor* tensor, + std::vector* shape); + + std::vector host_page_aligned_regions_; torch::Tensor key_cache_; torch::Tensor value_cache_; std::vector key_cache_shape_; diff --git a/xllm/core/framework/kv_cache/kv_cache_tensor_role.h b/xllm/core/framework/kv_cache/kv_cache_tensor_role.h index 296ef752c9..608f4449eb 100644 --- a/xllm/core/framework/kv_cache/kv_cache_tensor_role.h +++ b/xllm/core/framework/kv_cache/kv_cache_tensor_role.h @@ -37,6 +37,7 @@ class KVCacheTensorRole { SCORE_STATE = 11, INDEX_KV_STATE = 12, INDEX_SCORE_STATE = 13, + SWA = 14, INVALID = -1, }; @@ -70,6 +71,8 @@ class KVCacheTensorRole { value_ = INDEX_KV_STATE; } else if (str == "INDEX_SCORE_STATE" || str == "index_score_state") { value_ = INDEX_SCORE_STATE; + } else if (str == "SWA" || str == "swa") { + value_ = SWA; } else { value_ = INVALID; } @@ -114,6 +117,8 @@ class KVCacheTensorRole { return "index_kv_state"; } else if (this->value_ == INDEX_SCORE_STATE) { return "index_score_state"; + } else if (this->value_ == SWA) { + return "swa"; } else { return "invalid"; } diff --git a/xllm/core/framework/kv_cache/kv_cache_utils.cpp b/xllm/core/framework/kv_cache/kv_cache_utils.cpp index 5fcfb918b4..6c2f58e64f 100644 --- a/xllm/core/framework/kv_cache/kv_cache_utils.cpp +++ b/xllm/core/framework/kv_cache/kv_cache_utils.cpp @@ -16,7 +16,12 @@ limitations under the License. #include "framework/kv_cache/kv_cache_utils.h" #include +#include +#include +#include +#include +#include #include #include "core/framework/config/kv_cache_config.h" @@ -25,13 +30,18 @@ limitations under the License. #include "platform/mlu/mlu_tensor_alloc.h" #endif #if defined(USE_NPU) -#include "acl/acl.h" +#include "acl/acl_rt.h" + +extern "C" aclError aclrtHostRegister(void* ptr, + uint64_t size, + aclrtHostRegisterType type, + void** dev_ptr); +extern "C" aclError aclrtHostUnregister(void* ptr); #endif namespace xllm { namespace { -#if defined(USE_NPU) size_t get_tensor_nbytes(const std::vector& dims, torch::ScalarType dtype) { size_t count = 1; @@ -51,6 +61,7 @@ size_t get_tensor_nbytes(const std::vector& dims, return count * elem_size; } +#if defined(USE_NPU) void free_acl_tensor(void* ptr) { if (ptr == nullptr) { return; @@ -369,4 +380,130 @@ aclFormat get_npu_kv_cache_format(const std::string& model_type) { } #endif +HostPageAlignedRegion::HostPageAlignedRegion(size_t bytes) { + if (bytes == 0) { + return; + } + size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + total_bytes = ((bytes + page_size - 1) / page_size) * page_size; + base_ptr = mmap(nullptr, + total_bytes, + PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, + -1, + 0); + CHECK(base_ptr != MAP_FAILED) + << "Failed to mmap host memory, size=" << total_bytes; + if (mlock(base_ptr, total_bytes) != 0) { + const int32_t err = errno; + munmap(base_ptr, total_bytes); + base_ptr = nullptr; + total_bytes = 0; + LOG(FATAL) << "mlock failed, errno=" << err << " (" << strerror(err) + << "). Run: ulimit -l unlimited"; + } + +#if defined(USE_NPU) + void* mapped_ptr = nullptr; + const aclError ret = aclrtHostRegister( + base_ptr, total_bytes, ACL_HOST_REGISTER_MAPPED, &mapped_ptr); + if (ret != ACL_SUCCESS) { + munlock(base_ptr, total_bytes); + munmap(base_ptr, total_bytes); + base_ptr = nullptr; + total_bytes = 0; + } + CHECK_EQ(ret, ACL_SUCCESS) << "aclrtHostRegister fail: " << ret; +#endif +} + +HostPageAlignedRegion::HostPageAlignedRegion( + HostPageAlignedRegion&& other) noexcept + : base_ptr(other.base_ptr), total_bytes(other.total_bytes) { + other.base_ptr = nullptr; + other.total_bytes = 0; +} + +namespace { + +void release_host_page_aligned_region(void*& base_ptr, size_t& total_bytes) { + if (base_ptr == nullptr || total_bytes == 0) { + base_ptr = nullptr; + total_bytes = 0; + return; + } +#if defined(USE_NPU) + aclrtHostUnregister(base_ptr); +#endif + munlock(base_ptr, total_bytes); + munmap(base_ptr, total_bytes); + base_ptr = nullptr; + total_bytes = 0; +} + +} // namespace + +HostPageAlignedRegion& HostPageAlignedRegion::operator=( + HostPageAlignedRegion&& other) noexcept { + if (this == &other) { + return *this; + } + release_host_page_aligned_region(base_ptr, total_bytes); + base_ptr = other.base_ptr; + total_bytes = other.total_bytes; + other.base_ptr = nullptr; + other.total_bytes = 0; + return *this; +} + +HostPageAlignedRegion::~HostPageAlignedRegion() { + release_host_page_aligned_region(base_ptr, total_bytes); +} + +int64_t scale_host_block_count(int64_t block_count, double host_blocks_factor) { + CHECK_GT(block_count, 0) << "block_count must be positive."; + const double factor = std::max(host_blocks_factor, 1.0); + return std::max(block_count, + static_cast(block_count * factor)); +} + +std::vector build_host_tensor_shape( + const std::vector& base_shape, + double host_blocks_factor) { + CHECK(!base_shape.empty()) << "base_shape must not be empty."; + std::vector host_shape = base_shape; + host_shape[0] = scale_host_block_count(host_shape[0], host_blocks_factor); + return host_shape; +} + +std::vector build_host_group_tensor_shape( + const std::vector& base_shape, + double host_blocks_factor, + int64_t layer_count) { + CHECK_GT(layer_count, 0) << "layer_count must be positive."; + std::vector host_shape = + build_host_tensor_shape(base_shape, host_blocks_factor); + host_shape.insert(host_shape.begin() + 1, layer_count); + return host_shape; +} + +void create_host_page_aligned_tensor(const std::vector& dims, + torch::ScalarType dtype, + torch::Tensor* tensor, + HostPageAlignedRegion* region) { + CHECK(tensor != nullptr) << "tensor must not be null."; + CHECK(region != nullptr) << "region must not be null."; + const size_t tensor_bytes = get_tensor_nbytes(dims, dtype); + CHECK_GT(tensor_bytes, static_cast(0)) + << "host cache tensor bytes must be positive."; + + *region = HostPageAlignedRegion(tensor_bytes); + std::memset(region->base_ptr, 0, region->total_bytes); + *tensor = torch::from_blob( + region->base_ptr, + dims, + torch::TensorOptions().dtype(dtype).device(torch::Device(torch::kCPU))); + CHECK(tensor->is_contiguous()) << "host cache tensor must be contiguous."; +} + } // namespace xllm diff --git a/xllm/core/framework/kv_cache/kv_cache_utils.h b/xllm/core/framework/kv_cache/kv_cache_utils.h index 0ded915f6b..4719a2a800 100644 --- a/xllm/core/framework/kv_cache/kv_cache_utils.h +++ b/xllm/core/framework/kv_cache/kv_cache_utils.h @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -52,6 +53,7 @@ struct KVCacheCreateOptions { PROPERTY(torch::ScalarType, dtype) = torch::kBFloat16; // ssm dtype for linear attention layers PROPERTY(torch::ScalarType, ssm_dtype) = torch::kBFloat16; + PROPERTY(double, host_blocks_factor) = 0.0; PROPERTY(int64_t, num_layers) = 0; // full attention interval for linear attention layers PROPERTY(int64_t, full_attention_interval) = 1; @@ -110,6 +112,21 @@ struct KVCacheTensor { bool sequence_scoped = false; }; +using BlockTypeTensorMap = std::map; + +struct HostPageAlignedRegion { + void* base_ptr = nullptr; + size_t total_bytes = 0; + + HostPageAlignedRegion() = default; + explicit HostPageAlignedRegion(size_t bytes); + HostPageAlignedRegion(const HostPageAlignedRegion&) = delete; + HostPageAlignedRegion& operator=(const HostPageAlignedRegion&) = delete; + HostPageAlignedRegion(HostPageAlignedRegion&& other) noexcept; + HostPageAlignedRegion& operator=(HostPageAlignedRegion&& other) noexcept; + ~HostPageAlignedRegion(); +}; + struct DeepSeekV4KVCacheTensors { torch::Tensor key_cache; torch::Tensor index_cache; @@ -149,6 +166,30 @@ LinearAttentionKVCacheTensors create_linear_attention_kv_cache_tensors( const KVCacheShape& kv_cache_shape, const KVCacheCreateOptions& create_options); +// Scale a device block count to the host block count using host_blocks_factor +// (clamped to >= 1.0 so the host pool is never smaller than the device pool). +int64_t scale_host_block_count(int64_t block_count, double host_blocks_factor); + +// Build a host tensor shape from a per-layer device shape by scaling dim 0 +// (block count) by host_blocks_factor. +std::vector build_host_tensor_shape( + const std::vector& base_shape, + double host_blocks_factor); + +// Build a grouped host tensor shape: scales dim 0 then inserts a layer +// dimension at index 1, yielding [host_blocks, layer_count, ...per_block_dims]. +std::vector build_host_group_tensor_shape( + const std::vector& base_shape, + double host_blocks_factor, + int64_t layer_count); + +// Allocate a page-aligned, mlock'd (and NPU-registered) host tensor over a +// HostPageAlignedRegion. The region owns the memory; the tensor is a view. +void create_host_page_aligned_tensor(const std::vector& dims, + torch::ScalarType dtype, + torch::Tensor* tensor, + HostPageAlignedRegion* region); + #if defined(USE_NPU) aclFormat get_npu_kv_cache_format(const std::string& model_type); diff --git a/xllm/core/framework/kv_cache/linear_attention_kv_cache_impl.cpp b/xllm/core/framework/kv_cache/linear_attention_kv_cache_impl.cpp index b892acb570..26466387bc 100644 --- a/xllm/core/framework/kv_cache/linear_attention_kv_cache_impl.cpp +++ b/xllm/core/framework/kv_cache/linear_attention_kv_cache_impl.cpp @@ -37,6 +37,36 @@ LinearAttentionKVCacheImpl::LinearAttentionKVCacheImpl( ssm_cache_shape_ = kv_cache_shape.ssm_cache_shape(); } +LinearAttentionKVCacheImpl::LinearAttentionKVCacheImpl( + const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count) + : KVCacheImpl() { + CHECK(type == BlockType::SINGLE) + << "LinearAttentionKVCacheImpl host cache only supports " + "BlockType::SINGLE."; + host_page_aligned_regions_.reserve(2); + if (kv_cache_shape.has_conv_cache_shape()) { + create_host_tensor( + build_host_group_tensor_shape(kv_cache_shape.conv_cache_shape(), + create_options.host_blocks_factor(), + layer_count), + create_options.dtype(), + &conv_cache_, + &conv_cache_shape_); + } + if (kv_cache_shape.has_ssm_cache_shape()) { + create_host_tensor( + build_host_group_tensor_shape(kv_cache_shape.ssm_cache_shape(), + create_options.host_blocks_factor(), + layer_count), + create_options.ssm_dtype(), + &ssm_cache_, + &ssm_cache_shape_); + } +} + torch::Tensor LinearAttentionKVCacheImpl::get_conv_cache() const { return conv_cache_; } @@ -45,6 +75,21 @@ torch::Tensor LinearAttentionKVCacheImpl::get_ssm_cache() const { return ssm_cache_; } +BlockTypeTensorMap LinearAttentionKVCacheImpl::get_block_type_tensors( + BlockType type) const { + BlockTypeTensorMap tensor_map; + if (type != BlockType::SINGLE) { + return tensor_map; + } + if (conv_cache_.defined() && conv_cache_.numel() > 0) { + tensor_map.emplace(KVCacheTensorRole::CONV, conv_cache_); + } + if (ssm_cache_.defined() && ssm_cache_.numel() > 0) { + tensor_map.emplace(KVCacheTensorRole::SSM, ssm_cache_); + } + return tensor_map; +} + bool LinearAttentionKVCacheImpl::empty() const { return !conv_cache_.defined() || !ssm_cache_.defined(); } diff --git a/xllm/core/framework/kv_cache/linear_attention_kv_cache_impl.h b/xllm/core/framework/kv_cache/linear_attention_kv_cache_impl.h index 04ae31a1bb..fc3a75880c 100644 --- a/xllm/core/framework/kv_cache/linear_attention_kv_cache_impl.h +++ b/xllm/core/framework/kv_cache/linear_attention_kv_cache_impl.h @@ -25,10 +25,16 @@ class LinearAttentionKVCacheImpl final : public KVCacheImpl { const LinearAttentionKVCacheTensors& tensors); LinearAttentionKVCacheImpl(const KVCacheShape& kv_cache_shape, const KVCacheCreateOptions& create_options); + LinearAttentionKVCacheImpl(const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options, + BlockType type, + int64_t layer_count); torch::Tensor get_conv_cache() const override; torch::Tensor get_ssm_cache() const override; + BlockTypeTensorMap get_block_type_tensors(BlockType type) const override; + bool empty() const override; std::vector> get_shapes() const override; diff --git a/xllm/core/framework/kv_cache_transfer/CMakeLists.txt b/xllm/core/framework/kv_cache_transfer/CMakeLists.txt index 834709a232..8f9433242e 100644 --- a/xllm/core/framework/kv_cache_transfer/CMakeLists.txt +++ b/xllm/core/framework/kv_cache_transfer/CMakeLists.txt @@ -41,8 +41,7 @@ cc_library( $<$:llm_data_dist_transfer.h> $<$:spec_kv_cache_transfer.h> kv_cache_store.h - # hierarchy temporarily disabled during the block-manager refactor - # hierarchy_kv_cache_transfer.h + hierarchy_kv_cache_transfer.h $<$,$,$>:mooncake_transfer_engine.h> $<$,$,$>:mooncake_kv_cache_transfer.h> $<$:mooncake_weight_transfer.h> @@ -51,8 +50,7 @@ cc_library( $<$:llm_data_dist_transfer.cpp> $<$:spec_kv_cache_transfer.cpp> kv_cache_store.cpp - # hierarchy temporarily disabled during the block-manager refactor - # hierarchy_kv_cache_transfer.cpp + hierarchy_kv_cache_transfer.cpp $<$,$,$>:mooncake_transfer_engine.cpp> $<$,$,$>:mooncake_kv_cache_transfer.cpp> $<$:mooncake_weight_transfer.cpp> @@ -61,6 +59,7 @@ cc_library( :kv_cache :push_route :xtensor + :platform $<$:graph> glog::glog $<$:llm_datadist> 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 2b10cd4cd4..a13fa2bb57 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 @@ -15,627 +15,469 @@ limitations under the License. #include "framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h" -#include -#include - -#include -#include +#include #include +#include -#include "framework/kv_cache_transfer/kv_cache_store.h" namespace xllm { +namespace { + +constexpr uint32_t TIMEOUT_MS = 60000; + +// 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; + +std::vector build_layer_batch_ranges( + int64_t num_layers, + uint32_t requested_batches) { + std::vector 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)}); + } + return ranges; +} -constexpr uint64_t MBUF_SIZE = 128 * 1024 * 1024; -constexpr uint32_t BATCH_COPY_MAX_SIZE = 4096; -constexpr uint32_t TIMEOUT_S = 60; // second -constexpr uint32_t TIMEOUT_MS = 60000; // millisecond +bool has_tensor(const torch::Tensor& tensor) { + return tensor.defined() && tensor.numel() > 0; +} + +BlockTypeTensorMap build_block_type_tensor_map(const KVCache& kv_cache, + BlockType type) { + BlockTypeTensorMap map; + + 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(); + + 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::SINGLE: + 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: + if (!has_tensor(swa_cache) || has_tensor(key_cache) || + has_tensor(value_cache) || has_tensor(index_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); + 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 HierarchyKVCacheTransfer::HierarchyKVCacheTransfer( const Options& options, const torch::Device& device, - std::vector* kv_caches_ptr) - : options_(options), device_(device), kv_caches_ptr_(kv_caches_ptr) { + std::vector* kv_caches_ptr, + const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options) + : options_(options), + device_(device), + 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."; + device_.set_device(); device_.init_device_context(); - h2d_threadpool_ = std::make_unique( + load_threadpool_ = std::make_unique( /*num_threads=*/2, /*init_func=*/[this]() mutable { device_.set_device(); }, /*cpu_binding=*/false, - /*pool_name=*/"HierarchyKVCacheTransfer.h2d"); - d2h_threadpool_ = std::make_unique( - /*num_threads=*/5, - /*init_func=*/[this]() mutable { device_.set_device(); }, - /*cpu_binding=*/false, - /*pool_name=*/"HierarchyKVCacheTransfer.d2h"); - for (int i = 0; i < h2d_threadpool_->size() + d2h_threadpool_->size(); i++) { + /*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)); } - if (options_.host_blocks_factor() > 1) { - create_page_aligned_host_cache(); + 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()) { - KVCacheStoreInitConfig config; - config.localhost_name = options_.store_local_hostname(); - config.protocol = options_.store_protocol(); - config.metadata_server = options_.store_metadata_server(); - config.master_server_address = options_.store_master_server_address(); - config.tp_rank = options_.tp_rank(); - config.total_size = page_aligned_data_size_; - config.tensor_data = page_aligned_data_; - - if (!KVCacheStore::get_instance().init(config, &host_kv_caches_)) { - LOG(FATAL) << "Init KVCacheStore fail!"; +void HierarchyKVCacheTransfer::build_device_block_type_map() { + device_kv_caches_.clear(); + device_block_type_layer_ids_.clear(); + + const std::vector kBlockTypes = {BlockType::KV, + BlockType::SINGLE, + BlockType::SWA, + BlockType::C4, + BlockType::C128}; + + for (int64_t layer_id = 0; + layer_id < static_cast(kv_caches_ptr_->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); + } } } } -HierarchyKVCacheTransfer::~HierarchyKVCacheTransfer() { - if (page_aligned_data_ != nullptr) { +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)) + .enable_xtensor(false) + .enable_raw_device_allocator(false) + .host_blocks_factor(options_.host_blocks_factor()); #if defined(USE_NPU) - aclrtHostUnregister(page_aligned_data_); + host_opts.enable_kv_cache_huge_page_allocator(false); #endif - munlock(page_aligned_data_, page_aligned_data_size_); - munmap(page_aligned_data_, page_aligned_data_size_); + + host_kv_caches_[block_type] = std::make_unique( + kv_cache_shape_, host_opts, block_type, layer_count); } + + LOG(INFO) << "HierarchyKVCacheTransfer: created host cache for " + << host_kv_caches_.size() << " block type groups."; +} + +HierarchyKVCacheTransfer::CopyPlan HierarchyKVCacheTransfer::build_copy_plan( + 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()) { + 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 BlockTypeTensorMap host_tensors = + host_cache->get_block_type_tensors(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); + } + + CHECK_GE(host_block_id, 0) << "host block id must be non-negative."; + + for (size_t layer_slot = 0; layer_slot < group_caches.size(); + ++layer_slot) { + const int64_t absolute_layer_id = layer_ids[layer_slot]; + if (absolute_layer_id < layer_batch_range.begin_layer || + absolute_layer_id >= layer_batch_range.end_layer) { + continue; + } + + BlockTypeTensorMap device_tensors = + build_block_type_tensor_map(*group_caches[layer_slot], 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 torch::Tensor& host_tensor = host_tensor_it->second; + CHECK_LT(host_block_id, host_tensor.size(0)) + << "host block id out of range."; + 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); + } else { + plan.src_tensors.emplace_back(device_block); + plan.dst_tensors.emplace_back(host_block_layer); + } + } + } + } + + return plan; } uint32_t HierarchyKVCacheTransfer::transfer_kv_blocks( - const uint64_t batch_id, + uint64_t batch_id, const std::vector& block_transfer_info) { 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) { + case TransferType::D2H2G: + return offload(block_transfer_info); case TransferType::H2D: { - h2d_threadpool_->schedule( + // Create and register the synchronizer synchronously, before scheduling + // the async copy. The scheduler issues transfer_kv_blocks before step, so + // by the time this returns the entry is in the map and the forward path's + // set_layer_synchronizer is guaranteed to find it. Registering inside the + // async load_from_host would race the forward lookup (it could miss the + // entry, skip the wait, and read KV before the H2D copy completed). + 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_."; + { + 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."; + } + layer_wise_load_synchronizer_[batch_id] = synchronizer; + } + load_threadpool_->schedule( [this, - batch_id = batch_id, + synchronizer, block_transfer_info = std::move(block_transfer_info)]() mutable { - Slice info_slice{block_transfer_info}; - h2d_batch_copy(batch_id, info_slice); + load_from_host(synchronizer, block_transfer_info); }); - break; - } - case TransferType::D2G: - return offload_kv_blocks(std::move(block_transfer_info)); - case TransferType::G2D: { - // TODO load_kv_blocks async - LOG(ERROR) << "Unsupport copy type G2D."; - break; + return 0; } default: - LOG(ERROR) << "Unsupport copy type: " - << uint32_t(block_transfer_info[0].transfer_type); - break; + LOG(ERROR) << "Unsupported transfer type: " + << static_cast(block_transfer_info[0].transfer_type); + return 0; } - return 0; } uint32_t HierarchyKVCacheTransfer::transfer_kv_blocks( - const uint64_t batch_id, + uint64_t /*batch_id*/, Slice& block_transfer_info) { CHECK(!block_transfer_info.empty()); - - switch (block_transfer_info[0].transfer_type) { - case TransferType::G2H: - return load_from_store(block_transfer_info); - default: - LOG(ERROR) << "Unsupport copy type: " - << uint32_t(block_transfer_info[0].transfer_type); - break; - } + LOG(ERROR) << "Slice-based transfer not supported in Phase 1 (no store)."; return 0; } -void HierarchyKVCacheTransfer::set_layer_synchronizer( - ModelInputParams& params) { -#if defined(USE_NPU) - { - std::lock_guard lock(mutex_); - if (layer_wise_load_synchronizer_.count(params.meta.batch_id) != 0) { - params.parallel.layer_wise_load_synchronizer = - layer_wise_load_synchronizer_[params.meta.batch_id]; - layer_wise_load_synchronizer_.erase(params.meta.batch_id); - uint32_t event_cnt = - params.parallel.layer_wise_load_synchronizer->get_event_size(); - params.parallel.layers_per_bacth_copy = - (options_.layers() + event_cnt - 1) / event_cnt; - } - } -#endif -} - -uint32_t HierarchyKVCacheTransfer::offload_kv_blocks( +uint32_t HierarchyKVCacheTransfer::offload( const std::vector& block_transfer_info) { if (block_transfer_info.empty()) { return 0; } - const int64_t num_layers = options_.layers(); - uint32_t max_blocks_per_batch = - BATCH_COPY_MAX_SIZE / (cache_tensor_cnt_ * num_layers); - uint32_t total_slice = - (block_transfer_info.size() + max_blocks_per_batch - 1) / - max_blocks_per_batch; - - Slice transfer_info_slice(block_transfer_info); - std::vector> futures; - futures.reserve(total_slice); - - for (size_t i = 0; i < block_transfer_info.size(); - i += max_blocks_per_batch) { - folly::Promise promise; - auto future = promise.getSemiFuture(); - auto slice = transfer_info_slice.slice( - i, std::min(i + max_blocks_per_batch, block_transfer_info.size())); - - d2h_threadpool_->schedule([this, - promise = std::move(promise), - slice = std::move(slice)]() mutable { - bool ret = d2h_batch_copy(slice); - auto success_cnt = offload_to_store(slice); - if (success_cnt != slice.size()) { - LOG(WARNING) << "KVCacheStore not all put success: " << success_cnt - << "/" << slice.size(); - } - promise.setValue(ret); - }); - - futures.emplace_back(std::move(future)); + if (batch_memcpy_ == nullptr) { + return block_transfer_info.size(); } - if (!futures.empty()) { - try { - // TODO(kangmeng): add timeout - auto all_results = folly::collect(futures).get(); - if (!std::all_of(all_results.begin(), all_results.end(), [](bool result) { - return result; - })) { - LOG(FATAL) << "Not all D2H copy returned true"; - } - } catch (const std::exception& e) { - LOG(FATAL) << "Future execution failed: " << e.what(); - } + Slice slice(block_transfer_info); + if (!offload_to_host(slice)) { + LOG(ERROR) << "Offload to host failed."; + return 0; } - return block_transfer_info.size(); } -bool HierarchyKVCacheTransfer::d2h_batch_copy( +bool HierarchyKVCacheTransfer::offload_to_host( Slice& block_transfer_info) { -#if defined(USE_NPU) - const int64_t num_layers = options_.layers(); - uint32_t num_batches = - block_transfer_info.size() * num_layers * cache_tensor_cnt_; - void** srcs = new void*[num_batches]; - void** dsts = new void*[num_batches]; - size_t* copy_size = new size_t[num_batches]; - aclrtMemcpyBatchAttr attrs[1] = {d2h_attrs_}; - size_t attrs_indexes[1] = {0}; - size_t fail_index; - uint32_t curr_index = 0; - - for (const auto& info : block_transfer_info) { - auto dst_k_cache = host_kv_caches_.at(info.dst_block_id).get_k_cache(); - auto dst_v_cache = host_kv_caches_.at(info.dst_block_id).get_v_cache(); - auto dst_index_cache = - host_kv_caches_.at(info.dst_block_id).get_index_cache(); - auto dst_index_cache_scale = - host_kv_caches_.at(info.dst_block_id).get_indexer_cache_scale(); - - for (int layer_id = 0; layer_id < num_layers; layer_id++) { - auto src_k_cache = kv_caches_ptr_->at(layer_id).get_k_cache(); - srcs[curr_index] = src_k_cache[info.src_block_id].data_ptr(); - dsts[curr_index] = dst_k_cache[layer_id].data_ptr(); - copy_size[curr_index] = cache_size_per_layer_[0]; - curr_index++; - - if (cache_size_per_layer_[1] != 0) { - auto src_v_cache = kv_caches_ptr_->at(layer_id).get_v_cache(); - srcs[curr_index] = src_v_cache[info.src_block_id].data_ptr(); - dsts[curr_index] = dst_v_cache[layer_id].data_ptr(); - copy_size[curr_index] = cache_size_per_layer_[1]; - curr_index++; - } - - if (cache_size_per_layer_[2] != 0) { - auto src_index_cache = kv_caches_ptr_->at(layer_id).get_index_cache(); - srcs[curr_index] = src_index_cache[info.src_block_id].data_ptr(); - dsts[curr_index] = dst_index_cache[layer_id].data_ptr(); - copy_size[curr_index] = cache_size_per_layer_[2]; - curr_index++; - } - - if (cache_size_per_layer_[3] != 0) { - auto src_index_cache_scale = - kv_caches_ptr_->at(layer_id).get_indexer_cache_scale(); - CHECK(src_index_cache_scale.has_value()) - << "index cache scale is required for D2H copy."; - CHECK(dst_index_cache_scale.has_value()) - << "host index cache scale is required for D2H copy."; - srcs[curr_index] = - src_index_cache_scale.value()[info.src_block_id].data_ptr(); - dsts[curr_index] = dst_index_cache_scale.value()[layer_id].data_ptr(); - copy_size[curr_index] = cache_size_per_layer_[3]; - curr_index++; - } - } + if (block_transfer_info.empty()) { + return true; } + CHECK(batch_memcpy_ != nullptr) << "batch memcpy must be initialized."; std::unique_ptr stream; copy_stream_.wait_dequeue(stream); - c10::StreamGuard streamGuard = stream->set_stream_guard(); - - // TODO(kangmeng): change to async API - aclError ret = aclrtMemcpyBatch(dsts, - copy_size, - srcs, - copy_size, - num_batches, - attrs, - attrs_indexes, - 1, - &fail_index); - if (ret != 0 || fail_index != SIZE_MAX) { - LOG(ERROR) << "aclrtMemcpyBatch error: " << ret - << ", fail_index:" << fail_index; - copy_stream_.enqueue(std::move(stream)); - return false; - } - - if (stream->synchronize() != 0) { - LOG(ERROR) << "d2h_batch_copy timeout!"; - copy_stream_.enqueue(std::move(stream)); - return false; + 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; + } + if (!batch_memcpy_->copy_d2h( + plan.src_tensors, plan.dst_tensors, stream.get())) { + success = false; + break; + } } - copy_stream_.enqueue(std::move(stream)); - - delete[] dsts; - delete[] srcs; - delete[] copy_size; -#endif - return true; + return success; } -bool HierarchyKVCacheTransfer::h2d_batch_copy( - const uint64_t batch_id, - Slice& block_transfer_info) { -#if defined(USE_NPU) - CHECK(block_transfer_info.size() < BATCH_COPY_MAX_SIZE / cache_tensor_cnt_) - << "h2d_batch_copy support copy blocks less than " - << BATCH_COPY_MAX_SIZE / cache_tensor_cnt_ << ", but got " - << block_transfer_info.size(); - +bool HierarchyKVCacheTransfer::load_from_host( + std::shared_ptr synchronizer, + const std::vector& block_transfer_info) { if (block_transfer_info.empty()) { return true; } - const int64_t num_layers = options_.layers(); - uint32_t layers_per_bacth_copy = - num_layers / options_.layers_wise_copy_batchs(); - uint32_t num_batches = block_transfer_info.size() * cache_tensor_cnt_; - while (num_batches * layers_per_bacth_copy > BATCH_COPY_MAX_SIZE) { - layers_per_bacth_copy--; - } - - uint32_t copy_cnt = - (num_layers + layers_per_bacth_copy - 1) / layers_per_bacth_copy; - auto synchronizer = std::make_shared(copy_cnt); - { - std::lock_guard lock(mutex_); - if (layer_wise_load_synchronizer_.count(batch_id) != 0) { - LOG(FATAL) << "Batch id already exists!"; - } - layer_wise_load_synchronizer_[batch_id] = synchronizer; - } - - aclrtMemcpyBatchAttr attrs[1] = {h2d_attrs_}; - size_t attrs_indexes[1] = {0}; + CHECK(synchronizer != nullptr) << "layer synchronizer must not be null."; + CHECK(batch_memcpy_ != nullptr) << "batch memcpy must be initialized."; std::unique_ptr stream; copy_stream_.wait_dequeue(stream); - c10::StreamGuard streamGuard = stream->set_stream_guard(); - aclError ret = 0; - - void** srcs = new void*[num_batches * layers_per_bacth_copy]; - void** dsts = new void*[num_batches * layers_per_bacth_copy]; - size_t* copy_size = new size_t[num_batches * layers_per_bacth_copy]; - - for (int index = 0; index < copy_cnt; index++) { - int layer_id = index * layers_per_bacth_copy; - size_t fail_index = 0; - uint32_t curr_index = 0; - uint32_t layer_cnt = 0; - - while (layer_id < (index + 1) * layers_per_bacth_copy && - layer_id < num_layers) { - auto dst_k_cache = kv_caches_ptr_->at(layer_id).get_k_cache(); - auto dst_v_cache = kv_caches_ptr_->at(layer_id).get_v_cache(); - auto dst_index_cache = kv_caches_ptr_->at(layer_id).get_index_cache(); - auto dst_index_cache_scale = - kv_caches_ptr_->at(layer_id).get_indexer_cache_scale(); - - for (const auto& info : block_transfer_info) { - auto src_k_cache = host_kv_caches_.at(info.src_block_id).get_k_cache(); - srcs[curr_index] = src_k_cache[layer_id].data_ptr(); - dsts[curr_index] = dst_k_cache[info.dst_block_id].data_ptr(); - copy_size[curr_index] = cache_size_per_layer_[0]; - curr_index++; - - if (cache_size_per_layer_[1] != 0) { - auto src_v_cache = - host_kv_caches_.at(info.src_block_id).get_v_cache(); - srcs[curr_index] = src_v_cache[layer_id].data_ptr(); - dsts[curr_index] = dst_v_cache[info.dst_block_id].data_ptr(); - copy_size[curr_index] = cache_size_per_layer_[1]; - curr_index++; - } - - if (cache_size_per_layer_[2] != 0) { - auto src_index_cache = - host_kv_caches_.at(info.src_block_id).get_index_cache(); - srcs[curr_index] = src_index_cache[layer_id].data_ptr(); - dsts[curr_index] = dst_index_cache[info.dst_block_id].data_ptr(); - copy_size[curr_index] = cache_size_per_layer_[2]; - curr_index++; - } - - if (cache_size_per_layer_[3] != 0) { - auto src_index_cache_scale = - host_kv_caches_.at(info.src_block_id).get_indexer_cache_scale(); - CHECK(src_index_cache_scale.has_value()) - << "host index cache scale is required for H2D copy."; - CHECK(dst_index_cache_scale.has_value()) - << "index cache scale is required for H2D copy."; - srcs[curr_index] = src_index_cache_scale.value()[layer_id].data_ptr(); - dsts[curr_index] = - dst_index_cache_scale.value()[info.dst_block_id].data_ptr(); - copy_size[curr_index] = cache_size_per_layer_[3]; - curr_index++; - } + bool success = true; + 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]); + if (plan.src_tensors.empty()) { + if (!synchronizer->record_stream(static_cast(range_idx), + stream.get())) { + success = false; + break; } - layer_id++; - layer_cnt++; + continue; } - - ret = aclrtMemcpyBatch(dsts, - copy_size, - srcs, - copy_size, - num_batches * layer_cnt, - attrs, - attrs_indexes, - 1, - &fail_index); - - if (ret != 0 || fail_index != SIZE_MAX) { - LOG(ERROR) << "aclrtMemcpyBatch error: " << ret - << ", fail_index:" << fail_index; - } else { - auto* event = synchronizer->get_event(index); - ret = aclrtRecordEvent(*event, stream->get_stream()->stream()); - if (ret != 0) { - LOG(ERROR) << "aclrtRecordEvent error: " << ret; - } + if (!batch_memcpy_->copy_h2d( + plan.src_tensors, plan.dst_tensors, stream.get())) { + success = false; + break; + } + if (!synchronizer->record_stream(static_cast(range_idx), + stream.get())) { + success = false; + break; } - auto* event_flag = synchronizer->get_event_flag(index); - event_flag->store(true, std::memory_order_release); - if (ret != 0) break; - } - - if (stream->synchronize() != 0) { - LOG(ERROR) << "h2d_batch_copy timeout!"; - copy_stream_.enqueue(std::move(stream)); - return false; } copy_stream_.enqueue(std::move(stream)); - - delete[] dsts; - delete[] srcs; - delete[] copy_size; -#endif - return true; -} - -uint32_t HierarchyKVCacheTransfer::offload_to_store( - Slice& block_transfer_info) { - if (!options_.enable_kvcache_store()) { - return block_transfer_info.size(); - } - - return KVCacheStore::get_instance().batch_put(block_transfer_info); -} - -uint32_t HierarchyKVCacheTransfer::load_from_store( - Slice& block_transfer_info) { - if (!options_.enable_kvcache_store()) { - return 0; + // 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) { + synchronizer->abort(); } - return KVCacheStore::get_instance().batch_get(block_transfer_info); + return success; } -void HierarchyKVCacheTransfer::create_page_aligned_host_cache() { - CHECK(kv_caches_ptr_->size() > 0) << "hbm kv cache size should > 0."; - CHECK(options_.host_blocks_factor() > 1) << "host_blocks_factor should > 1."; - - std::vector> tensor_shapes = - kv_caches_ptr_->at(0).get_shapes(); - tensor_shapes.resize(4); - - CHECK(!tensor_shapes[0].empty()) << "k cache should not be empty!"; - -#if defined(USE_NPU) - int32_t device_id = device_.index(); - h2d_attrs_.dstLoc.id = device_id; - h2d_attrs_.dstLoc.type = aclrtMemLocationType::ACL_MEM_LOCATION_TYPE_DEVICE; - h2d_attrs_.srcLoc.id = device_id; - h2d_attrs_.srcLoc.type = aclrtMemLocationType::ACL_MEM_LOCATION_TYPE_HOST; - memset(h2d_attrs_.rsv, 0, 16); - - d2h_attrs_.dstLoc.id = device_id; - d2h_attrs_.dstLoc.type = aclrtMemLocationType::ACL_MEM_LOCATION_TYPE_HOST; - d2h_attrs_.srcLoc.id = device_id; - d2h_attrs_.srcLoc.type = aclrtMemLocationType::ACL_MEM_LOCATION_TYPE_DEVICE; - memset(d2h_attrs_.rsv, 0, 16); -#endif - - cache_size_per_layer_.resize(4, 0); - cache_size_per_layer_[0] = - kv_caches_ptr_->at(0).get_k_cache()[0].numel() * - kv_caches_ptr_->at(0).get_k_cache()[0].element_size(); - - if (!tensor_shapes[1].empty()) { - cache_size_per_layer_[1] = - kv_caches_ptr_->at(0).get_v_cache()[0].numel() * - kv_caches_ptr_->at(0).get_v_cache()[0].element_size(); - cache_tensor_cnt_++; - } - - if (!tensor_shapes[2].empty()) { - cache_size_per_layer_[2] = - kv_caches_ptr_->at(0).get_index_cache()[0].numel() * - kv_caches_ptr_->at(0).get_index_cache()[0].element_size(); - cache_tensor_cnt_++; - } - - if (!tensor_shapes[3].empty()) { - auto index_cache_scale = kv_caches_ptr_->at(0).get_indexer_cache_scale(); - CHECK(index_cache_scale.has_value()) - << "index cache scale shape exists but tensor is missing."; - cache_size_per_layer_[3] = index_cache_scale.value()[0].numel() * - index_cache_scale.value()[0].element_size(); - cache_tensor_cnt_++; - } - - uint64_t num_blocks = tensor_shapes[0][0] * options_.host_blocks_factor(); - std::vector cache_size_per_tensor(4, 0); - - for (size_t i = 0; i < tensor_shapes.size(); i++) { - if (!tensor_shapes[i].empty()) { - tensor_shapes[i][0] = options_.layers(); - cache_size_per_tensor[i] = cache_size_per_layer_[i] * options_.layers(); - page_aligned_data_size_ += num_blocks * cache_size_per_tensor[i]; - } - } - - size_t page_size = sysconf(_SC_PAGESIZE); - page_aligned_data_size_ = - ((page_aligned_data_size_ + page_size - 1) / page_size) * page_size; - - page_aligned_data_ = mmap(nullptr, - page_aligned_data_size_, - PROT_READ | PROT_WRITE, - MAP_PRIVATE | MAP_ANONYMOUS, - -1, - 0); - - if (page_aligned_data_ == MAP_FAILED) { - LOG(FATAL) << "Failed to allocate aligned memory pool!"; - } - - if (mlock(page_aligned_data_, page_aligned_data_size_) != 0) { - int err = errno; - munmap(page_aligned_data_, page_aligned_data_size_); - uint64_t limit_kb = (page_aligned_data_size_ + 1023) / 1024; - LOG(FATAL) << "Failed to lock memory pool! mlock errno=" << err << " (" - << strerror(err) << "), size=" << page_aligned_data_size_ - << " bytes (~" << (page_aligned_data_size_ >> 20) << " MiB). " - << "Run: ulimit -l " << limit_kb << " (or ulimit -l unlimited) " - << "then restart the process."; - } - -#if defined(USE_NPU) - auto ret = aclrtHostRegister(page_aligned_data_, - page_aligned_data_size_, - aclrtHostRegisterType::ACL_HOST_REGISTER_MAPPED, - &mapped_ptr_); - if (ret != 0) { - LOG(FATAL) << "aclrtHostRegister fail: " << ret; - } -#endif - - size_t current_offset = 0; - auto k_options = torch::TensorOptions() - .dtype(kv_caches_ptr_->at(0).get_k_cache().scalar_type()) - .device(torch::kCPU); - torch::ScalarType v_dtype = - kv_caches_ptr_->at(0).get_v_cache().defined() - ? kv_caches_ptr_->at(0).get_v_cache().scalar_type() - : kv_caches_ptr_->at(0).get_k_cache().scalar_type(); - torch::ScalarType index_dtype = - kv_caches_ptr_->at(0).get_index_cache().defined() - ? kv_caches_ptr_->at(0).get_index_cache().scalar_type() - : kv_caches_ptr_->at(0).get_k_cache().scalar_type(); - auto v_options = torch::TensorOptions().dtype(v_dtype).device(torch::kCPU); - auto index_options = - torch::TensorOptions().dtype(index_dtype).device(torch::kCPU); - std::optional first_index_cache_scale = - kv_caches_ptr_->at(0).get_indexer_cache_scale(); - auto index_scale_options = - torch::TensorOptions() - .dtype(first_index_cache_scale.has_value() - ? first_index_cache_scale->scalar_type() - : kv_caches_ptr_->at(0).get_k_cache().scalar_type()) - .device(torch::kCPU); - host_kv_caches_.reserve(num_blocks); - - for (size_t i = 0; i < num_blocks; ++i) { - torch::Tensor key_cache, value_cache, index_cache; - std::optional index_cache_scale; - void* k_tensor_ptr = - static_cast(page_aligned_data_) + current_offset; - key_cache = torch::from_blob(k_tensor_ptr, tensor_shapes[0], k_options); - current_offset += cache_size_per_tensor[0]; - - if (!tensor_shapes[1].empty()) { - void* v_tensor_ptr = - static_cast(page_aligned_data_) + current_offset; - value_cache = torch::from_blob(v_tensor_ptr, tensor_shapes[1], v_options); - current_offset += cache_size_per_tensor[1]; - } - - if (!tensor_shapes[2].empty()) { - void* index_tensor_ptr = - static_cast(page_aligned_data_) + current_offset; - index_cache = - torch::from_blob(index_tensor_ptr, tensor_shapes[2], index_options); - current_offset += cache_size_per_tensor[2]; - } - - if (!tensor_shapes[3].empty()) { - void* index_scale_tensor_ptr = - static_cast(page_aligned_data_) + current_offset; - index_cache_scale = torch::from_blob( - index_scale_tensor_ptr, tensor_shapes[3], index_scale_options); - current_offset += cache_size_per_tensor[3]; - } - - if (index_cache.defined()) { - host_kv_caches_.emplace_back( - IndexedKVCacheTensors{KVCacheTensors{key_cache, value_cache}, - index_cache, - index_cache_scale}); - continue; - } - - host_kv_caches_.emplace_back(KVCacheTensors{key_cache, value_cache}); - } - - LOG(INFO) << "host k cache shape: " - << host_kv_caches_[0].get_k_cache().sizes(); - LOG(INFO) << "host v cache shape: " - << host_kv_caches_[0].get_v_cache().sizes(); - LOG(INFO) << "host index cache shape: " - << host_kv_caches_[0].get_index_cache().sizes(); - if (host_kv_caches_[0].get_indexer_cache_scale().has_value()) { - LOG(INFO) << "host index cache scale shape: " - << host_kv_caches_[0].get_indexer_cache_scale().value().sizes(); +void HierarchyKVCacheTransfer::set_layer_synchronizer( + ModelInputParams& params) { + 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); } - - LOG(INFO) << "Host block init finish, total size: " << num_blocks; } } // namespace xllm 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 b2a4d555a5..43e7c3303d 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 @@ -15,34 +15,55 @@ limitations under the License. #pragma once -// clang-format off -#if defined(USE_NPU) -#include "acl/acl_rt.h" -#include "platform/npu/npu_layer_synchronizer.h" -#endif -// clang-format on - #include +#include #include +#include +#include +#include +#include "common/macros.h" #include "common/types.h" +#include "framework/block/block.h" #include "framework/kv_cache/kv_cache.h" +#include "framework/kv_cache/kv_cache_utils.h" #include "framework/model/model_input_params.h" +#include "platform/batch_memcpy.h" #include "platform/device.h" +#include "platform/layer_synchronizer.h" #include "util/blockingconcurrentqueue.h" #include "util/threadpool.h" namespace xllm { class HierarchyKVCacheTransfer { public: + struct LayerBatchRange { + int64_t begin_layer = 0; + int64_t end_layer = 0; + }; + + struct CopyPlan { + std::vector src_tensors; + std::vector dst_tensors; + }; + + 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); PROPERTY(uint32_t, layers); PROPERTY(double, host_blocks_factor) = 0.0; PROPERTY(uint32_t, layers_wise_copy_batchs) = 1; + PROPERTY(bool, enable_mla) = false; PROPERTY(bool, enable_kvcache_store) = false; - PROPERTY(std::string, store_protocol) = "tcp"; + PROPERTY(std::string, store_protocol) = "rdma"; PROPERTY(std::string, store_master_server_address) = ""; PROPERTY(std::string, store_metadata_server) = ""; PROPERTY(std::string, store_local_hostname) = ""; @@ -50,8 +71,10 @@ class HierarchyKVCacheTransfer { HierarchyKVCacheTransfer(const Options& options, const torch::Device& device, - std::vector* kv_caches_ptr); - ~HierarchyKVCacheTransfer(); + std::vector* kv_caches_ptr, + const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& create_options); + ~HierarchyKVCacheTransfer() = default; uint32_t transfer_kv_blocks( const uint64_t batch_id, @@ -63,49 +86,38 @@ class HierarchyKVCacheTransfer { void set_layer_synchronizer(ModelInputParams& params); private: - void create_page_aligned_host_cache(); - - uint32_t offload_kv_blocks( + void build_device_block_type_map(); + void create_host_cache(); + CopyPlan build_copy_plan( + 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, const std::vector& block_transfer_info); - bool d2h_batch_copy(Slice& block_transfer_info); - bool h2d_batch_copy(const uint64_t batch_id, - Slice& block_transfer_info); - - uint32_t offload_to_store(Slice& block_transfer_info); - uint32_t load_from_store(Slice& block_transfer_info); - private: - // options Options options_; - // device to run the model on Device device_; - // working thread for data copy - std::unique_ptr h2d_threadpool_; - std::unique_ptr d2h_threadpool_; - // copy streams only can be used in h2d_threadpool_ and d2h_threadpool_ + std::unique_ptr load_threadpool_; moodycamel::BlockingConcurrentQueue> copy_stream_; - std::vector* kv_caches_ptr_; - std::vector host_kv_caches_; - - void* page_aligned_data_ = nullptr; - size_t page_aligned_data_size_ = 0; - - std::vector cache_size_per_layer_; - uint32_t cache_tensor_cnt_ = 1; - -#if defined(USE_NPU) - void* mapped_ptr_ = nullptr; + 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_; - aclrtMemcpyBatchAttr h2d_attrs_; - aclrtMemcpyBatchAttr d2h_attrs_; + std::unique_ptr batch_memcpy_; mutable std::mutex mutex_; - std::unordered_map> + std::unordered_map> layer_wise_load_synchronizer_; -#endif }; } // namespace xllm diff --git a/xllm/core/framework/model/model_input_params.h b/xllm/core/framework/model/model_input_params.h index 1ed4263e2a..97e34b5b0f 100644 --- a/xllm/core/framework/model/model_input_params.h +++ b/xllm/core/framework/model/model_input_params.h @@ -27,6 +27,8 @@ limitations under the License. #include #include "common/types.h" +#include "framework/block/block.h" +#include "platform/layer_synchronizer.h" #if defined(USE_NPU) #include "platform/npu/npu_layer_synchronizer.h" #endif @@ -629,16 +631,19 @@ struct AttentionInput { }; enum class TransferType : uint8_t { - G2H = 0, // global memory(KVCache store) to host memory(DRAM) - H2D = 1, // host memory(DRAM) to device memory(HBM) - D2G = 2, // host memory(DRAM) to global memory(KVCache store) - G2D = 3 // global memory(KVCache store) to device memory(HBM) + G2H = 0, // global memory(KVCache store) to host memory(DRAM) + H2D = 1, // host memory(DRAM) to device memory(HBM) + D2G = 2, // device memory(HBM) to global memory(KVCache store) + G2D = 3, // global memory(KVCache store) to device memory(HBM) + D2H2G = 4, // device memory(HBM) to host memory(DRAM) to global + // memory(KVCache store) }; struct BlockTransferInfo { int32_t src_block_id = -1; int32_t dst_block_id = -1; uint8_t hash_key[XXH3_128BITS_HASH_VALUE_LEN]; + BlockType block_type = BlockType::KV; TransferType transfer_type; BlockTransferInfo(int32_t src_block_id, int32_t dst_block_id) { @@ -649,14 +654,19 @@ struct BlockTransferInfo { BlockTransferInfo(int32_t src_id, int32_t dst_id, const uint8_t* key, - TransferType type) - : src_block_id(src_id), dst_block_id(dst_id), transfer_type(type) { + TransferType type, + BlockType btype = BlockType::KV) + : src_block_id(src_id), + dst_block_id(dst_id), + block_type(btype), + transfer_type(type) { memcpy(hash_key, key, XXH3_128BITS_HASH_VALUE_LEN); } BlockTransferInfo(const BlockTransferInfo& other) : src_block_id(other.src_block_id), dst_block_id(other.dst_block_id), + block_type(other.block_type), transfer_type(other.transfer_type) { memcpy(hash_key, other.hash_key, XXH3_128BITS_HASH_VALUE_LEN); } @@ -664,6 +674,7 @@ struct BlockTransferInfo { BlockTransferInfo(BlockTransferInfo&& other) : src_block_id(other.src_block_id), dst_block_id(other.dst_block_id), + block_type(other.block_type), transfer_type(other.transfer_type) { memcpy(hash_key, other.hash_key, XXH3_128BITS_HASH_VALUE_LEN); @@ -674,6 +685,7 @@ struct BlockTransferInfo { BlockTransferInfo& operator=(const BlockTransferInfo& other) { src_block_id = other.src_block_id; dst_block_id = other.dst_block_id; + block_type = other.block_type; transfer_type = other.transfer_type; memcpy(hash_key, other.hash_key, XXH3_128BITS_HASH_VALUE_LEN); return *this; @@ -682,6 +694,7 @@ struct BlockTransferInfo { BlockTransferInfo& operator=(BlockTransferInfo&& other) { src_block_id = other.src_block_id; dst_block_id = other.dst_block_id; + block_type = other.block_type; transfer_type = other.transfer_type; memcpy(hash_key, other.hash_key, XXH3_128BITS_HASH_VALUE_LEN); @@ -801,9 +814,10 @@ struct ParallelInput { std::shared_ptr layer_synchronizer = nullptr; #elif defined(USE_NPU) std::shared_ptr layer_synchronizer = nullptr; +#endif uint32_t layers_per_bacth_copy = std::numeric_limits::max(); - std::shared_ptr layer_wise_load_synchronizer = - nullptr; + std::shared_ptr layer_wise_load_synchronizer = nullptr; +#if defined(USE_NPU) std::vector query_start_loc; std::vector has_initial_state; #endif @@ -839,9 +853,9 @@ struct ParallelInput { #if defined(USE_NPU) || defined(USE_MLU) || defined(USE_DCU) out.layer_synchronizer = layer_synchronizer; #endif -#if defined(USE_NPU) out.layers_per_bacth_copy = layers_per_bacth_copy; out.layer_wise_load_synchronizer = layer_wise_load_synchronizer; +#if defined(USE_NPU) out.query_start_loc = query_start_loc; out.has_initial_state = has_initial_state; #endif @@ -1016,7 +1030,6 @@ struct ModelInputParams { } bool synchronize_layer(uint32_t layer_idx) const { -#if defined(USE_NPU) if (parallel.layer_wise_load_synchronizer != nullptr && layer_idx % parallel.layers_per_bacth_copy == 0) { if (!parallel.layer_wise_load_synchronizer->synchronize_layer( @@ -1024,9 +1037,6 @@ struct ModelInputParams { return false; } } -#else - (void)layer_idx; -#endif return true; } diff --git a/xllm/core/framework/request/sequence_kv_state.cpp b/xllm/core/framework/request/sequence_kv_state.cpp index 0b0762c730..7f087f6198 100644 --- a/xllm/core/framework/request/sequence_kv_state.cpp +++ b/xllm/core/framework/request/sequence_kv_state.cpp @@ -55,6 +55,24 @@ void KVCacheState::incr_kv_cache_tokens_num(size_t num) { slice_window_pos_ += num; } +void KVCacheState::incr_kv_cache_tokens_num_up_to(size_t new_target) { + const size_t capacity = current_max_tokens_capacity(); + // Drift check: the counter was already advanced past capacity by an earlier + // call. Failing here loudly names the producer instead of silently letting + // the drift propagate to batch_input_builder's CHECK. + CHECK_LE(kv_cache_tokens_num_, capacity) + << "kv_cache_tokens_num_ drifted past capacity: " << kv_cache_tokens_num_ + << " > " << capacity; + CHECK_LE(new_target, capacity) + << "incr_kv_cache_tokens_num_up_to target above capacity: " << new_target + << " > " << capacity; + if (new_target > kv_cache_tokens_num_) { + const size_t delta = new_target - kv_cache_tokens_num_; + kv_cache_tokens_num_ = new_target; + slice_window_pos_ += delta; + } +} + Slice KVCacheState::blocks(BlockType type) const { const auto it = composite_blocks_.find(type); return it == composite_blocks_.end() ? Slice(empty_blocks()) diff --git a/xllm/core/framework/request/sequence_kv_state.h b/xllm/core/framework/request/sequence_kv_state.h index 8f83ba1101..e5ce0a1403 100644 --- a/xllm/core/framework/request/sequence_kv_state.h +++ b/xllm/core/framework/request/sequence_kv_state.h @@ -34,6 +34,13 @@ class KVCacheState { size_t kv_cache_tokens_num() const; void set_kv_cache_tokens_num(size_t num); void incr_kv_cache_tokens_num(size_t num); + // Advance kv_cache_tokens_num_ to `new_target` if it grows the counter. No-op + // when `new_target <= kv_cache_tokens_num_`. CHECK-fails when `new_target` + // exceeds current_max_tokens_capacity() -- callers must clamp before calling. + // CHECK-fails when the counter is already past capacity (drift detection). + // Used by the host-cache H2D restore path in HierarchyBlockManagerPool where + // the "up to" semantics naturally arise from mismatched host/device counters. + void incr_kv_cache_tokens_num_up_to(size_t new_target); // Blocks held under `type`; empty slice when the type is absent. Slice blocks(BlockType type) const; diff --git a/xllm/core/platform/CMakeLists.txt b/xllm/core/platform/CMakeLists.txt index 5efcbe4a66..fea84140ba 100644 --- a/xllm/core/platform/CMakeLists.txt +++ b/xllm/core/platform/CMakeLists.txt @@ -27,6 +27,8 @@ cc_library( shared_vmm_allocator.h vmm_torch_allocator.h sleepable_allocator.h + batch_memcpy.h + layer_synchronizer.h $<$:mlu/mlu_layer_synchronizer.h> $<$:mlu/mlu_tensor_alloc.h> $<$:dcu/dcu_layer_synchronizer.h> @@ -43,6 +45,8 @@ cc_library( vmm_api.cpp shared_vmm_allocator.cpp sleepable_allocator.cpp + batch_memcpy.cpp + layer_synchronizer.cpp $<$:mlu/mlu_layer_synchronizer.cpp> $<$:mlu/mlu_tensor_alloc.cpp> $<$:cuda_profiler.cpp> diff --git a/xllm/core/platform/batch_memcpy.cpp b/xllm/core/platform/batch_memcpy.cpp new file mode 100644 index 0000000000..7c75e595f8 --- /dev/null +++ b/xllm/core/platform/batch_memcpy.cpp @@ -0,0 +1,37 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "platform/batch_memcpy.h" + +#include "platform/device.h" + +#if defined(USE_NPU) +#include "platform/npu/npu_batch_memcpy.h" +#endif + +namespace xllm { + +std::unique_ptr create_batch_memcpy(const Device& device) { +#if defined(USE_NPU) + auto memcpy = std::make_unique(); + memcpy->init(device.index()); + return memcpy; +#else + (void)device; + return nullptr; +#endif +} + +} // namespace xllm diff --git a/xllm/core/platform/batch_memcpy.h b/xllm/core/platform/batch_memcpy.h new file mode 100644 index 0000000000..5ef1b13543 --- /dev/null +++ b/xllm/core/platform/batch_memcpy.h @@ -0,0 +1,47 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include + +#include "platform/stream.h" + +namespace xllm { + +class Device; + +class BatchMemcpy { + public: + virtual ~BatchMemcpy() = default; + + virtual void init(int32_t device_id) = 0; + + virtual bool copy_h2d(const std::vector& src_tensors, + const std::vector& dst_tensors, + Stream* stream) = 0; + + virtual bool copy_d2h(const std::vector& src_tensors, + const std::vector& dst_tensors, + Stream* stream) = 0; +}; + +std::unique_ptr create_batch_memcpy(const Device& device); + +} // namespace xllm diff --git a/xllm/core/platform/layer_synchronizer.cpp b/xllm/core/platform/layer_synchronizer.cpp new file mode 100644 index 0000000000..fa36482d8d --- /dev/null +++ b/xllm/core/platform/layer_synchronizer.cpp @@ -0,0 +1,73 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "platform/layer_synchronizer.h" + +#if defined(USE_NPU) +#include "platform/npu/npu_layer_synchronizer.h" +#endif + +namespace xllm { + +#if defined(USE_NPU) + +class NPULayerSynchronizerAdapter final : public LayerSynchronizer { + public: + explicit NPULayerSynchronizerAdapter(int64_t num_layers) + : impl_(num_layers) {} + + bool synchronize_layer(int64_t layer_index) override { + return impl_.synchronize_layer(layer_index); + } + + bool record_stream(int64_t layer_index, Stream* stream) override { + if (stream == nullptr) { + return false; + } + const aclError ret = aclrtRecordEvent(*impl_.get_event(layer_index), + stream->get_stream()->stream()); + if (ret != ACL_SUCCESS) { + return false; + } + impl_.get_event_flag(layer_index)->store(true, std::memory_order_release); + return true; + } + + void abort() override { impl_.abort(); } + + uint32_t size() const override { + return const_cast(impl_).get_event_size(); + } + + private: + NPULayerSynchronizerImpl impl_; +}; + +std::shared_ptr create_layer_synchronizer( + int64_t num_layers) { + return std::make_shared(num_layers); +} + +#else + +std::shared_ptr create_layer_synchronizer( + int64_t num_layers) { + (void)num_layers; + return nullptr; +} + +#endif + +} // namespace xllm diff --git a/xllm/core/platform/layer_synchronizer.h b/xllm/core/platform/layer_synchronizer.h new file mode 100644 index 0000000000..199dd074a1 --- /dev/null +++ b/xllm/core/platform/layer_synchronizer.h @@ -0,0 +1,42 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +#include "platform/stream.h" + +namespace xllm { + +class LayerSynchronizer { + public: + virtual ~LayerSynchronizer() = default; + + virtual bool synchronize_layer(int64_t layer_index) = 0; + virtual bool record_stream(int64_t layer_index, Stream* stream) = 0; + // Force every layer's wait to unblock and report failure. Called when a copy + // fails so a forward thread spinning in synchronize_layer does not hang + // forever; synchronize_layer returns false after abort so the caller aborts + // the forward instead of reading not-yet-copied KV cache. + virtual void abort() = 0; + virtual uint32_t size() const = 0; +}; + +std::shared_ptr create_layer_synchronizer( + int64_t num_layers); + +} // namespace xllm diff --git a/xllm/core/platform/npu/CMakeLists.txt b/xllm/core/platform/npu/CMakeLists.txt index da8a1990eb..cf901f57f1 100644 --- a/xllm/core/platform/npu/CMakeLists.txt +++ b/xllm/core/platform/npu/CMakeLists.txt @@ -5,10 +5,12 @@ cc_library( platform_npu HDRS npu_layer_synchronizer.h + npu_batch_memcpy.h device_capture_lock.h acl_graph_task_update_context.h SRCS npu_layer_synchronizer.cpp + npu_batch_memcpy.cpp DEPS torch_npu glog::glog diff --git a/xllm/core/platform/npu/npu_batch_memcpy.cpp b/xllm/core/platform/npu/npu_batch_memcpy.cpp new file mode 100644 index 0000000000..b78145345a --- /dev/null +++ b/xllm/core/platform/npu/npu_batch_memcpy.cpp @@ -0,0 +1,127 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "platform/npu/npu_batch_memcpy.h" + +#include + +#include + +namespace xllm { +namespace npu { + +void NPUBatchMemcpy::init(int32_t device_id) { + if (initialized_) { + return; + } + device_id_ = device_id; + + memset(&h2d_attr_, 0, sizeof(h2d_attr_)); + h2d_attr_.dstLoc.id = device_id; + h2d_attr_.dstLoc.type = aclrtMemLocationType::ACL_MEM_LOCATION_TYPE_DEVICE; + h2d_attr_.srcLoc.id = device_id; + h2d_attr_.srcLoc.type = aclrtMemLocationType::ACL_MEM_LOCATION_TYPE_HOST; + + memset(&d2h_attr_, 0, sizeof(d2h_attr_)); + d2h_attr_.dstLoc.id = device_id; + d2h_attr_.dstLoc.type = aclrtMemLocationType::ACL_MEM_LOCATION_TYPE_HOST; + d2h_attr_.srcLoc.id = device_id; + d2h_attr_.srcLoc.type = aclrtMemLocationType::ACL_MEM_LOCATION_TYPE_DEVICE; + + initialized_ = true; +} + +bool NPUBatchMemcpy::copy_h2d(const std::vector& src_tensors, + const std::vector& dst_tensors, + Stream* stream) { + return copy(src_tensors, dst_tensors, h2d_attr_, stream); +} + +bool NPUBatchMemcpy::copy_d2h(const std::vector& src_tensors, + const std::vector& dst_tensors, + Stream* stream) { + return copy(src_tensors, dst_tensors, d2h_attr_, stream); +} + +bool NPUBatchMemcpy::copy(const std::vector& src_tensors, + const std::vector& dst_tensors, + const aclrtMemcpyBatchAttr& attr, + Stream* stream) { + CHECK(initialized_) << "NPUBatchMemcpy not initialized."; + CHECK(stream != nullptr) << "Stream must not be null."; + CHECK_EQ(src_tensors.size(), dst_tensors.size()) + << "src and dst tensor count mismatch."; + + if (src_tensors.empty()) { + return true; + } + + const size_t count = src_tensors.size(); + + c10::StreamGuard guard = stream->set_stream_guard(); + + aclrtMemcpyBatchAttr attrs[1] = {attr}; + size_t attrs_indexes[1] = {0}; + + // aclrtMemcpyBatch caps a single call at kMaxBatchCopyCount descriptors, so a + // large transfer (e.g. a many-block prefix reload) is split into fixed-size + // chunks issued onto the same stream in order. The stream is synchronized + // once after all chunks are enqueued rather than per chunk. + std::vector srcs; + std::vector dsts; + std::vector sizes; + srcs.reserve(std::min(count, kMaxBatchCopyCount)); + dsts.reserve(std::min(count, kMaxBatchCopyCount)); + sizes.reserve(std::min(count, kMaxBatchCopyCount)); + + for (size_t offset = 0; offset < count; offset += kMaxBatchCopyCount) { + const size_t chunk = std::min(kMaxBatchCopyCount, count - offset); + srcs.clear(); + dsts.clear(); + sizes.clear(); + for (size_t i = 0; i < chunk; ++i) { + srcs.emplace_back(src_tensors[offset + i].data_ptr()); + dsts.emplace_back(dst_tensors[offset + i].data_ptr()); + sizes.emplace_back(static_cast(src_tensors[offset + i].nbytes())); + } + + size_t fail_index = SIZE_MAX; + aclError ret = aclrtMemcpyBatch(dsts.data(), + sizes.data(), + srcs.data(), + sizes.data(), + chunk, + attrs, + attrs_indexes, + /*attrsCount=*/1, + &fail_index); + if (ret != ACL_SUCCESS || fail_index != SIZE_MAX) { + // fail_index is chunk-local; report the position within the full batch. + LOG(ERROR) << "aclrtMemcpyBatch failed: ret=" << ret << ", fail_index=" + << (fail_index == SIZE_MAX ? fail_index : offset + fail_index); + return false; + } + } + + if (stream->synchronize() != 0) { + LOG(ERROR) << "BatchMemcpy stream synchronize timeout."; + return false; + } + + return true; +} + +} // namespace npu +} // namespace xllm diff --git a/xllm/core/platform/npu/npu_batch_memcpy.h b/xllm/core/platform/npu/npu_batch_memcpy.h new file mode 100644 index 0000000000..3bdc9e05dd --- /dev/null +++ b/xllm/core/platform/npu/npu_batch_memcpy.h @@ -0,0 +1,64 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +#include +#include +#include + +#include "common/macros.h" +#include "platform/batch_memcpy.h" +#include "platform/stream.h" + +namespace xllm { +namespace npu { + +class NPUBatchMemcpy final : public BatchMemcpy { + public: + NPUBatchMemcpy() = default; + ~NPUBatchMemcpy() override = default; + + void init(int32_t device_id) override; + + bool copy_h2d(const std::vector& src_tensors, + const std::vector& dst_tensors, + Stream* stream) override; + + bool copy_d2h(const std::vector& src_tensors, + const std::vector& dst_tensors, + Stream* stream) override; + + private: + static constexpr size_t kMaxBatchCopyCount = 4096; + + DISALLOW_COPY_AND_ASSIGN(NPUBatchMemcpy); + + bool copy(const std::vector& src_tensors, + const std::vector& dst_tensors, + const aclrtMemcpyBatchAttr& attr, + Stream* stream); + + bool initialized_ = false; + int32_t device_id_ = -1; + aclrtMemcpyBatchAttr h2d_attr_{}; + aclrtMemcpyBatchAttr d2h_attr_{}; +}; + +} // namespace npu +} // namespace xllm diff --git a/xllm/core/platform/npu/npu_layer_synchronizer.cpp b/xllm/core/platform/npu/npu_layer_synchronizer.cpp index 33c4f6144a..95a897684f 100644 --- a/xllm/core/platform/npu/npu_layer_synchronizer.cpp +++ b/xllm/core/platform/npu/npu_layer_synchronizer.cpp @@ -48,7 +48,14 @@ std::atomic* NPULayerSynchronizerImpl::get_event_flag( } bool NPULayerSynchronizerImpl::synchronize_layer(const int64_t layer_index) { - while (!event_record_flags_[layer_index].load(std::memory_order_acquire)); + while (!event_record_flags_[layer_index].load(std::memory_order_acquire)) { + if (aborted_.load(std::memory_order_acquire)) { + return false; + } + } + if (aborted_.load(std::memory_order_acquire)) { + return false; + } auto ret = aclrtSynchronizeEventWithTimeout(events_[layer_index], timeout_); if (ret != ACL_SUCCESS) { LOG(ERROR) << "Synchronize event failed: " << ret; @@ -69,4 +76,13 @@ bool NPULayerSynchronizerImpl::record_event(const int64_t layer_index, return true; } +void NPULayerSynchronizerImpl::abort() { + // Set aborted_ before raising the flags so any thread released from the spin + // observes the abort and reports failure rather than syncing a stale event. + aborted_.store(true, std::memory_order_release); + for (size_t i = 0; i < event_record_flags_.size(); ++i) { + event_record_flags_[i].store(true, std::memory_order_release); + } +} + } // namespace xllm diff --git a/xllm/core/platform/npu/npu_layer_synchronizer.h b/xllm/core/platform/npu/npu_layer_synchronizer.h index d9d6914097..776f4baaaa 100644 --- a/xllm/core/platform/npu/npu_layer_synchronizer.h +++ b/xllm/core/platform/npu/npu_layer_synchronizer.h @@ -32,11 +32,14 @@ class NPULayerSynchronizerImpl { std::atomic* get_event_flag(const int64_t layer_index); bool synchronize_layer(const int64_t layer_index); bool record_event(const int64_t layer_index, const int32_t device_index); + // Unblock all pending synchronize_layer spins and make them report failure. + void abort(); uint32_t get_event_size() { return events_.size(); }; private: std::vector events_; std::vector> event_record_flags_; + std::atomic aborted_{false}; const int32_t timeout_; }; diff --git a/xllm/core/runtime/params_utils.cpp b/xllm/core/runtime/params_utils.cpp index d7b4737124..8b90f54a05 100644 --- a/xllm/core/runtime/params_utils.cpp +++ b/xllm/core/runtime/params_utils.cpp @@ -277,7 +277,9 @@ uint64_t proto_to_block_transfer_info( pb_block_transfer_info.transfer_infos(i).dst_block_id(), reinterpret_cast( pb_block_transfer_info.transfer_infos(i).hash_key().data()), - TransferType(pb_block_transfer_info.transfer_type())); + TransferType(pb_block_transfer_info.transfer_type()), + static_cast( + pb_block_transfer_info.transfer_infos(i).block_type())); } return pb_block_transfer_info.batch_id(); @@ -302,6 +304,8 @@ bool block_transfer_info_to_proto( pb_cache.set_src_block_id(info.src_block_id); pb_cache.set_dst_block_id(info.dst_block_id); pb_cache.set_hash_key(info.hash_key, XXH3_128BITS_HASH_VALUE_LEN); + pb_cache.set_block_type( + static_cast(static_cast(info.block_type))); *pb_block_transfer_info->mutable_transfer_infos()->Add() = std::move(pb_cache); diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 6390495833..5df9559e25 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -395,6 +395,7 @@ bool WorkerImpl::allocate_kv_cache_storage(const KVCacheShape& kv_cache_shape, // 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); + init_hierarchy_kv_cache_transfer(kv_cache_shape, create_options); #if defined(USE_CUDA) || defined(USE_DCU) refresh_cuda_block_copy_runtime_state(); @@ -408,8 +409,6 @@ bool WorkerImpl::allocate_kv_cache(const KVCacheShape& kv_cache_shape) { return false; } - // hierarchy temporarily disabled during the block-manager refactor - // init_hierarchy_kv_cache_transfer(); status_ = Status::READY; return true; } @@ -439,9 +438,6 @@ bool WorkerImpl::allocate_kv_cache_with_transfer( context_.get_model_args().model_type(), options_.model_id()); - // hierarchy temporarily disabled during the block-manager refactor - // init_hierarchy_kv_cache_transfer(); - status_ = Status::READY; return true; } @@ -472,8 +468,6 @@ bool WorkerImpl::allocate_kv_cache_with_transfer( kv_cache_transfer_->register_kv_cache(kv_caches_, kv_cache_shape, dtype_); #endif - // hierarchy temporarily disabled during the block-manager refactor - // init_hierarchy_kv_cache_transfer(); status_ = Status::READY; return true; } @@ -1167,10 +1161,9 @@ folly::SemiFuture> WorkerImpl::step_async( threadpool_.schedule([this, input = std::move(input_on_device), promise = std::move(promise)]() mutable { - // hierarchy temporarily disabled during the block-manager refactor - // if (hierarchy_kv_cache_transfer_ != nullptr) { - // hierarchy_kv_cache_transfer_->set_layer_synchronizer(input.input_params); - // } + if (hierarchy_kv_cache_transfer_ != nullptr) { + hierarchy_kv_cache_transfer_->set_layer_synchronizer(input.input_params); + } // run the model on the given input in working thread if (!enable_schedule_overlap()) { @@ -1760,20 +1753,22 @@ folly::SemiFuture WorkerImpl::pull_kv_blocks_async( } uint32_t WorkerImpl::transfer_kv_blocks( - const uint64_t /*batch_id*/, - const std::vector& /*block_transfer_info*/) { - // hierarchy temporarily disabled during the block-manager refactor. - LOG(FATAL) << "hierarchy kv cache transfer is disabled during the " - "block-manager refactor."; + const uint64_t batch_id, + const std::vector& block_transfer_info) { + if (hierarchy_kv_cache_transfer_ != nullptr) { + return hierarchy_kv_cache_transfer_->transfer_kv_blocks( + batch_id, block_transfer_info); + } return 0; } uint32_t WorkerImpl::transfer_kv_blocks( - const uint64_t /*batch_id*/, - Slice& /*block_transfer_info*/) { - // hierarchy temporarily disabled during the block-manager refactor. - LOG(FATAL) << "hierarchy kv cache transfer is disabled during the " - "block-manager refactor."; + const uint64_t batch_id, + Slice& block_transfer_info) { + if (hierarchy_kv_cache_transfer_ != nullptr) { + return hierarchy_kv_cache_transfer_->transfer_kv_blocks( + batch_id, block_transfer_info); + } return 0; } @@ -1804,6 +1799,32 @@ int64_t WorkerImpl::get_active_activation_memory() { // transfer_options, device_, &kv_caches_); // } // } +void WorkerImpl::init_hierarchy_kv_cache_transfer( + const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& kv_cache_create_options) { + if (options_.host_blocks_factor() > 1.0) { + CHECK(!kv_caches_.empty()) << "kv_caches is not initialized."; + CHECK(hierarchy_kv_cache_transfer_ == nullptr) + << "Hierarchy KV cache transfer is already initialized."; + HierarchyKVCacheTransfer::Options transfer_options; + transfer_options + .tp_rank(options_.dp_size() > 1 + ? options_.node_rank() % options_.dp_size() + : options_.node_rank()) + .tp_size(options_.world_size() / options_.dp_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(false); + hierarchy_kv_cache_transfer_ = + std::make_unique(transfer_options, + device_, + &kv_caches_, + kv_cache_shape, + kv_cache_create_options); + } +} void WorkerImpl::prepare_mla_prefixcache_inputs( ModelInputParams& input_params) { const bool has_prefixcache_metadata = diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h index 50fae27196..62f8b97266 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -26,8 +26,7 @@ limitations under the License. #include "forward_params.h" #include "framework/eplb/eplb_executor.h" #include "framework/kv_cache/kv_cache_shape.h" -// hierarchy temporarily disabled during the block-manager refactor -// #include "framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h" +#include "framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h" #include "framework/kv_cache_transfer/kv_cache_store.h" #include "framework/kv_cache_transfer/kv_cache_transfer.h" #include "framework/model/causal_lm.h" @@ -214,8 +213,9 @@ class WorkerImpl { // Only used for deepseek chunked prefill ops on npu device void prepare_mla_prefixcache_inputs(ModelInputParams& input_params); - // hierarchy temporarily disabled during the block-manager refactor - // void init_hierarchy_kv_cache_transfer(); + void init_hierarchy_kv_cache_transfer( + const KVCacheShape& kv_cache_shape, + const KVCacheCreateOptions& kv_cache_create_options); bool can_prepare_npu_graph_decode_input( const ModelInputParams& input_params) const; @@ -335,8 +335,7 @@ class WorkerImpl { InstanceRole instance_role_ = InstanceRole::DEFAULT; std::shared_ptr kv_cache_transfer_; - // hierarchy temporarily disabled during the block-manager refactor - // std::unique_ptr hierarchy_kv_cache_transfer_; + std::unique_ptr hierarchy_kv_cache_transfer_; std::unique_ptr worker_rendezvous_; #if defined(USE_CUDA) || defined(USE_DCU) diff --git a/xllm/core/scheduler/chunked_prefill_scheduler.cpp b/xllm/core/scheduler/chunked_prefill_scheduler.cpp index 8e8d30acea..cead04ccb3 100644 --- a/xllm/core/scheduler/chunked_prefill_scheduler.cpp +++ b/xllm/core/scheduler/chunked_prefill_scheduler.cpp @@ -813,6 +813,14 @@ std::vector ChunkedPrefillScheduler::prepare_batch() { if (!is_batches_empty) { // only update the scheduling latency when there are requests to process COUNTER_ADD(scheduling_latency_seconds, timer.elapsed_seconds()); + // Dispatch the per-step KV transfers collected during allocate/deallocate: + // H2D loads recorded for host prefix-cache hits, and D2H offloads queued + // when sequences were deallocated. HierarchyBlockManagerPool overrides + // these; the default BlockManagerPool no-ops. Without this call the offload + // queue is never drained and its device/host blocks leak. + kv_cache_manager_->transfer_blocks(batches); + } else { + kv_cache_manager_->transfer_blocks(); } GAUGE_SET(num_pending_requests, diff --git a/xllm/core/scheduler/fixed_steps_scheduler.cpp b/xllm/core/scheduler/fixed_steps_scheduler.cpp index 0740e312c2..b1e2134e04 100644 --- a/xllm/core/scheduler/fixed_steps_scheduler.cpp +++ b/xllm/core/scheduler/fixed_steps_scheduler.cpp @@ -297,6 +297,9 @@ std::vector FixedStepsScheduler::prepare_batch() { if (!batches[0].empty()) { // only update the scheduling latency when there are requests to process COUNTER_ADD(scheduling_latency_seconds, timer.elapsed_seconds()); + kv_cache_manager_->transfer_blocks(batches); + } else { + kv_cache_manager_->transfer_blocks(); } GAUGE_SET(num_pending_requests, diff --git a/xllm/core/scheduler/pd_ooc_scheduler.cpp b/xllm/core/scheduler/pd_ooc_scheduler.cpp index a793b1a630..63d6b18f82 100644 --- a/xllm/core/scheduler/pd_ooc_scheduler.cpp +++ b/xllm/core/scheduler/pd_ooc_scheduler.cpp @@ -400,6 +400,9 @@ std::vector PDOOCScheduler::prepare_batch() { if (!is_batches_empty) { // only update the scheduling latency when there are requests to process COUNTER_ADD(scheduling_latency_seconds, timer.elapsed_seconds()); + kv_cache_manager_->transfer_blocks(batches); + } else { + kv_cache_manager_->transfer_blocks(); } GAUGE_SET(num_pending_requests, diff --git a/xllm/core/scheduler/prefill_only_scheduler.cpp b/xllm/core/scheduler/prefill_only_scheduler.cpp index 537bf57cb6..25f7cb19f9 100644 --- a/xllm/core/scheduler/prefill_only_scheduler.cpp +++ b/xllm/core/scheduler/prefill_only_scheduler.cpp @@ -696,6 +696,9 @@ std::vector PrefillOnlyScheduler::prepare_batch() { if (!is_batches_empty) { // only update the scheduling latency when there are requests to process COUNTER_ADD(scheduling_latency_seconds, timer.elapsed_seconds()); + kv_cache_manager_->transfer_blocks(batches); + } else { + kv_cache_manager_->transfer_blocks(); } GAUGE_SET(num_pending_requests, diff --git a/xllm/proto/worker.proto b/xllm/proto/worker.proto index 9f5d800e96..1dea2df059 100644 --- a/xllm/proto/worker.proto +++ b/xllm/proto/worker.proto @@ -105,12 +105,26 @@ enum TransferType { G2H = 0; H2D = 1; D2G = 2; + G2D = 3; + D2H2G = 4; +} + +// Mirrors xllm::BlockType (framework/block/block.h). Kept in wire order so +// numeric values round-trip losslessly across RPC. +enum BlockType { + KV = 0; + SWA = 1; + C4 = 2; + C128 = 3; + SINGLE = 4; + LINEAR = 5; } message BlockTransferInfo { int32 src_block_id = 1; int32 dst_block_id = 2; bytes hash_key = 3; + BlockType block_type = 4; } message BlockTransferInfos {