Support new DeepGEMM format in per token group quant (#7146)

2025-06-13 17:00:22 +08:00
parent aa46ed34d2
commit 5c66c4424f
5 changed files with 92 additions and 44 deletions
--- a/sgl-kernel/python/sgl_kernel/gemm.py
+++ b/sgl-kernel/python/sgl_kernel/gemm.py
@@ -90,9 +90,10 @@ def sgl_per_token_group_quant_fp8(
    eps: float,
    fp8_min: float,
    fp8_max: float,
+    scale_ue8m0: bool,
 ) -> None:
    torch.ops.sgl_kernel.sgl_per_token_group_quant_fp8.default(
-        input, output_q, output_s, group_size, eps, fp8_min, fp8_max
+        input, output_q, output_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0
    )