[googs1025]

Problem

For a KV block that does not divide the physical page, available_size() multiplies the free-page count by a theoretical blocks-per-page value. That can admit more blocks than the actual page IDs hold. With two 14 MiB pages and 6 MiB blocks it reports 4, but the page ranges contain only IDs 0, 1, and 3. An allocation of 4 maps three blocks, then fails and rolls back.

Reproduce

On a CUDA GPU, set KVCACHED_PAGE_SIZE_MB=14 KVCACHED_PAGE_PREALLOC_ENABLED=false and a fresh KVCACHED_IPC_NAME. In a fresh kvcached process:

from kvcached import vmm_ops
from kvcached.integration.sglang import interfaces as kvi

mib = 1024 * 1024
kvi.init_kvcached(device="cuda:0")
raw = vmm_ops.create_kv_tensors(28*mib, 1, "cuda:0", 1,
    num_kv_buffers=1, group_id=500, unified_pool=True, page_size=14*mib)
manager = kvi.get_kv_cache_manager(5, 1, 6*mib, 1,
    reserve_null_block=False, num_kv_buffers=1, group_id=500)
manager.wait_ready(timeout=15)
print(manager.available_size(), manager.alloc(4))

Before the fix this prints 4 None; the allocator logs a partial 3/4 rollback. Stop the process and delete only its idle IPC segment after the run.

Fix

For non-divisible geometries, a native snapshot sums the block capacity of the reserved pages and the next physically affordable free page IDs in alloc_page() order. The snapshot also clips each page to the exclusive max_block_id used by InternalPage::init() and checks transaction failure while holding the allocator lock. The divisible fast path stays count-based.

Validation

On an A10, the same probe reports capacity 3 after the fix: alloc(4) is rejected before partial mapping and alloc(3) succeeds. On the current main base, 243 related CPU tests pass, including the new capacity and failure-classification regressions.

[Copilot]

🟡 Changes recommended

Critical issues remain: the snapshot can miss a fatal background mapping failure due to a race and overcount usable blocks at ID boundaries, allowing allocation to fail after partially succeeding.

2 open issues

• Check transaction failure under the snapshot lock · New
• Clip snapshot ranges to max_block_id · New

What changed in this PR

This PR makes KV-cache capacity forecasts use actual allocatable page IDs for non-divisible block/page sizes, reducing allocations that fail after partial mapping.

Changes:

• Add a native capacity snapshot and expose it to Python.
• Use the snapshot for non-divisible geometries while retaining the divisible fast path and fatal-error classification.
• Add capacity and failure-classification regressions and update test doubles.

File	Description
tests/​test_page_allocator_snapshot.py	Tests native capacity using actual page IDs.
tests/​test_observability.py	Adds capacity forecasting to the allocator fake.
tests/​test_block_count_consistency.py	Tests capacity forecasts and native failure classification.
tests/​test_alloc_rollback.py	Updates the rollback-test allocator fake.
kvcached/​kv_cache_manager.py	Uses native snapshots for non-divisible geometries.
csrc/​torch_bindings.cpp	Exposes the capacity query to Python.
csrc/​page_allocator.cpp	Sums capacity across reserved and affordable free pages.
csrc/​inc/​page_allocator.hpp	Declares and documents the snapshot API.

_{Review effort: Balanced}

[googs1025]

Fixed in c0456a0: the native snapshot now clips each page at max_block_id, and available_size() forwards that bound. Added regression coverage for the 14 MiB / 6 MiB, max_block_id=3 case; 243 related CPU tests pass.

[shipiyouniao]

T4 base/head checks confirm the 14 MiB / 6 MiB capacity and max_block_id=3 fixes. Native builds, the new snapshot test, and CI pass.

LGTM