Add fp8_e8m0 scale type support

TensorLib/Dtype.lean

@@ -39,6 +39,7 @@ inductive Dtype where
 | float8_e4m3
 | float8_e3m4
 | float8_e5m2
+| float8_e8m0
 | float16
 | bfloat16
 | float32
@@ -65,6 +66,7 @@ def gen : Gen Dtype := Gen.elements [
   float8_e4m3,
   float8_e3m4,
   float8_e5m2,
+  float8_e8m0,
   float16,
   bfloat16,
   float32,
@@ -90,14 +92,15 @@ instance : ToString Dtype where
   | float8_e4m3 => "float8_e4m3fn"
   | float8_e3m4 => "float8_e3m4"
   | float8_e5m2 => "float8_e5m2" -- no fn since e5m2 has infinity
+  | float8_e8m0 => "float8_e8m0"
   | float16 => "float16"
   | bfloat16 => "bfloat16"
   | float32 => "float32"
   | float64 => "float64"
 
 
 def isOneByte (x : Dtype) : Bool := match x with
-| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 => true
+| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 | float8_e8m0 => true
 | _ => false
 
 def isMultiByte (x : Dtype) : Bool := ! x.isOneByte
@@ -134,15 +137,15 @@ def intMax (x : Dtype) : Int := match x with
 
 -- Added float16 and bfloat16 so bitwise op know to reject it
 def isFloat (x : Dtype) : Bool := match x with
-| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 => true
+| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 | float8_e8m0 => true
 | _ => false
 
 --! Number of bytes used by each element of the given dtype
 def itemsize (x : Dtype) : Nat := match x with
 | float64 | int64 | uint64 => 8
 | float32 | int32 | uint32 => 4
 | bfloat16 | float16 | int16 | uint16 => 2
-| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 => 1
+| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 | float8_e8m0 => 1
 
 -- Previously this was inline in join with a recursive swap,
 -- but adding more fp8 types made the match too large. Lean needs to prove
@@ -235,11 +238,16 @@ private def joinOrdered (x y : Dtype) : Option Dtype :=
   | _, _ => none
 
 def join (x y : Dtype) : Option Dtype :=
-  if x = y then x else if x.itemsize > y.itemsize then joinOrdered y x else joinOrdered x y
+  if x = y then x
+  -- e8m0 is a scale type so it doesnt promote with any other type.
+  else if x == .float8_e8m0 || y == .float8_e8m0 then none
+  else if x.itemsize > y.itemsize then joinOrdered y x
+  else joinOrdered x y
 
 
 -- Can we cast from one dtype to another without losing information
 def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
+| .bool, .float8_e8m0 => false
 | .bool, _ => true
 | _, .bool => false
 | .int8, .int8
@@ -316,6 +324,7 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with
 | .float8_e3m4, .float32
 | .float8_e3m4, .float64 => true
 | .float8_e3m4, _ => false
+| .float8_e8m0, _ => false
 | .float32, .float32
 | .float32, .float64 => true
 | .float32, _ => false
@@ -349,6 +358,16 @@ OverflowError: Python integer 128 out of bounds for int8
 
 Float types have named safe nat upper bounds.
 -/
+
+-- Maximum representable Fp32 value for each MX compute dtype.
+-- Used by quantizeMX to compute the E8M0 scale: m = fp8Max / amax.
+-- These are the actual format maxima, not the largest safe integer (see maxSafeNat).
+def fp8Max (dtype : Dtype) : Option Float32 := match dtype with
+  | .float8_e4m3 => some 448.0
+  | .float8_e5m2 => some 57344.0
+  | .float8_e3m4 => some 15.5
+  | _ => none
+
 private def maxSafeNat : Dtype -> Option Nat
 | .bool => none
 | .uint8 => some 0xFF
@@ -362,6 +381,7 @@ private def maxSafeNat : Dtype -> Option Nat
 | .float8_e4m3 => maxSafeNatForFloat8e4m3
 | .float8_e3m4 => maxSafeNatForFloat8e3m4
 | .float8_e5m2 => maxSafeNatForFloat8e5m2
+| .float8_e8m0 => none
 | .float16 => maxSafeNatForFloat16
 | .bfloat16 => maxSafeNatForBFloat16
 | .float32 => maxSafeNatForFloat32
@@ -384,6 +404,7 @@ private def minSafeInt : Dtype -> Option Int
 | .float8_e4m3 => some (-maxSafeNatForFloat8e4m3)
 | .float8_e3m4 => some (-maxSafeNatForFloat8e3m4)
 | .float8_e5m2 => some (-maxSafeNatForFloat8e5m2)
+| .float8_e8m0 => none
 | .float16 => some (-maxSafeNatForFloat16)
 | .bfloat16 => some (-maxSafeNatForBFloat16)
 | .float32 => some (-maxSafeNatForFloat32)
@@ -424,11 +445,18 @@ def decodeFloat8E3M4 (arr : ByteArray) : Err Float32 :=
 private def encodeFloat8E3M4 (f : Float32) : ByteArray :=
   ByteArray.mk #[f.toFloat8E3M4Bits]
 
+-- Decode 1-byte fp8_e8m0 to Float32.
+-- Centralizes the size check so callers don't need inline guards.
+def decodeFloat8E8M0 (arr : ByteArray) : Err Float32 :=
+  if arr.size != 1 then .error "decoder: expected 1 byte for float8_e8m0"
+  else .ok (arr.data[0]!.toFloat32FromFloat8E8M0)
+
 -- Dispatch fp8 decode by dtype
 private def decodeFloat8 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with
   | .float8_e4m3 => decodeFloat8E4M3 arr
   | .float8_e5m2 => decodeFloat8E5M2 arr
   | .float8_e3m4 => decodeFloat8E3M4 arr
+  | .float8_e8m0 => decodeFloat8E8M0 arr
   | _ => .error "decoder: expected float8 type"
 
 -- Dispatch fp8 encode by dtype
@@ -451,6 +479,7 @@ def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype
 | .float8_e4m3 => encodeFloat8E4M3 n.toFloat32
 | .float8_e3m4 => encodeFloat8E3M4 n.toFloat32
 | .float8_e5m2 => encodeFloat8E5M2 n.toFloat32
+| .float8_e8m0 => panic! "byteArrayOfNatOverflow not meaningful for float8_e8m0 (scale-only type)"
 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits
 | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits
 | .float32 => toLEByteArray n.toFloat32
@@ -562,6 +591,7 @@ private def byteArrayOfIntOverflow (dtype : Dtype) (n : Int) : ByteArray := matc
 | .float8_e4m3 => encodeFloat8E4M3 n.toFloat32
 | .float8_e3m4 => encodeFloat8E3M4 n.toFloat32
 | .float8_e5m2 => encodeFloat8E5M2 n.toFloat32
+| .float8_e8m0 => panic! "byteArrayOfIntOverflow not meaningful for float8_e8m0 (scale-only type)"
 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits
 | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits
 | .float32 => toLEByteArray n.toFloat32
@@ -719,6 +749,7 @@ def add (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
     let x <- decodeFloat8E5M2 x
     let y <- decodeFloat8E5M2 y
     return encodeFloat8E5M2 (x + y)
+  | .float8_e8m0 => .error "Arithmetic: addition not supported for float8_e8m0 (scale-only type)"
   | .float16
   | .bfloat16 => do
     let x <- dtype.decodeFloat16OrBFloat16 x
@@ -754,6 +785,7 @@ def sub (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
     let x <- decodeFloat8E5M2 x
     let y <- decodeFloat8E5M2 y
     return encodeFloat8E5M2 (x - y)
+  | .float8_e8m0 => .error "Arithmetic: subtraction not supported for float8_e8m0 (scale-only type)"
   | .float16
   | .bfloat16 => do
     let x <- dtype.decodeFloat16OrBFloat16 x
@@ -790,6 +822,7 @@ def mul (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
     let x <- decodeFloat8E5M2 x
     let y <- decodeFloat8E5M2 y
     return encodeFloat8E5M2 (x * y)
+  | .float8_e8m0 => .error "Arithmetic: multiplication not supported for float8_e8m0 (scale-only type)"
   | .float16
   | .bfloat16 => do
     let x <- dtype.decodeFloat16OrBFloat16 x
@@ -826,6 +859,7 @@ def div (dtype : Dtype) (x y : ByteArray) : Err ByteArray :=
     let x <- decodeFloat8E5M2 x
     let y <- decodeFloat8E5M2 y
     return encodeFloat8E5M2 (x / y)
+  | .float8_e8m0 => .error "Arithmetic: division not supported for float8_e8m0 (scale-only type)"
   | .float16
   | .bfloat16 => do
     let x <- dtype.decodeFloat16OrBFloat16 x
@@ -861,6 +895,7 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do
   | .float8_e5m2 => do
     let f <- decodeFloat8E5M2 x
     return encodeFloat8E5M2 f.abs
+  | .float8_e8m0 => .error "Absolute value not supported for float8_e8m0 (scale-only type)"
   | .float16
   | .bfloat16 => do
     let x <- dtype.decodeFloat16OrBFloat16 x
@@ -900,6 +935,7 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with
 | float8_e5m2 => do
   let f <- decodeFloat8E5M2 x
   return f == 0
+| .float8_e8m0 => return false
 | float16
 | bfloat16 => do
   let f <- dtype.decodeFloat16OrBFloat16 x
@@ -918,6 +954,8 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with
 def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err ByteArray :=
   if fromDtype == toDtype then return data else
     match fromDtype, toDtype with
+    | .float8_e8m0, _ => .error "castOverflow not supported for float8_e8m0 (scale-only type)"
+    | _, .float8_e8m0 => .error "castOverflow not supported for float8_e8m0 (scale-only type)"
     -- For floats use isZero so -0 is correctly handled.
     -- A raw byte check would treat -0.0 as !0 since sign bit is nonzero
     | _, bool =>
@@ -1049,6 +1087,26 @@ def logicalNot : Dtype -> ByteArray -> Err Bool := isZero
 #guard Dtype.float32.isZero! $ toLEByteArray (-0.0 : Float32)
 #guard Dtype.float64.isZero! $ toLEByteArray (-0.0 : Float)
 
+-- Round a Float32 value to the nearest representable value in the given compute dtype.
+-- Does this by encoding to the dtype's bit pattern then decoding back to Float32.
+-- This captures the element rounding error introduced by quantization.
+-- Returns Err because not all dtypes are valid compute dtypes (e.g. float8_e8m0 is scale-only).
+def roundToComputeDtype (v : Float32) (dtype : Dtype) : Err Float32 := match dtype with
+  -- encode fp32 -> fp8 bits, then decode fp8 bits -> fp32
+  | .float8_e4m3 => decodeFloat8E4M3 (ByteArray.mk #[v.toFloat8E4M3Bits])
+  | .float8_e5m2 => decodeFloat8E5M2 (ByteArray.mk #[v.toFloat8E5M2Bits])
+  | .float8_e3m4 => decodeFloat8E3M4 (ByteArray.mk #[v.toFloat8E3M4Bits])
+  -- TODO: add fp8_e2m5 when the PR is merged
+  -- encode fp32 -> fp16 bits, then decode fp16 bits -> fp32
+  | .float16  => byteArrayToFloat16  .float16  (toLEByteArray v.toFloat16Bits)
+  -- encode fp32 -> bf16 bits, then decode bf16 bits -> fp32
+  | .bfloat16 => byteArrayToBFloat16 .bfloat16 (toLEByteArray v.toBFloat16Bits)
+  -- float32 round-trip is identity (no precision loss)
+  | .float32  => .ok v
+  -- float8_e8m0 is a scale-only type, not a compute dtype
+  | .float8_e8m0 => .error "roundToComputeDtype: float8_e8m0 is a scale-only type"
+  | _ => .error s!"roundToComputeDtype: unsupported dtype {dtype}"
+
 private def logicalBinop (f : Bool -> Bool -> Bool) (t1 : Dtype) (x1 : ByteArray) (t2 : Dtype) (x2 : ByteArray) : Err Bool := do
   let z1 <- t1.nonZero x1
   let z2 <- t2.nonZero x2
@@ -1108,6 +1166,7 @@ private def liftFloatUnop (f32 : Float32 -> Err Float32) (f64 : Float -> Err Flo
     let f <- decodeFloat8E3M4 data
     let x <- f32 f
     return encodeFloat8E3M4 x
+  | .float8_e8m0 => throw "float operations not supported for float8_e8m0 (scale-only type)"
   | .float16 | .bfloat16 => do
     let f <- decodeFloat16OrBFloat16 dtype data
     let x <- f32 f
@@ -1163,7 +1222,7 @@ def tanh : Dtype -> ByteArray -> Err ByteArray :=
 def tanh! (dtype : Dtype) (data : ByteArray) : ByteArray := get! $ tanh dtype data
 
 private def shift (f : UInt64 -> UInt64 -> UInt64) (dtype : Dtype) (bits : ByteArray) (shiftAmount : ByteArray) : Err ByteArray := match dtype with
-| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 => throw "shifts not supported at float type"
+| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 | .float8_e8m0 => throw "shifts not supported at float type"
 | .bool => throw "In NumPy, bool shifts are cast to int64. This seems arbitrary so please cast (e.g. with astype) before you shift."
 | .uint64 | .int64 | .uint32 | .int32 | .uint16 | .int16 | .uint8 | .int8 =>
   let k := dtype.itemsize
@@ -1186,6 +1245,7 @@ def rightShift : Dtype -> ByteArray -> ByteArray -> Err ByteArray :=
 def rightShift! (dtype : Dtype) (bits : ByteArray) (shiftAmount : ByteArray) : ByteArray :=
   get! $ rightShift dtype bits shiftAmount
 
+
 section Bitwise
 
 open scoped Iterator.PairLockStep
@@ -1308,12 +1368,14 @@ private def canCastLosslessRoundTrip (fromDtype : Dtype) (data : ByteArray) (toD
   | .error _ => false
 
 private def canCastLosslessIntRoundTrip (fromDtype : Dtype) (n : Int) (toDtype : Dtype) : Bool :=
-  let res := do
-    let n <- fromDtype.byteArrayOfInt n
-    return canCastLosslessRoundTrip fromDtype n toDtype
-  match res with
-  | .ok b => b
-  | .error _ => false
+  if fromDtype == .float8_e8m0 || toDtype == .float8_e8m0 then false
+  else
+    let res := do
+      let n <- fromDtype.byteArrayOfInt n
+      return canCastLosslessRoundTrip fromDtype n toDtype
+    match res with
+    | .ok b => b
+    | .error _ => false
 
 #guard
   let fromDtype := Dtype.int8
@@ -1377,8 +1439,9 @@ warning: declaration uses 'sorry'
 -/
 #guard_msgs in
 example (fromDtype toDtype : Dtype) (n : Nat) :
-  canCastLosslessIntRoundTrip fromDtype 0 toDtype &&
-  canCastLosslessIntRoundTrip fromDtype 1 toDtype
+  fromDtype == .float8_e8m0 || toDtype == .float8_e8m0 ||
+  (canCastLosslessIntRoundTrip fromDtype 0 toDtype &&
+  canCastLosslessIntRoundTrip fromDtype 1 toDtype)
   := by plausible
 
 /--
@@ -1390,7 +1453,7 @@ warning: declaration uses 'sorry'
 -- One dtype should always go back and forth
 -- skip values outside dtypes range since they cannot be encoded in the first place.
 example (dtype : Dtype) (n : Nat) :
-  if n > dtype.maxSafeNat.getD n then true else canCastLosslessIntRoundTrip dtype n dtype := by plausible
+  dtype == .float8_e8m0 || (if n > dtype.maxSafeNat.getD n then true else canCastLosslessIntRoundTrip dtype n dtype) := by plausible
 
 /--
 info: Unable to find a counter-example

TensorLib/Float.lean

@@ -652,7 +652,35 @@ def _root_.Float32.toFloat8E3M4Bits (f : Float32) : UInt8 :=
 #guard (Float32.ofBits 0x3D000000).toFloat8E3M4Bits == (2 : UInt8)      -- 0.03125
 #guard (Float32.ofBits 0x3D800000).toFloat8E3M4Bits == (4 : UInt8)      -- 0.0625
 -- Negative overflow
-#guard (Float32.ofBits 0xC1800000).toFloat8E3M4Bits == (240 : UInt8)    -- -16.0 → -inf
+#guard (Float32.ofBits 0xC1800000).toFloat8E3M4Bits == (240 : UInt8)    -- -16.0 -> -inf
+
+
+-- Decoder for fp8_e8m0 (scale type)
+-- Reference: http://kib.kiev.ua/x86docs/Third-Parties/OCP/OCP_Microscaling%20Formats%20(MX)%20v1.0%20Spec_Final.pdf
+-- e8m0 is 8 bits unsigned bias exp (bias = 127), 0 mant bits
+-- Every value is a power of 2: 2 ^ (byte - 127)
+-- 0xFF = NaN; no Inf, no 0, no subnormals
+def _root_.UInt8.toFloat32FromFloat8E8M0 (bits: UInt8) : Float32 :=
+  -- case NaN
+  if bits == 0xFF then
+    -- byte 255 is NaN encoding acc to OCP
+    Float32.ofBits 0x7FC00000
+  else if bits == 0 then
+    -- Byte 0: 2^(-127) is a fp32 subnormal (below fp32's min normal 2^-126)
+    -- fp32 subnormal: sign=0, exp=0, mant=1<<22 gives 2^(-126) × 0.5 = 2^(-127)
+    Float32.ofBits 0x00400000
+  else
+    -- 2 ^ (byte - 127): construct fp32 bit pattern with sin = 0, exp = byte, mant = 0
+    -- fp32 value = 2 ^ (exp - 127) which is the value we want
+    Float32.ofBits (bits.toUInt32 <<< 23)
+
+-- E8M0 decode tests (verified against OCP MX spec)
+#guard (127 : UInt8).toFloat32FromFloat8E8M0 == 1.0 -- 2^(127-127) = 2^0 = 1.0
+#guard (128 : UInt8).toFloat32FromFloat8E8M0 == 2.0 -- 2^(128-127) = 2^1 = 2.0
+#guard (126 : UInt8).toFloat32FromFloat8E8M0 == 0.5 -- 2^(126-127) = 2^(-1) = 0.5
+#guard (254 : UInt8).toFloat32FromFloat8E8M0 == Float32.ofBits 0x7F000000 -- 2^127 (largest value)
+#guard (0 : UInt8).toFloat32FromFloat8E8M0 == Float32.ofBits 0x00400000 -- byte 0: 2^-127 as fp32 subnormal (0x00400000). Naïve bits<<23 would give +0.
+#guard (0xFF : UInt8).toFloat32FromFloat8E8M0.toBits == 0x7FC00000           -- byte 255 = NaN
 
 section Test
 

TensorLib/Npy.lean

@@ -114,7 +114,7 @@ def dtypeNameToNpyString (t : TensorLib.Dtype) : String := match t with
 -- float8_e3m4 serializes as "V1" in ml_dtypes, same as e4m3.
 -- The npy format cannot distinguish between fp8 subtypes that use V1.
 -- Reading "<V1" defaults to e4m3
-| .float8_e4m3 | .float8_e3m4 => "V1"
+| .float8_e4m3 | .float8_e3m4 | .float8_e8m0 => "V1"
 | .float8_e5m2 => "f1"
 | .float16 => "f2"
 | .bfloat16 => "V2"
@@ -433,8 +433,8 @@ end Save
 -- unreachable normally since toNpy blocks e3m4 before reaching
 -- save!, and parseFile maps V1 to e4m3. Guards against hand-constructed Ndarrays.
 def Ndarray.save! (arr : Ndarray) (file : System.FilePath) : IO Unit :=
-  if arr.header.descr.name == .float8_e3m4 then
-    throw $ IO.userError "float8_e3m4 cannot be saved to npy: format uses V1 which is indistinguishable from float8_e4m3"
+  if arr.header.descr.name == .float8_e3m4 || arr.header.descr.name == .float8_e8m0 then
+    throw $ IO.userError "float8_e3m4/float8_e8m0 cannot be saved to npy: format uses V1 which is indistinguishable from float8_e4m3"
   else
     IO.FS.writeBinFile file arr.toByteArray!
 
@@ -448,6 +448,7 @@ def Ndarray.save! (arr : Ndarray) (file : System.FilePath) : IO Unit :=
 -- Known limitation: e3m4 cannot round-trip through npy (reads back as e4m3)
 #guard Npy.Dtype.fromNpyString "<V1" != .ok { name := .float8_e3m4, order := .littleEndian }
 #guard Npy.Dtype.dtypeNameToNpyString .float8_e4m3 == "V1"
+#guard Npy.Dtype.dtypeNameToNpyString .float8_e8m0 == "V1"
 
 end Npy
 end TensorLib