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

Support triton_kernels for GPT-OSS on SM120 (#19718)

Browse Source 
Co-authored-by: amittell 1388680+amittell@users.noreply.github.com
This commit is contained in:
Brayden Zhong
2026-03-03 17:14:01 -05:00
committed by GitHub
parent 1135e214b3
commit 9305f0e58d
2 changed files with 41 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
python/sglang/srt/layers/quantization/mxfp4.py
View File

@@ -37,12 +37,12 @@ from sglang.srt.layers.quantization.base_config import (
from sglang.srt.layers.quantization.utils import is_layer_skipped
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
is_cuda,
is_flashinfer_available,
is_gfx95_supported,
is_hip,
is_sm90_supported,
is_sm100_supported,
is_sm120_supported,
is_triton_kernels_available,
mxfp_supported,
next_power_of_2,
@@ -52,8 +52,6 @@ from sglang.srt.utils import (
from sglang.srt.utils.common import get_bool_env_var
from sglang.srt.utils.custom_op import register_custom_op
_is_sm100_supported = is_cuda() and is_sm100_supported()
_is_sm90_supported = is_cuda() and is_sm90_supported()
has_triton_kernels = is_triton_kernels_available()
@@ -140,23 +138,40 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
from triton_kernels.tensor_details import layout
value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout(
mx_axis=1
)
scale_layout, scale_layout_opts = layout.make_default_matmul_mxfp4_w_scale_layout(
mx_axis=1, num_warps=num_warps
)
if _is_sm100_supported:
if is_sm120_supported():
# SM120 (Blackwell desktop) doesn't support persistent kernels / TMA block layout
# Use StridedLayout and disable persistent kernels to avoid assertion errors
from triton_kernels.tensor_details.layout import StridedLayout
value_layout = StridedLayout
value_layout_opts = {}
scale_layout = StridedLayout
scale_layout_opts = {}
constraints = {
"is_persistent": True,
"epilogue_subtile": 1,
}
opt_flags.update_opt_flags_constraints(constraints)
elif _is_sm90_supported:
constraints = {
"split_k": 1,
"is_persistent": False,
"num_stages": 1,
}
opt_flags.update_opt_flags_constraints(constraints)
else:
value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout(
mx_axis=1
)
scale_layout, scale_layout_opts = (
layout.make_default_matmul_mxfp4_w_scale_layout(
mx_axis=1, num_warps=num_warps
)
)
if is_sm100_supported():
constraints = {
"is_persistent": True,
"epilogue_subtile": 1,
}
opt_flags.update_opt_flags_constraints(constraints)
elif is_sm90_supported():
constraints = {
"split_k": 1,
}
opt_flags.update_opt_flags_constraints(constraints)
# transpose the tensor so that the quantization axis is on dim1
quant_tensor = quant_tensor.transpose(-2, -1)
scale = scale.transpose(-2, -1)
@@ -324,7 +339,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
# pad the intermediate size to be a multiple of 2 * mxfp4_block
# for to hold non-uniform sharded tensor as well as swizzling
intermediate_size_per_partition_after_pad = intermediate_size_per_partition
if _is_sm100_supported:
if is_sm100_supported():
if self.use_flashinfer:
intermediate_size_per_partition_after_pad = round_up(
intermediate_size_per_partition, 256

10
python/sglang/srt/server_args.py
View File

@@ -1486,10 +1486,16 @@ class ServerArgs:
self.dtype = "bfloat16"
if self.moe_runner_backend == "auto":
if is_blackwell_supported() and is_mxfp4_quant_format:
if is_sm100_supported() and is_mxfp4_quant_format:
self.moe_runner_backend = "flashinfer_mxfp4"
logger.warning(
"Detected Blackwell and MXFP4 quantization format for GPT-OSS model, enabling FlashInfer MXFP4 MOE kernel."
"Detected SM100 and MXFP4 quantization format for GPT-OSS model, enabling FlashInfer MXFP4 MOE kernel."
)
elif is_sm120_supported() and is_mxfp4_quant_format:
# trtllm-gen only supports SM100
self.moe_runner_backend = "triton_kernel"
logger.warning(
"Detected SM120 and MXFP4 quantization format for GPT-OSS model, enabling triton_kernel MOE kernel."
)
elif (
is_hip() and get_bool_env_var("SGLANG_USE_AITER")