[MRv2 Feature] MR v2 weight offloading support (#51413)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

tests/basic_correctness/test_cpu_offload.py

@@ -8,8 +8,10 @@ from ..utils import compare_two_settings
 
 @pytest.mark.parametrize("disable_pin_memory", [False, True])
 @pytest.mark.parametrize("disable_uva", [False, True])
-def test_cpu_offload(disable_pin_memory, disable_uva):
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
+def test_cpu_offload(disable_pin_memory, disable_uva, use_v2_model_runner):
     env_vars = {
+        "VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner)),
         "VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY": str(int(disable_pin_memory)),
         "VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": str(int(disable_uva)),
     }
@@ -24,6 +26,6 @@ def test_cpu_offload(disable_pin_memory, disable_uva):
         model="hmellor/tiny-random-LlamaForCausalLM",
         arg1=[],
         arg2=args,
-        env1=None,
+        env1=env_vars,
         env2=env_vars,
     )

tests/basic_correctness/test_prefetch_offload.py

@@ -2,10 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Test prefetch offloading correctness with Llama model."""
 
+import pytest
+
 from ..utils import compare_two_settings
 
 
-def test_prefetch_offload_llama():
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
+def test_prefetch_offload_llama(use_v2_model_runner):
     """Test prefetch CPU offloading with Llama-3.2-1B-Instruct.
 
     Compares outputs between:
@@ -30,4 +33,6 @@ def test_prefetch_offload_llama():
             "down_proj",
         ],
         [],  # Baseline: no offloading
+        env1={"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))},
+        env2={"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))},
     )

vllm/v1/worker/gpu/model_runner.py

@@ -46,6 +46,11 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     initialize_mamba_ssu_backend,
 )
 from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.offloader import (
+    create_offloader,
+    get_offloader,
+    set_offloader,
+)
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.encoder_budget import (
     MultiModalBudget,
@@ -292,6 +297,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.eplb = EPLBController(self.parallel_config, self.device)
         self.routed_experts_capturer: RoutedExpertsCapturer | None = None
 
+        set_offloader(create_offloader(self.vllm_config.offload_config))
+
     def update_max_model_len(self, max_model_len: int) -> None:
         self.max_model_len = max_model_len
         self.req_states.max_model_len = max_model_len
@@ -413,6 +420,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 device=self.device,
             )
 
+        get_offloader().post_init()
+
     def get_model(self) -> nn.Module:
         return self.model