feat: restore D2H offload collection.

Author: Kang-Meng

xllm/core/framework/block/hierarchy_block_manager_pool.cpp

@@ -54,21 +54,101 @@ HierarchyBlockManagerPool::HierarchyBlockManagerPool(
 void HierarchyBlockManagerPool::deallocate(Sequence* sequence) {
   DCHECK(sequence != nullptr);
   int32_t dp_rank = BlockManagerPool::get_dp_rank(sequence);
+  // Publish device KV blocks into the device prefix cache first, so that
+  // offload-eligible blocks reach ref_count == 2 (held only by the sequence
+  // state and the prefix-cache node). deallocate_for_sequence below re-runs
+  // this cache step; it is idempotent (hashes recomputed from token ids find
+  // the nodes inserted here), so moved-out blocks do not corrupt the cache.
   BlockManagerPool::cache(sequence);
 
-  // Release host blocks if any
+  collect_offload_pairs(sequence, dp_rank);
+
+  // Release the host blocks still held by the sequence. Blocks moved into the
+  // offload queue are now invalid in this vector and are skipped by deallocate;
+  // their host ids stay reserved (held by the queue) until the D2H copy
+  // completes and the offload callback caches + frees them.
   auto host_blocks = sequence->host_kv_state().blocks(BlockType::KV);
   if (!host_blocks.empty()) {
     host_block_managers_[dp_rank]->deallocate(host_blocks);
   }
 
-  // Release device blocks via the composite (includes prefix cache flush)
+  // Release device blocks via the composite (includes prefix cache flush).
+  // Offloaded device blocks were moved out above, so the KV leaf skips them;
+  // the offload callback releases them once the copy is done.
   auto* composite =
       static_cast<CompositeBlockManager*>(block_managers_[dp_rank].get());
   composite->deallocate_for_sequence(sequence);
   sequence->reset();
 }
 
+void HierarchyBlockManagerPool::collect_offload_pairs(Sequence* sequence,
+                                                      int32_t dp_rank) {
+  if (!options_.enable_prefix_cache()) {
+    return;
+  }
+
+  std::vector<Block>* device_blocks =
+      sequence->kv_state().mutable_blocks(BlockType::KV);
+  std::vector<Block>* host_blocks =
+      sequence->host_kv_state().mutable_blocks(BlockType::KV);
+  if (device_blocks == nullptr || device_blocks->empty()) {
+    return;
+  }
+
+  const size_t block_size = options_.block_size();
+  const size_t cached_host_block_num =
+      sequence->host_kv_state().kv_cache_tokens_num() / block_size;
+  const size_t cached_device_block_num =
+      sequence->kv_state().kv_cache_tokens_num() / block_size;
+
+  const size_t host_block_num =
+      host_blocks == nullptr ? 0 : host_blocks->size();
+  // Host already holds at least as many blocks as the device computed: nothing
+  // new to offload.
+  if (host_block_num >= device_blocks->size()) {
+    return;
+  }
+
+  // Allocate the host blocks needed to receive the device blocks that have no
+  // host counterpart yet.
+  const size_t needed_block_num = cached_device_block_num > host_block_num
+                                      ? cached_device_block_num - host_block_num
+                                      : 0;
+  if (needed_block_num != 0) {
+    std::vector<Block> new_host_blocks =
+        host_block_managers_[dp_rank]->allocate(needed_block_num);
+    if (new_host_blocks.size() != needed_block_num) {
+      // Host pool exhausted; skip offload this round rather than partially
+      // copy.
+      return;
+    }
+    sequence->add_host_blocks(BlockType::KV, new_host_blocks);
+    host_blocks = sequence->host_kv_state().mutable_blocks(BlockType::KV);
+  }
+  if (host_blocks == nullptr) {
+    return;
+  }
+
+  // Only offload blocks that are fully computed on device. In-batch prefix
+  // cache insertion may register blocks before they are computed, so bound the
+  // offload range by cached_device_block_num to avoid copying uncomputed data.
+  const size_t offload_end_block_num = std::min(
+      {cached_device_block_num, host_blocks->size(), device_blocks->size()});
+  for (size_t i = cached_host_block_num; i < offload_end_block_num; i++) {
+    // ref_count == 2 means the block is held only by this sequence and the
+    // prefix-cache node, i.e. it is uniquely owned and safe to offload. Beam
+    // forks (shared blocks) have ref_count > 2 and are skipped.
+    if (device_blocks->at(i).ref_count() != 2) {
+      continue;
+    }
+    host_blocks->at(i).set_hash_value(
+        device_blocks->at(i).get_immutable_hash_value());
+    auto block_pair = std::make_shared<OffloadBlockPair>(
+        std::move(device_blocks->at(i)), std::move(host_blocks->at(i)));
+    offload_block_pair_queues_[dp_rank].enqueue(std::move(block_pair));
+  }
+}
+
 bool HierarchyBlockManagerPool::allocate(Sequence* sequence,
                                          size_t num_tokens,
                                          size_t max_copy_in_blocks_num) {
@@ -97,9 +177,20 @@ bool HierarchyBlockManagerPool::allocate(Sequence* sequence,
   }
   auto hbm_blocks = sequence->kv_state().blocks(BlockType::KV);
   auto host_blocks = sequence->host_kv_state().blocks(BlockType::KV);
-  for (size_t i = hbm_cache_token_num / options_.block_size();
-       i <
-       max_copy_in_blocks_num + (hbm_cache_token_num / options_.block_size());
+  // H2D copies host block i -> device block i, so i must index both vectors.
+  // The host prefix match (host_cache_token_num) is computed over the full
+  // prompt and can exceed the device blocks allocated for this (possibly
+  // chunked) num_tokens, so clamp the copy range to the blocks that actually
+  // exist on both sides to avoid out-of-bounds reads.
+  const size_t hbm_block_begin = hbm_cache_token_num / options_.block_size();
+  const size_t copy_block_limit =
+      std::min(hbm_blocks.size(), host_blocks.size());
+  if (hbm_block_begin + max_copy_in_blocks_num > copy_block_limit) {
+    max_copy_in_blocks_num = copy_block_limit > hbm_block_begin
+                                 ? copy_block_limit - hbm_block_begin
+                                 : 0;
+  }
+  for (size_t i = hbm_block_begin; i < max_copy_in_blocks_num + hbm_block_begin;
        i++) {
     load_block_transfer_infos_[dp_rank].emplace_back(
         BlockTransferInfo(host_blocks[i].id(),
@@ -111,9 +202,7 @@ bool HierarchyBlockManagerPool::allocate(Sequence* sequence,
   size_t target_hbm_cache_token_num =
       max_copy_in_blocks_num == 0
           ? hbm_cache_token_num
-          : (max_copy_in_blocks_num +
-             (hbm_cache_token_num / options_.block_size())) *
-                options_.block_size();
+          : (max_copy_in_blocks_num + hbm_block_begin) * options_.block_size();
 
   sequence->kv_state().incr_kv_cache_tokens_num(target_hbm_cache_token_num -
                                                 hbm_cache_token_num);
@@ -139,16 +228,27 @@ bool HierarchyBlockManagerPool::allocate(Sequence* sequence,
     auto hbm_blocks = sequence->kv_state().blocks(BlockType::KV);
     auto host_blocks = sequence->host_kv_state().blocks(BlockType::KV);
 
-    for (size_t i = hbm_cache_token_num / options_.block_size();
-         i < host_cache_token_num / options_.block_size();
-         i++) {
+    // H2D copies host block i -> device block i. host_cache_token_num is the
+    // host prefix match over the full prompt and can exceed the device blocks
+    // allocated for this (chunked) num_tokens, so clamp to the blocks present
+    // on both sides to avoid out-of-bounds reads on hbm_blocks.
+    const size_t copy_block_end =
+        std::min({host_cache_token_num / options_.block_size(),
+                  hbm_blocks.size(),
+                  host_blocks.size()});
+    const size_t hbm_block_begin = hbm_cache_token_num / options_.block_size();
+    for (size_t i = hbm_block_begin; i < copy_block_end; i++) {
       load_block_transfer_infos_[dp_rank].emplace_back(
           BlockTransferInfo(host_blocks[i].id(),
                             hbm_blocks[i].id(),
                             host_blocks[i].get_immutable_hash_value(),
                             TransferType::H2D));
     }
-    sequence->kv_state().incr_kv_cache_tokens_num(host_cache_token_num -
+    const size_t target_hbm_cache_token_num =
+        copy_block_end > hbm_block_begin
+            ? copy_block_end * options_.block_size()
+            : hbm_cache_token_num;
+    sequence->kv_state().incr_kv_cache_tokens_num(target_hbm_cache_token_num -
                                                   hbm_cache_token_num);
   }
   return true;

xllm/core/framework/block/hierarchy_block_manager_pool.h

@@ -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/kv_cache_transfer/hierarchy_kv_cache_transfer.cpp

@@ -24,6 +24,12 @@ namespace {
 
 constexpr uint32_t TIMEOUT_MS = 60000;
 
+// Streams reserved for concurrent D2H offload callers. D2H runs synchronously
+// on the RemoteWorker copy threadpool (4 threads, see remote_worker.h); reserve
+// one stream per such thread so concurrent offloads never block on the stream
+// queue.
+constexpr size_t kOffloadStreamCount = 4;
+
 std::vector<HierarchyKVCacheTransfer::LayerBatchRange> build_layer_batch_ranges(
     int64_t num_layers,
     uint32_t requested_batches) {
@@ -133,13 +139,12 @@ HierarchyKVCacheTransfer::HierarchyKVCacheTransfer(
       /*init_func=*/[this]() mutable { device_.set_device(); },
       /*cpu_binding=*/false,
       /*pool_name=*/"HierarchyKVCacheTransfer.load");
-  offload_threadpool_ = std::make_unique<ThreadPool>(
-      /*num_threads=*/5,
-      /*init_func=*/[this]() mutable { device_.set_device(); },
-      /*cpu_binding=*/false,
-      /*pool_name=*/"HierarchyKVCacheTransfer.offload");
-  for (int i = 0; i < load_threadpool_->size() + offload_threadpool_->size();
-       ++i) {
+  // D2H offload runs synchronously on the caller (RemoteWorker copy thread) so
+  // its copied-block count can be returned to the scheduler; it is not posted
+  // to a local pool. Size the shared stream pool to cover the H2D load threads
+  // plus the concurrent D2H callers.
+  const size_t num_streams = load_threadpool_->size() + kOffloadStreamCount;
+  for (size_t i = 0; i < num_streams; ++i) {
     copy_stream_.enqueue(device_.get_stream_from_pool(TIMEOUT_MS));
   }
 
@@ -300,11 +305,27 @@ uint32_t HierarchyKVCacheTransfer::transfer_kv_blocks(
     case TransferType::D2H2G:
       return offload(block_transfer_info);
     case TransferType::H2D: {
+      // Create and register the synchronizer synchronously, before scheduling
+      // the async copy. The scheduler issues transfer_kv_blocks before step, so
+      // by the time this returns the entry is in the map and the forward path's
+      // set_layer_synchronizer is guaranteed to find it. Registering inside the
+      // async load_from_host would race the forward lookup (it could miss the
+      // entry, skip the wait, and read KV before the H2D copy completed).
+      auto synchronizer = create_layer_synchronizer(
+          static_cast<int64_t>(layer_batch_ranges_.size()));
+      if (synchronizer == nullptr) {
+        LOG(ERROR) << "Failed to create layer synchronizer.";
+        return 0;
+      }
+      {
+        std::lock_guard<std::mutex> lock(mutex_);
+        layer_wise_load_synchronizer_[batch_id] = synchronizer;
+      }
       load_threadpool_->schedule(
           [this,
-           batch_id,
+           synchronizer,
            block_transfer_info = std::move(block_transfer_info)]() mutable {
-            load_from_host(batch_id, block_transfer_info);
+            load_from_host(synchronizer, block_transfer_info);
           });
       return 0;
     }
@@ -369,23 +390,13 @@ bool HierarchyKVCacheTransfer::offload_to_host(
 }
 
 bool HierarchyKVCacheTransfer::load_from_host(
-    uint64_t batch_id,
+    std::shared_ptr<LayerSynchronizer> synchronizer,
     const std::vector<BlockTransferInfo>& block_transfer_info) {
   if (block_transfer_info.empty()) {
     return true;
   }
 
-  auto synchronizer = create_layer_synchronizer(
-      static_cast<int64_t>(layer_batch_ranges_.size()));
-  if (synchronizer == nullptr) {
-    LOG(ERROR) << "Failed to create layer synchronizer.";
-    return false;
-  }
-  {
-    std::lock_guard<std::mutex> lock(mutex_);
-    layer_wise_load_synchronizer_[batch_id] = synchronizer;
-  }
-
+  CHECK(synchronizer != nullptr) << "layer synchronizer must not be null.";
   CHECK(batch_memcpy_ != nullptr) << "batch memcpy must be initialized.";
 
   std::unique_ptr<Stream> stream;
@@ -416,6 +427,12 @@ bool HierarchyKVCacheTransfer::load_from_host(
   }
 
   copy_stream_.enqueue(std::move(stream));
+  // On failure some ranges were never recorded; abort the synchronizer so a
+  // forward thread spinning on those layers unblocks and reports failure
+  // (aborting the forward) instead of hanging or reading uncopied KV cache.
+  if (!success) {
+    synchronizer->abort();
+  }
   return success;
 }
 

xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.h

@@ -96,15 +96,14 @@ class HierarchyKVCacheTransfer {
   uint32_t offload(const std::vector<BlockTransferInfo>& block_transfer_info);
   bool offload_to_host(Slice<BlockTransferInfo>& block_transfer_info);
   bool load_from_host(
-      uint64_t batch_id,
+      std::shared_ptr<LayerSynchronizer> synchronizer,
       const std::vector<BlockTransferInfo>& block_transfer_info);
 
  private:
   Options options_;
   Device device_;
 
   std::unique_ptr<ThreadPool> load_threadpool_;
-  std::unique_ptr<ThreadPool> offload_threadpool_;
   moodycamel::BlockingConcurrentQueue<std::unique_ptr<Stream>> copy_stream_;
 
   std::vector<xllm::KVCache>* kv_caches_ptr_ = nullptr;

xllm/core/platform/layer_synchronizer.cpp

@@ -45,6 +45,8 @@ class NPULayerSynchronizerAdapter final : public LayerSynchronizer {
     return true;
   }
 
+  void abort() override { impl_.abort(); }
+
   uint32_t size() const override {
     return const_cast<NPULayerSynchronizerImpl&>(impl_).get_event_size();
   }

xllm/core/platform/layer_synchronizer.h

@@ -28,6 +28,11 @@ class LayerSynchronizer {
 
   virtual bool synchronize_layer(int64_t layer_index) = 0;
   virtual bool record_stream(int64_t layer_index, Stream* stream) = 0;
+  // Force every layer's wait to unblock and report failure. Called when a copy
+  // fails so a forward thread spinning in synchronize_layer does not hang
+  // forever; synchronize_layer returns false after abort so the caller aborts
+  // the forward instead of reading not-yet-copied KV cache.
+  virtual void abort() = 0;
   virtual uint32_t size() const = 0;
 };
 

xllm/core/platform/npu/npu_layer_synchronizer.cpp

@@ -48,7 +48,14 @@ std::atomic<bool>* NPULayerSynchronizerImpl::get_event_flag(
 }
 
 bool NPULayerSynchronizerImpl::synchronize_layer(const int64_t layer_index) {
-  while (!event_record_flags_[layer_index].load(std::memory_order_acquire));
+  while (!event_record_flags_[layer_index].load(std::memory_order_acquire)) {
+    if (aborted_.load(std::memory_order_acquire)) {
+      return false;
+    }
+  }
+  if (aborted_.load(std::memory_order_acquire)) {
+    return false;
+  }
   auto ret = aclrtSynchronizeEventWithTimeout(events_[layer_index], timeout_);
   if (ret != ACL_SUCCESS) {
     LOG(ERROR) << "Synchronize event failed: " << ret;
@@ -69,4 +76,13 @@ bool NPULayerSynchronizerImpl::record_event(const int64_t layer_index,
   return true;
 }
 
+void NPULayerSynchronizerImpl::abort() {
+  // Set aborted_ before raising the flags so any thread released from the spin
+  // observes the abort and reports failure rather than syncing a stale event.
+  aborted_.store(true, std::memory_order_release);
+  for (size_t i = 0; i < event_record_flags_.size(); ++i) {
+    event_record_flags_[i].store(true, std::memory_order_release);
+  }
+}
+
 }  // namespace xllm

xllm/core/platform/npu/npu_layer_synchronizer.h

@@ -32,11 +32,14 @@ class NPULayerSynchronizerImpl {
   std::atomic<bool>* get_event_flag(const int64_t layer_index);
   bool synchronize_layer(const int64_t layer_index);
   bool record_event(const int64_t layer_index, const int32_t device_index);
+  // Unblock all pending synchronize_layer spins and make them report failure.
+  void abort();
   uint32_t get_event_size() { return events_.size(); };
 
  private:
   std::vector<aclrtEvent> events_;
   std::vector<std::atomic<bool>> event_record_flags_;
+  std::atomic<bool> aborted_{false};
   const int32_t timeout_;
 };
 

xllm/core/scheduler/chunked_prefill_scheduler.cpp

@@ -808,6 +808,14 @@ std::vector<Batch> ChunkedPrefillScheduler::prepare_batch() {
   if (!is_batches_empty) {
     // only update the scheduling latency when there are requests to process
     COUNTER_ADD(scheduling_latency_seconds, timer.elapsed_seconds());
+    // Dispatch the per-step KV transfers collected during allocate/deallocate:
+    // H2D loads recorded for host prefix-cache hits, and D2H offloads queued
+    // when sequences were deallocated. HierarchyBlockManagerPool overrides
+    // these; the default BlockManagerPool no-ops. Without this call the offload
+    // queue is never drained and its device/host blocks leak.
+    kv_cache_manager_->transfer_blocks(batches);
+  } else {
+    kv_cache_manager_->transfer_blocks();
   }
 
   GAUGE_SET(num_pending_requests,