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

[Performance] Force split_k=1 for MXFP4 Triton kernels on Hopper (#16014)

Browse Source 
Signed-off-by: Mohammad Miadh Angkad <MAngkad.BSDSBA2027@aim.edu>
This commit is contained in:
Mohammad Miadh Angkad
2026-01-08 23:52:22 +08:00
committed by GitHub
parent c8dc4d2d1f
commit 05ab110e02
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/mxfp4.py
+7
View File
@@ -42,6 +42,7 @@ from sglang.srt.utils import (
is_flashinfer_available, is_flashinfer_available,
is_gfx95_supported, is_gfx95_supported,
is_hip, is_hip,
is_sm90_supported,
is_sm100_supported, is_sm100_supported,
is_triton_kernels_available, is_triton_kernels_available,
log_info_on_rank0, log_info_on_rank0,
@@ -53,6 +54,7 @@ from sglang.srt.utils import (
from sglang.srt.utils.custom_op import register_custom_op from sglang.srt.utils.custom_op import register_custom_op
_is_sm100_supported = is_cuda() and is_sm100_supported() _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() has_triton_kernels = is_triton_kernels_available()
@@ -108,6 +110,11 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
"epilogue_subtile": 1, "epilogue_subtile": 1,
} }
opt_flags.update_opt_flags_constraints(constraints) 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 # transpose the tensor so that the quantization axis is on dim1
quant_tensor = quant_tensor.transpose(-2, -1) quant_tensor = quant_tensor.transpose(-2, -1)
scale = scale.transpose(-2, -1) scale = scale.transpose(-2, -1)