huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-22 15:32:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

update get xpu memory api

Browse Source 
Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
This commit is contained in:
Wang, Yi A 2025-03-12 00:37:17 -07:00
parent cdc5380f2b
commit 468da545af
1 changed files with 2 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
server/text_generation_server/utils/import_utils.py
View File

@ -18,15 +18,9 @@ def get_cuda_free_memory(device, memory_fraction):
def get_xpu_free_memory(device, memory_fraction):
total_memory = torch.xpu.get_device_properties(device).total_memory
device_id = device.index
total_free_memory, total_xpu_memory = torch.xpu.mem_get_info(device)
memory_fraction = float(os.getenv("XPU_MEMORY_FRACTION", "1.0"))
free_memory = max(
0,
int(
total_memory * 0.9 * memory_fraction - torch.xpu.memory_reserved(device_id)
),
)
free_memory = max(0, total_free_memory - (1 - memory_fraction) * total_xpu_memory)
return free_memory