[MRv2 Feature] MR v2 weight offloading support (#51413)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 7, 2026 at 15:20 UTC
a671679e9f04742aecb2c9b35cd513a024c1321f
3 files changed
+19
-3
tests/basic_correctness/test_cpu_offload.py
+4
-2
@@ -8,8 +8,10 @@ from ..utils import compare_two_settings
8
9
@pytest.mark.parametrize("disable_pin_memory", [False, True])
10
@pytest.mark.parametrize("disable_uva", [False, True])
11
-def test_cpu_offload(disable_pin_memory, disable_uva):
11
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
12
+def test_cpu_offload(disable_pin_memory, disable_uva, use_v2_model_runner):
13
env_vars = {
14
+ "VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner)),
15
"VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY": str(int(disable_pin_memory)),
16
"VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": str(int(disable_uva)),
17
}
@@ -24,6 +26,6 @@ def test_cpu_offload(disable_pin_memory, disable_uva):
26
model="hmellor/tiny-random-LlamaForCausalLM",
27
arg1=[],
28
arg2=args,
27
- env1=None,
29
+ env1=env_vars,
30
env2=env_vars,
31
)
tests/basic_correctness/test_prefetch_offload.py
+6
-1
@@ -2,10 +2,13 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Test prefetch offloading correctness with Llama model."""
4
5
+import pytest
6
+
7
from ..utils import compare_two_settings
8
9
8
-def test_prefetch_offload_llama():
10
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
11
+def test_prefetch_offload_llama(use_v2_model_runner):
12
"""Test prefetch CPU offloading with Llama-3.2-1B-Instruct.
13
14
Compares outputs between:
@@ -30,4 +33,6 @@ def test_prefetch_offload_llama():
33
"down_proj",
34
],
35
[], # Baseline: no offloading
36
+ env1={"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))},
37
+ env2={"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))},
38
)
vllm/v1/worker/gpu/model_runner.py
+9
@@ -46,6 +46,11 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
46
initialize_mamba_ssu_backend,
47
)
48
from vllm.model_executor.model_loader import get_model_loader
49
+from vllm.model_executor.offloader import (
50
+ create_offloader,
51
+ get_offloader,
52
+ set_offloader,
53
+)
54
from vllm.multimodal import MULTIMODAL_REGISTRY
55
from vllm.multimodal.encoder_budget import (
56
MultiModalBudget,
@@ -292,6 +297,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
297
self.eplb = EPLBController(self.parallel_config, self.device)
298
self.routed_experts_capturer: RoutedExpertsCapturer | None = None
299
300
+ set_offloader(create_offloader(self.vllm_config.offload_config))
301
+
302
def update_max_model_len(self, max_model_len: int) -> None:
303
self.max_model_len = max_model_len
304
self.req_states.max_model_len = max_model_len
@@ -413,6 +420,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
420
device=self.device,
421
)
422
423
+ get_offloader().post_init()
424
+
425
def get_model(self) -> nn.Module:
426
return self.model
427