@@ -9,6 +9,8 @@
     unsupported_before,
 )
 
+# Keep device-API gating on individual tests: system-level NVML queries in
+# this module remain supported on platforms with partial device API support.
 pytestmark = skip_if_nvml_unsupported
 
 import array
 
@@ -5,6 +5,8 @@
 
 from cuda_python_test_helpers.arch_check import skip_if_nvml_device_apis_unsupported, skip_if_nvml_unsupported
 
+# Keep device-API gating on individual tests so the pure event conversion and
+# wrapping tests still run on platforms with partial device API support.
 pytestmark = skip_if_nvml_unsupported
 
 import helpers
 
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import ctypes
+import functools
 import multiprocessing as mp
 import sys
 
@@ -1506,37 +1507,46 @@ def __init__(self, size):
     assert ("release", NEW_HANDLE) in calls
 
 
+def _vmm_allocate_and_close(mr, requested_size, grow):
+    buf = mr.allocate(requested_size)
+    if grow:
+        buf = mr.modify_allocation(buf, 2 * buf.size)
+    aligned_size = buf.size
+    buf.close()
+    return aligned_size
+
+
+@functools.cache
+def _warm_up_vmm_allocate_and_close(device_id, grow):
+    device = Device(device_id)
+    mr = VirtualMemoryResource(
+        device,
+        config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"),
+    )
+    return _vmm_allocate_and_close(mr, 8 * 1024 * 1024, grow)
+
+
 @pytest.mark.thread_unsafe(reason="cuMemGetInfo measures process-wide free memory")
 @pytest.mark.parametrize("grow", [False, True], ids=["allocate", "grow"])
 def test_vmm_allocate_close_does_not_leak(init_cuda, grow):
     device = Device()
     if not device.properties.virtual_memory_management_supported:
         pytest.skip("Virtual memory management is not supported on this device")
 
+    aligned_size = _warm_up_vmm_allocate_and_close(device.device_id, grow)
     mr = VirtualMemoryResource(
         device,
         config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"),
     )
     requested_size = 8 * 1024 * 1024
 
-    def allocate_and_close():
-        buf = mr.allocate(requested_size)
-        if grow:
-            buf = mr.modify_allocation(buf, 2 * buf.size)
-        aligned_size = buf.size
-        buf.close()
-        return aligned_size
-
-    aligned_size = allocate_and_close()  # Warm up and learn the aligned allocation size.
-
     baseline = handle_return(driver.cuMemGetInfo())[0]
     for _ in range(8):
-        allocate_and_close()
+        _vmm_allocate_and_close(mr, requested_size, grow)
     free = handle_return(driver.cuMemGetInfo())[0]
 
-    # The broken path leaks aligned_size per iteration. Allow one allocation's
-    # worth of driver bookkeeping/caching while still detecting repeated leaks.
-    assert baseline - free < 2 * aligned_size
+    # A leak would cost at least one allocation per iteration.
+    assert baseline - free < aligned_size
 
 
 def test_vmm_allocator_rdma_unsupported_exception():
Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,8 @@`
`9`	`9`	`unsupported_before,`
`10`	`10`	`)`
`11`	`11`
	`12`	`+# Keep device-API gating on individual tests: system-level NVML queries in`
	`13`	`+# this module remain supported on platforms with partial device API support.`
`12`	`14`	`pytestmark = skip_if_nvml_unsupported`
`13`	`15`
`14`	`16`	`import array`