[diffusion] fix: fix UMA detection (#17113)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR
2026-01-15 22:29:46 +08:00
committed by GitHub
parent e7df8bdc5c
commit c9a45b7e3c

View File

@@ -124,13 +124,8 @@ class CudaPlatformBase(Platform):
if empty_cache:
torch.cuda.empty_cache()
# Orin, Thor, Spark
# SM 8.7 is Orin, 11.0 is Thor, 12.1 is Spark
SHARED_SYSMEM_DEVICE_MEM_SMS = (87, 110, 121)
capability = cls.get_device_capability(device_id)
sm = capability.to_int() if capability else 0
if sm in SHARED_SYSMEM_DEVICE_MEM_SMS:
device_props = torch.cuda.get_device_properties(device_id)
if device_props.is_integrated:
free_gpu_memory = psutil.virtual_memory().available
else:
free_gpu_memory, _ = torch.cuda.mem_get_info(device_id)