@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6831650c40
Fork 0

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 7, 17:45 UTC

[offloader] v2: Hide weight onloading latency via prefetching (#29941)

Signed-off-by: Ming Yang <minos.future@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

Ming Yang committed Feb 25, 2026 at 17:20 UTC 6831650c40ac3a34f049e285d9ad6b87daddbe00
20 files changed +1550 -131
.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh new
+57
@@ -0,0 +1,57 @@
1 +#!/usr/bin/env bash
2 +set -euxo pipefail
3 +
4 +# Nightly e2e test for prefetch offloading with a MoE model.
5 +# Runs DeepSeek-V2-Lite with prefetch offloading of MoE expert weights
6 +# and validates GSM8K accuracy matches baseline (no offloading).
7 +#
8 +# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
9 +THRESHOLD=${1:-0.25}
10 +NUM_Q=${2:-1319}
11 +PORT=${3:-8030}
12 +OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
13 +mkdir -p "${OUT_DIR}"
14 +
15 +wait_for_server() {
16 + local port=$1
17 + timeout 600 bash -c '
18 + until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
19 + sleep 1
20 + done'
21 +}
22 +
23 +MODEL="deepseek-ai/DeepSeek-V2-Lite"
24 +
25 +cleanup() {
26 + if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
27 + kill "${SERVER_PID}" 2>/dev/null || true
28 + for _ in {1..20}; do
29 + kill -0 "${SERVER_PID}" 2>/dev/null || break
30 + sleep 0.5
31 + done
32 + kill -9 "${SERVER_PID}" 2>/dev/null || true
33 + fi
34 +}
35 +trap cleanup EXIT
36 +
37 +vllm serve "$MODEL" \
38 + --max-model-len 2048 \
39 + --offload-group-size 8 \
40 + --offload-num-in-group 2 \
41 + --offload-prefetch-step 1 \
42 + --offload-params w13_weight w2_weight \
43 + --port "$PORT" &
44 +SERVER_PID=$!
45 +wait_for_server "$PORT"
46 +
47 +TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
48 +OUT="${OUT_DIR}/${TAG}_prefetch_offload.json"
49 +python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
50 +python3 - <<PY
51 +import json; acc=json.load(open('${OUT}'))['accuracy']
52 +print(f"${MODEL} prefetch_offload: accuracy {acc:.3f}")
53 +assert acc >= ${THRESHOLD}, f"${MODEL} prefetch_offload accuracy {acc}"
54 +PY
55 +
56 +cleanup
57 +SERVER_PID=
.buildkite/test_areas/e2e_integration.yaml
+9
@@ -28,3 +28,12 @@ steps:
28 working_dir: "/vllm-workspace"
29 commands:
30 - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
31 +
32 +- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
33 + timeout_in_minutes: 60
34 + device: h100
35 + optional: true
36 + num_devices: 1
37 + working_dir: "/vllm-workspace"
38 + commands:
39 + - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030
tests/basic_correctness/test_prefetch_offload.py new
+33
@@ -0,0 +1,33 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Test prefetch offloading correctness with Llama model."""
4 +
5 +from ..utils import compare_two_settings
6 +
7 +
8 +def test_prefetch_offload_llama():
9 + """Test prefetch CPU offloading with Llama-3.2-1B-Instruct.
10 +
11 + Compares outputs between:
12 + 1. Baseline (no offloading)
13 + 2. Prefetch offloading (group_size=8, num_in_group=2, prefetch_step=1)
14 +
15 + This tests prefetching-based offloading on a dense model.
16 + """
17 + compare_two_settings(
18 + "meta-llama/Llama-3.2-1B-Instruct",
19 + [
20 + # Prefetch offloading configuration
21 + "--offload-group-size",
22 + "8",
23 + "--offload-num-in-group",
24 + "2",
25 + "--offload-prefetch-step",
26 + "1",
27 + # Selective offloading: only MLP weights
28 + "--offload-params",
29 + "gate_up_proj",
30 + "down_proj",
31 + ],
32 + [], # Baseline: no offloading
33 + )
vllm/compilation/cuda_graph.py
+14
@@ -17,6 +17,7 @@ from vllm.config import CUDAGraphMode, VllmConfig
17 from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
18 from vllm.forward_context import BatchDescriptor, get_forward_context
19 from vllm.logger import init_logger
20 +from vllm.model_executor.offloader.base import get_offloader
21 from vllm.platforms import current_platform
22 from vllm.utils.torch_utils import current_stream, weak_ref_tensors
23
@@ -265,6 +266,11 @@ class CUDAGraphWrapper:
266 set_graph_pool_id(self.graph_pool)
267 else:
268 set_graph_pool_id(current_platform.graph_pool_handle())
269 +
270 + # Sync offloader's copy stream before capture.
271 + # Ensure any pre-capture prefetches from offloader are complete.
272 + get_offloader().sync_prev_onload()
273 +
274 # mind-exploding: carefully manage the reference and memory.
275 with torch.cuda.graph(
276 cudagraph,
@@ -273,6 +279,11 @@ class CUDAGraphWrapper:
279 ):
280 # `output` is managed by pytorch's cudagraph pool
281 output = self.runnable(*args, **kwargs)
282 + # Join offloader's copy stream after forward to avoid
283 + # unjoined stream error. The last layer's start_prefetch
284 + # forks copy_stream, but wait_prefetch only happens in
285 + # the next forward pass.
286 + get_offloader().join_after_forward()
287 if self.cudagraph_options.weak_ref_output:
288 # by converting it to weak ref,
289 # the original `output` will immediately be released
@@ -305,5 +316,8 @@ class CUDAGraphWrapper:
316 f"got {new_input_addresses}"
317 )
318
319 + # Sync offloader before replay - ensures any external dependencies
320 + # from pre-capture prefetches are satisfied.
321 + get_offloader().sync_prev_onload()
322 entry.cudagraph.replay()
323 return entry.output
vllm/config/__init__.py
+11
@@ -24,6 +24,12 @@ from vllm.config.model import (
24 )
25 from vllm.config.multimodal import MultiModalConfig
26 from vllm.config.observability import ObservabilityConfig
27 +from vllm.config.offload import (
28 + OffloadBackend,
29 + OffloadConfig,
30 + PrefetchOffloadConfig,
31 + UVAOffloadConfig,
32 +)
33 from vllm.config.parallel import EPLBConfig, ParallelConfig
34 from vllm.config.pooler import PoolerConfig
35 from vllm.config.profiler import ProfilerConfig
@@ -85,6 +91,11 @@ __all__ = [
91 "MultiModalConfig",
92 # From vllm.config.observability
93 "ObservabilityConfig",
94 + # From vllm.config.offload
95 + "OffloadBackend",
96 + "OffloadConfig",
97 + "PrefetchOffloadConfig",
98 + "UVAOffloadConfig",
99 # From vllm.config.parallel
100 "EPLBConfig",
101 "ParallelConfig",
vllm/config/cache.py
+7 -9
@@ -100,17 +100,15 @@ class CacheConfig:
100 load a 13B model with BF16 weight, which requires at least 26GB GPU memory.
101 Note that this requires fast CPU-GPU interconnect, as part of the model is
102 loaded from CPU memory to GPU memory on the fly in each model forward pass.
103 +
104 + DEPRECATED: This field is deprecated and will be removed in v0.16.
105 + Please use OffloadConfig.uva.cpu_offload_gb instead.
106 """
107 cpu_offload_params: set[str] = Field(default_factory=set)
105 - """ The set of parameter name segments to target for CPU offloading.
106 - Unmatched parameters are not offloaded. If this set is empty, parameters
107 - are offloaded non-selectively until the memory limit defined by
108 - `cpu_offload_gb` is reached.
109 - Examples:
110 - - For parameter name "mlp.experts.w2_weight":
111 - - "experts" or "experts.w2_weight" will match.
112 - - "expert" or "w2" will NOT match (must be exact segments).
113 - This allows distinguishing parameters like "w2_weight" and "w2_weight_scale".
108 + """The set of parameter name segments to target for CPU offloading.
109 +
110 + DEPRECATED: This field is deprecated and will be removed in v0.16.
111 + Please use OffloadConfig.uva.cpu_offload_params instead.
112 """
113 calculate_kv_scales: bool = False
114 """This enables dynamic calculation of `k_scale` and `v_scale` when
vllm/config/offload.py new
+153
@@ -0,0 +1,153 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Configuration for model weight offloading."""
4 +
5 +import warnings
6 +from typing import Literal
7 +
8 +from pydantic import Field, model_validator
9 +
10 +from vllm.config.utils import config
11 +
12 +OffloadBackend = Literal["auto", "uva", "prefetch"]
13 +
14 +
15 +@config
16 +class UVAOffloadConfig:
17 + """Configuration for UVA (Unified Virtual Addressing) CPU offloading.
18 +
19 + Uses zero-copy access from CPU-pinned memory. Simple but requires
20 + fast CPU-GPU interconnect.
21 + """
22 +
23 + cpu_offload_gb: float = Field(default=0, ge=0)
24 + """The space in GiB to offload to CPU, per GPU. Default is 0, which means
25 + no offloading. Intuitively, this argument can be seen as a virtual way to
26 + increase the GPU memory size. For example, if you have one 24 GB GPU and
27 + set this to 10, virtually you can think of it as a 34 GB GPU. Then you can
28 + load a 13B model with BF16 weight, which requires at least 26GB GPU memory.
29 + Note that this requires fast CPU-GPU interconnect, as part of the model is
30 + loaded from CPU memory to GPU memory on the fly in each model forward pass.
31 + This uses UVA (Unified Virtual Addressing) for zero-copy access.
32 + """
33 +
34 + cpu_offload_params: set[str] = Field(default_factory=set)
35 + """The set of parameter name segments to target for CPU offloading.
36 + Unmatched parameters are not offloaded. If this set is empty, parameters
37 + are offloaded non-selectively until the memory limit defined by
38 + `cpu_offload_gb` is reached.
39 + Examples:
40 + - For parameter name "mlp.experts.w2_weight":
41 + - "experts" or "experts.w2_weight" will match.
42 + - "expert" or "w2" will NOT match (must be exact segments).
43 + This allows distinguishing parameters like "w2_weight" and "w2_weight_scale".
44 + """
45 +
46 +
47 +@config
48 +class PrefetchOffloadConfig:
49 + """Configuration for prefetch-based CPU offloading.
50 +
51 + Groups layers and uses async H2D prefetch to hide transfer latency.
52 + """
53 +
54 + offload_group_size: int = Field(default=0, ge=0)
55 + """Group every N layers together. Offload last `offload_num_in_group`
56 + layers of each group. Default is 0 (disabled).
57 + Example: group_size=8, num_in_group=2 offloads layers 6,7,14,15,22,23,...
58 + Unlike cpu_offload_gb, this uses explicit async prefetching to hide transfer
59 + latency.
60 + """
61 +
62 + offload_num_in_group: int = Field(default=1, ge=1)
63 + """Number of layers to offload per group.
64 + Must be <= offload_group_size. Default is 1."""
65 +
66 + offload_prefetch_step: int = Field(default=1, ge=0)
67 + """Number of layers to prefetch ahead.
68 + Higher values hide more latency but use more GPU memory. Default is 1."""
69 +
70 + offload_params: set[str] = Field(default_factory=set)
71 + """The set of parameter name segments to target for prefetch offloading.
72 + Unmatched parameters are not offloaded. If this set is empty, ALL
73 + parameters of each offloaded layer are offloaded.
74 + Uses segment matching: "w13_weight" matches "mlp.experts.w13_weight"
75 + but not "mlp.experts.w13_weight_scale".
76 + """
77 +
78 +
79 +@config
80 +class OffloadConfig:
81 + """Configuration for model weight offloading to reduce GPU memory usage."""
82 +
83 + offload_backend: OffloadBackend = "auto"
84 + """The backend for weight offloading. Options:
85 + - "auto": Selects based on which sub-config has non-default values
86 + (prefetch if offload_group_size > 0, uva if cpu_offload_gb > 0).
87 + - "uva": UVA (Unified Virtual Addressing) zero-copy offloading.
88 + - "prefetch": Async prefetch with group-based layer offloading.
89 + """
90 +
91 + uva: UVAOffloadConfig = Field(default_factory=UVAOffloadConfig)
92 + """Parameters for UVA offloading backend."""
93 +
94 + prefetch: PrefetchOffloadConfig = Field(default_factory=PrefetchOffloadConfig)
95 + """Parameters for prefetch offloading backend."""
96 +
97 + @model_validator(mode="after")
98 + def validate_offload_config(self) -> "OffloadConfig":
99 + """Validate offload configuration constraints."""
100 + if self.offload_backend == "prefetch" or self.prefetch.offload_group_size > 0:
101 + if self.prefetch.offload_num_in_group > self.prefetch.offload_group_size:
102 + raise ValueError(
103 + f"offload_num_in_group ({self.prefetch.offload_num_in_group})"
104 + f" must be <= offload_group_size"
105 + f" ({self.prefetch.offload_group_size})"
106 + )
107 + if self.prefetch.offload_prefetch_step < 1:
108 + raise ValueError(
109 + f"offload_prefetch_step"
110 + f" ({self.prefetch.offload_prefetch_step})"
111 + f" must be >= 1 when prefetch offloading is enabled"
112 + f" (offload_group_size > 0)"
113 + )
114 +
115 + # Warn if both backends have non-default values
116 + uva_active = self.uva.cpu_offload_gb > 0
117 + prefetch_active = self.prefetch.offload_group_size > 0
118 + if self.offload_backend == "uva" and prefetch_active:
119 + warnings.warn(
120 + "Prefetch offload fields are set but offload_backend='uva'. "
121 + "Prefetch settings will be ignored.",
122 + stacklevel=2,
123 + )
124 + elif self.offload_backend == "prefetch" and uva_active:
125 + warnings.warn(
126 + "UVA offload fields are set but offload_backend='prefetch'. "
127 + "UVA settings will be ignored.",
128 + stacklevel=2,
129 + )
130 + elif self.offload_backend == "auto" and uva_active and prefetch_active:
131 + warnings.warn(
132 + "Both UVA and prefetch offload fields are set with "
133 + "offload_backend='auto'. Prefetch backend will be selected. "
134 + "Set offload_backend explicitly to suppress this warning.",
135 + stacklevel=2,
136 + )
137 + return self
138 +
139 + def compute_hash(self) -> str:
140 + """
141 + Provide a hash that uniquely identifies all the offload configs.
142 +
143 + All fields are included because PrefetchOffloader patches module
144 + forwards and inserts custom ops (wait_prefetch, start_prefetch)
145 + into the computation graph. Changing any offload setting can
146 + alter which layers are hooked and how prefetch indices are
147 + computed, so the compilation cache must distinguish them.
148 + """
149 + from vllm.config.utils import get_hash_factors, hash_factors
150 +
151 + factors = get_hash_factors(self, ignored_factors=set())
152 + hash_str = hash_factors(factors)
153 + return hash_str
vllm/config/vllm.py
+7
@@ -37,6 +37,7 @@ from .load import LoadConfig
37 from .lora import LoRAConfig
38 from .model import ModelConfig
39 from .observability import ObservabilityConfig
40 +from .offload import OffloadConfig
41 from .parallel import ParallelConfig
42 from .profiler import ProfilerConfig
43 from .scheduler import SchedulerConfig
@@ -259,6 +260,8 @@ class VllmConfig:
260 """Device configuration."""
261 load_config: LoadConfig = Field(default_factory=LoadConfig)
262 """Load configuration."""
263 + offload_config: OffloadConfig = Field(default_factory=OffloadConfig)
264 + """Model weight offloading configuration."""
265 attention_config: AttentionConfig = Field(default_factory=AttentionConfig)
266 """Attention configuration."""
267 kernel_config: KernelConfig = Field(default_factory=KernelConfig)
@@ -361,6 +364,10 @@ class VllmConfig:
364 vllm_factors.append(self.load_config.compute_hash())
365 else:
366 vllm_factors.append("None")
367 + if self.offload_config:
368 + vllm_factors.append(self.offload_config.compute_hash())
369 + else:
370 + vllm_factors.append("None")
371 if self.attention_config:
372 vllm_factors.append(self.attention_config.compute_hash())
373 else:
vllm/engine/arg_utils.py
+57 -8
@@ -48,12 +48,15 @@ from vllm.config import (
48 ModelConfig,
49 MultiModalConfig,
50 ObservabilityConfig,
51 + OffloadConfig,
52 ParallelConfig,
53 PoolerConfig,
54 + PrefetchOffloadConfig,
55 ProfilerConfig,
56 SchedulerConfig,
57 SpeculativeConfig,
58 StructuredOutputsConfig,
59 + UVAOffloadConfig,
60 VllmConfig,
61 WeightTransferConfig,
62 get_attr_docs,
@@ -439,8 +442,13 @@ class EngineArgs:
442 disable_sliding_window: bool = ModelConfig.disable_sliding_window
443 disable_cascade_attn: bool = ModelConfig.disable_cascade_attn
444 swap_space: float = CacheConfig.swap_space
442 - cpu_offload_gb: float = CacheConfig.cpu_offload_gb
443 - cpu_offload_params: set[str] = get_field(CacheConfig, "cpu_offload_params")
445 + offload_backend: str = OffloadConfig.offload_backend
446 + cpu_offload_gb: float = UVAOffloadConfig.cpu_offload_gb
447 + cpu_offload_params: set[str] = get_field(UVAOffloadConfig, "cpu_offload_params")
448 + offload_group_size: int = PrefetchOffloadConfig.offload_group_size
449 + offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group
450 + offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step
451 + offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params")
452 gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization
453 kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
454 max_num_batched_tokens: int | None = None
@@ -948,10 +956,6 @@ class EngineArgs:
956 cache_group.add_argument(
957 "--prefix-caching-hash-algo", **cache_kwargs["prefix_caching_hash_algo"]
958 )
951 - cache_group.add_argument("--cpu-offload-gb", **cache_kwargs["cpu_offload_gb"])
952 - cache_group.add_argument(
953 - "--cpu-offload-params", **cache_kwargs["cpu_offload_params"]
954 - )
959 cache_group.add_argument(
960 "--calculate-kv-scales", **cache_kwargs["calculate_kv_scales"]
961 )
@@ -977,6 +981,37 @@ class EngineArgs:
981 "--kv-offloading-backend", **cache_kwargs["kv_offloading_backend"]
982 )
983
984 + # Model weight offload related configs
985 + offload_kwargs = get_kwargs(OffloadConfig)
986 + uva_kwargs = get_kwargs(UVAOffloadConfig)
987 + prefetch_kwargs = get_kwargs(PrefetchOffloadConfig)
988 + offload_group = parser.add_argument_group(
989 + title="OffloadConfig",
990 + description=OffloadConfig.__doc__,
991 + )
992 + offload_group.add_argument(
993 + "--offload-backend", **offload_kwargs["offload_backend"]
994 + )
995 + offload_group.add_argument("--cpu-offload-gb", **uva_kwargs["cpu_offload_gb"])
996 + offload_group.add_argument(
997 + "--cpu-offload-params", **uva_kwargs["cpu_offload_params"]
998 + )
999 + offload_group.add_argument(
1000 + "--offload-group-size",
1001 + **prefetch_kwargs["offload_group_size"],
1002 + )
1003 + offload_group.add_argument(
1004 + "--offload-num-in-group",
1005 + **prefetch_kwargs["offload_num_in_group"],
1006 + )
1007 + offload_group.add_argument(
1008 + "--offload-prefetch-step",
1009 + **prefetch_kwargs["offload_prefetch_step"],
1010 + )
1011 + offload_group.add_argument(
1012 + "--offload-params", **prefetch_kwargs["offload_params"]
1013 + )
1014 +
1015 # Multimodal related configs
1016 multimodal_kwargs = get_kwargs(MultiModalConfig)
1017 multimodal_group = parser.add_argument_group(
@@ -1466,8 +1501,6 @@ class EngineArgs:
1501 sliding_window=sliding_window,
1502 enable_prefix_caching=self.enable_prefix_caching,
1503 prefix_caching_hash_algo=self.prefix_caching_hash_algo,
1469 - cpu_offload_gb=self.cpu_offload_gb,
1470 - cpu_offload_params=self.cpu_offload_params,
1504 calculate_kv_scales=self.calculate_kv_scales,
1505 kv_sharing_fast_prefill=self.kv_sharing_fast_prefill,
1506 mamba_cache_dtype=self.mamba_cache_dtype,
@@ -1825,6 +1858,21 @@ class EngineArgs:
1858 compilation_config.max_cudagraph_capture_size = (
1859 self.max_cudagraph_capture_size
1860 )
1861 +
1862 + offload_config = OffloadConfig(
1863 + offload_backend=self.offload_backend,
1864 + uva=UVAOffloadConfig(
1865 + cpu_offload_gb=self.cpu_offload_gb,
1866 + cpu_offload_params=self.cpu_offload_params,
1867 + ),
1868 + prefetch=PrefetchOffloadConfig(
1869 + offload_group_size=self.offload_group_size,
1870 + offload_num_in_group=self.offload_num_in_group,
1871 + offload_prefetch_step=self.offload_prefetch_step,
1872 + offload_params=self.offload_params,
1873 + ),
1874 + )
1875 +
1876 config = VllmConfig(
1877 model_config=model_config,
1878 cache_config=cache_config,
@@ -1832,6 +1880,7 @@ class EngineArgs:
1880 scheduler_config=scheduler_config,
1881 device_config=device_config,
1882 load_config=load_config,
1883 + offload_config=offload_config,
1884 attention_config=attention_config,
1885 kernel_config=kernel_config,
1886 lora_config=lora_config,
vllm/entrypoints/llm.py
+21
@@ -170,6 +170,19 @@ class LLM:
170 the model weights. This virtually increases the GPU memory space
171 you can use to hold the model weights, at the cost of CPU-GPU data
172 transfer for every forward pass.
173 + offload_group_size: Prefetch offloading: Group every N layers
174 + together. Offload last `offload_num_in_group` layers of each group.
175 + Default is 0 (disabled).
176 + offload_num_in_group: Prefetch offloading: Number of layers to
177 + offload per group. Default is 1.
178 + offload_prefetch_step: Prefetch offloading: Number of layers to
179 + prefetch ahead. Higher values hide more latency but use more GPU
180 + memory. Default is 1.
181 + offload_params: Prefetch offloading: Set of parameter name segments
182 + to selectively offload. Only parameters whose names contain one of
183 + these segments will be offloaded (e.g., {"gate_up_proj", "down_proj"}
184 + for MLP weights, or {"w13_weight", "w2_weight"} for MoE expert
185 + weights). If None or empty, all parameters are offloaded.
186 enforce_eager: Whether to enforce eager execution. If True, we will
187 disable CUDA graph and always execute the model in eager mode.
188 If False, we will use CUDA graph and eager execution in hybrid.
@@ -224,6 +237,10 @@ class LLM:
237 gpu_memory_utilization: float = 0.9,
238 swap_space: float = 4,
239 cpu_offload_gb: float = 0,
240 + offload_group_size: int = 0,
241 + offload_num_in_group: int = 1,
242 + offload_prefetch_step: int = 1,
243 + offload_params: set[str] | None = None,
244 enforce_eager: bool = False,
245 enable_return_routed_experts: bool = False,
246 disable_custom_all_reduce: bool = False,
@@ -333,6 +350,10 @@ class LLM:
350 kv_cache_memory_bytes=kv_cache_memory_bytes,
351 swap_space=swap_space,
352 cpu_offload_gb=cpu_offload_gb,
353 + offload_group_size=offload_group_size,
354 + offload_num_in_group=offload_num_in_group,
355 + offload_prefetch_step=offload_prefetch_step,
356 + offload_params=offload_params or set(),
357 enforce_eager=enforce_eager,
358 enable_return_routed_experts=enable_return_routed_experts,
359 disable_custom_all_reduce=disable_custom_all_reduce,
vllm/model_executor/models/utils.py
+14 -105
@@ -9,11 +9,9 @@ from typing import Any, Literal, Protocol, overload
9
10 import torch
11 import torch.nn as nn
12 -from torch.func import functional_call
12 from torch.nn.modules.module import register_module_module_registration_hook
13 from transformers import PretrainedConfig
14
16 -import vllm.envs as envs
15 from vllm.config import VllmConfig
16 from vllm.distributed import (
17 get_tensor_model_parallel_rank,
@@ -31,14 +29,11 @@ from vllm.model_executor.models.interfaces import supports_any_eagle
29 from vllm.multimodal import NestedTensors
30 from vllm.sequence import IntermediateTensors
31 from vllm.utils.math_utils import cdiv
34 -from vllm.utils.mem_utils import format_gib
32 from vllm.utils.platform_utils import (
33 is_pin_memory_available,
37 - is_uva_available,
34 )
35 from vllm.utils.torch_utils import (
36 direct_register_custom_op,
41 - get_accelerator_view_from_cpu_tensor,
37 )
38
39 logger = init_logger(__name__)
@@ -612,98 +607,6 @@ class PPMissingLayer(torch.nn.Identity):
607 return args[0] if args else next(iter(kwargs.values()))
608
609
615 -_CPU_OFFLOAD_BYTES = 0
616 -_CPU_OFFLOAD_MAX_BYTES = 0
617 -_CPU_OFFLOAD_PARAMS = set()
618 -
619 -
620 -def set_cpu_offload_max_bytes(max_bytes: int) -> None:
621 - global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES
622 - _CPU_OFFLOAD_BYTES = 0
623 - _CPU_OFFLOAD_MAX_BYTES = max_bytes
624 -
625 -
626 -def set_cpu_offload_params(params: set[str]) -> None:
627 - global _CPU_OFFLOAD_PARAMS
628 - _CPU_OFFLOAD_PARAMS = params
629 -
630 -
631 -def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
632 - if (params := next(module.parameters(), None)) is None:
633 - return module
634 -
635 - device = params.device
636 -
637 - if device == torch.device("cpu"):
638 - return module
639 -
640 - global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES
641 - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES:
642 - return module
643 -
644 - pin_memory = (
645 - is_pin_memory_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY
646 - )
647 - uva_offloading = is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA
648 -
649 - # offload parameters to CPU
650 - # use pin_memory if possible, which helps cudagraph capture speed
651 - offloaded_parameters = False
652 - for name, p in module.named_parameters():
653 - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES:
654 - # we use per-parameter offloading
655 - # one module might have some parameters offloaded and some not
656 - break
657 -
658 - if _CPU_OFFLOAD_PARAMS:
659 - # Check if parameter belongs to the offloading set
660 - # Add dots here to ensure we match full segments only
661 - # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight" but not
662 - # "mlp.experts.w2_weight_scale"
663 - should_offload = any(
664 - f".{param}." in f".{name}." for param in _CPU_OFFLOAD_PARAMS
665 - )
666 - if not should_offload:
667 - continue
668 -
669 - cpu_data = p.data.to(device="cpu")
670 - if pin_memory:
671 - cpu_data = cpu_data.pin_memory()
672 -
673 - if not uva_offloading:
674 - p.data = cpu_data
675 - else:
676 - p.data = get_accelerator_view_from_cpu_tensor(cpu_data)
677 - p._vllm_is_uva_offloaded = True
678 -
679 - _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size()
680 - offloaded_parameters = True
681 -
682 - if offloaded_parameters and not uva_offloading:
683 - original_forward = module.forward
684 -
685 - def forward(*args, **kwargs):
686 - module.forward = original_forward
687 - device_state = {
688 - # here we blindly call `to(device)`
689 - # if the parameter is already on the device, it will be a no-op
690 - k: v.to(device, non_blocking=True)
691 - for k, v in module.state_dict().items()
692 - }
693 -
694 - # set `tie_weights=False` as tied weights in original model
695 - # become untied when calling .to(device) individually
696 - output = functional_call(
697 - module, device_state, args=args, kwargs=kwargs, tie_weights=False
698 - )
699 - module.forward = forward
700 - return output
701 -
702 - module.forward = forward
703 -
704 - return module
705 -
706 -
610 def make_layers(
611 num_hidden_layers: int,
612 layer_fn: LayerFn,
@@ -711,25 +614,31 @@ def make_layers(
614 ) -> tuple[int, int, torch.nn.ModuleList]:
615 """Make a list of layers with the given layer function, taking
616 pipeline parallelism into account.
617 +
618 + Args:
619 + num_hidden_layers: Total number of hidden layers in the model.
620 + layer_fn: Function to create a layer given its index.
621 + prefix: Prefix for layer names.
622 +
623 + Returns:
624 + Tuple of (start_layer, end_layer, modules).
625 """
626 from vllm.distributed.parallel_state import get_pp_group
627 from vllm.distributed.utils import get_pp_indices
628 + from vllm.model_executor.offloader import get_offloader
629
630 start_layer, end_layer = get_pp_indices(
631 num_hidden_layers, get_pp_group().rank_in_group, get_pp_group().world_size
632 )
633 +
634 modules = torch.nn.ModuleList(
635 [PPMissingLayer() for _ in range(start_layer)]
723 - + [
724 - maybe_offload_to_cpu(layer_fn(prefix=f"{prefix}.{idx}"))
725 - for idx in range(start_layer, end_layer)
726 - ]
636 + + get_offloader().wrap_modules(
637 + layer_fn(prefix=f"{prefix}.{idx}") for idx in range(start_layer, end_layer)
638 + )
639 + [PPMissingLayer() for _ in range(end_layer, num_hidden_layers)]
640 )
729 - if _CPU_OFFLOAD_MAX_BYTES > 0:
730 - logger.info(
731 - "Total CPU offloaded parameters: %s GBs", format_gib(_CPU_OFFLOAD_BYTES)
732 - )
641 +
642 return start_layer, end_layer, modules
643
644
vllm/model_executor/offloader/__init__.py new
+23
@@ -0,0 +1,23 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Model parameter offloading infrastructure."""
4 +
5 +from vllm.model_executor.offloader.base import (
6 + BaseOffloader,
7 + NoopOffloader,
8 + create_offloader,
9 + get_offloader,
10 + set_offloader,
11 +)
12 +from vllm.model_executor.offloader.prefetch import PrefetchOffloader
13 +from vllm.model_executor.offloader.uva import UVAOffloader
14 +
15 +__all__ = [
16 + "BaseOffloader",
17 + "NoopOffloader",
18 + "UVAOffloader",
19 + "PrefetchOffloader",
20 + "create_offloader",
21 + "get_offloader",
22 + "set_offloader",
23 +]
vllm/model_executor/offloader/base.py new
+145
@@ -0,0 +1,145 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Adapted from
4 +# https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py
5 +"""Base classes for model parameter offloading."""
6 +
7 +from abc import ABC, abstractmethod
8 +from collections.abc import Generator
9 +from typing import TYPE_CHECKING
10 +
11 +import torch.nn as nn
12 +
13 +from vllm.logger import init_logger
14 +
15 +if TYPE_CHECKING:
16 + from vllm.config import OffloadConfig
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +"""
22 +class relation:
23 +
24 +BaseOffloader (ABC)
25 + * implemented by: UVAOffloader
26 + * implemented by: PrefetchOffloader
27 + * uses: _ModuleOffloader
28 + * uses: _BaseParamOffloader (ABC)
29 + * implemented by: _CpuParamOffloader
30 +"""
31 +
32 +
33 +class BaseOffloader(ABC):
34 + """Base class for model parameter offloading strategies.
35 +
36 + Offloaders control how model parameters are stored and loaded during
37 + inference. Different strategies trade memory for compute/transfer time.
38 + """
39 +
40 + @abstractmethod
41 + def wrap_modules(
42 + self,
43 + modules_generator: Generator[nn.Module, None, None],
44 + ) -> list[nn.Module]:
45 + """Wrap modules with offloading logic.
46 +
47 + Args:
48 + modules_generator: Generator yielding modules to potentially offload.
49 +
50 + Returns:
51 + List of modules, potentially with offloading hooks installed.
52 + """
53 + pass
54 +
55 + def post_init(self):
56 + """Called after model construction completes.
57 +
58 + Offloaders can use this to:
59 + - Finalize parameter storage
60 + - Start initial prefetching
61 + - Allocate shared resources
62 + """
63 + return
64 +
65 + def sync_prev_onload(self) -> None: # noqa: B027
66 + """Sync previous onload operations. Override in subclasses."""
67 + pass
68 +
69 + def join_after_forward(self) -> None: # noqa: B027
70 + """Join streams after forward. Override in subclasses."""
71 + pass
72 +
73 + def _wait_for_layer(self, layer_idx: int) -> None: # noqa: B027
74 + """Wait for layer prefetch. Override in subclasses."""
75 + pass
76 +
77 + def _start_prefetch(self, layer_idx: int) -> None: # noqa: B027
78 + """Start layer prefetch. Override in subclasses."""
79 + pass
80 +
81 +
82 +class NoopOffloader(BaseOffloader):
83 + """No-op offloader that returns modules as-is without any offloading."""
84 +
85 + def wrap_modules(
86 + self,
87 + modules_generator: Generator[nn.Module, None, None],
88 + ) -> list[nn.Module]:
89 + """Return modules unchanged."""
90 + return list(modules_generator)
91 +
92 +
93 +# Global singleton offloader instance (defaults to no-op).
94 +_instance: BaseOffloader = NoopOffloader()
95 +
96 +
97 +def get_offloader() -> BaseOffloader:
98 + """Get the global offloader instance."""
99 + return _instance
100 +
101 +
102 +def set_offloader(instance: BaseOffloader) -> None:
103 + """Set the global offloader instance."""
104 + global _instance
105 + _instance = instance
106 + logger.info("Offloader set to %s", type(instance).__name__)
107 +
108 +
109 +def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader:
110 + """Create an offloader based on the offload configuration.
111 +
112 + Uses the explicit ``offload_backend`` selector. When set to ``"auto"``,
113 + selects prefetch if ``offload_group_size > 0``, UVA if
114 + ``cpu_offload_gb > 0``, otherwise noop.
115 + """
116 + from vllm.model_executor.offloader.prefetch import PrefetchOffloader
117 + from vllm.model_executor.offloader.uva import UVAOffloader
118 +
119 + backend = offload_config.offload_backend
120 + uva = offload_config.uva
121 + prefetch = offload_config.prefetch
122 +
123 + if backend == "auto":
124 + if prefetch.offload_group_size > 0:
125 + backend = "prefetch"
126 + elif uva.cpu_offload_gb > 0:
127 + backend = "uva"
128 + else:
129 + return NoopOffloader()
130 +
131 + if backend == "prefetch":
132 + return PrefetchOffloader(
133 + group_size=prefetch.offload_group_size,
134 + num_in_group=prefetch.offload_num_in_group,
135 + prefetch_step=prefetch.offload_prefetch_step,
136 + offload_params=prefetch.offload_params,
137 + mode="cpu",
138 + )
139 + elif backend == "uva":
140 + return UVAOffloader(
141 + cpu_offload_max_bytes=int(uva.cpu_offload_gb * 1024**3),
142 + cpu_offload_params=uva.cpu_offload_params,
143 + )
144 + else:
145 + return NoopOffloader()
vllm/model_executor/offloader/prefetch.py new
+704
@@ -0,0 +1,704 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Adapted from
4 +# https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py
5 +"""Prefetch-based CPU offloading with async prefetching.
6 +
7 +Uses static buffers and event-based stream forking for torch.compile +
8 +CUDA graph compatibility. Events allow the copy stream to join CUDA
9 +graph captures, ensuring H2D copies are properly captured.
10 +"""
11 +
12 +from abc import ABC, abstractmethod
13 +from collections.abc import Generator
14 +from dataclasses import dataclass
15 +from typing import Any
16 +
17 +import torch
18 +import torch.nn as nn
19 +
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
25 +
26 +logger = init_logger(__name__)
27 +
28 +
29 +@dataclass
30 +class ParamInfo:
31 + """Metadata about an offloaded parameter."""
32 +
33 + name: str
34 + shape: tuple[int, ...]
35 + stride: tuple[int, ...]
36 + dtype: torch.dtype
37 +
38 + @property
39 + def key(self) -> tuple[str, tuple[int, ...], tuple[int, ...], torch.dtype]:
40 + """Unique key for buffer pool grouping.
41 +
42 + Includes parameter name to prevent different parameters with the same
43 + shape from sharing buffers within the same layer. Parameters with the
44 + same name across different layers will share buffers (via slots).
45 +
46 + Includes stride because parameters with same shape but different
47 + strides need separate buffers to preserve memory layout.
48 + """
49 + return (self.name, self.shape, self.stride, self.dtype)
50 +
51 + @property
52 + def num_bytes(self) -> int:
53 + """Size in bytes."""
54 + numel = 1
55 + for dim in self.shape:
56 + numel *= dim
57 + return numel * torch.finfo(self.dtype).bits // 8
58 +
59 +
60 +class StaticBufferPool:
61 + """Pre-allocated GPU buffer pool for offloaded parameters.
62 +
63 + Allocates slot_capacity copies of each unique parameter
64 + (name, shape, stride, dtype), allowing for double/triple buffering
65 + during prefetch.
66 +
67 + Buffer slots are reused circularly: layer N uses slot (N % slot_capacity).
68 +
69 + The key includes parameter name to prevent different parameters within
70 + the same layer from sharing buffers. Parameters with the same name
71 + across different layers share buffers via the slot mechanism.
72 + """
73 +
74 + def __init__(
75 + self,
76 + param_infos: list[ParamInfo],
77 + slot_capacity: int,
78 + device: torch.device,
79 + ):
80 + self.slot_capacity = slot_capacity
81 + self.total_bytes = 0
82 + self._device = device
83 +
84 + # Group by (shape, stride, dtype) - only allocate unique combinations
85 + unique_params: dict[tuple, ParamInfo] = {}
86 + for info in param_infos:
87 + if info.key not in unique_params:
88 + unique_params[info.key] = info
89 +
90 + # Allocate buffers: key -> list of tensors (one per slot)
91 + self._buffers: dict[tuple, list[torch.Tensor]] = {}
92 + for key, info in unique_params.items():
93 + slot_tensors = []
94 + for _ in range(slot_capacity):
95 + # Use empty_strided to preserve parameter's memory layout
96 + buf = torch.empty_strided(
97 + size=info.shape,
98 + stride=info.stride,
99 + dtype=info.dtype,
100 + device=device,
101 + )
102 + slot_tensors.append(buf)
103 + self.total_bytes += info.num_bytes
104 + self._buffers[key] = slot_tensors
105 +
106 + logger.debug(
107 + "[StaticBufferPool] Allocated %d unique (name, shape, stride, dtype), "
108 + "%d slots each, total %.4f GB",
109 + len(unique_params),
110 + slot_capacity,
111 + self.total_bytes / 1e9,
112 + )
113 +
114 + def get_buffer(
115 + self,
116 + name: str,
117 + shape: tuple[int, ...],
118 + stride: tuple[int, ...],
119 + dtype: torch.dtype,
120 + slot_idx: int,
121 + ) -> torch.Tensor:
122 + """Get a static buffer for the given name/shape/stride/dtype/slot."""
123 + key = (name, shape, stride, dtype)
124 + return self._buffers[key][slot_idx % self.slot_capacity]
125 +
126 +
127 +class PrefetchOffloader(BaseOffloader):
128 + """Prefetching-based offloader with group-based layer selection.
129 +
130 + Groups layers and uses async H2D prefetch to hide transfer latency.
131 + Uses static buffers and stream synchronization for torch.compile and
132 + CUDA graph compatibility.
133 +
134 + Args:
135 + group_size: Group every N layers together.
136 + num_in_group: Offload this many layers per group (last N of each group).
137 + prefetch_step: Number of layers to prefetch ahead.
138 + mode: Offload mode ("cpu" is currently supported).
139 + """
140 +
141 + def __init__(
142 + self,
143 + group_size: int,
144 + num_in_group: int,
145 + prefetch_step: int,
146 + offload_params: set[str] | None = None,
147 + mode: str = "cpu",
148 + ):
149 + self.group_size = group_size
150 + self.num_in_group = num_in_group
151 + self.prefetch_step = prefetch_step
152 + self.offload_params = offload_params or set()
153 + self.mode = mode
154 +
155 + # Copy stream for async H2D transfers
156 + self.copy_stream = torch.cuda.Stream()
157 +
158 + # Module offloaders and buffer pool (populated in wrap_modules/post_init)
159 + self.module_offloaders: list[_ModuleOffloader] = []
160 + self.buffer_pool: StaticBufferPool | None = None
161 + self.total_offloaded_bytes = 0
162 +
163 + def wrap_modules(
164 + self,
165 + modules_generator: Generator[nn.Module, None, None],
166 + ) -> list[nn.Module]:
167 + """Wrap modules with prefetch offloading logic."""
168 + assert len(self.module_offloaders) == 0, (
169 + "wrap_modules should only be called once"
170 + )
171 +
172 + all_modules = []
173 + offload_modules = []
174 +
175 + for module_index, module in enumerate(modules_generator):
176 + all_modules.append(module)
177 +
178 + # Select layers to offload based on group pattern
179 + # Offload last num_in_group layers of each group_size
180 + if module_index % self.group_size >= self.group_size - self.num_in_group:
181 + if self.offload_params:
182 + whitelist = [
183 + name
184 + for name, _ in module.named_parameters()
185 + if any(f".{p}." in f".{name}." for p in self.offload_params)
186 + ]
187 + else:
188 + whitelist = [name for name, _ in module.named_parameters()]
189 +
190 + if not whitelist:
191 + continue # skip layers with no matching params
192 +
193 + offload_modules.append(module)
194 + self.module_offloaders.append(
195 + _ModuleOffloader(
196 + mode=self.mode,
197 + module=module,
198 + copy_stream=self.copy_stream,
199 + whitelist_param_names=whitelist,
200 + layer_idx=len(self.module_offloaders),
201 + )
202 + )
203 +
204 + for index, module in enumerate(offload_modules):
205 + self._hook_module_forward(index, module)
206 +
207 + return all_modules
208 +
209 + def _hook_module_forward(self, index: int, module: nn.Module):
210 + """Hook module's forward with torch.compile-compatible sync."""
211 + original_forward = module.forward
212 +
213 + def forward(*args, **kwargs):
214 + # Temporarily restore original forward to avoid recursion
215 + module.forward = original_forward
216 +
217 + # Wait for this layer's prefetch to complete
218 + # mutates_args on input_tensor creates data dependency for torch.compile
219 + input_tensor = args[0] if args else kwargs.get("hidden_states")
220 + torch.ops.vllm.wait_prefetch(input_tensor, index)
221 +
222 + # No parameter swapping needed - parameters already point to
223 + # GPU static buffers (set in assign_static_buffer)
224 + output = original_forward(*args, **kwargs)
225 +
226 + # Start prefetch for next layer (circular)
227 + # mutates_args on output_tensor creates ordering dependency
228 + next_index = (index + self.prefetch_step) % len(self.module_offloaders)
229 + # Handle tuple output (e.g., (hidden_states, residual))
230 + if isinstance(output, tuple):
231 + torch.ops.vllm.start_prefetch(output[0], next_index)
232 + else:
233 + torch.ops.vllm.start_prefetch(output, next_index)
234 +
235 + # No explicit offload needed - static buffers are reused implicitly
236 +
237 + # Restore hooked forward
238 + module.forward = forward
239 + return output
240 +
241 + module.forward = forward
242 +
243 + def _wait_for_layer(self, layer_idx: int):
244 + """Called by custom op - wait for copy to complete.
245 +
246 + Synchronization strategy:
247 + - During CUDA graph capture: use event-based wait (graph-compatible)
248 + - Outside capture (warmup/eager): use wait_stream (more robust)
249 +
250 + During capture, we skip wait for pre-capture prefetches because:
251 + 1. sync_before_graph_capture() ensures pre-capture work is complete
252 + 2. We can't wait on pre-capture events during capture (isolation error)
253 + """
254 + offloader = self.module_offloaders[layer_idx]
255 +
256 + if torch.cuda.is_current_stream_capturing():
257 + # During capture, skip wait for pre-capture prefetches.
258 + # sync_before_graph_capture() ensures pre-capture work is complete.
259 + if not offloader._prefetch_in_capture:
260 + return
261 + # Event-based wait for in-capture prefetches (graph-compatible)
262 + torch.cuda.current_stream().wait_event(offloader._copy_done_event)
263 + # Mark that this prefetch has been waited on (joined).
264 + offloader._prefetch_in_capture = False
265 + else:
266 + if offloader._event_valid_for_eager:
267 + # Use per-layer event to only wait for THIS layer's copy,
268 + # allowing other layers' prefetches to run concurrently.
269 + torch.cuda.current_stream().wait_event(offloader._copy_done_event)
270 + else:
271 + # Event not usable (unrecorded or recorded during capture).
272 + # Fall back to wait_stream to drain all copy_stream work.
273 + torch.cuda.current_stream().wait_stream(self.copy_stream)
274 +
275 + def sync_prev_onload(self):
276 + """Sync previous onload operations.
277 +
278 + Ensures any H2D copies in flight on copy_stream complete before
279 + the compute stream continues. Call this before CUDA graph
280 + capture/replay or when synchronization is needed.
281 + """
282 + torch.cuda.current_stream().wait_stream(self.copy_stream)
283 +
284 + def _start_prefetch(self, layer_idx: int):
285 + """Called by custom op - start async copy to static buffer."""
286 + offloader = self.module_offloaders[layer_idx]
287 + offloader.start_onload_to_static()
288 +
289 + def join_after_forward(self):
290 + """Join copy_stream after model forward completes.
291 +
292 + Call this after the model forward pass but before CUDA graph capture
293 + ends. This ensures copy_stream is rejoined for any prefetches started
294 + during the forward pass.
295 +
296 + We join ALL layers that have _prefetch_in_capture=True, meaning their
297 + prefetch was started during capture but not yet waited on (joined).
298 + This handles both full and piecewise cudagraph modes correctly:
299 + - Full mode: joins layers 0..prefetch_step-1 (prefetched by last layers)
300 + - Piecewise mode: joins only layers prefetched by THIS subgraph's layers
301 + """
302 + if not self.module_offloaders:
303 + return
304 + # Join all layers whose prefetch was started in capture but not waited on
305 + for offloader in self.module_offloaders:
306 + if offloader._prefetch_in_capture:
307 + torch.cuda.current_stream().wait_event(offloader._copy_done_event)
308 + offloader._prefetch_in_capture = False
309 +
310 + def post_init(self):
311 + """Allocate static buffer pool and start initial prefetches.
312 +
313 + Note: Parameters have already been offloaded to CPU during wrap_modules()
314 + (in _CpuParamOffloader.__init__), so GPU memory is available for the
315 + static buffer pool.
316 + """
317 + # Sync CPU storage with current param.data BEFORE collecting param info.
318 + # This is needed because process_weights_after_loading may have:
319 + # 1. Transformed weights (quantization, transpose, etc.)
320 + # 2. Created new CPU tensors via device_loading_context
321 + # Our _cpu_storage would be stale otherwise.
322 + for offloader in self.module_offloaders:
323 + offloader.sync_cpu_storage()
324 +
325 + # Collect parameter info (now using synced CPU storage)
326 + param_infos: list[ParamInfo] = []
327 + device: torch.device | None = None
328 +
329 + for offloader in self.module_offloaders:
330 + param_infos.extend(offloader.get_param_infos())
331 + if device is None:
332 + device = offloader.device
333 +
334 + if device is None:
335 + # No modules to offload
336 + return
337 +
338 + # Allocate static buffer pool
339 + self.buffer_pool = StaticBufferPool(
340 + param_infos=param_infos,
341 + slot_capacity=self.prefetch_step,
342 + device=device,
343 + )
344 +
345 + # Assign buffer slots and point parameters to GPU buffers
346 + for idx, offloader in enumerate(self.module_offloaders):
347 + slot_idx = idx % self.prefetch_step
348 + offloader.assign_buffer_slot(self.buffer_pool, slot_idx)
349 +
350 + # Collect offloaded bytes
351 + for offloader in self.module_offloaders:
352 + offloader.post_init()
353 + self.total_offloaded_bytes += offloader.offloaded_bytes
354 +
355 + logger.info_once(
356 + f"[PrefetchOffloader] Initialized {len(self.module_offloaders)} modules. "
357 + f"Total GPU memory saved: {self.total_offloaded_bytes / 1e9:.4f} GB, "
358 + f"Static buffer pool: {self.buffer_pool.total_bytes / 1e9:.4f} GB "
359 + f"(group_size={self.group_size}, num_in_group={self.num_in_group}, "
360 + f"prefetch_step={self.prefetch_step}, mode={self.mode})"
361 + )
362 +
363 + # Start initial prefetches
364 + for i in range(min(self.prefetch_step, len(self.module_offloaders))):
365 + self.module_offloaders[i].start_onload_to_static()
366 +
367 +
368 +class _ModuleOffloader:
369 + """Manages offloading for a single module.
370 +
371 + Uses static buffers from a shared pool instead of dynamic allocation.
372 + """
373 +
374 + def __init__(
375 + self,
376 + mode: str,
377 + module: nn.Module,
378 + copy_stream: torch.cuda.Stream,
379 + whitelist_param_names: list[str],
380 + layer_idx: int,
381 + ):
382 + self.mode = mode
383 + self.module = module
384 + self.device = next(module.parameters()).device
385 + self.copy_stream = copy_stream
386 + self.layer_idx = layer_idx
387 + self.offloaded_bytes = 0
388 +
389 + # Event to signal when H2D copy to static buffer is complete.
390 + # Used for per-layer synchronization (both eager and capture modes).
391 + self._copy_done_event = torch.cuda.Event()
392 +
393 + # Track whether _copy_done_event is valid for eager-mode wait_event.
394 + # False when: (1) never recorded, or (2) last recorded during a
395 + # cudagraph capture (events become invalid after capture ends).
396 + # In these cases we fall back to wait_stream.
397 + self._event_valid_for_eager = False
398 +
399 + # Track if last prefetch was started during CUDA graph capture.
400 + # Used to skip wait_event during capture for pre-capture prefetches.
401 + self._prefetch_in_capture = False
402 +
403 + assert self.device != torch.device("cpu"), (
404 + "Module parameters should not already be on CPU "
405 + "(offloader handles CPU placement)"
406 + )
407 +
408 + # Buffer pool and slot (assigned in assign_buffer_slot)
409 + self._buffer_pool: StaticBufferPool | None = None
410 + self._buffer_slot_idx: int = 0
411 +
412 + param_dict = dict(self.module.named_parameters())
413 + assert all(name in param_dict for name in whitelist_param_names), (
414 + f"Whitelist params {whitelist_param_names} not found in module params "
415 + f"{list(param_dict.keys())}"
416 + )
417 +
418 + self._param_offloaders = {
419 + name: _BaseParamOffloader.create(mode, module=module, param_name=name)
420 + for name in whitelist_param_names
421 + }
422 +
423 + def post_init(self):
424 + """Collect total offloaded bytes (offloading already done in __init__)."""
425 + for param_offloader in self._param_offloaders.values():
426 + param_offloader.post_init()
427 + self.offloaded_bytes += param_offloader.offloaded_bytes
428 +
429 + def sync_cpu_storage(self):
430 + """Sync CPU storage with current param.data.
431 +
432 + Called after process_weights_after_loading to ensure _cpu_storage
433 + contains the final processed weights, not stale pre-loading data.
434 + """
435 + for param_offloader in self._param_offloaders.values():
436 + param_offloader.sync_cpu_storage()
437 +
438 + def get_param_infos(self) -> list[ParamInfo]:
439 + """Get parameter metadata for buffer pool allocation.
440 +
441 + Note: sync_cpu_storage() must be called before this method to ensure
442 + _cpu_storage reflects the final processed weights (after quantization).
443 + """
444 + infos = []
445 + for name, offloader in self._param_offloaders.items():
446 + cpu_storage = offloader._cpu_storage
447 + assert cpu_storage is not None, "CPU storage not initialized"
448 + infos.append(
449 + ParamInfo(
450 + name=name,
451 + shape=tuple(cpu_storage.shape),
452 + stride=tuple(cpu_storage.stride()),
453 + dtype=cpu_storage.dtype,
454 + )
455 + )
456 + return infos
457 +
458 + def assign_buffer_slot(self, pool: StaticBufferPool, slot_idx: int):
459 + """Assign this module to a buffer slot in the pool.
460 +
461 + Also assigns static GPU buffers to each parameter offloader,
462 + which moves the parameter data to point to the GPU buffer.
463 + """
464 + self._buffer_pool = pool
465 + self._buffer_slot_idx = slot_idx
466 +
467 + # Assign static buffers to parameters
468 + # Use CPU storage shape/stride/dtype since param.data is now empty
469 + for name, offloader in self._param_offloaders.items():
470 + cpu_storage = offloader._cpu_storage
471 + assert cpu_storage is not None, "CPU storage not initialized"
472 + buffer = pool.get_buffer(
473 + name=name,
474 + shape=tuple(cpu_storage.shape),
475 + stride=tuple(cpu_storage.stride()),
476 + dtype=cpu_storage.dtype,
477 + slot_idx=slot_idx,
478 + )
479 + offloader.assign_static_buffer(buffer)
480 +
481 + def start_onload_to_static(self):
482 + """Start async copy from CPU storage to GPU buffer.
483 +
484 + Uses event-based forking to join copy_stream to CUDA graph capture.
485 + This ensures H2D copies are properly captured when recording a graph.
486 +
487 + IMPORTANT: We must wait for the compute stream before copying, because
488 + the previous layer's forward may still be using the buffer (GPU ops are
489 + async). Without this sync, we could overwrite the buffer while it's
490 + being read.
491 + """
492 + assert self._buffer_pool is not None, "Buffer pool not assigned"
493 +
494 + # Track if this prefetch is being captured (for _wait_for_layer logic)
495 + self._prefetch_in_capture = torch.cuda.is_current_stream_capturing()
496 +
497 + # Fork: record event on compute stream, copy_stream waits on it
498 + # This joins copy_stream to any active CUDA graph capture
499 + fork_event = torch.cuda.Event()
500 + torch.cuda.current_stream().record_event(fork_event)
501 + self.copy_stream.wait_event(fork_event)
502 +
503 + with torch.cuda.stream(self.copy_stream):
504 + for name, offloader in self._param_offloaders.items():
505 + cpu_storage = offloader._cpu_storage
506 + gpu_buffer = offloader._gpu_buffer
507 + assert cpu_storage is not None, "CPU storage not initialized"
508 + assert gpu_buffer is not None, "GPU buffer not assigned"
509 + assert not is_pin_memory_available() or cpu_storage.is_pinned(), (
510 + f"CPU storage for {name} is not pinned! "
511 + "non_blocking=True H2D copy from non-pinned memory "
512 + "causes stream synchronization that breaks "
513 + "event-based fork synchronization."
514 + )
515 + gpu_buffer.copy_(cpu_storage, non_blocking=True)
516 +
517 + # Record completion event for _wait_for_layer to use
518 + self._copy_done_event.record(self.copy_stream)
519 + # Event is only valid for eager wait_event if recorded outside capture.
520 + # Events recorded during capture become invalid after capture ends.
521 + self._event_valid_for_eager = not torch.cuda.is_current_stream_capturing()
522 +
523 +
524 +class _BaseParamOffloader(ABC):
525 + """Base class for parameter offloading strategies."""
526 +
527 + # CPU storage for offloaded parameters (set by subclasses)
528 + _cpu_storage: torch.Tensor | None
529 + # GPU buffer reference (set by subclasses when using static buffers)
530 + _gpu_buffer: torch.Tensor | None
531 +
532 + @staticmethod
533 + def create(mode: str, **kwargs) -> "_BaseParamOffloader":
534 + """Factory method to create appropriate offloader for mode."""
535 + if mode == "cpu":
536 + return _CpuParamOffloader(**kwargs)
537 + else:
538 + raise ValueError(f"Unknown offload mode: {mode}")
539 +
540 + def __init__(self, module: nn.Module, param_name: str):
541 + self._module = module
542 + self._param_name = param_name
543 + self.offloaded_bytes = 0
544 + self._cpu_storage = None
545 + self._gpu_buffer = None
546 +
547 + @property
548 + def _param(self) -> nn.Parameter:
549 + """Get the parameter being offloaded.
550 +
551 + Supports dotted names (e.g. 'self_attn.qkv_proj.weight') by
552 + traversing the module hierarchy.
553 + """
554 + obj: Any = self._module
555 + for attr in self._param_name.split("."):
556 + obj = getattr(obj, attr)
557 + return obj
558 +
559 + def post_init(self):
560 + """Initialize offloading (move parameter to storage)."""
561 + return
562 +
563 + @abstractmethod
564 + def sync_cpu_storage(self) -> None:
565 + """Sync CPU storage with current param.data.
566 +
567 + Called after process_weights_after_loading to update _cpu_storage
568 + with the final processed weights.
569 + """
570 + pass
571 +
572 + @abstractmethod
573 + def assign_static_buffer(self, gpu_buffer: torch.Tensor) -> None:
574 + """Point parameter data to GPU static buffer."""
575 + pass
576 +
577 +
578 +class _CpuParamOffloader(_BaseParamOffloader):
579 + """Offload parameter to pinned CPU memory.
580 +
581 + Uses GPU static buffers as the actual parameter, with CPU storage
582 + kept separately. This ensures torch.compile sees GPU tensors at trace time.
583 +
584 + The offloading happens in two phases:
585 + 1. __init__() - copies GPU data to CPU, frees GPU memory immediately
586 + 2. assign_static_buffer() - points param.data to GPU static buffer
587 + """
588 +
589 + def __init__(self, module: nn.Module, param_name: str):
590 + super().__init__(module, param_name)
591 + self._cpu_storage: torch.Tensor | None = None
592 + self._gpu_buffer: torch.Tensor | None = None # Store reference to GPU buffer
593 +
594 + # Offload to CPU immediately to free GPU memory during model loading
595 + self._offload_to_cpu_internal()
596 +
597 + def _offload_to_cpu_internal(self):
598 + """Copy parameter data to pinned CPU storage and free GPU memory.
599 +
600 + This replaces param.data with CPU storage, allowing weight loading
601 + to continue writing to CPU memory. GPU memory is freed when the
602 + original GPU tensor is garbage collected.
603 + """
604 + param = self._param
605 + pin_memory = is_pin_memory_available()
606 +
607 + # Create pinned CPU storage and copy current GPU data
608 + self._cpu_storage = torch.empty_strided(
609 + size=param.data.size(),
610 + stride=param.data.stride(),
611 + dtype=param.data.dtype,
612 + layout=param.data.layout,
613 + device="cpu",
614 + pin_memory=pin_memory,
615 + )
616 + self._cpu_storage.copy_(param.data)
617 +
618 + self.offloaded_bytes = (
619 + self._cpu_storage.numel() * self._cpu_storage.element_size()
620 + )
621 +
622 + # Point param.data to CPU storage - this allows weight loading to work
623 + # and frees GPU memory when the original GPU tensor is garbage collected
624 + param.data = self._cpu_storage
625 +
626 + def _update_cpu_storage_from_param(self) -> None:
627 + """Update _cpu_storage from current param.data, ensuring pinned memory.
628 +
629 + After process_weights_after_loading, device_loading_context creates
630 + non-pinned CPU tensors via `p.data = p.data.to("cpu")`. Using
631 + non-pinned memory with `copy_(src, non_blocking=True)` causes CUDA to
632 + perform a stream synchronization before the copy, breaking the
633 + event-based fork synchronization and potentially allowing the copy
634 + to overwrite the GPU buffer while the compute stream still reads it.
635 +
636 + This method ensures _cpu_storage always uses pinned memory when
637 + available, re-pinning if necessary.
638 + """
639 + param = self._param
640 +
641 + if param.data.device.type == "cpu":
642 + if is_pin_memory_available() and not param.data.is_pinned():
643 + pinned = torch.empty_strided(
644 + size=param.data.size(),
645 + stride=param.data.stride(),
646 + dtype=param.data.dtype,
647 + layout=param.data.layout,
648 + device="cpu",
649 + pin_memory=True,
650 + )
651 + pinned.copy_(param.data)
652 + self._cpu_storage = pinned
653 + else:
654 + self._cpu_storage = param.data
655 + else:
656 + # param.data is on GPU - copy to existing CPU storage
657 + assert self._cpu_storage is not None
658 + self._cpu_storage.copy_(param.data)
659 +
660 + def assign_static_buffer(self, gpu_buffer: torch.Tensor) -> None:
661 + """Point parameter data to GPU static buffer.
662 +
663 + This is called after weight loading AND process_weights_after_loading
664 + complete. At this point:
665 + - param.data may have been replaced by device_loading_context
666 + (which creates new CPU tensors after quantization processing)
667 + - We need to update _cpu_storage to point to current param.data
668 + so that prefetch copies the processed weights, not stale data
669 + - Then point param.data to the GPU buffer for torch.compile
670 + """
671 + assert self._cpu_storage is not None, (
672 + "_offload_to_cpu_internal() must be called before assign_static_buffer()"
673 + )
674 +
675 + # Get current parameter (may have been replaced by
676 + # process_weights_after_loading)
677 + param = self._param
678 +
679 + # Update _cpu_storage to current param.data. This is critical because:
680 + # 1. process_weights_after_loading may transform weights (quantization)
681 + # 2. device_loading_context creates NEW CPU tensors when moving back
682 + # 3. Our old _cpu_storage would have pre-processed or stale data
683 + self._update_cpu_storage_from_param()
684 +
685 + # Store reference to GPU buffer for use in start_onload
686 + self._gpu_buffer = gpu_buffer
687 +
688 + # Point parameter to static GPU buffer - this is what torch.compile sees
689 + param.data = gpu_buffer
690 +
691 + def sync_cpu_storage(self) -> None:
692 + """Sync CPU storage with current param.data.
693 +
694 + Called after process_weights_after_loading to update _cpu_storage
695 + with the final processed weights. This is critical because:
696 + 1. process_weights_after_loading may transform weights (quantization)
697 + 2. device_loading_context creates NEW CPU tensors when moving back
698 + 3. Our old _cpu_storage would have pre-processed or stale data
699 + """
700 + self._update_cpu_storage_from_param()
701 +
702 + def post_init(self):
703 + """No-op: offloading done in offload_to_cpu/assign_static_buffer."""
704 + pass
vllm/model_executor/offloader/prefetch_ops.py new
+94
@@ -0,0 +1,94 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Custom ops for prefetch offloader torch.compile + CUDA graph compatibility.
4 +
5 +These ops use mutates_args to create data dependencies that prevent
6 +the compiler from reordering prefetch/sync operations.
7 +"""
8 +
9 +from __future__ import annotations
10 +
11 +import torch
12 +
13 +from vllm.model_executor.offloader.base import get_offloader
14 +from vllm.utils.torch_utils import direct_register_custom_op
15 +
16 +# --- wait_prefetch op ---
17 +
18 +
19 +def _wait_prefetch_impl(
20 + input_tensor: torch.Tensor,
21 + layer_idx: int,
22 +) -> None:
23 + """Wait for prefetch of layer_idx to complete.
24 +
25 + Synchronizes the compute stream with the copy stream to ensure
26 + the prefetched weights are ready for use.
27 +
28 + Args:
29 + input_tensor: Input to the layer (e.g., hidden_states) - declared
30 + as mutated to create data dependency for torch.compile.
31 + layer_idx: Index of the layer to wait for.
32 + """
33 + get_offloader()._wait_for_layer(layer_idx)
34 +
35 +
36 +def _wait_prefetch_fake(
37 + input_tensor: torch.Tensor,
38 + layer_idx: int,
39 +) -> None:
40 + """Fake implementation for torch.compile tracing."""
41 + return
42 +
43 +
44 +# --- start_prefetch op ---
45 +
46 +
47 +def _start_prefetch_impl(
48 + output_tensor: torch.Tensor,
49 + layer_idx: int,
50 +) -> None:
51 + """Start async prefetch of layer_idx weights.
52 +
53 + Initiates H2D copy on the copy stream for the specified layer.
54 +
55 + Args:
56 + output_tensor: Output from forward - declared as mutated to
57 + prevent torch.compile from reordering this op before the
58 + computation that produces output_tensor.
59 + layer_idx: Index of the layer to prefetch.
60 + """
61 + get_offloader()._start_prefetch(layer_idx)
62 +
63 +
64 +def _start_prefetch_fake(
65 + output_tensor: torch.Tensor,
66 + layer_idx: int,
67 +) -> None:
68 + """Fake implementation for torch.compile tracing."""
69 + return
70 +
71 +
72 +def register_prefetch_offloader_ops() -> None:
73 + """Register custom ops for prefetch offloader.
74 +
75 + Must be called before the ops are used. This is typically done
76 + at module import time.
77 + """
78 + direct_register_custom_op(
79 + op_name="wait_prefetch",
80 + op_func=_wait_prefetch_impl,
81 + mutates_args=["input_tensor"],
82 + fake_impl=_wait_prefetch_fake,
83 + )
84 +
85 + direct_register_custom_op(
86 + op_name="start_prefetch",
87 + op_func=_start_prefetch_impl,
88 + mutates_args=["output_tensor"],
89 + fake_impl=_start_prefetch_fake,
90 + )
91 +
92 +
93 +# Register ops at module import time
94 +register_prefetch_offloader_ops()
vllm/model_executor/offloader/uva.py new
+140
@@ -0,0 +1,140 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""UVA-based CPU offloading using Unified Virtual Addressing."""
4 +
5 +from collections.abc import Generator
6 +
7 +import torch
8 +import torch.nn as nn
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
14 +from vllm.utils.mem_utils import format_gib
15 +from vllm.utils.platform_utils import is_pin_memory_available, is_uva_available
16 +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
17 +
18 +logger = init_logger(__name__)
19 +
20 +
21 +class UVAOffloader(BaseOffloader):
22 + """Offloader using Unified Virtual Addressing (UVA) for zero-copy access.
23 +
24 + This offloader moves parameters to pinned CPU memory and creates CUDA views
25 + using UVA. The GPU can then directly access the CPU memory without explicit
26 + transfers, at the cost of PCIe bandwidth (slower than GPU memory).
27 +
28 + When UVA is disabled via env var, falls back to a functional_call-based
29 + approach that moves parameters on-demand.
30 +
31 + Args:
32 + cpu_offload_max_bytes: Maximum bytes to offload to CPU.
33 + cpu_offload_params: Set of parameter name segments to selectively
34 + offload. If empty, all parameters are eligible up to the byte limit.
35 + """
36 +
37 + def __init__(
38 + self,
39 + cpu_offload_max_bytes: int,
40 + cpu_offload_params: set[str] | None = None,
41 + ):
42 + self.cpu_offload_max_bytes = cpu_offload_max_bytes
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 + )
50 + self.uva_offloading = (
51 + is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA
52 + )
53 +
54 + def wrap_modules(
55 + self,
56 + modules_generator: Generator[nn.Module, None, None],
57 + ) -> list[nn.Module]:
58 + """Wrap modules with UVA offloading."""
59 + modules = [self._maybe_offload_to_cpu(module) for module in modules_generator]
60 + if self.cpu_offload_bytes > 0:
61 + logger.info(
62 + "Total CPU offloaded parameters: %s",
63 + format_gib(self.cpu_offload_bytes),
64 + )
65 + return modules
66 +
67 + def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module:
68 + """Offload module parameters to CPU using UVA if budget allows."""
69 + if (params := next(module.parameters(), None)) is None:
70 + return module
71 +
72 + device = params.device
73 +
74 + if device == torch.device("cpu"):
75 + return module
76 +
77 + if self.cpu_offload_bytes >= self.cpu_offload_max_bytes:
78 + return module
79 +
80 + # offload parameters to CPU
81 + # use pin_memory if possible, which helps cudagraph capture speed
82 + offloaded_parameters = False
83 + for name, p in module.named_parameters():
84 + if self.cpu_offload_bytes >= self.cpu_offload_max_bytes:
85 + # we use per-parameter offloading
86 + # one module might have some parameters offloaded and some not
87 + break
88 +
89 + if self.cpu_offload_params:
90 + # Check if parameter belongs to the offloading set
91 + # Add dots here to ensure we match full segments only
92 + # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight"
93 + # but not "mlp.experts.w2_weight_scale"
94 + should_offload = any(
95 + f".{param}." in f".{name}." for param in self.cpu_offload_params
96 + )
97 + if not should_offload:
98 + continue
99 +
100 + cpu_data = p.data.to(device="cpu")
101 + if self.pin_memory:
102 + cpu_data = cpu_data.pin_memory()
103 +
104 + if not self.uva_offloading:
105 + p.data = cpu_data
106 + else:
107 + p.data = get_accelerator_view_from_cpu_tensor(cpu_data)
108 + p._vllm_is_uva_offloaded = True
109 +
110 + self.cpu_offload_bytes += p.data.numel() * p.data.element_size()
111 + offloaded_parameters = True
112 +
113 + if offloaded_parameters and not self.uva_offloading:
114 + original_forward = module.forward
115 +
116 + def forward(*args, **kwargs):
117 + module.forward = original_forward
118 + device_state = {
119 + # here we blindly call `to(device)`
120 + # if the parameter is already on the device,
121 + # it will be a no-op
122 + k: v.to(device, non_blocking=True)
123 + for k, v in module.state_dict().items()
124 + }
125 +
126 + # set `tie_weights=False` as tied weights in original model
127 + # become untied when calling .to(device) individually
128 + output = functional_call(
129 + module,
130 + device_state,
131 + args=args,
132 + kwargs=kwargs,
133 + tie_weights=False,
134 + )
135 + module.forward = forward
136 + return output
137 +
138 + module.forward = forward
139 +
140 + return module
vllm/v1/worker/gpu/cudagraph_utils.py
+18
@@ -12,6 +12,7 @@ from vllm.config import VllmConfig
12 from vllm.config.compilation import CUDAGraphMode
13 from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
14 from vllm.forward_context import BatchDescriptor, set_forward_context
15 +from vllm.model_executor.offloader.base import get_offloader
16 from vllm.utils.math_utils import cdiv
17 from vllm.v1.kv_cache_interface import KVCacheConfig
18 from vllm.v1.worker.gpu.attn_utils import (
@@ -189,6 +190,11 @@ class CudaGraphManager:
190 # Capture the graph.
191 assert num_tokens not in self.graphs
192 graph = torch.cuda.CUDAGraph()
193 +
194 + # Sync offloader's copy stream before capture.
195 + # Ensure any pre-capture prefetches from offloader are complete.
196 + get_offloader().sync_prev_onload()
197 +
198 with (
199 set_forward_context(
200 attn_metadata=attn_metadata,
@@ -205,6 +211,11 @@ class CudaGraphManager:
211 positions=positions,
212 inputs_embeds=inputs_embeds,
213 )
214 + # Join offloader's copy stream after forward to avoid unjoined
215 + # stream error. The last layer's start_prefetch forks copy_stream,
216 + # but wait_prefetch only happens in the next forward pass.
217 + get_offloader().join_after_forward()
218 +
219 if self.use_aux_hidden_state_outputs:
220 hidden_states, aux_hidden_states = model_output
221 else:
@@ -329,6 +340,13 @@ class CudaGraphManager:
340 self, num_tokens: int
341 ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
342 assert num_tokens in self.graphs, f"No cudagraph for {num_tokens} tokens"
343 + # Sync offloader before replay - needed when transitioning from
344 + # eager/piecewise to full cudagraph (e.g., prefill → decode).
345 + # The previous eager iteration's start_prefetch may have queued
346 + # H2D copies on copy_stream that the graph's captured events
347 + # cannot see. Without this, replay could overwrite static buffers
348 + # while those copies are still in flight.
349 + get_offloader().sync_prev_onload()
350 self.graphs[num_tokens].replay()
351 assert self.hidden_states is not None
352 hidden_states = self.hidden_states[:num_tokens]
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+17
@@ -7,6 +7,7 @@ import torch
7
8 from vllm.config import VllmConfig
9 from vllm.config.compilation import CUDAGraphMode
10 +from vllm.model_executor.offloader.base import get_offloader
11 from vllm.v1.kv_cache_interface import KVCacheConfig
12 from vllm.v1.worker.gpu.block_table import BlockTables
13 from vllm.v1.worker.gpu.cudagraph_utils import (
@@ -115,6 +116,11 @@ class EagleCudaGraphManager:
116 ) -> None:
117 assert num_tokens not in self.graphs
118 graph = torch.cuda.CUDAGraph()
119 +
120 + # Sync offloader's copy stream before capture.
121 + # Ensure any pre-capture prefetches from offloader are complete.
122 + get_offloader().sync_prev_onload()
123 +
124 with torch.cuda.graph(graph, self.pool):
125 generate_fn(
126 num_reqs,
@@ -124,6 +130,10 @@ class EagleCudaGraphManager:
130 num_tokens_across_dp,
131 CUDAGraphMode.NONE,
132 )
133 + # Join offloader's copy stream after forward to avoid unjoined
134 + # stream error. The last layer's start_prefetch forks copy_stream,
135 + # but wait_prefetch only happens in the next forward pass.
136 + get_offloader().join_after_forward()
137 self.graphs[num_tokens] = graph
138
139 def _capture_piecewise_graph(
@@ -171,4 +181,11 @@ class EagleCudaGraphManager:
181
182 def run_fullgraph(self, num_tokens: int) -> None:
183 assert num_tokens in self.graphs
184 + # Sync offloader before replay - needed when transitioning from
185 + # eager/piecewise to full cudagraph (e.g., prefill → decode).
186 + # The previous eager iteration's start_prefetch may have queued
187 + # H2D copies on copy_stream that the graph's captured events
188 + # cannot see. Without this, replay could overwrite static buffers
189 + # while those copies are still in flight.
190 + get_offloader().sync_prev_onload()
191 self.graphs[num_tokens].replay()
vllm/v1/worker/gpu_model_runner.py
+13 -9
@@ -81,6 +81,11 @@ from vllm.model_executor.models.interfaces_base import (
81 is_pooling_model,
82 is_text_generation_model,
83 )
84 +from vllm.model_executor.offloader import (
85 + create_offloader,
86 + get_offloader,
87 + set_offloader,
88 +)
89 from vllm.multimodal import MULTIMODAL_REGISTRY
90 from vllm.multimodal.encoder_budget import MultiModalBudget
91 from vllm.multimodal.inputs import (
@@ -378,6 +383,7 @@ class GPUModelRunner(
383 self.vllm_config = vllm_config
384 self.model_config = vllm_config.model_config
385 self.cache_config = vllm_config.cache_config
386 + self.offload_config = vllm_config.offload_config
387 self.compilation_config = vllm_config.compilation_config
388 self.lora_config = vllm_config.lora_config
389 self.load_config = vllm_config.load_config
@@ -386,14 +392,6 @@ class GPUModelRunner(
392 self.speculative_config = vllm_config.speculative_config
393 self.observability_config = vllm_config.observability_config
394
389 - from vllm.model_executor.models.utils import (
390 - set_cpu_offload_max_bytes,
391 - set_cpu_offload_params,
392 - )
393 -
394 - set_cpu_offload_max_bytes(int(self.cache_config.cpu_offload_gb * 1024**3))
395 - set_cpu_offload_params(self.cache_config.cpu_offload_params)
396 -
395 model_config = self.model_config
396 cache_config = self.cache_config
397 scheduler_config = self.scheduler_config
@@ -749,6 +747,10 @@ class GPUModelRunner(
747 pin_memory=self.pin_memory,
748 )
749
750 + # Model weight offloader
751 + # Make sure this is called before any get_offloader call
752 + set_offloader(create_offloader(self.offload_config))
753 +
754 # Ephemeral state transferred between execute_model() and sample_tokens().
755 self.execute_model_state: ExecuteModelState | None = None
756 self.kv_connector_output: KVConnectorOutput | None = None
@@ -4342,6 +4344,8 @@ class GPUModelRunner(
4344 self.model, self.vllm_config, CUDAGraphMode.NONE, self.device
4345 )
4346
4347 + get_offloader().post_init()
4348 +
4349 def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None:
4350 """Extract Eagle3 auxiliary layer indices from speculative config.
4351
@@ -5780,7 +5784,7 @@ class GPUModelRunner(
5784 if block_sizes != [self.cache_config.block_size] or kernel_block_sizes != [
5785 self.cache_config.block_size
5786 ]:
5783 - assert self.cache_config.cpu_offload_gb == 0, (
5787 + assert self.offload_config.uva.cpu_offload_gb == 0, (
5788 "Cannot re-initialize the input batch when CPU weight "
5789 "offloading is enabled. See https://github.com/vllm-project/vllm/pull/18298 " # noqa: E501
5790 "for more details."
vllm/v1/worker/gpu_ubatch_wrapper.py
+13
@@ -20,6 +20,7 @@ from vllm.forward_context import (
20 override_forward_context,
21 )
22 from vllm.logger import init_logger
23 +from vllm.model_executor.offloader.base import get_offloader
24 from vllm.platforms import current_platform
25 from vllm.sequence import IntermediateTensors
26 from vllm.utils.import_utils import has_deep_gemm
@@ -239,6 +240,11 @@ class UBatchWrapper:
240 set_graph_pool_id(self.graph_pool)
241 else:
242 set_graph_pool_id(current_platform.graph_pool_handle())
243 +
244 + # Sync offloader's copy stream before capture.
245 + # Ensure any pre-capture prefetches from offloader are complete.
246 + get_offloader().sync_prev_onload()
247 +
248 with torch.cuda.graph(
249 cudagraph_metadata.cudagraph,
250 stream=compute_stream,
@@ -250,6 +256,10 @@ class UBatchWrapper:
256 sorted_results = [value for position, value in sorted(results)]
257 result = torch.cat(sorted_results, dim=0)
258 cudagraph_metadata.outputs = result
259 + # Join offloader's copy stream after forward to avoid unjoined
260 + # stream error. The last layer's start_prefetch forks copy_stream,
261 + # but wait_prefetch only happens in the next forward pass.
262 + get_offloader().join_after_forward()
263 self.cudagraphs[num_tokens] = cudagraph_metadata
264 return cudagraph_metadata.outputs
265
@@ -461,6 +471,9 @@ class UBatchWrapper:
471 and cudagraph_runtime_mode is CUDAGraphMode.FULL
472 ):
473 cudagraph_metadata = self.cudagraphs[num_tokens]
474 + # Sync offloader before replay - ensures any external dependencies
475 + # from pre-capture prefetches are satisfied.
476 + get_offloader().sync_prev_onload()
477 cudagraph_metadata.cudagraph.replay()
478 return cudagraph_metadata.outputs
479 else: