A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 07:30 UTC
[Feature][Perf] Support Selective CPU Weight Offloading (#34535)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Feb 13, 2026 at 23:02 UTC
b37b679770aade27f33d20c93bf467c6a7fba65d
4 files changed
+44
-2
vllm/config/cache.py
+11
@@ -101,6 +101,17 @@ class CacheConfig:
101
Note that this requires fast CPU-GPU interconnect, as part of the model is
102
loaded from CPU memory to GPU memory on the fly in each model forward pass.
103
"""
104
+ cpu_offload_params: set[str] = Field(default_factory=set)
105
+ """ The set of parameter name segments to target for CPU offloading.
106
+ Unmatched parameters are not offloaded. If this set is empty, parameters
107
+ are offloaded non-selectively until the memory limit defined by
108
+ `cpu_offload_gb` is reached.
109
+ Examples:
110
+ - For parameter name "mlp.experts.w2_weight":
111
+ - "experts" or "experts.w2_weight" will match.
112
+ - "expert" or "w2" will NOT match (must be exact segments).
113
+ This allows distinguishing parameters like "w2_weight" and "w2_weight_scale".
114
+ """
115
calculate_kv_scales: bool = False
116
"""This enables dynamic calculation of `k_scale` and `v_scale` when
117
kv_cache_dtype is fp8. If `False`, the scales will be loaded from the model
vllm/engine/arg_utils.py
+5
@@ -434,6 +434,7 @@ class EngineArgs:
434
disable_cascade_attn: bool = ModelConfig.disable_cascade_attn
435
swap_space: float = CacheConfig.swap_space
436
cpu_offload_gb: float = CacheConfig.cpu_offload_gb
437
+ cpu_offload_params: set[str] = get_field(CacheConfig, "cpu_offload_params")
438
gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization
439
kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes
440
max_num_batched_tokens: int | None = None
@@ -942,6 +943,9 @@ class EngineArgs:
943
"--prefix-caching-hash-algo", **cache_kwargs["prefix_caching_hash_algo"]
944
)
945
cache_group.add_argument("--cpu-offload-gb", **cache_kwargs["cpu_offload_gb"])
946
+ cache_group.add_argument(
947
+ "--cpu-offload-params", **cache_kwargs["cpu_offload_params"]
948
+ )
949
cache_group.add_argument(
950
"--calculate-kv-scales", **cache_kwargs["calculate_kv_scales"]
951
)
@@ -1453,6 +1457,7 @@ class EngineArgs:
1457
enable_prefix_caching=self.enable_prefix_caching,
1458
prefix_caching_hash_algo=self.prefix_caching_hash_algo,
1459
cpu_offload_gb=self.cpu_offload_gb,
1460
+ cpu_offload_params=self.cpu_offload_params,
1461
calculate_kv_scales=self.calculate_kv_scales,
1462
kv_sharing_fast_prefill=self.kv_sharing_fast_prefill,
1463
mamba_cache_dtype=self.mamba_cache_dtype,
vllm/model_executor/models/utils.py
+23
-1
@@ -31,6 +31,7 @@ from vllm.model_executor.models.interfaces import supports_any_eagle
31
from vllm.multimodal import NestedTensors
32
from vllm.sequence import IntermediateTensors
33
from vllm.utils.math_utils import cdiv
34
+from vllm.utils.mem_utils import format_gib
35
from vllm.utils.platform_utils import (
36
is_pin_memory_available,
37
is_uva_available,
@@ -613,6 +614,7 @@ class PPMissingLayer(torch.nn.Identity):
614
615
_CPU_OFFLOAD_BYTES = 0
616
_CPU_OFFLOAD_MAX_BYTES = 0
617
+_CPU_OFFLOAD_PARAMS = set()
618
619
620
def set_cpu_offload_max_bytes(max_bytes: int) -> None:
@@ -621,6 +623,11 @@ def set_cpu_offload_max_bytes(max_bytes: int) -> None:
623
_CPU_OFFLOAD_MAX_BYTES = max_bytes
624
625
626
+def set_cpu_offload_params(params: set[str]) -> None:
627
+ global _CPU_OFFLOAD_PARAMS
628
+ _CPU_OFFLOAD_PARAMS = params
629
+
630
+
631
def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
632
if (params := next(module.parameters(), None)) is None:
633
return module
@@ -642,12 +649,23 @@ def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module:
649
# offload parameters to CPU
650
# use pin_memory if possible, which helps cudagraph capture speed
651
offloaded_parameters = False
645
- for p in module.parameters():
652
+ for name, p in module.named_parameters():
653
if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES:
654
# we use per-parameter offloading
655
# one module might have some parameters offloaded and some not
656
break
657
658
+ if _CPU_OFFLOAD_PARAMS:
659
+ # Check if parameter belongs to the offloading set
660
+ # Add dots here to ensure we match full segments only
661
+ # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight" but not
662
+ # "mlp.experts.w2_weight_scale"
663
+ should_offload = any(
664
+ f".{param}." in f".{name}." for param in _CPU_OFFLOAD_PARAMS
665
+ )
666
+ if not should_offload:
667
+ continue
668
+
669
cpu_data = p.data.to(device="cpu")
670
if pin_memory:
671
cpu_data = cpu_data.pin_memory()
@@ -708,6 +726,10 @@ def make_layers(
726
]
727
+ [PPMissingLayer() for _ in range(end_layer, num_hidden_layers)]
728
)
729
+ if _CPU_OFFLOAD_MAX_BYTES > 0:
730
+ logger.info(
731
+ "Total CPU offloaded parameters: %s GBs", format_gib(_CPU_OFFLOAD_BYTES)
732
+ )
733
return start_layer, end_layer, modules
734
735
vllm/v1/worker/gpu_model_runner.py
+5
-1
@@ -345,9 +345,13 @@ class GPUModelRunner(
345
self.speculative_config = vllm_config.speculative_config
346
self.observability_config = vllm_config.observability_config
347
348
- from vllm.model_executor.models.utils import set_cpu_offload_max_bytes
348
+ from vllm.model_executor.models.utils import (
349
+ set_cpu_offload_max_bytes,
350
+ set_cpu_offload_params,
351
+ )
352
353
set_cpu_offload_max_bytes(int(self.cache_config.cpu_offload_gb * 1024**3))
354
+ set_cpu_offload_params(self.cache_config.cpu_offload_params)
355
356
model_config = self.model_config
357
cache_config = self.cache_config