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 15, 17:15 UTC
[V1] Fully Transparent Implementation of CPU Offloading (#15354)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
Mar 31, 2025 at 20:22 UTC
555aa21905a1d725b44a29ea8cbebf218ff14558
12 files changed
+148
-25
CMakeLists.txt
+1
@@ -234,6 +234,7 @@ set(VLLM_EXT_SRC
234
"csrc/activation_kernels.cu"
235
"csrc/layernorm_kernels.cu"
236
"csrc/layernorm_quant_kernels.cu"
237
+ "csrc/cuda_view.cu"
238
"csrc/quantization/gptq/q_gemm.cu"
239
"csrc/quantization/compressed_tensors/int8_quant_kernels.cu"
240
"csrc/quantization/fp8/common.cu"
csrc/cuda_view.cu
new
+39
@@ -0,0 +1,39 @@
1
+#include <torch/all.h>
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.
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();
12
+
13
+ // Get a device pointer corresponding to the pinned host memory
14
+ 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
+ // from_blob signature: from_blob(void *data, IntArrayRef sizes, ..., Deleter,
26
+ // const TensorOptions &) Provide a no-op deleter. The CPU tensor holds the
27
+ // memory, so we don't free it here.
28
+ auto deleter = [](void*) {
29
+ // no-op, since the memory is owned by the original CPU tensor
30
+ };
31
+
32
+ torch::Tensor cuda_tensor =
33
+ torch::from_blob(device_ptr, sizes, strides, deleter, options);
34
+
35
+ TORCH_CHECK(cuda_tensor.device().is_cuda(),
36
+ "Resulting tensor is not on CUDA device");
37
+
38
+ return cuda_tensor;
39
+}
csrc/ops.h
+2
@@ -119,6 +119,8 @@ void advance_step_flashinfer(
119
torch::Tensor& paged_kv_indices, torch::Tensor& paged_kv_indptr,
120
torch::Tensor& paged_kv_last_page_len, torch::Tensor& block_table_bounds);
121
122
+torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
123
+
124
#ifndef USE_ROCM
125
torch::Tensor aqlm_gemm(const torch::Tensor& input, const torch::Tensor& codes,
126
const torch::Tensor& codebooks,
csrc/torch_bindings.cpp
+4
@@ -31,6 +31,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
31
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
32
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
33
34
+ ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
35
+ ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
36
+ &get_cuda_view_from_cpu_tensor);
37
+
38
// Attention ops
39
// Compute the attention between an input query and the cached
40
// keys/values using PagedAttention.
tests/basic_correctness/test_cpu_offload.py
-7
@@ -1,15 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-import pytest
4
-
3
from ..utils import compare_two_settings
4
5
8
-@pytest.fixture(scope="function", autouse=True)
9
-def use_v0_only(monkeypatch):
10
- monkeypatch.setenv('VLLM_USE_V1', '0')
11
-
12
-
6
def test_cpu_offload():
7
compare_two_settings("meta-llama/Llama-3.2-1B-Instruct", [],
8
["--cpu-offload-gb", "1"])
tests/kernels/test_uva.py
new
+61
@@ -0,0 +1,61 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import pytest
3
+import torch
4
+
5
+from vllm.utils import get_cuda_view_from_cpu_tensor, is_uva_available
6
+
7
+CUDA_DEVICES = [
8
+ f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
9
+]
10
+
11
+
12
+@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
13
+@pytest.mark.parametrize("device", CUDA_DEVICES)
14
+def test_cpu_write(device):
15
+ torch.set_default_device(device)
16
+ cpu_tensor = torch.zeros(10,
17
+ 10,
18
+ device="cpu",
19
+ pin_memory=True,
20
+ dtype=torch.int32)
21
+ cuda_view = get_cuda_view_from_cpu_tensor(cpu_tensor)
22
+ assert cuda_view.device.type == "cuda"
23
+
24
+ assert cuda_view[0, 0] == 0
25
+ assert cuda_view[2, 3] == 0
26
+ assert cuda_view[4, 5] == 0
27
+
28
+ cpu_tensor[0, 0] = 1
29
+ cpu_tensor[2, 3] = 2
30
+ cpu_tensor[4, 5] = -1
31
+
32
+ cuda_view.mul_(2)
33
+ assert cuda_view[0, 0] == 2
34
+ assert cuda_view[2, 3] == 4
35
+ assert cuda_view[4, 5] == -2
36
+
37
+
38
+@pytest.mark.skipif(not is_uva_available(), reason="UVA is not available.")
39
+@pytest.mark.parametrize("device", CUDA_DEVICES)
40
+def test_gpu_write(device):
41
+ torch.set_default_device(device)
42
+ cpu_tensor = torch.zeros(10,
43
+ 10,
44
+ device="cpu",
45
+ pin_memory=True,
46
+ dtype=torch.int32)
47
+ cuda_view = get_cuda_view_from_cpu_tensor(cpu_tensor)
48
+ assert cuda_view.device.type == "cuda"
49
+
50
+ assert cuda_view[0, 0] == 0
51
+ assert cuda_view[2, 3] == 0
52
+ assert cuda_view[4, 5] == 0
53
+
54
+ cuda_view[0, 0] = 1
55
+ cuda_view[2, 3] = 2
56
+ cuda_view[4, 5] = -1
57
+ cuda_view.mul_(2)
58
+
59
+ assert cpu_tensor[0, 0] == 2
60
+ assert cpu_tensor[2, 3] == 4
61
+ assert cpu_tensor[4, 5] == -2
\ No newline at end of file
tests/quantization/test_cpu_offload.py
-7
@@ -10,13 +10,6 @@ from tests.quantization.utils import is_quant_method_supported
10
from ..utils import compare_two_settings
11
12
13
-@pytest.fixture(scope="function", autouse=True)
14
-def use_v0_only(monkeypatch):
15
- # Fall back to V0 if cpu offloading is enabled.
16
- # Fixture is required to that baseline uses V0.
17
- monkeypatch.setenv('VLLM_USE_V1', '0')
18
-
19
-
13
@pytest.mark.skipif(not is_quant_method_supported("fp8"),
14
reason="fp8 is not supported on this GPU type.")
15
def test_cpu_offload_fp8():
vllm/config.py
+3
-2
@@ -3562,9 +3562,10 @@ class VllmConfig:
3562
3563
if self.cache_config is not None and \
3564
self.cache_config.cpu_offload_gb > 0 and \
3565
- self.compilation_config.level != CompilationLevel.NO_COMPILATION:
3565
+ self.compilation_config.level != CompilationLevel.NO_COMPILATION \
3566
+ and not envs.VLLM_USE_V1:
3567
logger.warning(
3567
- "CPU offload is not supported with `torch.compile` yet."
3568
+ "CPU offload is not supported with `torch.compile` in v0 yet."
3569
" Disabling `torch.compile`.")
3570
self.compilation_config.level = CompilationLevel.NO_COMPILATION
3571
vllm/engine/arg_utils.py
-6
@@ -1595,12 +1595,6 @@ class EngineArgs:
1595
recommend_to_remove=False)
1596
return False
1597
1598
- # No CPU offloading yet.
1599
- if self.cpu_offload_gb != EngineArgs.cpu_offload_gb:
1600
- _raise_or_fallback(feature_name="--cpu-offload-gb",
1601
- recommend_to_remove=False)
1602
- return False
1603
-
1598
# Only Fp16 and Bf16 dtypes since we only support FA.
1599
V1_SUPPORTED_DTYPES = [torch.bfloat16, torch.float16]
1600
if model_config.dtype not in V1_SUPPORTED_DTYPES:
vllm/model_executor/models/utils.py
+18
-3
@@ -10,12 +10,14 @@ import torch.nn as nn
10
from torch.func import functional_call
11
from transformers import PretrainedConfig
12
13
+import vllm.envs as envs
14
from vllm.config import VllmConfig
15
from vllm.logger import init_logger
16
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
17
from vllm.multimodal import MultiModalPlaceholderMap, NestedTensors
18
from vllm.sequence import IntermediateTensors
18
-from vllm.utils import is_pin_memory_available
19
+from vllm.utils import (get_cuda_view_from_cpu_tensor, is_pin_memory_available,
20
+ is_uva_available)
21
22
logger = init_logger(__name__)
23
@@ -505,6 +507,14 @@ def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
507
return module
508
509
pin_memory = is_pin_memory_available()
510
+ uva_available = is_uva_available()
511
+
512
+ if envs.VLLM_USE_V1:
513
+ assert uva_available, ("V1 CPU offloading requires"
514
+ " uva (pin memory) support")
515
+ uva_offloading = True
516
+ else:
517
+ uva_offloading = False
518
519
# offload parameters to CPU
520
# use pin_memory if possible, which helps cudagraph capture speed
@@ -523,11 +533,16 @@ def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
533
device='cpu',
534
pin_memory=pin_memory)
535
cpu_data.copy_(p.data)
526
- p.data = cpu_data
536
+ if not uva_offloading:
537
+ p.data = cpu_data
538
+ else:
539
+ # keep the cpu data alive
540
+ p._vllm_offloaded_cpu_data = cpu_data
541
+ p.data = get_cuda_view_from_cpu_tensor(cpu_data)
542
_CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size()
543
offloaded_parameters = True
544
530
- if offloaded_parameters:
545
+ if offloaded_parameters and not uva_offloading:
546
original_forward = module.forward
547
548
def forward(*args, **kwargs):
vllm/utils.py
+16
@@ -795,6 +795,14 @@ def is_pin_memory_available() -> bool:
795
return current_platform.is_pin_memory_available()
796
797
798
+@cache
799
+def is_uva_available() -> bool:
800
+ """Check if Unified Virtual Addressing (UVA) is available."""
801
+ # UVA requires pinned memory.
802
+ # TODO: Add more requirements for UVA if needed.
803
+ return is_pin_memory_available()
804
+
805
+
806
class DeviceMemoryProfiler:
807
808
def __init__(self, device: Optional[torch.types.Device] = None):
@@ -1645,6 +1653,14 @@ def weak_ref_tensors(
1653
raise ValueError("Invalid type for tensors")
1654
1655
1656
+def get_cuda_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:
1657
+ """
1658
+ Get a CUDA view of a CPU tensor using Unified Virtual Addressing (UVA).
1659
+ """
1660
+ assert cpu_tensor.is_pinned(), "CPU tensor must be pinned"
1661
+ return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)
1662
+
1663
+
1664
def is_in_doc_build() -> bool:
1665
try:
1666
from sphinx.ext.autodoc.mock import _MockModule
vllm/v1/worker/gpu_model_runner.py
+4
@@ -69,6 +69,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
69
self.prompt_adapter_config = vllm_config.prompt_adapter_config
70
self.observability_config = vllm_config.observability_config
71
72
+ from vllm.model_executor.models.utils import set_cpu_offload_max_bytes
73
+ set_cpu_offload_max_bytes(
74
+ int(self.cache_config.cpu_offload_gb * 1024**3))
75
+
76
model_config = self.model_config
77
cache_config = self.cache_config
78
scheduler_config = self.scheduler_config