[aijanai]

Overview

Fixes #65378. On integrated (UMA) GPUs the extensible KV cache (#56492) sized itself from raw cudaMemGetInfo, which does not count reclaimable page cache as free, so it undersized the KV cache by up to 84% on a DGX Spark (GB10) and to zero blocks in disaggregated serving. This reads free memory through one helper with the same psutil correction the classic sizing path already applies (#35920).

Claims

• On GB10 the extensible KV cache now gets the KV cache the classic path finds: 1302 blocks / 67.6 GiB instead of 216 blocks / 11.2 GiB for Qwen3.8-27B NVFP4 at gpu_memory_utilization=0.8, and no longer fails with "No available memory for the cache blocks" in P/D.
• Discrete GPUs are unchanged: the helper returns torch.accelerator.get_memory_info as is when is_integrated_gpu() is false.
• No new configuration; --no-enable-extensible-kv-cache is no longer needed as a workaround on UMA parts.

Validation

All runs on one GB10 (DGX Spark, 119.7 GiB unified memory, CC 12.1) or one RTX PRO 6000 (discrete, 96 GB), image vllm/vllm-openai:nightly at 40941f2 (the extensible code it ships is identical to this branch's base), with this PR's mem_utils.py and extensible_kv_cache.py copied over the installed package for the "with fix" rows.

End to end, vllm serve nvidia/Qwen3.8-27B-NVFP4 --quantization modelopt_fp4 --kv-cache-dtype fp8 --max-model-len 4096 --gpu-memory-utilization 0.8 --max-num-seqs 32 --speculative-config '{"method":"mtp","num_speculative_tokens":3}', V2 model runner, extensible KV cache at its default:

Device	Load strategy	Fix	Extensible KV cache after warmup	Classic path (Available KV cache memory)
RTX PRO 6000 (discrete)	prefetch	no	560 blocks, 29.1 GiB, starts	52.95 GiB
GB10	prefetch	no	216 blocks, 11.2 GiB (42,130 tokens)	71.04 GiB
GB10	default	no	302 blocks, 15.7 GiB	71.55 GiB
GB10	prefetch	yes	1302 blocks, 67.6 GiB (253,952 tokens)	70.80 GiB

Before the server started on GB10, cudaMemGetInfo reported 59-81 GiB free while psutil.virtual_memory().available reported 114-116 GiB: the difference is page cache, from this load and earlier runs on the node. In disaggregated P/D (NIXL, same model, MTP k=3) the larger transient peak (7.28 GiB vs 1.71 GiB) took the same deficit below zero: both workers failed with 0 KV cache blocks (0.0 GiB) fit within the requested 95.71 GiB / ValueError: No available memory for the cache blocks.

Device allocations reclaim the page cache. After filling the page cache by reading the checkpoint, committing through the extensible cache's own _VirtualBuffer (cuMemCreate + cuMemMap) in 1 GiB steps reached 106 GiB, 1.8x the 59.2 GiB cudaMemGetInfo reported free, with no error; psutil "cached" fell from 56.4 to 4.1 GiB in step with the commits. So the psutil figure is the one an allocation can actually obtain.

New test fails before the fix, passes after (GB10):

$ pytest tests/v1/worker/test_extensible_kv_cache.py::test_measured_sizing_counts_reclaimable_memory_on_integrated_gpu
# nightly 40941f2c4b as shipped
>       assert num_blocks == 80
E       assert 0 == 80
tests/v1/worker/test_extensible_kv_cache.py:285: AssertionError
1 failed in 3.87s
# with this PR
1 passed in 3.89s

tests/utils_/test_mem_utils.py and tests/v1/worker/test_extensible_kv_cache.py with the fix on GB10: 26 passed, 1 failed. The failure, test_memory_profiling (asserts 256 MiB of non-torch memory, measures ~269-273 MB), also fails on the unmodified nightly with the upstream test file, 2 out of 2 runs, so it is a pre-existing UMA flake unrelated to this change. Both files are already run by CI (pytest -v -s v1/worker, pytest -v -s utils_ in .buildkite/test_areas/misc.yaml).

Details

Root cause. MemorySnapshot.measure() (vllm/utils/mem_utils.py) replaces cudaMemGetInfo's free memory with psutil.virtual_memory().available on integrated GPUs, and the extensible cache receives that as init_free_memory. But committable_blocks() and measure_kv_cache_blocks() in vllm/v1/worker/extensible_kv_cache.py re-read raw torch.accelerator.get_memory_info. The sizing in measure_kv_cache_bytes() is min(requested - non_kv_used, free + committed - margin) - transient_peak with non_kv_used = init_free - free - committed: with the raw figure, page cache counts as engine memory in the first term and is missing from the headroom in the second. In the P/D failure the headroom term alone was already negative: 10.46 + 0.73 - 5.46 - 7.28 = -1.55 GiB. Warmup was capped by the same raw read (only 0.73 GiB committed).

Fix. Move the existing correction into device_memory_info(device) in mem_utils and use it from MemorySnapshot.measure() and the three extensible call sites (committable_blocks, measure_kv_cache_blocks, and the free-memory read when releasing for a weight update). The one remaining get_memory_info call in the extensible cache reads only the total. Main-code change is ~30 lines; the rest is tests.

Alternative considered. Turning the extensible cache off on integrated GPUs in _resolve_extensible_kv_cache would also avoid the failure, but would hide the inconsistency and drop the feature on UMA parts. The reclaim measurement above shows the corrected figure is safe to size from, so this PR fixes the measurement instead.

Tests. The new tests patch current_platform.is_integrated_gpu, psutil and torch.accelerator.get_memory_info: CI has no integrated GPU, and these stand in for the hardware/OS readings rather than for vLLM objects. One covers the helper (integrated vs discrete), one covers the extensible sizing end to end through measure_kv_cache_blocks with the readings from this report.

Related. #55828 (open) sizes the KV cache from the process's own memory via NVML with a fallback to the device-wide delta; on GB10 that fallback would hit the same raw reading, so the two changes should agree on device_memory_info. #65204 / #65225 (granule rounding with offload connectors) is a different cause with the same error message.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[njhill]

Thanks @aijanai! fyi I added a small code simplification.

[njhill]

/ci run

[github-actions]

✅ @aijanai, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[github-actions]

✅ Triggered Buildkite CI #99614 for commit 7bce154a3825.

ℹ️ This PR is 13 commits behind upstream main. Merged upstream main into this branch as 7bce154a3825 before starting CI. Pull the branch before pushing more changes.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aijanai.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[njhill]

/ci retry

[github-actions]

✅ The previous CI build is still running: https://buildkite.com/vllm/ci/builds/99614

[njhill]

/ci retry

[github-actions]

✅ No failed, timed-out, or expired jobs need retrying in Buildkite CI #99614.

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99713 for commit 9e768ea65cd3.

ℹ️ This PR is 5 commits behind upstream main. Merged upstream main into this branch as 9e768ea65cd3 before starting CI. Pull the branch before pushing more changes.