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 3, 03:00 UTC
[ROCm][CI] add warmup to mem_util test before measurement (#43236)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
May 21, 2026 at 23:17 UTC
35d0141a0b68a188777e277e372f211098419f58
1 file changed
+7
tests/utils_/test_mem_utils.py
+7
@@ -20,6 +20,13 @@ def test_memory_profiling():
20
# 512 MiB allocation outside of this instance
21
handle1 = lib.cudaMalloc(512 * 1024 * 1024)
22
23
+ # Warm up PyTorch's CUDA/ROCm context so that its internal initialization
24
+ # overhead (streams, cuBLAS handles, etc.) is included in the baseline and
25
+ # does not inflate non-torch increase which is larger on ROCm than on CUDA
26
+ _warmup = torch.zeros(1, device="cuda")
27
+ del _warmup
28
+ torch.accelerator.empty_cache()
29
+
30
baseline_snapshot = MemorySnapshot()
31
32
# load weights