diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py index c60a0e1b1..24e716cf9 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py @@ -17,6 +17,7 @@ import torch from sglang.srt.layers.attention.linear.kernels.kernel_backend import ( LinearAttnKernelBase, ) +from sglang.srt.utils import is_cuda logger = logging.getLogger(__name__) @@ -54,7 +55,7 @@ def _get_flashinfer_gdn_kernels(): _flashinfer_gated_delta_rule_mtp_bf16 = gated_delta_rule_mtp_bf16 _flashinfer_gated_delta_rule_decode = gated_delta_rule_decode_pretranspose _flashinfer_gdn_available = ( - torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9 + is_cuda() and torch.cuda.get_device_capability()[0] >= 9 ) if _flashinfer_gdn_available: logger.info("FlashInfer GDN kernels loaded successfully") diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index e29f7efec..f7eda3c36 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -14,7 +14,7 @@ from sglang.srt.layers.dp_attention import ( get_attention_dp_size, is_dp_attention_enabled, ) -from sglang.srt.utils import is_npu +from sglang.srt.utils import is_cuda, is_npu _is_npu = is_npu() @@ -288,7 +288,7 @@ def initialize_moe_config(server_args: ServerArgs): DEEPEP_CONFIG = server_args.deepep_config or "" IS_TBO_ENABLED = server_args.enable_two_batch_overlap IS_SBO_ENABLED = server_args.enable_single_batch_overlap - if IS_SBO_ENABLED and torch.cuda.is_available(): + if IS_SBO_ENABLED and is_cuda(): if torch.cuda.get_device_capability()[0] == 9: raise ValueError( "SBO (single batch overlap) is not supported on SM90 GPUs with latest sgl-deep-gemm wheel. Please try removing --enable-single-batch-overlap argument." diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f2f30f4c1..cfe1237f2 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3399,7 +3399,7 @@ class ServerArgs: if ( decode == "flashinfer" and self.mamba_ssm_dtype != "bfloat16" - and torch.cuda.is_available() + and is_cuda() and torch.cuda.get_device_capability()[0] >= 10 ): raise ValueError( @@ -3415,7 +3415,7 @@ class ServerArgs: cuda_major = int(cuda_version.split(".")[0]) if cuda_version is not None else 0 if ( prefill == "flashinfer" - and torch.cuda.is_available() + and is_cuda() and torch.cuda.get_device_capability()[0] >= 10 and cuda_major < 13 ):