[Intel GPU] calculate free memory based on allocated memory for XPU (#32044)

Signed-off-by: P V R K Jyothendra Varma <polisetty.v.r.k.jyothendra.varma@intel.com>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Polisetty V R K Jyothendra Varma
2026-07-24 08:38:52 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent 433429b16a
commit 2f823a2eee
2 changed files with 6 additions and 34 deletions
@@ -103,23 +103,9 @@ class XpuPlatform(Platform):
if empty_cache:
torch.xpu.empty_cache()
# Use mem_get_info() with a sanity cap to avoid KV-cache over-allocation
# on drivers that incorrectly return total memory as free memory.
# Consistent with the fallback: free = max(0, total - allocated).
try:
free_gpu_memory, total_gpu_memory = torch.xpu.mem_get_info(device_id)
used_memory = float(torch.xpu.memory_allocated(device_id))
free_gpu_memory = min(
float(free_gpu_memory),
max(0.0, float(total_gpu_memory) - used_memory),
)
except Exception:
# Fallback for devices/drivers that do not support querying free memory
used_memory = float(torch.xpu.memory_allocated(device_id))
total_gpu_memory = float(
torch.xpu.get_device_properties(device_id).total_memory
)
free_gpu_memory = max(0.0, total_gpu_memory - used_memory)
used_memory = torch.xpu.memory_allocated(device_id)
total_gpu_memory = torch.xpu.get_device_properties(device_id).total_memory
free_gpu_memory = total_gpu_memory - used_memory
if distributed:
import torch.distributed as dist
+3 -17
View File
@@ -434,23 +434,9 @@ def get_available_gpu_memory(
if empty_cache:
empty_device_cache(torch.xpu)
# Use mem_get_info() with a sanity cap to avoid KV-cache over-allocation
# on drivers that incorrectly return total memory as free memory.
# Consistent with the fallback: free = max(0, total - allocated).
try:
free_gpu_memory, total_gpu_memory = torch.xpu.mem_get_info(gpu_id)
used_memory = float(torch.xpu.memory_allocated(gpu_id))
free_gpu_memory = min(
float(free_gpu_memory),
max(0.0, float(total_gpu_memory) - used_memory),
)
except Exception:
# Fallback for devices/drivers that do not support querying free memory
used_memory = float(torch.xpu.memory_allocated(gpu_id))
total_gpu_memory = float(
torch.xpu.get_device_properties(gpu_id).total_memory
)
free_gpu_memory = max(0.0, total_gpu_memory - used_memory)
used_memory = torch.xpu.memory_allocated(gpu_id)
total_gpu_memory = torch.xpu.get_device_properties(gpu_id).total_memory
free_gpu_memory = total_gpu_memory - used_memory
elif device == "hpu":
num_gpus = torch.hpu.device_count()