[Bugfix] Respect VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY in prefetch offloader

[he-yufeng]

Purpose

Fixes #37672

The prefetch CPU offloader (--offload-backend prefetch) ignores VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY and unconditionally allocates pinned memory. On unified-memory systems like GH200, pinned memory consumes GPU memory, causing OOM during model loading even when the env var is set to 1.

The UVA offloader already handles this correctly (added in #32993). This PR applies the same pattern to the prefetch backend and consolidates the duplicated check into a shared helper.

Changes

1. base.py — new should_pin_memory() helper that combines is_pin_memory_available() with the env var check, eliminating duplication across files

2. prefetch.py — three places in _CpuParamOffloader now use the shared helper:

   • _offload_to_cpu_internal() — pin_memory flag when allocating initial CPU storage
   • _update_cpu_storage_from_param() — skips re-pinning when the env var is set
   • start_onload_to_static() assertion — no longer fires when pinning is intentionally disabled

3. uva.py — self.pin_memory assignment now uses the same shared helper (functionally identical, just DRY)

4. __init__.py — exports should_pin_memory for external use

Other files in vllm/model_executor/offloader/ (prefetch_ops.py) don't deal with pin memory at all, so no changes needed.

Test Plan

• Existing tests/basic_correctness/test_cpu_offload.py covers VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY for the UVA backend
• The env var handling is identical between UVA and prefetch — both just gate pin_memory on the same boolean via the shared helper
• Verified no other references to the raw pattern remain in the offloader directory

[gemini-code-assist]

Code Review

This pull request correctly extends the handling of the VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY environment variable to the prefetch offloader, aligning its behavior with the UVA offloader. The changes prevent unconditional pinned memory allocation, which can cause out-of-memory errors on unified-memory systems. While the logic is correct, I've identified a code duplication issue across three locations. I've added a high-severity comment suggesting a refactoring to a shared helper function to improve maintainability and prevent future bugs.

[gemini-code-assist]

The logic to decide whether to use pinned memory, is_pin_memory_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY, is now duplicated in three places within this file. This creates a maintenance risk, as future changes to this logic might not be applied consistently across all locations, potentially re-introducing bugs.

To improve maintainability, I recommend extracting this logic into a single, module-level helper function.

For example:

def _should_pin_memory() -> bool:
    """Returns True if memory should be pinned for offloading."""
    return (is_pin_memory_available()
            and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY)

This function can then be used in all three modified locations:

1. In _CpuParamOffloader._offload_to_cpu_internal (this location):

   pin_memory = _should_pin_memory()

2. In _CpuParamOffloader._update_cpu_storage_from_param:

   if _should_pin_memory() and not param.data.is_pinned():
       ...

3. In PrefetchOffloader.start_onload_to_static:

   assert not _should_pin_memory() or cpu_storage.is_pinned(), (
       ...
   )

This refactoring will make the code cleaner, less error-prone, and easier to understand and modify in the future.

[ehfd]

@gemini-code-assist This is not just limited to prefetch.py. You need to take a look at all files in vllm/model_executor/offloader.

[gemini-code-assist]

You are absolutely correct, @he-yufeng. My apologies for the oversight. I will expand the scope of the refactoring to include all relevant files within the vllm/model_executor/offloader directory to ensure consistency in handling the VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY environment variable.

[he-yufeng]

Good point, thanks for the catch. I went through every file in vllm/model_executor/offloader/:

• base.py — added a shared should_pin_memory() helper that consolidates the is_pin_memory_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY check
• prefetch.py — all 3 call sites now use the helper
• uva.py — switched to the same helper (was already correct, just DRY'd it up)
• prefetch_ops.py — no pin memory usage, nothing to change
• __init__.py — just exports, added should_pin_memory

Also noted there's one more instance in model_loader/utils.py (outside the offloader dir) that could be migrated to the helper in a follow-up if desired.

[ehfd]

@gemini-code-assist Change the above review to not just be limited to vllm/model_executor/offloader/prefetch.py. You need to take a look at all files in vllm/model_executor/offloader.

[gemini-code-assist]

Thanks for the clarification, @ehfd. I've reviewed the changes in vllm/model_executor/offloader/prefetch.py. However, based on the provided diff, vllm/model_executor/offloader/prefetch.py is the only file within the vllm/model_executor/offloader directory that has been modified in this pull request. Therefore, my review is limited to the changes in this specific file.

[ehfd]

@gemini-code-assist Then, the solution is to retain the current implementation, since it's exactly the same way that works on vllm/model_executor/offloader/uva.py

[gemini-code-assist]

I have noted your comment regarding the implementation's consistency with vllm/model_executor/offloader/uva.py.

[ehfd]

Please check https://docs.vllm.ai/en/latest/contributing/#linting @he-yufeng

[PythonicMind2]

Thanks for the fix!
Pinned memory should definitely be respected in the prefetch path too, not just in the main offloading path. The intent is clear and the bug is real.

prefetch_offloader.py, line 84 concern
The variable name pin is too terse. Since this is a negation of a 'disable' flag it's easy to get the logic backwards on a quick read. Suggest:

use_pin_memory = not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY
return torch.empty(size, dtype=torch.uint8, pin_memory=use_pin_memory)

Also is it safe to read envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY every call, or should we cache this at __init__ time to avoid overhead in tight loops?

tests/worker/test_prefetch_offloader.py concern
No test added. This bug was silently ignoring the env var for the prefetch path. Please add a parametrized test covering both VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY=1 and =0, asserting buf.is_pinned() accordingly.

prefetch_offloader.py, line 12 nit
The new envs import is at the bottom of the import block. Please keep it sorted with the other vllm imports alphabetically.

Once the test is added and naming cleaned up this looks good to go, nice catch on the prefetch path being overlooked.

[he-yufeng]

Ran ruff check and ruff format on the offloader directory — all passing now. Thanks for the heads up.

[ehfd]

@youkaichao @mgoin @22quinn @LucasWilkinson @zou3519 @ywang96 @hmellor

Follow-up to #32993 and #29941.

[he-yufeng]

ping on this — GH200 users are still hitting the OOM.

[ehfd]

It is definitely needed. Else GH200 doesn't work properly with unified memory at all.