[wont merge] fp4

auto_round/data_type/nvfp.py

@@ -150,10 +150,50 @@ def cast_to_ue5m3(tensor):
 
 def cast_to_ue5m3_ste(x):
     fp4 = (cast_to_ue5m3(x).to(x.dtype) - x).detach() + x
-
     return fp4
 
 
+def ref_fp4_quant(x, global_scale, block_size=16, v=0, max_scale=1.0):
+    assert (not isinstance(global_scale, torch.Tensor)) or global_scale.dtype == torch.float32
+    assert x.ndim == 2
+    m, n = x.shape
+    if isinstance(max_scale, torch.Tensor):
+        max_scale = max_scale.unsqueeze(dim=-1).to(x.device)
+    vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32) * max_scale
+    scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX))
+    scale = torch.clip(scale, 0, FLOAT8_UE5M3_MAX)
+    scale = cast_to_ue5m3_ste(scale).to(torch.float32)
+    output_scale = get_reciprocal(scale * get_reciprocal(global_scale))
+
+    scaled_x = x.to(torch.float32) * output_scale + v
+    clipped_x = torch.clamp(scaled_x, -6.0, 6.0)
+    return (cast_to_fp4_ste(clipped_x) * get_reciprocal(output_scale)).reshape(m, n), output_scale
+
+
+@register_dtype("fp4_v2_with_global_scale")
+def full_quant(tensor, bits=4, group_size=16, v=0, max_scale=1.0, **kwargs):
+    assert (group_size == 32 or group_size == 16)
+    orig_dtype = tensor.dtype
+    tensor, orig_shape, pad_len = reshape_pad_tensor_by_group_size(tensor, group_size)
+    tensor_amax = tensor.abs().max().to(torch.float32)
+    global_scale = FLOAT8_UE5M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
+    qdq_res, output_scale = ref_fp4_quant(tensor, global_scale, group_size, v)
+    qdq_res = revert_tensor_by_pad(qdq_res, orig_shape=orig_shape, pad_len=pad_len)
+    return qdq_res.to(orig_dtype), output_scale, None
+
+
+@register_dtype("fp4_v2")
+def full_quant_no_global_scale(tensor, bits=4, group_size=32, v=0, max_scale=1.0, **kwargs):
+    assert (group_size == 32 or group_size == 16)
+    orig_dtype = tensor.dtype
+    tensor, orig_shape, pad_len = reshape_pad_tensor_by_group_size(tensor, group_size)
+    global_scale = 1.0
+    qdq_res, output_scale = ref_fp4_quant(tensor, global_scale, group_size, v, max_scale)
+    qdq_res = revert_tensor_by_pad(qdq_res, orig_shape=orig_shape, pad_len=pad_len)
+    return qdq_res.to(orig_dtype), output_scale, None
+
+
+
 if __name__ == "__main__":
     test = torch.tensor(
         [0.0, 2 ** (-17), (2 ** -14) * 0.875, 2 ** -14, 2 ** -13, 2 ** -6,