feat: add composite host cache support for MTP.

tests/core/framework/block/composite_block_manager_test.cpp

@@ -321,6 +321,28 @@ TEST(CompositeBlockManagerTest, FailedGrowthRollsBackNewBlocks) {
   manager.deallocate_for_sequence(&seq);
 }
 
+TEST(CompositeBlockManagerTest, FreshSequenceEvictsSwaPrefixForCapacity) {
+  BlockManager::Options opts = MakeCompositeOptions(/*base_num_blocks=*/256,
+                                                    kBaseBlockSize,
+                                                    /*window_size=*/12,
+                                                    /*max_seqs_per_batch=*/4);
+  CompositeBlockManager manager(build_composite_leaves(opts));
+  const std::vector<int32_t> first_prompt(kMaxTokensPerBatch, 1);
+  Sequence first = MakeTestSequence(0, first_prompt);
+  ASSERT_TRUE(manager.allocate_sequence(&first, first_prompt.size()));
+  first.kv_state().incr_kv_cache_tokens_num(first_prompt.size());
+  manager.deallocate_for_sequence(&first);
+  first.reset();
+
+  const std::vector<int32_t> second_prompt(kMaxTokensPerBatch, 2);
+  Sequence second = MakeTestSequence(1, second_prompt);
+  EXPECT_TRUE(manager.allocate_sequence(&second, second_prompt.size()));
+  EXPECT_EQ(SwaBlocks(second).size(),
+            ExpectedSwaLogicalBlocks(second_prompt.size()));
+
+  manager.deallocate_for_sequence(&second);
+}
+
 TEST(CompositeBlockManagerTest, DeallocateToleratesRolledBackEmptySequence) {
   BlockManager::Options opts = MakeCompositeOptions(/*base_num_blocks=*/128,
                                                     kBaseBlockSize,

tests/core/framework/block/hierarchy_block_manager_pool_test.cpp

@@ -314,6 +314,87 @@ TEST(HierarchyBlockManagerPoolTest, TypedLayoutHasPerDpRankLeaves) {
   }
 }
 
+TEST(HierarchyBlockManagerPoolTest,
+     RuntimeSizedTypedLayoutAllocatesColdPromptChunks) {
+  BlockManagerPool::Options options = make_typed_cache_options();
+  options.num_blocks(6400)
+      .max_tokens_per_batch(25000)
+      .max_seqs_per_batch(8)
+      .host_num_blocks_by_type({{BlockType::SWA, 856},
+                                {BlockType::C4, 6400},
+                                {BlockType::C128, 200}});
+  HierarchyBlockManagerPool pool(options,
+                                 /*engine=*/nullptr,
+                                 /*dp_size=*/1);
+
+  std::vector<int32_t> tokens(36025, 97);
+  Sequence sequence = make_test_sequence(/*index=*/0, tokens);
+
+  pool.allocate_shared(&sequence);
+  const HostCacheRestorePoint selected = pool.select_host_cache_restore(
+      &sequence, std::numeric_limits<size_t>::max());
+  pool.trim_host_cache(&sequence, selected);
+  EXPECT_TRUE(pool.allocate(&sequence, /*num_tokens=*/16384));
+  EXPECT_GE(sequence.kv_state().current_max_tokens_capacity(), 16384u);
+
+  sequence.kv_state().set_kv_cache_tokens_num(16384);
+  EXPECT_TRUE(pool.allocate(&sequence, /*num_tokens=*/32768));
+  EXPECT_GE(sequence.kv_state().current_max_tokens_capacity(), 32768u);
+  const Slice<Block> swa_blocks = sequence.kv_state().blocks(BlockType::SWA);
+  ASSERT_EQ(swa_blocks.size(), 256u);
+  EXPECT_FALSE(swa_blocks[0].is_valid());
+  EXPECT_TRUE(swa_blocks[127].is_valid());
+  EXPECT_TRUE(swa_blocks[128].is_valid());
+
+  sequence.kv_state().set_kv_cache_tokens_num(32768);
+  EXPECT_TRUE(pool.allocate(&sequence, /*num_tokens=*/36025));
+  EXPECT_GE(sequence.kv_state().current_max_tokens_capacity(), 36025u);
+
+  pool.deallocate(&sequence);
+}
+
+TEST(HierarchyBlockManagerPoolTest,
+     RuntimeSizedTypedLayoutRestoresColdPromptTail) {
+  BlockManagerPool::Options options = make_typed_cache_options();
+  options.num_blocks(6400)
+      .max_tokens_per_batch(25000)
+      .max_seqs_per_batch(8)
+      .host_num_blocks_by_type({{BlockType::SWA, 856},
+                                {BlockType::C4, 6400},
+                                {BlockType::C128, 200}});
+  HierarchyBlockManagerPool pool(options,
+                                 /*engine=*/nullptr,
+                                 /*dp_size=*/1);
+
+  std::vector<int32_t> tokens(36025, 101);
+  auto& host_leaves =
+      HierarchyPoolTestPeer::mutable_host_block_managers(pool).front();
+  seed_host_prefix(host_leaves.at(BlockType::SWA).leaf.get(), tokens);
+  seed_host_prefix(host_leaves.at(BlockType::C4).leaf.get(), tokens);
+  seed_host_prefix(host_leaves.at(BlockType::C128).leaf.get(), tokens);
+
+  Sequence restored = make_test_sequence(/*index=*/0, tokens);
+  ASSERT_TRUE(allocate_with_host_cache_budget(
+      &pool,
+      &restored,
+      /*num_tokens=*/tokens.size(),
+      /*max_copy_units=*/std::numeric_limits<size_t>::max()));
+  EXPECT_EQ(restored.kv_state().kv_cache_tokens_num(), 32768u);
+
+  const Slice<Block> swa_blocks = restored.kv_state().blocks(BlockType::SWA);
+  ASSERT_EQ(swa_blocks.size(), 282u);
+  EXPECT_FALSE(swa_blocks[254].is_valid());
+  EXPECT_TRUE(swa_blocks[255].is_valid());
+  EXPECT_TRUE(swa_blocks.back().is_valid());
+  const size_t valid_swa_blocks = static_cast<size_t>(std::count_if(
+      swa_blocks.begin(), swa_blocks.end(), [](const Block& block) {
+        return block.is_valid();
+      }));
+  EXPECT_EQ(valid_swa_blocks, 27u);
+
+  pool.deallocate(&restored);
+}
+
 TEST(HierarchyBlockManagerPoolTest,
      DecodeTypedLayoutKeepsSwaHostLeafForOffloadOnly) {
   BlockManagerPool::Options options = make_typed_cache_options();

tests/core/framework/kv_cache_transfer/CMakeLists.txt

@@ -27,6 +27,7 @@ cc_test(
     kv_transfer_completion_test.cpp
   DEPS
     :kv_transfer_completion
+    :kv_cache_transfer
     GTest::gtest_main
 )
 

tests/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer_test.cpp

@@ -46,15 +46,22 @@ class HierarchyKVCacheTransferTestPeer final {
   static void set_layer_batch_ranges(
       HierarchyKVCacheTransfer* transfer,
       std::vector<HierarchyKVCacheTransfer::LayerBatchRange> ranges) {
-    transfer->layer_batch_ranges_ = std::move(ranges);
+    transfer->participant_states_.at(CacheParticipant::TARGET)
+        .layer_batch_ranges = std::move(ranges);
   }
 
   static bool load_from_host(
       HierarchyKVCacheTransfer* transfer,
       std::shared_ptr<LayerSynchronizer> synchronizer,
       const std::vector<BlockTransferInfo>& block_transfer_info) {
-    return transfer->load_from_host(std::move(synchronizer),
-                                    block_transfer_info);
+    auto transaction =
+        std::make_shared<HierarchyKVCacheTransfer::LoadTransaction>();
+    transaction->synchronizers[CacheParticipant::TARGET] =
+        std::move(synchronizer);
+    transaction->required_participant_mask =
+        HierarchyKVCacheTransfer::participant_mask(CacheParticipant::TARGET);
+    return transfer->load_from_host(
+        CacheParticipant::TARGET, transaction, block_transfer_info);
   }
 };
 

tests/core/framework/kv_cache_transfer/kv_transfer_completion_test.cpp

@@ -17,15 +17,138 @@ limitations under the License.
 
 #include <gtest/gtest.h>
 
+#include <array>
 #include <chrono>
 #include <future>
 #include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "core/framework/kv_cache_transfer/kv_cache_store.h"
 
 namespace xllm {
+
+class KVCacheStoreTestPeer final {
+ public:
+  static void set_config(KVCacheStore* store,
+                         const KVCacheStoreInitConfig& config) {
+    store->config_ = config;
+  }
+
+  static void set_components(KVCacheStore* store,
+                             std::vector<HostCacheComponentSchema> components) {
+    store->components_ = std::move(components);
+  }
+
+  static std::string build_component_key(
+      const KVCacheStore& store,
+      const HostCacheComponentSchema& component,
+      const BlockTransferInfo& block_info) {
+    return store.build_component_key(component, block_info);
+  }
+
+  static size_t required_component_count(const KVCacheStore& store,
+                                         BlockType block_type) {
+    return store.required_components(block_type).size();
+  }
+};
+
 namespace {
 
 using namespace std::chrono_literals;
 
+BlockTransferInfo make_block_info(BlockType block_type) {
+  std::array<uint8_t, XXH3_128BITS_HASH_VALUE_LEN> hash_key{};
+  for (size_t index = 0; index < hash_key.size(); ++index) {
+    hash_key[index] = static_cast<uint8_t>(index + 1);
+  }
+  return BlockTransferInfo(/*src_id=*/3,
+                           /*dst_id=*/7,
+                           hash_key.data(),
+                           TransferType::D2H2G,
+                           block_type);
+}
+
+HostCacheComponentSchema make_component(CacheParticipant participant,
+                                        const std::string& model_identity,
+                                        BlockType block_type = BlockType::KV) {
+  HostCacheComponentSchema component;
+  component.participant = participant;
+  component.block_type = block_type;
+  component.model_identity = model_identity;
+  component.schema_fingerprint = "schema-fingerprint";
+  component.tp_rank = 0;
+  component.tp_size = 8;
+  return component;
+}
+
+TEST(KVCacheStoreKeyTest, SeparatesTargetAndDraftComponents) {
+  KVCacheStore store;
+  KVCacheStoreInitConfig config;
+  config.model_id = "deepseek-v4";
+  KVCacheStoreTestPeer::set_config(&store, config);
+  const BlockTransferInfo block_info = make_block_info(BlockType::C128);
+  HostCacheComponentSchema target = make_component(
+      CacheParticipant::TARGET, "deepseek-v4-target", BlockType::C128);
+  HostCacheComponentSchema draft = target;
+  draft.participant = CacheParticipant::DRAFT;
+
+  const std::string target_key =
+      KVCacheStoreTestPeer::build_component_key(store, target, block_info);
+  const std::string draft_key =
+      KVCacheStoreTestPeer::build_component_key(store, draft, block_info);
+
+  EXPECT_EQ(target_key.rfind("xllm-kv-v3:", 0), 0u);
+  EXPECT_EQ(draft_key.rfind("xllm-kv-v3:", 0), 0u);
+  EXPECT_NE(target_key, draft_key);
+}
+
+TEST(KVCacheStoreKeyTest, SeparatesParticipantModelAndSchemaIdentity) {
+  KVCacheStore store;
+  KVCacheStoreInitConfig config;
+  config.model_id = "deepseek-v4";
+  KVCacheStoreTestPeer::set_config(&store, config);
+  const BlockTransferInfo block_info = make_block_info(BlockType::SWA);
+  HostCacheComponentSchema first = make_component(
+      CacheParticipant::DRAFT, "draft-revision-a", BlockType::SWA);
+  HostCacheComponentSchema second = first;
+  second.model_identity = "draft-revision-b";
+  HostCacheComponentSchema third = first;
+  third.schema_fingerprint = "different-schema";
+
+  const std::string first_key =
+      KVCacheStoreTestPeer::build_component_key(store, first, block_info);
+  const std::string second_key =
+      KVCacheStoreTestPeer::build_component_key(store, second, block_info);
+  const std::string third_key =
+      KVCacheStoreTestPeer::build_component_key(store, third, block_info);
+
+  EXPECT_NE(first_key, second_key);
+  EXPECT_NE(first_key, third_key);
+}
+
+TEST(KVCacheStoreKeyTest, ExpandsLogicalTypeToEveryRequiredParticipant) {
+  KVCacheStore store;
+  std::vector<HostCacheComponentSchema> components;
+  components.emplace_back(
+      make_component(CacheParticipant::TARGET, "target", BlockType::C4));
+  components.emplace_back(
+      make_component(CacheParticipant::DRAFT, "draft", BlockType::C4));
+  components.emplace_back(
+      make_component(CacheParticipant::TARGET, "target", BlockType::SWA));
+  KVCacheStoreTestPeer::set_components(&store, std::move(components));
+
+  EXPECT_EQ(
+      KVCacheStoreTestPeer::required_component_count(store, BlockType::C4), 2u);
+  EXPECT_EQ(
+      KVCacheStoreTestPeer::required_component_count(store, BlockType::SWA),
+      1u);
+  EXPECT_EQ(
+      KVCacheStoreTestPeer::required_component_count(store, BlockType::C128),
+      0u);
+}
+
 TEST(KVTransferCompletionTest, WaitsForEveryTransfer) {
   folly::Promise<bool> first_promise;
   folly::Promise<bool> second_promise;

tests/core/runtime/CMakeLists.txt

@@ -116,7 +116,9 @@ if(USE_MLU)
       GTest::gtest_main
       torch_mlu
   )
+endif()
 
+if(USE_NPU OR USE_MLU)
   cc_test(
     NAME
       mtp_host_offload_test
@@ -129,8 +131,10 @@ if(USE_MLU)
       :kv_cache
       :platform
       GTest::gtest_main
-      torch_mlu
   )
+  if(USE_MLU)
+    target_link_libraries(mtp_host_offload_test PRIVATE torch_mlu)
+  endif()
 endif()
 
 if(USE_CUDA)