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

feat: Add Triton fallback option and SM120 MoE configs for FP8 models (#9251)

Browse Source
This commit is contained in:
Martin Vit
2025-08-21 04:33:15 +02:00
committed by GitHub
parent eb19ccadae
commit 7cd2ee06d7
3 changed files with 312 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
python/sglang/srt/layers/quantization/fp8_utils.py
View File

@@ -53,6 +53,7 @@ if _is_cuda:
from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm
use_vllm_cutlass_w8a8_fp8_kernel = get_bool_env_var("USE_VLLM_CUTLASS_W8A8_FP8_KERNEL")
use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL")
# Input scaling factors are no longer optional in _scaled_mm starting
# from pytorch 2.5. Allocating a dummy tensor to pass as input_scale
@@ -592,7 +593,7 @@ def apply_fp8_linear(
cutlass_compatible_b = (
weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
)
if not cutlass_compatible_b:
if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel:
# Massage the input to be 2D
qinput = qinput.view(-1, qinput.shape[-1])
output = triton_scaled_mm(
@@ -735,14 +736,25 @@ def apply_fp8_linear(
assert (
weight_scale.numel() == weight.shape[1]
), "cutlass w8a8 fp8 sgl-kernel only supports per-channel scale"
output = fp8_scaled_mm(
qinput,
weight,
x_scale,
weight_scale,
out_dtype=input.dtype,
bias=bias,
cutlass_compatible_b = (
weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
)
if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel:
# Massage the input to be 2D
qinput = qinput.view(-1, qinput.shape[-1])
output = triton_scaled_mm(
qinput, weight, x_scale, weight_scale, input.dtype, bias
)
else:
output = fp8_scaled_mm(
qinput,
weight,
x_scale,
weight_scale,
out_dtype=input.dtype,
bias=bias,
)
return output.view(*output_shape)
except (ImportError, NameError, AttributeError):
pass