[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))},
     )