Add sub-byte data types: float4_e2m1fn, float6_e2m3fn, float6_e3m2fn #181

sergey-kozub · 2024-09-02T09:39:22Z

This PR adds MX (microscaling) floating point types support.

F4e2m1, F6e2m3, F6e3m2 types are proposed in OpenCompute MX Specification.

These types have the following notable features:

No nan encoding, only finite values are supported;
No inf encoding, similar to the existing 8-bit types with fn suffix;
Sub-byte padded bit encoding, similar to the existing int2 and int4 types.

float4_e2m1fn
- Exponent bias: 1
- Maximum stored exponent value: 3 (binary 11)
- Maximum unbiased exponent value: 3 - 1 = 2
- Minimum stored exponent value: 1 (binary 01)
- Minimum unbiased exponent value: 1 − 1 = 0
- Has Positive and Negative zero
- Doesn't have infinity
- Doesn't have NaNs

Additional details:
- Zeros (+/-): S.00.0
- Max normal number: S.11.1 = ±2^(2) x (1 + 0.5) = ±6
- Min normal number: S.01.0 = ±2^(0) = ±1
- Min subnormal number: S.00.1 = ±2^(0) x 0.5 = ±0.5

float6_e2m3fn
- Exponent bias: 1
- Maximum stored exponent value: 3 (binary 11)
- Maximum unbiased exponent value: 3 - 1 = 2
- Minimum stored exponent value: 1 (binary 01)
- Minimum unbiased exponent value: 1 − 1 = 0
- Has Positive and Negative zero
- Doesn't have infinity
- Doesn't have NaNs

Additional details:
- Zeros (+/-): S.00.000
- Max normal number: S.11.111 = ±2^(2) x (1 + 0.875) = ±7.5
- Min normal number: S.01.000 = ±2^(0) = ±1
- Max subnormal number: S.00.111 = ±2^(0) x 0.875 = ±0.875
- Min subnormal number: S.00.001 = ±2^(0) x 0.125 = ±0.125

float6_e3m2fn
- Exponent bias: 3
- Maximum stored exponent value: 7 (binary 111)
- Maximum unbiased exponent value: 7 - 3 = 4
- Minimum stored exponent value: 1 (binary 001)
- Minimum unbiased exponent value: 1 − 3 = −2
- Has Positive and Negative zero
- Doesn't have infinity
- Doesn't have NaNs

Additional details:
- Zeros (+/-): S.000.00
- Max normal number: S.111.11 = ±2^(4) x (1 + 0.75) = ±28
- Min normal number: S.001.00 = ±2^(-2) = ±0.25
- Max subnormal number: S.000.11 = ±2^(-2) x 0.75 = ±0.1875
- Min subnormal number: S.000.01 = ±2^(-2) x 0.25 = ±0.0625

Related PRs:

PR-95392 [APFloat] Add APFloat support for FP4 data type
PR-94735 [APFloat] Add APFloat support for FP6 data types

sergey-kozub · 2024-09-02T09:41:32Z

Note: a small unrelated change in "_finfo.py" removes unreadable boilerplate and replaces it with (faster) dict lookups for instantiating "finfo" objects.

hawkinsp · 2024-09-10T18:24:52Z

I'm trying to understand the relationship between these types and the MX types. From my quick read of the MX spec, all of the types it defines are block-scaled formats, which these types are not?

Can you say more about the relationship and the use case for these?

sergey-kozub · 2024-09-10T18:36:24Z

I'm trying to understand the relationship between these types and the MX types. From my quick read of the MX spec, all of the types it defines are block-scaled formats, which these types are not?

The MXFP8 type is a pair of tensors (e.g., 1st could have the E5M2 type, 2nd - the E8M0 type with 32x less elements).

