feat: add revisioned instance memory limits

[shipiyouniao]

Baseline

Built on current main, including the coherent allocator accounting from #443.

Summary

• add a revisioned instance-level memory limit, split deterministically across all registered KV pools by their original capacity
• apply each pool share immediately through the existing KVCacheManager.resize() path, including while the instance is idle
• report explicit applied, deferred, stale, conflict, or unavailable status together with revision and byte accounting
• reuse the existing in_shrink, free, trim, and reclaimed-page behavior for downward convergence and later regrowth
• preserve active mappings when a lower limit cannot be applied immediately
• align limits to physical page bundles and clamp growth to the original KV reservation

Relationship to resize

This PR no longer introduces a separate physical-page limit or a second shrink state machine.

create_kv_tensors() owns the CUDA virtual-address reservation. resize() only changes allocator page bookkeeping and physical backing: shrinking moves page ids to the reclaimed list and unmaps releasable pages, while growing reuses those ids without rebuilding the KV tensors. In the current allocator, num_total_pages_ already bounds maximum physical residency.

The new code therefore adds the missing control protocol around the existing mechanism:

• revision and acknowledgement
• immediate application while idle
• observable deferred convergence
• one instance budget split across multiple pools

Ownership boundary

The caller owns quota policy, fairness, admission, and cross-instance arbitration. kvcached stores and applies the assigned limit next to the existing resize state machine.

Validation

Check	Result
Focused memory-limit and resize tests	11 passed
MyPy on changed source files	passed
Ruff check	passed
Test classification	valid: 29 total, 17 CPU, 7 GPU, 5 integration

The focused tests cover immediate idle shrink, deferred shrink without active-page revocation, convergence after usage drops, regrowth within the original reservation, maximum-capacity clamping, zero-limit drain, revision semantics, and multi-pool splitting.

All 11 GitHub CPU, MyPy, and pre-commit checks passed on the rewritten head.

Related design discussion: #375.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[Copilot]

Pull request overview

This PR adds an instance-level, revisioned physical-GPU memory cap that can be assigned by an external integration and then enforced locally by each live KV pool (both in Python and in the native page allocator), including support for shrinking without revoking in-use mappings.

Changes:

• Added a provider-owned, revisioned per-pool physical mapping limit in KVCacheManager, and updated available_size() to treat an active limit as the allocator-local “fast path” (no GPU free-memory query).
• Introduced a new kvcached.control registry (weakrefs) to split/apply an instance-wide cap across all registered pools, and wired vLLM/SGLang integrations to register/clear pools on lifecycle events.
• Extended the C++ PageAllocator + PyTorch bindings to support physical page limits, remaining-capacity queries, and a native per-GPU “growth guard” using a PCI-bus-ID file lock and a final free-memory headroom check.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/test_physical_memory_limit.py	Adds focused unit coverage for revisioning, alignment, deferred shrink semantics, and control-plane split behavior.
kvcached/kv_cache_manager.py	Implements per-pool physical limit state, revision handling, and available_size() behavior under active limits.
kvcached/integration/vllm/interfaces.py	Registers KV managers as physical-memory pools and clears them on shutdown.
kvcached/integration/sglang/interfaces.py	Registers KV managers as physical-memory pools and clears them on shutdown.
kvcached/control.py	New control-plane module for weakref pool registry + deterministic limit splitting and application.
csrc/torch_bindings.cpp	Exposes new PageAllocator limit and accounting APIs to Python via pybind.
csrc/page_allocator.cpp	Enforces a physical page limit in foreground alloc and background prealloc; adds limit/accounting helpers.
csrc/inc/page_allocator.hpp	Declares the new physical-limit APIs and internal helpers on PageAllocator.
csrc/inc/gpu_vmm.hpp	Adds a portable (CUDA/HIP) PCI bus ID query wrapper for per-GPU locking.
csrc/inc/allocator.hpp	Extends allocator state to track bytes-per-offset for guarded mapping.
csrc/allocator.cpp	Adds the cross-process “physical growth guard” and overflow-checked transaction sizing for VMM mappings.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[RixinLiu]

