[P/D disagg] Add HiRadixCache (hierarchical KV cache) for decode disaggregation - #26560
huanpengchu wants to merge 7 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for hierarchical caching (HiRadixCache) on the decode server in a disaggregated prefill/decode setup. It adds new configuration arguments, implements a two-phase prefetch and load-back mechanism in the decode preallocation queue, and adds compatibility checks to prevent its use with SWA or Mamba/SSM models. The review feedback highlights critical issues that need to be addressed: a potential lock leak in the HiCache path when matching prefixes, a state-tracking bug where node.backuped is not reset to False after freeing host memory (which could cause silent data loss during subsequent evictions), and a recommendation to use match_result.last_host_node directly to avoid potential attribute errors.
| if node.backuped: | ||
| self.cache_controller.mem_pool_host.free(node.host_value) | ||
| node.host_value = None |
There was a problem hiding this comment.
When restoring an evicted node, if it was backed up, its host value is freed and set to None. However, node.backuped is not reset to False.
If this node is evicted again in the future, the eviction logic might see node.backuped == True and skip backing it up to host memory, leading to silent data loss or crashes when attempting to load it back later. We must set node.backuped = False here.
| if node.backuped: | |
| self.cache_controller.mem_pool_host.free(node.host_value) | |
| node.host_value = None | |
| if node.backuped: | |
| self.cache_controller.mem_pool_host.free(node.host_value) | |
| node.host_value = None | |
| node.backuped = False |
| if new_node.backuped: | ||
| self.cache_controller.mem_pool_host.free(new_node.host_value) | ||
| new_node.host_value = None |
There was a problem hiding this comment.
Similarly to the above, new_node.backuped must be reset to False when its host value is freed to prevent skipping backup on subsequent evictions.
| if new_node.backuped: | |
| self.cache_controller.mem_pool_host.free(new_node.host_value) | |
| new_node.host_value = None | |
| if new_node.backuped: | |
| self.cache_controller.mem_pool_host.free(new_node.host_value) | |
| new_node.host_value = None | |
| new_node.backuped = False |
Add prefetch-based HiCache integration in the decode preallocation queue: - Two-phase prefix matching (read-only match + L2->L1 load back + lock) - L3->L2 async prefetch with timeout-based wait - Batch eviction before allocation to avoid per-request eviction overhead - Budget pre-check to skip wasteful DMA when request cannot fit - Permanent stuck detection and abort for oversized requests - HiCache abort cleanup at all failure points New CLI flags: --disaggregation-decode-enable-hicache --disaggregation-decode-hicache-ratio --disaggregation-decode-hicache-size --disaggregation-decode-hicache-write-policy Also fixes HiRadixCache bugs: - match_prefix_helper: break at evicted node boundary to prevent gaps - insert(): free host_value when restoring evicted nodes to prevent leaks - Remove noisy debug logs in init_load_back and check_prefetch_progress - Add SWA/SSM model incompatibility check for decode radix cache
Log L1 GPU hit, L2 host hit, L2 load_back tokens, and total prefix length after _load_back_and_lock completes.
backuped is a computed property (host_value is not None), so setting host_value = None already makes backuped return False. The explicit assignment caused AttributeError at runtime.
When a running request is retracted due to KV pool exhaustion, reset_for_retract() sets req.last_node = None. When the request is later resumed via resume_retracted_reqs(), _pre_alloc() is called without prefix matching, so last_node stays None. This causes cache_unfinished_req() to crash on dec_lock_ref(None). Fix: before _pre_alloc, point last_node to root_node and lock it so the subsequent dec_lock_ref in cache_unfinished_req is balanced.
When insert() restores an evicted+backuped node with new GPU data, preserve the existing host_value instead of freeing it. The host backup contains valid KV data (same tokens = same KV) and keeping it lets the node enter the optimal [GPU + Host] dual state directly, avoiding unnecessary write_backup cycles and host memory pool churn.
…lt directly - Move HiCache server_args handling from _handle_pd_disaggregation() in server_args.py to handle_pd_disaggregation() in pd_disaggregation_hook.py (aligned with upstream structure) - Remove duplicate _handle_pd_disaggregation method from server_args.py - Use match_result.last_host_node instead of req.last_host_node for clearer data flow in prefetch trigger
816d352 to
52fbe16
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces hierarchical cache (HiRadixCache) support for the decode server in prefill-decode (PD) disaggregation mode, including CLI configuration arguments, compatibility checks, and two-phase prefetching/load-back logic in the preallocation queue. Feedback highlights a bug where deferred requests fail to re-trigger prefetch due to an unreset prefetch_triggered_time, a missing implementation to free host_value in insert() as described in the PR description, and a missing validation check to prevent silent misconfiguration when enabling hicache without radix cache.
Prevent silent misconfiguration when --disaggregation-decode-enable-hicache is set without --disaggregation-decode-enable-radix-cache.
16e7086 to
2a4d1c1
Compare
ShangmingCai
left a comment
ShangmingCai
left a comment
There was a problem hiding this comment.
Could you check #26227? That PR is almost ready.
|
Thanks @huanpengchu. Closing this because it has had no updates in 100 days. Reopen it if the work is still relevant. Some directories moved recently, so an older branch may need retargeting: |
2 participants
Summary
In PD disaggregation, the decode worker can now use a hierarchical radix cache (L1 GPU + L2 host memory + Optional L3) to cache and reuse shared prefixes across requests. When a request arrives, the decode server first checks L1 (GPU radix cache), then L2 (host memory), and only requests the delta KV from prefill on a full miss.
This extends the existing decode radix cache (
--disaggregation-decode-enable-radix-cache, PR #19746) with a host-memory backing store that dramatically increases effective cache capacity.Enabled with
--disaggregation-decode-enable-hicacheon the decode server (requires--disaggregation-decode-enable-radix-cache).Motivation
The GPU-only decode radix cache is limited by GPU memory capacity. In a TP4/DP4 decode setup with
mem-fraction-static=0.85, each rank can hold ~162K tokens (~10 unique 16K-token prefixes). When the working set exceeds this (e.g., 50 distinct prefix groups), GPU cache hit rate drops below 20% and most of the benefit is lost.HiCache solves this by spilling evicted KV entries to host memory (L2) and loading them back via PCIe DMA on demand, providing a much larger effective cache without requiring additional GPU memory.
Main Changes
decode.py)hiradix_cache.py)match_prefix_helper: break at evicted node boundary to prevent gapsinsert(): freehost_valuewhen restoring evicted nodes to prevent leaksdecode.py)last_nodefor retracted requests so cache tree structure is preservedserver_args.py)--disaggregation-decode-enable-hicache--disaggregation-decode-hicache-ratio(L2/L1 size ratio, default 3.0)--disaggregation-decode-hicache-size(explicit L2 size in GB)--disaggregation-decode-hicache-write-policy(write_throughorwrite_back)Interface
Enable HiCache on the decode worker:
The
--disaggregation-decode-hicache-ratiocontrols L2 host memory size as a multiple of GPU KV capacity. With ratio=5.0 and ~40GB GPU KV, each rank allocates ~200GB host memory. For multi-DP setups, total host memory = ratio × GPU_KV × num_DP_ranks.Benchmark
Setup
generated-shared-prefix, 16K shared prefix + 4.5K suffix, 350 output tokens, concurrency 8，"10 groups × 50 prompts" means all 50 requests within each group share the same 16K-token prefix.Results — W1: High prefix reuse (10 groups × 50 prompts)
HiCache vs Radix-only: TTFT median -41%, throughput +5%
Results — W2: Low prefix reuse (50 groups × 10 prompts)
HiCache vs Radix-only: TTFT median -26%, throughput +3%
In the low-reuse scenario (50 prefix groups), GPU-only radix cache is worse than baseline because eviction churn dominates. HiCache's L2 host store absorbs the overflow and maintains a positive hit rate.
Cross-node PD: where HiCache shines most
The benchmarks above are single-node (prefill and decode on the same machine, RDMA over local InfiniBand). In production cross-node PD deployments, the benefit of HiCache is expected to be significantly larger .
Test Plan
CI States
Latest PR Test (Base): ❌ Missing
run-cilabel -- add it to run CI tests.Latest PR Test (Extra): ❌ Blocked --
run-ciis required first.