refactor: host prefix cache — device↔host KV-cache transfer framework (4/n).

tests/core/framework/batch/batch_test.cpp

@@ -2056,16 +2056,16 @@ 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;
   stopping_checker.set_max_generated_tokens(8);
 
   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);
 

tests/core/framework/kv_cache/kv_cache_test.cpp

@@ -22,9 +22,13 @@ limitations under the License.
 #include <vector>
 
 #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<int64_t> 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<int64_t>(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<int64_t>(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

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<uint32_t>::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<HierarchyBlockManagerPool>(options, this, dp_size_);
+  } else {
+    kv_cache_manager_ = std::make_unique<BlockManagerPool>(options, dp_size_);
   }
-  kv_cache_manager_ = std::make_unique<BlockManagerPool>(options, dp_size_);
 
   // init kv cache for each worker in parallel
   std::vector<folly::SemiFuture<bool>> futures;

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

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.

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<BlockManager> single_leaf =
         std::make_unique<SingleBlockManager>(
             /*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<ConcurrentBlockManagerImpl>(std::move(single_leaf));
     }

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)

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<BlockManager> maybe_concurrent(
     std::unique_ptr<BlockManager> 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<ConcurrentBlockManagerImpl>(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);