[Intel XPU]Add xpu support for get_device_memory_capacity (#13895)

This commit is contained in:
gaopengff
2025-11-27 12:55:52 +08:00
committed by GitHub
parent 7cb04dc0e5
commit 077ca70ee4

View File

@@ -1697,6 +1697,15 @@ def get_cpu_memory_capacity():
return float(numa_mem // (1 << 20))
def get_xpu_memory_capacity():
try:
if torch.xpu.is_available():
return torch.xpu.mem_get_info()[1] // 1024 // 1024 # unit: MB
raise ValueError("No GPU memory values found.")
except AttributeError:
raise RuntimeError("torch.xpu is not available.")
def get_device_memory_capacity(device: str = None):
if is_cuda():
gpu_mem = get_nvgpu_memory_capacity()
@@ -1708,6 +1717,8 @@ def get_device_memory_capacity(device: str = None):
gpu_mem = get_npu_memory_capacity()
elif device == "cpu":
gpu_mem = get_cpu_memory_capacity()
elif device == "xpu":
gpu_mem = get_xpu_memory_capacity()
else:
# GPU memory is not known yet or no GPU is available.
gpu_mem = None