[Dispatch]: dispatch per_token_group_quant_fp8_to_vllm (#1257)

Author: wenxian1

src/flag_gems/ops/per_token_group_quant_fp8.py

@@ -21,6 +21,7 @@ def _per_token_group_quant_fp8(
     eps,
     fp8_min,
     fp8_max,
+    scale_ue8m0,
     BLOCK: tl.constexpr,
 ):
     groups_per_row = y_num_columns // group_size
@@ -39,6 +40,10 @@ def _per_token_group_quant_fp8(
     y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
     _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
     y_s = _absmax / fp8_max
+
+    if scale_ue8m0:
+        y_s = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.abs(y_s), 1e-10))))
+
     y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
 
     tl.store(y_q_ptr + cols, y_q, mask=mask)
@@ -57,6 +62,7 @@ def _per_token_group_quant_fp8_colmajor(
     eps,
     fp8_min,
     fp8_max,
+    scale_ue8m0,
     BLOCK: tl.constexpr,
 ):
     groups_per_row = y_num_columns // group_size
@@ -75,6 +81,10 @@ def _per_token_group_quant_fp8_colmajor(
     y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
     _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
     y_s = _absmax / fp8_max
+
+    if scale_ue8m0:
+        y_s = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.abs(y_s), 1e-10))))
+
     y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
 
     tl.store(y_q_ptr + cols, y_q, mask=mask)
@@ -87,6 +97,7 @@ def per_token_group_quant_fp8(
     eps: float = 1e-10,
     dtype: Optional[torch.dtype] = None,
     column_major_scales: bool = False,
+    scale_ue8m0: bool = False,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     # dtype: The dype of output tensor. Note that only `torch.float8_e4m3fn`
     fp8_dtype = SUPPORTED_FP8_DTYPE if dtype is None else dtype
@@ -126,6 +137,7 @@ def per_token_group_quant_fp8(
             eps,
             fp8_min=fp8_min,
             fp8_max=fp8_max,
+            scale_ue8m0=scale_ue8m0,
             BLOCK=BLOCK,
             num_warps=num_warps,
             num_stages=num_stages,
@@ -141,6 +153,7 @@ def per_token_group_quant_fp8(
             eps,
             fp8_min=fp8_min,
             fp8_max=fp8_max,
+            scale_ue8m0=scale_ue8m0,
             BLOCK=BLOCK,
             num_warps=num_warps,
             num_stages=num_stages,

src/flag_gems/patches/patch_vllm_all.py

@@ -338,6 +338,32 @@ def custom_get_scheduler_metadata(
     )
 
 
+def custom_per_token_group_fp8_quant(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+    group_size: int,
+    eps: float,
+    fp8_min: float,
+    fp8_max: float,
+    scale_ue8m0: bool = False,
+):
+    from flag_gems.ops import per_token_group_quant_fp8
+
+    column_major_scales = output_s.stride(0) < output_s.stride(1)
+
+    x_q, x_s = per_token_group_quant_fp8(
+        x=input,
+        group_size=group_size,
+        eps=eps,
+        column_major_scales=column_major_scales,
+        scale_ue8m0=scale_ue8m0,
+    )
+
+    output_q.copy_(x_q)
+    output_s.copy_(x_s)
+
+
 def apply_gems_patches_to_vllm(verbose=True):
     import vllm  # noqa: F401
     from vllm.attention.ops.paged_attn import PagedAttention
@@ -376,3 +402,10 @@ def apply_gems_patches_to_vllm(verbose=True):
         "CUDA",
         verbose,
     )
+    patch_vllm_lib(
+        "_C",
+        "per_token_group_fp8_quant",
+        custom_per_token_group_fp8_quant,
+        "CUDA",
+        verbose,
+    )

tests/test_special_ops.py

@@ -1265,7 +1265,9 @@ def test_accuracy_contiguous(shape, dtype):
     gems_assert_equal(res_out, ref_out)
 
 
-def native_per_token_group_quant_fp8(x, group_size, eps=1e-10, dtype=None):
+def native_per_token_group_quant_fp8(
+    x, group_size, eps=1e-10, dtype=None, scale_ue8m0=False
+):
     if dtype is None:
         dtype = flag_gems.SUPPORTED_FP8_DTYPE
 
@@ -1281,6 +1283,9 @@ def native_per_token_group_quant_fp8(x, group_size, eps=1e-10, dtype=None):
     x_ = x.reshape(x.numel() // group_size, group_size)
     amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32)
     x_s = amax / fp8_max
+    if scale_ue8m0:
+        min_val = torch.tensor(1e-10, dtype=x_s.dtype, device=x_s.device)
+        x_s = torch.exp2(torch.ceil(torch.log2(torch.maximum(x_s.abs(), min_val))))
     x_q = (x_ / x_s).clamp(min=fp8_min, max=fp8_max).to(dtype)
     x_q = x_q.reshape(x.shape)
     x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,))
@@ -1294,14 +1299,21 @@ def native_per_token_group_quant_fp8(x, group_size, eps=1e-10, dtype=None):
 @pytest.mark.parametrize("dtype", FP8_QUANT_SHAPES["DTYPES"])
 @pytest.mark.parametrize("d", FP8_QUANT_SHAPES["D"])
 @pytest.mark.parametrize("num_tokens", FP8_QUANT_SHAPES["NUM_TOKENS"])
-def test_accuracy_per_token_group_quant_fp8(num_tokens, d, dtype, group_size, seed):
+@pytest.mark.parametrize("scale_ue8m0", [True, False])
+def test_accuracy_per_token_group_quant_fp8(
+    num_tokens, d, dtype, group_size, seed, scale_ue8m0
+):
     torch.manual_seed(seed)
     x = torch.rand(num_tokens, d, dtype=dtype, device=flag_gems.device)
     ref_x = to_reference(x)
 
-    ref_out, ref_scale = native_per_token_group_quant_fp8(ref_x, group_size)
+    ref_out, ref_scale = native_per_token_group_quant_fp8(
+        ref_x, group_size, scale_ue8m0=scale_ue8m0
+    )
     with flag_gems.use_gems():
-        out, scale = flag_gems.per_token_group_quant_fp8(x, group_size)
+        out, scale = flag_gems.per_token_group_quant_fp8(
+            x, group_size, scale_ue8m0=scale_ue8m0
+        )
 
     gems_assert_close(scale, ref_scale, dtype=torch.float32)