minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix cuda graph capturing error in sm120 mxfp8 triton path (#19835)

Browse Source
This commit is contained in:
Hanlin Bi
2026-03-29 01:59:24 -07:00
committed by GitHub
parent fec9961a1f
commit aa9177152e
2 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/fp8_kernel.py
+3 -1
View File
@@ -47,6 +47,8 @@ from sglang.srt.utils.patch_torch import register_fake_if_exists
_is_hip = is_hip() _is_hip = is_hip()
_is_cuda = is_cuda() _is_cuda = is_cuda()
_is_cpu = is_cpu() _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 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
if _is_cuda: if _is_cuda:
@@ -1299,7 +1301,7 @@ def mxfp8_block_scaled_matmul_triton(
SM120: 1, SM100: 4. SM120: 1, SM100: 4.
""" """
if num_stages is None: 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 M, K = a.shape
N, K_b = b.shape N, K_b = b.shape
assert K == K_b assert K == K_b
python/sglang/srt/layers/quantization/fp8_utils.py
+4 -2
View File
@@ -50,6 +50,8 @@ logger = logging.getLogger(__name__)
_is_hip = is_hip() _is_hip = is_hip()
_is_cuda = is_cuda() _is_cuda = is_cuda()
_is_fp8_fnuz = is_fp8_fnuz() _is_fp8_fnuz = is_fp8_fnuz()
_is_sm100_supported = is_sm100_supported()
_is_sm120_supported = is_sm120_supported()
_is_gfx95_supported = is_gfx95_supported() _is_gfx95_supported = is_gfx95_supported()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -863,7 +865,7 @@ def triton_mxfp8_blockscaled_linear(
bias: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None,
output_dtype: Optional[torch.dtype] = None, output_dtype: Optional[torch.dtype] = None,
) -> torch.Tensor: ) -> 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).") raise RuntimeError("MXFP8 dense linear requires Blackwell GPUs (SM100/SM120).")
input_2d = input.view(-1, input.shape[-1]).contiguous() input_2d = input.view(-1, input.shape[-1]).contiguous()
@@ -915,7 +917,7 @@ def triton_mxfp8_blockscaled_linear(
a_scale_packed = _pack_mxfp8_scales(x_scale_u8) a_scale_packed = _pack_mxfp8_scales(x_scale_u8)
b_scale_packed = _pack_mxfp8_scales(weight_scale) 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( output = mxfp8_block_scaled_matmul_triton(
q_input, q_input,
a_scale_packed, a_scale_packed,