Proper support of such MX type (where the value has two different primitive types) is way too complicated, but we could instead use two values. This way a dot op with scaled inputs (what we're actually interested in) could be represented as a custom call with four input tensors.

So, in order to implement MXFP8, we need E8M0 primitive type in XLA (and E5M2/E4M3 already exist). For MXFP4, we need both E8M0 and E2M1. Adding FP6 types (E2M3 and E3M2) just for completeness, they are very similar and will unblock us in the future. All of these types are described in the MX spec: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf

hawkinsp · 2024-09-10T19:12:25Z

README.md

@@ -66,6 +70,39 @@ A `bfloat16` number is a single-precision float truncated at 16 bits.

 Exponent: 8, Mantissa: 7, exponent bias: 127. IEEE 754, with NaN and inf.

+### `float4_e2m1`


Fix names to include fn suffix.

I'm actually having trouble finding a great definition of the f and n suffixes even in the LLVM discussion that added them: I don't suppose you have a link to the definition?

In particular, I'm not sure if n should appear in the name, given that n also appears in the suffix of FP8 types with a single NaN, but these have no NaN. So I'm a bit unclear what the suffix means.

f means "finite", n means "special NaN representation" (e.g. non-IEEE)
I saw this somewhere in the comments, will post a link once I find it.

Fixed the type name.

Oh, actually, it's in this same file, below:

F is for "finite" (no infinities), N for with special NaN encoding, UZ for unsigned zero.

I guess one could say that "no NaN encoding" is a "special NaN encoding".

Also, LLVM APFloat.cpp has these types with "FN" suffix:
https://github.com/llvm/llvm-project/blob/5537ae87b3a87b3abeb4e6983cecd9b103648243/llvm/lib/Support/APFloat.cpp#L150

We could probably change the suffix, but we need to be consistent across the repositories.

We should agree with LLVM, so that works for me.

Not sure if you didn't push the fix yet, the headers are still suffix-less.

hawkinsp · 2024-09-10T19:12:51Z

README.md

+Microscaling format, 4 bits (encoding: `0bSEEM`) using byte storage (higher 4
+bits are unused). NaN representation is undefined.
+
+Possible values: [0, 0.5, 1, 1.5, 2, 3, 4, 6]


I'd probably stick backticks around the values

What about the negative values?

Added backticks around the values (here and below).

Changed to "Possible absolute values" to keep the list short.

hawkinsp · 2024-09-10T19:16:09Z

ml_dtypes/_finfo.py

+    obj.epsneg = 0.125
+    obj.machep = -3
+    obj.negep = -3
+    obj.max = float6_e2m3fn(7.5)


I'd personally be tempted to specify these as bit patterns (float.fromhex("0x1234.1"), IIRC)

hawkinsp · 2024-09-10T19:21:52Z

ml_dtypes/_src/dtypes.cc

@@ -366,6 +423,14 @@ bool Initialize() {
  success &= RegisterTwoWayCustomCast<float8_e3m4, float8_e4m3fn, float>();
  success &= RegisterTwoWayCustomCast<float8_e3m4, float8_e5m2, float>();
  success &= RegisterTwoWayCustomCast<float8_e3m4, float8_e4m3, float>();
+


This is getting unwieldy. This is just covering all-pairs of extension types, I think? I suspect this can be factored better with some template trickery.

If nothing else, the function you added called RegisterCustomCastsWithBfloat16AndFloat8Types could just be used everywhere here and you call it once for each type?

Probably possible to do better than that with some template cunning.

Added some templates to reduce boilerplate in this file.

ml_dtypes/include/mxfloat.h

hawkinsp

Looks good, other than the clang-format failure.

…U) (#2581) This is a proposal to add MX (microscaling) floating point types to StableHLO. Related links: - StableHLO [PR#2582](#2582) Add MX floating point types (f4E2M1FN, f6E2M3FN, f6E3M2FN, f8E8M0FNU) - LLVM [PR#95392](llvm/llvm-project#95392) [APFloat] Add APFloat support for FP4 data type - LLVM [PR#94735](llvm/llvm-project#94735) [APFloat] Add APFloat support for FP6 data types - LLVM [PR#107127](llvm/llvm-project#107127) [APFloat] Add APFloat support for E8M0 type - LLVM [PR#108877](llvm/llvm-project#108877) [MLIR] Add f4E2M1FN type - LLVM [PR#107999](llvm/llvm-project#107999) [MLIR] Add f6E2M3FN type - LLVM [PR#105573](llvm/llvm-project#105573) [MLIR] Add f6E3M2FN type - LLVM [PR#111028](llvm/llvm-project#111028) [MLIR] Add f8E8M0FNU type - JAX-ML [PR#181](jax-ml/ml_dtypes#181) Add sub-byte data types: float4_e2m1fn, float6_e2m3fn, float6_e3m2fn - JAX-ML [PR#166](jax-ml/ml_dtypes#181) Add float8_e8m0_fnu (E8M0) OCP MX scale format

jakevdp assigned hawkinsp Sep 2, 2024

jakevdp requested a review from hawkinsp September 2, 2024 14:21

sergey-kozub force-pushed the mxfloat branch 3 times, most recently from e51626b to 9bdf962 Compare September 10, 2024 19:21

hawkinsp reviewed Sep 10, 2024

View reviewed changes

sergey-kozub force-pushed the mxfloat branch 3 times, most recently from f6f5f43 to 3d633fc Compare September 10, 2024 19:58

hawkinsp mentioned this pull request Sep 10, 2024

Add float8_e8m0_fnu (E8M0) OCP MX scale format #166

Merged

sergey-kozub force-pushed the mxfloat branch 2 times, most recently from a6aaef4 to 499967c Compare September 11, 2024 07:33

Add sub-byte data types: float4_e2m1fn, float6_e2m3fn, float6_e3m2fn

b68531f

sergey-kozub force-pushed the mxfloat branch from 499967c to b68531f Compare September 11, 2024 08:19

hawkinsp reviewed Sep 11, 2024

View reviewed changes

hawkinsp approved these changes Sep 12, 2024

View reviewed changes

hawkinsp added the pull ready label Sep 12, 2024

copybara-service bot merged commit 40e66e5 into jax-ml:main Sep 12, 2024
12 of 13 checks passed

This was referenced Sep 13, 2024

Support unpacked FP4E2M1 #164

Closed

Add MXFP4 (FP4 E2M1) and MXFP6 (FP6 E3M2/E2M3) #116

Closed

sergey-kozub mentioned this pull request Oct 7, 2024

[RFC] Microscaling data types (f4E2M1FN, f6E2M3FN, f6E3M2FN, f8E8M0FNU) openxla/stablehlo#2581

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add sub-byte data types: float4_e2m1fn, float6_e2m3fn, float6_e3m2fn #181

Add sub-byte data types: float4_e2m1fn, float6_e2m3fn, float6_e3m2fn #181

sergey-kozub commented Sep 2, 2024

sergey-kozub commented Sep 2, 2024

hawkinsp commented Sep 10, 2024

sergey-kozub commented Sep 10, 2024 •

edited

Loading

hawkinsp Sep 10, 2024

sergey-kozub Sep 10, 2024

sergey-kozub Sep 10, 2024

sergey-kozub Sep 10, 2024

sergey-kozub Sep 10, 2024

hawkinsp Sep 10, 2024

hawkinsp Sep 10, 2024

sergey-kozub Sep 10, 2024

hawkinsp Sep 10, 2024

hawkinsp Sep 10, 2024

sergey-kozub Sep 10, 2024

hawkinsp Sep 10, 2024

sergey-kozub Sep 11, 2024

hawkinsp Sep 10, 2024

sergey-kozub Sep 11, 2024

hawkinsp left a comment

		@@ -66,6 +70,39 @@ A `bfloat16` number is a single-precision float truncated at 16 bits.

		Exponent: 8, Mantissa: 7, exponent bias: 127. IEEE 754, with NaN and inf.

		### `float4_e2m1`

Add sub-byte data types: float4_e2m1fn, float6_e2m3fn, float6_e3m2fn #181

Add sub-byte data types: float4_e2m1fn, float6_e2m3fn, float6_e3m2fn #181

Conversation

sergey-kozub commented Sep 2, 2024

sergey-kozub commented Sep 2, 2024

hawkinsp commented Sep 10, 2024

sergey-kozub commented Sep 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hawkinsp left a comment

Choose a reason for hiding this comment

sergey-kozub commented Sep 10, 2024 •

edited

Loading