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

[jit kernel] Support per_token_group_quant_8bit jit kernel (#18905)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-02-20 21:01:05 +08:00
committed by GitHub
parent 38a69652e6
commit 7d953440ec
6 changed files with 811 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
python/sglang/srt/layers/quantization/fp8_kernel.py
View File

@@ -63,6 +63,10 @@ if _is_cuda:
enable_sgl_per_token_group_quant_8bit = False
from sglang.jit_kernel.per_token_group_quant_8bit import (
per_token_group_quant_8bit as sgl_per_token_group_quant_8bit_jit,
)
if _is_hip:
_has_vllm = False
if _use_aiter:
@@ -501,19 +505,31 @@ def sglang_per_token_group_quant_fp8(
if x.shape[0] > 0:
# Temporary
if enable_sgl_per_token_group_quant_8bit:
sgl_per_token_group_quant_8bit(
x,
x_q,
x_s,
group_size,
eps,
fp8_min,
fp8_max,
scale_ue8m0,
fuse_silu_and_mul,
masked_m,
enable_v2=enable_v2,
)
if enable_v2:
sgl_per_token_group_quant_8bit(
x,
x_q,
x_s,
group_size,
eps,
fp8_min,
fp8_max,
scale_ue8m0,
fuse_silu_and_mul,
masked_m,
enable_v2=True,
)
else:
sgl_per_token_group_quant_8bit_jit(
input=x,
output_q=x_q,
output_s=x_s,
group_size=group_size,
eps=eps,
fp8_min=fp8_min,
fp8_max=fp8_max,
scale_ue8m0=scale_ue8m0,
)
else:
assert not enable_v2
sgl_per_token_group_quant_fp8(