From c39076feff40b5638a2b911414bc09b8e6e953bc Mon Sep 17 00:00:00 2001
From: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Date: Fri, 7 Aug 2026 18:49:45 -0400
Subject: [PATCH] [CI Test] Add specific unit test for mrv2 offloading (#51440)

Signed-off-by: yewentao256 <zhyanwentao@126.com>
---
 tests/basic_correctness/test_cpu_offload.py   | 61 +++++++++++++++++--
 .../test_prefetch_offload.py                  |  7 +--
 2 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/tests/basic_correctness/test_cpu_offload.py b/tests/basic_correctness/test_cpu_offload.py
index c9ca342818..f3ea25e58a 100644
--- a/tests/basic_correctness/test_cpu_offload.py
+++ b/tests/basic_correctness/test_cpu_offload.py
@@ -3,15 +3,22 @@
 
 import pytest
 
+import vllm.envs as envs
+from vllm.model_executor.offloader import (
+    PrefetchOffloader,
+    UVAOffloader,
+    get_offloader,
+    set_offloader,
+)
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
 from ..utils import compare_two_settings
 
 
 @pytest.mark.parametrize("disable_pin_memory", [False, True])
 @pytest.mark.parametrize("disable_uva", [False, True])
-@pytest.mark.parametrize("use_v2_model_runner", [False, True])
-def test_cpu_offload(disable_pin_memory, disable_uva, use_v2_model_runner):
+def test_cpu_offload(disable_pin_memory, disable_uva):
     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)),
     }
@@ -26,6 +33,52 @@ def test_cpu_offload(disable_pin_memory, disable_uva, use_v2_model_runner):
         model="hmellor/tiny-random-LlamaForCausalLM",
         arg1=[],
         arg2=args,
-        env1=env_vars,
+        env1=None,
         env2=env_vars,
     )
+
+
+@pytest.mark.parametrize(
+    ("offload_kwargs", "offloader_type"),
+    [
+        ({"cpu_offload_gb": 1}, UVAOffloader),
+        (
+            {
+                "offload_group_size": 1,
+                "offload_num_in_group": 1,
+                "offload_prefetch_step": 1,
+            },
+            PrefetchOffloader,
+        ),
+    ],
+)
+def test_mrv2_weight_offloading(
+    vllm_runner, monkeypatch, offload_kwargs, offloader_type
+):
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    envs.disable_envs_cache()
+    original_offloader = get_offloader()
+
+    try:
+        with vllm_runner(
+            "hmellor/tiny-random-LlamaForCausalLM",
+            enforce_eager=True,
+            gpu_memory_utilization=0.02,
+            max_model_len=128,
+            max_num_seqs=1,
+            **offload_kwargs,
+        ) as vllm_model:
+            engine_core = vllm_model.llm.llm_engine.engine_core.engine_core
+            model_runner = engine_core.model_executor.driver_worker.worker.model_runner
+            assert isinstance(model_runner, GPUModelRunner)
+
+            offloader = get_offloader()
+            assert isinstance(offloader, offloader_type)
+            if isinstance(offloader, UVAOffloader):
+                assert offloader.cpu_offload_bytes > 0
+            else:
+                assert offloader.total_offloaded_bytes > 0
+                assert offloader.buffer_pool is not None
+    finally:
+        set_offloader(original_offloader)
+        envs.disable_envs_cache()
diff --git a/tests/basic_correctness/test_prefetch_offload.py b/tests/basic_correctness/test_prefetch_offload.py
index 125d364188..498887024e 100644
--- a/tests/basic_correctness/test_prefetch_offload.py
+++ b/tests/basic_correctness/test_prefetch_offload.py
@@ -2,13 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Test prefetch offloading correctness with Llama model."""
 
-import pytest
-
 from ..utils import compare_two_settings
 
 
-@pytest.mark.parametrize("use_v2_model_runner", [False, True])
-def test_prefetch_offload_llama(use_v2_model_runner):
+def test_prefetch_offload_llama():
     """Test prefetch CPU offloading with Llama-3.2-1B-Instruct.
 
     Compares outputs between:
@@ -33,6 +30,4 @@ def test_prefetch_offload_llama(use_v2_model_runner):
             "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))},
     )