From cfb9c574d3b0b830c6b3d5324075ecdeb9821402 Mon Sep 17 00:00:00 2001
From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Date: Fri, 3 Jul 2026 04:42:17 +0800
Subject: [PATCH] Fix UE8M0 scale rounding for DeepGEMM (#29956)

---
 python/sglang/srt/layers/quantization/fp8_utils.py      | 7 ++++++-
 test/manual/quant/test_block_fp8_deep_gemm_blackwell.py | 7 ++++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py
index ab7d91b27..2c8339497 100755
--- a/python/sglang/srt/layers/quantization/fp8_utils.py
+++ b/python/sglang/srt/layers/quantization/fp8_utils.py
@@ -1499,7 +1499,12 @@ def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
 
 # COPIED FROM DeepGEMM
 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(
diff --git a/test/manual/quant/test_block_fp8_deep_gemm_blackwell.py b/test/manual/quant/test_block_fp8_deep_gemm_blackwell.py
index 833c23e7c..1d138ee41 100644
--- a/test/manual/quant/test_block_fp8_deep_gemm_blackwell.py
+++ b/test/manual/quant/test_block_fp8_deep_gemm_blackwell.py
@@ -46,7 +46,12 @@ def per_block_quant_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
 
 def ceil_to_ue8m0(x: torch.Tensor):
     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]: