Fix UE8M0 scale rounding for DeepGEMM (#29956)

This commit is contained in:
Mohammad Miadh Angkad
2026-07-02 13:42:17 -07:00
committed by GitHub
parent bc25abb786
commit cfb9c574d3
2 changed files with 12 additions and 2 deletions
@@ -1499,7 +1499,12 @@ def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
# COPIED FROM DeepGEMM # COPIED FROM DeepGEMM
def ceil_to_ue8m0(x: torch.Tensor): def ceil_to_ue8m0(x: torch.Tensor):
return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) bits = x.abs().float().view(torch.int32)
exp = (bits >> 23) & 0xFF
mantissa = bits & 0x7FFFFF
exp = exp + (mantissa != 0).to(torch.int32)
exp = exp.clamp(1, 254)
return (exp << 23).view(torch.float32)
def channel_quant_to_tensor_quant( def channel_quant_to_tensor_quant(
@@ -46,7 +46,12 @@ def per_block_quant_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
def ceil_to_ue8m0(x: torch.Tensor): def ceil_to_ue8m0(x: torch.Tensor):
assert x.view(-1).amax().item() > 0 assert x.view(-1).amax().item() > 0
return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) bits = x.abs().float().view(torch.int32)
exp = (bits >> 23) & 0xFF
mantissa = bits & 0x7FFFFF
exp = exp + (mantissa != 0).to(torch.int32)
exp = exp.clamp(1, 254)
return (exp << 23).view(torch.float32)
def per_token_group_quant_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: def per_token_group_quant_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: