[offloader] v2: Hide weight onloading latency via prefetching

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",
+        ],
+    )

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

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,

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,

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:

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,105 +506,45 @@ 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,
     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
     )
+
     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
 
 

vllm/model_executor/offloader/__init__.py

@@ -0,0 +1,23 @@
+# 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,
+    create_offloader,
+    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",
+    "create_offloader",
+    "get_offloader",
+    "set_offloader",
+]