A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 13:15 UTC
[CI] Stabilize test_cpu_offloading by waiting for async offload before cache reset (#37335)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 18, 2026 at 00:26 UTC
ce2ef42fd3ae16476e66f910ee776e23c95d5b81
1 file changed
+46
-7
tests/v1/kv_offload/test_cpu_offloading.py
+46
-7
@@ -22,6 +22,17 @@ if current_platform.is_cuda():
22
elif current_platform.is_rocm():
23
ATTN_BACKENDS = ["TRITON_ATTN"]
24
25
+# Maximum time (seconds) to wait for the async CPU offload transfer
26
+# to complete before giving up.
27
+_RESET_CACHE_TIMEOUT = 30 if current_platform.is_rocm() else 10
28
+
29
+# ZMQ poll timeout (ms) for the first event.
30
+_FIRST_EVENT_POLL_MS = 10_000 if current_platform.is_rocm() else 1000
31
+
32
+# Hard ceiling (seconds) on how long get_new_cpu_stored_events may loop,
33
+# to prevent hangs if non-CPU events keep arriving indefinitely.
34
+_EVENT_DRAIN_TIMEOUT = 60
35
+
36
37
class MockSubscriber:
38
"""Helper class to receive and verify published events"""
@@ -47,9 +58,10 @@ class MockSubscriber:
58
poller = zmq.Poller()
59
poller.register(self.sub, zmq.POLLIN)
60
50
- timeout = 1000 # 1 second
51
- while True:
52
- events = dict(poller.poll(timeout))
61
+ poll_ms = _FIRST_EVENT_POLL_MS
62
+ deadline = time.monotonic() + _EVENT_DRAIN_TIMEOUT
63
+ while time.monotonic() < deadline:
64
+ events = dict(poller.poll(poll_ms))
65
66
if events.get(self.sub) != zmq.POLLIN:
67
return cpu_stored_events
@@ -63,13 +75,32 @@ class MockSubscriber:
75
for event in event_batch.events:
76
if isinstance(event, BlockStored) and event.medium == "CPU":
77
cpu_stored_events.append(event)
66
- timeout = 100
78
+ poll_ms = 100
79
+
80
+ return cpu_stored_events
81
82
def close(self):
83
"""Clean up resources"""
84
self.sub.close()
85
86
87
+def _wait_for_prefix_cache_reset(llm: LLM) -> None:
88
+ """Wait for async offload transfers to finish so prefix cache can reset.
89
+
90
+ The GPU-to-CPU offload runs on a CUDA stream asynchronously. While blocks
91
+ are still held by the offload worker, ``reset_prefix_cache`` returns
92
+ ``False``. Retry with a short sleep until it succeeds or we time out.
93
+ """
94
+ deadline = time.monotonic() + _RESET_CACHE_TIMEOUT
95
+ while not llm.reset_prefix_cache():
96
+ if time.monotonic() > deadline:
97
+ raise TimeoutError(
98
+ "reset_prefix_cache did not succeed within "
99
+ f"{_RESET_CACHE_TIMEOUT}s - async offload may be stuck"
100
+ )
101
+ time.sleep(0.1)
102
+
103
+
104
def _latency_test(llm: LLM, subscriber: MockSubscriber):
105
sampling_params = SamplingParams(max_tokens=1)
106
@@ -95,10 +126,16 @@ def _latency_test(llm: LLM, subscriber: MockSubscriber):
126
gpu_hit_time = time.time() - start_time
127
total_gpu_hit_time += gpu_hit_time
128
98
- # reset prefix cache to avoid GPU hit.
99
- llm.reset_prefix_cache()
129
+ # Wait for the async CPU offload to finish, then reset prefix cache
130
+ # so the next generate() must reload from CPU rather than GPU.
131
+ _wait_for_prefix_cache_reset(llm)
132
101
- assert subscriber.get_new_cpu_stored_events()
133
+ # Verify CPU stored events arrived (offload is done before we
134
+ # attempt to load from CPU).
135
+ assert subscriber.get_new_cpu_stored_events(), (
136
+ f"No CPU stored events received on iteration {i}; "
137
+ "async offload may not have completed in time"
138
+ )
139
140
# run generation again - this should trigger loading from CPU
141
start_time = time.time()
@@ -185,6 +222,8 @@ def test_cpu_offloading(cpu_block_size: int, attn_backend: str) -> None:
222
kv_events_config=kv_events_config,
223
kv_transfer_config=kv_transfer_config,
224
attention_config={"backend": attn_backend},
225
+ # ROCm: batch size 1 to reduce variability
226
+ **({"max_num_seqs": 1} if current_platform.is_rocm() else {}),
227
)
228
229
events_endpoint = events_endpoint.replace("*", "127.0.0.1")