diff --git a/python/sglang/check_env.py b/python/sglang/check_env.py index 516e90237..a9a8a62b2 100644 --- a/python/sglang/check_env.py +++ b/python/sglang/check_env.py @@ -194,22 +194,12 @@ class GPUEnv(BaseEnv): """ Get CUDA driver version. """ - versions = set() - try: - output = subprocess.check_output( - [ - "nvidia-smi", - "--query-gpu=driver_version", - "--format=csv,noheader,nounits", - ] - ) - versions = set(output.decode().strip().split("\n")) - if len(versions) == 1: - return {"CUDA Driver Version": versions.pop()} - else: - return {"CUDA Driver Versions": ", ".join(sorted(versions))} - except subprocess.SubprocessError: + from sglang.srt.utils.common import get_nvidia_driver_version_str + + ver = get_nvidia_driver_version_str() + if ver is None: return {"CUDA Driver Version": "Not Available"} + return {"CUDA Driver Version": ver} def get_topology(self): """ diff --git a/python/sglang/cli/killall.py b/python/sglang/cli/killall.py index 39a10f053..1c886e1eb 100755 --- a/python/sglang/cli/killall.py +++ b/python/sglang/cli/killall.py @@ -77,19 +77,11 @@ def _run_smi(query, query_type="gpu"): def _get_smi_version(): - """Return nvidia-smi driver version and CUDA version, or None on failure.""" - try: - out = subprocess.check_output( - [ - "nvidia-smi", - "--query-gpu=driver_version", - "--format=csv,noheader,nounits", - ], - text=True, - timeout=10, - ) - driver = out.strip().splitlines()[0].strip() if out.strip() else "unknown" - except (subprocess.SubprocessError, FileNotFoundError, IndexError): + """Return nvidia-smi driver version and GPU name, or None on failure.""" + from sglang.srt.utils.common import get_nvidia_driver_version_str + + driver = get_nvidia_driver_version_str() + if driver is None: return None try: out = subprocess.check_output( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ce4cc7df4..6b4759f5c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -42,6 +42,7 @@ from sglang.srt.utils.common import ( get_device_name, get_device_sm, get_int_env_var, + get_nvidia_driver_version, get_quantization_config, human_readable_int, is_blackwell_supported, @@ -1757,10 +1758,21 @@ class ServerArgs: and is_triton_kernels_available() and self.quantization is None ): - self.moe_runner_backend = "triton_kernel" - logger.warning( - "Detected GPT-OSS model, enabling triton_kernels MOE kernel." - ) + # The triton_kernels package segfaults on Blackwell (B200) + # with NVIDIA driver >= 595. Fall back to triton backend. + if is_blackwell_supported() and get_nvidia_driver_version() >= ( + 595, + ): + self.moe_runner_backend = "triton" + logger.warning( + "Detected GPT-OSS model on Blackwell with driver >= 595, " + "using triton MOE kernel to avoid triton_kernels SIGSEGV." + ) + else: + self.moe_runner_backend = "triton_kernel" + logger.warning( + "Detected GPT-OSS model, enabling triton_kernels MOE kernel." + ) if self.moe_runner_backend == "triton_kernel": assert ( diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index d7ea34c76..74bb6962e 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3399,6 +3399,41 @@ def is_triton_kernels_available() -> bool: return importlib.util.find_spec("triton_kernels") is not None +@lru_cache(maxsize=1) +def get_nvidia_driver_version() -> tuple: + """Return the NVIDIA driver version as a tuple of ints, e.g. (595, 58, 3). + Returns (0,) on failure.""" + version_str = get_nvidia_driver_version_str() + if version_str is None: + return (0,) + try: + return tuple(int(x) for x in version_str.split(".")) + except ValueError: + return (0,) + + +@lru_cache(maxsize=1) +def get_nvidia_driver_version_str() -> str: + """Return the NVIDIA driver version string, e.g. '595.58.03'. + Returns None on failure.""" + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + check=True, + timeout=10, + ) + version_str = result.stdout.strip().split("\n")[0].strip() + return version_str if version_str else None + except (subprocess.CalledProcessError, FileNotFoundError, ValueError): + return None + + def check_cuda_result(raw_output): import cuda.bindings.runtime as cuda_rt