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

Fix SM120 triton_kernels MXFP4 block_k for GPT-OSS (#20040)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-03-07 02:53:08 +08:00
committed by GitHub
parent de1a0afcbc
commit 759700c808
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -139,8 +139,9 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
from triton_kernels.tensor_details import layout
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
# SM120 desktop Blackwell does not support the persistent/TMA MXFP4 path.
# This MXFP4 path uses StridedLayout and the non-persistent kernel with
# block_k=128 so the selected tile stays within the per-block shared-memory budget.
from triton_kernels.tensor_details.layout import StridedLayout
value_layout = StridedLayout
@@ -149,6 +150,7 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
scale_layout_opts = {}
constraints = {
"is_persistent": False,
"block_k": 128,
"num_stages": 1,
}
opt_flags.update_opt_flags_constraints(constraints)