[XPU] C++ implementation for get_memory_info (#47134)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Qiming Zhang committed
Jun 30, 2026 at 06:34 UTC
91055efd363a42b79a7236a485b5079704e76ec6
1 file changed
+72
vllm/platforms/xpu.py
+72
@@ -28,6 +28,78 @@ else:
28
logger = init_logger(__name__)
29
30
31
+def get_mem_info_wrapper(
32
+ device: int | str | torch.device | None = None,
33
+) -> tuple[int, int]:
34
+ """
35
+ Get memory info for a device, compatible with torch.accelerator.get_memory_info API.
36
+
37
+ Args:
38
+ device: Device specification. Can be:
39
+ - None: Use current XPU device
40
+ - int: Device index
41
+ - str: Device string (e.g., "xpu:0", "xpu")
42
+ - torch.device: Device object
43
+
44
+ Returns:
45
+ Tuple[int, int]: (free_memory, total_memory) in bytes
46
+ """
47
+ # Handle None - use current device
48
+ if device is None:
49
+ device = torch.xpu.current_device()
50
+
51
+ # Handle torch.device objects
52
+ elif isinstance(device, torch.device):
53
+ if device.type != "xpu":
54
+ raise RuntimeError(f"Expected 'xpu' device, got '{device.type}'")
55
+ # If device index is not specified, use current device
56
+ device = (
57
+ device.index if device.index is not None else torch.xpu.current_device()
58
+ )
59
+
60
+ # Handle string device specifications (e.g., "xpu:0", "xpu")
61
+ elif isinstance(device, str):
62
+ if not device.startswith("xpu"):
63
+ raise RuntimeError(f"Expected 'xpu' device string, got '{device}'")
64
+ # Parse device string
65
+ parts = device.split(":")
66
+ if len(parts) == 1:
67
+ # "xpu" -> use current device
68
+ device = torch.xpu.current_device()
69
+ elif len(parts) == 2:
70
+ # "xpu:0" -> use index 0
71
+ try:
72
+ device = int(parts[1])
73
+ except ValueError as err:
74
+ raise RuntimeError(
75
+ f"Invalid device index: '{device}', expected integer after ':'"
76
+ ) from err
77
+ else:
78
+ raise RuntimeError(f"Invalid device string format: '{device}'")
79
+
80
+ # At this point, device should be an int
81
+ if isinstance(device, int):
82
+ # bounds check
83
+ device_count = torch.xpu.device_count()
84
+ if not (0 <= device < device_count):
85
+ raise ValueError(
86
+ f"Invalid device index {device}, must be in range [0, {device_count})"
87
+ )
88
+
89
+ elif not isinstance(device, int):
90
+ raise TypeError(
91
+ f"device must be int, str, torch.device, or None, got {type(device)}"
92
+ )
93
+
94
+ # Call the underlying C++ implementation
95
+ free, total = torch.ops._C_cache_ops.getMemoryInfo(device)
96
+
97
+ return free, total
98
+
99
+
100
+torch.accelerator.get_memory_info = get_mem_info_wrapper
101
+
102
+
103
class XPUPlatform(Platform):
104
_enum = PlatformEnum.XPU
105
device_name: str = "xpu"