fix(kv-cache): size pages via get_block_range in _get_num_alloced_blocks

[SuperMarioYL]

PR Draft — Round 1

Title: fix(kv-cache): size pages via get_block_range in _get_num_alloced_blocks

Base: main ← Branch: fix/page-capacity-get-block-range

---

Summary

KVCacheManager._get_num_alloced_blocks and available_size sized every page with
InternalPage.get_num_blocks(page_size, block_mem_size) — the theoretical
page_size // block_mem_size. But get_page_occupancy already uses the
boundary-aware InternalPage.get_block_range(page_id, page_size, block_mem_size),
and the comment there explains why: blocks straddling a page boundary belong to
neither page, so a page's capacity comes from its own get_block_range, not from
page_size // block_mem_size.

When block_mem_size does not evenly divide page_size (HYBRID_LINEAR / Mamba GDN
per-block state — the exact case the _alloc parking comment at
kv_cache_manager.py:335 names), some page ids yield zero usable blocks while
get_num_blocks reports one or more. Those 0-block pages are parked in
full_pages (so they're not re-handed-out) but were still counted by
get_num_blocks, inflating _get_num_alloced_blocks.

That inflation is observable, not cosmetic: free()'s lazy-shrink completion gate
_get_num_alloced_blocks() <= target_num_blocks reads the inflated count, so an
operator's kvctl limit shrink never crosses the threshold and physical GPU memory
is not returned (the elastic-shrink-stalls failure).

Fix: add a module-level _page_capacity(page_id) helper that mirrors
get_page_occupancy's get_block_range usage, and use it in
_get_num_alloced_blocks (per tracked page id in full_pages / avail_pages) so a
parked 0-block page contributes 0 instead of the theoretical count. The helper is
module-level (not a staticmethod) so it is unit-testable without the compiled
kvcached.vmm_ops extension or a GPU, matching the _get_max_cached_blocks /
_make_cache_key idiom.

Scope: available_size's free_pages term could not be made per-id precise —
the C++ PageAllocator exposes only a count of free pages, not their ids, so
the boundary-aware _page_capacity (which depends on page_id) can't be applied
there without a C++ change. The avail_blocks term in available_size was already
boundary-aware (num_avail_blocks is maintained from page.num_free_blocks(), which
InternalPage::init populates from get_block_range). An inline note documents the
remaining upper-bound term and the precise-fix prerequisite.

Test

tests/test_block_count_consistency.py (registered in tests/manifests/cpu.txt)
builds a real KVCacheManager via object.__new__ (per test_resize_reserved_order.py)
with page_size=4, block_mem_size=3 (3 does not divide 4) and a stub InternalPage
exposing the real static get_num_blocks / get_block_range formulas. It asserts
_get_num_alloced_blocks() returns the get_block_range-based count, not the
inflated get_num_blocks count. The test is red on main (concrete value
failures assert 1 == 0 / assert 2 == 1, not import errors, and does not mock the
function under fix) and green on this branch.

CPU suite: 142 passed (was 138; +4 from the new test); the 10 errors are
pre-existing macOS /dev/shm environment failures in unrelated files.

Why not just available_size

Addressed above — available_size's count-based free-page term is an upper bound
until the allocator exposes free page ids; the load-bearing accounting
(_get_num_alloced_blocks, used by the shrink gate) is fully fixed here.