[Platform] Replace torch.cuda.mem_get_info with torch.accelerator.get_memory_info (#44825)

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>

tests/basic_correctness/test_mem.py

@@ -23,7 +23,7 @@ def test_python_error():
     error happening from the C++ side.
     """
     allocator = get_mem_allocator_instance()
-    total_bytes = current_platform.mem_get_info()[1]
+    total_bytes = torch.accelerator.get_memory_info()[1]
     alloc_bytes = int(total_bytes * 0.7)
     tensors = []
     with allocator.use_memory_pool():
@@ -64,9 +64,9 @@ def test_basic_cumem():
     output = x + y + z
     assert torch.allclose(output, torch.ones_like(output) * 3)
 
-    free_bytes = current_platform.mem_get_info()[0]
+    free_bytes = torch.accelerator.get_memory_info()[0]
     allocator.sleep()
-    free_bytes_after_sleep = current_platform.mem_get_info()[0]
+    free_bytes_after_sleep = torch.accelerator.get_memory_info()[0]
     assert free_bytes_after_sleep > free_bytes
     allocator.wake_up()
 
@@ -99,9 +99,9 @@ def test_cumem_with_cudagraph():
     with torch.cuda.graph(model_graph):
         y = model(x)
 
-    free_bytes = current_platform.mem_get_info()[0]
+    free_bytes = torch.accelerator.get_memory_info()[0]
     allocator.sleep()
-    free_bytes_after_sleep = current_platform.mem_get_info()[0]
+    free_bytes_after_sleep = torch.accelerator.get_memory_info()[0]
     assert free_bytes_after_sleep > free_bytes
     allocator.wake_up()
 
@@ -132,7 +132,7 @@ def test_cumem_with_cudagraph():
     ],
 )
 def test_end_to_end(model: str):
-    free, total = current_platform.mem_get_info()
+    free, total = torch.accelerator.get_memory_info()
     used_bytes_baseline = total - free  # in case other process is running
     llm = LLM(model, enable_sleep_mode=True)
     prompt = "How are you?"
@@ -144,7 +144,7 @@ def test_end_to_end(model: str):
     # test sleep level 1 here.
     llm.sleep(level=1)
 
-    free_gpu_bytes_after_sleep, total = current_platform.mem_get_info()
+    free_gpu_bytes_after_sleep, total = torch.accelerator.get_memory_info()
     used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
     # now the memory usage is mostly cudagraph memory pool,
     # and it should be less than the model weights (1B model, 2GiB weights)
@@ -164,7 +164,7 @@ def test_end_to_end(model: str):
     llm.sleep(level=1)
     llm.wake_up(tags=["weights"])
 
-    free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()
+    free_gpu_bytes_wake_up_w, total = torch.accelerator.get_memory_info()
     used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
 
     # should just reallocate memory for weights (1B model, ~2GiB weights)
@@ -181,7 +181,7 @@ def test_end_to_end(model: str):
 @create_new_process_for_each_test()
 def test_deep_sleep():
     model = "hmellor/tiny-random-LlamaForCausalLM"
-    free, total = current_platform.mem_get_info()
+    free, total = torch.accelerator.get_memory_info()
     used_bytes_baseline = total - free  # in case other process is running
     llm = LLM(model, enable_sleep_mode=True)
     prompt = "How are you?"
@@ -191,13 +191,13 @@ def test_deep_sleep():
     # Put the engine to deep sleep
     llm.sleep(level=2)
 
-    free_gpu_bytes_after_sleep, total = current_platform.mem_get_info()
+    free_gpu_bytes_after_sleep, total = torch.accelerator.get_memory_info()
     used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline
     assert used_bytes < 3 * GiB_bytes
 
     llm.wake_up(tags=["weights"])
     llm.collective_rpc("reload_weights")
-    free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()
+    free_gpu_bytes_wake_up_w, total = torch.accelerator.get_memory_info()
     used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
     assert used_bytes < 4 * GiB_bytes
 
@@ -213,7 +213,7 @@ def test_deep_sleep():
 def test_deep_sleep_async():
     async def test():
         model = "hmellor/tiny-random-LlamaForCausalLM"
-        free, total = current_platform.mem_get_info()
+        free, total = torch.accelerator.get_memory_info()
         used_bytes_baseline = total - free  # in case other process is running
         engine_args = AsyncEngineArgs(
             model=model,
@@ -232,7 +232,7 @@ def test_deep_sleep_async():
 
         await llm.wake_up(tags=["weights"])
         await llm.collective_rpc("reload_weights")
-        free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()
+        free_gpu_bytes_wake_up_w, total = torch.accelerator.get_memory_info()
         used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline
         assert used_bytes < 4 * GiB_bytes
 

tests/kernels/moe/test_moe.py

@@ -406,7 +406,7 @@ def test_fused_moe_int64_overflow(workspace_init):
     Reproduces the scenario from PR #34279.
     """
     # ~12 GB GPU memory needed for intermediate caches
-    free_mem = torch.cuda.mem_get_info()[0]
+    free_mem = torch.accelerator.get_memory_info()[0]
     if free_mem < 12 * 1024**3:
         pytest.skip("Insufficient GPU memory for overflow test")
 

tests/models/multimodal/generation/test_memory_leak.py

@@ -83,7 +83,7 @@ def _ru_maxrss_bytes() -> int | None:
 
 def _gpu_used_bytes() -> int:
     torch.accelerator.synchronize()
-    free_bytes, total_bytes = current_platform.mem_get_info()
+    free_bytes, total_bytes = torch.accelerator.get_memory_info()
     return int(total_bytes - free_bytes)
 
 

tests/utils_/test_mem_utils.py

@@ -36,7 +36,7 @@ def test_memory_profiling():
     weights_memory = 128 * 1024 * 1024 * 4  # 512 MiB
 
     def measure_current_non_torch():
-        free, total = torch.cuda.mem_get_info()
+        free, total = torch.accelerator.get_memory_info()
         current_used = total - free
         current_torch = torch.accelerator.memory_reserved()
         current_non_torch = current_used - current_torch
@@ -81,8 +81,9 @@ def test_memory_snapshot_uses_psutil_on_integrated_gpu():
     with (
         patch("vllm.utils.mem_utils.current_platform") as mock_platform,
         patch("vllm.utils.mem_utils.psutil") as mock_psutil,
+        patch("torch.accelerator") as mock_accelerator,
     ):
-        mock_platform.mem_get_info.return_value = (
+        mock_accelerator.get_memory_info.return_value = (
             mock_cuda_free,
             mock_cuda_total,
         )
@@ -90,8 +91,8 @@ def test_memory_snapshot_uses_psutil_on_integrated_gpu():
         mock_platform.memory_stats.return_value = {
             "allocated_bytes.all.peak": 0,
         }
-        mock_platform.memory_reserved.return_value = 0
-        mock_platform.current_device = lambda: "cuda:0"
+        mock_accelerator.memory_reserved.return_value = 0
+        mock_accelerator.current_device = lambda: "cuda:0"
 
         mock_vmem = MagicMock()
         mock_vmem.available = mock_psutil_available
@@ -105,24 +106,25 @@ def test_memory_snapshot_uses_psutil_on_integrated_gpu():
 
 
 def test_memory_snapshot_uses_cuda_on_discrete_gpu():
-    """On discrete GPUs, free_memory should come from CUDA mem_get_info."""
+    """On discrete GPUs, free_memory should come from accelerator  get_memory_info."""
     mock_cuda_free = 70 * 1024**3
     mock_cuda_total = 80 * 1024**3
 
     with (
         patch("vllm.utils.mem_utils.current_platform") as mock_platform,
         patch("vllm.utils.mem_utils.psutil") as mock_psutil,
+        patch("torch.accelerator") as mock_accelerator,
     ):
-        mock_platform.mem_get_info.return_value = (
+        mock_accelerator.get_memory_info.return_value = (
             mock_cuda_free,
             mock_cuda_total,
         )
         mock_platform.is_integrated_gpu.return_value = False
-        mock_platform.memory_stats.return_value = {
+        mock_accelerator.memory_stats.return_value = {
             "allocated_bytes.all.peak": 0,
         }
-        mock_platform.memory_reserved.return_value = 0
-        mock_platform.current_device = lambda: "cuda:0"
+        mock_accelerator.memory_reserved.return_value = 0
+        mock_accelerator.current_device = lambda: "cuda:0"
 
         snapshot = MemorySnapshot(device="cuda:0")
 

tests/v1/kv_connector/unit/test_rixl_gpu_mem_diag.py

@@ -29,9 +29,8 @@ def _gpu_snapshot(tag: str, prev_alloc: float = 0.0) -> dict:
     torch.accelerator.synchronize()
     alloc = torch.accelerator.memory_allocated()
     reserved = torch.accelerator.memory_reserved()
-    # mem_get_info is not available on torch.accelerator
     try:
-        drv_free, drv_total = torch.cuda.mem_get_info()
+        drv_free, drv_total = torch.accelerator.get_memory_info()
         drv_used = drv_total - drv_free
         drv_pct = drv_used / drv_total * 100
     except Exception:

tests/v1/sample/test_logprobs.py

@@ -1285,7 +1285,7 @@ def test_token_logprobs_large_batch_int64_row_offset():
     batch_size = 2**31 // vocab_size + 64  # batch_size * vocab_size > 2**31
     # logits (the large input) plus small logprob/rank outputs; ~1 GB headroom.
     required_bytes = batch_size * vocab_size * 4 + (1 << 30)
-    if torch.cuda.mem_get_info()[0] < required_bytes:
+    if torch.accelerator.get_memory_info()[0] < required_bytes:
         pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory")
 
     logits = torch.randn(batch_size, vocab_size, device=device, dtype=torch.float32)

tests/v1/sample/test_topk_topp_sampler.py

@@ -426,7 +426,7 @@ class TestTritonTopkTopp:
         # logits is modified in place; the only extra device memory is the
         # per-SM scratch buffer (~num_sm * vocab), so allow ~1 GB of headroom.
         required_bytes = batch_size * vocab_size * 4 + (1 << 30)
-        if torch.cuda.mem_get_info()[0] < required_bytes:
+        if torch.accelerator.get_memory_info()[0] < required_bytes:
             pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory")
 
         logits = torch.randn(

tools/pre_commit/check_torch_cuda.py

@@ -8,11 +8,12 @@ import regex as re
 # Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
 # --------------------------------------------------------------------------- #
 _TORCH_CUDA_PATTERNS = [
-    r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|set_device|device\()\b",
+    r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|mem_get_info|set_device|device\()\b",
     r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b",
     r"\bwith\storch\.cuda\.device\b",
     # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally
     r"\bcuda_device_count_stateless\(\)\b",
+    r"\bcurrent_platform\.mem_get_info\(\)\b",
 ]
 
 ALLOWED_FILES = {

vllm/model_executor/models/gemma4_mm.py

@@ -63,7 +63,6 @@ from vllm.multimodal.processing.processor import (
     PromptUpdate,
     PromptUpdateDetails,
 )
-from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
@@ -1274,7 +1273,7 @@ class Gemma4ForConditionalGeneration(
         # pass has already allocated activations we should account for.
         last_hidden_states_map: dict[int, torch.Tensor] = {}
         for patches, items in buckets.items():
-            free, total = current_platform.mem_get_info()
+            free, total = torch.accelerator.get_memory_info()
             max_batch_size = min(
                 len(items),
                 self._encoder_chunk(
@@ -1382,7 +1381,7 @@ class Gemma4ForConditionalGeneration(
             fc_list = list(frame_counts)
 
         total_frames = pixel_values.shape[0]
-        free, total = current_platform.mem_get_info()
+        free, total = torch.accelerator.get_memory_info()
         max_batch_size = min(
             total_frames,
             self._encoder_chunk(

vllm/platforms/cpu.py

@@ -92,11 +92,6 @@ class CpuPlatform(Platform):
 
         return meminfo.total_memory
 
-    @classmethod
-    def mem_get_info(cls) -> tuple[int, int]:
-        meminfo = get_memory_node_info()
-        return meminfo.available_memory, meminfo.total_memory
-
     @classmethod
     def set_device(cls, device: torch.device) -> None:
         """

vllm/utils/mem_utils.py

@@ -143,7 +143,7 @@ class MemorySnapshot:
             "allocated_bytes.all.peak", 0
         )
 
-        self.free_memory, self.total_memory = current_platform.mem_get_info(device)
+        self.free_memory, self.total_memory = torch.accelerator.get_memory_info(device)
         if current_platform.is_integrated_gpu(device.index):
             # On UMA (Unified Memory Architecture) platforms where CPU and
             # GPU share physical memory (e.g. GH200, DGX Spark, Jetson Orin),

vllm/v1/worker/cpu/shm.py

@@ -38,6 +38,14 @@ class _StreamPlaceholder:
         pass
 
 
+from vllm.utils.cpu_resource_utils import get_memory_node_info
+
+
+def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]:
+    meminfo = get_memory_node_info()
+    return meminfo.available_memory, meminfo.total_memory
+
+
 torch.Event = _EventPlaceholder
 torch.cuda.Event = _EventPlaceholder
 torch.cuda.Stream = _StreamPlaceholder
@@ -46,6 +54,7 @@ torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()
 torch.accelerator.synchronize = noop
 torch.accelerator.empty_cache = noop
 torch.Tensor.pin_memory = fake_pin_memory
+torch.accelerator.get_memory_info = get_memory_info
 
 # Patch vLLM torch utils
 import vllm.utils.torch_utils as torch_utils

vllm/v1/worker/gpu/model_runner.py

@@ -701,7 +701,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         start_time = time.perf_counter()
         gc.collect()
         torch.accelerator.empty_cache()
-        start_free_gpu_memory = torch.cuda.mem_get_info()[0]
+        start_free_gpu_memory = torch.accelerator.get_memory_info()[0]
 
         with self.maybe_setup_dummy_loras(self.lora_config):
             attn_states = self.cudagraph_manager.capture(
@@ -720,7 +720,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 self.speculator.capture(attn_states)
 
         end_time = time.perf_counter()
-        end_free_gpu_memory = torch.cuda.mem_get_info()[0]
+        end_free_gpu_memory = torch.accelerator.get_memory_info()[0]
         elapsed_time = end_time - start_time
         cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory
         # This usually takes 5~20 seconds.

vllm/v1/worker/gpu/spec_decode/eagle/utils.py

@@ -20,7 +20,7 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:
     # Use the faster GPU path when there is plenty of headroom;
     # otherwise compare on CPU.
     w = draft.weight
-    if w.is_cuda and torch.cuda.mem_get_info(w.device)[0] < w.numel() * 2:
+    if w.is_cuda and torch.accelerator.get_memory_info(w.device)[0] < w.numel() * 2:
         return torch.equal(w.cpu(), target.weight.cpu())
     return torch.equal(w, target.weight)
 

vllm/v1/worker/gpu_model_runner.py

@@ -6527,7 +6527,7 @@ class GPUModelRunner(
                     mem_samples: list[int] = []
 
                     for i, desc in enumerate(profile_descs):
-                        mem_before = torch.cuda.mem_get_info()[0]
+                        mem_before = torch.accelerator.get_memory_info()[0]
                         self._warmup_and_capture(
                             desc,
                             cudagraph_runtime_mode=mode,
@@ -6541,7 +6541,7 @@ class GPUModelRunner(
                             ),
                         )
                         torch.accelerator.synchronize()
-                        free_after = torch.cuda.mem_get_info()[0]
+                        free_after = torch.accelerator.get_memory_info()[0]
                         mem_samples.append(mem_before - free_after)
 
                     first_capture = mem_samples[0]
@@ -6563,10 +6563,10 @@ class GPUModelRunner(
                     )
 
                 if encoder_cudagraph_manager is not None:
-                    mem_before = torch.cuda.mem_get_info()[0]
+                    mem_before = torch.accelerator.get_memory_info()[0]
                     encoder_cudagraph_manager.capture(graph_pool=encoder_profiling_pool)
                     torch.accelerator.synchronize()
-                    free_after = torch.cuda.mem_get_info()[0]
+                    free_after = torch.accelerator.get_memory_info()[0]
                     encoder_memory_estimate = max(mem_before - free_after, 0)
 
                     logger.debug(
@@ -6632,7 +6632,7 @@ class GPUModelRunner(
         with self._freeze_gc(), graph_capture(device=self.device):
             torch.accelerator.synchronize()
             torch.accelerator.empty_cache()
-            start_free_gpu_memory = torch.cuda.mem_get_info()[0]
+            start_free_gpu_memory = torch.accelerator.get_memory_info()[0]
 
             for (
                 runtime_mode,
@@ -6650,7 +6650,7 @@ class GPUModelRunner(
                 self.encoder_cudagraph_manager.capture(graph_pool=encoder_graph_pool)
 
             torch.accelerator.synchronize()
-            end_free_gpu_memory = torch.cuda.mem_get_info()[0]
+            end_free_gpu_memory = torch.accelerator.get_memory_info()[0]
 
         # Disable cudagraph capturing globally, so any unexpected cudagraph
         # capturing will be detected and raise an error after here.

vllm/v1/worker/gpu_worker.py

@@ -172,7 +172,7 @@ class Worker(WorkerBase):
 
     def sleep(self, level: int = 1) -> None:
         torch.accelerator.synchronize()
-        free_bytes_before_sleep = current_platform.mem_get_info()[0]
+        free_bytes_before_sleep = torch.accelerator.get_memory_info()[0]
 
         # Save the buffers before level 2 sleep
         if level == 2:
@@ -187,7 +187,7 @@ class Worker(WorkerBase):
         torch.accelerator.synchronize()
         deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)
         while True:
-            free_bytes_after_sleep, total = current_platform.mem_get_info()
+            free_bytes_after_sleep, total = torch.accelerator.get_memory_info()
             freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
             if freed_bytes >= 0 or time.monotonic() >= deadline:
                 break
@@ -459,8 +459,8 @@ class Worker(WorkerBase):
             )
 
             # Profile CUDA graph memory if graphs will be captured.
-            # Skip on ROCm/HIP/XPU as graph pool handles and mem_get_info behave
-            # differently and can produce incorrect/negative estimates.
+            # Skip on ROCm/HIP/XPU as graph pool handles and get_memory_info
+            # behave differently and can produce incorrect/negative estimates.
             cudagraph_memory_estimate = 0
             if (
                 current_platform.is_cuda()

vllm/v1/worker/xpu_model_runner.py

@@ -45,7 +45,6 @@ def _torch_cuda_wrapper():
     torch.cuda.default_stream = torch.xpu.current_stream
     torch.cuda.current_stream = torch.xpu.current_stream
     torch.cuda.stream = torch.xpu.stream
-    torch.cuda.mem_get_info = torch.xpu.mem_get_info
     torch.cuda.Event = torch.Event
     torch.cuda.set_stream = torch.xpu.set_stream
     if supports_xpu_graph():