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 8, 07:30 UTC
[Feature] Support CPU Offloading without Pytorch Pinned Memory that leads to doubled allocation (#32993)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Wei Zhao committed
Feb 13, 2026 at 11:11 UTC
59d53066d8daed5c2e39e3bce38ac308bc80a9ae
6 files changed
+129
-64
csrc/cuda_view.cu
+50
-25
@@ -2,33 +2,58 @@
2
#include <torch/cuda.h>
3
#include <cuda_runtime.h>
4
5
-// This function assumes that `cpu_tensor` is a CPU tensor allocated with pinned
6
-// memory, and that UVA (Unified Virtual Addressing) is enabled.
5
+// This function assumes that `cpu_tensor` is a CPU tensor,
6
+// and that UVA (Unified Virtual Addressing) is enabled.
7
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
8
TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
9
10
- // Get raw host pointer from CPU tensor
11
- void* host_ptr = cpu_tensor.data_ptr();
10
+ // handle empty tensor
11
+ if (cpu_tensor.numel() == 0) {
12
+ return torch::empty(cpu_tensor.sizes(),
13
+ cpu_tensor.options().device(torch::kCUDA));
14
+ }
15
+
16
+ if (cpu_tensor.is_pinned()) {
17
+ // If CPU tensor is pinned, directly get the device pointer.
18
+ void* host_ptr = const_cast<void*>(cpu_tensor.data_ptr());
19
+ void* device_ptr = nullptr;
20
+ cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
21
+ TORCH_CHECK(err == cudaSuccess,
22
+ "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
23
+
24
+ return torch::from_blob(
25
+ device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(),
26
+ [base = cpu_tensor](void*) {}, // keep cpu tensor alive
27
+ cpu_tensor.options().device(torch::kCUDA));
28
+ }
29
+
30
+ // If CPU tensor is not pinned, allocate a new pinned memory buffer.
31
+ torch::Tensor contiguous_cpu = cpu_tensor.contiguous();
32
+ size_t nbytes = contiguous_cpu.nbytes();
33
+
34
+ void* host_ptr = nullptr;
35
+ cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
36
+ if (err != cudaSuccess) {
37
+ AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err));
38
+ }
39
+
40
+ err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes,
41
+ cudaMemcpyDefault);
42
+ if (err != cudaSuccess) {
43
+ cudaFreeHost(host_ptr);
44
+ AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err));
45
+ }
46
13
- // Get a device pointer corresponding to the pinned host memory
47
void* device_ptr = nullptr;
15
- cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
16
- TORCH_CHECK(err == cudaSuccess,
17
- "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
18
-
19
- // We'll use the same sizes, strides, and dtype as the CPU tensor.
20
- // TODO: check if layout is respected.
21
- auto sizes = cpu_tensor.sizes();
22
- auto strides = cpu_tensor.strides();
23
- auto options = cpu_tensor.options().device(torch::kCUDA);
24
-
25
- // use default no-op deleter, since the memory is owned by the original CPU
26
- // tensor
27
- torch::Tensor cuda_tensor =
28
- torch::from_blob(device_ptr, sizes, strides, options);
29
-
30
- TORCH_CHECK(cuda_tensor.device().is_cuda(),
31
- "Resulting tensor is not on CUDA device");
32
-
33
- return cuda_tensor;
34
-}
48
+ err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
49
+ if (err != cudaSuccess) {
50
+ cudaFreeHost(host_ptr);
51
+ AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
52
+ }
53
+
54
+ auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
55
+
56
+ return torch::from_blob(device_ptr, contiguous_cpu.sizes(),
57
+ contiguous_cpu.strides(), deleter,
58
+ contiguous_cpu.options().device(torch::kCUDA));
59
+}
\ No newline at end of file
tests/basic_correctness/test_cpu_offload.py
+21
-2
@@ -1,10 +1,29 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import pytest
5
+
6
from ..utils import compare_two_settings
7
8
7
-def test_cpu_offload():
9
+@pytest.mark.parametrize("disable_pin_memory", [False, True])
10
+@pytest.mark.parametrize("disable_uva", [False, True])
11
+def test_cpu_offload(disable_pin_memory, disable_uva):
12
+ env_vars = {
13
+ "VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY": str(int(disable_pin_memory)),
14
+ "VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": str(int(disable_uva)),
15
+ }
16
+
17
+ args = ["--cpu-offload-gb", "1"]
18
+
19
+ # cuda graph only works with UVA offloading
20
+ if disable_uva:
21
+ args.append("--enforce-eager")
22
+
23
compare_two_settings(
9
- "hmellor/tiny-random-LlamaForCausalLM", [], ["--cpu-offload-gb", "1"]
24
+ model="hmellor/tiny-random-LlamaForCausalLM",
25
+ arg1=[],
26
+ arg2=args,
27
+ env1=None,
28
+ env2=env_vars,
29
)
vllm/envs.py
+10
@@ -230,6 +230,8 @@ if TYPE_CHECKING:
230
VLLM_USE_V2_MODEL_RUNNER: bool = False
231
VLLM_LOG_MODEL_INSPECTION: bool = False
232
VLLM_DEBUG_MFU_METRICS: bool = False
233
+ VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False
234
+ VLLM_WEIGHT_OFFLOADING_DISABLE_UVA: bool = False
235
VLLM_DISABLE_LOG_LOGO: bool = False
236
VLLM_LORA_DISABLE_PDL: bool = False
237
@@ -1542,6 +1544,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
1544
"VLLM_DEBUG_MFU_METRICS": lambda: bool(
1545
int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0"))
1546
),
1547
+ # Disable using pytorch's pin memory for CPU offloading.
1548
+ "VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY": lambda: bool(
1549
+ int(os.getenv("VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY", "0"))
1550
+ ),
1551
+ # Disable using UVA (Unified Virtual Addressing) for CPU offloading.
1552
+ "VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": lambda: bool(
1553
+ int(os.getenv("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA", "0"))
1554
+ ),
1555
# Disable logging of vLLM logo at server startup time.
1556
"VLLM_DISABLE_LOG_LOGO": lambda: bool(int(os.getenv("VLLM_DISABLE_LOG_LOGO", "0"))),
1557
# Disable PDL for LoRA, as enabling PDL with LoRA on SM100 causes
vllm/model_executor/model_loader/utils.py
+23
-17
@@ -11,6 +11,7 @@ import torch
11
from torch import nn
12
from typing_extensions import assert_never
13
14
+import vllm.envs as envs
15
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
16
from vllm.logger import init_logger
17
from vllm.model_executor.layers.attention import Attention, MLAAttention
@@ -25,6 +26,7 @@ from vllm.model_executor.model_loader.reload import (
26
from vllm.model_executor.models.interfaces import SupportsQuant
27
from vllm.tracing import instrument
28
from vllm.utils.platform_utils import is_pin_memory_available
29
+from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
30
31
logger = init_logger(__name__)
32
@@ -111,7 +113,8 @@ def process_weights_after_loading(
113
):
114
# TODO(lucas): see if there is a way to unify the signatures
115
# of process_weights_after_loading
114
- module.process_weights_after_loading(model_config.dtype)
116
+ with device_loading_context(module, target_device):
117
+ module.process_weights_after_loading(model_config.dtype)
118
119
# Needed for torchao model reloading via model.reload_weights
120
# @kylesayrs @jerryzh168 this can be removed if callers move to `reload_weights`
@@ -127,38 +130,41 @@ def device_loading_context(module: torch.nn.Module, target_device: torch.device)
130
return
131
132
original_device_states: dict[str, torch.device] = {}
133
+ uva_offloaded_parameters: list[str] = []
134
135
# Store original device states and move parameters to GPU if they're on CPU
136
for name, p in module.named_parameters():
137
if p.device.type == "cpu":
138
original_device_states[name] = p.device
139
p.data = p.data.to(target_device)
140
+ if getattr(p, "_vllm_is_uva_offloaded", False):
141
+ uva_offloaded_parameters.append(name)
142
# Parameters already on target device are not touched
143
144
try:
145
yield module
146
147
finally:
148
+ use_pin_memory = (
149
+ is_pin_memory_available()
150
+ and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY
151
+ )
152
# Restore parameters to their original devices, ignoring new parameters
143
- pin_memory = is_pin_memory_available()
153
for name, p in module.named_parameters():
154
if name in original_device_states:
155
original_device: torch.device = original_device_states[name]
147
- if original_device.type == "cpu":
148
- # `torch.empty_like` does not support `pin_memory` argument
149
- cpu_data = torch.empty_strided(
150
- size=p.data.size(),
151
- stride=p.data.stride(),
152
- dtype=p.data.dtype,
153
- layout=p.data.layout,
154
- device="cpu",
155
- pin_memory=pin_memory,
156
- )
157
- cpu_data.copy_(p.data)
158
- p.data = cpu_data
159
- else:
160
- p.data = p.data.to(original_device)
161
- # New parameters or parameters already on target device are untouched
156
+ p.data = p.data.to(original_device)
157
+
158
+ # parameter is UVA offloaded, but was replaced with a new device tensor
159
+ # re-offload it to CPU using UVA
160
+ if name in uva_offloaded_parameters and not getattr(
161
+ p, "_vllm_is_uva_offloaded", False
162
+ ):
163
+ cpu_data = p.data.to(device="cpu")
164
+ if use_pin_memory:
165
+ cpu_data = cpu_data.pin_memory()
166
+ p.data = get_accelerator_view_from_cpu_tensor(cpu_data)
167
+ p._vllm_is_uva_offloaded = True
168
169
170
_MODEL_ARCH_BY_HASH = dict[int, tuple[type[nn.Module], str]]()
vllm/model_executor/models/utils.py
+17
-18
@@ -13,6 +13,7 @@ from torch.func import functional_call
13
from torch.nn.modules.module import register_module_module_registration_hook
14
from transformers import PretrainedConfig
15
16
+import vllm.envs as envs
17
from vllm.config import VllmConfig
18
from vllm.distributed import (
19
get_tensor_model_parallel_rank,
@@ -633,11 +634,10 @@ def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
634
if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES:
635
return module
636
636
- pin_memory = is_pin_memory_available()
637
- uva_available = is_uva_available()
638
-
639
- assert uva_available, "V1 CPU offloading requires uva (pin memory) support"
640
- uva_offloading = True
637
+ pin_memory = (
638
+ is_pin_memory_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY
639
+ )
640
+ uva_offloading = is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA
641
642
# offload parameters to CPU
643
# use pin_memory if possible, which helps cudagraph capture speed
@@ -648,22 +648,16 @@ def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
648
# one module might have some parameters offloaded and some not
649
break
650
651
- # `torch.empty_like` does not support `pin_memory` argument
652
- cpu_data = torch.empty_strided(
653
- size=p.data.size(),
654
- stride=p.data.stride(),
655
- dtype=p.data.dtype,
656
- layout=p.data.layout,
657
- device="cpu",
658
- pin_memory=pin_memory,
659
- )
660
- cpu_data.copy_(p.data)
651
+ cpu_data = p.data.to(device="cpu")
652
+ if pin_memory:
653
+ cpu_data = cpu_data.pin_memory()
654
+
655
if not uva_offloading:
656
p.data = cpu_data
657
else:
664
- # keep the cpu data alive
665
- p._vllm_offloaded_cpu_data = cpu_data
658
p.data = get_accelerator_view_from_cpu_tensor(cpu_data)
659
+ p._vllm_is_uva_offloaded = True
660
+
661
_CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size()
662
offloaded_parameters = True
663
@@ -678,7 +672,12 @@ def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
672
k: v.to(device, non_blocking=True)
673
for k, v in module.state_dict().items()
674
}
681
- output = functional_call(module, device_state, args=args, kwargs=kwargs)
675
+
676
+ # set `tie_weights=False` as tied weights in original model
677
+ # become untied when calling .to(device) individually
678
+ output = functional_call(
679
+ module, device_state, args=args, kwargs=kwargs, tie_weights=False
680
+ )
681
module.forward = forward
682
return output
683
vllm/utils/torch_utils.py
+8
-2
@@ -678,12 +678,18 @@ def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tens
678
"""
679
Get an accelerator view of a CPU tensor using Unified Virtual Addressing (UVA).
680
"""
681
- assert cpu_tensor.is_pinned(), "CPU tensor must be pinned"
681
from vllm.platforms import current_platform
682
683
if current_platform.is_xpu():
684
+ assert cpu_tensor.is_pinned(), "CPU tensor must be pinned"
685
return torch.ops._C.get_xpu_view_from_cpu_tensor(cpu_tensor)
686
- return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)
686
+ elif current_platform.is_cuda():
687
+ return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)
688
+ else:
689
+ raise ValueError(
690
+ f"`get_accelerator_view_from_cpu_tensor` is currently "
691
+ f"not supported in: {current_platform.device_name}"
692
+ )
693
694
695
# Helper function used in testing.