Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Fridge003 merged 8 commits into from
Mar 29, 2026
+7 −3
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
1a2fe3a
fix cuda graph error
Mar 3, 2026
eb83055
cleanup
Mar 3, 2026
73959dc
cleanyo
Mar 3, 2026
743a4a1
Merge branch 'main' into fix_value
b8zhong Mar 9, 2026
55eb9e4
Merge branch 'main' into fix_value
wolfcomos Mar 9, 2026
7c28962
Merge branch 'main' into fix_value
wolfcomos Mar 11, 2026
52ea8cf
Merge branch 'main' into fix_value
b8zhong Mar 11, 2026
cf01741
Merge branch 'main' into fix_value
Fridge003 Mar 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion python/sglang/srt/layers/quantization/fp8_kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
6 changes: 4 additions & 2 deletions python/sglang/srt/layers/quantization/fp8_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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,
Expand Down
Loading