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

[perf] experimental enhance fp8 per-tensor quant (#5370)

Browse Source
This commit is contained in:
JieXin Liang
2025-04-15 03:35:43 +08:00
committed by GitHub
parent e9fc2ac7b6
commit bdde237562
4 changed files with 178 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -168,13 +168,13 @@ def input_to_float8(
"""This function quantizes input values to float8 values with tensor-wise quantization."""
finfo = torch.finfo(dtype)
min_val, max_val = x.aminmax()
amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
amax = torch.maximum(min_val.abs(), max_val.abs()).float().clamp(min=1e-12)
fp8_max = finfo.max
if _is_hip:
dtype = torch.float8_e4m3fnuz
fp8_max = 224.0
scale = fp8_max / amax
x_scl_sat = (x * scale).clamp(min=-fp8_max, max=fp8_max)
x_scl_sat = (x.float() * scale).clamp(min=-fp8_max, max=fp8_max)
return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()
@@ -213,7 +213,11 @@ def block_quant_to_tensor_quant(
for j in range(n_tiles):
x_dq_block_tiles[j][i][:, :] = x_dq_block_tiles[j][i] * x_s[j][i]
x_q_tensor, scale = input_to_float8(x_dq_block, dtype=x_q_block.dtype)
x_q_tensor, scale = (
sgl_scaled_fp8_quant(x_dq_block)
if _is_cuda
else input_to_float8(x_dq_block, dtype=x_q_block.dtype)
)
return x_q_tensor, scale
@@ -222,7 +226,11 @@ def channel_quant_to_tensor_quant(
x_s: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
x_dq_channel = x_q_channel.to(torch.float32) * x_s
x_q_tensor, scale = input_to_float8(x_dq_channel, dtype=x_q_channel.dtype)
x_q_tensor, scale = (
sgl_scaled_fp8_quant(x_dq_channel)
if _is_cuda
else input_to_float8(x_dq_channel, dtype=x_q_channel.dtype)
)
return x_q_tensor, scale