fix: decouple virtual KV capacity from shared-GPU profiling

[shipiyouniao]

Summary

Fix the shared-GPU startup capacity failures described in #193 without serializing colocated vLLM or SGLang engines.

When kvcached is enabled, logical KV capacity is derived from stable device geometry and process-local memory costs. It is no longer derived from whole-device free-memory deltas that can change while another process is starting, serving, or exiting.

The known startup free-memory guard becomes a targeted warning in this mode. Explicit kv_cache_memory_bytes remains authoritative, kvcached-disabled behavior is unchanged, and unrelated model or CUDA failures still propagate.

Fixes #193.

Root cause

kvcached reserves CUDA virtual address space for KV tensors, but model weights, CUDA contexts, profiling workloads, activations, communication buffers, CUDA Graph capture, and physical KV page mappings still consume physical GPU memory.

vLLM computes a requested budget from the whole device:

requested_memory = ceil(total_memory * gpu_memory_utilization)

It then profiles non-KV costs and calculates the KV remainder:

available_kv_cache_memory_bytes = (
    requested_memory
    - profile_result.non_kv_cache_memory
    - cudagraph_memory_estimate_applied
)

This calculation explicitly assumes that other processes sharing the GPU do not change their memory use during profiling:

# Here we assume that the other processes using the same
# GPU did not change their memory usage during the profiling.

That assumption does not hold for colocated engines. MemorySnapshot observes whole-device free memory, while PyTorch allocator counters are process-local. A peer's allocation or release can therefore be attributed to the worker being profiled.

Observed failure modes

Initial free-memory guard

If a peer has already consumed memory before the initial snapshot, startup can fail with:

Free memory on device ... on startup is less than desired GPU memory utilization

The previous kvcached patch bypassed this guard but replaced requested_memory with current free memory. The subsequent profiling calculation remained sensitive to peer activity.

Peer allocation during profiling

If another process allocates memory during profiling, that growth can be counted as the current worker's non_torch_increase. The resulting KV capacity may become zero or negative:

requested_memory
- non_kv_cache_memory
- cudagraph_memory_estimate_applied
<= 0

This can produce No available memory for the cache blocks, or a positive capacity too small for one request at max_model_len.

Peer release during profiling

The inverse failure is also possible. An unfixed colocated startup observed free memory increase from 8.57 GiB to 12.91 GiB during determine_available_memory(). That violated vLLM's assertion:

assert self.init_snapshot.free_memory >= free_gpu_memory

EngineCore then terminated with Error in memory profiling. The snapshot is device-wide, so the peer does not need to be in the same process or container to invalidate it. See the vLLM 0.22.1 check.

Why startup serialization is insufficient

Serializing two initializers prevents simultaneous profiling, but it does not make whole-device deltas process-local:

1. Engine A finishes initialization and releases the lock.
2. A starts serving and changes its physical KV mappings.
3. Engine B acquires the lock and profiles while A continues changing GPU memory.

The same issue remains when engines A and B are serving and engine C starts later. A fully correct lock would require quiescing every running peer, which would stall traffic and make cold-start latency grow with the number of instances.

Implementation

vLLM

For supported vLLM releases, the patch:

1. Mirrors vLLM's ceil(total_memory * gpu_memory_utilization) formula to establish a stable logical budget.

2. Downgrades only the known whole-device startup guard to a warning when kvcached is enabled.

3. Reuses the budget stored in self.requested_memory instead of recomputing it from mutable free memory.

4. Runs the required profile workload, but sizes KV from process-local costs:

   requested_memory
   - weights_memory
   - torch_peak_increase
   

5. Ignores device-wide non_torch_increase, after_profile.free_memory, the cross-process free-memory assertion, and the CUDA Graph estimate for KV sizing.

6. Still runs profile_cudagraph_memory() when supported and records self.cudagraph_memory_estimate for vLLM's later comparison/logging, but does not fold the estimate into peak_activation_memory or subtract it from logical KV capacity.