Thanks! Before comparing designs I walked the current resize() path end to end, so we're working from the same baseline. Short version: the three downward behaviors you describe are already implemented today. What is missing is the trigger and the acknowledgement.

What happens today when the limit is lowered

kvctl limit writes shared memory
        ↓                                                       page_allocator.cpp:801
_alloc reads it                                          kv_cache_manager.py:267 
        ↓
resize() → PageAllocator::resize()
        ↓  rejected: target below current in-use         page_allocator.cpp:323
in_shrink = True, target = N                             kv_cache_manager.py:449
        ↓
available_size() drops **its new-page term to 0**            kv_cache_manager.py:464
        ↓
_alloc gate refuses anything needing a new page          kv_cache_manager.py:271
        ↓
free() re-checks after every release                     kv_cache_manager.py:399
        ↓  once usage is at or below N
real resize; reserved pages unmapped                     page_allocator.cpp:382

Two notes on this path:

"Never revoke active mappings" is structural, the shrink path only ever draws from free_page_list_ and reserved_page_list_, never from the in-use set.

"Block additional physical growth until mapped usage falls below the cap" is also achieved by available_size()

One correction on page states

There is no "unused" page category distinct from reserved. A page is in exactly one of four states:

State	Where	Has page id	Has physical memory	Counted in num_free_pages_
free	free_page_list_	yes	no	yes
reserved (prealloc)	reserved_page_list_	yes	yes	yes
in use	handed to KVCacheManager	yes	yes	no
reclaimed	reclaimed_page_list_	yes	no	no

The allocator's own accounting says the same (page_allocator.cpp:706-713): physical memory is in-use + reserved, and free_page_list_ contributes nothing. So "unmap unused or reserved pages" collapses to a single set.
It also gives an exact invariant — mapped pages = num_total_pages_ − |free_page_list_|, which means num_total_pages is the upper bound of physical page usage, so it already bounds physical residency.

Gap 1: the limit is only polled from the allocation path

An instance with no traffic never notices that its limit changed, and "an idle instance hands memory back" is the main thing this feature is for.

Gap 2: no revision, no acknowledgement (You mentioned in the comment)

On hard limit vs immediate hard limit

Agreed, and I think that framing is right: a hard limit on new physical growth, not an immediate on. CPU offloading as a separate follow-up for the stronger convergence case sounds right too.

Question

Could you enumerate what the proposed cap covers that resize() plus the lazy-shrink state machine does not? I can see the trigger and the revisioning; I could not identify a difference in the downward behavior itself.

[shipiyouniao]

Thanks for tracing the existing path end to end. I agree with your conclusion about the downward mechanics: resize() plus the existing lazy-shrink state machine already provides the behavior we need for draining without revoking active mappings, so the PR should reuse that behavior rather than introduce a parallel shrink state machine.

The remaining distinction is the contract being limited. resize() reduces logical/virtual pool capacity by changing num_total_pages_. The proposed cap is intended to leave that virtual capacity intact and limit only concurrent physical residency, so an instance can regrow within its existing address space after the assigned cap increases.

After your walkthrough, I think the focused delta should be:

• a revisioned external update and acknowledgement, including updates applied while an instance is idle;
• a physical-residency budget that is checked consistently by foreground mapping and background preallocation without reducing virtual page capacity;
• reuse of the existing trim/free/lazy-convergence behavior for downward updates;
• no second shrink state machine.

The cross-process physical-growth guard is orthogonal and materially broadens the map path, so I will remove it from this PR and treat it separately.

I will revise the implementation around that narrower boundary after the allocator counter/snapshot prerequisite in #443 is available. Thanks, this was a useful correction to the current shape of the PR.

[shipiyouniao]

Reworked in 891921a following the discussion above.

The revised implementation is now stacked on #443, keeps virtual pool capacity unchanged, limits mapped physical residency in both foreground allocation and preallocation, and reuses the existing free/trim behavior for downward convergence. The cross-process growth guard and the parallel shrink design have been removed.

Validation on the revised head:

