From d971f228987b7a27ad3b266e4a2508d29b126554 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sat, 15 Nov 2025 17:31:04 +0800 Subject: [PATCH] Super tiny expose transform_scale_ue8m0 API for RL frameworks (#13323) --- python/sglang/srt/layers/quantization/fp8_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index ed62ebb01..5cbde4c3e 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -460,7 +460,7 @@ def _requant_weight_ue8m0( weight_block_size=weight_block_size, ) - out_s = _transform_scale_ue8m0(out_s, mn=out_w.shape[-2]) + out_s = transform_scale_ue8m0(out_s, mn=out_w.shape[-2]) return out_w, out_s @@ -492,11 +492,11 @@ def quant_weight_ue8m0( def transform_scale_ue8m0_inplace(param, mn): - param.data = _transform_scale_ue8m0(param.data, mn=mn) + param.data = transform_scale_ue8m0(param.data, mn=mn) # NOTE copy and modified from DeepGEMM -def _transform_scale_ue8m0(sf, mn): +def transform_scale_ue8m0(sf, mn): import deep_gemm.utils.layout sf = sf.index_select(-2, torch.arange(mn, device=sf.device) // 128) @@ -507,7 +507,7 @@ def _transform_scale_ue8m0(sf, mn): def inverse_transform_scale_ue8m0(sf_packed, mn): sf_fp32 = _inverse_transform_scale_ue8m0_impl(sf_packed) # Can call consistency check every time since this is only called on startup - sf_packed_recreated = _transform_scale_ue8m0(sf_fp32, mn=mn) + sf_packed_recreated = transform_scale_ue8m0(sf_fp32, mn=mn) assert torch.all( sf_packed == sf_packed_recreated ), f"{sf_packed=} {sf_packed_recreated}"