feat: restore D2H offload collection.

Author: Kang-Meng

xllm/core/distributed_runtime/llm_engine.cpp

@@ -40,8 +40,7 @@ limitations under the License.
 #include "core/framework/config/parallel_config.h"
 #include "core/framework/config/service_config.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"
@@ -566,19 +565,12 @@ 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) {
+    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/hierarchy_block_manager_pool.cpp

@@ -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_;

xllm/core/framework/block/hierarchy_block_manager_pool.h

@@ -29,6 +29,7 @@ class KVCacheTensorRole {
     CONV = 3,
     SSM = 4,
     INDEX_SCALE = 5,
+    SWA = 6,
     INVALID = -1,
   };
 
@@ -46,6 +47,8 @@ class KVCacheTensorRole {
       value_ = CONV;
     } else if (str == "SSM" || str == "ssm") {
       value_ = SSM;
+    } else if (str == "SWA" || str == "swa") {
+      value_ = SWA;
     } else {
       value_ = INVALID;
     }
@@ -74,6 +77,8 @@ class KVCacheTensorRole {
       return "conv";
     } else if (this->value_ == SSM) {
       return "ssm";
+    } else if (this->value_ == SWA) {
+      return "swa";
     } else {
       return "invalid";
     }

xllm/core/framework/kv_cache/kv_cache_tensor_role.h

@@ -393,7 +393,7 @@ HostPageAlignedRegion::HostPageAlignedRegion(size_t bytes) {
   CHECK(base_ptr != MAP_FAILED)
       << "Failed to mmap host memory, size=" << total_bytes;
   if (mlock(base_ptr, total_bytes) != 0) {
-    const int err = errno;
+    const int32_t err = errno;
     munmap(base_ptr, total_bytes);
     base_ptr = nullptr;
     total_bytes = 0;

xllm/core/framework/kv_cache/kv_cache_utils.cpp

@@ -19,6 +19,7 @@ limitations under the License.
 #include <torch/torch.h>
 
 #include <cstdint>
+#include <map>
 #include <memory>
 #include <optional>
 #include <string>
@@ -50,6 +51,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;
@@ -106,6 +108,21 @@ struct KVCacheTensor {
   torch::Tensor tensor;
 };
 
+using BlockTypeTensorMap = std::map<KVCacheTensorRole::Value, torch::Tensor>;
+
+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;

xllm/core/framework/kv_cache/kv_cache_utils.h

@@ -41,8 +41,7 @@ cc_library(
     $<$<BOOL:${USE_NPU}>:llm_data_dist_transfer.h>
     $<$<BOOL:${USE_NPU}>: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
     $<$<OR:$<BOOL:${USE_NPU}>,$<BOOL:${USE_MLU}>,$<BOOL:${USE_DCU}>>:mooncake_transfer_engine.h>
     $<$<OR:$<BOOL:${USE_NPU}>,$<BOOL:${USE_MLU}>,$<BOOL:${USE_DCU}>>:mooncake_kv_cache_transfer.h>
     $<$<BOOL:${USE_NPU}>:mooncake_weight_transfer.h>
@@ -51,8 +50,7 @@ cc_library(
     $<$<BOOL:${USE_NPU}>:llm_data_dist_transfer.cpp>
     $<$<BOOL:${USE_NPU}>: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
     $<$<OR:$<BOOL:${USE_NPU}>,$<BOOL:${USE_MLU}>,$<BOOL:${USE_DCU}>>:mooncake_transfer_engine.cpp>
     $<$<OR:$<BOOL:${USE_NPU}>,$<BOOL:${USE_MLU}>,$<BOOL:${USE_DCU}>>:mooncake_kv_cache_transfer.cpp>
     $<$<BOOL:${USE_NPU}>:mooncake_weight_transfer.cpp>
@@ -61,6 +59,7 @@ cc_library(
     :kv_cache
     :push_route
     :xtensor
+    :platform
     $<$<BOOL:${USE_NPU}>:graph>
     glog::glog
     $<$<BOOL:${USE_NPU}>:llm_datadist>