diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index cb3204e3d..5afbcd66a 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1725,6 +1725,39 @@ def get_device_sm(): return 0 +def _cuda_mem_fallback(reason: str) -> int: + """Fallback to torch.cuda.mem_get_info() and return total GPU memory in MiB. + + Queries all visible CUDA devices and returns the minimum total memory, + consistent with the nvidia-smi path that takes min(memory_values). + + Returns the total memory in MiB, or raises RuntimeError if CUDA is + unavailable or mem_get_info() fails. + """ + if not torch.cuda.is_available(): + raise RuntimeError(reason) + try: + device_count = torch.cuda.device_count() + if device_count == 0: + # Include the original failure reason for diagnostics + raise RuntimeError(f"{reason} No CUDA devices found via torch.cuda.") + memory_values = [] + for i in range(device_count): + total = torch.cuda.mem_get_info(i)[1] // 1024 // 1024 # unit: MiB + memory_values.append(total) + result = min(memory_values) + logger.warning( + f"{reason} Falling back to torch.cuda.mem_get_info(). " + f"Reported total GPU memory per device (MiB): {memory_values}, " + f"using min: {result} MiB." + ) + return result + except (RuntimeError, ValueError, OSError) as e: + raise RuntimeError( + f"{reason} torch.cuda.mem_get_info() fallback also failed: {e}" + ) from e + + def get_nvgpu_memory_capacity(): try: # Run nvidia-smi and capture the output @@ -1736,7 +1769,9 @@ def get_nvgpu_memory_capacity(): ) if result.returncode != 0: - raise RuntimeError(f"nvidia-smi error: {result.stderr.strip()}") + return _cuda_mem_fallback( + f"nvidia-smi failed (exit code {result.returncode}: {result.stderr.strip()})." + ) # Parse the output to extract memory values memory_values = [ @@ -1746,20 +1781,17 @@ def get_nvgpu_memory_capacity(): ] if not memory_values: - # Fallback to torch.cuda.mem_get_info() when failed to get memory capacity from nvidia-smi, + # Fallback when nvidia-smi returns no parseable values, # typically in NVIDIA MIG mode. - if torch.cuda.is_available(): - logger.warning( - "Failed to get GPU memory capacity from nvidia-smi, falling back to torch.cuda.mem_get_info()." - ) - return torch.cuda.mem_get_info()[1] // 1024 // 1024 # unit: MB - raise ValueError("No GPU memory values found.") + return _cuda_mem_fallback( + "Failed to get GPU memory capacity from nvidia-smi." + ) # Return the minimum memory value return min(memory_values) except FileNotFoundError: - raise RuntimeError( + return _cuda_mem_fallback( "nvidia-smi not found. Ensure NVIDIA drivers are installed and accessible." )