From fb42e5219edcbce66fb1e758c004e610e618f70a Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Tue, 30 Jun 2026 14:39:52 +0800 Subject: [PATCH] [Platform] Replace `torch.cuda.mem_get_info` with `torch.accelerator.get_memory_info` (#44825) Signed-off-by: Kunshang Ji Signed-off-by: Kunshang Ji --- tests/basic_correctness/test_mem.py | 26 +++++++++---------- tests/kernels/moe/test_moe.py | 2 +- .../multimodal/generation/test_memory_leak.py | 2 +- tests/utils_/test_mem_utils.py | 20 +++++++------- .../unit/test_rixl_gpu_mem_diag.py | 3 +-- tests/v1/sample/test_logprobs.py | 2 +- tests/v1/sample/test_topk_topp_sampler.py | 2 +- tools/pre_commit/check_torch_cuda.py | 3 ++- vllm/model_executor/models/gemma4_mm.py | 5 ++-- vllm/platforms/cpu.py | 5 ---- vllm/utils/mem_utils.py | 2 +- vllm/v1/worker/cpu/shm.py | 9 +++++++ vllm/v1/worker/gpu/model_runner.py | 4 +-- vllm/v1/worker/gpu/spec_decode/eagle/utils.py | 2 +- vllm/v1/worker/gpu_model_runner.py | 12 ++++----- vllm/v1/worker/gpu_worker.py | 8 +++--- vllm/v1/worker/xpu_model_runner.py | 1 - 17 files changed, 56 insertions(+), 52 deletions(-) diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index 2c9a99c500d..c0f8a592223 100644 --- a/tests/basic_correctness/test_mem.py +++ b/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 diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index f8b98c82a24..69c50cbb11f 100644 --- a/tests/kernels/moe/test_moe.py +++ b/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") diff --git a/tests/models/multimodal/generation/test_memory_leak.py b/tests/models/multimodal/generation/test_memory_leak.py index 5ee505257c1..45eac5b80ab 100644 --- a/tests/models/multimodal/generation/test_memory_leak.py +++ b/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) diff --git a/tests/utils_/test_mem_utils.py b/tests/utils_/test_mem_utils.py index 861e73c7ded..421aec3e9b1 100644 --- a/tests/utils_/test_mem_utils.py +++ b/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") diff --git a/tests/v1/kv_connector/unit/test_rixl_gpu_mem_diag.py b/tests/v1/kv_connector/unit/test_rixl_gpu_mem_diag.py index c3adc05e3ef..2371b5555bc 100644 --- a/tests/v1/kv_connector/unit/test_rixl_gpu_mem_diag.py +++ b/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: diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 5ed0a476279..ec150272792 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/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) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 8d906e83f2d..8a3d313f1d5 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/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( diff --git a/tools/pre_commit/check_torch_cuda.py b/tools/pre_commit/check_torch_cuda.py index bfbb36ffbff..aec7b85d59c 100644 --- a/tools/pre_commit/check_torch_cuda.py +++ b/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 = { diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 30c379d86c1..5e405c44621 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/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( diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index c5d7ec2fe71..369e07dd256 100644 --- a/vllm/platforms/cpu.py +++ b/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: """ diff --git a/vllm/utils/mem_utils.py b/vllm/utils/mem_utils.py index 3894742c6be..b0ac4b16e47 100644 --- a/vllm/utils/mem_utils.py +++ b/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), diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py index d691ada90b2..e21e3712975 100644 --- a/vllm/v1/worker/cpu/shm.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c9f4362a6fb..f94d96c8330 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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. diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index ed441b380f0..11961ceef4d 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/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) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index ff1eba09fd0..a772a897a16 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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. diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 07c3615edcb..2947101d977 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/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() diff --git a/vllm/v1/worker/xpu_model_runner.py b/vllm/v1/worker/xpu_model_runner.py index f93d0439557..05cbb6bc958 100644 --- a/vllm/v1/worker/xpu_model_runner.py +++ b/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():