[XPU] sleep mode: fix KV cache release test (#57485)
Signed-off-by: Yan Ma <yan.ma@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Yan Ma committed
Sep 18, 2026 at 23:49 UTC
4991f976698d8ebf1a3b6d743c959b52bac689cc
1 file changed
+17
-8
tests/basic_correctness/test_mem.py
+17
-8
@@ -128,31 +128,40 @@ def test_release_kv_cache_memory_preserves_generation(full_sleep, monkeypatch):
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(worker):
132
- return mapped_usage(get_mem_allocator_instance())
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
134
- mapped_before = llm.collective_rpc(get_mapped_bytes)[0]
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()
138
- assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before
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()
144
- mapped_after = llm.collective_rpc(get_mapped_bytes)[0]
145
- assert mapped_before - mapped_after >= kv_cache_memory_bytes * 0.99
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)
151
- assert llm.collective_rpc(get_mapped_bytes)[0] == 0
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")
155
- assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before
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