@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fb42e5219e

[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>

Kunshang Ji committed Jun 30, 2026 at 14:39 UTC fb42e5219edcbce66fb1e758c004e610e618f70a

17 files changed +56 -52

tests/basic_correctness/test_mem.py

+13 -13

@@ -23,7 +23,7 @@ def test_python_error():

23 error happening from the C++ side.

24 """

25 allocator = get_mem_allocator_instance()

26 - total_bytes = current_platform.mem_get_info()[1]

26 + total_bytes = torch.accelerator.get_memory_info()[1]

27 alloc_bytes = int(total_bytes * 0.7)

28 tensors = []

29 with allocator.use_memory_pool():

@@ -64,9 +64,9 @@ def test_basic_cumem():

64 output = x + y + z

65 assert torch.allclose(output, torch.ones_like(output) * 3)

67 - free_bytes = current_platform.mem_get_info()[0]

67 + free_bytes = torch.accelerator.get_memory_info()[0]

68 allocator.sleep()

69 - free_bytes_after_sleep = current_platform.mem_get_info()[0]

69 + free_bytes_after_sleep = torch.accelerator.get_memory_info()[0]

70 assert free_bytes_after_sleep > free_bytes

71 allocator.wake_up()

@@ -99,9 +99,9 @@ def test_cumem_with_cudagraph():

99 with torch.cuda.graph(model_graph):

100 y = model(x)

101

102 - free_bytes = current_platform.mem_get_info()[0]

102 + free_bytes = torch.accelerator.get_memory_info()[0]

103 allocator.sleep()

104 - free_bytes_after_sleep = current_platform.mem_get_info()[0]

104 + free_bytes_after_sleep = torch.accelerator.get_memory_info()[0]

105 assert free_bytes_after_sleep > free_bytes

106 allocator.wake_up()

107

@@ -132,7 +132,7 @@ def test_cumem_with_cudagraph():

132 ],

133 )

134 def test_end_to_end(model: str):

135 - free, total = current_platform.mem_get_info()

135 + free, total = torch.accelerator.get_memory_info()

136 used_bytes_baseline = total - free # in case other process is running

137 llm = LLM(model, enable_sleep_mode=True)

138 prompt = "How are you?"

@@ -144,7 +144,7 @@ def test_end_to_end(model: str):

144 # test sleep level 1 here.

145 llm.sleep(level=1)

146

147 - free_gpu_bytes_after_sleep, total = current_platform.mem_get_info()

147 + free_gpu_bytes_after_sleep, total = torch.accelerator.get_memory_info()

148 used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline

149 # now the memory usage is mostly cudagraph memory pool,

150 # and it should be less than the model weights (1B model, 2GiB weights)

@@ -164,7 +164,7 @@ def test_end_to_end(model: str):

164 llm.sleep(level=1)

165 llm.wake_up(tags=["weights"])

166

167 - free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()

167 + free_gpu_bytes_wake_up_w, total = torch.accelerator.get_memory_info()

168 used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline

169

170 # should just reallocate memory for weights (1B model, ~2GiB weights)

@@ -181,7 +181,7 @@ def test_end_to_end(model: str):

181 @create_new_process_for_each_test()

182 def test_deep_sleep():

183 model = "hmellor/tiny-random-LlamaForCausalLM"

184 - free, total = current_platform.mem_get_info()

184 + free, total = torch.accelerator.get_memory_info()

185 used_bytes_baseline = total - free # in case other process is running

186 llm = LLM(model, enable_sleep_mode=True)

187 prompt = "How are you?"

@@ -191,13 +191,13 @@ def test_deep_sleep():

191 # Put the engine to deep sleep

192 llm.sleep(level=2)

193

194 - free_gpu_bytes_after_sleep, total = current_platform.mem_get_info()

194 + free_gpu_bytes_after_sleep, total = torch.accelerator.get_memory_info()

195 used_bytes = total - free_gpu_bytes_after_sleep - used_bytes_baseline

196 assert used_bytes < 3 * GiB_bytes

197

198 llm.wake_up(tags=["weights"])

199 llm.collective_rpc("reload_weights")

200 - free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()

200 + free_gpu_bytes_wake_up_w, total = torch.accelerator.get_memory_info()

201 used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline

202 assert used_bytes < 4 * GiB_bytes

203

@@ -213,7 +213,7 @@ def test_deep_sleep():

213 def test_deep_sleep_async():

214 async def test():

215 model = "hmellor/tiny-random-LlamaForCausalLM"

216 - free, total = current_platform.mem_get_info()

216 + free, total = torch.accelerator.get_memory_info()

217 used_bytes_baseline = total - free # in case other process is running

218 engine_args = AsyncEngineArgs(

219 model=model,

@@ -232,7 +232,7 @@ def test_deep_sleep_async():

232

233 await llm.wake_up(tags=["weights"])

234 await llm.collective_rpc("reload_weights")

235 - free_gpu_bytes_wake_up_w, total = current_platform.mem_get_info()

235 + free_gpu_bytes_wake_up_w, total = torch.accelerator.get_memory_info()

236 used_bytes = total - free_gpu_bytes_wake_up_w - used_bytes_baseline

237 assert used_bytes < 4 * GiB_bytes

238

tests/kernels/moe/test_moe.py

+1 -1

@@ -406,7 +406,7 @@ def test_fused_moe_int64_overflow(workspace_init):

406 Reproduces the scenario from PR #34279.

407 """

408 # ~12 GB GPU memory needed for intermediate caches

409 - free_mem = torch.cuda.mem_get_info()[0]

409 + free_mem = torch.accelerator.get_memory_info()[0]

410 if free_mem < 12 * 1024**3:

411 pytest.skip("Insufficient GPU memory for overflow test")

412

tests/models/multimodal/generation/test_memory_leak.py

+1 -1

@@ -83,7 +83,7 @@ def _ru_maxrss_bytes() -> int | None:

84 def _gpu_used_bytes() -> int:

85 torch.accelerator.synchronize()

86 - free_bytes, total_bytes = current_platform.mem_get_info()

86 + free_bytes, total_bytes = torch.accelerator.get_memory_info()

87 return int(total_bytes - free_bytes)

tests/utils_/test_mem_utils.py

+11 -9

@@ -36,7 +36,7 @@ def test_memory_profiling():

36 weights_memory = 128 * 1024 * 1024 * 4 # 512 MiB

38 def measure_current_non_torch():

39 - free, total = torch.cuda.mem_get_info()

39 + free, total = torch.accelerator.get_memory_info()

40 current_used = total - free

41 current_torch = torch.accelerator.memory_reserved()

42 current_non_torch = current_used - current_torch

@@ -81,8 +81,9 @@ def test_memory_snapshot_uses_psutil_on_integrated_gpu():

81 with (

82 patch("vllm.utils.mem_utils.current_platform") as mock_platform,

83 patch("vllm.utils.mem_utils.psutil") as mock_psutil,

84 + patch("torch.accelerator") as mock_accelerator,

85 ):

85 - mock_platform.mem_get_info.return_value = (

86 + mock_accelerator.get_memory_info.return_value = (

87 mock_cuda_free,

88 mock_cuda_total,

89 )

@@ -90,8 +91,8 @@ def test_memory_snapshot_uses_psutil_on_integrated_gpu():

91 mock_platform.memory_stats.return_value = {

92 "allocated_bytes.all.peak": 0,

93 }

93 - mock_platform.memory_reserved.return_value = 0

94 - mock_platform.current_device = lambda: "cuda:0"

94 + mock_accelerator.memory_reserved.return_value = 0

95 + mock_accelerator.current_device = lambda: "cuda:0"

97 mock_vmem = MagicMock()

98 mock_vmem.available = mock_psutil_available

@@ -105,24 +106,25 @@ def test_memory_snapshot_uses_psutil_on_integrated_gpu():

106

107

108 def test_memory_snapshot_uses_cuda_on_discrete_gpu():

108 - """On discrete GPUs, free_memory should come from CUDA mem_get_info."""

109 + """On discrete GPUs, free_memory should come from accelerator get_memory_info."""

110 mock_cuda_free = 70 * 1024**3

111 mock_cuda_total = 80 * 1024**3

112

113 with (

114 patch("vllm.utils.mem_utils.current_platform") as mock_platform,

115 patch("vllm.utils.mem_utils.psutil") as mock_psutil,

116 + patch("torch.accelerator") as mock_accelerator,

117 ):

116 - mock_platform.mem_get_info.return_value = (

118 + mock_accelerator.get_memory_info.return_value = (

119 mock_cuda_free,

120 mock_cuda_total,

121 )

122 mock_platform.is_integrated_gpu.return_value = False

121 - mock_platform.memory_stats.return_value = {

123 + mock_accelerator.memory_stats.return_value = {

124 "allocated_bytes.all.peak": 0,

125 }

124 - mock_platform.memory_reserved.return_value = 0

125 - mock_platform.current_device = lambda: "cuda:0"

126 + mock_accelerator.memory_reserved.return_value = 0

127 + mock_accelerator.current_device = lambda: "cuda:0"

128

129 snapshot = MemorySnapshot(device="cuda:0")

130

tests/v1/kv_connector/unit/test_rixl_gpu_mem_diag.py

+1 -2

@@ -29,9 +29,8 @@ def _gpu_snapshot(tag: str, prev_alloc: float = 0.0) -> dict:

29 torch.accelerator.synchronize()

30 alloc = torch.accelerator.memory_allocated()

31 reserved = torch.accelerator.memory_reserved()

32 - # mem_get_info is not available on torch.accelerator

32 try:

34 - drv_free, drv_total = torch.cuda.mem_get_info()

33 + drv_free, drv_total = torch.accelerator.get_memory_info()

34 drv_used = drv_total - drv_free

35 drv_pct = drv_used / drv_total * 100

36 except Exception:

tests/v1/sample/test_logprobs.py

+1 -1

@@ -1285,7 +1285,7 @@ def test_token_logprobs_large_batch_int64_row_offset():

1285 batch_size = 2**31 // vocab_size + 64 # batch_size * vocab_size > 2**31

1286 # logits (the large input) plus small logprob/rank outputs; ~1 GB headroom.

1287 required_bytes = batch_size * vocab_size * 4 + (1 << 30)

1288 - if torch.cuda.mem_get_info()[0] < required_bytes:

1288 + if torch.accelerator.get_memory_info()[0] < required_bytes:

1289 pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory")

1290

1291 logits = torch.randn(batch_size, vocab_size, device=device, dtype=torch.float32)

tests/v1/sample/test_topk_topp_sampler.py

+1 -1

@@ -426,7 +426,7 @@ class TestTritonTopkTopp:

426 # logits is modified in place; the only extra device memory is the

427 # per-SM scratch buffer (~num_sm * vocab), so allow ~1 GB of headroom.

428 required_bytes = batch_size * vocab_size * 4 + (1 << 30)

429 - if torch.cuda.mem_get_info()[0] < required_bytes:

429 + if torch.accelerator.get_memory_info()[0] < required_bytes:

430 pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory")

431

432 logits = torch.randn(

tools/pre_commit/check_torch_cuda.py

+2 -1

@@ -8,11 +8,12 @@ import regex as re

8 # Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`

9 # --------------------------------------------------------------------------- #

10 _TORCH_CUDA_PATTERNS = [

12 r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b",

13 r"\bwith\storch\.cuda\.device\b",

14 # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally

15 r"\bcuda_device_count_stateless\(\)\b",

16 + r"\bcurrent_platform\.mem_get_info\(\)\b",

17 ]

19 ALLOWED_FILES = {

vllm/model_executor/models/gemma4_mm.py

+2 -3

@@ -63,7 +63,6 @@ from vllm.multimodal.processing.processor import (

63 PromptUpdate,

64 PromptUpdateDetails,

65 )

66 -from vllm.platforms import current_platform

66 from vllm.sequence import IntermediateTensors

67 from vllm.utils.tensor_schema import TensorSchema, TensorShape

@@ -1274,7 +1273,7 @@ class Gemma4ForConditionalGeneration(

1273 # pass has already allocated activations we should account for.

1274 last_hidden_states_map: dict[int, torch.Tensor] = {}

1275 for patches, items in buckets.items():

1277 - free, total = current_platform.mem_get_info()

1276 + free, total = torch.accelerator.get_memory_info()

1277 max_batch_size = min(

1278 len(items),

1279 self._encoder_chunk(

@@ -1382,7 +1381,7 @@ class Gemma4ForConditionalGeneration(

1381 fc_list = list(frame_counts)

1382

1383 total_frames = pixel_values.shape[0]

1385 - free, total = current_platform.mem_get_info()

1384 + free, total = torch.accelerator.get_memory_info()

1385 max_batch_size = min(

1386 total_frames,

1387 self._encoder_chunk(

vllm/platforms/cpu.py

-5

@@ -92,11 +92,6 @@ class CpuPlatform(Platform):

93 return meminfo.total_memory

95 - @classmethod

96 - def mem_get_info(cls) -> tuple[int, int]:

97 - meminfo = get_memory_node_info()

98 - return meminfo.available_memory, meminfo.total_memory

99 -

95 @classmethod

96 def set_device(cls, device: torch.device) -> None:

97 """

vllm/utils/mem_utils.py

+1 -1

@@ -143,7 +143,7 @@ class MemorySnapshot:

143 "allocated_bytes.all.peak", 0

144 )

145

146 - self.free_memory, self.total_memory = current_platform.mem_get_info(device)

146 + self.free_memory, self.total_memory = torch.accelerator.get_memory_info(device)

147 if current_platform.is_integrated_gpu(device.index):

148 # On UMA (Unified Memory Architecture) platforms where CPU and

149 # GPU share physical memory (e.g. GH200, DGX Spark, Jetson Orin),

vllm/v1/worker/cpu/shm.py

@@ -38,6 +38,14 @@ class _StreamPlaceholder:

38 pass

41 +from vllm.utils.cpu_resource_utils import get_memory_node_info

42 +

43 +

44 +def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]:

45 + meminfo = get_memory_node_info()

46 + return meminfo.available_memory, meminfo.total_memory

47 +

48 +

49 torch.Event = _EventPlaceholder

50 torch.cuda.Event = _EventPlaceholder

51 torch.cuda.Stream = _StreamPlaceholder

@@ -46,6 +54,7 @@ torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()

54 torch.accelerator.synchronize = noop

55 torch.accelerator.empty_cache = noop

56 torch.Tensor.pin_memory = fake_pin_memory

57 +torch.accelerator.get_memory_info = get_memory_info

59 # Patch vLLM torch utils

60 import vllm.utils.torch_utils as torch_utils

vllm/v1/worker/gpu/model_runner.py

+2 -2

@@ -701,7 +701,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

701 start_time = time.perf_counter()

702 gc.collect()

703 torch.accelerator.empty_cache()

704 - start_free_gpu_memory = torch.cuda.mem_get_info()[0]

704 + start_free_gpu_memory = torch.accelerator.get_memory_info()[0]

705

706 with self.maybe_setup_dummy_loras(self.lora_config):

707 attn_states = self.cudagraph_manager.capture(

@@ -720,7 +720,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

720 self.speculator.capture(attn_states)

721

722 end_time = time.perf_counter()

723 - end_free_gpu_memory = torch.cuda.mem_get_info()[0]

723 + end_free_gpu_memory = torch.accelerator.get_memory_info()[0]

724 elapsed_time = end_time - start_time

725 cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory

726 # This usually takes 5~20 seconds.

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

+1 -1

@@ -20,7 +20,7 @@ def _should_share(eagle: nn.Module, flag: str, draft, target) -> bool:

20 # Use the faster GPU path when there is plenty of headroom;

21 # otherwise compare on CPU.

22 w = draft.weight

23 - if w.is_cuda and torch.cuda.mem_get_info(w.device)[0] < w.numel() * 2:

23 + if w.is_cuda and torch.accelerator.get_memory_info(w.device)[0] < w.numel() * 2:

24 return torch.equal(w.cpu(), target.weight.cpu())

25 return torch.equal(w, target.weight)

vllm/v1/worker/gpu_model_runner.py

+6 -6

@@ -6527,7 +6527,7 @@ class GPUModelRunner(

6527 mem_samples: list[int] = []

6528

6529 for i, desc in enumerate(profile_descs):

6530 - mem_before = torch.cuda.mem_get_info()[0]

6530 + mem_before = torch.accelerator.get_memory_info()[0]

6531 self._warmup_and_capture(

6532 desc,

6533 cudagraph_runtime_mode=mode,

@@ -6541,7 +6541,7 @@ class GPUModelRunner(

6541 ),

6542 )

6543 torch.accelerator.synchronize()

6544 - free_after = torch.cuda.mem_get_info()[0]

6544 + free_after = torch.accelerator.get_memory_info()[0]

6545 mem_samples.append(mem_before - free_after)

6546

6547 first_capture = mem_samples[0]

@@ -6563,10 +6563,10 @@ class GPUModelRunner(

6563 )

6564

6565 if encoder_cudagraph_manager is not None:

6566 - mem_before = torch.cuda.mem_get_info()[0]

6566 + mem_before = torch.accelerator.get_memory_info()[0]

6567 encoder_cudagraph_manager.capture(graph_pool=encoder_profiling_pool)

6568 torch.accelerator.synchronize()

6569 - free_after = torch.cuda.mem_get_info()[0]

6569 + free_after = torch.accelerator.get_memory_info()[0]

6570 encoder_memory_estimate = max(mem_before - free_after, 0)

6571

6572 logger.debug(

@@ -6632,7 +6632,7 @@ class GPUModelRunner(

6632 with self._freeze_gc(), graph_capture(device=self.device):

6633 torch.accelerator.synchronize()

6634 torch.accelerator.empty_cache()

6635 - start_free_gpu_memory = torch.cuda.mem_get_info()[0]

6635 + start_free_gpu_memory = torch.accelerator.get_memory_info()[0]

6636

6637 for (

6638 runtime_mode,

@@ -6650,7 +6650,7 @@ class GPUModelRunner(

6650 self.encoder_cudagraph_manager.capture(graph_pool=encoder_graph_pool)

6651

6652 torch.accelerator.synchronize()

6653 - end_free_gpu_memory = torch.cuda.mem_get_info()[0]

6653 + end_free_gpu_memory = torch.accelerator.get_memory_info()[0]

6654

6655 # Disable cudagraph capturing globally, so any unexpected cudagraph

6656 # capturing will be detected and raise an error after here.

vllm/v1/worker/gpu_worker.py

+4 -4

@@ -172,7 +172,7 @@ class Worker(WorkerBase):

172

173 def sleep(self, level: int = 1) -> None:

174 torch.accelerator.synchronize()

175 - free_bytes_before_sleep = current_platform.mem_get_info()[0]

175 + free_bytes_before_sleep = torch.accelerator.get_memory_info()[0]

176

177 # Save the buffers before level 2 sleep

178 if level == 2:

@@ -187,7 +187,7 @@ class Worker(WorkerBase):

187 torch.accelerator.synchronize()

188 deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)

189 while True:

190 - free_bytes_after_sleep, total = current_platform.mem_get_info()

190 + free_bytes_after_sleep, total = torch.accelerator.get_memory_info()

191 freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep

192 if freed_bytes >= 0 or time.monotonic() >= deadline:

193 break

@@ -459,8 +459,8 @@ class Worker(WorkerBase):

459 )

460

461 # Profile CUDA graph memory if graphs will be captured.

462 - # Skip on ROCm/HIP/XPU as graph pool handles and mem_get_info behave

463 - # differently and can produce incorrect/negative estimates.

462 + # Skip on ROCm/HIP/XPU as graph pool handles and get_memory_info

463 + # behave differently and can produce incorrect/negative estimates.

464 cudagraph_memory_estimate = 0

465 if (

466 current_platform.is_cuda()

vllm/v1/worker/xpu_model_runner.py

-1

@@ -45,7 +45,6 @@ def _torch_cuda_wrapper():

45 torch.cuda.default_stream = torch.xpu.current_stream

46 torch.cuda.current_stream = torch.xpu.current_stream

47 torch.cuda.stream = torch.xpu.stream

48 - torch.cuda.mem_get_info = torch.xpu.mem_get_info

48 torch.cuda.Event = torch.Event

49 torch.cuda.set_stream = torch.xpu.set_stream

50 if supports_xpu_graph():