[rishabhsinha17]

Fixes #437. Validate get_kv_cache_manager()'s num_blocks against the capacity actually reserved by alloc_kv_cache(), raising a ValueError that names both numbers instead of aborting the process later inside FTensor::map. num_blocks=None now derives the created capacity.

Credit to @rob-9 for the diagnosis on #437: alloc_kv_cache() warned and clamped the created tensors (65,536 requested -> 45,056 on his L4) while get_kv_cache_manager() was configured independently with the original 65,536, so the manager exposed 512 page ids over an FTensor holding 352 compound pages, and the first map past the reservation failed cuMemUnmap and hit CHECK_GPU's std::abort. His control run with the manager at 45,056 completed all 400 churn iterations.

Changes

• alloc_kv_cache() records per group_id the created capacity; get_kv_cache_manager() refuses num_blocks beyond it (ValueError naming both numbers) and derives it for num_blocks=None. The per-layer capacity is computed in the caller's block geometry (ftensor_bytes_per_layer // num_kv_buffers // (block_size * cell_size)), consistent for MHA, MLA, and HYBRID_LINEAR units.
• When no allocation is recorded in the process (vLLM's engine/worker split: the manager lives in EngineCore, alloc_kv_cache runs in the worker), an explicit num_blocks is trusted unchanged, so the integration flow is unaffected.
• 12 CPU tests (tests/test_vllm_ftensor_capacity_guard.py), including the SIGABRT in FTensor::map() — mem_unmap fails with "invalid argument" under repeated large alloc/free churn #437 reproducer geometry and MLA single-buffer units. The stub fixture restores the kvcached module graph on teardown; the earlier revision that didn't broke later ElasticBlockPool tests on the CI 3.10 job when the file ran early in the manifest.

Scope: no csrc changes, deliberately

The C++ half of #437 (bounds-check page ids in FTensor::map/unmap, stop aborting) is left to open PR #418, which already adds FTensor::validate_offset_() on those paths and converts the CHECK_GPU aborts into std::runtime_error. Duplicating that here would conflict with #418/#455/#465, which all touch the same csrc region. This PR is the Python-side guard that stops the invalid configuration before any page id is handed out; it composes with #418.

Re-checked 2026-08-27 with git merge-tree --write-tree against the current heads: merges clean with #455 and #465; with #418 the only conflict is both PRs appending test filenames to tests/manifests/cpu.txt (trivial union; #418's interfaces.py hunks are confined to init_kvcached() and don't overlap this PR's).

Validation

GPU (this PR's target failure), 1x NVIDIA L40S (driver 570.124.06), kvcached editable install, torch 2.11.0+cu130: reproducer from #437 scaled to this GPU's capacity (46GB creates 90,880 blocks for the issue's 16-layer fp16 geometry, so the issue's 65,536 doesn't clamp here; 131,072 keeps the same structure of one number passed to both calls).

• main (60cad94): SIGABRT, same signature as the issue:

[kvcached][WARNING] interfaces.py:406 Requested 131072 blocks, but only 90880 blocks are available.
[repro_437] get_kv_cache_manager(num_blocks=131072)
iter 0 ok
csrc/ftensor.cpp:112 gpu_vmm::mem_unmap(vaddr, page_size_) failed in CUDA driver (1): invalid argument
Aborted (core dumped)  [exit 134]

• this branch, same run: the ValueError, before any page id exists:

ValueError: num_blocks=131072 exceeds the capacity of the KV cache tensors created by
alloc_kv_cache() for group 0: 90880 blocks of 16384 bytes (2977955840 reserved bytes per
layer, 2 KV buffers). A manager configured beyond the created tensors maps pages outside
the reserved virtual range and aborts in FTensor::map (issue #437). Pass num_blocks=None
to derive the capacity, or pass the clamped block count logged by alloc_kv_cache().

• this branch, num_blocks=None: derives 90,880 and all 400 churn iterations complete (COMPLETED WITHOUT ABORT).

CPU: 12 new tests pass; full CPU manifest 188 passed under a CI-matching env (3.10, CPU torch, pytest 8.4). tools/check_test_classification.py valid; ruff/isort clean; mypy --python-version 3.10 introduces no new errors vs main.

One unrelated observation from the same runs: at interpreter exit every configuration, including main with a valid num_blocks=90880 control, core-dumps in FTensor/Page static destructors with unknown CUDA driver error (4) (driver already deinitialized; the standalone reproducer never calls shutdown_kvcached()). Pre-existing teardown behavior, orthogonal to this fix; #418's abort-to-exception conversion likely covers it.

[rob-9]

great! i'll give this a review when i can.

[rob-9]

This still misses the supported manager-first order. If the manager is created for 65,536 blocks before alloc_kv_cache() clamps the tensors to 45,056, it is never revalidated and the #437 mismatch remains. Could alloc_kv_cache() check any existing manager's requested capacity before the tensors become available for mapping?

[rishabhsinha17]

This still misses the supported manager-first order. If the manager is created for 65,536 blocks before alloc_kv_cache() clamps the tensors to 45,056, it is never revalidated and the #437 mismatch remains. Could alloc_kv_cache() check any existing manager's requested capacity before the tensors become available for mapping?

Done in 519c763. alloc_kv_cache() now looks up any manager already registered for its group_id and validates it before create_kv_tensors(), so the tensors never become mappable for a manager that cannot address them. The check has to precede creation: the manager's _post_init thread polls kv_tensors_created() every 1 ms and starts mapping as soon as it flips.

• kvcached/integration/vllm/interfaces.py:435-456: the per-group geometry record is built first, _validate_registered_managers(group_id, created_capacity) runs, then create_kv_tensors(), then the record is stored.
• interfaces.py:670 _validate_registered_managers(): iterates get_registered_kv_cache_pools(integration="vllm") from kvcached.pool_registry, where get_kv_cache_manager() already registers every manager it builds (weak references, so a collected manager from a torn-down engine cannot veto a new allocation), filters on group_id, and runs the shared validator on manager.num_blocks / block_mem_size / num_layers / num_kv_buffers.
• interfaces.py:610 _validate_manager_against_created_tensors(): the one validator both orders use. Same ValueError naming both numbers (65,536 requested vs 45,056 created); nothing is shrunk silently.

Tests in tests/test_vllm_ftensor_capacity_guard.py, TestManagerFirstOrder (L311): a manager for 65,536 followed by an alloc whose device budget clamps to 45,056 raises with both numbers, create_kv_tensors is never called and no record is written (L314); a matching manager lets the alloc proceed and records the geometry (L353); alloc-first end to end with num_blocks=None (L369); a manager for another group and a collected manager are not compared (L383, L392).

[rob-9]

Could we also check that num_layers and num_kv_buffers match the allocation? In contiguous mode, these values determine the mapping stride. Right now num_layers is recorded but never checked, so the block count can pass while the mapping offsets are still invalid.

[rishabhsinha17]

Could we also check that num_layers and num_kv_buffers match the allocation? In contiguous mode, these values determine the mapping stride. Right now num_layers is recorded but never checked, so the block count can pass while the mapping offsets are still invalid.

Done in 519c763. The record now also carries the num_kv_buffers the FTensors were created with (compound_num_kv_buffers: 2 for MHA/GQA, 1 for MLA and HYBRID_LINEAR), and both values are checked in both orders:

• interfaces.py:640-652 in _validate_manager_against_created_tensors(): num_layers and num_kv_buffers must equal the record exactly, checked before the capacity check. The ValueError names both values, e.g. num_layers=32 does not match num_layers=16 of the KV cache tensors alloc_kv_cache() creates for group 0, and states why: FTensorAllocator::create_kv_tensors sizes the compound page as kPageSize * num_layers * num_kv_buffers (csrc/allocator.cpp:142) and PageAllocator::map_pages strides offsets by page_size_ * num_layers_ * num_kv_buffers_ (csrc/page_allocator.cpp:677), so a mismatch maps at the wrong offsets even when the block count fits.
• interfaces.py:692 _resolve_manager_num_blocks() now takes num_layers and runs the validator for get_kv_cache_manager() (alloc-first order, call at L751). num_blocks=None derivation goes through the same check, so it cannot bypass a layout mismatch.
• The manager-first order gets the same checks through _validate_registered_managers(), which reads manager.num_layers / manager.num_kv_buffers.

The check is unconditional rather than contiguous-only: in the per-layer layout the two values still size the PageAllocator's physical-memory accounting, and every in-tree caller already passes the same values to both calls (coordinator path: group_size and _get_kv_cache_params; v0.8 path: layer count and spec-derived buffer count), so no working configuration trips it.

Tests: num_layers and num_kv_buffers mismatches at the resolver (L232, L239), through get_kv_cache_manager() before construction (L280, L289), through alloc_kv_cache() before create_kv_tensors (L332, L341), and None derivation still checking layout (L247).

[rishabhsinha17]

Addressed both review comments in 519c763, a second commit on top of 933b2ab (no rebase; upstream/main is unchanged at 60cad94).

• Manager-first order: alloc_kv_cache() validates any manager already registered for its group_id (via kvcached.pool_registry) against the geometry it is about to create, before create_kv_tensors() makes the tensors mappable. Same ValueError naming both numbers; nothing is shrunk silently.
• num_layers and num_kv_buffers are now recorded and checked in both orders; the ValueError names the mismatched values and the compound page stride they set.
• One shared validator, _validate_manager_against_created_tensors(), backs both get_kv_cache_manager() and alloc_kv_cache().

Tests: 12 new in tests/test_vllm_ftensor_capacity_guard.py (24 total), including driving alloc_kv_cache() under the stubs with a device budget that reproduces the 65,536 -> 45,056 clamp. Full CPU manifest 200 passed locally (the 10 macOS-only shared-memory errors are identical on main); ruff/isort clean; mypy --python-version 3.10 unchanged vs main; tools/check_test_classification.py valid.

[rob-9]

both addressed, lgtm

[shipiyouniao]

The manager-first check and the layer/buffer geometry checks address rob-9's findings. I did not find another blocking issue in this pass.

Validation:

• All 24 capacity-guard tests passed on head 281ec39.
• I also tested the clean merge with main 73fa57c on one NVIDIA T4, with both contiguous and non-contiguous layouts. The native sources were unchanged from the already-built main backend.
• Requesting 131,072 blocks over tensors with capacity 29,824 was rejected before manager construction; incorrect layer/buffer counts were also rejected. num_blocks=None derived 29,824 correctly.
• Each layout completed 400 allocation/free cycles with CUDA readback. Each also hit 10 injected native physical-allocation failures and recovered on all 10 retries, then shut down cleanly. These are single-device allocator tests, not TP/MPS serving tests.

One non-blocking description update: #418 has merged, and the test count is now 24 rather than 12. Please refresh the old dependency/validation paragraphs so readers do not think this still waits on an open #418 or that the teardown behavior was verified by that older run.

LGTM for this Python-side guard. @RixinLiu, could you take the final look and handle the merge?