[GirasoleY]

Overview

Functionality

• SimpleCPUOffloadConnector allocates each rank's CPU KV caches as views into one host mapping sized to the resolved KV blocks, rounded up to the selected page size. cpu_bytes_to_use_per_rank remains the cache budget. Registration is split into chunks of at most 256 GiB, the limit HiSparse already uses, without bisecting a KV block or backing page.
• The new cpu_hugepage_block_size key ("1GB" or "2MB") backs that mapping with explicit HugeTLB pages; without it, ordinary pages are used.
• The mapping inherits the worker's NUMA memory policy, including vLLM's NUMA binding options or an external launcher policy.
• A HugeTLB pool or container hugetlb limit that's too small fails at startup with an error naming the shortfall, never falling back to 4 KiB pages.
• The mapping is excluded from fork() (MADV_DONTFORK), so child processes can't make pinned pages copy-on-write.

{"kv_connector": "SimpleCPUOffloadConnector", "kv_role": "kv_both",
 "kv_connector_extra_config": {"cpu_bytes_to_use_per_rank": 137438953472,
                               "cpu_hugepage_block_size": "1GB"}}

Performance evidence from the earlier revision

• In vllm serve (TP=2, 64 GiB per rank, ordinary pages), engine init takes 19.0 s vs 20.2 s before.
• With 1 GiB pages, pinning 128 GiB takes 0.19 s instead of 2.5 s on ordinary pages.
• GPU↔CPU copy bandwidth is unchanged at about 42 GB/s.

Known limitations

