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

[Kang-Meng]

Background

Previously HierarchyBlockManagerPool was selected when host_blocks_factor > 1 but the
actual D2H/H2D transfer path was disabled: blocks were never offloaded to host and host
prefix-cache hits were never reloaded to device. This PR re-enables the full offload→reload
pipeline for BlockType::KV.

---

Changes

Commit 1 — allocate host KV cache as page-aligned pinned memory in KVCache impls (ae8ba1a1)

Ensures host-side tensors used for D2H/H2D DMA are allocated as pinned, page-aligned memory:

• HostPageAlignedRegion: on NPU builds, calls
  aclrtHostRegister(..., ACL_HOST_REGISTER_MAPPED, ...) immediately after mmap+mlock
  so the host buffer is mapped into device address space. Destructor calls
  aclrtHostUnregister before munlock/munmap. Move-assignment operator cleaned up (was
  leaking if this held existing memory).
• New utility functions in kv_cache_utils: scale_host_block_count (clamps factor to
  ≥ 1.0), build_host_tensor_shape (scales dim-0 by factor),
  build_host_group_tensor_shape (inserts a layer dimension at index 1, yielding
  [host_blocks, layer_count, ...]), and create_host_page_aligned_tensor (allocates
  region + creates torch::from_blob view).
• KVCacheImpl and all impl subclasses (Standard, DeepSeekV4, Indexed, LinearAttention)
  gain a host constructor (KVCacheShape, KVCacheCreateOptions, BlockType, layer_count)
  that allocates the grouped host tensor for its block type, and a
  get_block_type_tensors(BlockType) virtual method queried by the transfer layer.
• Unit test in kv_cache_test.cpp covers host tensor allocation and shape contract for
  each impl.

Commit 2 — restore D2H offload collection (cbcbc2f1)

Restores the D2H offload trigger (removed in commit 1 to keep the diff reviewable) and
fixes a race in the H2D path:

• collect_offload_pairs(): extracted from deallocate(). Allocates host blocks to match
  not-yet-offloaded device blocks (aborts without partial copy if the host pool is
  exhausted), then enqueues OffloadBlockPairs for blocks with ref_count == 2 (uniquely
  owned; beam-shared blocks with ref_count > 2 are skipped). Bounded by
  kv_cache_tokens_num so uncomputed in-batch blocks are never copied.
• H2D race fix: LayerSynchronizer is now created and inserted into
  layer_wise_load_synchronizer_ synchronously inside transfer_kv_blocks(), before
  the async copy task is scheduled. Previously the synchronizer was created inside the
  threadpool task; the forward's set_layer_synchronizer() could look up the map before
  the task ran, miss the entry, skip the wait, and read KV before H2D completed.
• offload_threadpool_ removed: D2H now runs synchronously on the caller (RemoteWorker
  copy thread) so the copied-block count can be returned to the scheduler. The shared
  stream pool is sized to load_threadpool_.size() + kOffloadStreamCount (4) to cover
  concurrent D2H callers.
• Two bounds guards on H2D copy loops: copy range clamped to
  min(hbm_blocks.size(), host_blocks.size()) to prevent out-of-bounds when a
  chunked-prefill allocation produces fewer device blocks than the host prefix match covers.

---

Test

Validated on A3 (Ascend 910) with Qwen3-32B TP2, MAX_MEM_UTIL=0.70 (device pool: 480
blocks), generated-shared-prefix workload (60 groups × 8 prompts, 2560-token shared
prefix, concurrency=4):

	factor=4 (host cache)	factor=0 (device only)
hit_rate	50.6%	50.7%
HIT TTFT	141 ms	147 ms
MISS TTFT	840 ms	974 ms
Mean TTFT	503 ms	572 ms
mem-fails	0	0

• Activation confirmed: HierarchyKVCacheTransfer: created host cache for 1 block type groups.
• Host-backed hit confirmed: repeated long prompt logs num_prefix_cache_tokens: 256,
  TTFT 74 ms → 45 ms, byte-consistent output.
• No device-block leak: 0 No enough memory to schedule events; idle recovery HTTP 200
  after full workload.
• Accuracy: finish_reason=stop, coherent French output, idle and under load.

[gemini-code-assist]

Code Review

This pull request re-enables and refactors the hierarchical KV cache transfer and block manager pool, which were temporarily disabled during a block-manager refactor. It introduces platform-independent abstractions like BatchMemcpy and LayerSynchronizer (with NPU implementations), updates various KVCacheImpl subclasses to support host cache allocation, and integrates these components back into WorkerImpl and LLMEngine. The code review feedback focuses on enforcing repository style guides, including reserving vector capacity, using fixed-width integers instead of plain int, annotating constant arguments with parameter names, avoiding relative include paths, avoiding auto for primitive types, and adding defensive bounds checks for tensor indexing.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.