perf(kv-cache): cache get_avail_physical_pages in available_size to skip per-alloc cudaMemGetInfo

[SuperMarioYL]

Summary

KVCacheManager.available_size() calls PageAllocator.get_avail_physical_pages() on every invocation, and that C++ method issues a cudaMemGetInfo driver call. available_size() runs once per allocation request (patches.py:792) and once per scheduler step (:927), so the driver call fires on every alloc/step. This PR caches the get_avail_physical_pages() result for 100 ms — the resize_watcher poll interval — and invalidates it on resize() and on the in_shrink toggle, so the driver call fires at most once per 100 ms window instead of once per call.

Motivation

benchmarks/bench_alloc/README.md documents a 12.5× allocation-throughput win attributed to dropping cudaMemGetInfo from available_size(). At HEAD (ac9680a) that drop did not happen: the call was relocated from available_size() into get_avail_physical_pages() (csrc/page_allocator.cpp:481), which available_size() still calls (kvcached/kv_cache_manager.py:502). This PR delivers the documented win by caching the call rather than relocating it. Refs #299 (the "Reduce CUDA Call Overhead in available_size" track item); this PR takes the narrower TTL-cache approach rather than the issue's proposed C++ rewrite.

Changes

• kvcached/kv_cache_manager.py: add a TTL cache (value + time.monotonic() timestamp, 100 ms) for get_avail_physical_pages() in available_size(). Invalidate on resize(), on the in_shrink toggle in free(), and on clear(). get_num_free_pages() and get_num_reserved_pages() stay uncached (cheap / atomic).
• tests/test_available_size_cache.py (new, CPU-only): asserts the call count drops from N-per-available_size() to 1-per-100 ms window, refetches after the TTL expires, and refetches after a resize.

Reviewer note — staleness bound

The physical free-page count can go stale for up to 100 ms between refreshes on a multitenant GPU where another process frees/allocates in that window. The staleness is bounded by the same 100 ms interval the allocator already tolerates for virtual free-page polling (the resize_watcher cadence), and an explicit resize() or in_shrink transition invalidates immediately, so no stale value is served across a resize.

Tests

python -m pytest tests/test_available_size_cache.py — 3 passed. The cache test is red on main (no cache → call count N) and green on this branch. Existing CPU tests touching available_size (test_prefix_cache, test_bestfit_page_selection, test_page_aware_eviction, test_alloc_rollback, test_observability) still pass. The new test is registered in tests/manifests/cpu.txt; python tools/check_test_classification.py passes.

[jeff3071]

Thanks for your contribution!

I think this change introduces a correctness issue. physical_free_pages now uses a cached value:

physical_free_pages = (
                self._get_cached_avail_physical_pages()
                + self.page_allocator.get_num_reserved_pages())

But the cached pages do not change after alloc that cause cached pages number is dirty.
That would cause available_size return wrong number.

However, the cache is not updated or invalidated after alloc() maps new physical pages. As a result, subsequent calls to available_size() within the TTL may use a stale physical-page count and overestimate the available KV-cache capacity.

This can be reproduced by changing n_blocks to 1024 in tests/test_kvcache_manager.py::test_basic_alloc_free:

def test_basic_alloc_free(setup_kvcache):
    # instantiate a kv cache manager with known size
    manager = setup_kvcache

    # initial available blocks
    initial_available = manager.available_size()

    # allocate some blocks
    # change this
    n_blocks = 1024 
    handle = manager.alloc(n_blocks)
    after_alloc = manager.available_size()
    assert after_alloc + n_blocks == initial_available

    # free the allocated blocks
    manager.free(handle)
    after_free = manager.available_size()
    assert after_free == initial_available

Result

>       assert after_alloc + n_blocks == initial_available
E       assert (20992 + 1024) == 21632

tests/test_kvcache_manager.py:99: AssertionError

[SuperMarioYL]

Thanks for catching this — you are right. The cache served stale physical-free data after alloc() mapped new pages (and after free()/clear() returned pages), so available_size() could overestimate KV-cache capacity within the TTL window. Fixed in 2395135 by dropping the cached value at each physical-pool mutation, consistent with the existing resize()/in_shrink invalidation:

• alloc_page() maps a new physical page → invalidate (in alloc)
• free_pages() returns pages to the driver → invalidate (in free and clear)

The per-scheduler-step hot path (available_size() with no pool mutation between calls) still collapses to one driver read per TTL window; only the page-map/return paths now force a re-read, which is the correct place to re-read since capacity actually changed there.

Since test_basic_alloc_free needs a GPU, I added a CPU-only regression test (tests/test_available_size_cache.py::test_available_size_refetches_after_alloc) that stubs the page allocator: alloc maps a page and available_size() must drop by one page's worth of blocks instead of replaying the pre-alloc cached value. It is red on the previous commit and green on this one.

[SuperMarioYL]

Rebased onto latest main to resolve a merge conflict in kv_cache_manager.py. Two upstream commits (#448, #414) added memory-limit tracking and refactored null-block reservation in the same __init__ region where this PR adds the avail-physical-pages TTL cache. Both sets of instance variables are kept; the cache invalidation in resize() now adopts upstream's >= 0 assertion. All 12 tests pass and mypy is clean.