• Registration uses bounded chunks to avoid oversized CUDA registrations ([Bugfix][KV Offload] Register the offload region in chunks #51081). pin_tensor, unpin_tensor, and the registration boundary calculation live in vllm/utils/host_memory.py, shared by HiSparse and CPU/disk offload as applicable.
• Each rank needs ceil(resolved cache bytes / page size) HugeTLB pages, reserved on the host and, on Kubernetes, requested as hugepages-<size>.
• Up-front population needs Linux 5.14+ for HugeTLB pages; ordinary pages fall back to first touch on older kernels.

Related: #45693 (same key name for OffloadingConnector), #58852, #59430, #51081.

Validation

The AMD measurements below describe the earlier revision. Current CUDA validation is recorded separately below. NUMA placement now follows the worker policy.

MI350X (gfx950), ROCm 7.2.3, Linux 6.8, vllm/vllm-openai-rocm:nightly-92044241a0 with this PR's three changed files overlaid; the image's copies were byte-identical to this PR's base.

Test	Result
pytest tests/v1/simple_kv_offload/test_host_buffer.py tests/v1/simple_kv_offload/test_worker.py	32 passed
8 ranks × 38 GiB (1 GiB pages), random store/load with byte-level checks, 120 s	0 mismatches; every rank's pages on its GPU's NUMA node
More 1 GiB pages than the host has (330 on 320)	Fails at mmap with a clear error
150 GiB soak, 10 min, 6,724 concurrent fork()s	0 mismatches
vllm serve (TP=2, 32 GiB of 1 GiB pages per rank): SIGKILL a worker, SIGKILL the engine core, SIGTERM, then restart; plus SIGKILL/SIGSEGV/exception at buffer level, including mid-populate, mid-pin, and with forked children still alive	Requests served; the HugeTLB pool and GPU memory return to baseline within 5 s each time, and the restarted engine serves

The 8-rank, soak and pool rows ran on the previous revision, which populated pages with 16 threads; this revision populates with one madvise call and was re-checked with the unit tests and the vllm serve crash/restart row.

The earlier AMD validation did not include CUDA or an offload-on/off output comparison. The subsequent CUDA runs below include a small-model output smoke comparison; a full model-accuracy evaluation remains outstanding. AMD CI already runs tests/v1/simple_kv_offload; the HugeTLB tests skip on hosts without a pool.

Review fixes (76365e976b, de6a707b43)

These AI-assisted fixes allocate only the resolved cache bytes, inherit NUMA policy, and preserve one registration per matching unregister. Shared helpers moved to vllm/utils/host_memory.py; their function ASTs are unchanged by the move. The independent HIP SDMA optimization is handled by #65559.

• An isolated reproduction showed that the original chunked pin_tensor left one registered range after HiSparse teardown for a tensor larger than 1 GiB. pin_tensor now registers exactly one range again.
• Pre-commit checks passed, including Ruff and mypy.
• tests/v1/worker/test_utils.py::test_hisparse_registration_chunks_end_between_host_blocks: collection blocked by missing FlashAttention CUDA extensions. The environment was not modified.
• Subsequent CUDA integration and output smoke results are recorded below. No before/after performance benchmark was run.

CUDA validation and alignment fix (26bcba935b, 2026-11-05)

Validated on an exclusive Slurm node with 4 NVIDIA GB200 GPUs, driver 580.126.09, CUDA 13.0 / PyTorch 2.14.0, and Linux 6.17 with 64 KiB OS pages. The isolated environment used this PR's Python source and native/generated dependencies from vllm/vllm-openai:nightly-40941f2c4ba96fba0a7e282a5355e89cd74d9ac5.

The first large HugeTLB run exposed a real bug: under NUMA binding, a registration ending inside a 2 MiB backing page returned cudaErrorInvalidValue. For example, registering 1 GiB - 64 KiB failed while a whole 1 GiB succeeded on the same mapping. The shared boundary helper now receives the actual backing-page size; HiSparse keeps its OS-page default. The previously failing four-rank scenario passed after this correction.

Check	Result
Targeted host-buffer, worker, mixed-layout, ordering, boundary and NUMA integration tests	16 passed; 1 skipped because this kernel has no 1 GiB HugeTLB support
Four concurrent HugeTLB ranks, 38/38/12/12 GiB, 120 s	21,792 byte-exact round trips, zero mismatches; all pages on the configured NUMA nodes
150 GiB ordinary-page soak, 600 s, concurrent forks	18,490 round trips and 36,980 forks, zero mismatches
Qwen3-0.6B NUMA integration	Restored 16 CPU KV blocks; generated token IDs matched the cold result
Qwen3-0.6B serving, TP=2, 32 GiB HugeTLB per rank	Three greedy prompts matched offload-disabled outputs; restart served successfully
Worker SIGKILL / engine SIGKILL / server SIGTERM	HugeTLB pool and GPU memory returned to baseline in 4.10 / 3.38 / 0.58 s
32 GiB buffer: SIGKILL during population, SIGKILL after partial registration, SIGSEGV, exception	All resources recovered in under 1 s; core dumps disabled
Final cleanup	HugeTLB configuration restored to zero pages, all GPUs at zero allocated MiB, container removed, Slurm allocation released

The ordinary-page soak started before the HugeTLB alignment correction; its OS-page arithmetic is unchanged, and ordinary-page focused tests passed again afterward. The CUDA run uses four ranks rather than the earlier eight-rank AMD setup. NUMA node 1 could reserve only about 25.9 GiB of 2 MiB pages, so its two ranks used 12 GiB each. This validates 2 MiB HugeTLB on GB200, not 1 GiB pages or all CUDA platforms. The output check is a three-prompt serving smoke test, not a full accuracy evaluation. No speedup is claimed from these runs.

This commit also extended the existing boundary test to HugeTLB page sizes. Pre-commit checks, including mypy, passed.

Simplification (2a6867aa5c)

• The registration chunk limit is now the shared HOST_REGISTER_CHUNK_BYTES (256 GiB) instead of 1 GiB. With a 1 GiB limit and 1 GiB pages, every boundary had to be a 1 GiB multiple that is also a block boundary, so any block size that does not divide 1 GiB (for example MLA: 576 × 2 × 64 = 73728 B) failed at startup. With the larger limit, a boundary exists at lcm(block, page), which is 9 GiB in that example. HiSparse behavior is unchanged.
• HostBuffer is now a function, alloc_host_buffer(), because torch.frombuffer keeps the mapping alive. It reuses the existing MADV_POPULATE_WRITE helper and its pre-5.14 fallback from shared_offload_region.py. cpu_hugepage_block_size accepts only the documented 2MB, 1GB and none.
• The docs section, comments and the populate/pin timing log are shorter.
• Tests check page rounding, HugeTLB pool consumption, fail-closed allocation and a GPU→CPU→GPU round trip on ordinary and HugeTLB pages. They no longer read smaps flags, kernel zeroing or private attributes. The single-fix pin/unpin test and the unrelated served_model_name change to the integration test were removed.
• Validation: pre-commit (ruff, mypy) passed. tests/v1/simple_kv_offload/test_host_buffer.py: 7 passed, 2 skipped (no free hugepages on the host). test_worker.py and tests/v1/worker/test_utils.py could not be collected locally because the FlashAttention CUDA extensions are missing; they rely on CI. The GB200 runs above predate this commit.

---

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.

Not a duplicate: the related PRs target OffloadingConnector and HiSparse.

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

[YannikHinteregger]

Thanks for the contribution! I left some inline comments and hope I provided enough context.

Could you please also include benchmarks showing total engine startup time before and after this change

[YannikHinteregger]

Could we keep one buffer but register it in chunks? A single registration of 512 GiB or more fails on CUDA (#51081) so large per-rank buffers would fail at startup.

We’re also working on moving CPU pinning into a background thread. That path needs smaller chunks with a short yield between calls because a long registration can block other CUDA calls and delay inference (CUDA docs)

[GirasoleY]

addressed in 76e2461

Currently left this in main thread, but can definitely move to a background thread if needed.

[YannikHinteregger]

Could you please mark this test to skip on systems where NUMA placement is unsupported or blocked

[GirasoleY]

addressed in 76e2461

[YannikHinteregger]

Switching from torch.zeros to mmap bypasses PyTorch’s allocator, which honors THP_MEM_ALLOC_ENABLE=1. Could we apply MADV_HUGEPAGE ourselves when this flag is enabled and explicit HugeTLB pages are disabled?

[GirasoleY]

addressed in 76e2461

[mergify]

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

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

[mergify]

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

[GirasoleY]

/ci run

[GirasoleY]

start up time before after for Qwen3-8b

[mergify]

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

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

[ivanium]

Thanks for the PR! @GirasoleY I made a bunch of small simplifications; please help double check