Repository navigation
[Feature] Add SLRU eviction policy & fix RadixCache hit_count bug - #18843
Conversation
Summary of ChangesHello @liubiyongge, I'm Gemini Code Assist1! 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 significantly improves cache management for RAG and batch processing by introducing the Segmented LRU (SLRU) eviction policy. This policy is designed to prevent 'scan' workloads from flushing valuable hot data, thereby reducing cache misses and tail latency. A crucial underlying bug in the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
The pull request introduces the SLRU (Segmented LRU) eviction policy to improve scan resistance in the KV cache and aims to fix a bug where hit_count was not being incremented. While the addition of the SLRU strategy and the increment logic in _match_prefix_helper are correct, the implementation of frequency tracking is incomplete. Specifically, the hit_count is not inherited during node splits, and it is not incremented during insertions. These omissions will lead to inaccurate frequency data, particularly for shared prefixes which are critical for the effectiveness of SLRU and LFU policies. Additionally, some user-facing error messages and help texts were not updated to include the new policy.
- Added `SLRUStrategy` (Segmented LRU) to `eviction_policy.py` for scan resistance. - Fixed a critical bug in `radix_cache.py` where `hit_count` was never incremented during match/insert. - Fixed a typo in `server_args.py` (`add_radix_eviction_policy_choices` -> `radix_eviction_policy`) causing AttributeError. - Registered 'slru' as a valid choice for `--radix-eviction-policy`. Benchmarks show SLRU reduces tail latency by 7.6x (550ms -> 73ms) under heavy RAG scan workloads compared to LRU.
f964fd4 to
6aa0547
Compare
hzh0425
left a comment
There was a problem hiding this comment.
Looks good
Could you add a ci accuracy test ?
|
/tag-and-rerun-ci |
- Create unit test to verify SLRU eviction policy works correctly - Test includes setup and execution phases to validate that: * High frequency access keys are retained in cache * Low frequency access keys are evicted when capacity is exceeded - Use realistic cache sizes to trigger actual eviction behavior - Validate through device indices presence in match results
0f9a616 to
c44212c
Compare
For partial matches (prefix_len < len(child.key)), only the newly created split node should have its hit_count incremented, not the child node. Previously, hit_count was incremented unconditionally before checking the match type.
Avoid self-referencing hit_count inflation where chunked requests increment hit_count on nodes they created in previous chunks. Add a new helper method that conditionally updates hit_count based on the chunked flag.
|
@liubiyongge could you help check the failde ci |
The test test_hicache_storage_file_backend.py::TestHiCacheStorageAccuracy::test_eval_accuracy failed in CI, but passed locally. Despite a careful code review, the issue remains unclear. |
…l-project#18843) Co-authored-by: zhangheng <hzh0425@apache.org>
…l-project#18843) Co-authored-by: zhangheng <hzh0425@apache.org>
…l-project#18843) Co-authored-by: zhangheng <hzh0425@apache.org>
PR sgl-project#18843 added the chunked guard so a chunked request does not inflate hit_count on radix nodes it created, but only the middle-chunk insert (scheduler.py) passed chunked=True. The final prefill chunk caches the unfinished request via batch_result_processor without the flag, so its prefix nodes were bumped once here and again by cache_finished_req, yielding hit_count=2. Pass chunked=True so each node is bumped exactly once on completion.
…ng it Revert the earlier chunked=True at the prefill-completion cache site in batch_result_processor: that branch runs for every request finishing prefill (chunked final chunk and non-chunked single prefill alike), so forcing chunked=True suppressed the legitimate hit_count bump on the general prefill path, not just chunked self-inserts. A request inserts its prompt prefix twice -- via cache_unfinished_req when prefill completes and via cache_finished_req on finish -- while decode nodes are inserted once, so prompt nodes settle at hit_count 2 and decode nodes at 1, identically for chunked and non-chunked prefill. sgl-project#18843's middle-chunk chunked=True still prevents the unbounded early-node inflation. Replace the old all==1 expectation with side-by-side chunked and non-chunked tests asserting the hit_count value set is exactly {1, 2}.
…l-project#18843) Co-authored-by: zhangheng <hzh0425@apache.org>
Motivation
In RAG (Retrieval-Augmented Generation) and batch processing scenarios, the system often faces "scan" workloads: large volumes of one-time requests (e.g., retrieving many documents). Under the default
LRUStrategy, these one-time requests flush valuable hot data (e.g., shared System Prompts, multi-turn dialog history) out of the KV cache. This leads to frequent re-computation (Cache Miss) and high tail latency.While implementing SLRU (Segmented LRU) to solve this, I discovered a critical bug in
RadixCache: thenode.hit_countwas never incremented duringmatch_prefixorinsertoperations. This rendered any frequency-based policies (like the existing LFU) ineffective.Modifications
New Strategy (
sglang/srt/mem_cache/eviction_policy.py):SLRUStrategy(Segmented LRU).hit_count < k(default 2). Evicted first.hit_count >= k. Protected from scan noise.Critical Bug Fix (
sglang/srt/managers/router/radix_cache.py):node.hit_countinmatch_prefix(on cache hit) andinsert(on new node creation). This fixes the underlying mechanism for LFU/SLRU.Bug Fix (
sglang/srt/server_args.py):add_radix_eviction_policy_choices->radix_eviction_policy. This fixes anAttributeErrorwhen accessingServerArgs.radix_eviction_policy.'slru'as a valid choice for the CLI argument--radix-eviction-policy.Accuracy Tests
This PR affects the cache eviction policy and memory management logic. It does not modify the model architecture, kernel implementation, or arithmetic precision. Therefore, the model output accuracy remains unaffected.
Benchmarking and Profiling
I performed a "Scan Attack" benchmark to verify Scan Resistance.
mem_fraction_static=0.78(Simulating memory pressure with ~11k token capacity).Results:
The results demonstrate that
SLRUeffectively isolates hot data from scan traffic, maintaining consistent low latency (TTFT).Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci