Summary of Changes

Hello @pansicheng, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant enhancement to the KV cache management system by integrating disk storage into the existing HiRadixCache. This "Hidisk" feature expands the total capacity of the KV cache, allowing for larger context windows or more concurrent requests by offloading less frequently accessed KV pairs to disk. The changes include a multi-tiered caching strategy with synchronous writes to memory and disk, and intelligent reads that leverage temporal locality, resulting in improved cache hit rates and overall throughput, as demonstrated by the provided benchmarks.

Highlights

• Disk-Backed Hierarchical Cache: Introduced a new 'Hidisk' feature that integrates disk storage with the existing HiRadixCache, significantly expanding the KV cache capacity. This allows for larger context windows and improved performance by offloading less frequently accessed KV pairs to disk.
• Multi-Tiered Cache Management: Implemented a sophisticated multi-tiered caching strategy. KV data is synchronously written to both memory and disk, and read operations intelligently fetch data from memory or disk into GPU memory, with a portion of disk data copied back to host memory to leverage temporal locality.
• Dedicated Disk I/O Control: Added a new HiCacheControllerDisk class with dedicated threads and queues for managing disk I/O operations (writing from device to disk, loading from disk to host, and loading from disk to device), ensuring efficient and asynchronous data transfers.
• Configurable Disk Cache: Introduced new command-line arguments (--hicache-use-disk, --hicache-disk-path, --hicache-disk-ratio, --hicache-disk-size) to enable and configure the disk-backed hierarchical cache, providing flexibility for deployment and tuning.
• Enhanced Memory Management and Eviction: Refactored the cache eviction and load-back policies to intelligently manage KV cache across device, host, and disk tiers. The system now proactively evicts data to lower tiers (host then disk) when GPU memory is constrained and loads data back as needed.
• Performance Benchmarking: Included a new benchmarking script (disk.sh) and updated the multi-turn benchmark to demonstrate the performance benefits of the disk-backed cache, showing improved cache hit rates and overall throughput compared to memory-only caching.

Code Review

This pull request introduces disk storage integration into the hierarchical KV cache, significantly expanding its capacity. The changes involve adding new classes and methods to manage disk operations, updating the cache controller and radix tree logic to handle the new disk layer, and adding configuration options via server arguments. The benchmark results provided demonstrate the potential performance benefits of this feature. The code is generally well-structured, using inheritance and helper methods to manage complexity. There are a few areas identified for potential improvement related to correctness (handling device indices with overlap schedule) and maintainability (clarity of certain logic flows and documentation).