From a3dfa47086e0c9ca2ae8c52df9747ff9d41e792e Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 2 Dec 2025 00:28:40 -0800 Subject: [PATCH 01/36] offloader v2 Signed-off-by: Ming Yang --- tests/basic_correctness/test_v2_offload.py | 31 +++ vllm/config/cache.py | 10 + vllm/engine/arg_utils.py | 15 ++ vllm/entrypoints/llm.py | 14 + vllm/model_executor/models/deepseek_v2.py | 25 ++ vllm/model_executor/models/utils.py | 27 +- vllm/model_executor/offloader/__init__.py | 21 ++ vllm/model_executor/offloader/base.py | 94 +++++++ vllm/model_executor/offloader/uva.py | 175 +++++++++++++ vllm/model_executor/offloader/v2.py | 290 +++++++++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 36 ++- 11 files changed, 733 insertions(+), 5 deletions(-) create mode 100644 tests/basic_correctness/test_v2_offload.py create mode 100644 vllm/model_executor/offloader/__init__.py create mode 100644 vllm/model_executor/offloader/base.py create mode 100644 vllm/model_executor/offloader/uva.py create mode 100644 vllm/model_executor/offloader/v2.py diff --git a/tests/basic_correctness/test_v2_offload.py b/tests/basic_correctness/test_v2_offload.py new file mode 100644 index 000000000000..bbe89a02def2 --- /dev/null +++ b/tests/basic_correctness/test_v2_offload.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test V2 offloading correctness with DeepSeek V2 model.""" + +from ..utils import compare_two_settings + + +def test_v2_offload_deepseek(): + """Test V2 CPU offloading with DeepSeek-V2-Lite. + + Compares outputs between: + 1. Baseline (no offloading) + 2. V2 offloading (group_size=8, num_in_group=2, prefetch_step=1) + + This tests the advanced offloading with prefetching on a MoE model. + """ + compare_two_settings( + "deepseek-ai/DeepSeek-V2-Lite", + [], # Baseline: no offloading + [ + # V2 offloading configuration + "--offload-group-size", + "8", + "--offload-num-in-group", + "2", + "--offload-prefetch-step", + "1", + # currently not compatible with torch.compile + "--enforce-eager", + ], + ) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 00530846fce0..d25a68c3e2e5 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -89,6 +89,16 @@ class CacheConfig: Note that this requires fast CPU-GPU interconnect, as part of the model is loaded from CPU memory to GPU memory on the fly in each model forward pass. """ + offload_group_size: int = Field(default=0, ge=0) + """Advanced CPU offloading: Group every N layers together. Offload last + `offload_num_in_group` layers of each group. Default is 0 (disabled). + Example: group_size=8, num_in_group=2 offloads layers 6,7,14,15,22,23,... + """ + offload_num_in_group: int = Field(default=1, ge=1) + """Advanced CPU offloading: Number of layers to offload per group. Default is 1.""" + offload_prefetch_step: int = Field(default=1, ge=0) + """Advanced CPU offloading: Number of layers to prefetch ahead. Higher values hide + more latency but use more GPU memory. Default is 1.""" calculate_kv_scales: bool = False """This enables dynamic calculation of `k_scale` and `v_scale` when kv_cache_dtype is fp8. If `False`, the scales will be loaded from the model diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 096217da4fe4..dbc621f2dea8 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -433,6 +433,9 @@ class EngineArgs: disable_cascade_attn: bool = ModelConfig.disable_cascade_attn swap_space: float = CacheConfig.swap_space cpu_offload_gb: float = CacheConfig.cpu_offload_gb + offload_group_size: int = CacheConfig.offload_group_size + offload_num_in_group: int = CacheConfig.offload_num_in_group + offload_prefetch_step: int = CacheConfig.offload_prefetch_step gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None @@ -896,6 +899,15 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--prefix-caching-hash-algo", **cache_kwargs["prefix_caching_hash_algo"] ) cache_group.add_argument("--cpu-offload-gb", **cache_kwargs["cpu_offload_gb"]) + cache_group.add_argument( + "--offload-group-size", **cache_kwargs["offload_group_size"] + ) + cache_group.add_argument( + "--offload-num-in-group", **cache_kwargs["offload_num_in_group"] + ) + cache_group.add_argument( + "--offload-prefetch-step", **cache_kwargs["offload_prefetch_step"] + ) cache_group.add_argument( "--calculate-kv-scales", **cache_kwargs["calculate_kv_scales"] ) @@ -1402,6 +1414,9 @@ def create_engine_config( enable_prefix_caching=self.enable_prefix_caching, prefix_caching_hash_algo=self.prefix_caching_hash_algo, cpu_offload_gb=self.cpu_offload_gb, + offload_group_size=self.offload_group_size, + offload_num_in_group=self.offload_num_in_group, + offload_prefetch_step=self.offload_prefetch_step, calculate_kv_scales=self.calculate_kv_scales, kv_sharing_fast_prefill=self.kv_sharing_fast_prefill, mamba_cache_dtype=self.mamba_cache_dtype, diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index c121fa71f019..47ceddde0afc 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -153,6 +153,14 @@ class LLM: the model weights. This virtually increases the GPU memory space you can use to hold the model weights, at the cost of CPU-GPU data transfer for every forward pass. + offload_group_size: Advanced CPU offloading: Group every N layers + together. Offload last `offload_num_in_group` layers of each group. + Default is 0 (disabled). + offload_num_in_group: Advanced CPU offloading: Number of layers to + offload per group. Default is 1. + offload_prefetch_step: Advanced CPU offloading: Number of layers to + prefetch ahead. Higher values hide more latency but use more GPU + memory. Default is 1. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. @@ -202,6 +210,9 @@ def __init__( gpu_memory_utilization: float = 0.9, swap_space: float = 4, cpu_offload_gb: float = 0, + offload_group_size: int = 0, + offload_num_in_group: int = 1, + offload_prefetch_step: int = 1, enforce_eager: bool = False, disable_custom_all_reduce: bool = False, hf_token: bool | str | None = None, @@ -317,6 +328,9 @@ def __init__( kv_cache_memory_bytes=kv_cache_memory_bytes, swap_space=swap_space, cpu_offload_gb=cpu_offload_gb, + offload_group_size=offload_group_size, + offload_num_in_group=offload_num_in_group, + offload_prefetch_step=offload_prefetch_step, enforce_eager=enforce_eager, disable_custom_all_reduce=disable_custom_all_reduce, hf_token=hf_token, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index a8eb4a69b6f2..6932871902fd 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1275,6 +1275,31 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config, prefix, topk_indices_buffer=topk_indices_buffer ), prefix=f"{prefix}.layers", + offloader_kwargs=dict( + # Extract the MLP submodule - for MoE layers, go deeper to the experts + submodule_accessor=lambda layer: ( + layer.mlp.experts + if isinstance(layer.mlp, DeepseekV2MoE) + else layer.mlp + ), + # Specify which parameters to offload + whitelist_param_names_creator=lambda module: ( + [ + # Core MoE expert weights + "w13_weight", + "w2_weight", + # NVFP4 quantization scales (if present) + *( + ["w13_blockscale_swizzled", "w2_blockscale_swizzled"] + if hasattr(module, "w13_blockscale_swizzled") + else [] + ), + ] + # Only offload from MoE experts (SharedFusedMoE/FusedMoE) + if hasattr(module, "w13_weight") + else [] + ), + ), ) if get_pp_group().is_last_rank: diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index f25ab9153a50..6d0aefbd8122 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -590,24 +590,43 @@ def make_layers( num_hidden_layers: int, layer_fn: LayerFn, prefix: str, + offloader_kwargs: dict | None = None, ) -> tuple[int, int, torch.nn.ModuleList]: """Make a list of layers with the given layer function, taking pipeline parallelism into account. + + Args: + num_hidden_layers: Total number of hidden layers in the model. + layer_fn: Function to create a layer given its index. + prefix: Prefix for layer names. + offloader_kwargs: Optional kwargs for offloader (submodule_accessor, + whitelist_param_names_creator). + + Returns: + Tuple of (start_layer, end_layer, modules). """ from vllm.distributed.parallel_state import get_pp_group from vllm.distributed.utils import get_pp_indices + from vllm.model_executor.offloader import get_offloader start_layer, end_layer = get_pp_indices( num_hidden_layers, get_pp_group().rank_in_group, get_pp_group().world_size ) + + logger.debug(f"{offloader_kwargs=}") + modules = torch.nn.ModuleList( [PPMissingLayer() for _ in range(start_layer)] - + [ - maybe_offload_to_cpu(layer_fn(prefix=f"{prefix}.{idx}")) - for idx in range(start_layer, end_layer) - ] + + get_offloader().wrap_modules( + ( + layer_fn(prefix=f"{prefix}.{idx}") + for idx in range(start_layer, end_layer) + ), + **(offloader_kwargs or {}), + ) + [PPMissingLayer() for _ in range(end_layer, num_hidden_layers)] ) + return start_layer, end_layer, modules diff --git a/vllm/model_executor/offloader/__init__.py b/vllm/model_executor/offloader/__init__.py new file mode 100644 index 000000000000..fa29ce76b6e6 --- /dev/null +++ b/vllm/model_executor/offloader/__init__.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model parameter offloading infrastructure.""" + +from vllm.model_executor.offloader.base import ( + BaseOffloader, + NoopOffloader, + get_offloader, + set_offloader, +) +from vllm.model_executor.offloader.uva import UVAOffloader +from vllm.model_executor.offloader.v2 import OffloaderV2 + +__all__ = [ + "BaseOffloader", + "NoopOffloader", + "UVAOffloader", + "OffloaderV2", + "get_offloader", + "set_offloader", +] diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py new file mode 100644 index 000000000000..8560c52e47b4 --- /dev/null +++ b/vllm/model_executor/offloader/base.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base classes for model parameter offloading.""" + +from abc import ABC, abstractmethod +from collections.abc import Callable, Generator + +import torch.nn as nn + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# Type aliases for clarity +_SubmoduleAccessor = Callable[[nn.Module], nn.Module] +_WhitelistParamNamesCreator = Callable[[nn.Module], list[str]] + + +class BaseOffloader(ABC): + """Base class for model parameter offloading strategies. + + Offloaders control how model parameters are stored and loaded during + inference. Different strategies trade memory for compute/transfer time. + """ + + @abstractmethod + def wrap_modules( + self, + modules_generator: Generator[nn.Module, None, None], + submodule_accessor: _SubmoduleAccessor | None = None, + whitelist_param_names_creator: _WhitelistParamNamesCreator | None = None, + ) -> list[nn.Module]: + """Wrap modules with offloading logic. + + Args: + modules_generator: Generator yielding modules to potentially offload. + submodule_accessor: Optional function to extract a submodule from + each module (e.g., lambda layer: layer.mlp.experts). + whitelist_param_names_creator: Optional function to get parameter + names to offload from a submodule (e.g., ["w13_weight", "w2_weight"]). + + Returns: + List of modules, potentially with offloading hooks installed. + """ + pass + + def post_init(self): + """Called after model construction completes. + + Offloaders can use this to: + - Finalize parameter storage + - Start initial prefetching + - Allocate shared resources + """ + pass + + @property + def forbid_copy_engine_usage(self) -> bool: + """Whether copy engine can be used (affects NCCL operations). + + Some offloading modes may conflict with CUDA copy engine usage + in distributed operations. + """ + return False + + +class NoopOffloader(BaseOffloader): + """No-op offloader that returns modules as-is without any offloading.""" + + def wrap_modules( + self, + modules_generator: Generator[nn.Module, None, None], + submodule_accessor: _SubmoduleAccessor | None = None, + whitelist_param_names_creator: _WhitelistParamNamesCreator | None = None, + ) -> list[nn.Module]: + """Return modules unchanged.""" + return list(modules_generator) + + +# Global singleton offloader instance +_instance: BaseOffloader | None = NoopOffloader() + + +def get_offloader() -> BaseOffloader: + """Get the global offloader instance.""" + assert _instance is not None, "Offloader instance is None" + logger.debug(f"{_instance=}") + return _instance + + +def set_offloader(instance: BaseOffloader) -> None: + """Set the global offloader instance.""" + global _instance + _instance = instance diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py new file mode 100644 index 000000000000..b72b292ac642 --- /dev/null +++ b/vllm/model_executor/offloader/uva.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""UVA-based CPU offloading using Unified Virtual Addressing.""" + +from collections.abc import Callable, Generator + +import torch +import torch.nn as nn +from torch.func import functional_call + +from vllm.model_executor.offloader.base import BaseOffloader +from vllm.utils.platform_utils import is_pin_memory_available, is_uva_available +from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor + + +class UVAOffloader(BaseOffloader): + """Offloader using Unified Virtual Addressing (UVA) for zero-copy access. + + This offloader moves parameters to pinned CPU memory and creates CUDA views + using UVA. The GPU can then directly access the CPU memory without explicit + transfers, at the cost of PCIe bandwidth (slower than GPU memory). + + Args: + cpu_offload_max_bytes: Maximum bytes to offload to CPU. + """ + + def __init__(self, cpu_offload_max_bytes: int): + assert is_uva_available(), "UVA offloading requires UVA (pin memory) support" + + self.cpu_offload_max_bytes = cpu_offload_max_bytes + self.cpu_offload_bytes = 0 + + def wrap_modules( + self, + modules_generator: Generator[nn.Module, None, None], + submodule_accessor: Callable[[nn.Module], nn.Module] | None = None, + whitelist_param_names_creator: Callable[[nn.Module], list[str]] | None = None, + ) -> list[nn.Module]: + """Wrap modules with UVA offloading. + + Note: UVA offloading operates at module level, so submodule_accessor + and whitelist_param_names_creator are ignored. + """ + return [self._maybe_offload_to_cpu(module) for module in modules_generator] + + def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: + """Offload module parameters to CPU using UVA if budget allows.""" + # Check if module has parameters + if (params := next(module.parameters(), None)) is None: + return module + + device = params.device + + # Skip if already on CPU + if device == torch.device("cpu"): + return module + + # Check budget + if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: + return module + + pin_memory = is_pin_memory_available() + + # Offload parameters to pinned CPU memory + offloaded_parameters = False + for p in module.parameters(): + if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: + # Per-parameter offloading: some params may be offloaded, others not + break + + # Create pinned CPU tensor + cpu_data = torch.empty_strided( + size=p.data.size(), + stride=p.data.stride(), + dtype=p.data.dtype, + layout=p.data.layout, + device="cpu", + pin_memory=pin_memory, + ) + cpu_data.copy_(p.data) + + # Keep CPU data alive and create CUDA view via UVA + p._vllm_offloaded_cpu_data = cpu_data + p.data = get_cuda_view_from_cpu_tensor(cpu_data) + + self.cpu_offload_bytes += p.data.numel() * p.data.element_size() + offloaded_parameters = True + + return module + + +# Backward compatibility: Global state for legacy set_cpu_offload_max_bytes() +_CPU_OFFLOAD_BYTES = 0 +_CPU_OFFLOAD_MAX_BYTES = 0 + + +def set_cpu_offload_max_bytes(max_bytes: int) -> None: + """Set maximum bytes to offload for legacy UVA offloading. + + Deprecated: Use UVAOffloader class directly. + """ + global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES + _CPU_OFFLOAD_BYTES = 0 + _CPU_OFFLOAD_MAX_BYTES = max_bytes + + +def maybe_offload_to_cpu(module: nn.Module) -> nn.Module: + """Offload module to CPU using UVA (legacy function). + + Deprecated: Use UVAOffloader class directly. + """ + if (params := next(module.parameters(), None)) is None: + return module + + device = params.device + + if device == torch.device("cpu"): + return module + + global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES + if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: + return module + + pin_memory = is_pin_memory_available() + uva_available = is_uva_available() + + assert uva_available, "V1 CPU offloading requires uva (pin memory) support" + uva_offloading = True + + # offload parameters to CPU + # use pin_memory if possible, which helps cudagraph capture speed + offloaded_parameters = False + for p in module.parameters(): + if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: + # we use per-parameter offloading + # one module might have some parameters offloaded and some not + break + + # `torch.empty_like` does not support `pin_memory` argument + cpu_data = torch.empty_strided( + size=p.data.size(), + stride=p.data.stride(), + dtype=p.data.dtype, + layout=p.data.layout, + device="cpu", + pin_memory=pin_memory, + ) + cpu_data.copy_(p.data) + if not uva_offloading: + p.data = cpu_data + else: + # keep the cpu data alive + p._vllm_offloaded_cpu_data = cpu_data + p.data = get_cuda_view_from_cpu_tensor(cpu_data) + _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size() + offloaded_parameters = True + + if offloaded_parameters and not uva_offloading: + original_forward = module.forward + + def forward(*args, **kwargs): + module.forward = original_forward + device_state = { + # here we blindly call `to(device)` + # if the parameter is already on the device, it will be a no-op + k: v.to(device, non_blocking=True) + for k, v in module.state_dict().items() + } + output = functional_call(module, device_state, args=args, kwargs=kwargs) + module.forward = forward + return output + + module.forward = forward + + return module diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/v2.py new file mode 100644 index 000000000000..fe0fe4f17704 --- /dev/null +++ b/vllm/model_executor/offloader/v2.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from +# https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py +"""OffloaderV2: Advanced CPU offloading with async prefetching.""" + +from abc import ABC, abstractmethod +from collections.abc import Callable, Generator + +import torch +import torch.nn as nn +from torch.func import functional_call + +from vllm.logger import init_logger +from vllm.model_executor.offloader.base import BaseOffloader +from vllm.utils.platform_utils import is_pin_memory_available + +logger = init_logger(__name__) + +# Type aliases +_SubmoduleAccessor = Callable[[nn.Module], nn.Module] +_WhitelistParamNamesCreator = Callable[[nn.Module], list[str]] + + +class OffloaderV2(BaseOffloader): + """Advanced offloader with group-based selection and async prefetching. + + Unlike UVA offloading which provides zero-copy access, V2 explicitly + manages parameter transfers with prefetching to hide latency. + + Args: + group_size: Group every N layers together. + num_in_group: Offload this many layers per group (last N of each group). + prefetch_step: Number of layers to prefetch ahead. + mode: Offload mode ("cpu" is currently supported). + """ + + def __init__( + self, + group_size: int, + num_in_group: int, + prefetch_step: int, + mode: str = "cpu", + ): + self.group_size = group_size + self.num_in_group = num_in_group + self.prefetch_step = prefetch_step + self.mode = mode + self.alt_stream = torch.cuda.Stream() + self.module_offloaders: list[_ModuleOffloader] = [] + self.total_offloaded_bytes = 0 + + def wrap_modules( + self, + modules_generator: Generator[nn.Module, None, None], + submodule_accessor: _SubmoduleAccessor | None = None, + whitelist_param_names_creator: _WhitelistParamNamesCreator | None = None, + ) -> list[nn.Module]: + """Wrap modules with V2 offloading and prefetching logic.""" + assert len(self.module_offloaders) == 0, ( + "wrap_modules should only be called once" + ) + + all_modules = [] + offload_submodules = [] + + for module_index, module in enumerate(modules_generator): + all_modules.append(module) + + # Select layers to offload based on group pattern + # Offload last num_in_group layers of each group_size + if module_index % self.group_size >= self.group_size - self.num_in_group: + submodule = submodule_accessor(module) if submodule_accessor else module + whitelist_param_names = ( + whitelist_param_names_creator(submodule) + if whitelist_param_names_creator + else [name for name, _ in submodule.named_parameters()] + ) + + offload_submodules.append(submodule) + self.module_offloaders.append( + _ModuleOffloader( + mode=self.mode, + module=submodule, + alt_stream=self.alt_stream, + whitelist_param_names=whitelist_param_names, + ) + ) + + # Hook forward passes for all offloaded submodules + for index, submodule in enumerate(offload_submodules): + self._hook_module_forward(index, submodule) + + return all_modules + + def _hook_module_forward(self, index: int, module: nn.Module): + """Hook module's forward to implement prefetch + execute + offload pattern.""" + original_forward = module.forward + + def forward(*args, **kwargs): + module.forward = original_forward + device_tensors = self.module_offloaders[index].wait_and_get_device_tensors() + output = functional_call(module, device_tensors, args=args, kwargs=kwargs) + next_index = (index + self.prefetch_step) % len(self.module_offloaders) + self.module_offloaders[next_index].start_onload() + self.module_offloaders[index].offload() + module.forward = forward + return output + + module.forward = forward + + def post_init(self): + """Initialize offloaders and start prefetching first N modules.""" + for offloader in self.module_offloaders: + offloader.post_init() + self.total_offloaded_bytes += offloader.offloaded_bytes + + logger.info_once( + f"[OffloaderV2] Initialized {len(self.module_offloaders)} modules. " + f"Total GPU memory saved: {self.total_offloaded_bytes / 1e9:.4f} GB " + f"(group_size={self.group_size}, num_in_group={self.num_in_group}, " + f"prefetch_step={self.prefetch_step}, mode={self.mode})" + ) + + for i in range(min(self.prefetch_step, len(self.module_offloaders))): + self.module_offloaders[i].start_onload() + + @property + def forbid_copy_engine_usage(self) -> bool: + """CPU mode may conflict with copy engine in some scenarios.""" + return self.mode == "cpu" + + +class _ModuleOffloader: + """Manages offloading for a single module. + + Responsibilities: + - Create parameter offloaders for each parameter + - Coordinate async loading via alternate CUDA stream + - Provide device tensors when needed + """ + + def __init__( + self, + mode: str, + module: nn.Module, + alt_stream: torch.cuda.Stream, + whitelist_param_names: list[str], + ): + self.mode = mode + self.module = module + self.device = next(module.parameters()).device + self.alt_stream = alt_stream + self.offloaded_bytes = 0 + + assert self.device != torch.device("cpu"), ( + "Module parameters should not already be on CPU " + "(offloader handles CPU placement)" + ) + + self._device_tensors: dict[str, torch.Tensor] | None = None + self._load_event: torch.cuda.Event | None = None + + param_dict = dict(self.module.named_parameters()) + assert all(name in param_dict for name in whitelist_param_names), ( + f"Whitelist params {whitelist_param_names} not found in module params " + f"{list(param_dict.keys())}" + ) + + self._param_offloaders = { + name: _BaseParamOffloader.create(mode, module=module, param_name=name) + for name in whitelist_param_names + } + + def post_init(self): + """Collect total offloaded bytes (offloading already done in __init__).""" + for param_offloader in self._param_offloaders.values(): + param_offloader.post_init() + self.offloaded_bytes += param_offloader.offloaded_bytes + + def start_onload(self): + """Start async loading in alternate CUDA stream.""" + # Synchronize with main stream before starting + self.alt_stream.wait_stream(torch.cuda.current_stream()) + + with torch.cuda.stream(self.alt_stream): + # Load all parameters to device + self._device_tensors = { + name: offloader.create_device_tensor() + for name, offloader in self._param_offloaders.items() + } + # Record completion event + self._load_event = torch.cuda.Event() + self._load_event.record() + + def offload(self): + """Free device tensors (offload from GPU memory).""" + self._device_tensors = None + self._load_event = None + + def wait_and_get_device_tensors(self) -> dict[str, torch.Tensor]: + """Wait for async loading to complete and return device tensors.""" + assert self._device_tensors is not None, ( + "Tensors not loaded (call start_onload first)" + ) + # Wait for loading event to complete + if self._load_event is not None: + self._load_event.wait() + return self._device_tensors + + +class _BaseParamOffloader(ABC): + """Base class for parameter offloading strategies.""" + + @staticmethod + def create(mode: str, **kwargs) -> "_BaseParamOffloader": + """Factory method to create appropriate offloader for mode.""" + if mode == "cpu": + return _CpuParamOffloader(**kwargs) + else: + raise ValueError(f"Unknown offload mode: {mode}") + + def __init__(self, module: nn.Module, param_name: str): + self._module = module + self._param_name = param_name + self.offloaded_bytes = 0 + + @property + def _param(self) -> nn.Parameter: + """Get the parameter being offloaded.""" + return getattr(self._module, self._param_name) + + def post_init(self): + """Initialize offloading (move parameter to storage).""" + return + + @abstractmethod + def create_device_tensor(self) -> torch.Tensor: + """Create device tensor from offloaded storage.""" + pass + + +class _CpuParamOffloader(_BaseParamOffloader): + """Offload parameter to pinned CPU memory.""" + + def __init__(self, module: nn.Module, param_name: str): + super().__init__(module, param_name) + + # Offload immediately to free GPU memory by moving param.data to CPU + self._move_param_to_cpu() + + def _move_param_to_cpu(self): + """Move parameter data to pinned CPU memory (modify param.data in-place).""" + param = self._param + pin_memory = is_pin_memory_available() + + # Calculate memory size + self.offloaded_bytes = param.data.numel() * param.data.element_size() + + # Create pinned CPU tensor with same layout + cpu_data = torch.empty_strided( + size=param.data.size(), + stride=param.data.stride(), + dtype=param.data.dtype, + layout=param.data.layout, + device="cpu", + pin_memory=pin_memory, + ) + cpu_data.copy_(param.data) + + logger.debug_once( + f"[OffloaderV2] Offloaded parameter '{self._param_name}': " + f"shape={tuple(param.shape)}, dtype={param.dtype}, " + f"size={self.offloaded_bytes / 1e9:.6f} GB, pinned={pin_memory}" + ) + + # Modify param.data in-place to point to CPU memory + # This keeps the parameter in the module but with CPU data + param.data = cpu_data + + def post_init(self): + """No-op: offloading already done in __init__.""" + pass + + def create_device_tensor(self) -> torch.Tensor: + """Load from CPU to GPU (async if pinned). + + Returns a CUDA copy of the parameter (which has CPU data). + """ + return self._param.to("cuda", non_blocking=True) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 1b250a8bd009..f849f4aa2837 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -278,9 +278,38 @@ def __init__( self.speculative_config = vllm_config.speculative_config self.observability_config = vllm_config.observability_config + # Set up offloader based on configuration + # For backward compatibility, still support legacy set_cpu_offload_max_bytes from vllm.model_executor.models.utils import set_cpu_offload_max_bytes + from vllm.model_executor.offloader import ( + NoopOffloader, + OffloaderV2, + UVAOffloader, + set_offloader, + ) - set_cpu_offload_max_bytes(int(self.cache_config.cpu_offload_gb * 1024**3)) + # Priority: V2 offloading if configured, else UVA, else noop + if self.cache_config.offload_group_size > 0: + # Use V2 offloading + offloader = OffloaderV2( + group_size=self.cache_config.offload_group_size, + num_in_group=self.cache_config.offload_num_in_group, + prefetch_step=self.cache_config.offload_prefetch_step, + mode="cpu", + ) + set_offloader(offloader) + elif self.cache_config.cpu_offload_gb > 0: + # Use UVA offloading (legacy) + offloader = UVAOffloader( + cpu_offload_max_bytes=int(self.cache_config.cpu_offload_gb * 1024**3) + ) + set_offloader(offloader) + # Also set legacy global state for backward compatibility + set_cpu_offload_max_bytes(int(self.cache_config.cpu_offload_gb * 1024**3)) + else: + # No offloading + set_offloader(NoopOffloader()) + set_cpu_offload_max_bytes(0) model_config = self.model_config cache_config = self.cache_config @@ -3614,6 +3643,11 @@ def load_model(self, eep_scale_up: bool = False) -> None: self.model, self.vllm_config, CUDAGraphMode.NONE, self.device ) + # Initialize offloader after model is loaded + from vllm.model_executor.offloader import get_offloader + + get_offloader().post_init() + def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None: """Extract Eagle3 auxiliary layer indices from speculative config. From 9bf409705f640ada5dfecd1afd3f31e8dc42d967 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 2 Dec 2025 22:01:41 -0800 Subject: [PATCH 02/36] cleanup Signed-off-by: Ming Yang --- vllm/model_executor/models/utils.py | 82 --------------- vllm/model_executor/offloader/__init__.py | 2 +- vllm/model_executor/offloader/base.py | 27 ++--- vllm/model_executor/offloader/uva.py | 119 +++++++--------------- vllm/model_executor/offloader/v2.py | 19 +--- 5 files changed, 52 insertions(+), 197 deletions(-) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 6d0aefbd8122..cfb8faead401 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -8,7 +8,6 @@ import torch import torch.nn as nn -from torch.func import functional_call from transformers import PretrainedConfig from vllm.config import VllmConfig @@ -30,11 +29,9 @@ from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import ( is_pin_memory_available, - is_uva_available, ) from vllm.utils.torch_utils import ( direct_register_custom_op, - get_cuda_view_from_cpu_tensor, ) logger = init_logger(__name__) @@ -509,83 +506,6 @@ def forward(self, *args, **kwargs): return args[0] if args else next(iter(kwargs.values())) -_CPU_OFFLOAD_BYTES = 0 -_CPU_OFFLOAD_MAX_BYTES = 0 - - -def set_cpu_offload_max_bytes(max_bytes: int) -> None: - global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES - _CPU_OFFLOAD_BYTES = 0 - _CPU_OFFLOAD_MAX_BYTES = max_bytes - - -def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module: - if (params := next(module.parameters(), None)) is None: - return module - - device = params.device - - if device == torch.device("cpu"): - return module - - global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: - return module - - pin_memory = is_pin_memory_available() - uva_available = is_uva_available() - - assert uva_available, "V1 CPU offloading requires uva (pin memory) support" - uva_offloading = True - - # offload parameters to CPU - # use pin_memory if possible, which helps cudagraph capture speed - offloaded_parameters = False - for p in module.parameters(): - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: - # we use per-parameter offloading - # one module might have some parameters offloaded and some not - break - - # `torch.empty_like` does not support `pin_memory` argument - cpu_data = torch.empty_strided( - size=p.data.size(), - stride=p.data.stride(), - dtype=p.data.dtype, - layout=p.data.layout, - device="cpu", - pin_memory=pin_memory, - ) - cpu_data.copy_(p.data) - if not uva_offloading: - p.data = cpu_data - else: - # keep the cpu data alive - p._vllm_offloaded_cpu_data = cpu_data - p.data = get_cuda_view_from_cpu_tensor(cpu_data) - _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size() - offloaded_parameters = True - - if offloaded_parameters and not uva_offloading: - original_forward = module.forward - - def forward(*args, **kwargs): - module.forward = original_forward - device_state = { - # here we blindly call `to(device)` - # if the parameter is already on the device, it will be a no-op - k: v.to(device, non_blocking=True) - for k, v in module.state_dict().items() - } - output = functional_call(module, device_state, args=args, kwargs=kwargs) - module.forward = forward - return output - - module.forward = forward - - return module - - def make_layers( num_hidden_layers: int, layer_fn: LayerFn, @@ -613,8 +533,6 @@ def make_layers( num_hidden_layers, get_pp_group().rank_in_group, get_pp_group().world_size ) - logger.debug(f"{offloader_kwargs=}") - modules = torch.nn.ModuleList( [PPMissingLayer() for _ in range(start_layer)] + get_offloader().wrap_modules( diff --git a/vllm/model_executor/offloader/__init__.py b/vllm/model_executor/offloader/__init__.py index fa29ce76b6e6..b02da6de1a9c 100644 --- a/vllm/model_executor/offloader/__init__.py +++ b/vllm/model_executor/offloader/__init__.py @@ -8,8 +8,8 @@ get_offloader, set_offloader, ) +from vllm.model_executor.offloader.offloader_v2 import OffloaderV2 from vllm.model_executor.offloader.uva import UVAOffloader -from vllm.model_executor.offloader.v2 import OffloaderV2 __all__ = [ "BaseOffloader", diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index 8560c52e47b4..c43feedd4b33 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from +# https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py """Base classes for model parameter offloading.""" from abc import ABC, abstractmethod @@ -11,11 +13,22 @@ logger = init_logger(__name__) -# Type aliases for clarity _SubmoduleAccessor = Callable[[nn.Module], nn.Module] _WhitelistParamNamesCreator = Callable[[nn.Module], list[str]] +""" +class relation: + +BaseOffloader (ABC) + * implemented by: UVAOffloader + * implemented by: OffloaderV2 + * uses: _ModuleOffloader + * uses: _BaseParamOffloader (ABC) + * implemented by: _CpuParamOffloader +""" + + class BaseOffloader(ABC): """Base class for model parameter offloading strategies. @@ -52,16 +65,7 @@ def post_init(self): - Start initial prefetching - Allocate shared resources """ - pass - - @property - def forbid_copy_engine_usage(self) -> bool: - """Whether copy engine can be used (affects NCCL operations). - - Some offloading modes may conflict with CUDA copy engine usage - in distributed operations. - """ - return False + return class NoopOffloader(BaseOffloader): @@ -84,7 +88,6 @@ def wrap_modules( def get_offloader() -> BaseOffloader: """Get the global offloader instance.""" assert _instance is not None, "Offloader instance is None" - logger.debug(f"{_instance=}") return _instance diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index b72b292ac642..59890c5baa09 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -45,30 +45,34 @@ def wrap_modules( def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: """Offload module parameters to CPU using UVA if budget allows.""" - # Check if module has parameters if (params := next(module.parameters(), None)) is None: return module device = params.device - # Skip if already on CPU if device == torch.device("cpu"): return module - # Check budget - if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: + global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES + if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: return module pin_memory = is_pin_memory_available() + uva_available = is_uva_available() + + assert uva_available, "V1 CPU offloading requires uva (pin memory) support" + uva_offloading = True - # Offload parameters to pinned CPU memory + # offload parameters to CPU + # use pin_memory if possible, which helps cudagraph capture speed offloaded_parameters = False for p in module.parameters(): - if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: - # Per-parameter offloading: some params may be offloaded, others not + if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: + # we use per-parameter offloading + # one module might have some parameters offloaded and some not break - # Create pinned CPU tensor + # `torch.empty_like` does not support `pin_memory` argument cpu_data = torch.empty_strided( size=p.data.size(), stride=p.data.stride(), @@ -78,13 +82,31 @@ def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: pin_memory=pin_memory, ) cpu_data.copy_(p.data) + if not uva_offloading: + p.data = cpu_data + else: + # keep the cpu data alive + p._vllm_offloaded_cpu_data = cpu_data + p.data = get_cuda_view_from_cpu_tensor(cpu_data) + _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size() + offloaded_parameters = True - # Keep CPU data alive and create CUDA view via UVA - p._vllm_offloaded_cpu_data = cpu_data - p.data = get_cuda_view_from_cpu_tensor(cpu_data) + if offloaded_parameters and not uva_offloading: + original_forward = module.forward + + def forward(*args, **kwargs): + module.forward = original_forward + device_state = { + # here we blindly call `to(device)` + # if the parameter is already on the device, it will be a no-op + k: v.to(device, non_blocking=True) + for k, v in module.state_dict().items() + } + output = functional_call(module, device_state, args=args, kwargs=kwargs) + module.forward = forward + return output - self.cpu_offload_bytes += p.data.numel() * p.data.element_size() - offloaded_parameters = True + module.forward = forward return module @@ -102,74 +124,3 @@ def set_cpu_offload_max_bytes(max_bytes: int) -> None: global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES _CPU_OFFLOAD_BYTES = 0 _CPU_OFFLOAD_MAX_BYTES = max_bytes - - -def maybe_offload_to_cpu(module: nn.Module) -> nn.Module: - """Offload module to CPU using UVA (legacy function). - - Deprecated: Use UVAOffloader class directly. - """ - if (params := next(module.parameters(), None)) is None: - return module - - device = params.device - - if device == torch.device("cpu"): - return module - - global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: - return module - - pin_memory = is_pin_memory_available() - uva_available = is_uva_available() - - assert uva_available, "V1 CPU offloading requires uva (pin memory) support" - uva_offloading = True - - # offload parameters to CPU - # use pin_memory if possible, which helps cudagraph capture speed - offloaded_parameters = False - for p in module.parameters(): - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: - # we use per-parameter offloading - # one module might have some parameters offloaded and some not - break - - # `torch.empty_like` does not support `pin_memory` argument - cpu_data = torch.empty_strided( - size=p.data.size(), - stride=p.data.stride(), - dtype=p.data.dtype, - layout=p.data.layout, - device="cpu", - pin_memory=pin_memory, - ) - cpu_data.copy_(p.data) - if not uva_offloading: - p.data = cpu_data - else: - # keep the cpu data alive - p._vllm_offloaded_cpu_data = cpu_data - p.data = get_cuda_view_from_cpu_tensor(cpu_data) - _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size() - offloaded_parameters = True - - if offloaded_parameters and not uva_offloading: - original_forward = module.forward - - def forward(*args, **kwargs): - module.forward = original_forward - device_state = { - # here we blindly call `to(device)` - # if the parameter is already on the device, it will be a no-op - k: v.to(device, non_blocking=True) - for k, v in module.state_dict().items() - } - output = functional_call(module, device_state, args=args, kwargs=kwargs) - module.forward = forward - return output - - module.forward = forward - - return module diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/v2.py index fe0fe4f17704..8c69a7248e10 100644 --- a/vllm/model_executor/offloader/v2.py +++ b/vllm/model_executor/offloader/v2.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from # https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py -"""OffloaderV2: Advanced CPU offloading with async prefetching.""" +"""OffloaderV2: CPU offloading with async prefetching.""" from abc import ABC, abstractmethod from collections.abc import Callable, Generator @@ -17,7 +17,6 @@ logger = init_logger(__name__) -# Type aliases _SubmoduleAccessor = Callable[[nn.Module], nn.Module] _WhitelistParamNamesCreator = Callable[[nn.Module], list[str]] @@ -87,7 +86,6 @@ def wrap_modules( ) ) - # Hook forward passes for all offloaded submodules for index, submodule in enumerate(offload_submodules): self._hook_module_forward(index, submodule) @@ -125,11 +123,6 @@ def post_init(self): for i in range(min(self.prefetch_step, len(self.module_offloaders))): self.module_offloaders[i].start_onload() - @property - def forbid_copy_engine_usage(self) -> bool: - """CPU mode may conflict with copy engine in some scenarios.""" - return self.mode == "cpu" - class _ModuleOffloader: """Manages offloading for a single module. @@ -180,16 +173,13 @@ def post_init(self): def start_onload(self): """Start async loading in alternate CUDA stream.""" - # Synchronize with main stream before starting self.alt_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.alt_stream): - # Load all parameters to device self._device_tensors = { name: offloader.create_device_tensor() for name, offloader in self._param_offloaders.items() } - # Record completion event self._load_event = torch.cuda.Event() self._load_event.record() @@ -203,7 +193,6 @@ def wait_and_get_device_tensors(self) -> dict[str, torch.Tensor]: assert self._device_tensors is not None, ( "Tensors not loaded (call start_onload first)" ) - # Wait for loading event to complete if self._load_event is not None: self._load_event.wait() return self._device_tensors @@ -245,8 +234,6 @@ class _CpuParamOffloader(_BaseParamOffloader): def __init__(self, module: nn.Module, param_name: str): super().__init__(module, param_name) - - # Offload immediately to free GPU memory by moving param.data to CPU self._move_param_to_cpu() def _move_param_to_cpu(self): @@ -254,10 +241,8 @@ def _move_param_to_cpu(self): param = self._param pin_memory = is_pin_memory_available() - # Calculate memory size self.offloaded_bytes = param.data.numel() * param.data.element_size() - # Create pinned CPU tensor with same layout cpu_data = torch.empty_strided( size=param.data.size(), stride=param.data.stride(), @@ -274,8 +259,6 @@ def _move_param_to_cpu(self): f"size={self.offloaded_bytes / 1e9:.6f} GB, pinned={pin_memory}" ) - # Modify param.data in-place to point to CPU memory - # This keeps the parameter in the module but with CPU data param.data = cpu_data def post_init(self): From 152af73378b5e81117183c23c77e8a3df54405d0 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 2 Dec 2025 22:28:40 -0800 Subject: [PATCH 03/36] cleanup: extract func; move imports Signed-off-by: Ming Yang --- vllm/model_executor/offloader/__init__.py | 2 ++ vllm/model_executor/offloader/base.py | 30 ++++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 44 +++++------------------ 3 files changed, 40 insertions(+), 36 deletions(-) diff --git a/vllm/model_executor/offloader/__init__.py b/vllm/model_executor/offloader/__init__.py index b02da6de1a9c..81cf93e62989 100644 --- a/vllm/model_executor/offloader/__init__.py +++ b/vllm/model_executor/offloader/__init__.py @@ -5,6 +5,7 @@ from vllm.model_executor.offloader.base import ( BaseOffloader, NoopOffloader, + create_offloader, get_offloader, set_offloader, ) @@ -16,6 +17,7 @@ "NoopOffloader", "UVAOffloader", "OffloaderV2", + "create_offloader", "get_offloader", "set_offloader", ] diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index c43feedd4b33..5643beda44a3 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -6,11 +6,15 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Generator +from typing import TYPE_CHECKING import torch.nn as nn from vllm.logger import init_logger +if TYPE_CHECKING: + from vllm.config import CacheConfig + logger = init_logger(__name__) _SubmoduleAccessor = Callable[[nn.Module], nn.Module] @@ -95,3 +99,29 @@ def set_offloader(instance: BaseOffloader) -> None: """Set the global offloader instance.""" global _instance _instance = instance + + +def create_offloader(cache_config: "CacheConfig") -> BaseOffloader: + """Create an offloader based on the cache configuration. + + Priority: V2 offloading if configured, else UVA, else noop. + """ + from vllm.model_executor.offloader.offloader_v2 import OffloaderV2 + from vllm.model_executor.offloader.uva import UVAOffloader + + if cache_config.offload_group_size > 0: + # Use V2 offloading + return OffloaderV2( + group_size=cache_config.offload_group_size, + num_in_group=cache_config.offload_num_in_group, + prefetch_step=cache_config.offload_prefetch_step, + mode="cpu", + ) + elif cache_config.cpu_offload_gb > 0: + # Use UVA offloading (legacy) + return UVAOffloader( + cpu_offload_max_bytes=int(cache_config.cpu_offload_gb * 1024**3) + ) + else: + # No offloading + return NoopOffloader() diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index f849f4aa2837..a86dd0937786 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -72,6 +72,11 @@ is_pooling_model, is_text_generation_model, ) +from vllm.model_executor.offloader import ( + create_offloader, + get_offloader, + set_offloader, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( BatchedTensorInputs, @@ -278,39 +283,6 @@ def __init__( self.speculative_config = vllm_config.speculative_config self.observability_config = vllm_config.observability_config - # Set up offloader based on configuration - # For backward compatibility, still support legacy set_cpu_offload_max_bytes - from vllm.model_executor.models.utils import set_cpu_offload_max_bytes - from vllm.model_executor.offloader import ( - NoopOffloader, - OffloaderV2, - UVAOffloader, - set_offloader, - ) - - # Priority: V2 offloading if configured, else UVA, else noop - if self.cache_config.offload_group_size > 0: - # Use V2 offloading - offloader = OffloaderV2( - group_size=self.cache_config.offload_group_size, - num_in_group=self.cache_config.offload_num_in_group, - prefetch_step=self.cache_config.offload_prefetch_step, - mode="cpu", - ) - set_offloader(offloader) - elif self.cache_config.cpu_offload_gb > 0: - # Use UVA offloading (legacy) - offloader = UVAOffloader( - cpu_offload_max_bytes=int(self.cache_config.cpu_offload_gb * 1024**3) - ) - set_offloader(offloader) - # Also set legacy global state for backward compatibility - set_cpu_offload_max_bytes(int(self.cache_config.cpu_offload_gb * 1024**3)) - else: - # No offloading - set_offloader(NoopOffloader()) - set_cpu_offload_max_bytes(0) - model_config = self.model_config cache_config = self.cache_config scheduler_config = self.scheduler_config @@ -627,6 +599,9 @@ def __init__( self.execute_model_state: ExecuteModelState | None = None self.kv_connector_output: KVConnectorOutput | None = None + # Model weight offloader + set_offloader(create_offloader(self.cache_config)) + def reset_mm_cache(self) -> None: if self.mm_budget: self.mm_budget.reset_cache() @@ -3643,9 +3618,6 @@ def load_model(self, eep_scale_up: bool = False) -> None: self.model, self.vllm_config, CUDAGraphMode.NONE, self.device ) - # Initialize offloader after model is loaded - from vllm.model_executor.offloader import get_offloader - get_offloader().post_init() def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None: From 67cd6cc0a453e6c830bf3af67c342f3f88834a08 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 2 Dec 2025 23:52:30 -0800 Subject: [PATCH 04/36] fix import Signed-off-by: Ming Yang --- vllm/model_executor/offloader/__init__.py | 2 +- vllm/model_executor/offloader/base.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/offloader/__init__.py b/vllm/model_executor/offloader/__init__.py index 81cf93e62989..a8031cb55ff4 100644 --- a/vllm/model_executor/offloader/__init__.py +++ b/vllm/model_executor/offloader/__init__.py @@ -9,8 +9,8 @@ get_offloader, set_offloader, ) -from vllm.model_executor.offloader.offloader_v2 import OffloaderV2 from vllm.model_executor.offloader.uva import UVAOffloader +from vllm.model_executor.offloader.v2 import OffloaderV2 __all__ = [ "BaseOffloader", diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index 5643beda44a3..402131084b56 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -106,8 +106,8 @@ def create_offloader(cache_config: "CacheConfig") -> BaseOffloader: Priority: V2 offloading if configured, else UVA, else noop. """ - from vllm.model_executor.offloader.offloader_v2 import OffloaderV2 from vllm.model_executor.offloader.uva import UVAOffloader + from vllm.model_executor.offloader.v2 import OffloaderV2 if cache_config.offload_group_size > 0: # Use V2 offloading From f43b5776e645bec1874045ab6abb55fc1e9b383b Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Wed, 3 Dec 2025 22:11:13 -0800 Subject: [PATCH 05/36] address comment: remove legacy code Signed-off-by: Ming Yang --- vllm/model_executor/offloader/uva.py | 22 +++------------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index 59890c5baa09..d1041a62500a 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -53,8 +53,7 @@ def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: if device == torch.device("cpu"): return module - global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: + if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: return module pin_memory = is_pin_memory_available() @@ -67,7 +66,7 @@ def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: # use pin_memory if possible, which helps cudagraph capture speed offloaded_parameters = False for p in module.parameters(): - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: + if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: # we use per-parameter offloading # one module might have some parameters offloaded and some not break @@ -88,7 +87,7 @@ def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: # keep the cpu data alive p._vllm_offloaded_cpu_data = cpu_data p.data = get_cuda_view_from_cpu_tensor(cpu_data) - _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size() + self.cpu_offload_bytes += p.data.numel() * p.data.element_size() offloaded_parameters = True if offloaded_parameters and not uva_offloading: @@ -109,18 +108,3 @@ def forward(*args, **kwargs): module.forward = forward return module - - -# Backward compatibility: Global state for legacy set_cpu_offload_max_bytes() -_CPU_OFFLOAD_BYTES = 0 -_CPU_OFFLOAD_MAX_BYTES = 0 - - -def set_cpu_offload_max_bytes(max_bytes: int) -> None: - """Set maximum bytes to offload for legacy UVA offloading. - - Deprecated: Use UVAOffloader class directly. - """ - global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES - _CPU_OFFLOAD_BYTES = 0 - _CPU_OFFLOAD_MAX_BYTES = max_bytes From 0ed35705f65f1d6b9a5387ad64e2c8b82127bacd Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sat, 20 Dec 2025 10:01:23 -0800 Subject: [PATCH 06/36] refactor: move offload config from CacheConfig to OffloadConfig Address review comments: 1. Add validation that offload_num_in_group <= offload_group_size 2. Move offload settings out of CacheConfig (not KV cache related) Changes: - Create new OffloadConfig in vllm/config/offload.py with all model weight offloading settings (cpu_offload_gb, offload_group_size, offload_num_in_group, offload_prefetch_step) - Add pydantic validator to ensure offload_num_in_group <= offload_group_size - Add offload_config field to VllmConfig - Update arg_utils.py to create OffloadConfig and use separate arg group - Update offloader/base.py to accept OffloadConfig instead of CacheConfig - Update gpu_model_runner.py to use offload_config - Mark cpu_offload_gb in CacheConfig as deprecated (kept for backward compat) Signed-off-by: Ming Yang --- vllm/config/__init__.py | 3 + vllm/config/cache.py | 13 +---- vllm/config/offload.py | 80 +++++++++++++++++++++++++++ vllm/config/vllm.py | 7 +++ vllm/engine/arg_utils.py | 51 +++++++++++------ vllm/model_executor/offloader/base.py | 18 +++--- vllm/v1/worker/gpu_model_runner.py | 5 +- 7 files changed, 138 insertions(+), 39 deletions(-) create mode 100644 vllm/config/offload.py diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 0e91dd57420a..7ccc7db2874d 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -22,6 +22,7 @@ ) from vllm.config.multimodal import MultiModalConfig from vllm.config.observability import ObservabilityConfig +from vllm.config.offload import OffloadConfig from vllm.config.parallel import EPLBConfig, ParallelConfig from vllm.config.pooler import PoolerConfig from vllm.config.profiler import ProfilerConfig @@ -77,6 +78,8 @@ "MultiModalConfig", # From vllm.config.observability "ObservabilityConfig", + # From vllm.config.offload + "OffloadConfig", # From vllm.config.parallel "EPLBConfig", "ParallelConfig", diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 2a0c948c0b7b..57af05ca3c4c 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -100,17 +100,10 @@ class CacheConfig: load a 13B model with BF16 weight, which requires at least 26GB GPU memory. Note that this requires fast CPU-GPU interconnect, as part of the model is loaded from CPU memory to GPU memory on the fly in each model forward pass. + + DEPRECATED: This field is deprecated and will be removed in a future + release. Please use OffloadConfig.cpu_offload_gb instead. """ - offload_group_size: int = Field(default=0, ge=0) - """Advanced CPU offloading: Group every N layers together. Offload last - `offload_num_in_group` layers of each group. Default is 0 (disabled). - Example: group_size=8, num_in_group=2 offloads layers 6,7,14,15,22,23,... - """ - offload_num_in_group: int = Field(default=1, ge=1) - """Advanced CPU offloading: Number of layers to offload per group. Default is 1.""" - offload_prefetch_step: int = Field(default=1, ge=0) - """Advanced CPU offloading: Number of layers to prefetch ahead. Higher values hide - more latency but use more GPU memory. Default is 1.""" calculate_kv_scales: bool = False """This enables dynamic calculation of `k_scale` and `v_scale` when kv_cache_dtype is fp8. If `False`, the scales will be loaded from the model diff --git a/vllm/config/offload.py b/vllm/config/offload.py new file mode 100644 index 000000000000..9d7f3950f19a --- /dev/null +++ b/vllm/config/offload.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Configuration for model weight offloading.""" + +from typing import Any + +from pydantic import Field, model_validator +from pydantic.dataclasses import dataclass + +from vllm.config.utils import config +from vllm.utils.hashing import safe_hash + + +@config +@dataclass +class OffloadConfig: + """Configuration for model weight offloading to CPU. + + This controls how model parameters are offloaded to CPU memory to reduce + GPU memory usage, at the cost of additional CPU-GPU transfers during + inference. + """ + + cpu_offload_gb: float = Field(default=0, ge=0) + """The space in GiB to offload to CPU, per GPU. Default is 0, which means + no offloading. Intuitively, this argument can be seen as a virtual way to + increase the GPU memory size. For example, if you have one 24 GB GPU and + set this to 10, virtually you can think of it as a 34 GB GPU. Then you can + load a 13B model with BF16 weight, which requires at least 26GB GPU memory. + Note that this requires fast CPU-GPU interconnect, as part of the model is + loaded from CPU memory to GPU memory on the fly in each model forward pass. + This uses UVA (Unified Virtual Addressing) for zero-copy access. + """ + + offload_group_size: int = Field(default=0, ge=0) + """Advanced CPU offloading (V2): Group every N layers together. Offload last + `offload_num_in_group` layers of each group. Default is 0 (disabled). + Example: group_size=8, num_in_group=2 offloads layers 6,7,14,15,22,23,... + Unlike cpu_offload_gb, this uses explicit async prefetching to hide transfer + latency. + """ + + offload_num_in_group: int = Field(default=1, ge=1) + """Advanced CPU offloading (V2): Number of layers to offload per group. + Must be <= offload_group_size. Default is 1.""" + + offload_prefetch_step: int = Field(default=1, ge=0) + """Advanced CPU offloading (V2): Number of layers to prefetch ahead. + Higher values hide more latency but use more GPU memory. Default is 1.""" + + @model_validator(mode="after") + def validate_offload_config(self) -> "OffloadConfig": + """Validate that offload_num_in_group <= offload_group_size.""" + if ( + self.offload_group_size > 0 + and self.offload_num_in_group > self.offload_group_size + ): + raise ValueError( + f"offload_num_in_group ({self.offload_num_in_group}) must be " + f"<= offload_group_size ({self.offload_group_size})" + ) + return self + + def compute_hash(self) -> str: + """ + WARNING: Whenever a new field is added to this config, + ensure that it is included in the factors list if + it affects the computation graph. + + Provide a hash that uniquely identifies all the configs + that affect the structure of the computation + graph from input ids/embeddings to the final hidden states, + excluding anything before input ids/embeddings and after + the final hidden states. + """ + # Offload settings don't affect the computation graph structure, + # only the memory layout and transfer patterns. + factors: list[Any] = [] + hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() + return hash_str diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 0439dc52e7e6..0571cd1dc450 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -38,6 +38,7 @@ from .lora import LoRAConfig from .model import ModelConfig from .observability import ObservabilityConfig +from .offload import OffloadConfig from .parallel import ParallelConfig from .profiler import ProfilerConfig from .scheduler import SchedulerConfig @@ -194,6 +195,8 @@ class VllmConfig: """Device configuration.""" load_config: LoadConfig = Field(default_factory=LoadConfig) """Load configuration.""" + offload_config: OffloadConfig = Field(default_factory=OffloadConfig) + """Model weight offloading configuration.""" attention_config: AttentionConfig = Field(default_factory=AttentionConfig) """Attention configuration.""" lora_config: LoRAConfig | None = None @@ -285,6 +288,10 @@ def compute_hash(self) -> str: vllm_factors.append(self.load_config.compute_hash()) else: vllm_factors.append("None") + if self.offload_config: + vllm_factors.append(self.offload_config.compute_hash()) + else: + vllm_factors.append("None") if self.attention_config: vllm_factors.append(self.attention_config.compute_hash()) else: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 865ca0ff27ce..a55fc8101255 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -48,6 +48,7 @@ ModelConfig, MultiModalConfig, ObservabilityConfig, + OffloadConfig, ParallelConfig, PoolerConfig, ProfilerConfig, @@ -434,10 +435,10 @@ class EngineArgs: disable_sliding_window: bool = ModelConfig.disable_sliding_window disable_cascade_attn: bool = ModelConfig.disable_cascade_attn swap_space: float = CacheConfig.swap_space - cpu_offload_gb: float = CacheConfig.cpu_offload_gb - offload_group_size: int = CacheConfig.offload_group_size - offload_num_in_group: int = CacheConfig.offload_num_in_group - offload_prefetch_step: int = CacheConfig.offload_prefetch_step + cpu_offload_gb: float = OffloadConfig.cpu_offload_gb + offload_group_size: int = OffloadConfig.offload_group_size + offload_num_in_group: int = OffloadConfig.offload_num_in_group + offload_prefetch_step: int = OffloadConfig.offload_prefetch_step gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None @@ -915,16 +916,6 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: cache_group.add_argument( "--prefix-caching-hash-algo", **cache_kwargs["prefix_caching_hash_algo"] ) - cache_group.add_argument("--cpu-offload-gb", **cache_kwargs["cpu_offload_gb"]) - cache_group.add_argument( - "--offload-group-size", **cache_kwargs["offload_group_size"] - ) - cache_group.add_argument( - "--offload-num-in-group", **cache_kwargs["offload_num_in_group"] - ) - cache_group.add_argument( - "--offload-prefetch-step", **cache_kwargs["offload_prefetch_step"] - ) cache_group.add_argument( "--calculate-kv-scales", **cache_kwargs["calculate_kv_scales"] ) @@ -947,6 +938,25 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--kv-offloading-backend", **cache_kwargs["kv_offloading_backend"] ) + # Model weight offload related configs + offload_kwargs = get_kwargs(OffloadConfig) + offload_group = parser.add_argument_group( + title="OffloadConfig", + description=OffloadConfig.__doc__, + ) + offload_group.add_argument( + "--cpu-offload-gb", **offload_kwargs["cpu_offload_gb"] + ) + offload_group.add_argument( + "--offload-group-size", **offload_kwargs["offload_group_size"] + ) + offload_group.add_argument( + "--offload-num-in-group", **offload_kwargs["offload_num_in_group"] + ) + offload_group.add_argument( + "--offload-prefetch-step", **offload_kwargs["offload_prefetch_step"] + ) + # Multimodal related configs multimodal_kwargs = get_kwargs(MultiModalConfig) multimodal_group = parser.add_argument_group( @@ -1396,10 +1406,6 @@ def create_engine_config( sliding_window=sliding_window, enable_prefix_caching=self.enable_prefix_caching, prefix_caching_hash_algo=self.prefix_caching_hash_algo, - cpu_offload_gb=self.cpu_offload_gb, - offload_group_size=self.offload_group_size, - offload_num_in_group=self.offload_num_in_group, - offload_prefetch_step=self.offload_prefetch_step, calculate_kv_scales=self.calculate_kv_scales, kv_sharing_fast_prefill=self.kv_sharing_fast_prefill, mamba_cache_dtype=self.mamba_cache_dtype, @@ -1730,6 +1736,14 @@ def create_engine_config( compilation_config.max_cudagraph_capture_size = ( self.max_cudagraph_capture_size ) + + offload_config = OffloadConfig( + cpu_offload_gb=self.cpu_offload_gb, + offload_group_size=self.offload_group_size, + offload_num_in_group=self.offload_num_in_group, + offload_prefetch_step=self.offload_prefetch_step, + ) + config = VllmConfig( model_config=model_config, cache_config=cache_config, @@ -1737,6 +1751,7 @@ def create_engine_config( scheduler_config=scheduler_config, device_config=device_config, load_config=load_config, + offload_config=offload_config, attention_config=attention_config, lora_config=lora_config, speculative_config=speculative_config, diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index 402131084b56..18e07f08ecd1 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -13,7 +13,7 @@ from vllm.logger import init_logger if TYPE_CHECKING: - from vllm.config import CacheConfig + from vllm.config import OffloadConfig logger = init_logger(__name__) @@ -101,26 +101,26 @@ def set_offloader(instance: BaseOffloader) -> None: _instance = instance -def create_offloader(cache_config: "CacheConfig") -> BaseOffloader: - """Create an offloader based on the cache configuration. +def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: + """Create an offloader based on the offload configuration. Priority: V2 offloading if configured, else UVA, else noop. """ from vllm.model_executor.offloader.uva import UVAOffloader from vllm.model_executor.offloader.v2 import OffloaderV2 - if cache_config.offload_group_size > 0: + if offload_config.offload_group_size > 0: # Use V2 offloading return OffloaderV2( - group_size=cache_config.offload_group_size, - num_in_group=cache_config.offload_num_in_group, - prefetch_step=cache_config.offload_prefetch_step, + group_size=offload_config.offload_group_size, + num_in_group=offload_config.offload_num_in_group, + prefetch_step=offload_config.offload_prefetch_step, mode="cpu", ) - elif cache_config.cpu_offload_gb > 0: + elif offload_config.cpu_offload_gb > 0: # Use UVA offloading (legacy) return UVAOffloader( - cpu_offload_max_bytes=int(cache_config.cpu_offload_gb * 1024**3) + cpu_offload_max_bytes=int(offload_config.cpu_offload_gb * 1024**3) ) else: # No offloading diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index cb85de0b64bb..e37dc058d219 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -282,6 +282,7 @@ def __init__( self.vllm_config = vllm_config self.model_config = vllm_config.model_config self.cache_config = vllm_config.cache_config + self.offload_config = vllm_config.offload_config self.compilation_config = vllm_config.compilation_config self.lora_config = vllm_config.lora_config self.load_config = vllm_config.load_config @@ -610,7 +611,7 @@ def __init__( self.layerwise_nvtx_hooks_registered = False # Model weight offloader - set_offloader(create_offloader(self.cache_config)) + set_offloader(create_offloader(self.offload_config)) def reset_mm_cache(self) -> None: if self.mm_budget: @@ -5108,7 +5109,7 @@ def may_reinitialize_input_batch( if block_sizes != [self.cache_config.block_size] or kernel_block_sizes != [ self.cache_config.block_size ]: - assert self.cache_config.cpu_offload_gb == 0, ( + assert self.offload_config.cpu_offload_gb == 0, ( "Cannot re-initialize the input batch when CPU weight " "offloading is enabled. See https://github.com/vllm-project/vllm/pull/18298 " # noqa: E501 "for more details." From b07eb3f18d18626ef6126509805e84a12673750f Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sat, 20 Dec 2025 10:19:22 -0800 Subject: [PATCH 07/36] refactor: generalize MoE detection for offloading Address review comment about generalizing MoE detection by searching for FusedMoE instances instead of hardcoding specific model classes. Changes: - Add find_fused_moe_submodule() helper to fused_moe/layer.py - Export the new helper from fused_moe/__init__.py - Replace hardcoded DeepseekV2MoE check with generic helper This makes the offloader configuration reusable across all MoE models (Qwen2MoE, Mixtral, etc.) without model-specific imports. Signed-off-by: Ming Yang --- .../layers/fused_moe/__init__.py | 2 ++ vllm/model_executor/layers/fused_moe/layer.py | 25 +++++++++++++++++++ vllm/model_executor/models/deepseek_v2.py | 11 ++++---- 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 8fee4038b60b..6b038a16d31b 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.layer import ( FusedMoE, FusedMoeWeightScaleSupported, + find_fused_moe_submodule, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEActivationFormat, @@ -55,6 +56,7 @@ def get_config() -> dict[str, Any] | None: "RoutingMethodType", "SharedFusedMoE", "activation_without_mul", + "find_fused_moe_submodule", "override_config", "get_config", ] diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index db97d6eb88ea..6eb9d7d6c181 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -2175,3 +2175,28 @@ def moe_forward_shared_fake( # Mark the FusedMoE weight_loader as supporting MoE-specific parameters # to avoid expensive runtime reflection in model loading code FusedMoE.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] + + +def find_fused_moe_submodule(module: torch.nn.Module) -> torch.nn.Module: + """Find a FusedMoE submodule for offloading, or return the module itself. + + Searches module attributes for instances of FusedMoE (or subclasses like + SharedFusedMoE). + + Args: + module: The module to search within (typically layer.mlp). + + Returns: + The first FusedMoE instance found, or the original module if none found. + """ + for attr_name in dir(module): + if attr_name.startswith("_"): + continue + try: + attr = getattr(module, attr_name, None) + except Exception: + continue + if isinstance(attr, FusedMoE): + return attr + + return module diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index bbe8c298972a..60940b2e79c9 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -49,7 +49,10 @@ from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.model_executor.layers.fused_moe import SharedFusedMoE +from vllm.model_executor.layers.fused_moe import ( + SharedFusedMoE, + find_fused_moe_submodule, +) from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -1276,11 +1279,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", offloader_kwargs=dict( # Extract the MLP submodule - for MoE layers, go deeper to the experts - submodule_accessor=lambda layer: ( - layer.mlp.experts - if isinstance(layer.mlp, DeepseekV2MoE) - else layer.mlp - ), + submodule_accessor=lambda layer: find_fused_moe_submodule(layer.mlp), # Specify which parameters to offload whitelist_param_names_creator=lambda module: ( [ From ac5fb492da4d1f00692da398725a92e20d21d2d3 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sat, 20 Dec 2025 10:41:50 -0800 Subject: [PATCH 08/36] cleanup: remove NVFP4 scale params from offload whitelist Remove w13_blockscale_swizzled and w2_blockscale_swizzled from the offload parameter whitelist - only offload the core expert weights. Signed-off-by: Ming Yang --- vllm/model_executor/models/deepseek_v2.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 60940b2e79c9..735d4e2910c0 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1286,12 +1286,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # Core MoE expert weights "w13_weight", "w2_weight", - # NVFP4 quantization scales (if present) - *( - ["w13_blockscale_swizzled", "w2_blockscale_swizzled"] - if hasattr(module, "w13_blockscale_swizzled") - else [] - ), ] # Only offload from MoE experts (SharedFusedMoE/FusedMoE) if hasattr(module, "w13_weight") From 80e764e812ebfb3770578eff160ca7370e743ae2 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sat, 20 Dec 2025 10:44:17 -0800 Subject: [PATCH 09/36] cleanup: remove dead code in UVAOffloader Remove redundant uva_offloading variable and dead code branches. Since UVA is required (asserted in __init__), the non-UVA code paths were never executed. Signed-off-by: Ming Yang --- vllm/model_executor/offloader/uva.py | 33 +++------------------------- 1 file changed, 3 insertions(+), 30 deletions(-) diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index d1041a62500a..14971a4f32b3 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -6,7 +6,6 @@ import torch import torch.nn as nn -from torch.func import functional_call from vllm.model_executor.offloader.base import BaseOffloader from vllm.utils.platform_utils import is_pin_memory_available, is_uva_available @@ -57,14 +56,9 @@ def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: return module pin_memory = is_pin_memory_available() - uva_available = is_uva_available() - - assert uva_available, "V1 CPU offloading requires uva (pin memory) support" - uva_offloading = True # offload parameters to CPU # use pin_memory if possible, which helps cudagraph capture speed - offloaded_parameters = False for p in module.parameters(): if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: # we use per-parameter offloading @@ -81,30 +75,9 @@ def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: pin_memory=pin_memory, ) cpu_data.copy_(p.data) - if not uva_offloading: - p.data = cpu_data - else: - # keep the cpu data alive - p._vllm_offloaded_cpu_data = cpu_data - p.data = get_cuda_view_from_cpu_tensor(cpu_data) + # keep the cpu data alive + p._vllm_offloaded_cpu_data = cpu_data + p.data = get_cuda_view_from_cpu_tensor(cpu_data) self.cpu_offload_bytes += p.data.numel() * p.data.element_size() - offloaded_parameters = True - - if offloaded_parameters and not uva_offloading: - original_forward = module.forward - - def forward(*args, **kwargs): - module.forward = original_forward - device_state = { - # here we blindly call `to(device)` - # if the parameter is already on the device, it will be a no-op - k: v.to(device, non_blocking=True) - for k, v in module.state_dict().items() - } - output = functional_call(module, device_state, args=args, kwargs=kwargs) - module.forward = forward - return output - - module.forward = forward return module From 5bc88d87b550bb969e7776fdc02a76e9b4c5e61e Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sun, 25 Jan 2026 09:54:55 -0800 Subject: [PATCH 10/36] [Core] V2 offloader: static buffers and custom ops for torch.compile This commit introduces V2 offloader improvements for torch.compile and CUDA graph compatibility: 1. Static GPU buffer pool with double/triple buffering support 2. Custom ops (wait_prefetch, start_prefetch) that create data dependencies via mutates_args to prevent torch.compile from reordering operations 3. Stream synchronization to ensure correct prefetch/compute ordering 4. Stride-preserving buffers to maintain parameter memory layout Key design decisions: - Parameters point to static GPU buffers (not CPU storage) so torch.compile sees GPU tensors during tracing - sync_cpu_storage() updates CPU storage after process_weights_after_loading to capture quantization-processed weights - Buffer slot keys include parameter name to prevent different parameters with same shape from sharing buffers within a layer Verified with lm_eval gsm8k accuracy >0.94 on DeepSeek-R1-0528-NVFP4-v2. Signed-off-by: Ming Yang --- tests/basic_correctness/test_v2_offload.py | 4 +- vllm/model_executor/offloader/v2.py | 459 +++++++++++++++++---- vllm/model_executor/offloader/v2_ops.py | 118 ++++++ 3 files changed, 510 insertions(+), 71 deletions(-) create mode 100644 vllm/model_executor/offloader/v2_ops.py diff --git a/tests/basic_correctness/test_v2_offload.py b/tests/basic_correctness/test_v2_offload.py index bbe89a02def2..acffb2240f50 100644 --- a/tests/basic_correctness/test_v2_offload.py +++ b/tests/basic_correctness/test_v2_offload.py @@ -25,7 +25,7 @@ def test_v2_offload_deepseek(): "2", "--offload-prefetch-step", "1", - # currently not compatible with torch.compile - "--enforce-eager", + # torch.compile is automatically disabled when V2 offloading is + # enabled (via enable_if in @support_torch_compile decorator) ], ) diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/v2.py index 8c69a7248e10..180f1a4b25d8 100644 --- a/vllm/model_executor/offloader/v2.py +++ b/vllm/model_executor/offloader/v2.py @@ -2,15 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from # https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py -"""OffloaderV2: CPU offloading with async prefetching.""" +"""OffloaderV2: CPU offloading with async prefetching. + +This version uses static buffers and stream synchronization (instead of +CUDA events) for torch.compile + CUDA graph compatibility. +""" from abc import ABC, abstractmethod from collections.abc import Callable, Generator +from dataclasses import dataclass import torch import torch.nn as nn -from torch.func import functional_call +# Import v2_ops to register custom ops at module load time +import vllm.model_executor.offloader.v2_ops # noqa: F401 from vllm.logger import init_logger from vllm.model_executor.offloader.base import BaseOffloader from vllm.utils.platform_utils import is_pin_memory_available @@ -21,11 +27,109 @@ _WhitelistParamNamesCreator = Callable[[nn.Module], list[str]] +@dataclass +class ParamInfo: + """Metadata about an offloaded parameter.""" + + name: str + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: torch.dtype + + @property + def key(self) -> tuple[str, tuple[int, ...], tuple[int, ...], torch.dtype]: + """Unique key for buffer pool grouping. + + Includes parameter name to prevent different parameters with the same + shape from sharing buffers within the same layer. Parameters with the + same name across different layers will share buffers (via slots). + + Includes stride because parameters with same shape but different + strides need separate buffers to preserve memory layout. + """ + return (self.name, self.shape, self.stride, self.dtype) + + @property + def num_bytes(self) -> int: + """Size in bytes.""" + numel = 1 + for dim in self.shape: + numel *= dim + return numel * torch.tensor([], dtype=self.dtype).element_size() + + +class StaticBufferPool: + """Pre-allocated GPU buffer pool for offloaded parameters. + + Allocates slot_capacity copies of each unique parameter + (name, shape, stride, dtype), allowing for double/triple buffering + during prefetch. + + Buffer slots are reused circularly: layer N uses slot (N % slot_capacity). + + The key includes parameter name to prevent different parameters within + the same layer from sharing buffers. Parameters with the same name + across different layers share buffers via the slot mechanism. + """ + + def __init__( + self, + param_infos: list[ParamInfo], + slot_capacity: int, + device: torch.device, + ): + self.slot_capacity = slot_capacity + self.total_bytes = 0 + self._device = device + + # Group by (shape, stride, dtype) - only allocate unique combinations + unique_params: dict[tuple, ParamInfo] = {} + for info in param_infos: + if info.key not in unique_params: + unique_params[info.key] = info + + # Allocate buffers: key -> list of tensors (one per slot) + self._buffers: dict[tuple, list[torch.Tensor]] = {} + for key, info in unique_params.items(): + slot_tensors = [] + for _ in range(slot_capacity): + # Use empty_strided to preserve parameter's memory layout + buf = torch.empty_strided( + size=info.shape, + stride=info.stride, + dtype=info.dtype, + device=device, + ) + slot_tensors.append(buf) + self.total_bytes += info.num_bytes + self._buffers[key] = slot_tensors + + logger.debug( + "[StaticBufferPool] Allocated %d unique (name, shape, stride, dtype), " + "%d slots each, total %.4f GB", + len(unique_params), + slot_capacity, + self.total_bytes / 1e9, + ) + + def get_buffer( + self, + name: str, + shape: tuple[int, ...], + stride: tuple[int, ...], + dtype: torch.dtype, + slot_idx: int, + ) -> torch.Tensor: + """Get a static buffer for the given name/shape/stride/dtype/slot.""" + key = (name, shape, stride, dtype) + return self._buffers[key][slot_idx % self.slot_capacity] + + class OffloaderV2(BaseOffloader): """Advanced offloader with group-based selection and async prefetching. - Unlike UVA offloading which provides zero-copy access, V2 explicitly - manages parameter transfers with prefetching to hide latency. + Uses static buffers and stream synchronization for torch.compile and + CUDA graph compatibility. Args: group_size: Group every N layers together. @@ -45,10 +149,20 @@ def __init__( self.num_in_group = num_in_group self.prefetch_step = prefetch_step self.mode = mode - self.alt_stream = torch.cuda.Stream() + + # Copy stream for async H2D transfers + self.copy_stream = torch.cuda.Stream() + + # Module offloaders and buffer pool (populated in wrap_modules/post_init) self.module_offloaders: list[_ModuleOffloader] = [] + self.buffer_pool: StaticBufferPool | None = None self.total_offloaded_bytes = 0 + # Register this instance for custom ops + from vllm.model_executor.offloader.v2_ops import set_offloader_instance + + set_offloader_instance(self) + def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], @@ -81,8 +195,9 @@ def wrap_modules( _ModuleOffloader( mode=self.mode, module=submodule, - alt_stream=self.alt_stream, + copy_stream=self.copy_stream, whitelist_param_names=whitelist_param_names, + layer_idx=len(self.module_offloaders), ) ) @@ -92,58 +207,132 @@ def wrap_modules( return all_modules def _hook_module_forward(self, index: int, module: nn.Module): - """Hook module's forward to implement prefetch + execute + offload pattern.""" + """Hook module's forward with torch.compile-compatible sync.""" original_forward = module.forward def forward(*args, **kwargs): + # Temporarily restore original forward to avoid recursion module.forward = original_forward - device_tensors = self.module_offloaders[index].wait_and_get_device_tensors() - output = functional_call(module, device_tensors, args=args, kwargs=kwargs) + + # Wait for this layer's prefetch to complete + input_tensor = args[0] if args else kwargs.get("hidden_states") + input_tensor = torch.ops.vllm.wait_prefetch(input_tensor, index) + + # Replace the first arg with the returned tensor to maintain dependency + if args: + args = (input_tensor,) + args[1:] + else: + kwargs["hidden_states"] = input_tensor + + # No parameter swapping needed - parameters already point to + # GPU static buffers (set in assign_static_buffer) + output = original_forward(*args, **kwargs) + + # Start prefetch for next layer (circular) + # Custom op returns output_tensor to create data dependency next_index = (index + self.prefetch_step) % len(self.module_offloaders) - self.module_offloaders[next_index].start_onload() - self.module_offloaders[index].offload() + # Handle tuple output (e.g., (hidden_states, residual)) + if isinstance(output, tuple): + output_tensor = torch.ops.vllm.start_prefetch(output[0], next_index) + output = (output_tensor,) + output[1:] + else: + output = torch.ops.vllm.start_prefetch(output, next_index) + + # No explicit offload needed - static buffers are reused implicitly + + # Restore hooked forward module.forward = forward return output module.forward = forward + def _wait_for_layer(self, layer_idx: int): + """Called by custom op - wait for copy stream to complete.""" + # wait_stream creates a CUDA graph dependency edge when captured + torch.cuda.current_stream().wait_stream(self.copy_stream) + + def _start_prefetch(self, layer_idx: int): + """Called by custom op - start async copy to static buffer.""" + offloader = self.module_offloaders[layer_idx] + offloader.start_onload_to_static() + def post_init(self): - """Initialize offloaders and start prefetching first N modules.""" + """Allocate static buffer pool and start initial prefetches. + + Note: Parameters have already been offloaded to CPU during wrap_modules() + (in _CpuParamOffloader.__init__), so GPU memory is available for the + static buffer pool. + """ + # Sync CPU storage with current param.data BEFORE collecting param info. + # This is needed because process_weights_after_loading may have: + # 1. Transformed weights (quantization, transpose, etc.) + # 2. Created new CPU tensors via device_loading_context + # Our _cpu_storage would be stale otherwise. + for offloader in self.module_offloaders: + offloader.sync_cpu_storage() + + # Collect parameter info (now using synced CPU storage) + param_infos: list[ParamInfo] = [] + device: torch.device | None = None + + for offloader in self.module_offloaders: + param_infos.extend(offloader.get_param_infos()) + if device is None: + device = offloader.device + + if device is None: + # No modules to offload + return + + # Allocate static buffer pool + self.buffer_pool = StaticBufferPool( + param_infos=param_infos, + slot_capacity=self.prefetch_step, + device=device, + ) + + # Assign buffer slots and point parameters to GPU buffers + for idx, offloader in enumerate(self.module_offloaders): + slot_idx = idx % self.prefetch_step + offloader.assign_buffer_slot(self.buffer_pool, slot_idx) + + # Collect offloaded bytes for offloader in self.module_offloaders: offloader.post_init() self.total_offloaded_bytes += offloader.offloaded_bytes logger.info_once( f"[OffloaderV2] Initialized {len(self.module_offloaders)} modules. " - f"Total GPU memory saved: {self.total_offloaded_bytes / 1e9:.4f} GB " + f"Total GPU memory saved: {self.total_offloaded_bytes / 1e9:.4f} GB, " + f"Static buffer pool: {self.buffer_pool.total_bytes / 1e9:.4f} GB " f"(group_size={self.group_size}, num_in_group={self.num_in_group}, " f"prefetch_step={self.prefetch_step}, mode={self.mode})" ) + # Start initial prefetches for i in range(min(self.prefetch_step, len(self.module_offloaders))): - self.module_offloaders[i].start_onload() + self.module_offloaders[i].start_onload_to_static() class _ModuleOffloader: """Manages offloading for a single module. - Responsibilities: - - Create parameter offloaders for each parameter - - Coordinate async loading via alternate CUDA stream - - Provide device tensors when needed + Uses static buffers from a shared pool instead of dynamic allocation. """ def __init__( self, mode: str, module: nn.Module, - alt_stream: torch.cuda.Stream, + copy_stream: torch.cuda.Stream, whitelist_param_names: list[str], + layer_idx: int, ): self.mode = mode self.module = module self.device = next(module.parameters()).device - self.alt_stream = alt_stream + self.copy_stream = copy_stream + self.layer_idx = layer_idx self.offloaded_bytes = 0 assert self.device != torch.device("cpu"), ( @@ -151,8 +340,9 @@ def __init__( "(offloader handles CPU placement)" ) - self._device_tensors: dict[str, torch.Tensor] | None = None - self._load_event: torch.cuda.Event | None = None + # Buffer pool and slot (assigned in assign_buffer_slot) + self._buffer_pool: StaticBufferPool | None = None + self._buffer_slot_idx: int = 0 param_dict = dict(self.module.named_parameters()) assert all(name in param_dict for name in whitelist_param_names), ( @@ -171,36 +361,92 @@ def post_init(self): param_offloader.post_init() self.offloaded_bytes += param_offloader.offloaded_bytes - def start_onload(self): - """Start async loading in alternate CUDA stream.""" - self.alt_stream.wait_stream(torch.cuda.current_stream()) - - with torch.cuda.stream(self.alt_stream): - self._device_tensors = { - name: offloader.create_device_tensor() - for name, offloader in self._param_offloaders.items() - } - self._load_event = torch.cuda.Event() - self._load_event.record() - - def offload(self): - """Free device tensors (offload from GPU memory).""" - self._device_tensors = None - self._load_event = None - - def wait_and_get_device_tensors(self) -> dict[str, torch.Tensor]: - """Wait for async loading to complete and return device tensors.""" - assert self._device_tensors is not None, ( - "Tensors not loaded (call start_onload first)" - ) - if self._load_event is not None: - self._load_event.wait() - return self._device_tensors + def sync_cpu_storage(self): + """Sync CPU storage with current param.data. + + Called after process_weights_after_loading to ensure _cpu_storage + contains the final processed weights, not stale pre-loading data. + """ + for param_offloader in self._param_offloaders.values(): + param_offloader.sync_cpu_storage() + + def get_param_infos(self) -> list[ParamInfo]: + """Get parameter metadata for buffer pool allocation. + + Note: sync_cpu_storage() must be called before this method to ensure + _cpu_storage reflects the final processed weights (after quantization). + """ + infos = [] + for name, offloader in self._param_offloaders.items(): + cpu_storage = offloader._cpu_storage + assert cpu_storage is not None, "CPU storage not initialized" + infos.append( + ParamInfo( + name=name, + shape=tuple(cpu_storage.shape), + stride=tuple(cpu_storage.stride()), + dtype=cpu_storage.dtype, + ) + ) + return infos + + def assign_buffer_slot(self, pool: StaticBufferPool, slot_idx: int): + """Assign this module to a buffer slot in the pool. + + Also assigns static GPU buffers to each parameter offloader, + which moves the parameter data to point to the GPU buffer. + """ + self._buffer_pool = pool + self._buffer_slot_idx = slot_idx + + # Assign static buffers to parameters + # Use CPU storage shape/stride/dtype since param.data is now empty + for name, offloader in self._param_offloaders.items(): + cpu_storage = offloader._cpu_storage + assert cpu_storage is not None, "CPU storage not initialized" + buffer = pool.get_buffer( + name=name, + shape=tuple(cpu_storage.shape), + stride=tuple(cpu_storage.stride()), + dtype=cpu_storage.dtype, + slot_idx=slot_idx, + ) + offloader.assign_static_buffer(buffer) + + def start_onload_to_static(self): + """Start async copy from CPU storage to GPU buffer. + + Uses the stored _gpu_buffer reference directly instead of accessing + param.data, to avoid any issues with parameter identity changes. + + IMPORTANT: We must wait for the compute stream before copying, because + the previous layer's forward may still be using the buffer (GPU ops are + async). Without this sync, we could overwrite the buffer while it's + being read. + """ + assert self._buffer_pool is not None, "Buffer pool not assigned" + + # Wait for compute stream to finish using the buffer before overwriting + self.copy_stream.wait_stream(torch.cuda.current_stream()) + + with torch.cuda.stream(self.copy_stream): + for name, offloader in self._param_offloaders.items(): + cpu_storage = offloader._cpu_storage + gpu_buffer = offloader._gpu_buffer + assert cpu_storage is not None, "CPU storage not initialized" + assert gpu_buffer is not None, "GPU buffer not assigned" + # Async copy from pinned CPU storage to GPU buffer + gpu_buffer.copy_(cpu_storage, non_blocking=True) class _BaseParamOffloader(ABC): """Base class for parameter offloading strategies.""" + # CPU storage for offloaded parameters (set by subclasses) + _cpu_storage: torch.Tensor | None + # GPU buffer reference (set by subclasses when using static buffers) + _gpu_buffer: torch.Tensor | None + @staticmethod def create(mode: str, **kwargs) -> "_BaseParamOffloader": """Factory method to create appropriate offloader for mode.""" @@ -213,6 +459,8 @@ def __init__(self, module: nn.Module, param_name: str): self._module = module self._param_name = param_name self.offloaded_bytes = 0 + self._cpu_storage = None + self._gpu_buffer = None @property def _param(self) -> nn.Parameter: @@ -224,26 +472,51 @@ def post_init(self): return @abstractmethod - def create_device_tensor(self) -> torch.Tensor: - """Create device tensor from offloaded storage.""" + def sync_cpu_storage(self) -> None: + """Sync CPU storage with current param.data. + + Called after process_weights_after_loading to update _cpu_storage + with the final processed weights. + """ + pass + + @abstractmethod + def assign_static_buffer(self, gpu_buffer: torch.Tensor) -> None: + """Point parameter data to GPU static buffer.""" pass class _CpuParamOffloader(_BaseParamOffloader): - """Offload parameter to pinned CPU memory.""" + """Offload parameter to pinned CPU memory. + + Uses GPU static buffers as the actual parameter, with CPU storage + kept separately. This ensures torch.compile sees GPU tensors at trace time. + + The offloading happens in two phases: + 1. __init__() - copies GPU data to CPU, frees GPU memory immediately + 2. assign_static_buffer() - points param.data to GPU static buffer + """ def __init__(self, module: nn.Module, param_name: str): super().__init__(module, param_name) - self._move_param_to_cpu() + self._cpu_storage: torch.Tensor | None = None + self._gpu_buffer: torch.Tensor | None = None # Store reference to GPU buffer + + # Offload to CPU immediately to free GPU memory during model loading + self._offload_to_cpu_internal() + + def _offload_to_cpu_internal(self): + """Copy parameter data to pinned CPU storage and free GPU memory. - def _move_param_to_cpu(self): - """Move parameter data to pinned CPU memory (modify param.data in-place).""" + This replaces param.data with CPU storage, allowing weight loading + to continue writing to CPU memory. GPU memory is freed when the + original GPU tensor is garbage collected. + """ param = self._param pin_memory = is_pin_memory_available() - self.offloaded_bytes = param.data.numel() * param.data.element_size() - - cpu_data = torch.empty_strided( + # Create pinned CPU storage and copy current GPU data + self._cpu_storage = torch.empty_strided( size=param.data.size(), stride=param.data.stride(), dtype=param.data.dtype, @@ -251,23 +524,71 @@ def _move_param_to_cpu(self): device="cpu", pin_memory=pin_memory, ) - cpu_data.copy_(param.data) + self._cpu_storage.copy_(param.data) - logger.debug_once( - f"[OffloaderV2] Offloaded parameter '{self._param_name}': " - f"shape={tuple(param.shape)}, dtype={param.dtype}, " - f"size={self.offloaded_bytes / 1e9:.6f} GB, pinned={pin_memory}" + self.offloaded_bytes = ( + self._cpu_storage.numel() * self._cpu_storage.element_size() ) - param.data = cpu_data + # Point param.data to CPU storage - this allows weight loading to work + # and frees GPU memory when the original GPU tensor is garbage collected + param.data = self._cpu_storage - def post_init(self): - """No-op: offloading already done in __init__.""" - pass + def assign_static_buffer(self, gpu_buffer: torch.Tensor) -> None: + """Point parameter data to GPU static buffer. + + This is called after weight loading AND process_weights_after_loading + complete. At this point: + - param.data may have been replaced by device_loading_context + (which creates new CPU tensors after quantization processing) + - We need to update _cpu_storage to point to current param.data + so that prefetch copies the processed weights, not stale data + - Then point param.data to the GPU buffer for torch.compile + """ + assert self._cpu_storage is not None, ( + "_offload_to_cpu_internal() must be called before assign_static_buffer()" + ) + + # Get current parameter (may have been replaced by + # process_weights_after_loading) + param = self._param + + # Update _cpu_storage to current param.data. This is critical because: + # 1. process_weights_after_loading may transform weights (quantization) + # 2. device_loading_context creates NEW CPU tensors when moving back + # 3. Our old _cpu_storage would have pre-processed or stale data + if param.data.device.type == "cpu": + # param.data is already on CPU - use it as our CPU storage + self._cpu_storage = param.data + else: + # param.data is on GPU - copy to our CPU storage + self._cpu_storage.copy_(param.data) - def create_device_tensor(self) -> torch.Tensor: - """Load from CPU to GPU (async if pinned). + # Store reference to GPU buffer for use in start_onload + self._gpu_buffer = gpu_buffer - Returns a CUDA copy of the parameter (which has CPU data). + # Point parameter to static GPU buffer - this is what torch.compile sees + param.data = gpu_buffer + + def sync_cpu_storage(self) -> None: + """Sync CPU storage with current param.data. + + Called after process_weights_after_loading to update _cpu_storage + with the final processed weights. This is critical because: + 1. process_weights_after_loading may transform weights (quantization) + 2. device_loading_context creates NEW CPU tensors when moving back + 3. Our old _cpu_storage would have pre-processed or stale data """ - return self._param.to("cuda", non_blocking=True) + param = self._param + + if param.data.device.type == "cpu": + # param.data is already on CPU - use it as our CPU storage + self._cpu_storage = param.data + else: + # param.data is on GPU - copy to existing CPU storage + assert self._cpu_storage is not None + self._cpu_storage.copy_(param.data) + + def post_init(self): + """No-op: offloading done in offload_to_cpu/assign_static_buffer.""" + pass diff --git a/vllm/model_executor/offloader/v2_ops.py b/vllm/model_executor/offloader/v2_ops.py new file mode 100644 index 000000000000..28e35ab35f50 --- /dev/null +++ b/vllm/model_executor/offloader/v2_ops.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Custom ops for V2 offloader torch.compile + CUDA graph compatibility. + +These ops use mutates_args to create data dependencies that prevent +the compiler from reordering prefetch/sync operations. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from vllm.utils.torch_utils import direct_register_custom_op + +if TYPE_CHECKING: + from vllm.model_executor.offloader.v2 import OffloaderV2 + +# Global reference to the offloader instance, set during OffloaderV2.__init__ +_offloader_instance: OffloaderV2 | None = None + + +def set_offloader_instance(offloader: OffloaderV2 | None) -> None: + """Set the global offloader instance for custom ops to use.""" + global _offloader_instance + _offloader_instance = offloader + + +# --- wait_prefetch op --- + + +def _wait_prefetch_impl( + input_tensor: torch.Tensor, + layer_idx: int, +) -> torch.Tensor: + """Wait for prefetch of layer_idx to complete. + + Synchronizes the compute stream with the copy stream to ensure + the prefetched weights are ready for use. + + Args: + input_tensor: Input to the layer (e.g., hidden_states) - returned + to create data dependency chain. + layer_idx: Index of the layer to wait for. + + Returns: + input_tensor unchanged, but creates data dependency for torch.compile. + """ + if _offloader_instance is not None: + _offloader_instance._wait_for_layer(layer_idx) + return input_tensor + + +def _wait_prefetch_fake( + input_tensor: torch.Tensor, + layer_idx: int, +) -> torch.Tensor: + """Fake implementation for torch.compile tracing.""" + return input_tensor + + +# --- start_prefetch op --- + + +def _start_prefetch_impl( + output_tensor: torch.Tensor, + layer_idx: int, +) -> torch.Tensor: + """Start async prefetch of layer_idx weights. + + Initiates H2D copy on the copy stream for the specified layer. + + Args: + output_tensor: Output from forward - returned to create ordering + dependency. This prevents torch.compile from reordering + this op before the computation that produces output_tensor. + layer_idx: Index of the layer to prefetch. + + Returns: + output_tensor unchanged, creating data dependency for torch.compile. + """ + if _offloader_instance is not None: + _offloader_instance._start_prefetch(layer_idx) + return output_tensor + + +def _start_prefetch_fake( + output_tensor: torch.Tensor, + layer_idx: int, +) -> torch.Tensor: + """Fake implementation for torch.compile tracing.""" + return output_tensor + + +def register_v2_offloader_ops() -> None: + """Register custom ops for V2 offloader. + + Must be called before the ops are used. This is typically done + at module import time. + """ + direct_register_custom_op( + op_name="wait_prefetch", + op_func=_wait_prefetch_impl, + mutates_args=["input_tensor"], + fake_impl=_wait_prefetch_fake, + ) + + direct_register_custom_op( + op_name="start_prefetch", + op_func=_start_prefetch_impl, + mutates_args=["output_tensor"], + fake_impl=_start_prefetch_fake, + ) + + +# Register ops at module import time +register_v2_offloader_ops() From 719af1bb73199acd94912dd529281246f49e79fe Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Mon, 26 Jan 2026 16:45:53 -0800 Subject: [PATCH 11/36] [Core] V2 offloader: CUDA graph capture support Add CUDA graph capture compatibility for the V2 offloader when used with torch.compile. This enables CPU weight offloading to work with piecewise and full cudagraph modes. Problem: When CUDA graph capture started, pre-capture prefetches on copy_stream caused errors: - cudaErrorStreamCaptureIsolation: dependency on uncaptured work - cudaErrorStreamCaptureUnjoined: forked stream not rejoined - cudaErrorInvalidValue: waiting on invalid event state Solution: Use event-based stream forking during capture and wait_stream outside capture. Key mechanisms: 1. Event-based fork/join during capture: - Record fork_event on compute stream before prefetch - copy_stream.wait_event(fork_event) joins it to capture - Record _copy_done_event after H2D copies complete - compute_stream.wait_event(_copy_done_event) creates dependency 2. Conditional synchronization: - _prefetch_in_capture flag tracks if prefetch was during capture - During capture: skip wait_event for pre-capture prefetches - Outside capture: use wait_stream for robustness 3. Capture boundary handling: - sync_before_graph_capture(): ensure pre-capture work complete - join_after_forward(): rejoin copy_stream after model forward - Reset flags between piecewise cudagraph subgraphs Integration points: - cuda_graph.py: CUDAGraphWrapper capture - cudagraph_utils.py: V1 worker capture - eagle_cudagraph.py: speculative decoding capture - gpu_ubatch_wrapper.py: micro-batch wrapper capture With this fix, relaxed capture mode (capture_error_mode="relaxed") is no longer needed and has been removed from all capture sites. Signed-off-by: Ming Yang --- tests/basic_correctness/test_v2_offload.py | 2 +- vllm/compilation/cuda_graph.py | 17 +++ vllm/model_executor/offloader/v2.py | 128 ++++++++++++++++-- vllm/model_executor/offloader/v2_ops.py | 29 ++++ vllm/v1/worker/gpu/cudagraph_utils.py | 16 +++ .../worker/gpu/spec_decode/eagle_cudagraph.py | 11 ++ vllm/v1/worker/gpu_ubatch_wrapper.py | 9 ++ 7 files changed, 203 insertions(+), 9 deletions(-) diff --git a/tests/basic_correctness/test_v2_offload.py b/tests/basic_correctness/test_v2_offload.py index acffb2240f50..071b702c3569 100644 --- a/tests/basic_correctness/test_v2_offload.py +++ b/tests/basic_correctness/test_v2_offload.py @@ -16,7 +16,6 @@ def test_v2_offload_deepseek(): """ compare_two_settings( "deepseek-ai/DeepSeek-V2-Lite", - [], # Baseline: no offloading [ # V2 offloading configuration "--offload-group-size", @@ -28,4 +27,5 @@ def test_v2_offload_deepseek(): # torch.compile is automatically disabled when V2 offloading is # enabled (via enable_if in @support_torch_compile decorator) ], + [], # Baseline: no offloading ) diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index 7ffa74d0d7e6..fb1dfe9b2e72 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -17,6 +17,10 @@ from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id from vllm.forward_context import BatchDescriptor, get_forward_context from vllm.logger import init_logger +from vllm.model_executor.offloader.v2_ops import ( + join_offloader_after_forward, + sync_offloader_before_capture, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import current_stream, weak_ref_tensors @@ -265,6 +269,11 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: set_graph_pool_id(self.graph_pool) else: set_graph_pool_id(current_platform.graph_pool_handle()) + + # Sync offloader's copy stream before capture. + # Ensure any pre-capture prefetches from offloader are complete. + sync_offloader_before_capture() + # mind-exploding: carefully manage the reference and memory. with torch.cuda.graph( cudagraph, @@ -273,6 +282,11 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: ): # `output` is managed by pytorch's cudagraph pool output = self.runnable(*args, **kwargs) + # Join offloader's copy stream after forward to avoid + # unjoined stream error. The last layer's start_prefetch + # forks copy_stream, but wait_prefetch only happens in + # the next forward pass. + join_offloader_after_forward() if self.cudagraph_options.weak_ref_output: # by converting it to weak ref, # the original `output` will immediately be released @@ -305,5 +319,8 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: f"got {new_input_addresses}" ) + # Sync offloader before replay - ensures any external dependencies + # from pre-capture prefetches are satisfied. + sync_offloader_before_capture() entry.cudagraph.replay() return entry.output diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/v2.py index 180f1a4b25d8..7336e4c1314f 100644 --- a/vllm/model_executor/offloader/v2.py +++ b/vllm/model_executor/offloader/v2.py @@ -4,8 +4,9 @@ # https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py """OffloaderV2: CPU offloading with async prefetching. -This version uses static buffers and stream synchronization (instead of -CUDA events) for torch.compile + CUDA graph compatibility. +This version uses static buffers and event-based stream forking for +torch.compile + CUDA graph compatibility. Events allow the copy stream +to join CUDA graph captures, ensuring H2D copies are properly captured. """ from abc import ABC, abstractmethod @@ -247,15 +248,108 @@ def forward(*args, **kwargs): module.forward = forward def _wait_for_layer(self, layer_idx: int): - """Called by custom op - wait for copy stream to complete.""" - # wait_stream creates a CUDA graph dependency edge when captured + """Called by custom op - wait for copy to complete. + + Synchronization strategy: + - During CUDA graph capture: use event-based wait (graph-compatible) + - Outside capture (warmup/eager): use wait_stream (more robust) + + During capture, we skip wait for pre-capture prefetches because: + 1. sync_before_graph_capture() ensures pre-capture work is complete + 2. We can't wait on pre-capture events during capture (isolation error) + """ + offloader = self.module_offloaders[layer_idx] + + if torch.cuda.is_current_stream_capturing(): + # During capture, skip wait for pre-capture prefetches. + # sync_before_graph_capture() ensures pre-capture work is complete. + if not offloader._prefetch_in_capture: + return + # Event-based wait for in-capture prefetches (graph-compatible) + torch.cuda.current_stream().wait_event(offloader._copy_done_event) + # Mark that this prefetch has been waited on (joined). + offloader._prefetch_in_capture = False + else: + # Outside capture: use wait_stream for robustness. + # Events used in previous captures can be in invalid state. + torch.cuda.current_stream().wait_stream(self.copy_stream) + + def sync_before_graph_capture(self): + """Sync copy stream before CUDA graph capture or replay. + + Pre-capture prefetches from warmup must complete before capture. + This method ensures those dependencies are satisfied. + + Call this: + 1. Before capturing a CUDA graph + 2. Before replaying a CUDA graph (if prefetches were issued outside) + + Also resets _prefetch_in_capture flags. This is critical for piecewise + cudagraph mode where multiple subgraphs are captured sequentially. + Each subgraph's prefetches must be tracked independently - we can't + wait on events recorded in a different subgraph's capture. + """ torch.cuda.current_stream().wait_stream(self.copy_stream) + # Reset flags so only prefetches started in THIS capture are tracked. + # This prevents cross-capture event waits which cause cudaErrorInvalidValue. + for offloader in self.module_offloaders: + offloader._prefetch_in_capture = False + def _start_prefetch(self, layer_idx: int): """Called by custom op - start async copy to static buffer.""" offloader = self.module_offloaders[layer_idx] offloader.start_onload_to_static() + def _join_copy_stream(self, layer_idx: int): + """Called by custom op - join copy_stream back to compute stream. + + Used after the last layer's start_prefetch to ensure copy_stream is + rejoined before CUDA graph capture ends. The start_prefetch forks + copy_stream into the capture, but the corresponding wait_prefetch + only happens in the next forward pass. During capture, this would + leave copy_stream unjoined, causing cudaErrorStreamCaptureUnjoined. + + During capture, we check if the prefetch was started during the same + capture. If not (e.g., from post_init or warmup), we skip the wait + because: + 1. sync_before_graph_capture() ensures pre-capture work is complete + 2. We can't wait on pre-capture events during capture (would cause + cudaErrorStreamCaptureIsolation) + """ + offloader = self.module_offloaders[layer_idx] + + # During capture, skip wait for pre-capture prefetches. + # sync_before_graph_capture() ensures pre-capture work is complete. + if ( + torch.cuda.is_current_stream_capturing() + and not offloader._prefetch_in_capture + ): + return + + torch.cuda.current_stream().wait_event(offloader._copy_done_event) + + def join_after_forward(self): + """Join copy_stream after model forward completes. + + Call this after the model forward pass but before CUDA graph capture + ends. This ensures copy_stream is rejoined for any prefetches started + during the forward pass. + + We join ALL layers that have _prefetch_in_capture=True, meaning their + prefetch was started during capture but not yet waited on (joined). + This handles both full and piecewise cudagraph modes correctly: + - Full mode: joins layers 0..prefetch_step-1 (prefetched by last layers) + - Piecewise mode: joins only layers prefetched by THIS subgraph's layers + """ + if not self.module_offloaders: + return + # Join all layers whose prefetch was started in capture but not waited on + for offloader in self.module_offloaders: + if offloader._prefetch_in_capture: + torch.cuda.current_stream().wait_event(offloader._copy_done_event) + offloader._prefetch_in_capture = False + def post_init(self): """Allocate static buffer pool and start initial prefetches. @@ -335,6 +429,15 @@ def __init__( self.layer_idx = layer_idx self.offloaded_bytes = 0 + # Event to signal when H2D copy to static buffer is complete. + # Used for CUDA graph compatible synchronization during capture. + # Outside capture, we use wait_stream instead (more robust). + self._copy_done_event = torch.cuda.Event() + + # Track if last prefetch was started during CUDA graph capture. + # Used to skip wait_event during capture for pre-capture prefetches. + self._prefetch_in_capture = False + assert self.device != torch.device("cpu"), ( "Module parameters should not already be on CPU " "(offloader handles CPU placement)" @@ -416,8 +519,8 @@ def assign_buffer_slot(self, pool: StaticBufferPool, slot_idx: int): def start_onload_to_static(self): """Start async copy from CPU storage to GPU buffer. - Uses the stored _gpu_buffer reference directly instead of accessing - param.data, to avoid any issues with parameter identity changes. + Uses event-based forking to join copy_stream to CUDA graph capture. + This ensures H2D copies are properly captured when recording a graph. IMPORTANT: We must wait for the compute stream before copying, because the previous layer's forward may still be using the buffer (GPU ops are @@ -426,8 +529,14 @@ def start_onload_to_static(self): """ assert self._buffer_pool is not None, "Buffer pool not assigned" - # Wait for compute stream to finish using the buffer before overwriting - self.copy_stream.wait_stream(torch.cuda.current_stream()) + # Track if this prefetch is being captured (for _wait_for_layer logic) + self._prefetch_in_capture = torch.cuda.is_current_stream_capturing() + + # Fork: record event on compute stream, copy_stream waits on it + # This joins copy_stream to any active CUDA graph capture + fork_event = torch.cuda.Event() + torch.cuda.current_stream().record_event(fork_event) + self.copy_stream.wait_event(fork_event) with torch.cuda.stream(self.copy_stream): for name, offloader in self._param_offloaders.items(): @@ -438,6 +547,9 @@ def start_onload_to_static(self): # Async copy from pinned CPU storage to GPU buffer gpu_buffer.copy_(cpu_storage, non_blocking=True) + # Record completion event for _wait_for_layer to use + self._copy_done_event.record(self.copy_stream) + class _BaseParamOffloader(ABC): """Base class for parameter offloading strategies.""" diff --git a/vllm/model_executor/offloader/v2_ops.py b/vllm/model_executor/offloader/v2_ops.py index 28e35ab35f50..509499260c78 100644 --- a/vllm/model_executor/offloader/v2_ops.py +++ b/vllm/model_executor/offloader/v2_ops.py @@ -27,6 +27,18 @@ def set_offloader_instance(offloader: OffloaderV2 | None) -> None: _offloader_instance = offloader +def sync_offloader_before_capture() -> None: + """Sync offloader's copy stream before CUDA graph capture. + + Call this before capturing or replaying CUDA graphs. This ensures + any pre-capture prefetch work is complete before the graph operations. + + Safe to call even if no offloader is active (no-op in that case). + """ + if _offloader_instance is not None: + _offloader_instance.sync_before_graph_capture() + + # --- wait_prefetch op --- @@ -93,6 +105,23 @@ def _start_prefetch_fake( return output_tensor +def join_offloader_after_forward() -> None: + """Join copy_stream after model forward completes. + + Call this after the model forward pass but before CUDA graph capture + ends. This ensures copy_stream is rejoined for any prefetches started + during the forward pass. + + The last layer prefetches a layer that won't have its wait_prefetch + called until the next forward pass. During capture, this leaves + copy_stream unjoined, causing cudaErrorStreamCaptureUnjoined. + + Safe to call even if no offloader is active (no-op in that case). + """ + if _offloader_instance is not None: + _offloader_instance.join_after_forward() + + def register_v2_offloader_ops() -> None: """Register custom ops for V2 offloader. diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 4cb1ff3b9ec7..6b95cdae6461 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -12,6 +12,10 @@ from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import graph_capture, is_global_first_rank from vllm.forward_context import set_forward_context +from vllm.model_executor.offloader.v2_ops import ( + join_offloader_after_forward, + sync_offloader_before_capture, +) from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_attn_metadata @@ -118,6 +122,11 @@ def capture_graph( # Capture the graph. assert num_tokens not in self.graphs graph = torch.cuda.CUDAGraph() + + # Sync offloader's copy stream before capture. + # Ensure any pre-capture prefetches from offloader are complete. + sync_offloader_before_capture() + with ( set_forward_context( attn_metadata, @@ -133,6 +142,10 @@ def capture_graph( positions=positions, inputs_embeds=inputs_embeds, ) + # Join offloader's copy stream after forward to avoid unjoined + # stream error. The last layer's start_prefetch forks copy_stream, + # but wait_prefetch only happens in the next forward pass. + join_offloader_after_forward() self.hidden_states[:num_tokens] = hidden_states self.graphs[num_tokens] = graph @@ -162,6 +175,9 @@ def capture( def run(self, num_tokens: int) -> torch.Tensor: assert num_tokens in self.graphs + # Sync offloader before replay - ensures any external dependencies + # from pre-capture prefetches are satisfied. + sync_offloader_before_capture() self.graphs[num_tokens].replay() assert self.hidden_states is not None return self.hidden_states[:num_tokens] diff --git a/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py b/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py index c4a511778270..4c2f7ffd1b67 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py @@ -6,6 +6,10 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.offloader.v2_ops import ( + join_offloader_after_forward, + sync_offloader_before_capture, +) from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.block_table import BlockTables @@ -88,6 +92,10 @@ def capture_graph( graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, self.pool): generate_fn(num_tokens, attn_metadata, num_tokens_across_dp) + # Join offloader's copy stream after forward to avoid unjoined + # stream error. The last layer's start_prefetch forks copy_stream, + # but wait_prefetch only happens in the next forward pass. + join_offloader_after_forward() self.graphs[num_tokens] = graph @torch.inference_mode() @@ -112,4 +120,7 @@ def capture( def run(self, num_tokens: int) -> None: assert num_tokens in self.graphs + # Sync offloader before replay - ensures any external dependencies + # from pre-capture prefetches are satisfied. + sync_offloader_before_capture() self.graphs[num_tokens].replay() diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index af09129e67b1..332ba7497779 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -20,6 +20,7 @@ override_forward_context, ) from vllm.logger import init_logger +from vllm.model_executor.offloader.v2_ops import sync_offloader_before_capture from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.import_utils import has_deep_gemm @@ -239,6 +240,11 @@ def _capture_ubatch_thread(results, ubatch_metadata): set_graph_pool_id(self.graph_pool) else: set_graph_pool_id(current_platform.graph_pool_handle()) + + # Sync offloader's copy stream before capture. + # Ensure any pre-capture prefetches from offloader are complete. + sync_offloader_before_capture() + with torch.cuda.graph( cudagraph_metadata.cudagraph, stream=compute_stream, @@ -456,6 +462,9 @@ def __call__(self, *args, **kwargs): and cudagraph_runtime_mode is CUDAGraphMode.FULL ): cudagraph_metadata = self.cudagraphs[num_tokens] + # Sync offloader before replay - ensures any external dependencies + # from pre-capture prefetches are satisfied. + sync_offloader_before_capture() cudagraph_metadata.cudagraph.replay() return cudagraph_metadata.outputs else: From 91c180eb8855206d8b4d688fc0834e07a2682db8 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Mon, 26 Jan 2026 17:18:37 -0800 Subject: [PATCH 12/36] clean up unnecessary code Signed-off-by: Ming Yang --- vllm/model_executor/offloader/v2.py | 39 ----------------------------- 1 file changed, 39 deletions(-) diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/v2.py index 7336e4c1314f..620616ec2de0 100644 --- a/vllm/model_executor/offloader/v2.py +++ b/vllm/model_executor/offloader/v2.py @@ -283,52 +283,14 @@ def sync_before_graph_capture(self): Call this: 1. Before capturing a CUDA graph 2. Before replaying a CUDA graph (if prefetches were issued outside) - - Also resets _prefetch_in_capture flags. This is critical for piecewise - cudagraph mode where multiple subgraphs are captured sequentially. - Each subgraph's prefetches must be tracked independently - we can't - wait on events recorded in a different subgraph's capture. """ torch.cuda.current_stream().wait_stream(self.copy_stream) - # Reset flags so only prefetches started in THIS capture are tracked. - # This prevents cross-capture event waits which cause cudaErrorInvalidValue. - for offloader in self.module_offloaders: - offloader._prefetch_in_capture = False - def _start_prefetch(self, layer_idx: int): """Called by custom op - start async copy to static buffer.""" offloader = self.module_offloaders[layer_idx] offloader.start_onload_to_static() - def _join_copy_stream(self, layer_idx: int): - """Called by custom op - join copy_stream back to compute stream. - - Used after the last layer's start_prefetch to ensure copy_stream is - rejoined before CUDA graph capture ends. The start_prefetch forks - copy_stream into the capture, but the corresponding wait_prefetch - only happens in the next forward pass. During capture, this would - leave copy_stream unjoined, causing cudaErrorStreamCaptureUnjoined. - - During capture, we check if the prefetch was started during the same - capture. If not (e.g., from post_init or warmup), we skip the wait - because: - 1. sync_before_graph_capture() ensures pre-capture work is complete - 2. We can't wait on pre-capture events during capture (would cause - cudaErrorStreamCaptureIsolation) - """ - offloader = self.module_offloaders[layer_idx] - - # During capture, skip wait for pre-capture prefetches. - # sync_before_graph_capture() ensures pre-capture work is complete. - if ( - torch.cuda.is_current_stream_capturing() - and not offloader._prefetch_in_capture - ): - return - - torch.cuda.current_stream().wait_event(offloader._copy_done_event) - def join_after_forward(self): """Join copy_stream after model forward completes. @@ -544,7 +506,6 @@ def start_onload_to_static(self): gpu_buffer = offloader._gpu_buffer assert cpu_storage is not None, "CPU storage not initialized" assert gpu_buffer is not None, "GPU buffer not assigned" - # Async copy from pinned CPU storage to GPU buffer gpu_buffer.copy_(cpu_storage, non_blocking=True) # Record completion event for _wait_for_layer to use From 20db2a1ce21dd976f4440f7e43398a719e184ef8 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Mon, 26 Jan 2026 17:42:03 -0800 Subject: [PATCH 13/36] [Core] V2 offloader: simplify API by using get_offloader() directly Remove wrapper functions from v2_ops.py and add methods to BaseOffloader so call sites use get_offloader() directly. Changes: - Add sync_prev_onload(), join_after_forward(), _wait_for_layer(), _start_prefetch() to BaseOffloader with no-op defaults - Remove sync_offloader_before_capture() and join_offloader_after_forward() wrapper functions from v2_ops.py - Update cuda_graph.py, cudagraph_utils.py, gpu_ubatch_wrapper.py, eagle_cudagraph.py to call get_offloader() methods directly - Rename sync_before_graph_capture to sync_prev_onload (more accurate since it's called before both capture and replay) Signed-off-by: Ming Yang --- vllm/compilation/cuda_graph.py | 11 ++-- vllm/model_executor/offloader/base.py | 16 ++++++ vllm/model_executor/offloader/v2.py | 18 ++----- vllm/model_executor/offloader/v2_ops.py | 51 ++----------------- vllm/v1/worker/gpu/cudagraph_utils.py | 11 ++-- .../worker/gpu/spec_decode/eagle_cudagraph.py | 9 ++-- vllm/v1/worker/gpu_ubatch_wrapper.py | 6 +-- 7 files changed, 38 insertions(+), 84 deletions(-) diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index fb1dfe9b2e72..7bada5e7ca3c 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -17,10 +17,7 @@ from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id from vllm.forward_context import BatchDescriptor, get_forward_context from vllm.logger import init_logger -from vllm.model_executor.offloader.v2_ops import ( - join_offloader_after_forward, - sync_offloader_before_capture, -) +from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.utils.torch_utils import current_stream, weak_ref_tensors @@ -272,7 +269,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: # Sync offloader's copy stream before capture. # Ensure any pre-capture prefetches from offloader are complete. - sync_offloader_before_capture() + get_offloader().sync_prev_onload() # mind-exploding: carefully manage the reference and memory. with torch.cuda.graph( @@ -286,7 +283,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: # unjoined stream error. The last layer's start_prefetch # forks copy_stream, but wait_prefetch only happens in # the next forward pass. - join_offloader_after_forward() + get_offloader().join_after_forward() if self.cudagraph_options.weak_ref_output: # by converting it to weak ref, # the original `output` will immediately be released @@ -321,6 +318,6 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None: # Sync offloader before replay - ensures any external dependencies # from pre-capture prefetches are satisfied. - sync_offloader_before_capture() + get_offloader().sync_prev_onload() entry.cudagraph.replay() return entry.output diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index 18e07f08ecd1..c00547113e67 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -71,6 +71,22 @@ def post_init(self): """ return + def sync_prev_onload(self) -> None: # noqa: B027 + """Sync previous onload operations. Override in subclasses.""" + pass + + def join_after_forward(self) -> None: # noqa: B027 + """Join streams after forward. Override in subclasses.""" + pass + + def _wait_for_layer(self, layer_idx: int) -> None: # noqa: B027 + """Wait for layer prefetch. Override in subclasses.""" + pass + + def _start_prefetch(self, layer_idx: int) -> None: # noqa: B027 + """Start layer prefetch. Override in subclasses.""" + pass + class NoopOffloader(BaseOffloader): """No-op offloader that returns modules as-is without any offloading.""" diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/v2.py index 620616ec2de0..8806bfa5f753 100644 --- a/vllm/model_executor/offloader/v2.py +++ b/vllm/model_executor/offloader/v2.py @@ -159,11 +159,6 @@ def __init__( self.buffer_pool: StaticBufferPool | None = None self.total_offloaded_bytes = 0 - # Register this instance for custom ops - from vllm.model_executor.offloader.v2_ops import set_offloader_instance - - set_offloader_instance(self) - def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], @@ -274,15 +269,12 @@ def _wait_for_layer(self, layer_idx: int): # Events used in previous captures can be in invalid state. torch.cuda.current_stream().wait_stream(self.copy_stream) - def sync_before_graph_capture(self): - """Sync copy stream before CUDA graph capture or replay. - - Pre-capture prefetches from warmup must complete before capture. - This method ensures those dependencies are satisfied. + def sync_prev_onload(self): + """Sync previous onload operations. - Call this: - 1. Before capturing a CUDA graph - 2. Before replaying a CUDA graph (if prefetches were issued outside) + Ensures any H2D copies in flight on copy_stream complete before + the compute stream continues. Call this before CUDA graph + capture/replay or when synchronization is needed. """ torch.cuda.current_stream().wait_stream(self.copy_stream) diff --git a/vllm/model_executor/offloader/v2_ops.py b/vllm/model_executor/offloader/v2_ops.py index 509499260c78..2f82fbfbadc2 100644 --- a/vllm/model_executor/offloader/v2_ops.py +++ b/vllm/model_executor/offloader/v2_ops.py @@ -8,37 +8,11 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import torch +from vllm.model_executor.offloader.base import get_offloader from vllm.utils.torch_utils import direct_register_custom_op -if TYPE_CHECKING: - from vllm.model_executor.offloader.v2 import OffloaderV2 - -# Global reference to the offloader instance, set during OffloaderV2.__init__ -_offloader_instance: OffloaderV2 | None = None - - -def set_offloader_instance(offloader: OffloaderV2 | None) -> None: - """Set the global offloader instance for custom ops to use.""" - global _offloader_instance - _offloader_instance = offloader - - -def sync_offloader_before_capture() -> None: - """Sync offloader's copy stream before CUDA graph capture. - - Call this before capturing or replaying CUDA graphs. This ensures - any pre-capture prefetch work is complete before the graph operations. - - Safe to call even if no offloader is active (no-op in that case). - """ - if _offloader_instance is not None: - _offloader_instance.sync_before_graph_capture() - - # --- wait_prefetch op --- @@ -59,8 +33,7 @@ def _wait_prefetch_impl( Returns: input_tensor unchanged, but creates data dependency for torch.compile. """ - if _offloader_instance is not None: - _offloader_instance._wait_for_layer(layer_idx) + get_offloader()._wait_for_layer(layer_idx) return input_tensor @@ -92,8 +65,7 @@ def _start_prefetch_impl( Returns: output_tensor unchanged, creating data dependency for torch.compile. """ - if _offloader_instance is not None: - _offloader_instance._start_prefetch(layer_idx) + get_offloader()._start_prefetch(layer_idx) return output_tensor @@ -105,23 +77,6 @@ def _start_prefetch_fake( return output_tensor -def join_offloader_after_forward() -> None: - """Join copy_stream after model forward completes. - - Call this after the model forward pass but before CUDA graph capture - ends. This ensures copy_stream is rejoined for any prefetches started - during the forward pass. - - The last layer prefetches a layer that won't have its wait_prefetch - called until the next forward pass. During capture, this leaves - copy_stream unjoined, causing cudaErrorStreamCaptureUnjoined. - - Safe to call even if no offloader is active (no-op in that case). - """ - if _offloader_instance is not None: - _offloader_instance.join_after_forward() - - def register_v2_offloader_ops() -> None: """Register custom ops for V2 offloader. diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 6b95cdae6461..c54fa09fff72 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -12,10 +12,7 @@ from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import graph_capture, is_global_first_rank from vllm.forward_context import set_forward_context -from vllm.model_executor.offloader.v2_ops import ( - join_offloader_after_forward, - sync_offloader_before_capture, -) +from vllm.model_executor.offloader.base import get_offloader from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_attn_metadata @@ -125,7 +122,7 @@ def capture_graph( # Sync offloader's copy stream before capture. # Ensure any pre-capture prefetches from offloader are complete. - sync_offloader_before_capture() + get_offloader().sync_prev_onload() with ( set_forward_context( @@ -145,7 +142,7 @@ def capture_graph( # Join offloader's copy stream after forward to avoid unjoined # stream error. The last layer's start_prefetch forks copy_stream, # but wait_prefetch only happens in the next forward pass. - join_offloader_after_forward() + get_offloader().join_after_forward() self.hidden_states[:num_tokens] = hidden_states self.graphs[num_tokens] = graph @@ -177,7 +174,7 @@ def run(self, num_tokens: int) -> torch.Tensor: assert num_tokens in self.graphs # Sync offloader before replay - ensures any external dependencies # from pre-capture prefetches are satisfied. - sync_offloader_before_capture() + get_offloader().sync_prev_onload() self.graphs[num_tokens].replay() assert self.hidden_states is not None return self.hidden_states[:num_tokens] diff --git a/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py b/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py index 4c2f7ffd1b67..355140cfcece 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py @@ -6,10 +6,7 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode -from vllm.model_executor.offloader.v2_ops import ( - join_offloader_after_forward, - sync_offloader_before_capture, -) +from vllm.model_executor.offloader.base import get_offloader from vllm.v1.attention.backend import AttentionMetadataBuilder from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.block_table import BlockTables @@ -95,7 +92,7 @@ def capture_graph( # Join offloader's copy stream after forward to avoid unjoined # stream error. The last layer's start_prefetch forks copy_stream, # but wait_prefetch only happens in the next forward pass. - join_offloader_after_forward() + get_offloader().join_after_forward() self.graphs[num_tokens] = graph @torch.inference_mode() @@ -122,5 +119,5 @@ def run(self, num_tokens: int) -> None: assert num_tokens in self.graphs # Sync offloader before replay - ensures any external dependencies # from pre-capture prefetches are satisfied. - sync_offloader_before_capture() + get_offloader().sync_prev_onload() self.graphs[num_tokens].replay() diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 332ba7497779..18e99107ad25 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -20,7 +20,7 @@ override_forward_context, ) from vllm.logger import init_logger -from vllm.model_executor.offloader.v2_ops import sync_offloader_before_capture +from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.import_utils import has_deep_gemm @@ -243,7 +243,7 @@ def _capture_ubatch_thread(results, ubatch_metadata): # Sync offloader's copy stream before capture. # Ensure any pre-capture prefetches from offloader are complete. - sync_offloader_before_capture() + get_offloader().sync_prev_onload() with torch.cuda.graph( cudagraph_metadata.cudagraph, @@ -464,7 +464,7 @@ def __call__(self, *args, **kwargs): cudagraph_metadata = self.cudagraphs[num_tokens] # Sync offloader before replay - ensures any external dependencies # from pre-capture prefetches are satisfied. - sync_offloader_before_capture() + get_offloader().sync_prev_onload() cudagraph_metadata.cudagraph.replay() return cudagraph_metadata.outputs else: From 2debd782f06a5731359728ed77cd6a7299ef4e49 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Mon, 26 Jan 2026 18:47:29 -0800 Subject: [PATCH 14/36] [Core] V2 offloader: fix CUDA graph bugs in Eagle and UBatch Fix three bugs in V2 offloader CUDA graph support: 1. EagleCudaGraphManager: add sync_prev_onload() before capture to ensure warmup prefetches complete before graph capture starts 2. UBatchWrapper: add join_after_forward() inside capture block to avoid unjoined stream error from last layer's start_prefetch 3. Config validation: require offload_prefetch_step >= 1 when V2 offloading is enabled to prevent ZeroDivisionError in StaticBufferPool Signed-off-by: Ming Yang --- vllm/config/offload.py | 21 +++++++++++-------- .../worker/gpu/spec_decode/eagle_cudagraph.py | 5 +++++ vllm/v1/worker/gpu_ubatch_wrapper.py | 4 ++++ 3 files changed, 21 insertions(+), 9 deletions(-) diff --git a/vllm/config/offload.py b/vllm/config/offload.py index 9d7f3950f19a..91e241866fec 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -50,15 +50,18 @@ class OffloadConfig: @model_validator(mode="after") def validate_offload_config(self) -> "OffloadConfig": - """Validate that offload_num_in_group <= offload_group_size.""" - if ( - self.offload_group_size > 0 - and self.offload_num_in_group > self.offload_group_size - ): - raise ValueError( - f"offload_num_in_group ({self.offload_num_in_group}) must be " - f"<= offload_group_size ({self.offload_group_size})" - ) + """Validate offload configuration constraints.""" + if self.offload_group_size > 0: + if self.offload_num_in_group > self.offload_group_size: + raise ValueError( + f"offload_num_in_group ({self.offload_num_in_group}) must be " + f"<= offload_group_size ({self.offload_group_size})" + ) + if self.offload_prefetch_step < 1: + raise ValueError( + f"offload_prefetch_step ({self.offload_prefetch_step}) must be " + f">= 1 when V2 offloading is enabled (offload_group_size > 0)" + ) return self def compute_hash(self) -> str: diff --git a/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py b/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py index 355140cfcece..05f5736664e2 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py @@ -87,6 +87,11 @@ def capture_graph( # Capture the graph. assert num_tokens not in self.graphs graph = torch.cuda.CUDAGraph() + + # Sync offloader's copy stream before capture. + # Ensure any pre-capture prefetches from offloader are complete. + get_offloader().sync_prev_onload() + with torch.cuda.graph(graph, self.pool): generate_fn(num_tokens, attn_metadata, num_tokens_across_dp) # Join offloader's copy stream after forward to avoid unjoined diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 18e99107ad25..6d8c18a7b1f1 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -256,6 +256,10 @@ def _capture_ubatch_thread(results, ubatch_metadata): sorted_results = [value for position, value in sorted(results)] result = torch.cat(sorted_results, dim=0) cudagraph_metadata.outputs = result + # Join offloader's copy stream after forward to avoid unjoined + # stream error. The last layer's start_prefetch forks copy_stream, + # but wait_prefetch only happens in the next forward pass. + get_offloader().join_after_forward() self.cudagraphs[num_tokens] = cudagraph_metadata return cudagraph_metadata.outputs From f536bc4c33f0c2d11816be96f58b7461d4d12093 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Thu, 19 Feb 2026 13:58:11 -0800 Subject: [PATCH 15/36] minor param order adjustmetn Signed-off-by: Ming Yang --- vllm/engine/arg_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 630d8cf75f39..2e874e2d3ec5 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -440,10 +440,10 @@ class EngineArgs: disable_cascade_attn: bool = ModelConfig.disable_cascade_attn swap_space: float = CacheConfig.swap_space cpu_offload_gb: float = OffloadConfig.cpu_offload_gb + cpu_offload_params: set[str] = get_field(OffloadConfig, "cpu_offload_params") offload_group_size: int = OffloadConfig.offload_group_size offload_num_in_group: int = OffloadConfig.offload_num_in_group offload_prefetch_step: int = OffloadConfig.offload_prefetch_step - cpu_offload_params: set[str] = get_field(OffloadConfig, "cpu_offload_params") gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None From 021acc1a892fc4e6fbc391230cfce37350d67e2f Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Thu, 19 Feb 2026 14:39:18 -0800 Subject: [PATCH 16/36] Keep cpu_offload_params on CacheConfig as deprecated field Adding back cpu_offload_params to CacheConfig to avoid breaking downstream code that constructs CacheConfig(cpu_offload_params=...) directly. The field defaults to an empty set and directs users to OffloadConfig.cpu_offload_params. Signed-off-by: Ming Yang --- vllm/config/cache.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 6009479ddf58..bdff700ef4d6 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -101,6 +101,12 @@ class CacheConfig: DEPRECATED: This field is deprecated and will be removed in a future release. Please use OffloadConfig.cpu_offload_gb instead. """ + cpu_offload_params: set[str] = Field(default_factory=set) + """The set of parameter name segments to target for CPU offloading. + + DEPRECATED: This field is deprecated and will be removed in a future + release. Please use OffloadConfig.cpu_offload_params instead. + """ calculate_kv_scales: bool = False """This enables dynamic calculation of `k_scale` and `v_scale` when kv_cache_dtype is fp8. If `False`, the scales will be loaded from the model From 0525c42a11bc890de85f515fabc28ae1faa201a7 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Thu, 19 Feb 2026 15:12:16 -0800 Subject: [PATCH 17/36] Fix OffloadConfig.compute_hash() returning constant hash OffloaderV2 patches module forwards and inserts custom ops (wait_prefetch, start_prefetch) into the computation graph, so offload settings must be part of the compilation cache key. The previous constant hash caused stale cache hits when switching between offloaded and non-offloaded runs of the same model. Signed-off-by: Ming Yang --- vllm/config/offload.py | 29 ++++++++++++++--------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/vllm/config/offload.py b/vllm/config/offload.py index aef484ea8aad..0b43f12b3840 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -2,13 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Configuration for model weight offloading.""" -from typing import Any - from pydantic import Field, model_validator from pydantic.dataclasses import dataclass from vllm.config.utils import config -from vllm.utils.hashing import safe_hash @config @@ -78,18 +75,20 @@ def validate_offload_config(self) -> "OffloadConfig": def compute_hash(self) -> str: """ - WARNING: Whenever a new field is added to this config, - ensure that it is included in the factors list if - it affects the computation graph. + Provide a hash that uniquely identifies all the offload configs. - Provide a hash that uniquely identifies all the configs - that affect the structure of the computation - graph from input ids/embeddings to the final hidden states, - excluding anything before input ids/embeddings and after - the final hidden states. + All fields are included because OffloaderV2 patches module + forwards and inserts custom ops (wait_prefetch, start_prefetch) + into the computation graph. Changing any offload setting can + alter which layers are hooked and how prefetch indices are + computed, so the compilation cache must distinguish them. """ - # Offload settings don't affect the computation graph structure, - # only the memory layout and transfer patterns. - factors: list[Any] = [] - hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() + # OffloaderV2 (offload_group_size > 0) patches module forwards + # and inserts custom ops (wait_prefetch, start_prefetch) into the + # computation graph, so all offload settings must be part of the + # cache key to avoid stale compilation cache hits. + from vllm.config.utils import get_hash_factors, hash_factors + + factors = get_hash_factors(self, ignored_factors=set()) + hash_str = hash_factors(factors) return hash_str From 65067062b9f7614be0a5061f7556dcbc15b1b656 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Thu, 19 Feb 2026 19:41:10 -0800 Subject: [PATCH 18/36] Fix OffloaderV2 accuracy bug: re-pin CPU storage after process_weights_after_loading After process_weights_after_loading, device_loading_context creates non-pinned CPU tensors via `p.data = p.data.to("cpu")`. When sync_cpu_storage() or assign_static_buffer() adopted these non-pinned tensors as _cpu_storage, subsequent non_blocking=True H2D copies would trigger an implicit CUDA stream synchronization (cudaMemcpyAsync from pageable memory). This broke the event-based fork synchronization between the compute and copy streams, allowing the copy to overwrite the GPU buffer while the compute stream was still reading from it. Fix by extracting _update_cpu_storage_from_param() that re-pins CPU tensors when needed, and adding an assertion in start_onload_to_static to catch non-pinned CPU storage early. Signed-off-by: Ming Yang --- vllm/model_executor/offloader/v2.py | 56 +++++++++++++++++++++-------- 1 file changed, 41 insertions(+), 15 deletions(-) diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/v2.py index 8806bfa5f753..b60d6029a59c 100644 --- a/vllm/model_executor/offloader/v2.py +++ b/vllm/model_executor/offloader/v2.py @@ -498,6 +498,11 @@ def start_onload_to_static(self): gpu_buffer = offloader._gpu_buffer assert cpu_storage is not None, "CPU storage not initialized" assert gpu_buffer is not None, "GPU buffer not assigned" + assert not is_pin_memory_available() or cpu_storage.is_pinned(), \ + f"CPU storage for {name} is not pinned! " \ + "non_blocking=True H2D copy from non-pinned memory " \ + "causes stream synchronization that breaks " \ + "event-based fork synchronization." gpu_buffer.copy_(cpu_storage, non_blocking=True) # Record completion event for _wait_for_layer to use @@ -599,6 +604,40 @@ def _offload_to_cpu_internal(self): # and frees GPU memory when the original GPU tensor is garbage collected param.data = self._cpu_storage + def _update_cpu_storage_from_param(self) -> None: + """Update _cpu_storage from current param.data, ensuring pinned memory. + + After process_weights_after_loading, device_loading_context creates + non-pinned CPU tensors via `p.data = p.data.to("cpu")`. Using + non-pinned memory with `copy_(src, non_blocking=True)` causes CUDA to + perform a stream synchronization before the copy, breaking the + event-based fork synchronization and potentially allowing the copy + to overwrite the GPU buffer while the compute stream still reads it. + + This method ensures _cpu_storage always uses pinned memory when + available, re-pinning if necessary. + """ + param = self._param + + if param.data.device.type == "cpu": + if is_pin_memory_available() and not param.data.is_pinned(): + pinned = torch.empty_strided( + size=param.data.size(), + stride=param.data.stride(), + dtype=param.data.dtype, + layout=param.data.layout, + device="cpu", + pin_memory=True, + ) + pinned.copy_(param.data) + self._cpu_storage = pinned + else: + self._cpu_storage = param.data + else: + # param.data is on GPU - copy to existing CPU storage + assert self._cpu_storage is not None + self._cpu_storage.copy_(param.data) + def assign_static_buffer(self, gpu_buffer: torch.Tensor) -> None: """Point parameter data to GPU static buffer. @@ -622,12 +661,7 @@ def assign_static_buffer(self, gpu_buffer: torch.Tensor) -> None: # 1. process_weights_after_loading may transform weights (quantization) # 2. device_loading_context creates NEW CPU tensors when moving back # 3. Our old _cpu_storage would have pre-processed or stale data - if param.data.device.type == "cpu": - # param.data is already on CPU - use it as our CPU storage - self._cpu_storage = param.data - else: - # param.data is on GPU - copy to our CPU storage - self._cpu_storage.copy_(param.data) + self._update_cpu_storage_from_param() # Store reference to GPU buffer for use in start_onload self._gpu_buffer = gpu_buffer @@ -644,15 +678,7 @@ def sync_cpu_storage(self) -> None: 2. device_loading_context creates NEW CPU tensors when moving back 3. Our old _cpu_storage would have pre-processed or stale data """ - param = self._param - - if param.data.device.type == "cpu": - # param.data is already on CPU - use it as our CPU storage - self._cpu_storage = param.data - else: - # param.data is on GPU - copy to existing CPU storage - assert self._cpu_storage is not None - self._cpu_storage.copy_(param.data) + self._update_cpu_storage_from_param() def post_init(self): """No-op: offloading done in offload_to_cpu/assign_static_buffer.""" From 55a9934d96ac95567faa8ef3c04a81a66a55b130 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Thu, 19 Feb 2026 19:52:49 -0800 Subject: [PATCH 19/36] pre-commit fixes Signed-off-by: Ming Yang --- vllm/model_executor/models/utils.py | 1 - vllm/model_executor/offloader/uva.py | 10 +++------- vllm/model_executor/offloader/v2.py | 9 +++++---- 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 537c55d80e46..05f061850218 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -12,7 +12,6 @@ from torch.nn.modules.module import register_module_module_registration_hook from transformers import PretrainedConfig -import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import ( get_tensor_model_parallel_rank, diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index ead4ca952b4e..2c6535b0fd49 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -48,8 +48,7 @@ def __init__( and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY ) self.uva_offloading = ( - is_uva_available() - and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA + is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA ) def wrap_modules( @@ -63,9 +62,7 @@ def wrap_modules( Note: UVA offloading operates at module level, so submodule_accessor and whitelist_param_names_creator are ignored. """ - modules = [ - self._maybe_offload_to_cpu(module) for module in modules_generator - ] + modules = [self._maybe_offload_to_cpu(module) for module in modules_generator] if self.cpu_offload_bytes > 0: logger.info( "Total CPU offloaded parameters: %s", @@ -101,8 +98,7 @@ def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight" # but not "mlp.experts.w2_weight_scale" should_offload = any( - f".{param}." in f".{name}." - for param in self.cpu_offload_params + f".{param}." in f".{name}." for param in self.cpu_offload_params ) if not should_offload: continue diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/v2.py index b60d6029a59c..e7d4c13201e9 100644 --- a/vllm/model_executor/offloader/v2.py +++ b/vllm/model_executor/offloader/v2.py @@ -498,11 +498,12 @@ def start_onload_to_static(self): gpu_buffer = offloader._gpu_buffer assert cpu_storage is not None, "CPU storage not initialized" assert gpu_buffer is not None, "GPU buffer not assigned" - assert not is_pin_memory_available() or cpu_storage.is_pinned(), \ - f"CPU storage for {name} is not pinned! " \ - "non_blocking=True H2D copy from non-pinned memory " \ - "causes stream synchronization that breaks " \ + assert not is_pin_memory_available() or cpu_storage.is_pinned(), ( + f"CPU storage for {name} is not pinned! " + "non_blocking=True H2D copy from non-pinned memory " + "causes stream synchronization that breaks " "event-based fork synchronization." + ) gpu_buffer.copy_(cpu_storage, non_blocking=True) # Record completion event for _wait_for_layer to use From 8478e5568d044f7655d06ec09a37518e37df05e5 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sun, 22 Feb 2026 11:06:14 -0800 Subject: [PATCH 20/36] Fix OffloaderV2 crash with dotted parameter names _BaseParamOffloader._param used plain getattr() which doesn't traverse dotted paths like 'self_attn.qkv_proj.weight'. Models that don't provide a whitelist_param_names_creator (e.g. Llama) hit this because named_parameters() yields fully-qualified names. Signed-off-by: Ming Yang --- vllm/model_executor/offloader/v2.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/v2.py index e7d4c13201e9..17ee301d409b 100644 --- a/vllm/model_executor/offloader/v2.py +++ b/vllm/model_executor/offloader/v2.py @@ -12,6 +12,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Generator from dataclasses import dataclass +from typing import Any import torch import torch.nn as nn @@ -535,8 +536,15 @@ def __init__(self, module: nn.Module, param_name: str): @property def _param(self) -> nn.Parameter: - """Get the parameter being offloaded.""" - return getattr(self._module, self._param_name) + """Get the parameter being offloaded. + + Supports dotted names (e.g. 'self_attn.qkv_proj.weight') by + traversing the module hierarchy. + """ + obj: Any = self._module + for attr in self._param_name.split("."): + obj = getattr(obj, attr) + return obj def post_init(self): """Initialize offloading (move parameter to storage).""" From 4520428385cf3fd6c23edd7a6027ba6fb876650f Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sun, 22 Feb 2026 11:28:17 -0800 Subject: [PATCH 21/36] Specify removal version for deprecated CacheConfig offload fields Per vllm deprecation policy, deprecated parameters must specify which version they will be removed in. Updated cpu_offload_gb and cpu_offload_params on CacheConfig to target v0.16 for removal. Signed-off-by: Ming Yang --- vllm/config/cache.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index bdff700ef4d6..85e6f5e4faaa 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -98,14 +98,14 @@ class CacheConfig: Note that this requires fast CPU-GPU interconnect, as part of the model is loaded from CPU memory to GPU memory on the fly in each model forward pass. - DEPRECATED: This field is deprecated and will be removed in a future - release. Please use OffloadConfig.cpu_offload_gb instead. + DEPRECATED: This field is deprecated and will be removed in v0.16. + Please use OffloadConfig.cpu_offload_gb instead. """ cpu_offload_params: set[str] = Field(default_factory=set) """The set of parameter name segments to target for CPU offloading. - DEPRECATED: This field is deprecated and will be removed in a future - release. Please use OffloadConfig.cpu_offload_params instead. + DEPRECATED: This field is deprecated and will be removed in v0.16. + Please use OffloadConfig.cpu_offload_params instead. """ calculate_kv_scales: bool = False """This enables dynamic calculation of `k_scale` and `v_scale` when From fcaa9daf440bb1eee6a69d6028ddd91866c41411 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sun, 22 Feb 2026 11:30:04 -0800 Subject: [PATCH 22/36] Use @config decorator alone for OffloadConfig The @config decorator already applies pydantic dataclass internally, so the explicit @dataclass was redundant. Aligns with how other config classes (e.g. CacheConfig) are defined. Signed-off-by: Ming Yang --- vllm/config/offload.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/config/offload.py b/vllm/config/offload.py index 0b43f12b3840..4f6b3db587cb 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -3,13 +3,11 @@ """Configuration for model weight offloading.""" from pydantic import Field, model_validator -from pydantic.dataclasses import dataclass from vllm.config.utils import config @config -@dataclass class OffloadConfig: """Configuration for model weight offloading to CPU. From 79ec7d56544f093fe7b885750163aaa6ca86a9cb Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sun, 22 Feb 2026 12:08:23 -0800 Subject: [PATCH 23/36] Add explicit offload_backend selector with nested sub-configs Introduce OffloadBackend ("auto"|"uva"|"prefetch"), UVAOffloadConfig, and PrefetchOffloadConfig as nested sub-configs of OffloadConfig. The --offload-backend CLI flag lets users explicitly choose a backend instead of relying on implicit field-based detection. Signed-off-by: Ming Yang --- vllm/config/__init__.py | 10 ++- vllm/config/cache.py | 4 +- vllm/config/offload.py | 93 +++++++++++++++++++++------ vllm/engine/arg_utils.py | 44 ++++++++----- vllm/model_executor/offloader/base.py | 33 ++++++---- vllm/v1/worker/gpu_model_runner.py | 2 +- 6 files changed, 136 insertions(+), 50 deletions(-) diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 6c668cdadd07..452fb046660a 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -24,7 +24,12 @@ ) from vllm.config.multimodal import MultiModalConfig from vllm.config.observability import ObservabilityConfig -from vllm.config.offload import OffloadConfig +from vllm.config.offload import ( + OffloadBackend, + OffloadConfig, + PrefetchOffloadConfig, + UVAOffloadConfig, +) from vllm.config.parallel import EPLBConfig, ParallelConfig from vllm.config.pooler import PoolerConfig from vllm.config.profiler import ProfilerConfig @@ -87,7 +92,10 @@ # From vllm.config.observability "ObservabilityConfig", # From vllm.config.offload + "OffloadBackend", "OffloadConfig", + "PrefetchOffloadConfig", + "UVAOffloadConfig", # From vllm.config.parallel "EPLBConfig", "ParallelConfig", diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 85e6f5e4faaa..2e25435da83d 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -99,13 +99,13 @@ class CacheConfig: loaded from CPU memory to GPU memory on the fly in each model forward pass. DEPRECATED: This field is deprecated and will be removed in v0.16. - Please use OffloadConfig.cpu_offload_gb instead. + Please use OffloadConfig.uva.cpu_offload_gb instead. """ cpu_offload_params: set[str] = Field(default_factory=set) """The set of parameter name segments to target for CPU offloading. DEPRECATED: This field is deprecated and will be removed in v0.16. - Please use OffloadConfig.cpu_offload_params instead. + Please use OffloadConfig.uva.cpu_offload_params instead. """ calculate_kv_scales: bool = False """This enables dynamic calculation of `k_scale` and `v_scale` when diff --git a/vllm/config/offload.py b/vllm/config/offload.py index 4f6b3db587cb..d92a3656eaed 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -2,18 +2,22 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Configuration for model weight offloading.""" +import warnings +from typing import Literal + from pydantic import Field, model_validator from vllm.config.utils import config +OffloadBackend = Literal["auto", "uva", "prefetch"] + @config -class OffloadConfig: - """Configuration for model weight offloading to CPU. +class UVAOffloadConfig: + """Configuration for UVA (Unified Virtual Addressing) CPU offloading. - This controls how model parameters are offloaded to CPU memory to reduce - GPU memory usage, at the cost of additional CPU-GPU transfers during - inference. + Uses zero-copy access from CPU-pinned memory. Simple but requires + fast CPU-GPU interconnect. """ cpu_offload_gb: float = Field(default=0, ge=0) @@ -39,36 +43,89 @@ class OffloadConfig: This allows distinguishing parameters like "w2_weight" and "w2_weight_scale". """ + +@config +class PrefetchOffloadConfig: + """Configuration for prefetch-based CPU offloading. + + Groups layers and uses async H2D prefetch to hide transfer latency. + """ + offload_group_size: int = Field(default=0, ge=0) - """Advanced CPU offloading (V2): Group every N layers together. Offload last - `offload_num_in_group` layers of each group. Default is 0 (disabled). + """Group every N layers together. Offload last `offload_num_in_group` + layers of each group. Default is 0 (disabled). Example: group_size=8, num_in_group=2 offloads layers 6,7,14,15,22,23,... Unlike cpu_offload_gb, this uses explicit async prefetching to hide transfer latency. """ offload_num_in_group: int = Field(default=1, ge=1) - """Advanced CPU offloading (V2): Number of layers to offload per group. + """Number of layers to offload per group. Must be <= offload_group_size. Default is 1.""" offload_prefetch_step: int = Field(default=1, ge=0) - """Advanced CPU offloading (V2): Number of layers to prefetch ahead. + """Number of layers to prefetch ahead. Higher values hide more latency but use more GPU memory. Default is 1.""" + +@config +class OffloadConfig: + """Configuration for model weight offloading to reduce GPU memory usage.""" + + offload_backend: OffloadBackend = "auto" + """The backend for weight offloading. Options: + - "auto": Selects based on which sub-config has non-default values + (prefetch if offload_group_size > 0, uva if cpu_offload_gb > 0). + - "uva": UVA (Unified Virtual Addressing) zero-copy offloading. + - "prefetch": Async prefetch with group-based layer offloading. + """ + + uva: UVAOffloadConfig = Field(default_factory=UVAOffloadConfig) + """Parameters for UVA offloading backend.""" + + prefetch: PrefetchOffloadConfig = Field(default_factory=PrefetchOffloadConfig) + """Parameters for prefetch offloading backend.""" + @model_validator(mode="after") def validate_offload_config(self) -> "OffloadConfig": """Validate offload configuration constraints.""" - if self.offload_group_size > 0: - if self.offload_num_in_group > self.offload_group_size: + if self.offload_backend == "prefetch" or self.prefetch.offload_group_size > 0: + if self.prefetch.offload_num_in_group > self.prefetch.offload_group_size: raise ValueError( - f"offload_num_in_group ({self.offload_num_in_group}) must be " - f"<= offload_group_size ({self.offload_group_size})" + f"offload_num_in_group ({self.prefetch.offload_num_in_group})" + f" must be <= offload_group_size" + f" ({self.prefetch.offload_group_size})" ) - if self.offload_prefetch_step < 1: + if self.prefetch.offload_prefetch_step < 1: raise ValueError( - f"offload_prefetch_step ({self.offload_prefetch_step}) must be " - f">= 1 when V2 offloading is enabled (offload_group_size > 0)" + f"offload_prefetch_step" + f" ({self.prefetch.offload_prefetch_step})" + f" must be >= 1 when prefetch offloading is enabled" + f" (offload_group_size > 0)" ) + + # Warn if both backends have non-default values + uva_active = self.uva.cpu_offload_gb > 0 + prefetch_active = self.prefetch.offload_group_size > 0 + if self.offload_backend == "uva" and prefetch_active: + warnings.warn( + "Prefetch offload fields are set but offload_backend='uva'. " + "Prefetch settings will be ignored.", + stacklevel=2, + ) + elif self.offload_backend == "prefetch" and uva_active: + warnings.warn( + "UVA offload fields are set but offload_backend='prefetch'. " + "UVA settings will be ignored.", + stacklevel=2, + ) + elif self.offload_backend == "auto" and uva_active and prefetch_active: + warnings.warn( + "Both UVA and prefetch offload fields are set with " + "offload_backend='auto'. Prefetch backend will be selected. " + "Set offload_backend explicitly to suppress this warning.", + stacklevel=2, + ) return self def compute_hash(self) -> str: @@ -81,10 +138,6 @@ def compute_hash(self) -> str: alter which layers are hooked and how prefetch indices are computed, so the compilation cache must distinguish them. """ - # OffloaderV2 (offload_group_size > 0) patches module forwards - # and inserts custom ops (wait_prefetch, start_prefetch) into the - # computation graph, so all offload settings must be part of the - # cache key to avoid stale compilation cache hits. from vllm.config.utils import get_hash_factors, hash_factors factors = get_hash_factors(self, ignored_factors=set()) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 2e874e2d3ec5..6c9f35eac026 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -51,10 +51,12 @@ OffloadConfig, ParallelConfig, PoolerConfig, + PrefetchOffloadConfig, ProfilerConfig, SchedulerConfig, SpeculativeConfig, StructuredOutputsConfig, + UVAOffloadConfig, VllmConfig, WeightTransferConfig, get_attr_docs, @@ -439,11 +441,12 @@ class EngineArgs: disable_sliding_window: bool = ModelConfig.disable_sliding_window disable_cascade_attn: bool = ModelConfig.disable_cascade_attn swap_space: float = CacheConfig.swap_space - cpu_offload_gb: float = OffloadConfig.cpu_offload_gb - cpu_offload_params: set[str] = get_field(OffloadConfig, "cpu_offload_params") - offload_group_size: int = OffloadConfig.offload_group_size - offload_num_in_group: int = OffloadConfig.offload_num_in_group - offload_prefetch_step: int = OffloadConfig.offload_prefetch_step + offload_backend: str = OffloadConfig.offload_backend + cpu_offload_gb: float = UVAOffloadConfig.cpu_offload_gb + cpu_offload_params: set[str] = get_field(UVAOffloadConfig, "cpu_offload_params") + offload_group_size: int = PrefetchOffloadConfig.offload_group_size + offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group + offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None @@ -978,24 +981,30 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: # Model weight offload related configs offload_kwargs = get_kwargs(OffloadConfig) + uva_kwargs = get_kwargs(UVAOffloadConfig) + prefetch_kwargs = get_kwargs(PrefetchOffloadConfig) offload_group = parser.add_argument_group( title="OffloadConfig", description=OffloadConfig.__doc__, ) offload_group.add_argument( - "--cpu-offload-gb", **offload_kwargs["cpu_offload_gb"] + "--offload-backend", **offload_kwargs["offload_backend"] ) + offload_group.add_argument("--cpu-offload-gb", **uva_kwargs["cpu_offload_gb"]) offload_group.add_argument( - "--offload-group-size", **offload_kwargs["offload_group_size"] + "--cpu-offload-params", **uva_kwargs["cpu_offload_params"] ) offload_group.add_argument( - "--offload-num-in-group", **offload_kwargs["offload_num_in_group"] + "--offload-group-size", + **prefetch_kwargs["offload_group_size"], ) offload_group.add_argument( - "--offload-prefetch-step", **offload_kwargs["offload_prefetch_step"] + "--offload-num-in-group", + **prefetch_kwargs["offload_num_in_group"], ) offload_group.add_argument( - "--cpu-offload-params", **offload_kwargs["cpu_offload_params"] + "--offload-prefetch-step", + **prefetch_kwargs["offload_prefetch_step"], ) # Multimodal related configs @@ -1846,11 +1855,16 @@ def create_engine_config( ) offload_config = OffloadConfig( - cpu_offload_gb=self.cpu_offload_gb, - offload_group_size=self.offload_group_size, - offload_num_in_group=self.offload_num_in_group, - offload_prefetch_step=self.offload_prefetch_step, - cpu_offload_params=self.cpu_offload_params, + offload_backend=self.offload_backend, + uva=UVAOffloadConfig( + cpu_offload_gb=self.cpu_offload_gb, + cpu_offload_params=self.cpu_offload_params, + ), + prefetch=PrefetchOffloadConfig( + offload_group_size=self.offload_group_size, + offload_num_in_group=self.offload_num_in_group, + offload_prefetch_step=self.offload_prefetch_step, + ), ) config = VllmConfig( diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index a0e840dcdecc..a930a71df7ee 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -120,25 +120,36 @@ def set_offloader(instance: BaseOffloader) -> None: def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: """Create an offloader based on the offload configuration. - Priority: V2 offloading if configured, else UVA, else noop. + Uses the explicit ``offload_backend`` selector. When set to ``"auto"``, + selects prefetch if ``offload_group_size > 0``, UVA if + ``cpu_offload_gb > 0``, otherwise noop. """ from vllm.model_executor.offloader.uva import UVAOffloader from vllm.model_executor.offloader.v2 import OffloaderV2 - if offload_config.offload_group_size > 0: - # Use V2 offloading + backend = offload_config.offload_backend + uva = offload_config.uva + prefetch = offload_config.prefetch + + if backend == "auto": + if prefetch.offload_group_size > 0: + backend = "prefetch" + elif uva.cpu_offload_gb > 0: + backend = "uva" + else: + return NoopOffloader() + + if backend == "prefetch": return OffloaderV2( - group_size=offload_config.offload_group_size, - num_in_group=offload_config.offload_num_in_group, - prefetch_step=offload_config.offload_prefetch_step, + group_size=prefetch.offload_group_size, + num_in_group=prefetch.offload_num_in_group, + prefetch_step=prefetch.offload_prefetch_step, mode="cpu", ) - elif offload_config.cpu_offload_gb > 0: - # Use UVA offloading (legacy) + elif backend == "uva": return UVAOffloader( - cpu_offload_max_bytes=int(offload_config.cpu_offload_gb * 1024**3), - cpu_offload_params=offload_config.cpu_offload_params, + cpu_offload_max_bytes=int(uva.cpu_offload_gb * 1024**3), + cpu_offload_params=uva.cpu_offload_params, ) else: - # No offloading return NoopOffloader() diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 333c2cdd008a..cd2606ad214c 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -5727,7 +5727,7 @@ def may_reinitialize_input_batch( if block_sizes != [self.cache_config.block_size] or kernel_block_sizes != [ self.cache_config.block_size ]: - assert self.offload_config.cpu_offload_gb == 0, ( + assert self.offload_config.uva.cpu_offload_gb == 0, ( "Cannot re-initialize the input batch when CPU weight " "offloading is enabled. See https://github.com/vllm-project/vllm/pull/18298 " # noqa: E501 "for more details." From 5fe6d7c1edff059269f321bfb1bbaf2e47775d90 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sun, 22 Feb 2026 12:13:14 -0800 Subject: [PATCH 24/36] Rename OffloaderV2 to PrefetchOffloader Rephrase "V2 offloading" as "prefetching-based offloading" throughout. The two backends (UVA and prefetch) serve different use cases rather than being iterations of the same approach. Renames: - OffloaderV2 -> PrefetchOffloader - v2.py -> prefetch.py - v2_ops.py -> prefetch_ops.py - test_v2_offload.py -> test_prefetch_offload.py Signed-off-by: Ming Yang --- ...v2_offload.py => test_prefetch_offload.py} | 16 +++++++------- vllm/config/offload.py | 2 +- vllm/model_executor/offloader/__init__.py | 4 ++-- vllm/model_executor/offloader/base.py | 6 +++--- .../offloader/{v2.py => prefetch.py} | 21 ++++++++++--------- .../offloader/{v2_ops.py => prefetch_ops.py} | 8 +++---- 6 files changed, 29 insertions(+), 28 deletions(-) rename tests/basic_correctness/{test_v2_offload.py => test_prefetch_offload.py} (50%) rename vllm/model_executor/offloader/{v2.py => prefetch.py} (97%) rename vllm/model_executor/offloader/{v2_ops.py => prefetch_ops.py} (92%) diff --git a/tests/basic_correctness/test_v2_offload.py b/tests/basic_correctness/test_prefetch_offload.py similarity index 50% rename from tests/basic_correctness/test_v2_offload.py rename to tests/basic_correctness/test_prefetch_offload.py index 071b702c3569..75d24bfdc5e0 100644 --- a/tests/basic_correctness/test_v2_offload.py +++ b/tests/basic_correctness/test_prefetch_offload.py @@ -1,31 +1,31 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Test V2 offloading correctness with DeepSeek V2 model.""" +"""Test prefetch offloading correctness with DeepSeek V2 model.""" from ..utils import compare_two_settings -def test_v2_offload_deepseek(): - """Test V2 CPU offloading with DeepSeek-V2-Lite. +def test_prefetch_offload_deepseek(): + """Test prefetch CPU offloading with DeepSeek-V2-Lite. Compares outputs between: 1. Baseline (no offloading) - 2. V2 offloading (group_size=8, num_in_group=2, prefetch_step=1) + 2. Prefetch offloading (group_size=8, num_in_group=2, prefetch_step=1) - This tests the advanced offloading with prefetching on a MoE model. + This tests prefetching-based offloading on a MoE model. """ compare_two_settings( "deepseek-ai/DeepSeek-V2-Lite", [ - # V2 offloading configuration + # Prefetch offloading configuration "--offload-group-size", "8", "--offload-num-in-group", "2", "--offload-prefetch-step", "1", - # torch.compile is automatically disabled when V2 offloading is - # enabled (via enable_if in @support_torch_compile decorator) + # torch.compile is automatically disabled when prefetch offloading + # is enabled (via enable_if in @support_torch_compile decorator) ], [], # Baseline: no offloading ) diff --git a/vllm/config/offload.py b/vllm/config/offload.py index d92a3656eaed..ae57781035b9 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -132,7 +132,7 @@ def compute_hash(self) -> str: """ Provide a hash that uniquely identifies all the offload configs. - All fields are included because OffloaderV2 patches module + All fields are included because PrefetchOffloader patches module forwards and inserts custom ops (wait_prefetch, start_prefetch) into the computation graph. Changing any offload setting can alter which layers are hooked and how prefetch indices are diff --git a/vllm/model_executor/offloader/__init__.py b/vllm/model_executor/offloader/__init__.py index a8031cb55ff4..a6522ff7c0a3 100644 --- a/vllm/model_executor/offloader/__init__.py +++ b/vllm/model_executor/offloader/__init__.py @@ -9,14 +9,14 @@ get_offloader, set_offloader, ) +from vllm.model_executor.offloader.prefetch import PrefetchOffloader from vllm.model_executor.offloader.uva import UVAOffloader -from vllm.model_executor.offloader.v2 import OffloaderV2 __all__ = [ "BaseOffloader", "NoopOffloader", "UVAOffloader", - "OffloaderV2", + "PrefetchOffloader", "create_offloader", "get_offloader", "set_offloader", diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index a930a71df7ee..ecbb82119b24 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -26,7 +26,7 @@ class relation: BaseOffloader (ABC) * implemented by: UVAOffloader - * implemented by: OffloaderV2 + * implemented by: PrefetchOffloader * uses: _ModuleOffloader * uses: _BaseParamOffloader (ABC) * implemented by: _CpuParamOffloader @@ -124,8 +124,8 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: selects prefetch if ``offload_group_size > 0``, UVA if ``cpu_offload_gb > 0``, otherwise noop. """ + from vllm.model_executor.offloader.prefetch import PrefetchOffloader from vllm.model_executor.offloader.uva import UVAOffloader - from vllm.model_executor.offloader.v2 import OffloaderV2 backend = offload_config.offload_backend uva = offload_config.uva @@ -140,7 +140,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: return NoopOffloader() if backend == "prefetch": - return OffloaderV2( + return PrefetchOffloader( group_size=prefetch.offload_group_size, num_in_group=prefetch.offload_num_in_group, prefetch_step=prefetch.offload_prefetch_step, diff --git a/vllm/model_executor/offloader/v2.py b/vllm/model_executor/offloader/prefetch.py similarity index 97% rename from vllm/model_executor/offloader/v2.py rename to vllm/model_executor/offloader/prefetch.py index 17ee301d409b..59e189f47fb7 100644 --- a/vllm/model_executor/offloader/v2.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -2,11 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from # https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py -"""OffloaderV2: CPU offloading with async prefetching. +"""Prefetch-based CPU offloading with async prefetching. -This version uses static buffers and event-based stream forking for -torch.compile + CUDA graph compatibility. Events allow the copy stream -to join CUDA graph captures, ensuring H2D copies are properly captured. +Uses static buffers and event-based stream forking for torch.compile + +CUDA graph compatibility. Events allow the copy stream to join CUDA +graph captures, ensuring H2D copies are properly captured. """ from abc import ABC, abstractmethod @@ -17,8 +17,8 @@ import torch import torch.nn as nn -# Import v2_ops to register custom ops at module load time -import vllm.model_executor.offloader.v2_ops # noqa: F401 +# Import prefetch_ops to register custom ops at module load time +import vllm.model_executor.offloader.prefetch_ops # noqa: F401 from vllm.logger import init_logger from vllm.model_executor.offloader.base import BaseOffloader from vllm.utils.platform_utils import is_pin_memory_available @@ -127,9 +127,10 @@ def get_buffer( return self._buffers[key][slot_idx % self.slot_capacity] -class OffloaderV2(BaseOffloader): - """Advanced offloader with group-based selection and async prefetching. +class PrefetchOffloader(BaseOffloader): + """Prefetching-based offloader with group-based layer selection. + Groups layers and uses async H2D prefetch to hide transfer latency. Uses static buffers and stream synchronization for torch.compile and CUDA graph compatibility. @@ -166,7 +167,7 @@ def wrap_modules( submodule_accessor: _SubmoduleAccessor | None = None, whitelist_param_names_creator: _WhitelistParamNamesCreator | None = None, ) -> list[nn.Module]: - """Wrap modules with V2 offloading and prefetching logic.""" + """Wrap modules with prefetch offloading logic.""" assert len(self.module_offloaders) == 0, ( "wrap_modules should only be called once" ) @@ -351,7 +352,7 @@ def post_init(self): self.total_offloaded_bytes += offloader.offloaded_bytes logger.info_once( - f"[OffloaderV2] Initialized {len(self.module_offloaders)} modules. " + f"[PrefetchOffloader] Initialized {len(self.module_offloaders)} modules. " f"Total GPU memory saved: {self.total_offloaded_bytes / 1e9:.4f} GB, " f"Static buffer pool: {self.buffer_pool.total_bytes / 1e9:.4f} GB " f"(group_size={self.group_size}, num_in_group={self.num_in_group}, " diff --git a/vllm/model_executor/offloader/v2_ops.py b/vllm/model_executor/offloader/prefetch_ops.py similarity index 92% rename from vllm/model_executor/offloader/v2_ops.py rename to vllm/model_executor/offloader/prefetch_ops.py index 2f82fbfbadc2..f3b4465f120e 100644 --- a/vllm/model_executor/offloader/v2_ops.py +++ b/vllm/model_executor/offloader/prefetch_ops.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Custom ops for V2 offloader torch.compile + CUDA graph compatibility. +"""Custom ops for prefetch offloader torch.compile + CUDA graph compatibility. These ops use mutates_args to create data dependencies that prevent the compiler from reordering prefetch/sync operations. @@ -77,8 +77,8 @@ def _start_prefetch_fake( return output_tensor -def register_v2_offloader_ops() -> None: - """Register custom ops for V2 offloader. +def register_prefetch_offloader_ops() -> None: + """Register custom ops for prefetch offloader. Must be called before the ops are used. This is typically done at module import time. @@ -99,4 +99,4 @@ def register_v2_offloader_ops() -> None: # Register ops at module import time -register_v2_offloader_ops() +register_prefetch_offloader_ops() From 08bab8a0541d3c1bf6d5cc9eaefd58ff4337e028 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sun, 22 Feb 2026 14:51:25 -0800 Subject: [PATCH 25/36] Move PrefetchOffloader parameterization from model code to CLI Add --offload-params CLI flag to PrefetchOffloadConfig so parameter selection is driven entirely from config, matching how UVA's --cpu-offload-params works. This removes model-specific Python callbacks (submodule_accessor, whitelist_param_names_creator) from the offloader interface, making prefetch offloading work out-of-the-box for all models without per-model code changes. - Add offload_params field to PrefetchOffloadConfig with segment matching - Add --offload-params CLI flag wired through EngineArgs - Replace callback-based wrap_modules with config-driven matching - Simplify BaseOffloader/NoopOffloader/UVAOffloader signatures - Remove offloader_kwargs from make_layers() and deepseek_v2.py - Delete find_fused_moe_submodule (no longer needed) - Add --offload-params to prefetch offload test Signed-off-by: Ming Yang --- .../test_prefetch_offload.py | 4 ++ vllm/config/offload.py | 8 ++++ vllm/engine/arg_utils.py | 5 +++ .../layers/fused_moe/__init__.py | 2 - vllm/model_executor/layers/fused_moe/layer.py | 25 ------------ vllm/model_executor/models/deepseek_v2.py | 20 +--------- vllm/model_executor/models/utils.py | 9 +---- vllm/model_executor/offloader/base.py | 14 +------ vllm/model_executor/offloader/prefetch.py | 40 ++++++++++--------- vllm/model_executor/offloader/uva.py | 10 +---- 10 files changed, 44 insertions(+), 93 deletions(-) diff --git a/tests/basic_correctness/test_prefetch_offload.py b/tests/basic_correctness/test_prefetch_offload.py index 75d24bfdc5e0..52510416d274 100644 --- a/tests/basic_correctness/test_prefetch_offload.py +++ b/tests/basic_correctness/test_prefetch_offload.py @@ -24,6 +24,10 @@ def test_prefetch_offload_deepseek(): "2", "--offload-prefetch-step", "1", + # Selective offloading: only MoE expert weights + "--offload-params", + "w13_weight", + "w2_weight", # torch.compile is automatically disabled when prefetch offloading # is enabled (via enable_if in @support_torch_compile decorator) ], diff --git a/vllm/config/offload.py b/vllm/config/offload.py index ae57781035b9..ad65e8acf35a 100644 --- a/vllm/config/offload.py +++ b/vllm/config/offload.py @@ -67,6 +67,14 @@ class PrefetchOffloadConfig: """Number of layers to prefetch ahead. Higher values hide more latency but use more GPU memory. Default is 1.""" + offload_params: set[str] = Field(default_factory=set) + """The set of parameter name segments to target for prefetch offloading. + Unmatched parameters are not offloaded. If this set is empty, ALL + parameters of each offloaded layer are offloaded. + Uses segment matching: "w13_weight" matches "mlp.experts.w13_weight" + but not "mlp.experts.w13_weight_scale". + """ + @config class OffloadConfig: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 6c9f35eac026..1aa52e2c787f 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -447,6 +447,7 @@ class EngineArgs: offload_group_size: int = PrefetchOffloadConfig.offload_group_size offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step + offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params") gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None @@ -1006,6 +1007,9 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--offload-prefetch-step", **prefetch_kwargs["offload_prefetch_step"], ) + offload_group.add_argument( + "--offload-params", **prefetch_kwargs["offload_params"] + ) # Multimodal related configs multimodal_kwargs = get_kwargs(MultiModalConfig) @@ -1864,6 +1868,7 @@ def create_engine_config( offload_group_size=self.offload_group_size, offload_num_in_group=self.offload_num_in_group, offload_prefetch_step=self.offload_prefetch_step, + offload_params=self.offload_params, ), ) diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 1c61ca212244..c6cb31b629a0 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -19,7 +19,6 @@ from vllm.model_executor.layers.fused_moe.layer import ( FusedMoE, FusedMoeWeightScaleSupported, - find_fused_moe_submodule, ) from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEActivationFormat, @@ -69,7 +68,6 @@ def get_config() -> dict[str, Any] | None: "SharedFusedMoE", "ZeroExpertFusedMoE", "activation_without_mul", - "find_fused_moe_submodule", "apply_moe_activation", "override_config", "get_config", diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 5639d28876b8..6cb3dae26736 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -1547,28 +1547,3 @@ def extra_repr(self) -> str: # Mark the FusedMoE weight_loader as supporting MoE-specific parameters # to avoid expensive runtime reflection in model loading code FusedMoE.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] - - -def find_fused_moe_submodule(module: torch.nn.Module) -> torch.nn.Module: - """Find a FusedMoE submodule for offloading, or return the module itself. - - Searches module attributes for instances of FusedMoE (or subclasses like - SharedFusedMoE). - - Args: - module: The module to search within (typically layer.mlp). - - Returns: - The first FusedMoE instance found, or the original module if none found. - """ - for attr_name in dir(module): - if attr_name.startswith("_"): - continue - try: - attr = getattr(module, attr_name, None) - except Exception: - continue - if isinstance(attr, FusedMoE): - return attr - - return module diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 86fbe3f998a6..3b3b7a1a37e5 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -47,10 +47,7 @@ from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.model_executor.layers.fused_moe import ( - SharedFusedMoE, - find_fused_moe_submodule, -) +from vllm.model_executor.layers.fused_moe import SharedFusedMoE from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -1130,21 +1127,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config, prefix, topk_indices_buffer=topk_indices_buffer ), prefix=f"{prefix}.layers", - offloader_kwargs=dict( - # Extract the MLP submodule - for MoE layers, go deeper to the experts - submodule_accessor=lambda layer: find_fused_moe_submodule(layer.mlp), - # Specify which parameters to offload - whitelist_param_names_creator=lambda module: ( - [ - # Core MoE expert weights - "w13_weight", - "w2_weight", - ] - # Only offload from MoE experts (SharedFusedMoE/FusedMoE) - if hasattr(module, "w13_weight") - else [] - ), - ), ) if get_pp_group().is_last_rank: diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 05f061850218..c55693bcff93 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -611,7 +611,6 @@ def make_layers( num_hidden_layers: int, layer_fn: LayerFn, prefix: str, - offloader_kwargs: dict | None = None, ) -> tuple[int, int, torch.nn.ModuleList]: """Make a list of layers with the given layer function, taking pipeline parallelism into account. @@ -620,8 +619,6 @@ def make_layers( num_hidden_layers: Total number of hidden layers in the model. layer_fn: Function to create a layer given its index. prefix: Prefix for layer names. - offloader_kwargs: Optional kwargs for offloader (submodule_accessor, - whitelist_param_names_creator). Returns: Tuple of (start_layer, end_layer, modules). @@ -637,11 +634,7 @@ def make_layers( modules = torch.nn.ModuleList( [PPMissingLayer() for _ in range(start_layer)] + get_offloader().wrap_modules( - ( - layer_fn(prefix=f"{prefix}.{idx}") - for idx in range(start_layer, end_layer) - ), - **(offloader_kwargs or {}), + layer_fn(prefix=f"{prefix}.{idx}") for idx in range(start_layer, end_layer) ) + [PPMissingLayer() for _ in range(end_layer, num_hidden_layers)] ) diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index ecbb82119b24..51ccd9b0edda 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -5,7 +5,7 @@ """Base classes for model parameter offloading.""" from abc import ABC, abstractmethod -from collections.abc import Callable, Generator +from collections.abc import Generator from typing import TYPE_CHECKING import torch.nn as nn @@ -17,9 +17,6 @@ logger = init_logger(__name__) -_SubmoduleAccessor = Callable[[nn.Module], nn.Module] -_WhitelistParamNamesCreator = Callable[[nn.Module], list[str]] - """ class relation: @@ -44,17 +41,11 @@ class BaseOffloader(ABC): def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], - submodule_accessor: _SubmoduleAccessor | None = None, - whitelist_param_names_creator: _WhitelistParamNamesCreator | None = None, ) -> list[nn.Module]: """Wrap modules with offloading logic. Args: modules_generator: Generator yielding modules to potentially offload. - submodule_accessor: Optional function to extract a submodule from - each module (e.g., lambda layer: layer.mlp.experts). - whitelist_param_names_creator: Optional function to get parameter - names to offload from a submodule (e.g., ["w13_weight", "w2_weight"]). Returns: List of modules, potentially with offloading hooks installed. @@ -94,8 +85,6 @@ class NoopOffloader(BaseOffloader): def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], - submodule_accessor: _SubmoduleAccessor | None = None, - whitelist_param_names_creator: _WhitelistParamNamesCreator | None = None, ) -> list[nn.Module]: """Return modules unchanged.""" return list(modules_generator) @@ -144,6 +133,7 @@ def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: group_size=prefetch.offload_group_size, num_in_group=prefetch.offload_num_in_group, prefetch_step=prefetch.offload_prefetch_step, + offload_params=prefetch.offload_params, mode="cpu", ) elif backend == "uva": diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index 59e189f47fb7..4cdb53fc1710 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -10,7 +10,7 @@ """ from abc import ABC, abstractmethod -from collections.abc import Callable, Generator +from collections.abc import Generator from dataclasses import dataclass from typing import Any @@ -25,9 +25,6 @@ logger = init_logger(__name__) -_SubmoduleAccessor = Callable[[nn.Module], nn.Module] -_WhitelistParamNamesCreator = Callable[[nn.Module], list[str]] - @dataclass class ParamInfo: @@ -146,11 +143,13 @@ def __init__( group_size: int, num_in_group: int, prefetch_step: int, + offload_params: set[str] | None = None, mode: str = "cpu", ): self.group_size = group_size self.num_in_group = num_in_group self.prefetch_step = prefetch_step + self.offload_params = offload_params or set() self.mode = mode # Copy stream for async H2D transfers @@ -164,8 +163,6 @@ def __init__( def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], - submodule_accessor: _SubmoduleAccessor | None = None, - whitelist_param_names_creator: _WhitelistParamNamesCreator | None = None, ) -> list[nn.Module]: """Wrap modules with prefetch offloading logic.""" assert len(self.module_offloaders) == 0, ( @@ -173,7 +170,7 @@ def wrap_modules( ) all_modules = [] - offload_submodules = [] + offload_modules = [] for module_index, module in enumerate(modules_generator): all_modules.append(module) @@ -181,26 +178,31 @@ def wrap_modules( # Select layers to offload based on group pattern # Offload last num_in_group layers of each group_size if module_index % self.group_size >= self.group_size - self.num_in_group: - submodule = submodule_accessor(module) if submodule_accessor else module - whitelist_param_names = ( - whitelist_param_names_creator(submodule) - if whitelist_param_names_creator - else [name for name, _ in submodule.named_parameters()] - ) - - offload_submodules.append(submodule) + if self.offload_params: + whitelist = [ + name + for name, _ in module.named_parameters() + if any(f".{p}." in f".{name}." for p in self.offload_params) + ] + else: + whitelist = [name for name, _ in module.named_parameters()] + + if not whitelist: + continue # skip layers with no matching params + + offload_modules.append(module) self.module_offloaders.append( _ModuleOffloader( mode=self.mode, - module=submodule, + module=module, copy_stream=self.copy_stream, - whitelist_param_names=whitelist_param_names, + whitelist_param_names=whitelist, layer_idx=len(self.module_offloaders), ) ) - for index, submodule in enumerate(offload_submodules): - self._hook_module_forward(index, submodule) + for index, module in enumerate(offload_modules): + self._hook_module_forward(index, module) return all_modules diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py index 2c6535b0fd49..c524e43cddae 100644 --- a/vllm/model_executor/offloader/uva.py +++ b/vllm/model_executor/offloader/uva.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """UVA-based CPU offloading using Unified Virtual Addressing.""" -from collections.abc import Callable, Generator +from collections.abc import Generator import torch import torch.nn as nn @@ -54,14 +54,8 @@ def __init__( def wrap_modules( self, modules_generator: Generator[nn.Module, None, None], - submodule_accessor: Callable[[nn.Module], nn.Module] | None = None, - whitelist_param_names_creator: Callable[[nn.Module], list[str]] | None = None, ) -> list[nn.Module]: - """Wrap modules with UVA offloading. - - Note: UVA offloading operates at module level, so submodule_accessor - and whitelist_param_names_creator are ignored. - """ + """Wrap modules with UVA offloading.""" modules = [self._maybe_offload_to_cpu(module) for module in modules_generator] if self.cpu_offload_bytes > 0: logger.info( From 0d18ec24cf3fb32e852bd9057f24e0c86f4e4abc Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Sun, 22 Feb 2026 15:09:24 -0800 Subject: [PATCH 26/36] Make prefetch custom ops return None instead of tensor mutates_args already creates data dependencies for torch.compile, so returning the tensor is unnecessary. Use in-place mutation semantics and simplify the call sites in _hook_module_forward. Signed-off-by: Ming Yang --- vllm/model_executor/offloader/prefetch.py | 16 ++++------ vllm/model_executor/offloader/prefetch_ops.py | 30 +++++++------------ 2 files changed, 16 insertions(+), 30 deletions(-) diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index 4cdb53fc1710..1bacb5dce7e3 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -215,28 +215,22 @@ def forward(*args, **kwargs): module.forward = original_forward # Wait for this layer's prefetch to complete + # mutates_args on input_tensor creates data dependency for torch.compile input_tensor = args[0] if args else kwargs.get("hidden_states") - input_tensor = torch.ops.vllm.wait_prefetch(input_tensor, index) - - # Replace the first arg with the returned tensor to maintain dependency - if args: - args = (input_tensor,) + args[1:] - else: - kwargs["hidden_states"] = input_tensor + torch.ops.vllm.wait_prefetch(input_tensor, index) # No parameter swapping needed - parameters already point to # GPU static buffers (set in assign_static_buffer) output = original_forward(*args, **kwargs) # Start prefetch for next layer (circular) - # Custom op returns output_tensor to create data dependency + # mutates_args on output_tensor creates ordering dependency next_index = (index + self.prefetch_step) % len(self.module_offloaders) # Handle tuple output (e.g., (hidden_states, residual)) if isinstance(output, tuple): - output_tensor = torch.ops.vllm.start_prefetch(output[0], next_index) - output = (output_tensor,) + output[1:] + torch.ops.vllm.start_prefetch(output[0], next_index) else: - output = torch.ops.vllm.start_prefetch(output, next_index) + torch.ops.vllm.start_prefetch(output, next_index) # No explicit offload needed - static buffers are reused implicitly diff --git a/vllm/model_executor/offloader/prefetch_ops.py b/vllm/model_executor/offloader/prefetch_ops.py index f3b4465f120e..d1f59b67b4ad 100644 --- a/vllm/model_executor/offloader/prefetch_ops.py +++ b/vllm/model_executor/offloader/prefetch_ops.py @@ -19,30 +19,26 @@ def _wait_prefetch_impl( input_tensor: torch.Tensor, layer_idx: int, -) -> torch.Tensor: +) -> None: """Wait for prefetch of layer_idx to complete. Synchronizes the compute stream with the copy stream to ensure the prefetched weights are ready for use. Args: - input_tensor: Input to the layer (e.g., hidden_states) - returned - to create data dependency chain. + input_tensor: Input to the layer (e.g., hidden_states) - declared + as mutated to create data dependency for torch.compile. layer_idx: Index of the layer to wait for. - - Returns: - input_tensor unchanged, but creates data dependency for torch.compile. """ get_offloader()._wait_for_layer(layer_idx) - return input_tensor def _wait_prefetch_fake( input_tensor: torch.Tensor, layer_idx: int, -) -> torch.Tensor: +) -> None: """Fake implementation for torch.compile tracing.""" - return input_tensor + return # --- start_prefetch op --- @@ -51,30 +47,26 @@ def _wait_prefetch_fake( def _start_prefetch_impl( output_tensor: torch.Tensor, layer_idx: int, -) -> torch.Tensor: +) -> None: """Start async prefetch of layer_idx weights. Initiates H2D copy on the copy stream for the specified layer. Args: - output_tensor: Output from forward - returned to create ordering - dependency. This prevents torch.compile from reordering - this op before the computation that produces output_tensor. + output_tensor: Output from forward - declared as mutated to + prevent torch.compile from reordering this op before the + computation that produces output_tensor. layer_idx: Index of the layer to prefetch. - - Returns: - output_tensor unchanged, creating data dependency for torch.compile. """ get_offloader()._start_prefetch(layer_idx) - return output_tensor def _start_prefetch_fake( output_tensor: torch.Tensor, layer_idx: int, -) -> torch.Tensor: +) -> None: """Fake implementation for torch.compile tracing.""" - return output_tensor + return def register_prefetch_offloader_ops() -> None: From e5d375e6cb6a2065f231be87979b6b57106d973a Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Mon, 23 Feb 2026 00:26:17 -0800 Subject: [PATCH 27/36] Switch prefetch offload test from DeepSeek-V2-Lite to Llama-3.2-1B-Instruct Signed-off-by: Ming Yang --- tests/basic_correctness/test_prefetch_offload.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/basic_correctness/test_prefetch_offload.py b/tests/basic_correctness/test_prefetch_offload.py index 52510416d274..6491c4a27000 100644 --- a/tests/basic_correctness/test_prefetch_offload.py +++ b/tests/basic_correctness/test_prefetch_offload.py @@ -1,21 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Test prefetch offloading correctness with DeepSeek V2 model.""" +"""Test prefetch offloading correctness with Llama model.""" from ..utils import compare_two_settings -def test_prefetch_offload_deepseek(): - """Test prefetch CPU offloading with DeepSeek-V2-Lite. +def test_prefetch_offload_llama(): + """Test prefetch CPU offloading with Llama-3.2-1B-Instruct. Compares outputs between: 1. Baseline (no offloading) 2. Prefetch offloading (group_size=8, num_in_group=2, prefetch_step=1) - This tests prefetching-based offloading on a MoE model. + This tests prefetching-based offloading on a dense model. """ compare_two_settings( - "deepseek-ai/DeepSeek-V2-Lite", + "meta-llama/Llama-3.2-1B-Instruct", [ # Prefetch offloading configuration "--offload-group-size", @@ -24,10 +24,10 @@ def test_prefetch_offload_deepseek(): "2", "--offload-prefetch-step", "1", - # Selective offloading: only MoE expert weights + # Selective offloading: only MLP weights "--offload-params", - "w13_weight", - "w2_weight", + "gate_up_proj", + "down_proj", # torch.compile is automatically disabled when prefetch offloading # is enabled (via enable_if in @support_torch_compile decorator) ], From 3c3ba9b4927ebcbc216f6a3fe2ce2376bbfcef6b Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Mon, 23 Feb 2026 10:30:41 -0800 Subject: [PATCH 28/36] Remove incorrect comment Signed-off-by: Ming Yang --- tests/basic_correctness/test_prefetch_offload.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/basic_correctness/test_prefetch_offload.py b/tests/basic_correctness/test_prefetch_offload.py index 6491c4a27000..498887024ee6 100644 --- a/tests/basic_correctness/test_prefetch_offload.py +++ b/tests/basic_correctness/test_prefetch_offload.py @@ -28,8 +28,6 @@ def test_prefetch_offload_llama(): "--offload-params", "gate_up_proj", "down_proj", - # torch.compile is automatically disabled when prefetch offloading - # is enabled (via enable_if in @support_torch_compile decorator) ], [], # Baseline: no offloading ) From 96301e3cf15944640eaa662546eda40ef1fc905b Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 24 Feb 2026 13:08:53 -0800 Subject: [PATCH 29/36] Use per-layer event sync in eager mode to enable H2D overlap In eager mode, _wait_for_layer was using wait_stream(copy_stream) which waits for ALL pending H2D copies, serializing every layer's prefetch. Switch to per-layer wait_event so each layer only waits for its own copy, allowing prefetch_step>1 to overlap H2D with compute from intervening layers. Events recorded during cudagraph capture become invalid after capture ends, so we track _event_valid_for_eager and fall back to wait_stream when the event was last recorded in a capture context. Signed-off-by: Ming Yang --- vllm/model_executor/offloader/prefetch.py | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index 1bacb5dce7e3..319de9ad9631 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -263,9 +263,14 @@ def _wait_for_layer(self, layer_idx: int): # Mark that this prefetch has been waited on (joined). offloader._prefetch_in_capture = False else: - # Outside capture: use wait_stream for robustness. - # Events used in previous captures can be in invalid state. - torch.cuda.current_stream().wait_stream(self.copy_stream) + if offloader._event_valid_for_eager: + # Use per-layer event to only wait for THIS layer's copy, + # allowing other layers' prefetches to run concurrently. + torch.cuda.current_stream().wait_event(offloader._copy_done_event) + else: + # Event not usable (unrecorded or recorded during capture). + # Fall back to wait_stream to drain all copy_stream work. + torch.cuda.current_stream().wait_stream(self.copy_stream) def sync_prev_onload(self): """Sync previous onload operations. @@ -382,10 +387,15 @@ def __init__( self.offloaded_bytes = 0 # Event to signal when H2D copy to static buffer is complete. - # Used for CUDA graph compatible synchronization during capture. - # Outside capture, we use wait_stream instead (more robust). + # Used for per-layer synchronization (both eager and capture modes). self._copy_done_event = torch.cuda.Event() + # Track whether _copy_done_event is valid for eager-mode wait_event. + # False when: (1) never recorded, or (2) last recorded during a + # cudagraph capture (events become invalid after capture ends). + # In these cases we fall back to wait_stream. + self._event_valid_for_eager = False + # Track if last prefetch was started during CUDA graph capture. # Used to skip wait_event during capture for pre-capture prefetches. self._prefetch_in_capture = False @@ -506,6 +516,9 @@ def start_onload_to_static(self): # Record completion event for _wait_for_layer to use self._copy_done_event.record(self.copy_stream) + # Event is only valid for eager wait_event if recorded outside capture. + # Events recorded during capture become invalid after capture ends. + self._event_valid_for_eager = not torch.cuda.is_current_stream_capturing() class _BaseParamOffloader(ABC): From 6c6600d96ed83fbf61bfbd620c6d0ec6bc590c20 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 24 Feb 2026 13:51:00 -0800 Subject: [PATCH 30/36] Add offload_params to LLM interface Signed-off-by: Ming Yang --- vllm/entrypoints/llm.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index f523c7bbe42c..c50a30f884fc 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -169,14 +169,19 @@ class LLM: the model weights. This virtually increases the GPU memory space you can use to hold the model weights, at the cost of CPU-GPU data transfer for every forward pass. - offload_group_size: Advanced CPU offloading: Group every N layers + offload_group_size: Prefetch offloading: Group every N layers together. Offload last `offload_num_in_group` layers of each group. Default is 0 (disabled). - offload_num_in_group: Advanced CPU offloading: Number of layers to + offload_num_in_group: Prefetch offloading: Number of layers to offload per group. Default is 1. - offload_prefetch_step: Advanced CPU offloading: Number of layers to + offload_prefetch_step: Prefetch offloading: Number of layers to prefetch ahead. Higher values hide more latency but use more GPU memory. Default is 1. + offload_params: Prefetch offloading: Set of parameter name segments + to selectively offload. Only parameters whose names contain one of + these segments will be offloaded (e.g., {"gate_up_proj", "down_proj"} + for MLP weights, or {"w13_weight", "w2_weight"} for MoE expert + weights). If None or empty, all parameters are offloaded. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. @@ -234,6 +239,7 @@ def __init__( offload_group_size: int = 0, offload_num_in_group: int = 1, offload_prefetch_step: int = 1, + offload_params: set[str] | None = None, enforce_eager: bool = False, enable_return_routed_experts: bool = False, disable_custom_all_reduce: bool = False, @@ -346,6 +352,7 @@ def _make_config(value: Any, cls: type[_R]) -> _R: offload_group_size=offload_group_size, offload_num_in_group=offload_num_in_group, offload_prefetch_step=offload_prefetch_step, + offload_params=offload_params or set(), enforce_eager=enforce_eager, enable_return_routed_experts=enable_return_routed_experts, disable_custom_all_reduce=disable_custom_all_reduce, From ff52cd45497269c4153a59f638b781b0b241c075 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 24 Feb 2026 13:52:59 -0800 Subject: [PATCH 31/36] Use torch.finfo instead of creating tensor to get element size Signed-off-by: Ming Yang --- vllm/model_executor/offloader/prefetch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py index 319de9ad9631..b43cb8b7d87f 100644 --- a/vllm/model_executor/offloader/prefetch.py +++ b/vllm/model_executor/offloader/prefetch.py @@ -54,7 +54,7 @@ def num_bytes(self) -> int: numel = 1 for dim in self.shape: numel *= dim - return numel * torch.tensor([], dtype=self.dtype).element_size() + return numel * torch.finfo(self.dtype).bits // 8 class StaticBufferPool: From 2a773855a8268218b975eec26d73908ce29e4324 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 24 Feb 2026 13:54:18 -0800 Subject: [PATCH 32/36] Initialize global offloader as None to fail loudly before set_offloader Signed-off-by: Ming Yang --- vllm/model_executor/offloader/base.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index 51ccd9b0edda..b5e839828eec 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -90,13 +90,17 @@ def wrap_modules( return list(modules_generator) -# Global singleton offloader instance -_instance: BaseOffloader | None = NoopOffloader() +# Global singleton offloader instance. +# Starts as None so that accidental use before set_offloader() fails loudly. +_instance: BaseOffloader | None = None def get_offloader() -> BaseOffloader: """Get the global offloader instance.""" - assert _instance is not None, "Offloader instance is None" + assert _instance is not None, ( + "Offloader not initialized. set_offloader() must be called " + "before get_offloader()." + ) return _instance From 20696cd61e18c579e0ad6be9c4808c4fa2a3903e Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 24 Feb 2026 14:01:09 -0800 Subject: [PATCH 33/36] Add nightly e2e test for prefetch offloading with DeepSeek-V2-Lite Runs DeepSeek-V2-Lite with prefetch offloading of MoE expert weights (w13_weight, w2_weight) and validates GSM8K accuracy matches baseline. Signed-off-by: Ming Yang --- .../deepseek_v2_lite_prefetch_offload.sh | 57 +++++++++++++++++++ .buildkite/test_areas/e2e_integration.yaml | 9 +++ 2 files changed, 66 insertions(+) create mode 100755 .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh diff --git a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh new file mode 100755 index 000000000000..dddf23f1f2fd --- /dev/null +++ b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +set -euxo pipefail + +# Nightly e2e test for prefetch offloading with a MoE model. +# Runs DeepSeek-V2-Lite with prefetch offloading of MoE expert weights +# and validates GSM8K accuracy matches baseline (no offloading). +# +# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT] +THRESHOLD=${1:-0.25} +NUM_Q=${2:-1319} +PORT=${3:-8030} +OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled} +mkdir -p "${OUT_DIR}" + +wait_for_server() { + local port=$1 + timeout 600 bash -c ' + until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do + sleep 1 + done' +} + +MODEL="deepseek-ai/DeepSeek-V2-Lite" + +cleanup() { + if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then + kill "${SERVER_PID}" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "${SERVER_PID}" 2>/dev/null || break + sleep 0.5 + done + kill -9 "${SERVER_PID}" 2>/dev/null || true + fi +} +trap cleanup EXIT + +vllm serve "$MODEL" \ + --max-model-len 2048 \ + --offload-group-size 8 \ + --offload-num-in-group 2 \ + --offload-prefetch-step 1 \ + --offload-params w13_weight w2_weight \ + --port "$PORT" & +SERVER_PID=$! +wait_for_server "$PORT" + +TAG=$(echo "$MODEL" | tr '/: \\n' '_____') +OUT="${OUT_DIR}/${TAG}_prefetch_offload.json" +python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}" +python3 - <= ${THRESHOLD}, f"${MODEL} prefetch_offload accuracy {acc}" +PY + +cleanup +SERVER_PID= diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index d95b73073d6a..285481ee18ec 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -28,3 +28,12 @@ steps: working_dir: "/vllm-workspace" commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 + +- label: DeepSeek V2-Lite Prefetch Offload Accuracy + timeout_in_minutes: 60 + device: h100 + optional: true + num_devices: 1 + working_dir: "/vllm-workspace" + commands: + - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030 From f537f21e7fc3408d15945e9113f64f0dcebc3e4b Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 24 Feb 2026 15:19:19 -0800 Subject: [PATCH 34/36] Fix sync_prev_onload comments in run_fullgraph Signed-off-by: Ming Yang --- vllm/v1/worker/gpu/cudagraph_utils.py | 8 ++++++-- vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 5f862cfdf848..d70a4c7ab18d 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -340,8 +340,12 @@ def run_fullgraph( self, num_tokens: int ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: assert num_tokens in self.graphs, f"No cudagraph for {num_tokens} tokens" - # Sync offloader before replay - ensures any external dependencies - # from pre-capture prefetches are satisfied. + # Sync offloader before replay - needed when transitioning from + # eager/piecewise to full cudagraph (e.g., prefill → decode). + # The previous eager iteration's start_prefetch may have queued + # H2D copies on copy_stream that the graph's captured events + # cannot see. Without this, replay could overwrite static buffers + # while those copies are still in flight. get_offloader().sync_prev_onload() self.graphs[num_tokens].replay() assert self.hidden_states is not None diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py index f01b613f42f3..eda8c37d53f4 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py @@ -181,7 +181,11 @@ def capture( def run_fullgraph(self, num_tokens: int) -> None: assert num_tokens in self.graphs - # Sync offloader before replay - ensures any external dependencies - # from pre-capture prefetches are satisfied. + # Sync offloader before replay - needed when transitioning from + # eager/piecewise to full cudagraph (e.g., prefill → decode). + # The previous eager iteration's start_prefetch may have queued + # H2D copies on copy_stream that the graph's captured events + # cannot see. Without this, replay could overwrite static buffers + # while those copies are still in flight. get_offloader().sync_prev_onload() self.graphs[num_tokens].replay() From 18b576eef254da1810400c8aec4d4f2b00487c07 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Tue, 24 Feb 2026 18:36:18 -0700 Subject: [PATCH 35/36] Update .buildkite/test_areas/e2e_integration.yaml Signed-off-by: Michael Goin --- .buildkite/test_areas/e2e_integration.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index 285481ee18ec..5b7f96bc7a26 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -29,7 +29,7 @@ steps: commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 -- label: DeepSeek V2-Lite Prefetch Offload Accuracy +- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100) timeout_in_minutes: 60 device: h100 optional: true From 3df4d98efe460af3fe845946fadc2459275e8bef Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Tue, 24 Feb 2026 21:37:50 -0800 Subject: [PATCH 36/36] Default global offloader to NoopOffloader and log on set Set _instance default to NoopOffloader() so get_offloader() always returns a valid instance. Log the offloader type in set_offloader() for visibility into which backend is active. Signed-off-by: Ming Yang --- vllm/model_executor/offloader/base.py | 10 +++------- vllm/v1/worker/gpu_model_runner.py | 1 + 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py index b5e839828eec..7c61b318b881 100644 --- a/vllm/model_executor/offloader/base.py +++ b/vllm/model_executor/offloader/base.py @@ -90,17 +90,12 @@ def wrap_modules( return list(modules_generator) -# Global singleton offloader instance. -# Starts as None so that accidental use before set_offloader() fails loudly. -_instance: BaseOffloader | None = None +# Global singleton offloader instance (defaults to no-op). +_instance: BaseOffloader = NoopOffloader() def get_offloader() -> BaseOffloader: """Get the global offloader instance.""" - assert _instance is not None, ( - "Offloader not initialized. set_offloader() must be called " - "before get_offloader()." - ) return _instance @@ -108,6 +103,7 @@ def set_offloader(instance: BaseOffloader) -> None: """Set the global offloader instance.""" global _instance _instance = instance + logger.info("Offloader set to %s", type(instance).__name__) def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 15479523e344..3ef5e5e6d528 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -707,6 +707,7 @@ def __init__( ) # Model weight offloader + # Make sure this is called before any get_offloader call set_offloader(create_offloader(self.offload_config)) # Ephemeral state transferred between execute_model() and sample_tokens().