[PD]: Support HiCache prefetching and pd-incremental transfer on decode side - #26227
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 implements hierarchical cache (HiCache) support for the disaggregated decode path, enabling KV cache restoration across L1 (device), L2 (host), and L3 (storage) layers. It introduces the DecodePrefixMatch and HiCacheRestoreResult structures to manage cache hit metadata and the state of local restoration. The core logic involves querying storage hit lengths during prefix matching and a multi-phase state machine in the decode manager to drive prefetching and load-back operations. Review feedback identifies critical performance bottlenecks caused by distributed synchronization (all_reduce) occurring within per-request loops, which should be batched or moved outside the loops. Additionally, the reviewer noted redundant prefix matching logic and potential throughput issues due to the serialization of L2-to-L1 restore operations.
| if self.scheduler.metrics_reporter.enable_metrics: | ||
| self.scheduler.metrics_collector.increment_transfer_failed_reqs() | ||
| continue | ||
| elif poll == KVPoll.Success: |
There was a problem hiding this comment.
Is there a reason why we change it from elif to if?
| if poll == KVPoll.Failed: | ||
| try: | ||
| decode_req.kv_receiver.failure_exception() | ||
| except Exception as e: | ||
| error_message += f" with exception {e}" |
There was a problem hiding this comment.
What could be the reason of HiCacheRestoreResult.FAILED? Should we append some error_message detail as well?
There was a problem hiding this comment.
Add a failed log in decode_hicache_mixin
ShangmingCai
left a comment
ShangmingCai
left a comment
There was a problem hiding this comment.
No other comments.
Co-authored-by: Shangming Cai <csmthu@gmail.com> Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
67027ae to
00c2ced
Compare
|
/rerun-test test/registered/disaggregation/test_disaggregation_decode_radix_cache.py |
|
Results for 🚀 |
| operation | ||
| ) | ||
| storage_hit_count_tensor = torch.tensor(storage_hit_count, dtype=torch.int) | ||
| self._all_reduce_attn_groups( |
There was a problem hiding this comment.
We could revisit this in the future if this per-request all_reduce causes a performance issue for large batch sizes.
|
CI has passed. |
…de side (sgl-project#26227) Co-authored-by: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Co-authored-by: Shangming Cai <csmthu@gmail.com> Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
…de side (#26227) Co-authored-by: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Co-authored-by: Shangming Cai <csmthu@gmail.com> Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
|
@hzh0425 I have a question: when the KV cache of the decode node is in L3, even if the cache hits in L3, it still needs to be pulled from L3 to the CPU, and then from the CPU to the GPU. The time spent in this process should be greater than the overhead of transferring from P to D, right? |
@weibingo Yeah, but the matched prefix KV fetching could be overlapped with the computing and transferring of the incremental part, so this design is better for multi-turn conversations and agentic use cases. |
The conditional re-match for requests whose reset_for_retract cleared last_node was introduced on main by PR #26227 together with the TODO(retraction) note in _pre_alloc. The merge of main into the stateless-scheduler branch kept the TODO but dropped the conditional, unconditionally passing tree_cache=None on the decode-radix path, so retracted-resumed requests lost prefix reuse. Restore the base behavior and its explanatory comment.
…de side (sgl-project#26227) Co-authored-by: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Co-authored-by: Shangming Cai <csmthu@gmail.com> Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
…de side (sgl-project#26227) Co-authored-by: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Co-authored-by: Shangming Cai <csmthu@gmail.com> Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
6 participants
Motivation
Roadmap: #21846 、 #21703
Still in progress, TODO:
Modifications
Accuracy Tests
Speed Tests and Profiling
Qwen32B + PD + HiCache + Mooncake
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ✅ Run #26766211033
Latest PR Test (Extra): ✅ Run #26766210099