refactor: abstract multi-tier kv cache transfer from WorkerImpl.

Author: Kang-Meng

xllm/core/distributed_runtime/worker_service.cpp

@@ -400,11 +400,8 @@ void WorkerService::TransferBlocks(
   std::vector<BlockTransferInfo> block_transfer_info;
   uint64_t batch_id = proto_to_block_transfer_info(*req, block_transfer_info);
 
-  if (batch_id == UNINITIALIZED_BATCH_ID) {
-    resp->set_success_cnt(worker_->transfer_kv_blocks(block_transfer_info));
-  } else {
-    worker_->transfer_kv_blocks(batch_id, std::move(block_transfer_info));
-  }
+  resp->set_success_cnt(
+      worker_->transfer_kv_blocks(batch_id, std::move(block_transfer_info)));
   return;
 }
 
@@ -480,7 +477,8 @@ void WorkerService::PrefetchFromStorage(
                                    std::min(i + options_.prefetch_bacth_size(),
                                             transfer_slice.size()));
 
-          auto success_cnt = worker_->prefetch_from_storage(current_slice);
+          auto success_cnt = worker_->transfer_kv_blocks(UNINITIALIZED_BATCH_ID,
+                                                         current_slice);
 
           if (success_cnt != current_slice.size() ||
               i + options_.prefetch_bacth_size() >= transfer_slice.size()) {

xllm/core/framework/kv_cache/CMakeLists.txt

@@ -15,6 +15,7 @@ cc_library(
     $<$<BOOL:${USE_NPU}>:llm_data_dist_transfer.h>
     $<$<BOOL:${USE_NPU}>:spec_kv_cache_transfer.h>
     kv_cache_store.h
+    multi_tier_kv_cache_transfer.h
   SRCS
     embedding_allocator.cpp
     $<$<BOOL:${USE_NPU}>:hccl_kv_cache_transfer.cpp>
@@ -23,6 +24,7 @@ cc_library(
     $<$<BOOL:${USE_NPU}>:llm_data_dist_transfer.cpp>
     $<$<BOOL:${USE_NPU}>:spec_kv_cache_transfer.cpp>
     kv_cache_store.cpp
+    multi_tier_kv_cache_transfer.cpp
   DEPS
     :common
     $<$<BOOL:${USE_NPU}>:graph>

xllm/core/framework/kv_cache/kv_cache.cpp

@@ -35,6 +35,41 @@ torch::Tensor KVCache::get_k_cache() const { return key_cache_; }
 torch::Tensor KVCache::get_v_cache() const { return value_cache_; }
 torch::Tensor KVCache::get_index_cache() const { return index_cache_; }
 
+std::vector<std::vector<int64_t>> KVCache::get_shapes() {
+  std::vector<std::vector<int64_t>> tensor_shapes(3);
+  if (key_cache_.defined()) {
+    std::vector<int64_t> shape;
+    auto sizes = key_cache_.sizes();
+    shape.resize(sizes.size());
+    for (int i = 0; i < sizes.size(); ++i) {
+      shape[i] = sizes[i];
+    }
+    tensor_shapes[0] = std::move(shape);
+  }
+
+  if (value_cache_.defined() && key_cache_.numel() != 0) {
+    std::vector<int64_t> shape;
+    auto sizes = value_cache_.sizes();
+    shape.resize(sizes.size());
+    for (int i = 0; i < sizes.size(); ++i) {
+      shape[i] = sizes[i];
+    }
+    tensor_shapes[1] = std::move(shape);
+  }
+
+  if (index_cache_.defined() && index_cache_.numel() != 0) {
+    std::vector<int64_t> shape;
+    auto sizes = index_cache_.sizes();
+    shape.resize(sizes.size());
+    for (int i = 0; i < sizes.size(); ++i) {
+      shape[i] = sizes[i];
+    }
+    tensor_shapes[2] = std::move(shape);
+  }
+
+  return tensor_shapes;
+}
+
 void KVCache::swap_blocks(torch::Tensor& src_tensor,
                           torch::Tensor& dst_tensor) {
   // batch select keys and values

xllm/core/framework/kv_cache/kv_cache.h

@@ -40,6 +40,8 @@ class KVCache final {
   torch::Tensor get_v_cache() const;
   torch::Tensor get_index_cache() const;
 
+  std::vector<std::vector<int64_t>> get_shapes();
+
   std::shared_ptr<XTensor> get_k_xtensor() const;
   std::shared_ptr<XTensor> get_v_xtensor() const;