7. Uses runtime feature detection for older supported workers that predate init_snapshot: after init_device(), it stores the same logical budget and measures the reset process-local PyTorch peak.

8. Leaves explicit kv_cache_memory_bytes and kvcached-disabled paths unchanged.

SGLang

For SGLang 0.5.11 and later, the patch wraps ModelRunner._profile_available_bytes():

logical_budget = ceil(total_memory * mem_fraction_static)
local_available = logical_budget - torch.cuda.memory_reserved(gpu_id)

The rank-local result is reduced with MIN across the current SGLang world group before pool sizing, preserving identical TP geometry. If any rank cannot query its local capacity, a failure sentinel makes every rank fall back to SGLang's original distributed profiler, so collective order cannot diverge.

Mamba reservation is applied after synchronization. Non-GPU devices and kvcached-disabled runs retain the original profiler. The patch is installed after elastic memory-pool aliases because importing ModelRunner caches those pool classes in module globals.

Temporary physical pressure

The vLLM and SGLang integrations reserve block 0 as a null block during post-initialization. If physical KV capacity is temporarily unavailable, reservation now sleeps and retries instead of terminating EngineCore. It also retries the race where capacity is observed but another engine takes the page before allocation.

The expected block identity remains fail-loud. CUDA allocation failures, unsupported layouts, invalid geometry, and unrelated runtime failures are not converted into warnings.

Validation

Hardware validation used one NVIDIA Tesla T4.

Scenario	Result
vLLM 0.24.0, PyTorch 2.11.0+cu129, Qwen2.5-0.5B-Instruct	CUDA Graph profiling and capture completed; the estimate was retained and logged as ignored for sizing; the service reached health and an OpenAI Chat request returned exactly OK
SGLang 0.5.15, PyTorch 2.11.0+cu129, sglang-kernel 0.4.4+cu129, Qwen2.5-3B-Instruct	All kvcached SGLang patches applied; the process-local logical capacity was reported; with CUDA Graph disabled for the T4 smoke run, the service reached health and an OpenAI Chat request returned exactly OK
Focused virtual-capacity and patch-order regressions	25 passed
Full Linux CPU manifest	167 passed on each of Python 3.9, 3.10, 3.11, 3.12, and 3.13
Static and native checks	Five mypy jobs, pre-commit, and the GPU-free C++ tests passed

The vLLM 0.24.0 run also verifies the worker contract required by compile_or_warm_up_model(): self.non_torch_memory is initialized without reintroducing the device-wide delta. Separate regressions verify that a large peer-driven non_torch_increase and the CUDA Graph estimate do not change logical KV capacity.

Known limitation

This patch does not make a permanently infeasible configuration runnable. At least one physical page and sufficient non-KV headroom must eventually become available.

While reserving the null block, _post_init_done is not set until block 0 is allocated. Temporary pressure can recover, but if no physical KV page ever becomes available, post-initialization keeps waiting instead of failing fast. The retry sleeps and does not busy-spin, but the engine does not become usable.

This PR does not introduce an arbitrary timeout because that would turn temporary pressure back into a timing-dependent startup failure. A follow-up can make the wait lifecycle-aware and cancellable, or connect it to an explicit startup deadline supplied by the caller.

Design boundary

This is a compatibility fix for the current kvcached integrations. It does not replace vLLM's complete memory profiler or define a shared-GPU physical-memory accounting model.

The selected boundary, rejected alternatives, known limitation, and criteria for reopening the larger design are documented in Discussion #454.

[Copilot]

Pull request overview

