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

Tiny support 3D tensors in inverse_transform_scale_ue8m0 (#14002)

Browse Source
This commit is contained in:
fzyzcjy
2025-11-27 20:20:45 +08:00
committed by GitHub
parent ab843ced31
commit 2bc8ee8b74
1 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -532,8 +532,13 @@ def _inverse_transform_scale_ue8m0_impl(sf_packed):
:param sf_packed: (scale_mn, scale_k/4) int32
:return: (scale_mn, scale_k), float32
"""
if len(sf_packed.shape) == 3:
return torch.stack(
[_inverse_transform_scale_ue8m0_impl(x) for x in sf_packed], dim=0
)
block_size = 128
assert len(sf_packed.shape) == 2
assert len(sf_packed.shape) == 2, f"{sf_packed.shape=}"
assert sf_packed.dtype == torch.int32
mn_repeat_128, k_div_4 = sf_packed.shape
@@ -547,7 +552,12 @@ def _inverse_transform_scale_ue8m0_impl(sf_packed):
# remove repeat
sf_reshaped = sf_fp32.view(mn, block_size, k)
sf_unrepeated = sf_reshaped[:, 0:1, :]
assert torch.all(sf_unrepeated == sf_reshaped)
if not torch.all(sf_unrepeated == sf_reshaped):
from sglang.srt.debug_utils.dumper import get_tensor_info
raise AssertionError(
f"sf_unrepeated != sf_reshaped ({get_tensor_info(sf_unrepeated)=} {get_tensor_info(sf_reshaped)=})"
)
sf_unrepeated = sf_unrepeated.squeeze(1).contiguous()
assert sf_unrepeated.shape == (mn, k)