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 11, 23:45 UTC
[ROCm][CI] Increase the memory threshold for test_deep_sleep_fp8_kvcache (#30104)
Signed-off-by: charlifu <charlifu@amd.com>
Charlie Fu committed
Dec 4, 2025 at 22:51 UTC
2c22c4ca2d88df503818741470a3ffc21f30a4b4
1 file changed
+7
-2
tests/basic_correctness/test_cumem.py
+7
-2
@@ -260,13 +260,18 @@ def test_deep_sleep_fp8_kvcache():
260
llm.sleep(level=2)
261
262
used_bytes = current_platform.get_current_memory_usage() - used_bytes_baseline
263
- assert used_bytes < 3 * GiB_bytes
263
+
264
+ # Rocm uses more memory for CudaGraphs, so we add 2 GiB more for the threshold
265
+ rocm_extra_mem_bytes = 2 * GiB_bytes if current_platform.is_rocm() else 0
266
+ mem_threshold_after_sleep = 3 * GiB_bytes + rocm_extra_mem_bytes
267
+ assert used_bytes < mem_threshold_after_sleep
268
269
llm.wake_up(tags=["weights"])
270
llm.collective_rpc("reload_weights")
271
272
used_bytes = current_platform.get_current_memory_usage() - used_bytes_baseline
269
- assert used_bytes < 4 * GiB_bytes
273
+ mem_threshold_after_wake_up = 4 * GiB_bytes + rocm_extra_mem_bytes
274
+ assert used_bytes < mem_threshold_after_wake_up
275
276
# now allocate kv cache and cuda graph memory
277
llm.wake_up(tags=["kv_cache"])