This PR adds cross-process startup serialization for vLLM’s shared-GPU memory profiling window when kvcached is enabled, preventing concurrent vLLM instances from corrupting each other’s CUDA-graph-based KV cache budgeting (related to #193).

Changes:

• Introduces a per-visible-GPU flock-based lock (configurable via env vars) to serialize vLLM’s profile-to-CUDA-graph capture window.
• Updates the vLLM EngineCore patch to lock _initialize_kv_caches when available, with a compatibility fallback to locking the full EngineCore init.
• Adds unit tests validating lock resource selection, patch scoping, and POSIX cross-process serialization; includes the new test in the CPU manifest.

Reviewed changes

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

File	Description
kvcached/integration/vllm/patches.py	Adds startup-locking primitives and integrates them into the EngineCore patching flow.
tests/test_vllm_startup_serialization.py	Adds unit tests covering resource selection and cross-process locking behavior (POSIX-only where needed).
tests/test_vllm_tp_world_size.py	Disables startup locking in this unit-test module to keep Windows test runs working.
tests/manifests/cpu.txt	Ensures the new serialization test runs in the CPU test manifest.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[shipiyouniao]

I replaced the previous startup-serialization approach with the lock-free implementation now in this PR.

The lock prevented two initializers from profiling at the same time, but it did not make vLLM's whole-device profiling assumption valid. After engine A released the lock, it could begin serving and change physical KV mappings while engine B profiled; the same problem remained for a later engine C. A fully correct lock would require quiescing every running peer, which would stall serving traffic and make cold starts increasingly expensive.

The new implementation instead separates logical KV sizing from transient whole-device free memory:

• vLLM receives a stable logical KV capacity derived from total memory and gpu_memory_utilization;
• the known startup free-memory guard becomes a targeted warning only when kvcached is enabled;
• existing explicit capacity and kvcached-disabled paths remain unchanged;
• temporary physical pressure during null-block reservation waits and retries, including the allocator race after an availability check;
• unrelated CUDA failures and invalid layouts/configuration still fail normally.

On a Tesla T4 with vLLM 0.18.0 and Qwen2.5-0.5B-Instruct, two C256 engines started simultaneously without a startup lock, received the same logical KV capacity, reached health, and both returned OK. Starting B while A served also completed successfully (A: 64/64 requests, 2.92 req/s, 875.69 output tok/s). A third engine under temporary physical KV pressure progressed about 3 seconds after a peer released pages and then returned OK.

The branch also passes 22 focused tests and the full 157-test Linux CPU manifest.

[Copilot]

Pull request overview

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

Suppressed comments (1)

kvcached/integration/vllm/patches.py:1958

• patch_worker_determine_available_memory() temporarily sets cache_config.kv_cache_memory_bytes, but the finally block always restores it to None. Even though this branch is entered only when it was None, restoring the prior value is safer in case vLLM mutates/normalizes the field internally or a future config type changes the default value.

                cache_config.kv_cache_memory_bytes = virtual_budget
                try:
                    available_memory = int(
                        original_determine(self, *args, **kwargs)
                    )

[shipiyouniao]

One limitation of the current null-block retry should be explicit: if physical KV capacity never becomes available, post-initialization remains waiting instead of failing fast.

Concretely, the null block is block 0, which the engine reserves as a sentinel before normal KV allocations begin. The upstream code attempts this reservation once; if _alloc(1) returns None, the identity check fails and startup raises RuntimeError immediately.

With this PR, logical KV sizing can complete even while the physical page allocator has no page available. During KVCacheManager._post_init(), _reserve_null_block() therefore behaves as follows:

1. If available_size() is zero, it sleeps for 10 ms and retries.
2. If capacity was observed but another colocated engine wins the allocation race, it also sleeps and retries.
3. _post_init_done is not set until block 0 has actually been reserved, so normal allocations remain blocked behind post-initialization.
4. Once another engine releases a physical page, the retry reserves [0], post-initialization completes, and the engine can continue. We validated this recovery path by releasing peer capacity while the new engine was waiting.

The side effect appears when the configuration is permanently infeasible, for example when colocated engines retain all physical KV pages and none can ever be released for the new engine. In that case the process does not busy-spin or crash, but it also never becomes usable: post-initialization keeps retrying, and callers waiting on _post_init_done remain blocked. This is what "waiting instead of failing fast" means here.

A fixed timeout would turn temporary pressure back into a timing-dependent startup failure, so this PR does not invent an arbitrary one. A durable follow-up should make the wait lifecycle-aware and cancellable, or connect it to an explicit startup deadline supplied by the caller, and expose the waiting state clearly for operators. The indefinite-wait behavior is therefore a real tradeoff that should be considered when reviewing this change.

[shipiyouniao]

I added another observed failure mode to the PR description. An unfixed colocated startup recorded free memory increasing from 8.57 GiB to 12.91 GiB during determine_available_memory(), so vLLM's init_snapshot.free_memory >= free_gpu_memory assertion failed and terminated EngineCore.

This is the inverse of a peer allocating memory during profiling: a peer released about 4.34 GiB while the worker was measuring whole-device memory. Both directions invalidate the same cross-process assumption. The explicit kv_cache_memory_bytes path used by this PR still runs profile_run() for compilation, but skips that whole-device delta and assertion.

[Copilot]

Pull request overview

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

Suppressed comments (1)

tests/test_sglang_autopatch_order.py:24

• The test derives patch ordering via ast.walk(tree), but ast.walk does not guarantee a source-order traversal. That makes this ordering assertion potentially unreliable across Python versions / AST changes, and it could accidentally include unrelated *Patch() calls elsewhere in the file. Parse the specific register_patches_with_versions([...]) list literal and read its elements in-order instead.

    patch_names = [
        node.func.id
        for node in ast.walk(tree)
        if isinstance(node, ast.Call)
        and isinstance(node.func, ast.Name)
        and node.func.id.endswith("Patch")
    ]

[RixinLiu]

LGTM. Two suggestions.

Take the direct path on the vLLM side too, like SGLang. The SGLang patch replaces _profile_available_bytes() outright, and this PR already contains the equivalent direct form for vLLM in the pre-0.11 fallback (profile_run() then return the budget).

available_kv_cache_memory_bytes should be requested_memory − weights_memory − torch_peak_increase, not requested_memory. Of the three fields in non_kv_cache_memory, only non_torch_increase reads device-wide state and can be changed by a peer; weights_memory comes from model_runner.model_memory_usage and torch_peak_increase from memory_stats(), both process-local.

Measured on an A100, enforce_eager, util 0.92, requested 36.33 GiB:

	Qwen2.5-0.5B	Qwen3.5-4B
weights_memory	0.93 GiB	8.61 GiB
torch_peak_increase	0.59 GiB	1.70 GiB
non_torch_increase	0.11 GiB	0.11 GiB
real KV capacity	34.70 GiB	25.90 GiB
this PR reports	36.33 GiB (+4.7%)	36.33 GiB (+40%)
requested − weights − torch_peak	34.81 GiB (+0.3%)	26.02 GiB (+0.5%)

[shipiyouniao]

Good catch. I changed the vLLM path to profile directly and calculate:

available_kv_cache_memory_bytes = requested_memory - weights_memory - torch_peak_increase

non_torch_increase is no longer used for sizing. An explicit kv_cache_memory_bytes still keeps the original vLLM path.

I also reran the simultaneous-start test on a T4 with two vLLM 0.18.0 C256 engines. Both produced the same values:

Field	Both engines
requested budget	14,073,377,588 bytes
weights	994,719,744 bytes
torch peak increase	586,687,488 bytes
available KV capacity	12,491,970,356 bytes

Both reached health and returned exactly OK. The regression test injects a large non_torch_increase and verifies that it has no effect on the result. Updated in 24b136d3.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

Suppressed comments (1)

kvcached/integration/vllm/patches.py:1946

• The direct profile path omits vLLM's process-local CUDA-graph estimate. Supported vLLM 0.18/0.22 releases call profile_cudagraph_memory() and subtract it when VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS is enabled; this calculation subtracts only weights and the forward-pass torch peak, so that opt-in configuration overstates KV capacity and can leave no room for graph capture. Preserve and deduct the enabled graph estimate without reintroducing whole-device deltas.

            torch_peak_increase = int(profile_result.torch_peak_increase)
            available_memory = (
                virtual_budget - weights_memory - torch_peak_increase
            )

[shipiyouniao]

One safety boundary is worth making explicit before merge.

The new vLLM formula is intentionally a per-engine logical KV capacity. It subtracts process-local weights, torch peak, and the applicable CUDA Graph estimate, but it does not subtract peer engines'' physical costs. With large colocated models, peer weights alone can be tens of GiB, so the aggregate logical capacities may overcommit physical memory by a large margin.

That overcommit is acceptable only if kvcached''s runtime physical admission remains authoritative and atomic across processes. The current scheduling path does re-read KVCacheManager.available_size(), but the check and the following multi-page cuMemCreate/map batch are not one cross-process transaction in upstream main or this branch. Two engines can therefore observe the same headroom, both pass the check, and race into physical growth. The CUDA driver will eventually reject the later allocation, and the normal result should be an allocation miss plus rollback, but pages allocated before that failure can temporarily consume the configured headroom. A concurrent Torch/CUDA allocation that does not participate in the same protocol can then receive a real OOM.

We observed that this is not merely a theoretical cold-start path in an 80B shared-GPU overload run. That build included a per-physical-GPU growth guard covering the headroom check and mapping transaction. Under load it recorded:

833 x Physical KV allocation miss; deferring scheduling
39,984 x Page 47 is not mapped

39,984 = 833 * 48: the older distributed compensation path attempted unmap on 48 worker/page targets for each rejected growth. Those noisy rollback messages are addressed separately by the transactional VMM work, but the important evidence here is that physical admission was exercised hundreds of times during serving even with the shared growth guard. No process-level OOM occurred and all 1,000 requests eventually completed.

This does not mean peer weights should be subtracted from the logical capacity: doing so would make capacity order-dependent again and defeat elastic overcommit. It means this PR increases the importance of a separate physical-growth transaction boundary.

The required lock is also different from the startup/profiling lock removed earlier. It would serialize only:

physical-GPU headroom snapshot
-> admission decision
-> one physical KV growth/map batch
-> rollback on failure

It would not serialize model loading, profiling, normal inference, or already-backed KV access.

I suggest documenting atomic runtime physical admission as a safety dependency/known limitation of logical overcommit, and handling the per-GPU growth transaction in a focused follow-up rather than mixing it into this profiling fix.

[RixinLiu]

Both changes look right. Two things.

Why subtract the CUDA Graph estimate? It comes from the same device-wide mem_get_info() reading as non_torch_increase, and is worse: a peer allocating during the second capture inflates per_graph and that error is multiplied by the graph count.

Could this run on the versions we support? Validation used vLLM 0.18 / SGLang 0.5.13.post1 on a T4; README and CI target 0.24.0 / 0.5.15.

[shipiyouniao]

Validated the currently supported versions and updated the PR accordingly.

Integration	Runtime validation	Result
vLLM 0.24.0	PyTorch 2.11.0+cu129, T4, Qwen2.5-0.5B-Instruct	CUDA Graph profiling and capture completed; the estimate was retained but ignored for logical sizing; health passed and the chat request returned exactly OK
SGLang 0.5.15	PyTorch 2.11.0+cu129, sglang-kernel 0.4.4+cu129, T4, Qwen2.5-3B-Instruct	All SGLang patches applied; process-local logical capacity was reported; health passed and the chat request returned exactly OK

The vLLM 0.24 run also exposed its compile_or_warm_up_model() contract for self.non_torch_memory. The patch now initializes that field to zero, preserving the object contract without reintroducing the peer-contaminated device-wide delta.

The final commit is aba5125. The 25 focused regressions, the 167-test CPU manifest on Python 3.9-3.13, five mypy jobs, pre-commit, and the GPU-free C++ tests all pass. Remote CI is green as well.

[RixinLiu]

Thanks! It looks better a lot!