feat: add prefix-cache-affinity dp rank routing.

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<int32_t> tokens = {1, 2, 3, 4, 5, 6};
+  std::vector<Block> blocks = manager->allocate(/*num_blocks=*/3);
+  ASSERT_EQ(blocks.size(), 3u);
+  manager->cache(
+      Slice<int32_t>(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<int32_t>(tokens));
+  const size_t matched2 =
+      manager->prefix_cache_match_length(Slice<int32_t>(tokens));
+  EXPECT_EQ(matched1, 3u);
+  EXPECT_EQ(matched2, matched1);
+
+  // A shorter prefix matches only its leading full blocks.
+  const std::vector<int32_t> short_tokens = {1, 2, 3, 4};
+  EXPECT_EQ(manager->prefix_cache_match_length(Slice<int32_t>(short_tokens)),
+            2u);
+
+  // A diverging prefix matches only the shared leading block.
+  const std::vector<int32_t> diverged = {1, 2, 9, 9};
+  EXPECT_EQ(manager->prefix_cache_match_length(Slice<int32_t>(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<int32_t> 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<size_t> 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<int32_t> 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<int32_t> 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<int32_t> 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

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);

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())

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(

xllm/core/framework/block/block_manager.h

@@ -96,6 +96,17 @@ class BlockManager {
                      const Slice<XXH3Key>& block_hashes = {}) = 0;
   virtual void cache(const std::vector<Block>& 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<int32_t>& /*token_ids*/,
+      const Slice<Block>& /*existed_shared_blocks*/ = {},
+      const MMData& /*mm_data*/ = MMData(),
+      const Slice<XXH3Key>& /*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;

xllm/core/framework/block/block_manager_impl.cpp

@@ -183,6 +183,18 @@ std::vector<Block> BlockManagerImpl::allocate_shared(
   return {};
 }
 
+size_t BlockManagerImpl::prefix_cache_match_length(
+    const Slice<int32_t>& token_ids,
+    const Slice<Block>& existed_shared_blocks,
+    const MMData& mm_data,
+    const Slice<XXH3Key>& 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<int32_t>& token_ids,
                              std::vector<Block>& blocks,
                              size_t existed_shared_blocks_num,

xllm/core/framework/block/block_manager_impl.h

@@ -56,6 +56,12 @@ class BlockManagerImpl : public BlockManager {
              const Slice<XXH3Key>& block_hashes = {}) override;
   void cache(const std::vector<Block>& blocks) override;
 
+  size_t prefix_cache_match_length(
+      const Slice<int32_t>& token_ids,
+      const Slice<Block>& existed_shared_blocks = {},
+      const MMData& mm_data = MMData(),
+      const Slice<XXH3Key>& block_hashes = {}) const override;
+
   size_t num_blocks_in_prefix_cache() const override {
     if (options_.enable_prefix_cache()) {
       CHECK(prefix_cache_);

xllm/core/framework/block/block_manager_pool.cpp

@@ -16,6 +16,7 @@ limitations under the License.
 #include "block_manager_pool.h"
 
 #include <algorithm>
+#include <cmath>
 #include <limits>
 
 #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<size_t> matched(num_managers, 0);
+  std::vector<size_t> free_blocks(num_managers, 0);
+  std::vector<bool> 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<size_t>::max();
+  size_t total_blocks = 0;
+  for (size_t i = 0; i < num_managers; ++i) {
+    auto* composite =
+        static_cast<CompositeBlockManager*>(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<int32_t>(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<double>(max_used - min_used) /
+                             static_cast<double>(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<size_t>(std::ceil(options_.cache_aware_match_threshold() *
+                                    static_cast<double>(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<int32_t>(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;