[CI] Split (H200 MIG/MI300) Basic Correctness into named jobs (#57054)
Signed-off-by: Thang Nguyen <thangnguyenvn647@gmail.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Codex <noreply@openai.com>
Thang Nguyen committed
Sep 25, 2026 at 20:05 UTC
da1aaec31fb366c0ded497c1fa4b79a2fcc99969
18 files changed
+461
-330
.buildkite/intel_jobs/basic_correctness_intel.yaml
+4
-3
@@ -18,11 +18,12 @@ steps:
18
VLLM_TEST_DEVICE: "xpu"
19
source_file_dependencies:
20
- vllm/
21
- - tests/basic_correctness/test_cumem.py
21
+ - tests/basic_correctness/memory/
22
+ - tests/basic_correctness/cpu_offload/
23
commands:
24
- >-
25
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
26
'cd tests &&
27
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
27
- pytest -v -s basic_correctness/test_cpu_offload.py &&
28
- pytest -v -s basic_correctness/test_mem.py'
28
+ pytest -v -s basic_correctness/cpu_offload &&
29
+ pytest -v -s basic_correctness/memory'
.buildkite/test-amd.yaml
+66
-9
@@ -1030,14 +1030,14 @@ steps:
1030
- tests/distributed/test_ray_v2_executor.py
1031
- tests/distributed/test_ray_v2_executor_e2e.py
1032
- tests/distributed/test_pipeline_parallel.py
1033
- - tests/basic_correctness/test_basic_correctness.py
1033
+ - tests/basic_correctness/models/test_basic_correctness.py
1034
- vllm/platforms/rocm.py
1035
commands:
1036
- export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
1037
- pytest -v -s distributed/test_ray_v2_executor.py
1038
- pytest -v -s distributed/test_ray_v2_executor_e2e.py
1039
- pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
1040
- - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
1040
+ - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/models/test_basic_correctness.py -k "ray"
1041
1042
- label: ":amd: (MI300) Distributed DP + EP" # TBD
1043
timeout_in_minutes: 50
@@ -2762,7 +2762,38 @@ steps:
2762
2763
#----------------------------------------------------- mi355 · basic_correctness -----------------------------------------------------#
2764
2765
-- label: ":amd: (MI355 DPX) Basic Correctness" # TBD
2765
+- label: ":amd: (MI355 DPX) Basic Correctness Models %N" # TBD
2766
+ timeout_in_minutes: 60
2767
+ parallelism: 2
2768
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
2769
+ dind: false
2770
+ agent_pool: mi355_dpx
2771
+ fast_check: true
2772
+ working_dir: "/vllm-workspace/tests"
2773
+ source_file_dependencies:
2774
+ - vllm/
2775
+ - tests/basic_correctness/
2776
+ optional: true
2777
+ commands:
2778
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
2779
+ - VLLM_TARGET_TEST_SUITE=MI355 pytest -v -s basic_correctness/models --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
2780
+
2781
+- label: ":amd: (MI355 DPX) Basic Correctness CuMem" # TBD
2782
+ timeout_in_minutes: 60
2783
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
2784
+ dind: false
2785
+ agent_pool: mi355_dpx
2786
+ fast_check: true
2787
+ working_dir: "/vllm-workspace/tests"
2788
+ source_file_dependencies:
2789
+ - vllm/
2790
+ - tests/basic_correctness/
2791
+ optional: true
2792
+ commands:
2793
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
2794
+ - pytest -v -s basic_correctness/memory/cumem
2795
+
2796
+- label: ":amd: (MI355 DPX) Basic Correctness Sleep Mode" # TBD
2797
timeout_in_minutes: 60
2798
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
2799
dind: false
@@ -2771,15 +2802,41 @@ steps:
2802
working_dir: "/vllm-workspace/tests"
2803
source_file_dependencies:
2804
- vllm/
2774
- - tests/basic_correctness/test_basic_correctness
2775
- - tests/basic_correctness/test_cpu_offload
2776
- - tests/basic_correctness/test_mem.py
2805
+ - tests/basic_correctness/
2806
+ optional: true
2807
+ commands:
2808
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
2809
+ - pytest -v -s basic_correctness/memory/sleep_mode
2810
+
2811
+- label: ":amd: (MI355 DPX) Basic Correctness CPU Offload" # TBD
2812
+ timeout_in_minutes: 60
2813
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
2814
+ dind: false
2815
+ agent_pool: mi355_dpx
2816
+ fast_check: true
2817
+ working_dir: "/vllm-workspace/tests"
2818
+ source_file_dependencies:
2819
+ - vllm/
2820
+ - tests/basic_correctness/
2821
+ optional: true
2822
+ commands:
2823
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
2824
+ - pytest -v -s basic_correctness/cpu_offload
2825
+
2826
+- label: ":amd: (MI355 DPX) Basic Correctness Prefetch Offload" # TBD
2827
+ timeout_in_minutes: 60
2828
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
2829
+ dind: false
2830
+ agent_pool: mi355_dpx
2831
+ fast_check: true
2832
+ working_dir: "/vllm-workspace/tests"
2833
+ source_file_dependencies:
2834
+ - vllm/
2835
+ - tests/basic_correctness/
2836
optional: true
2837
commands:
2838
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
2780
- - pytest -v -s basic_correctness/test_mem.py
2781
- - VLLM_TARGET_TEST_SUITE=MI355 pytest -v -s basic_correctness/test_basic_correctness.py
2782
- - pytest -v -s basic_correctness/test_cpu_offload.py
2839
+ - pytest -v -s basic_correctness/prefetch_offload
2840
2841
#-------------------------------------------------------- mi355 · benchmarks ---------------------------------------------------------#
2842
.buildkite/test_areas/basic_correctness.yaml
+75
-10
@@ -1,25 +1,90 @@
1
group: Basic Correctness
2
-depends_on:
2
+depends_on:
3
- image-build
4
steps:
5
-- label: ":nvidia: (H200 MIG 18GB) Basic Correctness"
5
+- label: ":nvidia: (H200 MIG 18GB) Basic Correctness Models %N"
6
key: basic-correctness
7
timeout_in_minutes: 68
8
device: h200_18gb
9
- source_file_dependencies:
9
+ parallelism: 2
10
+ source_file_dependencies: &basic-correctness-deps
11
- vllm/
12
- "!vllm/distributed/kv_transfer/"
12
- - tests/basic_correctness/test_basic_correctness
13
- - tests/basic_correctness/test_cpu_offload
14
- - tests/basic_correctness/test_mem.py
13
+ - tests/basic_correctness/
14
commands:
15
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
17
- - pytest -v -s basic_correctness/test_mem.py
18
- - pytest -v -s basic_correctness/test_basic_correctness.py
19
- - pytest -v -s basic_correctness/test_cpu_offload.py
16
+ - pytest -v -s basic_correctness/models --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
17
mirror:
18
amd:
22
- label: ":amd: (MI355 DPX) Basic Correctness"
19
+ label: ":amd: (MI355 DPX) Basic Correctness Models %N"
20
+ dind: false
21
+ device: mi355_dpx
22
+ timeout_in_minutes: 60
23
+ depends_on:
24
+ - image-build-amd
25
+
26
+- label: ":nvidia: (H200 MIG 18GB) Basic Correctness CuMem"
27
+ key: basic-correctness-cumem
28
+ timeout_in_minutes: 68
29
+ device: h200_18gb
30
+ source_file_dependencies: *basic-correctness-deps
31
+ commands:
32
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
33
+ - pytest -v -s basic_correctness/memory/cumem
34
+ mirror:
35
+ amd:
36
+ label: ":amd: (MI355 DPX) Basic Correctness CuMem"
37
+ dind: false
38
+ device: mi355_dpx
39
+ timeout_in_minutes: 60
40
+ depends_on:
41
+ - image-build-amd
42
+
43
+- label: ":nvidia: (H200 MIG 18GB) Basic Correctness Sleep Mode"
44
+ key: basic-correctness-sleep-mode
45
+ timeout_in_minutes: 68
46
+ device: h200_18gb
47
+ source_file_dependencies: *basic-correctness-deps
48
+ commands:
49
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
50
+ - pytest -v -s basic_correctness/memory/sleep_mode
51
+ mirror:
52
+ amd:
53
+ label: ":amd: (MI355 DPX) Basic Correctness Sleep Mode"
54
+ dind: false
55
+ device: mi355_dpx
56
+ timeout_in_minutes: 60
57
+ depends_on:
58
+ - image-build-amd
59
+
60
+- label: ":nvidia: (H200 MIG 18GB) Basic Correctness CPU Offload"
61
+ key: basic-correctness-cpu-offload
62
+ timeout_in_minutes: 68
63
+ device: h200_18gb
64
+ source_file_dependencies: *basic-correctness-deps
65
+ commands:
66
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
67
+ - pytest -v -s basic_correctness/cpu_offload
68
+ mirror:
69
+ amd:
70
+ label: ":amd: (MI355 DPX) Basic Correctness CPU Offload"
71
+ dind: false
72
+ device: mi355_dpx
73
+ timeout_in_minutes: 60
74
+ depends_on:
75
+ - image-build-amd
76
+
77
+- label: ":nvidia: (H200 MIG 18GB) Basic Correctness Prefetch Offload"
78
+ key: basic-correctness-prefetch-offload
79
+ timeout_in_minutes: 68
80
+ device: h200_18gb
81
+ source_file_dependencies: *basic-correctness-deps
82
+ commands:
83
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
84
+ - pytest -v -s basic_correctness/prefetch_offload
85
+ mirror:
86
+ amd:
87
+ label: ":amd: (MI355 DPX) Basic Correctness Prefetch Offload"
88
dind: false
89
device: mi355_dpx
90
timeout_in_minutes: 60
.buildkite/test_areas/distributed.yaml
+2
-2
@@ -525,11 +525,11 @@ steps:
525
- tests/distributed/test_ray_v2_executor.py
526
- tests/distributed/test_ray_v2_executor_e2e.py
527
- tests/distributed/test_pipeline_parallel.py
528
- - tests/basic_correctness/test_basic_correctness.py
528
+ - tests/basic_correctness/models/test_basic_correctness.py
529
commands:
530
- export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
531
- export NCCL_CUMEM_HOST_ENABLE=0
532
- pytest -v -s distributed/test_ray_v2_executor.py
533
- pytest -v -s distributed/test_ray_v2_executor_e2e.py
534
- pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
535
- - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
535
+ - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/models/test_basic_correctness.py -k "ray"
tests/basic_correctness/cpu_offload/__init__.py
tests/basic_correctness/cpu_offload/test_cpu_offload.py
renamed
+1
-1
@@ -15,7 +15,7 @@ from vllm.model_executor.offloader import (
15
)
16
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
17
18
-from ..utils import compare_two_settings
18
+from ...utils import compare_two_settings
19
20
21
@pytest.mark.parametrize("disable_pin_memory", [False, True])
tests/basic_correctness/memory/__init__.py
tests/basic_correctness/memory/cumem/__init__.py
tests/basic_correctness/memory/cumem/test_cumem.py
new
+302
@@ -0,0 +1,302 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+import vllm.device_allocator.cumem as cumem
8
+from vllm import LLM, SamplingParams
9
+from vllm.device_allocator import get_mem_allocator_instance
10
+from vllm.platforms import current_platform
11
+
12
+from ....utils import create_new_process_for_each_test
13
+
14
+DEVICE_TYPE = current_platform.device_type
15
+
16
+
17
+def mapped_usage(allocator) -> int:
18
+ """Bytes the allocator still has mapped on the device.
19
+
20
+ `torch.accelerator.get_memory_info()` is a device-wide `cudaMemGetInfo` /
21
+ `hipMemGetInfo` query, so its delta also carries whatever else on the GPU
22
+ allocates or frees while we measure. The allocator's own bookkeeping is
23
+ process-local and exact.
24
+ """
25
+ return sum(
26
+ data.handle[1]
27
+ for data in allocator.pointer_to_data.values()
28
+ if not data.is_asleep
29
+ )
30
+
31
+
32
+def _wake_up_with_poisoned_mappings(allocator, byte_value: int = 0xA5) -> None:
33
+ """Wake discarded allocations with deterministic nonzero contents."""
34
+ original_create_and_map = cumem.create_and_map
35
+
36
+ def create_and_map_with_poison(handle) -> None:
37
+ original_create_and_map(handle)
38
+ _, size, ptr, _ = handle
39
+ cumem.libcudart.cudaMemset(ptr, byte_value, size)
40
+
41
+ cumem.create_and_map = create_and_map_with_poison
42
+ try:
43
+ allocator.wake_up()
44
+ finally:
45
+ cumem.create_and_map = original_create_and_map
46
+
47
+
48
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
49
+def test_python_error():
50
+ """Test if Python error occurs when there's low-level
51
+ error happening from the C++ side.
52
+ """
53
+ allocator = get_mem_allocator_instance()
54
+ total_bytes = torch.accelerator.get_memory_info()[1]
55
+ alloc_bytes = int(total_bytes * 0.7)
56
+ tensors = []
57
+ with allocator.use_memory_pool():
58
+ # allocate 70% of the total memory
59
+ x = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
60
+ tensors.append(x)
61
+ # release the memory
62
+ allocator.sleep(offload_tags=())
63
+
64
+ # allocate more memory than the total memory
65
+ y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
66
+ tensors.append(y)
67
+ with pytest.raises(RuntimeError):
68
+ # when the allocator is woken up, it should raise an error
69
+ # because we don't have enough memory
70
+ allocator.wake_up()
71
+
72
+
73
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
74
+def test_basic_cumem():
75
+ # some tensors from default memory pool
76
+ shape = (1024, 1024)
77
+ x = torch.empty(shape, device=DEVICE_TYPE)
78
+ x.zero_()
79
+
80
+ # some tensors from custom memory pool
81
+ allocator = get_mem_allocator_instance()
82
+ with allocator.use_memory_pool():
83
+ # custom memory pool
84
+ y = torch.empty(shape, device=DEVICE_TYPE)
85
+ y.zero_()
86
+ y += 1
87
+ z = torch.empty(shape, device=DEVICE_TYPE)
88
+ z.zero_()
89
+ z += 2
90
+
91
+ # they can be used together
92
+ output = x + y + z
93
+ assert torch.allclose(output, torch.ones_like(output) * 3)
94
+
95
+ # Track the allocator's own mapped bytes, not device-wide free memory.
96
+ mapped_bytes = mapped_usage(allocator)
97
+ assert mapped_bytes >= y.nbytes + z.nbytes
98
+ allocator.sleep()
99
+ assert mapped_usage(allocator) == 0
100
+ allocator.wake_up()
101
+ assert mapped_usage(allocator) == mapped_bytes
102
+
103
+ # they can be used together
104
+ output = x + y + z
105
+ assert torch.allclose(output, torch.ones_like(output) * 3)
106
+
107
+
108
+@pytest.mark.parametrize("full_sleep", [0, 1, 2], ids=["kv-only", "sleep-1", "sleep-2"])
109
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
110
+def test_release_kv_cache_memory_preserves_generation(full_sleep, monkeypatch):
111
+ """Rejected release and KV-only/full sleep cycles preserve greedy tokens."""
112
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
113
+ kv_cache_memory_bytes = 256 * 1024 * 1024
114
+ llm = LLM(
115
+ "Qwen/Qwen3-0.6B",
116
+ enable_sleep_mode=True,
117
+ enforce_eager=True,
118
+ kv_cache_memory_bytes=kv_cache_memory_bytes,
119
+ max_model_len=1024,
120
+ max_num_seqs=4,
121
+ gpu_memory_utilization=0.05,
122
+ )
123
+ prompt = "How are you?"
124
+ sampling_params = SamplingParams(temperature=0, max_tokens=10)
125
+ expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
126
+
127
+ def get_mapped_bytes_by_tag(worker):
128
+ allocator = get_mem_allocator_instance()
129
+ mapped_bytes = mapped_usage(allocator)
130
+ kv_cache_bytes = sum(
131
+ data.handle[1]
132
+ for data in allocator.pointer_to_data.values()
133
+ if data.tag == "kv_cache" and not data.is_asleep
134
+ )
135
+ return mapped_bytes, kv_cache_bytes
136
+
137
+ mapped_before, kv_cache_before = llm.collective_rpc(get_mapped_bytes_by_tag)[0]
138
+ # Utility RPC transports the engine error as a plain Exception.
139
+ with pytest.raises(Exception, match="requires a completed pause"):
140
+ llm.release_kv_cache_memory()
141
+ assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == mapped_before
142
+ assert not llm.llm_engine.is_sleeping()
143
+ assert llm.generate(prompt, sampling_params)[0].outputs[0].token_ids == expected
144
+
145
+ llm.sleep(level=0)
146
+ llm.release_kv_cache_memory()
147
+ mapped_after, kv_cache_after = llm.collective_rpc(get_mapped_bytes_by_tag)[0]
148
+ assert kv_cache_before > 0
149
+ assert kv_cache_after == 0
150
+ assert mapped_before - mapped_after >= kv_cache_before
151
+ assert mapped_after > 0
152
+ assert llm.llm_engine.is_sleeping()
153
+
154
+ if full_sleep:
155
+ llm.sleep(level=full_sleep)
156
+ assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == 0
157
+ llm.wake_up(tags=None if full_sleep else ["kv_cache"])
158
+ if full_sleep == 2:
159
+ llm.collective_rpc("reload_weights")
160
+ assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == mapped_before
161
+ assert not llm.llm_engine.is_sleeping()
162
+ actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
163
+ assert actual == expected
164
+
165
+
166
+@pytest.mark.parametrize("first_level", [1, 2])
167
+@pytest.mark.parametrize("second_level", [1, 2])
168
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
169
+def test_sleep_with_only_weights_asleep(first_level, second_level, monkeypatch):
170
+ """Repeated sleep after KV-only wake preserves mappings and recoverability."""
171
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
172
+ llm = LLM(
173
+ "Qwen/Qwen3-0.6B",
174
+ enable_sleep_mode=True,
175
+ enforce_eager=True,
176
+ kv_cache_memory_bytes=256 * 1024 * 1024,
177
+ max_model_len=1024,
178
+ max_num_seqs=4,
179
+ gpu_memory_utilization=0.05,
180
+ )
181
+ prompt = "How are you?"
182
+ sampling_params = SamplingParams(temperature=0, max_tokens=10)
183
+ expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
184
+
185
+ def get_mapped_bytes(worker):
186
+ return mapped_usage(get_mem_allocator_instance())
187
+
188
+ llm.sleep(level=first_level)
189
+ llm.wake_up(tags=["kv_cache"])
190
+ mapped_before = llm.collective_rpc(get_mapped_bytes)[0]
191
+ assert mapped_before > 0
192
+ llm.sleep(level=second_level)
193
+ assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before
194
+ llm.wake_up(tags=["weights"])
195
+ if first_level == 2:
196
+ llm.collective_rpc("reload_weights")
197
+ assert not llm.llm_engine.is_sleeping()
198
+ actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
199
+ assert actual == expected
200
+
201
+
202
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
203
+def test_discard_tags():
204
+ """Test that discard(tags) selectively frees GPU memory for specific
205
+ tags while keeping other tags mapped and usable."""
206
+ allocator = get_mem_allocator_instance()
207
+
208
+ with allocator.use_memory_pool("weights"):
209
+ weights = torch.ones(1024, 1024, device=DEVICE_TYPE)
210
+
211
+ with allocator.use_memory_pool("kv_cache"):
212
+ kv = torch.ones(512, 512, device=DEVICE_TYPE)
213
+
214
+ mapped_bytes = mapped_usage(allocator)
215
+
216
+ # Discard kv_cache only — weights should remain valid
217
+ allocator.discard("kv_cache")
218
+
219
+ mapped_bytes_after_discard = mapped_usage(allocator)
220
+ assert mapped_bytes - mapped_bytes_after_discard >= kv.nbytes
221
+ assert mapped_bytes_after_discard >= weights.nbytes
222
+
223
+ # Weights are still usable
224
+ assert torch.allclose(weights, torch.ones_like(weights))
225
+
226
+ # Wake up and verify kv_cache is remapped; discarded contents are undefined.
227
+ allocator.wake_up()
228
+ assert kv.shape == (512, 512)
229
+
230
+ # Full sleep/wake cycle still works after discard
231
+ allocator.sleep(offload_tags="weights")
232
+ allocator.wake_up()
233
+ assert torch.allclose(weights, torch.ones_like(weights))
234
+
235
+
236
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
237
+@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
238
+def test_level2_discards_ordinary_tensor_with_weights_tag():
239
+ """Reproduce the level-2 variant for an ordinary tensor in weights."""
240
+ allocator = get_mem_allocator_instance()
241
+
242
+ with allocator.use_memory_pool("weights"):
243
+ fake_weight = torch.full((4096,), 0x44, dtype=torch.uint8, device=DEVICE_TYPE)
244
+ ordinary_tensor = torch.full(
245
+ (4096,), 0x55, dtype=torch.uint8, device=DEVICE_TYPE
246
+ )
247
+
248
+ pointers = (fake_weight.data_ptr(), ordinary_tensor.data_ptr())
249
+ allocator.sleep(offload_tags=())
250
+ _wake_up_with_poisoned_mappings(allocator)
251
+ torch.accelerator.synchronize()
252
+
253
+ assert (fake_weight.data_ptr(), ordinary_tensor.data_ptr()) == pointers
254
+ assert torch.all(fake_weight == 0xA5)
255
+ assert torch.all(ordinary_tensor == 0xA5)
256
+
257
+
258
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
259
+@pytest.mark.skipif(current_platform.is_xpu(), reason="CUDA graph not supported on XPU")
260
+def test_cumem_with_cudagraph():
261
+ allocator = get_mem_allocator_instance()
262
+ with allocator.use_memory_pool():
263
+ weight = torch.eye(1024, device=DEVICE_TYPE)
264
+ with allocator.use_memory_pool(tag="discard"):
265
+ cache = torch.empty(1024, 1024, device=DEVICE_TYPE)
266
+
267
+ def model(x):
268
+ out = x @ weight
269
+ cache[: out.size(0)].copy_(out)
270
+ return out + 1
271
+
272
+ x = torch.empty(128, 1024, device=DEVICE_TYPE)
273
+
274
+ # warmup
275
+ model(x)
276
+
277
+ # capture cudagraph
278
+ model_graph = torch.cuda.CUDAGraph()
279
+ with torch.cuda.graph(model_graph):
280
+ y = model(x)
281
+
282
+ mapped_bytes = mapped_usage(allocator)
283
+ assert mapped_bytes >= weight.nbytes + cache.nbytes
284
+ allocator.sleep()
285
+ assert mapped_usage(allocator) == 0
286
+ allocator.wake_up()
287
+ assert mapped_usage(allocator) == mapped_bytes
288
+
289
+ # after waking up, the content in the weight tensor
290
+ # should be restored, but the content in the cache tensor
291
+ # should be discarded
292
+
293
+ # this operation is also compatible with cudagraph
294
+
295
+ x.random_()
296
+ model_graph.replay()
297
+
298
+ # cache content is as expected
299
+ assert torch.allclose(x, cache[: x.size(0)])
300
+
301
+ # output content is as expected
302
+ assert torch.allclose(y, x + 1)
tests/basic_correctness/memory/sleep_mode/__init__.py
tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py
renamed
+3
-297
@@ -9,301 +9,10 @@ import torch
9
import vllm.device_allocator.cumem as cumem
10
import vllm.envs as envs
11
from vllm import LLM, AsyncEngineArgs, AsyncLLMEngine, SamplingParams
12
-from vllm.device_allocator import get_mem_allocator_instance
12
from vllm.platforms import current_platform
13
from vllm.utils.mem_constants import GiB_bytes
14
16
-from ..utils import create_new_process_for_each_test, requires_fp8
17
-
18
-DEVICE_TYPE = current_platform.device_type
19
-
20
-
21
-def mapped_usage(allocator) -> int:
22
- """Bytes the allocator still has mapped on the device.
23
-
24
- `torch.accelerator.get_memory_info()` is a device-wide `cudaMemGetInfo` /
25
- `hipMemGetInfo` query, so its delta also carries whatever else on the GPU
26
- allocates or frees while we measure. The allocator's own bookkeeping is
27
- process-local and exact.
28
- """
29
- return sum(
30
- data.handle[1]
31
- for data in allocator.pointer_to_data.values()
32
- if not data.is_asleep
33
- )
34
-
35
-
36
-def _wake_up_with_poisoned_mappings(allocator, byte_value: int = 0xA5) -> None:
37
- """Wake discarded allocations with deterministic nonzero contents."""
38
- original_create_and_map = cumem.create_and_map
39
-
40
- def create_and_map_with_poison(handle) -> None:
41
- original_create_and_map(handle)
42
- _, size, ptr, _ = handle
43
- cumem.libcudart.cudaMemset(ptr, byte_value, size)
44
-
45
- cumem.create_and_map = create_and_map_with_poison
46
- try:
47
- allocator.wake_up()
48
- finally:
49
- cumem.create_and_map = original_create_and_map
50
-
51
-
52
-@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
53
-def test_python_error():
54
- """Test if Python error occurs when there's low-level
55
- error happening from the C++ side.
56
- """
57
- allocator = get_mem_allocator_instance()
58
- total_bytes = torch.accelerator.get_memory_info()[1]
59
- alloc_bytes = int(total_bytes * 0.7)
60
- tensors = []
61
- with allocator.use_memory_pool():
62
- # allocate 70% of the total memory
63
- x = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
64
- tensors.append(x)
65
- # release the memory
66
- allocator.sleep(offload_tags=())
67
-
68
- # allocate more memory than the total memory
69
- y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
70
- tensors.append(y)
71
- with pytest.raises(RuntimeError):
72
- # when the allocator is woken up, it should raise an error
73
- # because we don't have enough memory
74
- allocator.wake_up()
75
-
76
-
77
-@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
78
-def test_basic_cumem():
79
- # some tensors from default memory pool
80
- shape = (1024, 1024)
81
- x = torch.empty(shape, device=DEVICE_TYPE)
82
- x.zero_()
83
-
84
- # some tensors from custom memory pool
85
- allocator = get_mem_allocator_instance()
86
- with allocator.use_memory_pool():
87
- # custom memory pool
88
- y = torch.empty(shape, device=DEVICE_TYPE)
89
- y.zero_()
90
- y += 1
91
- z = torch.empty(shape, device=DEVICE_TYPE)
92
- z.zero_()
93
- z += 2
94
-
95
- # they can be used together
96
- output = x + y + z
97
- assert torch.allclose(output, torch.ones_like(output) * 3)
98
-
99
- # Track the allocator's own mapped bytes, not device-wide free memory.
100
- mapped_bytes = mapped_usage(allocator)
101
- assert mapped_bytes >= y.nbytes + z.nbytes
102
- allocator.sleep()
103
- assert mapped_usage(allocator) == 0
104
- allocator.wake_up()
105
- assert mapped_usage(allocator) == mapped_bytes
106
-
107
- # they can be used together
108
- output = x + y + z
109
- assert torch.allclose(output, torch.ones_like(output) * 3)
110
-
111
-
112
-@pytest.mark.parametrize("full_sleep", [0, 1, 2], ids=["kv-only", "sleep-1", "sleep-2"])
113
-@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
114
-def test_release_kv_cache_memory_preserves_generation(full_sleep, monkeypatch):
115
- """Rejected release and KV-only/full sleep cycles preserve greedy tokens."""
116
- monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
117
- kv_cache_memory_bytes = 256 * 1024 * 1024
118
- llm = LLM(
119
- "Qwen/Qwen3-0.6B",
120
- enable_sleep_mode=True,
121
- enforce_eager=True,
122
- kv_cache_memory_bytes=kv_cache_memory_bytes,
123
- max_model_len=1024,
124
- max_num_seqs=4,
125
- gpu_memory_utilization=0.05,
126
- )
127
- prompt = "How are you?"
128
- sampling_params = SamplingParams(temperature=0, max_tokens=10)
129
- expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
130
-
131
- def get_mapped_bytes_by_tag(worker):
132
- allocator = get_mem_allocator_instance()
133
- mapped_bytes = mapped_usage(allocator)
134
- kv_cache_bytes = sum(
135
- data.handle[1]
136
- for data in allocator.pointer_to_data.values()
137
- if data.tag == "kv_cache" and not data.is_asleep
138
- )
139
- return mapped_bytes, kv_cache_bytes
140
-
141
- mapped_before, kv_cache_before = llm.collective_rpc(get_mapped_bytes_by_tag)[0]
142
- # Utility RPC transports the engine error as a plain Exception.
143
- with pytest.raises(Exception, match="requires a completed pause"):
144
- llm.release_kv_cache_memory()
145
- assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == mapped_before
146
- assert not llm.llm_engine.is_sleeping()
147
- assert llm.generate(prompt, sampling_params)[0].outputs[0].token_ids == expected
148
-
149
- llm.sleep(level=0)
150
- llm.release_kv_cache_memory()
151
- mapped_after, kv_cache_after = llm.collective_rpc(get_mapped_bytes_by_tag)[0]
152
- assert kv_cache_before > 0
153
- assert kv_cache_after == 0
154
- assert mapped_before - mapped_after >= kv_cache_before
155
- assert mapped_after > 0
156
- assert llm.llm_engine.is_sleeping()
157
-
158
- if full_sleep:
159
- llm.sleep(level=full_sleep)
160
- assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == 0
161
- llm.wake_up(tags=None if full_sleep else ["kv_cache"])
162
- if full_sleep == 2:
163
- llm.collective_rpc("reload_weights")
164
- assert llm.collective_rpc(get_mapped_bytes_by_tag)[0][0] == mapped_before
165
- assert not llm.llm_engine.is_sleeping()
166
- actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
167
- assert actual == expected
168
-
169
-
170
-@pytest.mark.parametrize("first_level", [1, 2])
171
-@pytest.mark.parametrize("second_level", [1, 2])
172
-@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
173
-def test_sleep_with_only_weights_asleep(first_level, second_level, monkeypatch):
174
- """Repeated sleep after KV-only wake preserves mappings and recoverability."""
175
- monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
176
- llm = LLM(
177
- "Qwen/Qwen3-0.6B",
178
- enable_sleep_mode=True,
179
- enforce_eager=True,
180
- kv_cache_memory_bytes=256 * 1024 * 1024,
181
- max_model_len=1024,
182
- max_num_seqs=4,
183
- gpu_memory_utilization=0.05,
184
- )
185
- prompt = "How are you?"
186
- sampling_params = SamplingParams(temperature=0, max_tokens=10)
187
- expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
188
-
189
- def get_mapped_bytes(worker):
190
- return mapped_usage(get_mem_allocator_instance())
191
-
192
- llm.sleep(level=first_level)
193
- llm.wake_up(tags=["kv_cache"])
194
- mapped_before = llm.collective_rpc(get_mapped_bytes)[0]
195
- assert mapped_before > 0
196
- llm.sleep(level=second_level)
197
- assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before
198
- llm.wake_up(tags=["weights"])
199
- if first_level == 2:
200
- llm.collective_rpc("reload_weights")
201
- assert not llm.llm_engine.is_sleeping()
202
- actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids
203
- assert actual == expected
204
-
205
-
206
-@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
207
-def test_discard_tags():
208
- """Test that discard(tags) selectively frees GPU memory for specific
209
- tags while keeping other tags mapped and usable."""
210
- allocator = get_mem_allocator_instance()
211
-
212
- with allocator.use_memory_pool("weights"):
213
- weights = torch.ones(1024, 1024, device=DEVICE_TYPE)
214
-
215
- with allocator.use_memory_pool("kv_cache"):
216
- kv = torch.ones(512, 512, device=DEVICE_TYPE)
217
-
218
- mapped_bytes = mapped_usage(allocator)
219
-
220
- # Discard kv_cache only — weights should remain valid
221
- allocator.discard("kv_cache")
222
-
223
- mapped_bytes_after_discard = mapped_usage(allocator)
224
- assert mapped_bytes - mapped_bytes_after_discard >= kv.nbytes
225
- assert mapped_bytes_after_discard >= weights.nbytes
226
-
227
- # Weights are still usable
228
- assert torch.allclose(weights, torch.ones_like(weights))
229
-
230
- # Wake up and verify kv_cache is remapped; discarded contents are undefined.
231
- allocator.wake_up()
232
- assert kv.shape == (512, 512)
233
-
234
- # Full sleep/wake cycle still works after discard
235
- allocator.sleep(offload_tags="weights")
236
- allocator.wake_up()
237
- assert torch.allclose(weights, torch.ones_like(weights))
238
-
239
-
240
-@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
241
-@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
242
-def test_level2_discards_ordinary_tensor_with_weights_tag():
243
- """Reproduce the level-2 variant for an ordinary tensor in weights."""
244
- allocator = get_mem_allocator_instance()
245
-
246
- with allocator.use_memory_pool("weights"):
247
- fake_weight = torch.full((4096,), 0x44, dtype=torch.uint8, device=DEVICE_TYPE)
248
- ordinary_tensor = torch.full(
249
- (4096,), 0x55, dtype=torch.uint8, device=DEVICE_TYPE
250
- )
251
-
252
- pointers = (fake_weight.data_ptr(), ordinary_tensor.data_ptr())
253
- allocator.sleep(offload_tags=())
254
- _wake_up_with_poisoned_mappings(allocator)
255
- torch.accelerator.synchronize()
256
-
257
- assert (fake_weight.data_ptr(), ordinary_tensor.data_ptr()) == pointers
258
- assert torch.all(fake_weight == 0xA5)
259
- assert torch.all(ordinary_tensor == 0xA5)
260
-
261
-
262
-@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
263
-@pytest.mark.skipif(current_platform.is_xpu(), reason="CUDA graph not supported on XPU")
264
-def test_cumem_with_cudagraph():
265
- allocator = get_mem_allocator_instance()
266
- with allocator.use_memory_pool():
267
- weight = torch.eye(1024, device=DEVICE_TYPE)
268
- with allocator.use_memory_pool(tag="discard"):
269
- cache = torch.empty(1024, 1024, device=DEVICE_TYPE)
270
-
271
- def model(x):
272
- out = x @ weight
273
- cache[: out.size(0)].copy_(out)
274
- return out + 1
275
-
276
- x = torch.empty(128, 1024, device=DEVICE_TYPE)
277
-
278
- # warmup
279
- model(x)
280
-
281
- # capture cudagraph
282
- model_graph = torch.cuda.CUDAGraph()
283
- with torch.cuda.graph(model_graph):
284
- y = model(x)
285
-
286
- mapped_bytes = mapped_usage(allocator)
287
- assert mapped_bytes >= weight.nbytes + cache.nbytes
288
- allocator.sleep()
289
- assert mapped_usage(allocator) == 0
290
- allocator.wake_up()
291
- assert mapped_usage(allocator) == mapped_bytes
292
-
293
- # after waking up, the content in the weight tensor
294
- # should be restored, but the content in the cache tensor
295
- # should be discarded
296
-
297
- # this operation is also compatible with cudagraph
298
-
299
- x.random_()
300
- model_graph.replay()
301
-
302
- # cache content is as expected
303
- assert torch.allclose(x, cache[: x.size(0)])
304
-
305
- # output content is as expected
306
- assert torch.allclose(y, x + 1)
15
+from ....utils import create_new_process_for_each_test, multi_gpu_test, requires_fp8
16
17
18
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
@@ -445,8 +154,8 @@ def _lora_logits_mapping_present(model) -> bool:
154
)
155
156
448
-@create_new_process_for_each_test()
449
-def test_deep_sleep_lora_tp2(num_gpus_available, monkeypatch):
157
+@multi_gpu_test(num_gpus=2)
158
+def test_deep_sleep_lora_tp2(monkeypatch):
159
"""Level-2 sleep/wake/reload with enable_lora=True and TP=2.
160
161
With TP > 1 the LoRA logits processor carries
@@ -455,9 +164,6 @@ def test_deep_sleep_lora_tp2(num_gpus_available, monkeypatch):
164
attribute allocated in the sleep-mode pool, so level-2 sleep destroys
165
its contents — it must be restored after reload.
166
"""
458
- if num_gpus_available < 2:
459
- pytest.skip("Requires at least 2 GPUs")
460
-
167
# Needed for apply_model to reach the multiproc TP workers below.
168
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
169
tests/basic_correctness/models/__init__.py
tests/basic_correctness/models/test_basic_correctness.py
renamed
+4
-4
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Compare the short outputs of HF and vLLM when using greedy sampling.
4
5
-Run `pytest tests/basic_correctness/test_basic_correctness.py`.
5
+Run `pytest tests/basic_correctness/models/test_basic_correctness.py`.
6
"""
7
8
import os
@@ -17,9 +17,9 @@ from vllm import LLM
17
from vllm.platforms import current_platform
18
from vllm.v1.engine.llm_engine import LLMEngine
19
20
-from ..conftest import HfRunner, VllmRunner
21
-from ..models.utils import check_outputs_equal
22
-from ..utils import multi_gpu_test
20
+from ...conftest import HfRunner, VllmRunner
21
+from ...models.utils import check_outputs_equal
22
+from ...utils import multi_gpu_test
23
24
ATTN_BACKEND = ["ROCM_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"]
25
tests/basic_correctness/prefetch_offload/__init__.py
tests/basic_correctness/prefetch_offload/test_prefetch_offload.py
renamed
+1
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Test prefetch offloading correctness with Llama model."""
4
5
-from ..utils import compare_two_settings
5
+from ...utils import compare_two_settings
6
7
8
def test_prefetch_offload_llama():
tests/conftest.py
+1
-1
@@ -928,7 +928,7 @@ class HfRunner:
928
# shutdown. This is helpful in cases where the HfRunner is
929
# initialized after significant GPU memory is already occupied,
930
# e.g. in
931
- # tests/basic_correctness/test_basic_correctness.py::test_models_distributed
931
+ # tests/basic_correctness/models/test_basic_correctness.py::test_models_distributed
932
# where vllm worker processes are still alive and holding GPU
933
# memory when hf_runner.__exit__ is called.
934
from tests.utils import record_gpu_memory_usage_stats
tests/quantization/test_cpu_offload.py
+1
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
# Expanded quantized model tests for CPU offloading
5
-# Base tests: tests/basic_correctness/test_cpu_offload.py
5
+# Base tests: tests/basic_correctness/cpu_offload/test_cpu_offload.py
6
7
import pytest
8
tests/v1/worker/test_sleep_mode_backend.py
+1
-1
@@ -4,7 +4,7 @@
4
5
These cover the registry/factory contract and capability flags. They do not
6
touch CUDA - the ``cumem`` suspend/resume path is exercised end-to-end on GPU
7
-in ``tests/basic_correctness/test_cumem.py``.
7
+in ``tests/basic_correctness/memory/``.
8
"""
9
10
import pytest