From c6872fc8fb9aca25ac6c035ea876ee34f021f923 Mon Sep 17 00:00:00 2001 From: Jinghong Li <53993397+AethoceSora@users.noreply.github.com> Date: Fri, 24 Apr 2026 00:26:04 +0800 Subject: [PATCH] Fix: fallback to torch API when NVML memory query is not supported (#23426) Co-authored-by: ronnie_zheng --- python/sglang/multimodal_gen/runtime/platforms/cuda.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index a2f6f3c95..32b8212a8 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -514,7 +514,10 @@ class NvmlCudaPlatform(CudaPlatformBase): def get_device_total_memory(cls, device_id: int = 0) -> int: physical_device_id = device_id_to_physical_device_id(device_id) handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) - return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) + try: + return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) + except pynvml.NVMLError_NotSupported: + return int(torch.cuda.get_device_properties(device_id).total_memory) @classmethod @with_nvml_context