From 05ab110e02baa5e876de6ff0108526efd8984dc3 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad Date: Thu, 8 Jan 2026 23:52:22 +0800 Subject: [PATCH] [Performance] Force split_k=1 for MXFP4 Triton kernels on Hopper (#16014) Signed-off-by: Mohammad Miadh Angkad --- python/sglang/srt/layers/quantization/mxfp4.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 274d7fe35..64a50e433 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -42,6 +42,7 @@ from sglang.srt.utils import ( is_flashinfer_available, is_gfx95_supported, is_hip, + is_sm90_supported, is_sm100_supported, is_triton_kernels_available, log_info_on_rank0, @@ -53,6 +54,7 @@ from sglang.srt.utils import ( 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() @@ -108,6 +110,11 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps): "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)