Repository navigation
fix: release the prefix cache when the KV memory limit is lowered - #552
Conversation
d61241f to
5b7aa88
Compare
There was a problem hiding this comment.
🟡 Changes recommended
Valid zero and preexisting resize targets can remain unapplied.
2 open findings
What changed in this PR
Enables deferred KV-limit reductions to reclaim prefix-cache pages during traffic and idle periods.
Changes:
- Adds shrink reclaimers and retirement-aware allocation.
- Integrates idle limit handling with vLLM and SGLang.
- Adds CPU coverage and updates memory-control documentation.
| File | Description |
|---|---|
kvcached/kv_cache_manager.py |
Coordinates deferred shrink reclamation. |
kvcached/integration/vllm/patches.py |
Adds prefix eviction and idle limit watcher. |
kvcached/integration/vllm/page_eviction.py |
Limits page-aware victim selection by page count. |
kvcached/integration/sglang/patches.py |
Adds SGLang prefix reclaimers and idle processing. |
kvcached/integration/sglang/autopatch.py |
Registers the scheduler idle patch. |
tests/test_prefix_cache_shrink.py |
Tests shrinking, traffic, idle, and parallel behavior. |
tests/manifests/cpu.txt |
Includes the new CPU test suite. |
examples/09_prefix_caching/README.md |
Documents memory-limit precedence. |
examples/02_memory_control/README.md |
Documents shrink behavior and exceptions. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
2a9e1a7 to
e6297cc
Compare
|
Thanks for contributing, @Jeffwan! We will take a closer look. |
|
Could you add real-GPU vLLM TP and PP validation before merging? This changes when cached pages are reclaimed, so the single-GPU results do not establish that release remains correctly ordered across workers. TP=2/PP=1 and TP=1/PP=2 would cover the two paths. Please lower the limit while batches are queued and prefix-cache hits are active, finish or cancel requests, then check that all ranks release only after the worker barrier and that mapped usage converges while idle. Please exercise both A delayed or failed rank during release should not lead to a false |
shipiyouniao
left a comment
There was a problem hiding this comment.
I found three shrink-path cases that need fixing, detailed inline. I reproduced them in CPU control-flow tests using the PR's fake page allocator and engine/cache doubles; these are not GPU results. The existing 16 tests passed, while the added cases produced four failing parameterizations and one async-only skip.
rishabhsinha17
left a comment
There was a problem hiding this comment.
this belongs in kvcached. the limit is a kvcached construct, kvctl limit and set_memory_limit, with no counterpart in vllm or sglang. only this layer can turn a lowered limit into prefix cache evictions, and without it kvctl limit on a prefix caching engine never converges while idle. our 0.5.20 validation saw the limit hold through live reclamation only (finding 5 in #509).
not ready at e6297cc, one blocker on top of shipiyouniao's three, on the same path as his first. a direct api call evicts the engine's prefix cache on the caller's thread. set_memory_limit enters _resize, which now runs the reclaimers (the callbacks that evict cached pages) inline, and the memory control README tells controllers to call it from a handler in the engine process, which runs off the engine thread. i reproduced it with your fixtures. fill the cache to 9 pages, call set_memory_limit for 4 pages from a thread named control-handler, and _evict_block_ids runs on control-handler while the status comes back applied. your watcher docstring already names the engine thread as the only one allowed to change the block pool, the sglang reclaimer evicts the radix cache (its prefix cache) the same way, and before this pr the call touched allocator and manager state only. one fix covers both findings. let set_memory_limit publish the shrink and hand eviction and the drain of retired pages to the engine thread, through the idle callback and wakeup the watcher already uses, and the idle hook on sglang.
everything else checked out. victims returns only pages whose whole occupancy is evictable and rechecks each page before selection, a running request keeps its blocks, and max_pages only caps #538's heap walk. _idle_state_callbacks exists on 0.17 through 0.30 and WAKEUP from 0.20, and vllm runs those callbacks on the engine thread while has_work is false, 0.17 to 0.19 skip the watcher. the C++ allocator retains the last published target, the seen check reports once per change, and kvctl limit 0 is the zero case shipiyouniao flagged.
your 16 tests pass at e6297cc on my mac, cpu manifest 1248 passed with the failure set identical to main, and the eviction and capacity guard files are green at your head. a10 numbers stay yours, happy to push the thread repro as a test if useful.
|
I have not tested TP/PP yet. I will make it later. |
|
/ci run |
Nightly correctness · 2026-10-09 · FAIL (14 of 20 cases)Run · kvcached e6297cc (merged into main at 0222845)
Setup of every caseOn request in a pull request (gpu-pr.yml) on 2x NVIDIA L4 (23034 MiB, driver 580.178.04); the SGLang and vLLM cases run at the same time, one GPU each. kvcached is installed from the pull request merged into main into each engine's official image; each case starts one server in it.
CasesEach kvcached case is compared with the native case of the same engine and model.
Each case:
Elastic tests
Logs: artifact, kept 30 days: serve.log and cmd.txt per case, summary.json with every request's tokens, setup.json, run.log, GPU state at the end. |
|
GPU CI result for e6297cc merged into main at 0222845 (run), on 2x L4: The kvcached cases match native's greedy tokens and
None of these cases lowers the KV memory limit, so this run shows that normal serving still works, but it does not exercise the path this PR changes. |
Address review on ovg-project#552: - resize() and set_memory_limit() only record a shrink, as a controller may call them from its own thread; alloc(), request completion and the idle wakeup evict on the engine thread - the vLLM idle wakeup also follows set_memory_limit() revisions and drains the retired pages - SGLang registers its reclaimers when the scheduler is built, before any idle iteration - a published resize target of 0 is applied; -1 means none (test fakes now return -1 like the C++ allocator) Signed-off-by: Jiaxin Shan <seedjeffwan@gmail.com>
A lower limit (kvctl limit / set_memory_limit) stayed deferred while cached prefixes held its pages, idle or under traffic: allocations reuse cached slots one LRU victim at a time and never empty a page, and a kvctl target was only applied by the next allocation. - resize() only records a shrink; the prefix cache evicts whole pages on the engine thread (alloc, request completion, idle), and pages retired for the shrink are not reallocated - idle engines apply new limits: a vLLM EngineCore wakeup that follows kvctl targets and set_memory_limit() revisions, and an SGLang scheduler idle hook (SGLang TP/PP > 1 unchanged) - a published resize target of 0 is applied; -1 means none - tests: tests/test_prefix_cache_shrink.py (CPU) and tests/e2e/limit.py (nightly GPU profile) Signed-off-by: Jiaxin Shan <seedjeffwan@gmail.com>
9cb8ad0 to
a84b6d8
Compare
|
@rishabhsinha17 thanks. eviction no longer runs inside I finished the PP/TP tests.
every non-cancelled request succeeded, and no run had a traceback, CUDA/NCCL error or worker restart before shutdown. |
shipiyouniao
left a comment
There was a problem hiding this comment.
LGTM from my side! We still need to land a few other PRs to sort out CI before merging this one. Thanks for your patience :)
RixinLiu
left a comment
There was a problem hiding this comment.
Overall LGTM. But I think there's a problem when kvctl limit and set_memory_limit() are both used. With the fakes in test_prefix_cache_shrink.py: kvctl limit to 12 pages (9 mapped), then set_memory_limit() to 4 pages, which stays deferred. When the idle callback runs, apply_resize_target() applies the old kvctl target again, so the shrink gets cancelled and memory_limit_state() says applied while 9 pages are still mapped.
But this problem is not related to this PR. No blocking from my side.
rishabhsinha17
left a comment
There was a problem hiding this comment.
LGTM at a84b6d8. _resize now only records the shrink. eviction runs in alloc() and reclaim_for_shrink(), called on the engine thread by the vllm idle callback and the sglang on_idle wrapper. no controller thread touches the block pool.
my repro passes. set_memory_limit for 4 pages from a thread named control-handler evicts nothing and returns deferred, the idle pass then evicts on the engine thread to applied, 4 mapped. the reclaim call put back into _resize turns test_set_memory_limit_leaves_eviction_to_the_engine_thread red. shipiyouniao's three hold, revision followed, reclaimer registered after Scheduler.init, zero reaches resize.
your file 23 passed, cpu manifest 1293 with the failure set identical to main, the #469, #511 and #538 files green. rixin's kvctl then api case is pre-existing, main's alloc() re-applies the kept kvctl target the same way, a two line follow-up applies a published target once, not every pass. a10 and a40 stay yours.
|
the kvctl then api case from rixin's note is now #558. pre-existing at main, alloc() re-applied the retained kvctl quota on every pass, so two lines in that block, _apply_resize_target at this head, apply a published quota once. |
|
Merged |
|
thanks for all your help! I will put more efforts other issues from my side. |
…nto _apply_resize_target ovg-project#552 extracted the kvctl block of _alloc_impl() into _apply_resize_target() and accepts a zero quota. The applied-quota check and its record move into that helper unchanged, so alloc() and the idle pass share them. Signed-off-by: Rishabh Sinha <rsinha17@terpmail.umd.edu>
…arget shipiyouniao's review of ovg-project#558 found that applying the quota once stalls a kvctl shrink the pool refused while reserved pages were the only pages in use. _resize() frees the reserved blocks before it enters the shrink, so that free() cannot complete it, and the next allocation skipped the retained quota. At main the next allocation applied it again and the shrink completed. The control published last is the one in force now. _apply_resize_target() applies the retained quota on every pass while kvctl published last, as at main, and skips it once resize() published a later target, which set_memory_limit() calls. A changed kvctl value is a new publication and takes the limit back. clear() resets nothing, since it publishes nothing. The -1 and zero handling from ovg-project#552 stays. Refs ovg-project#552 Signed-off-by: Rishabh Sinha <rsinha17@terpmail.umd.edu>

Lowering a kvcached limit did not release prefix-cache pages: the shrink stayed
deferredwhile idle and under traffic, and akvctl limitwas only applied at the next allocation.When a resize is deferred, kvcached now asks the prefix cache to evict whole pages until the instance fits, and applies
kvctllimits while the engine is idle (vLLM EngineCore wakeup, SGLang scheduler idle hook). Pages of running requests are never revoked. SGLang with TP/PP > 1 is unchanged.A10, Qwen3-4B, MiB mapped after
kvctl limit: