[GodRishUniverse]

Overview

On ROCm, hipMemGetInfo counts every allocation as at least 2 MB, so after warmup the extensible KV cache (#56492) sizes itself from free memory that looks a few GB smaller than it is. This PR adds that memory back once, at the final sizing, using amdgpu's exact VRAM count. Only unpartitioned discrete ROCm GPUs are affected.

Root cause. The amdgpu driver rounds each allocation up to 2 MB in the free memory it reports, on purpose (commit 1ac354beecfd, "drm/amdgpu: Pessimistic availability based on rounded up allocations"), unless the GPU's memory is kept in system memory (apu_prefer_gtt). Its own allocation check uses exact sizes. The many small kernel code objects loaded during warmup therefore hide a few GB. The sysfs counter mem_info_vram_used counts exact sizes for the whole device.

Claims

• On an unpartitioned discrete ROCm GPU, the extensible KV cache sizes itself from exact device memory use. On Llama 3.1 8B / R9700 the KV cache went from 51,712 to 76,288 tokens, with no OOM and no failed requests (one run per config).
• The correction is never more than what is physically free and never negative. Other processes' allocations are not added back.
• Integrated GPUs, partitioned GPUs, a missing or mismatched sysfs count, and all non-ROCm platforms keep the current behaviour (no-op platform hooks).

Validation

Unit tests (the sysfs lookup and correction arithmetic need no GPU; the GPU tests skip off ROCm):

pytest tests/utils_/test_rocm_vram_lookup.py tests/v1/worker/test_extensible_kv_cache.py tests/v1/worker/test_gpu_worker.py tests/v1/worker/test_rocm_device_memory.py

On a ROCm GPU (Radeon AI PRO R9700, gfx1201). --optional adds tests that allocate all but 256 MB of the corrected free memory, both as 64 MB / 2 MB torch blocks and as one ExtensibleTensor chunk (the defragment=True case):

pytest tests/v1/worker/test_rocm_device_memory.py --optional

Serving comparison, main vs this PR, with a fresh compile cache for each run:

vllm serve meta-llama/Llama-3.1-8B-Instruct --max-model-len 8192 --gpu-memory-utilization 0.97

vllm bench serve --backend openai --endpoint /v1/completions --model meta-llama/Llama-3.1-8B-Instruct \
  --dataset-name random --random-input-len 512 --random-output-len 128 --ignore-eos \
  --request-rate inf --seed 0 --max-concurrency 128 --num-prompts 1280
# repeated with --max-concurrency 256 --num-prompts 2560

Llama 3.1 8B, TP1, Radeon AI PRO R9700 (32 GB), --gpu-memory-utilization 0.97. One run per config, so there is no variability figure:

	main	this PR
KV cache	51,712 tokens	76,288 tokens (+48%)
Output tok/s, concurrency 128 / 256	918.9 / 905.5	1,054.0 / 1,026.9
Median TTFT, concurrency 128 / 256	3,934 / 25,195 ms	1,084 / 17,110 ms
Most memory used / allowed	24.6 / 31.1 GB	27.8 / 31.1 GB
Failed requests / OOM	0 / 0	0 / 0

Not yet done:

• KV cache size across several starts.
• Multi-GPU, CDNA (MI300/MI355) and spec decode. I only have access to an R9700.

Limitations.

• gpu_memory_utilization now bounds exact use on unpartitioned discrete ROCm GPUs.
• Several instances on one GPU: with the extensible cache on, request_memory skips the startup check. An earlier instance's rounding is already in a later instance's baseline and is never added back, so a later instance gets a smaller cache, depending on start order. This doesn't cause OOM.
• KFD's page-table reserve and RAS pages are not in the exact count (about 80 MB on R9700). The 3 GB ROCm margin floor ([Bugfix][ROCm] Raise the extensible KV cache margin floor to 3 GiB #65418) covers them.
• Tested on R9700 (gfx1201) only. CDNA should behave the same according to the driver code, but this is unverified. In AMD CI, tests/v1/worker runs on mi355_dpx (partitioned, so it falls back), but other jobs run on mi300_1 / mi355_1. I don't know whether those pools are partitioned. If they aren't, this path runs there untested, and the correction could be limited to validated GPUs.
• The V1 runner and --no-enable-extensible-kv-cache are unchanged.

Not a duplicate. I found no other open PR correcting ROCm's 2 MB-rounded free memory. #65413 (merged) corrects UMA free memory and is the fallback inside the new hook. #55828 (open) changes which memory counts as this process's own and only takes effect on NVML. No runtime overlap and either merge order works.

Related: #56492 (extensible KV cache), #65413 (UMA free-memory correction), #65418 (ROCm 3 GiB margin floor), #55828 (process-scoped KV sizing), #60121 (asynchronous VRAM release on ROCm).

Note: From my understanding, amdsmi_get_gpu_memory_usage reads the same mem_info_vram_used so amdsmi and sysfs should be the same but if it can't be read then it will ask KFD and will again use 2MB rounded values that hipMemGetInfo reports.

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.

AI assistance was used. I have reviewed the changes.

[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.

[mergify]

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

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

[mergify]

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

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

[mergify]

Documentation preview: https://vllm--63401.org.readthedocs.build/en/63401/

[mergify]

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

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

[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]

The approach and math look sound: the correction is capped at what's physically free, the erase and read-order races I traced at the final reading only shrink it, and non-ROCm platforms return early. This review is from reading the code only; I haven't run it. Main requests, with details inline:

1. Move the ROCm logic out of common files. 129 of the 254 source lines are in common files but only do anything on ROCm.
2. Fix the race at the final reading, then report results across several starts.
3. Widen validation beyond R9700, TP1 and one model.
4. Rework the tests to follow tests/AGENTS.md.

Other points:

• Multi-instance limitation: the description says a later instance's start-up check uses rounded free memory, but request_memory skips that check when the extensible cache is on. From reading the code, the first instance's rounding is already in the second's baseline and never gets added back, so the second instance silently gets a smaller cache, depending on start order. Please correct the description. --gpu-memory-utilization now meaning exact use on ROCm probably deserves a docs note too.
• Overlapping PRs: #65413 rewrites the same get_memory_info lines and adds its own device_memory_info correction, and #55828 changes the non_kv_used calculation this correction feeds into. Please agree a merge order and check whether the free-memory corrections can share one mechanism, though #65413 corrects every reading and this PR only the final one.
• PR description: please use the template sections (Overview, Claims, Validation, Details) and give the run count and variability for the benchmark table.

[njhill]

The baseline samples for 200 ms to wait out erases, but this reading comes straight after empty_cache(), just as warmup's cached memory is freed. Memory still being erased counts as used, so the correction shrinks and can reach 0, silently falling back to main's size. _min_across_ranks then applies the unluckiest rank's result to every rank.

Could both readings share one sampling helper? If that lives behind the platform hook, free memory needs re-reading after sampling, so the hook may be better off returning the corrected value. Please also report KV cache size across several starts, including multi-GPU.

[GodRishUniverse]

I am running tests on R9700 TP4 right now, I don't have access to MI355x or MI300 right now @njhill to test it on CDNA.

[njhill]

Two validation gaps:

• Single allocation: this only allocates 64 MB and 2 MB chunks. defragment=True (registers_kv_cache_memory) remaps each segment as one driver allocation, which is the case kernel commit 1ac354beecfd's rounding is meant to help. Could you cover it?
• CDNA and multi-GPU: only R9700 TP1 is tested so far, and [Bugfix][ROCm] Raise the extensible KV cache margin floor to 3 GiB #65418 just fixed OOMs on MI355 spec decode. This removes about 3.4 GB of unintended slack wherever it activates. The 3 GiB floor should still cover that since KFD checks exact sizes, but it hasn't been shown on CDNA. Could you do an unpartitioned MI300 or MI355 run with spec decode? Otherwise, consider limiting the correction to validated GPUs.

Also, the description says AMD CI falls back because it's partitioned. The worker tests do use mi355_dpx, but .buildkite/test-amd.yaml also has mi355_1 and mi300_1 pools. If those are unpartitioned, this path runs in CI on hardware that hasn't been tested. Do you know which mode they use?

[GodRishUniverse]

@njhill I don't have access to MI300 or MI355 at the moment. If you have access, then I'd appreciate if you can run it. I'll cc @dllehr-amd just in case, if he is able to run on MI.

I have not run CI

Also, the description says AMD CI falls back because it's partitioned.

This is my agent's assumption

[njhill]

Possible flake: the child process's own rounding also counts toward the correction, so 32 MB may be tight.

[GodRishUniverse]

Made the upper limit as 256 MB, does that work @njhill ?

[mergify]

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

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

[twu3202]

Re the merge order with #55828 (mine): the two PRs overlap in the diff but not at runtime.

• [Core] Size the KV cache, including the extensible one, from this process's own device memory #55828's process-scoped path only runs where get_process_memory_usage() returns a number, and only NvmlCudaPlatform implements it. On ROCm it returns None, so measure_kv_cache_blocks keeps the device-wide free-memory delta, and your correction feeds into non_kv_used and the headroom exactly as it does in this PR.
• Only ROCm implements get_device_memory_used(), so on CUDA your correction is 0.
• In the diff, both PRs add a Platform method at the same spot in interface.py, and both edit right after the same free-memory reading in measure_kv_cache_blocks (now device_memory_info after [Bugfix] Size the extensible KV cache from UMA-corrected free memory #65413).

So either order works for me, and whichever lands second resolves a small conflict. #55828 is already rebased onto #65413.

On sharing one mechanism: #55828 isn't a free-memory correction. It only changes which memory counts as this process's own use and leaves the free-memory reading alone, so I don't think it has anything to share with this PR or #65413.

[GodRishUniverse]

Converting to draft while I rebase and move the ROCm logic behind platform hooks. Will re-request review when ready.

[GodRishUniverse]

@njhill
Gemma 4 31B-it, TP4 on R9700, --max-model-len 4096, --gpu-memory-utilization 0.97, 3 starts per side, fresh compile cache each start. Main is 25e7c5728d.

At --max-model-len 8192, main failed to start (KV cache too small for one request) while the PR run did go ahead and start.

	main	PR
KV cache	4,359 tokens	32,744 tokens
Memory added back per GPU	n/a	5.98 GiB
Output tok/s	121.1–121.3	482.8–483.7
Median TTFT	126 s	17.9 s
Median TPOT	54.5 ms	104.3–104.5 ms
Failed requests / OOM	0 / 0	0 / 0
Peak VRAM, worst GPU	19.3 GiB	25.5 GiB

Commands:

HIP_VISIBLE_DEVICES=0,1,2,3 vllm serve google/gemma-4-31B-it --served-model-name gemma4 \
  --tensor-parallel-size 4 --max-model-len 4096 --gpu-memory-utilization 0.97 --port 8000

vllm bench serve --backend openai --endpoint /v1/completions --model gemma4 \
  --tokenizer google/gemma-4-31B-it --dataset-name random --random-input-len 512 \
  --random-output-len 128 --ignore-eos --request-rate inf --seed 0 \
  --max-concurrency 128 --num-prompts 1280

[GodRishUniverse]

Re the merge order with #55828 (mine): the two PRs overlap in the diff but not at runtime.

* [[Core] Size the KV cache, including the extensible one, from this process's own device memory #55828](https://github.com/vllm-project/vllm/pull/55828)'s process-scoped path only runs where `get_process_memory_usage()` returns a number, and only `NvmlCudaPlatform` implements it. On ROCm it returns `None`, so `measure_kv_cache_blocks` keeps the device-wide free-memory delta, and your correction feeds into `non_kv_used` and the headroom exactly as it does in this PR.

* Only ROCm implements `get_device_memory_used()`, so on CUDA your correction is 0.

* In the diff, both PRs add a `Platform` method at the same spot in `interface.py`, and both edit right after the same free-memory reading in `measure_kv_cache_blocks` (now `device_memory_info` after [[Bugfix] Size the extensible KV cache from UMA-corrected free memory #65413](https://github.com/vllm-project/vllm/pull/65413)).

So either order works for me, and whichever lands second resolves a small conflict. #55828 is already rebased onto #65413.

On sharing one mechanism: #55828 isn't a free-memory correction. It only changes which memory counts as this process's own use and leaves the free-memory reading alone, so I don't think it has anything to share with this PR or #65413.

I think #65413 got merged. Yeah apologies, I meant that the PR is overlapping in diff but runtime is completely separate. The edits now should reduce overlap now as I changed my edits and scoped it to a new file.