diff --git a/tests/core/framework/block/block_manager_test.cpp b/tests/core/framework/block/block_manager_test.cpp index bd9b4ca14b..d743be005e 100644 --- a/tests/core/framework/block/block_manager_test.cpp +++ b/tests/core/framework/block/block_manager_test.cpp @@ -469,4 +469,178 @@ TEST(BlockManagerPoolTest, EXPECT_EQ(pool->num_blocks_in_prefix_cache()[0], 1u); } +// prefix_cache_match_length() reports the leading prefix-hit length in blocks +// and must leave the cache (size, LRU, ref counts -> utilization) untouched. +TEST(BlockManagerTest, PrefixCacheMatchLengthIsReadOnly) { + BlockManager::Options options; + options.num_blocks(16).block_size(2).enable_prefix_cache(true); + // Leak intentionally: the prefix cache keeps the cached blocks referenced at + // teardown, which would otherwise trip the free-list check in + // ~BlockManagerImpl. + auto* manager = new BlockManagerImpl(options); + + // Three full blocks: [1,2] [3,4] [5,6]. + const std::vector tokens = {1, 2, 3, 4, 5, 6}; + std::vector blocks = manager->allocate(/*num_blocks=*/3); + ASSERT_EQ(blocks.size(), 3u); + manager->cache( + Slice(tokens), blocks, /*existed_shared_blocks_num=*/0); + ASSERT_EQ(manager->num_blocks_in_prefix_cache(), 3u); + + const size_t cached_before = manager->num_blocks_in_prefix_cache(); + const double util_before = manager->kv_cache_utilization(); + + // Full hit, and repeated probing is idempotent. + const size_t matched1 = + manager->prefix_cache_match_length(Slice(tokens)); + const size_t matched2 = + manager->prefix_cache_match_length(Slice(tokens)); + EXPECT_EQ(matched1, 3u); + EXPECT_EQ(matched2, matched1); + + // A shorter prefix matches only its leading full blocks. + const std::vector short_tokens = {1, 2, 3, 4}; + EXPECT_EQ(manager->prefix_cache_match_length(Slice(short_tokens)), + 2u); + + // A diverging prefix matches only the shared leading block. + const std::vector diverged = {1, 2, 9, 9}; + EXPECT_EQ(manager->prefix_cache_match_length(Slice(diverged)), 1u); + + // Probing mutated nothing. + EXPECT_EQ(manager->num_blocks_in_prefix_cache(), cached_before); + EXPECT_DOUBLE_EQ(manager->kv_cache_utilization(), util_before); +} + +// With cache-aware routing on, a request sharing a cached prefix is routed to +// the rank holding that prefix even when another rank has more free blocks. +TEST(BlockManagerPoolTest, CacheAwareDpRoutingPrefersLongestPrefixRank) { + ScopedValue max_seqs_guard( + &SchedulerConfig::get_instance().max_seqs_per_batch(), 2); + + BlockManagerPool::Options options; + options.num_blocks(32) + .host_num_blocks(0) + .block_size(2) + .enable_prefix_cache(true) + .enable_cache_aware_dp(true) + .cache_aware_match_threshold(0.5) + // Disable the imbalance guard so this test isolates affinity behavior. + .cache_aware_imbalance_threshold(1.0); + // Leak: prefix cache holds block refs at teardown (see test above). + auto* pool = new BlockManagerPool(options, /*dp_size=*/2); + + // First allocation ties on free blocks across ranks, so it lands on rank 0; + // publishing its blocks seeds rank 0's prefix cache. + Sequence seq_a = MakeSequence(0, /*prompt_tokens=*/{1, 2, 3, 4, 5, 6}); + ASSERT_TRUE(pool->allocate(&seq_a)); + ASSERT_EQ(seq_a.dp_rank(), 0); + pool->cache(&seq_a, /*num_tokens=*/6); + ASSERT_GT(pool->num_blocks_in_prefix_cache()[0], 0u); + + const std::vector cached_before = pool->num_blocks_in_prefix_cache(); + + // Rank 1 has more free blocks now, but the shared prefix wins: route to 0. + Sequence seq_b = MakeSequence(1, /*prompt_tokens=*/{1, 2, 3, 4, 5, 6}); + pool->allocate_shared(&seq_b); + EXPECT_EQ(seq_b.dp_rank(), 0); + + // Routing only probes the prefix cache; it must not change cache contents. + EXPECT_EQ(pool->num_blocks_in_prefix_cache(), cached_before); + + // A request with no cached prefix falls back to the rank with more free + // blocks (rank 1, since rank 0 still holds seq_a's blocks). + Sequence seq_c = MakeSequence(2, /*prompt_tokens=*/{7, 8, 9, 10}); + pool->allocate_shared(&seq_c); + EXPECT_EQ(seq_c.dp_rank(), 1); +} + +// Cold start: a rank holding a tiny shared prefix (one block out of a long +// request) must NOT attract the request. The match-fraction threshold keeps +// routing load-balanced until a meaningful prefix accumulates. +TEST(BlockManagerPoolTest, CacheAwareDpRoutingIgnoresTinyPrefixOnColdStart) { + ScopedValue max_seqs_guard( + &SchedulerConfig::get_instance().max_seqs_per_batch(), 2); + + BlockManagerPool::Options options; + options.num_blocks(64) + .host_num_blocks(0) + .block_size(2) + .enable_prefix_cache(true) + .enable_cache_aware_dp(true) + .cache_aware_match_threshold(0.5) + // Isolate the match-fraction guard from the imbalance guard. + .cache_aware_imbalance_threshold(1.0); + auto* pool = new BlockManagerPool(options, /*dp_size=*/2); + + // Seed rank 0 with a single cached block [1,2]. + Sequence seq_a = MakeSequence(0, /*prompt_tokens=*/{1, 2}); + ASSERT_TRUE(pool->allocate(&seq_a)); + ASSERT_EQ(seq_a.dp_rank(), 0); + pool->cache(&seq_a, /*num_tokens=*/2); + ASSERT_GT(pool->num_blocks_in_prefix_cache()[0], 0u); + + // A long request shares only that one leading block (1/5 = 20% < 50%), so it + // is routed by free blocks (rank 1) instead of being herded to rank 0. + Sequence seq_b = + MakeSequence(1, /*prompt_tokens=*/{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}); + pool->allocate_shared(&seq_b); + EXPECT_EQ(seq_b.dp_rank(), 1); +} + +// A popular prefix that fully matches still yields to load balancing once the +// holding rank becomes too loaded relative to the emptiest rank. +TEST(BlockManagerPoolTest, CacheAwareDpRoutingRebalancesWhenOverloaded) { + ScopedValue max_seqs_guard( + &SchedulerConfig::get_instance().max_seqs_per_batch(), 2); + + BlockManagerPool::Options options; + options.num_blocks(32) + .host_num_blocks(0) + .block_size(2) + .enable_prefix_cache(true) + .enable_cache_aware_dp(true) + .cache_aware_match_threshold(0.5) + // Tight imbalance budget: rank 0 holding a few blocks already trips it. + .cache_aware_imbalance_threshold(0.05); + auto* pool = new BlockManagerPool(options, /*dp_size=*/2); + + Sequence seq_a = MakeSequence(0, /*prompt_tokens=*/{1, 2, 3, 4, 5, 6}); + ASSERT_TRUE(pool->allocate(&seq_a)); + ASSERT_EQ(seq_a.dp_rank(), 0); + pool->cache(&seq_a, /*num_tokens=*/6); + + // Same full prefix, but rank 0 is now too loaded vs rank 1, so the imbalance + // guard routes the request to the least-loaded rank (rank 1). + Sequence seq_b = MakeSequence(1, /*prompt_tokens=*/{1, 2, 3, 4, 5, 6}); + pool->allocate_shared(&seq_b); + EXPECT_EQ(seq_b.dp_rank(), 1); +} + +// With cache-aware routing off, routing ignores the prefix and follows free +// blocks, so the same shared-prefix request lands on the emptier rank. +TEST(BlockManagerPoolTest, CacheAwareDpRoutingDisabledUsesMaxFreeRank) { + ScopedValue max_seqs_guard( + &SchedulerConfig::get_instance().max_seqs_per_batch(), 2); + + BlockManagerPool::Options options; + options.num_blocks(32) + .host_num_blocks(0) + .block_size(2) + .enable_prefix_cache(true) + .enable_cache_aware_dp(false); + auto* pool = new BlockManagerPool(options, /*dp_size=*/2); + + Sequence seq_a = MakeSequence(0, /*prompt_tokens=*/{1, 2, 3, 4, 5, 6}); + ASSERT_TRUE(pool->allocate(&seq_a)); + ASSERT_EQ(seq_a.dp_rank(), 0); + pool->cache(&seq_a, /*num_tokens=*/6); + + // Same prefix, but routing ignores it and picks the rank with more free + // blocks (rank 1). + Sequence seq_b = MakeSequence(1, /*prompt_tokens=*/{1, 2, 3, 4, 5, 6}); + pool->allocate_shared(&seq_b); + EXPECT_EQ(seq_b.dp_rank(), 1); +} + } // namespace xllm diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index a560f85c1c..92912256a0 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -58,6 +58,12 @@ DECLARE_bool(enable_prefix_cache); DECLARE_bool(enable_in_batch_prefix_cache); +DECLARE_bool(enable_prefix_cache_aware_dp_routing); + +DECLARE_double(prefix_cache_aware_dp_match_threshold); + +DECLARE_double(prefix_cache_aware_dp_imbalance_threshold); + DECLARE_int64(max_encoder_cache_size); DECLARE_uint32(xxh3_128bits_seed); diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index afe50b826d..34b11442ba 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -509,6 +509,13 @@ bool LLMEngine::allocate_kv_cache(const KVCacheCapacity& kv_cache_cap) { ::xllm::KVCacheConfig::get_instance().enable_xtensor() ? false : options_.enable_prefix_cache()) + .enable_cache_aware_dp(::xllm::KVCacheConfig::get_instance() + .enable_prefix_cache_aware_dp_routing()) + .cache_aware_match_threshold(::xllm::KVCacheConfig::get_instance() + .prefix_cache_aware_dp_match_threshold()) + .cache_aware_imbalance_threshold( + ::xllm::KVCacheConfig::get_instance() + .prefix_cache_aware_dp_imbalance_threshold()) .enable_disagg_pd(options_.enable_disagg_pd()) .enable_kvcache_store(options_.enable_kvcache_store()) .enable_xtensor(::xllm::KVCacheConfig::get_instance().enable_xtensor()) diff --git a/xllm/core/distributed_runtime/vlm_engine.cpp b/xllm/core/distributed_runtime/vlm_engine.cpp index b5819efea0..3945f7ff2b 100644 --- a/xllm/core/distributed_runtime/vlm_engine.cpp +++ b/xllm/core/distributed_runtime/vlm_engine.cpp @@ -310,6 +310,13 @@ bool VLMEngine::allocate_kv_cache(const KVCacheCapacity& kv_cache_cap) { .block_size(block_size) .enable_linear_state(has_linear_attention_layers(args_)) .enable_prefix_cache(options_.enable_prefix_cache()) + .enable_cache_aware_dp(::xllm::KVCacheConfig::get_instance() + .enable_prefix_cache_aware_dp_routing()) + .cache_aware_match_threshold(::xllm::KVCacheConfig::get_instance() + .prefix_cache_aware_dp_match_threshold()) + .cache_aware_imbalance_threshold( + ::xllm::KVCacheConfig::get_instance() + .prefix_cache_aware_dp_imbalance_threshold()) .enable_disagg_pd(options_.enable_disagg_pd()) .hasher_type(BlockHasherType::MM) .max_concurrent_requests( diff --git a/xllm/core/framework/block/block_manager.h b/xllm/core/framework/block/block_manager.h index 47ab3e48d5..2bc0f971ad 100644 --- a/xllm/core/framework/block/block_manager.h +++ b/xllm/core/framework/block/block_manager.h @@ -96,6 +96,17 @@ class BlockManager { const Slice& block_hashes = {}) = 0; virtual void cache(const std::vector& blocks) = 0; + // Read-only probe of the prefix-cache hit length (in blocks) for `token_ids`. + // Mirrors allocate_shared's arguments but mutates no state. Non-prefix leaves + // keep the base default of 0. + virtual size_t prefix_cache_match_length( + const Slice& /*token_ids*/, + const Slice& /*existed_shared_blocks*/ = {}, + const MMData& /*mm_data*/ = MMData(), + const Slice& /*block_hashes*/ = {}) const { + return 0; + } + virtual size_t num_blocks_in_prefix_cache() const = 0; virtual size_t num_free_blocks() const = 0; virtual size_t num_used_blocks() const = 0; diff --git a/xllm/core/framework/block/block_manager_impl.cpp b/xllm/core/framework/block/block_manager_impl.cpp index bb80bb2ed4..ad7751f7b2 100644 --- a/xllm/core/framework/block/block_manager_impl.cpp +++ b/xllm/core/framework/block/block_manager_impl.cpp @@ -183,6 +183,18 @@ std::vector BlockManagerImpl::allocate_shared( return {}; } +size_t BlockManagerImpl::prefix_cache_match_length( + const Slice& token_ids, + const Slice& existed_shared_blocks, + const MMData& mm_data, + const Slice& block_hashes) const { + if (!options_.enable_prefix_cache() || !prefix_cache_) { + return 0; + } + return prefix_cache_->match_length( + token_ids, existed_shared_blocks, mm_data, block_hashes); +} + void BlockManagerImpl::cache(const Slice& token_ids, std::vector& blocks, size_t existed_shared_blocks_num, diff --git a/xllm/core/framework/block/block_manager_impl.h b/xllm/core/framework/block/block_manager_impl.h index fcab21570f..e22053d7be 100644 --- a/xllm/core/framework/block/block_manager_impl.h +++ b/xllm/core/framework/block/block_manager_impl.h @@ -56,6 +56,12 @@ class BlockManagerImpl : public BlockManager { const Slice& block_hashes = {}) override; void cache(const std::vector& blocks) override; + size_t prefix_cache_match_length( + const Slice& token_ids, + const Slice& existed_shared_blocks = {}, + const MMData& mm_data = MMData(), + const Slice& block_hashes = {}) const override; + size_t num_blocks_in_prefix_cache() const override { if (options_.enable_prefix_cache()) { CHECK(prefix_cache_); diff --git a/xllm/core/framework/block/block_manager_pool.cpp b/xllm/core/framework/block/block_manager_pool.cpp index 546da8f390..ccf42aa1b0 100644 --- a/xllm/core/framework/block/block_manager_pool.cpp +++ b/xllm/core/framework/block/block_manager_pool.cpp @@ -16,6 +16,7 @@ limitations under the License. #include "block_manager_pool.h" #include +#include #include #include "block_manager_impl.h" @@ -112,12 +113,102 @@ int32_t BlockManagerPool::get_manager_with_max_free_blocks() const { return max_index; } +int32_t BlockManagerPool::get_cache_aware_dp_rank(Sequence* sequence) const { + const size_t block_size = options_.block_size(); + CHECK_GT(block_size, 0u); + const size_t needed_blocks = + (sequence->num_tokens() + block_size - 1) / block_size; + const size_t num_managers = block_managers_.size(); + + // One pass collects, per rank: the prefix-cache hit length, free blocks, + // whether the whole request fits (free + evictable-but-unmatched cache), the + // least-loaded rank (load-balance baseline), and the used-block spread. + std::vector matched(num_managers, 0); + std::vector free_blocks(num_managers, 0); + std::vector can_fit(num_managers, false); + int32_t least_loaded_rank = 0; + size_t max_free = 0; + size_t max_used = 0; + size_t min_used = std::numeric_limits::max(); + size_t total_blocks = 0; + for (size_t i = 0; i < num_managers; ++i) { + auto* composite = + static_cast(block_managers_[i].get()); + matched[i] = composite->prefix_match_length_for_sequence(sequence); + free_blocks[i] = composite->num_free_blocks(); + const size_t used = composite->num_used_blocks(); + const size_t cached = composite->num_blocks_in_prefix_cache(); + total_blocks = std::max(total_blocks, composite->num_total_blocks()); + + // Matched blocks are reused (not evicted); only the rest of the cache could + // be reclaimed to make room for the request. + const size_t evictable = cached > matched[i] ? cached - matched[i] : 0; + const size_t available = free_blocks[i] + evictable; + const size_t net_needed = + needed_blocks > matched[i] ? needed_blocks - matched[i] : 0; + can_fit[i] = net_needed <= available; + + if (i == 0 || free_blocks[i] > max_free) { + max_free = free_blocks[i]; + least_loaded_rank = static_cast(i); + } + max_used = std::max(max_used, used); + min_used = std::min(min_used, used); + } + + // Guard 1 (popular-prefix overload): when ranks are already too imbalanced, + // ignore affinity and refill the least-loaded rank. Without this, a prefix + // shared by every request would pin all traffic onto a single rank. + if (total_blocks > 0 && max_used > min_used) { + const double imbalance = static_cast(max_used - min_used) / + static_cast(total_blocks); + if (imbalance > options_.cache_aware_imbalance_threshold()) { + return least_loaded_rank; + } + } + + // Guard 2 (cold-start tiny-prefix herding): only honor affinity when the hit + // covers a meaningful fraction of the request. A 1-block prefix shared at + // cold start does not qualify, so requests spread by load instead. + const size_t min_match_blocks = + static_cast(std::ceil(options_.cache_aware_match_threshold() * + static_cast(needed_blocks))); + + int32_t best_rank = -1; + size_t best_matched = 0; + size_t best_free = 0; + for (size_t i = 0; i < num_managers; ++i) { + if (!can_fit[i] || matched[i] == 0 || matched[i] < min_match_blocks) { + continue; + } + const bool better = + best_rank < 0 || matched[i] > best_matched || + (matched[i] == best_matched && free_blocks[i] > best_free); + if (better) { + best_rank = static_cast(i); + best_matched = matched[i]; + best_free = free_blocks[i]; + } + } + if (best_rank >= 0) { + return best_rank; + } + + // No rank offers a meaningful, fitting prefix -> balance by free blocks. + return least_loaded_rank; +} + int32_t BlockManagerPool::get_dp_rank(Sequence* sequence) const { int32_t dp_rank; if (sequence->dp_rank() >= 0) { dp_rank = sequence->dp_rank(); } else { - dp_rank = get_manager_with_max_free_blocks(); + if (options_.enable_prefix_cache() && options_.enable_cache_aware_dp() && + block_managers_.size() > 1) { + dp_rank = get_cache_aware_dp_rank(sequence); + } else { + dp_rank = get_manager_with_max_free_blocks(); + } sequence->set_dp_rank(dp_rank); } return dp_rank; diff --git a/xllm/core/framework/block/block_manager_pool.h b/xllm/core/framework/block/block_manager_pool.h index 2e4b3fd05b..ab5189c22c 100644 --- a/xllm/core/framework/block/block_manager_pool.h +++ b/xllm/core/framework/block/block_manager_pool.h @@ -52,6 +52,19 @@ class BlockManagerPool : public KVCacheManager { // Hasher type bound to the engine (TEXT for LLM, MM for VLM). PROPERTY(BlockHasherType, hasher_type) = BlockHasherType::TEXT; PROPERTY(uint32_t, num_single_blocks) = 0; + // When true (and prefix cache is on with dp_size > 1), a new sequence is + // routed to the dp rank that can fit the whole prefill request and has the + // longest prefix-cache hit, instead of the rank with the most free blocks. + PROPERTY(bool, enable_cache_aware_dp) = false; + // Minimum fraction of the prefill request's blocks that must hit a rank's + // prefix cache before cache-affinity prefers that rank. Below this, routing + // balances by free blocks. Prevents a tiny shared prefix (e.g. a one-block + // system prompt) from herding every request onto one rank at cold start. + PROPERTY(double, cache_aware_match_threshold) = 0.5; + // Maximum tolerated (max_used - min_used) / total_blocks before cache + // affinity is suspended in favor of the least-loaded rank. Prevents a + // popular shared prefix from pinning all traffic to a single rank. + PROPERTY(double, cache_aware_imbalance_threshold) = 0.1; }; explicit BlockManagerPool(const Options& options, int32_t dp_size = 1); @@ -104,6 +117,11 @@ class BlockManagerPool : public KVCacheManager { protected: int32_t get_manager_with_max_free_blocks() const; + // Prefix-cache-affinity routing: pick the dp rank by the ordered key + // (can_fit_whole_prefill, prefix_match_blocks, free_blocks), preferring a + // rank that fits the request and shares the longest prefix. Falls back to the + // longest-prefix rank when none can fit. + int32_t get_cache_aware_dp_rank(Sequence* sequence) const; int32_t get_dp_rank(Sequence* sequence) const; bool process_beam_search(Sequence* sequence, bool need_swap = false); diff --git a/xllm/core/framework/block/composite_block_manager.cpp b/xllm/core/framework/block/composite_block_manager.cpp index 14cd9af6dc..d8b3a377ae 100644 --- a/xllm/core/framework/block/composite_block_manager.cpp +++ b/xllm/core/framework/block/composite_block_manager.cpp @@ -242,6 +242,25 @@ void CompositeBlockManager::allocate_shared_for_sequence(Sequence* seq) { seq->add_shared_blocks(BlockType::KV, std::move(shared)); } +size_t CompositeBlockManager::prefix_match_length_for_sequence( + Sequence* seq) const { + if (seq == nullptr) { + return 0; + } + auto it = leaves_.find(BlockType::KV); + if (it == leaves_.end() || !it->second.supports_prefix_cache) { + return 0; + } + BlockManager& kv_leaf = *it->second.leaf; + seq->update_block_hashes(static_cast(kv_leaf.block_size()), + kv_leaf.options().hasher_type()); + KVCacheState& kv_state = seq->kv_state(); + const auto existed = kv_state.blocks(BlockType::KV) + .slice(0, kv_state.shared_blocks_num(BlockType::KV)); + return kv_leaf.prefix_cache_match_length( + seq->tokens(), existed, seq->mm_data(), seq->block_hashes()); +} + void CompositeBlockManager::cache_for_sequence(Sequence* seq) { if (seq == nullptr) { return; diff --git a/xllm/core/framework/block/composite_block_manager.h b/xllm/core/framework/block/composite_block_manager.h index b8211439a6..1736951261 100644 --- a/xllm/core/framework/block/composite_block_manager.h +++ b/xllm/core/framework/block/composite_block_manager.h @@ -55,6 +55,12 @@ class CompositeBlockManager : public BlockManager { bool allocate_sequence(Sequence* seq, size_t num_tokens); void deallocate_for_sequence(Sequence* seq); void allocate_shared_for_sequence(Sequence* seq); + // Read-only counterpart of allocate_shared_for_sequence: returns how many + // prefix-cache blocks the KV leaf would share for `seq` without allocating + // or mutating any block-manager state (it does refresh the sequence's cached + // block hashes, which a subsequent allocation needs anyway). Returns 0 when + // there is no prefix-capable KV leaf. Used by cache-aware DP routing. + size_t prefix_match_length_for_sequence(Sequence* seq) const; void cache_for_sequence(Sequence* seq); void cache_for_sequence(Sequence* seq, size_t num_tokens); diff --git a/xllm/core/framework/block/concurrent_block_manager_impl.cpp b/xllm/core/framework/block/concurrent_block_manager_impl.cpp index 34a5b5cd52..37a4ac6936 100644 --- a/xllm/core/framework/block/concurrent_block_manager_impl.cpp +++ b/xllm/core/framework/block/concurrent_block_manager_impl.cpp @@ -109,6 +109,16 @@ void ConcurrentBlockManagerImpl::reset_prefix_cache() { inner_->reset_prefix_cache(); } +size_t ConcurrentBlockManagerImpl::prefix_cache_match_length( + const Slice& token_ids, + const Slice& existed_shared_blocks, + const MMData& mm_data, + const Slice& block_hashes) const { + std::lock_guard lock(mutex_); + return inner_->prefix_cache_match_length( + token_ids, existed_shared_blocks, mm_data, block_hashes); +} + size_t ConcurrentBlockManagerImpl::num_blocks_in_prefix_cache() const { std::lock_guard lock(mutex_); return inner_->num_blocks_in_prefix_cache(); diff --git a/xllm/core/framework/block/concurrent_block_manager_impl.h b/xllm/core/framework/block/concurrent_block_manager_impl.h index c6fb6cf14c..b5e9584830 100644 --- a/xllm/core/framework/block/concurrent_block_manager_impl.h +++ b/xllm/core/framework/block/concurrent_block_manager_impl.h @@ -63,6 +63,12 @@ class ConcurrentBlockManagerImpl : public BlockManager { void reset_prefix_cache() override; + size_t prefix_cache_match_length( + const Slice& token_ids, + const Slice& existed_shared_blocks = {}, + const MMData& mm_data = MMData(), + const Slice& block_hashes = {}) const override; + size_t num_blocks_in_prefix_cache() const override; size_t num_free_blocks() const override; size_t num_used_blocks() const override; diff --git a/xllm/core/framework/config/kv_cache_config.cpp b/xllm/core/framework/config/kv_cache_config.cpp index c5bb1c929e..59334ff29e 100644 --- a/xllm/core/framework/config/kv_cache_config.cpp +++ b/xllm/core/framework/config/kv_cache_config.cpp @@ -48,6 +48,28 @@ DEFINE_bool(enable_in_batch_prefix_cache, "Whether to cache admitted prefill full blocks into the prefix " "cache so that later requests in the same batch can share them."); +DEFINE_bool(enable_prefix_cache_aware_dp_routing, + false, + "When data parallel and prefix cache are enabled, route a new " + "sequence to the dp rank that can hold the whole prefill request " + "and has the longest prefix-cache hit, instead of the rank with " + "the most free blocks."); + +DEFINE_double(prefix_cache_aware_dp_match_threshold, + 0.5, + "Cache-aware dp routing: minimum fraction of a request's blocks " + "that must hit a rank's prefix cache before affinity prefers it. " + "Below this, routing balances by free blocks. Prevents a tiny " + "shared prefix from herding every request onto one rank."); + +DEFINE_double( + prefix_cache_aware_dp_imbalance_threshold, + 0.1, + "Cache-aware dp routing: maximum tolerated (max_used - min_used) " + "/ total_blocks before affinity is suspended in favor of the " + "least-loaded rank. Prevents a popular shared prefix from pinning " + "all traffic onto a single rank."); + DEFINE_int64(max_linear_state_cache_slots, 0, "Maximum active linear-attention state cache slots. 0 derives an " @@ -75,6 +97,9 @@ void KVCacheConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(kv_cache_dtype); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_prefix_cache); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_in_batch_prefix_cache); + XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_prefix_cache_aware_dp_routing); + XLLM_CONFIG_ASSIGN_FROM_FLAG(prefix_cache_aware_dp_match_threshold); + XLLM_CONFIG_ASSIGN_FROM_FLAG(prefix_cache_aware_dp_imbalance_threshold); XLLM_CONFIG_ASSIGN_FROM_FLAG(max_linear_state_cache_slots); XLLM_CONFIG_ASSIGN_FROM_FLAG(xxh3_128bits_seed); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_xtensor); @@ -88,6 +113,9 @@ void KVCacheConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(kv_cache_dtype); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_prefix_cache); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_in_batch_prefix_cache); + XLLM_CONFIG_ASSIGN_FROM_JSON(enable_prefix_cache_aware_dp_routing); + XLLM_CONFIG_ASSIGN_FROM_JSON(prefix_cache_aware_dp_match_threshold); + XLLM_CONFIG_ASSIGN_FROM_JSON(prefix_cache_aware_dp_imbalance_threshold); XLLM_CONFIG_ASSIGN_FROM_JSON(max_linear_state_cache_slots); XLLM_CONFIG_ASSIGN_FROM_JSON(xxh3_128bits_seed); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_xtensor); @@ -109,6 +137,12 @@ void KVCacheConfig::append_config_json( config_json, default_config, enable_prefix_cache); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, enable_in_batch_prefix_cache); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, enable_prefix_cache_aware_dp_routing); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, prefix_cache_aware_dp_match_threshold); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, prefix_cache_aware_dp_imbalance_threshold); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, max_linear_state_cache_slots); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/kv_cache_config.h b/xllm/core/framework/config/kv_cache_config.h index d43fdafb35..900520131d 100644 --- a/xllm/core/framework/config/kv_cache_config.h +++ b/xllm/core/framework/config/kv_cache_config.h @@ -47,6 +47,9 @@ class KVCacheConfig final { "kv_cache_dtype", "enable_prefix_cache", "enable_in_batch_prefix_cache", + "enable_prefix_cache_aware_dp_routing", + "prefix_cache_aware_dp_match_threshold", + "prefix_cache_aware_dp_imbalance_threshold", "max_linear_state_cache_slots", "xxh3_128bits_seed", "enable_xtensor", @@ -66,6 +69,12 @@ class KVCacheConfig final { PROPERTY(bool, enable_in_batch_prefix_cache) = false; + PROPERTY(bool, enable_prefix_cache_aware_dp_routing) = false; + + PROPERTY(double, prefix_cache_aware_dp_match_threshold) = 0.5; + + PROPERTY(double, prefix_cache_aware_dp_imbalance_threshold) = 0.1; + PROPERTY(int64_t, max_linear_state_cache_slots) = 0; PROPERTY(uint32_t, xxh3_128bits_seed) = 1024; diff --git a/xllm/core/framework/prefix_cache/prefix_cache.cpp b/xllm/core/framework/prefix_cache/prefix_cache.cpp index b1903a4573..e8395c5891 100644 --- a/xllm/core/framework/prefix_cache/prefix_cache.cpp +++ b/xllm/core/framework/prefix_cache/prefix_cache.cpp @@ -103,6 +103,58 @@ std::vector PrefixCache::match(const Slice& token_ids, return blocks; } +size_t PrefixCache::match_length(const Slice& token_ids, + const Slice& existed_shared_blocks, + const MMData& mm_data, + const Slice& block_hashes) const { + // allign tokens to block boundary + const size_t n_tokens = round_down(token_ids.size(), block_size_); + if (n_tokens == 0) { + return 0; + } + + const size_t n_blocks = n_tokens / block_size_; + const size_t start_block = existed_shared_blocks.size(); + // The already-shared blocks are guaranteed hits; start counting from there. + size_t matched = start_block; + + // Read-only counterpart of match()'s match_block lambda: it only probes the + // hash table and never touches the LRU list or ref counts. + auto probe = [&](const XXH3Key& token_hash_key) -> bool { + return cached_blocks_.find(token_hash_key) != cached_blocks_.end(); + }; + + // Fast path: precomputed chained hash covers every matchable block. + if (block_hashes.size() >= n_blocks) { + for (size_t b = start_block; b < n_blocks; ++b) { + if (!probe(block_hashes[b])) { + break; + } + ++matched; + } + return matched; + } + + // Fallback: compute the chained hash on the fly. + XXH3Key token_hash_key = + existed_shared_blocks.empty() + ? XXH3Key{} + : XXH3Key{existed_shared_blocks.back().get_immutable_hash_value()}; + auto hasher = + BlockHasher::create(hasher_type_, mm_data, start_block * block_size_); + for (size_t b = start_block; b < n_blocks; ++b) { + const size_t i = b * block_size_; + const uint8_t* pre_hash_value = (b == 0) ? nullptr : token_hash_key.data; + hasher->compute( + token_ids, i, i + block_size_, pre_hash_value, token_hash_key); + if (!probe(token_hash_key)) { + break; + } + ++matched; + } + return matched; +} + size_t PrefixCache::insert(const Slice& token_ids, std::vector& blocks, size_t existed_shared_blocks_num, diff --git a/xllm/core/framework/prefix_cache/prefix_cache.h b/xllm/core/framework/prefix_cache/prefix_cache.h index 976e56f59d..fdb761069c 100644 --- a/xllm/core/framework/prefix_cache/prefix_cache.h +++ b/xllm/core/framework/prefix_cache/prefix_cache.h @@ -71,6 +71,15 @@ class PrefixCache { const MMData& mm_data = MMData(), const Slice& block_hashes = {}); + // Read-only probe of how many leading blocks of `token_ids` would be matched + // by `match()` (including `existed_shared_blocks`), without mutating the LRU + // order, hit/total metrics, or block ref counts. Used by cache-aware DP + // routing to score a rank before committing a sequence to it. + virtual size_t match_length(const Slice& token_ids, + const Slice& existed_shared_blocks = {}, + const MMData& mm_data = MMData(), + const Slice& block_hashes = {}) const; + // insert the token ids and blocks into the prefix tree // and set hash key to the corresponding block // return the length of new inserted tokens.