diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 24a1691d225d..016c836fe303 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -47,6 +47,8 @@ _is_hip = is_hip() _is_cuda = is_cuda() _is_cpu = is_cpu() +_is_sm100_supported = is_sm100_supported() +_is_sm120_supported = is_sm120_supported() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_cuda: @@ -1299,7 +1301,7 @@ def mxfp8_block_scaled_matmul_triton( SM120: 1, SM100: 4. """ if num_stages is None: - num_stages = 1 if is_sm120_supported() else (4 if is_sm100_supported() else 1) + num_stages = 1 if _is_sm120_supported else (4 if _is_sm100_supported else 1) M, K = a.shape N, K_b = b.shape assert K == K_b diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 3f26b736a6a9..c0677afb0574 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -51,6 +51,8 @@ _is_hip = is_hip() _is_cuda = is_cuda() _is_fp8_fnuz = is_fp8_fnuz() +_is_sm100_supported = is_sm100_supported() +_is_sm120_supported = is_sm120_supported() _is_gfx95_supported = is_gfx95_supported() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip @@ -883,7 +885,7 @@ def triton_mxfp8_blockscaled_linear( bias: Optional[torch.Tensor] = None, output_dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: - if not (_is_cuda and (is_sm100_supported() or is_sm120_supported())): + if not (_is_cuda and (_is_sm100_supported or _is_sm120_supported)): raise RuntimeError("MXFP8 dense linear requires Blackwell GPUs (SM100/SM120).") input_2d = input.view(-1, input.shape[-1]).contiguous() @@ -935,7 +937,7 @@ def triton_mxfp8_blockscaled_linear( a_scale_packed = _pack_mxfp8_scales(x_scale_u8) b_scale_packed = _pack_mxfp8_scales(weight_scale) - num_stages = 1 if is_sm120_supported() else (4 if is_sm100_supported() else 1) + num_stages = 1 if _is_sm120_supported else (4 if _is_sm100_supported else 1) output = mxfp8_block_scaled_matmul_triton( q_input, a_scale_packed,