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

[quantization] Add tuned Triton tile configs for channelwise FP8 GEMM… (#34331)

Browse Source 
Co-authored-by: runkai <runkai@amazon.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
Runkai Tao
2026-08-14 11:11:02 +08:00
committed by GitHub
co-authored by runkai Xiaoyu Zhang
parent 85cdf1178d
commit e1c4db9621
7 changed files with 574 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/quantization/configs/N=24576,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json
+103
View File
@@ -0,0 +1,103 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 16,
"BLOCK_SIZE_K": 1024,
"num_warps": 2,
"num_stages": 2
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 16,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 16,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 16,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"32": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 512,
"num_warps": 4,
"num_stages": 3
},
"48": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 5
},
"64": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 5
},
"96": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 8,
"num_stages": 3
},
"128": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 3
},
"192": {
"BLOCK_SIZE_M": 256,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 8,
"num_stages": 1
},
"256": {
"BLOCK_SIZE_M": 256,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 8,
"num_stages": 1
},
"512": null,
"1024": null,
"2048": null,
"4096": {
"BLOCK_SIZE_M": 256,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 2
},
"8192": {
"BLOCK_SIZE_M": 256,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 2
}
}
python/sglang/kernels/ops/quantization/configs/N=4096,K=12288,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json
+121
View File
@@ -0,0 +1,121 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"32": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 16,
"BLOCK_SIZE_K": 512,
"num_warps": 4,
"num_stages": 5
},
"48": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 5
},
"64": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 8,
"num_stages": 5
},
"96": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 8,
"num_stages": 3
},
"128": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 8,
"num_stages": 4
},
"192": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 3
},
"256": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"num_warps": 8,
"num_stages": 3
},
"512": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 4
},
"1024": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 3
},
"2048": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 3
},
"4096": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 3
},
"8192": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 2
}
}
python/sglang/kernels/ops/quantization/configs/N=4096,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json
+121
View File
@@ -0,0 +1,121 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 1
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 1
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 16,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 4
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 1
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 16,
"BLOCK_SIZE_K": 256,
"num_warps": 2,
"num_stages": 4
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 512,
"num_warps": 4,
"num_stages": 3
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 3
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 3
},
"96": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 3
},
"128": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 3
},
"192": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 3
},
"256": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 3
},
"512": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 4
},
"1024": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"num_warps": 4,
"num_stages": 4
},
"2048": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 3
},
"4096": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"num_warps": 4,
"num_stages": 2
},
"8192": {
"BLOCK_SIZE_M": 256,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 2
}
}
python/sglang/kernels/ops/quantization/configs/N=6144,K=4096,device_name=NVIDIA_L40S,dtype=fp8_w8a8_channelwise.json
+121
View File
@@ -0,0 +1,121 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 2
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 512,
"num_warps": 4,
"num_stages": 2
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 1024,
"num_warps": 4,
"num_stages": 3
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 16,
"BLOCK_SIZE_K": 512,
"num_warps": 4,
"num_stages": 4
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 2,
"num_stages": 3
},
"48": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 8,
"num_stages": 3
},
"64": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 3
},
"96": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 256,
"num_warps": 2,
"num_stages": 2
},
"128": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 2
},
"192": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 2
},
"256": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"num_warps": 4,
"num_stages": 2
},
"512": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"num_warps": 4,
"num_stages": 2
},
"1024": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 4
},
"2048": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"num_warps": 4,
"num_stages": 2
},
"4096": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 128,
"num_warps": 8,
"num_stages": 3
},
"8192": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"num_warps": 4,
"num_stages": 2
}
}
python/sglang/kernels/ops/quantization/fp8_kernel.py
+65 -2
View File
@@ -1263,6 +1263,57 @@ def get_w8a8_block_fp8_configs(
return None
@functools.lru_cache
def get_w8a8_channelwise_fp8_configs(N: int, K: int) -> Optional[Dict[int, Any]]:
"""
Return tuned Triton configurations for the per-token/per-channel w8a8 fp8
GEMM (`triton_scaled_mm`), for one weight shape on the current device.
The return value maps an irregular grid of token counts M to
`triton_scaled_mm` tile configs; the closest M in the grid should be picked
for a given batch. A null entry means CUTLASS was faster at that M, so the
grid keeps the point (to stop a neighbouring M snapping onto it) but the
caller must fall back. A missing file / shape likewise means "keep the
default".
"""
# Intentional: the tuned lookup (host-side device name + file I/O) isn't
# traceable, so under torch.compile return None and fall back to CUTLASS --
# same as mainline's block-FP8 get_w8a8_block_fp8_configs.
if torch._dynamo.is_compiling():
return None
device_name = get_device_name().replace(" ", "_")
json_file_name = (
f"N={N},K={K},device_name={device_name},dtype=fp8_w8a8_channelwise.json"
)
config_file_path = os.path.join(
os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name
)
if not os.path.exists(config_file_path):
return None
with open(config_file_path) as f:
log_info_on_rank0(
logger,
f"Using configuration from {config_file_path} for W8A8 channelwise FP8 GEMM.",
)
return {int(key): val for key, val in json.load(f).items()}
def get_w8a8_channelwise_fp8_config(N: int, K: int, M: int) -> Optional[Dict[str, int]]:
"""Tuned `triton_scaled_mm` config for this GEMM, or None to keep the default.
N / K are the weight's output / contraction dims and M the token count. None
means no tuned config applies -- the caller keeps its existing kernel choice
-- either because this device / shape was never tuned, or the nearest tuned
M is a point where CUTLASS won. The caller gates on the feature flag.
"""
configs = get_w8a8_channelwise_fp8_configs(N, K)
if not configs:
return None
return configs[min(configs.keys(), key=lambda x: abs(x - M))]
def select_w8a8_block_fp8_matmul_kernel(M, N, META):
return _w8a8_block_fp8_matmul
@@ -2124,6 +2175,8 @@ def triton_scaled_mm(
block_size_n: int = 32,
block_size_k: int = 32,
use_heuristic=True,
num_warps: Optional[int] = None,
num_stages: Optional[int] = None,
) -> torch.Tensor:
M, K = input.shape
N = weight.shape[1]
@@ -2163,14 +2216,23 @@ def triton_scaled_mm(
tile_shape = (64, 128, 128)
else:
tile_shape = (128, 128, 128)
block_size_m, block_size_n, block_size_k = tile_shape
block_size_m, block_size_n, block_size_k = tile_shape
# else: use the block_size_{m,n,k} the caller passed (e.g. a tuned config).
block_size_sa = 1 if has_scalar(scale_a) else block_size_m
block_size_sb = 1 if has_scalar(scale_b) else block_size_n
accumulator_dtype = tl.float32 if input.is_floating_point() else tl.int32
# num_warps / num_stages are triton.jit launch kwargs, and triton rejects an
# explicit None -- pass them only when the caller set them, so the default
# path keeps triton's own defaults.
launch_kwargs = {}
if num_warps is not None:
launch_kwargs["num_warps"] = num_warps
if num_stages is not None:
launch_kwargs["num_stages"] = num_stages
# A = input, B = weight, C = result
# A = M x K, B = K x N, C = M x N
scaled_mm_kernel[grid](
@@ -2195,6 +2257,7 @@ def triton_scaled_mm(
BLOCK_SIZE_K=block_size_k,
BLOCK_SIZE_SCALE_A=block_size_sa,
BLOCK_SIZE_SCALE_B=block_size_sb,
**launch_kwargs,
)
return result.to(out_dtype)
python/sglang/srt/environ.py
+9
View File
@@ -847,6 +847,15 @@ class Envs:
SGLANG_QUANT_ALLOW_DOWNCASTING = EnvBool(False)
SGLANG_FP8_IGNORED_LAYERS = EnvStr("")
SGLANG_FP4_IGNORED_LAYERS = EnvStr("")
# On by default; set SGLANG_ENABLE_FP8_GEMM_CONFIG_TUNE=0 as a kill switch.
# Consults the tuned per-(N, K, M) Triton tile config table in
# apply_fp8_linear. When a tuned config exists for this GPU / weight shape /
# token count, run the Triton w8a8 FP8 GEMM with it; otherwise keep the
# default CUTLASS path. Only takes effect on a GPU with a matching
# dtype=fp8_w8a8_channelwise config JSON under
# kernels/ops/quantization/configs/ (currently L40S), so it is a no-op on
# any other GPU / untuned shape even when enabled.
SGLANG_ENABLE_FP8_GEMM_CONFIG_TUNE = EnvBool(True)
# ===================================================================
# Humming quantization
python/sglang/srt/layers/quantization/fp8_utils.py
+34
View File
@@ -11,6 +11,7 @@ from sglang.kernels.ops.quantization.fp8_kernel import (
sglang_per_token_group_quant_fp8,
sglang_per_token_group_quant_fp8_row_padded,
)
from sglang.srt.environ import envs
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil
from sglang.srt.runtime_context import get_exec, get_parallel
@@ -23,6 +24,7 @@ from sglang.kernels.ops.quantization.fp8_kernel import (
fp8_dtype,
fp8_max,
fp8_min,
get_w8a8_channelwise_fp8_config,
is_fp8_fnuz,
per_token_group_quant_fp8,
scaled_fp8_quant,
@@ -1788,6 +1790,12 @@ def apply_fp8_linear(
use_cutlass_channelwise_gemm = (
channelwise_cutlass and cutlass_compatible_b and not use_triton_w8a8_fp8_kernel
)
# Consider a tuned Triton tile only where the shape would otherwise go to
# CUTLASS (that is the path the offline sweep tuned against). On by default;
# SGLANG_ENABLE_FP8_GEMM_CONFIG_TUNE=0 is the kill switch.
use_tuned_triton_channelwise = (
use_cutlass_channelwise_gemm and envs.SGLANG_ENABLE_FP8_GEMM_CONFIG_TUNE.get()
)
native_scalar_a_scale = use_cutlass_channelwise_gemm and (
_is_sm90_supported or _is_sm100_supported or _is_sm120_supported
)
@@ -1866,12 +1874,38 @@ def apply_fp8_linear(
)
if channelwise_cutlass:
# A tuned config exists only for shapes where tuned Triton beat the
# CUTLASS dispatch on this GPU; otherwise this is None and the backend
# choice below is unchanged. weight is [K, N] here.
tuned_config = (
get_w8a8_channelwise_fp8_config(
N=weight.shape[1], K=weight.shape[0], M=qinput.shape[0]
)
if use_tuned_triton_channelwise
else None
)
if not use_cutlass_channelwise_gemm:
# Massage the input to be 2D
qinput = qinput.view(-1, qinput.shape[-1])
output = triton_scaled_mm(
qinput, weight, x_scale, weight_scale, output_dtype, bias
)
elif tuned_config is not None:
qinput = qinput.view(-1, qinput.shape[-1])
output = triton_scaled_mm(
qinput,
weight,
x_scale,
weight_scale,
output_dtype,
bias,
block_size_m=tuned_config["BLOCK_SIZE_M"],
block_size_n=tuned_config["BLOCK_SIZE_N"],
block_size_k=tuned_config["BLOCK_SIZE_K"],
use_heuristic=False,
num_warps=tuned_config["num_warps"],
num_stages=tuned_config["num_stages"],
)
else:
output = fp8_scaled_mm(
qinput,