A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 04:45 UTC
[Bugfix] Respect VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY in prefetch offloader (#37699)
Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Yufeng He committed
Apr 15, 2026 at 11:43 UTC
bcc2306cefa4179c548d3e638e7a22a88d281733
4 files changed
+23
-11
vllm/model_executor/offloader/__init__.py
+2
@@ -8,6 +8,7 @@ from vllm.model_executor.offloader.base import (
8
create_offloader,
9
get_offloader,
10
set_offloader,
11
+ should_pin_memory,
12
)
13
from vllm.model_executor.offloader.prefetch import PrefetchOffloader
14
from vllm.model_executor.offloader.uva import UVAOffloader
@@ -20,4 +21,5 @@ __all__ = [
21
"create_offloader",
22
"get_offloader",
23
"set_offloader",
24
+ "should_pin_memory",
25
]
vllm/model_executor/offloader/base.py
+14
@@ -10,7 +10,9 @@ from typing import TYPE_CHECKING
10
11
import torch.nn as nn
12
13
+import vllm.envs as envs
14
from vllm.logger import init_logger
15
+from vllm.utils.platform_utils import is_pin_memory_available
16
17
if TYPE_CHECKING:
18
from vllm.config import OffloadConfig
@@ -18,6 +20,18 @@ if TYPE_CHECKING:
20
logger = init_logger(__name__)
21
22
23
+def should_pin_memory() -> bool:
24
+ """Check if pinned memory should be used for weight offloading.
25
+
26
+ Combines the platform capability check with the user override env var.
27
+ On unified-memory systems (e.g. GH200) pinned memory eats into GPU
28
+ memory, so users can disable it via VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY.
29
+ """
30
+ return (
31
+ is_pin_memory_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY
32
+ )
33
+
34
+
35
"""
36
class relation:
37
vllm/model_executor/offloader/prefetch.py
+4
-5
@@ -20,8 +20,7 @@ import torch.nn as nn
20
# Import prefetch_ops to register custom ops at module load time
21
import vllm.model_executor.offloader.prefetch_ops # noqa: F401
22
from vllm.logger import init_logger
23
-from vllm.model_executor.offloader.base import BaseOffloader
24
-from vllm.utils.platform_utils import is_pin_memory_available
23
+from vllm.model_executor.offloader.base import BaseOffloader, should_pin_memory
24
25
logger = init_logger(__name__)
26
@@ -528,7 +527,7 @@ class _ModuleOffloader:
527
gpu_buffer = offloader._gpu_buffer
528
assert cpu_storage is not None, "CPU storage not initialized"
529
assert gpu_buffer is not None, "GPU buffer not assigned"
531
- assert not is_pin_memory_available() or cpu_storage.is_pinned(), (
530
+ assert not should_pin_memory() or cpu_storage.is_pinned(), (
531
f"CPU storage for {name} is not pinned! "
532
"non_blocking=True H2D copy from non-pinned memory "
533
"causes stream synchronization that breaks "
@@ -629,7 +628,7 @@ class _CpuParamOffloader(_BaseParamOffloader):
628
original GPU tensor is garbage collected.
629
"""
630
param = self._param
632
- pin_memory = is_pin_memory_available()
631
+ pin_memory = should_pin_memory()
632
633
# Create pinned CPU storage and copy current GPU data
634
self._cpu_storage = torch.empty_strided(
@@ -666,7 +665,7 @@ class _CpuParamOffloader(_BaseParamOffloader):
665
param = self._param
666
667
if param.data.device.type == "cpu":
669
- if is_pin_memory_available() and not param.data.is_pinned():
668
+ if should_pin_memory() and not param.data.is_pinned():
669
pinned = torch.empty_strided(
670
size=param.data.size(),
671
stride=param.data.stride(),
vllm/model_executor/offloader/uva.py
+3
-6
@@ -10,9 +10,9 @@ from torch.func import functional_call
10
11
import vllm.envs as envs
12
from vllm.logger import init_logger
13
-from vllm.model_executor.offloader.base import BaseOffloader
13
+from vllm.model_executor.offloader.base import BaseOffloader, should_pin_memory
14
from vllm.utils.mem_utils import format_gib
15
-from vllm.utils.platform_utils import is_pin_memory_available, is_uva_available
15
+from vllm.utils.platform_utils import is_uva_available
16
from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
17
18
logger = init_logger(__name__)
@@ -43,10 +43,7 @@ class UVAOffloader(BaseOffloader):
43
self.cpu_offload_bytes = 0
44
self.cpu_offload_params = cpu_offload_params or set()
45
46
- self.pin_memory = (
47
- is_pin_memory_available()
48
- and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY
49
- )
46
+ self.pin_memory = should_pin_memory()
47
self.uva_offloading = (
48
is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA
49
)