• all 11 GitHub checks passed;
• 127 CPU tests passed;
• 3 native T4 allocator tests passed, including a concurrent in-flight map during a non-blocking limit update.

The PR description has been updated with the current contract and dependency.

[Copilot]

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated no new comments.

[RixinLiu]

Hi @shipiyouniao , After reviewing this pr, i still think we don't need a new mechanism, please see my below reasons:

1. resize() never touches the virtual address space.

The reservation is made once in create_kv_tensors() and released in exactly one place. resize() only moves page ids between free_page_list_, reserved_page_list_ and reclaimed_page_list_, adjusts num_total_pages_, and calls unmap_pages(), which drops the physical backing but leaves the address range reserved. Growing back reuses the same ids out of reclaimed_page_list_.

Your claim "reduces logical/virtual pool capacity by changing num_total_pages_" is accurate, but num_total_pages_ is just bookkeeping. Lowering it doesn't make any address go away, so raising it doesn't need anything to have been preserved.

To me, it feels kind of inconsistent, if we have two mechanisms, one limits on the num of virtual pages, and the other limits on the num of physical pages. I think a very detailed case that explicitly illustrate why we must have a new mechanism instead of just enhancing current resize mechanism will make me understand more.

2. What I do want out of this PR, pointed at resize() instead.
The control protocol is the valuable half here, and kvctl limit badly needs it:

• Revision and acknowledgement. Today kvctl limit writes total_size into MemInfoStruct and stops. No revision field, no ack, no way for the caller to find out what happened. Your revisioned update is the right shape
• Apply while idle. The resize target is only consumed inside _alloc (kv_cache_manager.py:267), so with no traffic, kvctl limit never takes effect at all. That is a real bug and your first bullet covers it.
• Report the state. in_shrink already means "accepted, draining", but nothing exports it. Your applied / deferred / stale / conflict vocabulary is exactly what's missing.
• Split one budget across pools. control.py splitting an instance budget by virtual capacity is a capability we don't have today. It's orthogonal to what gets limited, so it sits on top of resize() unchanged.

[shipiyouniao]

Thanks, I agree with this correction. I had conflated allocator page capacity with CUDA virtual-address reservation. Since the address space is reserved by create_kv_tensors(), while resize() only moves page ids and changes physical backing, preserving a separate physical_page_limit_ does not provide a distinct capability in the current allocator.

I will revise this PR to enhance the existing resize() path instead of maintaining a second limiting mechanism. The focused scope will be:

• revisioned, idempotent updates with acknowledgement;
• immediate application while the instance is idle;
• applied, deferred, stale, and conflict state reporting;
• splitting an instance budget across registered pools;
• reuse of the existing in_shrink, release, trim, and regrowth behavior.

The separate physical-page limit, pending-map accounting, and duplicate allocation/preallocation gates will be removed. I will also add focused coverage for idle shrink, deferred shrink without revoking active mappings, regrowth, and multi-pool budget splitting. Thanks for walking through the existing semantics in detail.

[shipiyouniao]

Reworked in 69697d43 following your latest review.

The separate physical-page limit and all duplicate C++ allocation/preallocation gates have been removed. The PR now contains only the revisioned control protocol around the existing resize() / in_shrink behavior:

• immediate application while idle;
• applied, deferred, stale, conflict, and unavailable acknowledgements;
• deterministic instance-budget splitting across registered pools;
• page-bundle alignment and clamping to the original reservation;
• deferred shrink without revoking active mappings, followed by normal regrowth through reclaimed page ids.

The final diff is one commit touching four Python/test files, with no C++ map or preallocation hot-path changes. All 11 GitHub CPU, MyPy, and pre-commit checks passed.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

[RixinLiu]

LGTM, thanks. I will merge it now, but perhaps you can add some document/instruction to tell people (who wanna build their own production layer upon these) how to use them.

[shipiyouniao]

Added the requested production integration guidance in #458. It documents the process-local control-handler boundary, revision and acknowledgement semantics, deferred convergence polling, multi-pool splitting, and the ownership boundary between kvcached and the external controller.