From 7caab1c1ed8de1430854cea4887fa24ada99ab51 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Tue, 4 Aug 2026 15:09:01 -0700 Subject: [PATCH 1/9] Add fp8_e8m0 scale type support --- TensorLib/Dtype.lean | 52 ++++++++++++++++++++++++++++++++++--------- TensorLib/Float.lean | 29 +++++++++++++++++++++++- TensorLib/Npy.lean | 7 +++--- TensorLib/Tensor.lean | 7 +++++- TensorLib/Test.lean | 28 ++++++++++++++++++++++- 5 files changed, 106 insertions(+), 17 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 7f2df3f..4eb9170 100644 --- a/TensorLib/Dtype.lean +++ b/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,6 +92,7 @@ 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" @@ -97,7 +100,7 @@ instance : ToString Dtype where 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,7 +137,7 @@ 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 @@ -142,7 +145,7 @@ 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,7 +238,11 @@ 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 @@ -316,6 +323,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 @@ -362,6 +370,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 +393,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 +434,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 +468,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 +580,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 +738,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 +774,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 +811,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 +848,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 +884,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 +924,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 +943,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 => @@ -1108,6 +1135,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 +1191,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 @@ -1308,12 +1336,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 true + 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 diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index 1bc96d8..f497238 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -652,7 +652,34 @@ 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 + +-- Encoder for fp8_e8m0 (scale type) +-- 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: fp32 exp=0, mant=0 = +0 (not 2^-127) section Test diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 5c8f085..53059f0 100644 --- a/TensorLib/Npy.lean +++ b/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" +| .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 " t.mapM (fun b => Dtype.decodeFloat8E5M2 b) | .float8_e4m3 => t.mapM (fun b => Dtype.decodeFloat8E4M3 b) | .float8_e3m4 => t.mapM (fun b => Dtype.decodeFloat8E3M4 b) + -- e8m0 is a scale-only type, but we include a decode case here to prevent + -- the default branch from misinterpreting 1-byte e8m0 data as multi-byte fp32/fp64. + | .float8_e8m0 => t.mapM (fun b => Dtype.decodeFloat8E8M0 b) | .float16 => t.mapM (fun b => Dtype.byteArrayToFloat16 .float16 b) | .bfloat16 => t.mapM (fun b => Dtype.byteArrayToBFloat16 .bfloat16 b) | _ => t.mapM ( fun b => Float32.ofLEByteArray b) @@ -612,6 +615,7 @@ def toFloat64Tree (arr : Tensor) : Err (Format.Tree Float) := do | .float8_e5m2 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E5M2 b; return f.toFloat) | .float8_e4m3 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E4M3 b; return f.toFloat) | .float8_e3m4 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E3M4 b; return f.toFloat) + | .float8_e8m0 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E8M0 b; return f.toFloat) | .float16 => t.mapM (fun b => do let f <- Dtype.byteArrayToFloat16 .float16 b; return f.toFloat) | .bfloat16 => t.mapM (fun b => do let f <- Dtype.byteArrayToBFloat16 .bfloat16 b; return f.toFloat) | .float32 => t.mapM (fun b => do let f <- Float32.ofLEByteArray b; return f.toFloat) @@ -681,7 +685,7 @@ def toNpy (arr : Tensor) : Err Npy.Ndarray := -- Our guard helps to surface an explicit error during write instead of allowing a -- silent round-trip corruption — without it, a user could save an e3m4 tensor, load it back, -- and get wrong values (interpreted as e4m3) with no indication anything went wrong. - if arr.dtype == .float8_e3m4 then .error "float8_e3m4 cannot be saved to npy: format uses V1 which is indistinguishable from float8_e4m3" + if arr.dtype == .float8_e3m4 || arr.dtype == .float8_e8m0 then .error "float8_e3m4/float8_e8m0 cannot be saved to npy: format uses V1 which is indistinguishable from float8_e4m3" else let arr := if arr.isTriviallyReshapable then arr else arr.copy let descr := Npy.Dtype.mk arr.dtype Npy.ByteOrder.littleEndian @@ -771,6 +775,7 @@ open TensorLib.Tensor.Format.Tree #guard match (Tensor.zeros .float8_e3m4 (Shape.mk [2])).toNpy with | .error _ => true | .ok _ => false -- toNpy accepts e4m3 (not blocked) #guard match (Tensor.zeros .float8_e4m3 (Shape.mk [2])).toNpy with | .ok _ => true | .error _ => false +#guard match (Tensor.zeros .float8_e8m0 (Shape.mk [2])).toNpy with | .error _ => true | .ok _ => false end Test diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index bf449d0..313f224 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -712,6 +712,31 @@ private def testFloat8E3M4EdgeCases : IO Bool := do return checks.all id +-- E8M0 decode boundary values from OCP MX spec §5.4.1 +-- Exponent range: -127 to 127, bias=127, no inf, no zero, NaN=0xFF +private def testFloat8E8M0EdgeCases : IO Bool := do + let mut checks : List Bool := [] + + -- Min value: byte 0 = 2^(-127) + let v <- IO.ofExcept (Dtype.decodeFloat8E8M0 (ByteArray.mk #[0])) + let pass := v == Float32.ofBits 0x00400000 + IO.println s!"fp8_e8m0 byte 0 (2^-127): {pass}" + checks := pass :: checks + + -- Identity: byte 127 = 2^0 = 1.0 + let v <- IO.ofExcept (Dtype.decodeFloat8E8M0 (ByteArray.mk #[127])) + let pass := v == 1.0 + IO.println s!"fp8_e8m0 byte 127 (1.0): {pass}" + checks := pass :: checks + + -- Max value: byte 254 = 2^127 + let v <- IO.ofExcept (Dtype.decodeFloat8E8M0 (ByteArray.mk #[254])) + let pass := v == Float32.ofBits 0x7F000000 + IO.println s!"fp8_e8m0 byte 254 (2^127): {pass}" + checks := pass :: checks + + return checks.all id + def runAllTests : IO Bool := do return (<- testTensorElementBV Dtype.uint16) && (<- testTensorElementBV Dtype.uint32) && @@ -719,7 +744,8 @@ def runAllTests : IO Bool := do (<- testBFloat16EdgeCases) && (<- testFloat8E4M3EdgeCases) && (<- testFloat8E5M2EdgeCases) && - (<- testFloat8E3M4EdgeCases) + (<- testFloat8E3M4EdgeCases) && + (<- testFloat8E8M0EdgeCases) end Test end TensorLib From 16fdd363f75e88a0ade776598787f9e81a611f20 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Thu, 6 Aug 2026 16:18:34 -0700 Subject: [PATCH 2/9] Added quantize, dequantize, guards, and test cases --- TensorLib/Dtype.lean | 13 +++- TensorLib/Tensor.lean | 175 ++++++++++++++++++++++++++++++++++++++++++ TensorLib/Test.lean | 74 +++++++++++++++++- 3 files changed, 260 insertions(+), 2 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 4eb9170..9964958 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -357,6 +357,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 @@ -441,7 +451,8 @@ def decodeFloat8E8M0 (arr : ByteArray) : Err Float32 := else .ok (arr.data[0]!.toFloat32FromFloat8E8M0) -- Dispatch fp8 decode by dtype -private def decodeFloat8 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with +-- change from private since I need to call this in anotehr file for quantizing +def decodeFloat8 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with | .float8_e4m3 => decodeFloat8E4M3 arr | .float8_e5m2 => decodeFloat8E5M2 arr | .float8_e3m4 => decodeFloat8E3M4 arr diff --git a/TensorLib/Tensor.lean b/TensorLib/Tensor.lean index 15aa23e..e16beb8 100644 --- a/TensorLib/Tensor.lean +++ b/TensorLib/Tensor.lean @@ -695,6 +695,179 @@ def toNpy (arr : Tensor) : Err Npy.Ndarray := let startIndex := 0 .ok { header, data, startIndex } +-- Dequantize an MX scaled tensor: v_i = decodeE8M0(scale) x fp32(qW_i), one scale per group +-- reconstructs the original fp32 values from a block scaled quantized tensor by multiplying each elemt by its group's decoded E8M0 scale +def dequantizeMX (qW : Tensor) (scales : Tensor) (groupSize: Nat) : Err Tensor := do + -- scales tensor must be E8M0 (the only MX scale format supported by tensorlib) + if scales.dtype != .float8_e8m0 then .error "dequantizeMX: scales must have dtype float8_e8m0" + else + -- dequantization is defined along the last dimension so we need atleast 1 + let lastDim <- match qW.shape.val.getLast? with + | none => .error "dequantizeMX: qW must have atleast one dimension" + | some d => .ok d + -- each group of groupSize elements shares one scale byte + if lastDim % groupSize != 0 then + .error "dequantizeMX: groupSize must divide the last dimension of qW" + else + -- expected scales shape same as qW but last dimension is divided by groupSize + let expectedScalesShape := TensorLib.Shape.mk (qW.shape.val.dropLast ++ [lastDim / groupSize]) + if scales.shape != expectedScalesShape then + .error "dequantizeMX: Scales shape does not match qW shape / groupSize" + else + -- flatten both tensors to lists of raw bytes, one byteArray per element + let qWElems <- qW.toList + let scElems <- scales.toList + -- split qW elements into consecutive groups of groupSize + -- each group corresponds to one scale byte in scElems + let groups := List.toChunks groupSize qWElems + -- zip each group of qW elements with its corresponding scale byte + -- for each pair: decode scale, decode each element, multiply (v_i = X * P_i) + let resultGroups <- (groups.zip scElems).mapM fun (group, scaleBytes) => do + -- decode E8M0 scale byte to fp32: X = 2^(byte - 127), 0xFF -> NaN + let X <- Dtype.decodeFloat8E8M0 scaleBytes + -- for each element in the group, decode to fp32 and multiply by scale + group.mapM fun elemBytes => do + let p <- Dtype.byteArrayToFloat32 qW.dtype elemBytes + -- OCP MX spec §5.1: v_i = X * P_i + let v := X * p + -- pack result back as fp32 bytes + Dtype.byteArrayOfFloat32 .float32 v + -- flatten result groups back to a single list of byte array + let flatElems := resultGroups.flatten + -- concatenate all bytes into a single byte array + let data := flatElems.foldl (fun acc bytes => acc.append bytes) (ByteArray.emptyWithCapacity (flatElems.length * Dtype.float32.itemsize)) + -- fp32 tensor with same shape as qW + return {dtype := .float32, shape := qW.shape, data := data} + +-- guards for tricky cases for dequantize +-- case 1: fractional scale (byte 126 = 0.5), groupSize 1 +-- 3.0 * 0.5 = 1.5, 5.0 * 0.5 = 2.5 +#guard + let qW := Tensor.ofFloat32List! [3.0, 5.0] + let scales := { dtype := .float8_e8m0, shape := TensorLib.Shape.mk [2], data := ByteArray.mk #[126, 126] : Tensor } + match Tensor.dequantizeMX qW scales 1 with + | .error _ => false + | .ok result => result.toFloat32Tree! == .root [1.5, 2.5] + +-- case 2: different scales per group +-- group 1: byte 128 = 2.0, so [2.0, 4.0] -> [4.0, 8.0] +-- group 2: byte 126 = 0.5, so [6.0, 8.0] -> [3.0, 4.0] +#guard + let qW := Tensor.ofFloat32List! [2.0, 4.0, 6.0, 8.0] + let scales := { dtype := .float8_e8m0, shape := TensorLib.Shape.mk [2], data := ByteArray.mk #[128, 126] : Tensor } + match Tensor.dequantizeMX qW scales 2 with + | .error _ => false + | .ok result => result.toFloat32Tree! == .root [4.0, 8.0, 3.0, 4.0] + +-- case 3: negative values, sign must be preserved +-- byte 128 = 2.0, so [-2.0, -4.0] -> [-4.0, -8.0] +#guard + let qW := Tensor.ofFloat32List! [-2.0, -4.0] + let scales := { dtype := .float8_e8m0, shape := TensorLib.Shape.mk [1], data := ByteArray.mk #[128] : Tensor } + match Tensor.dequantizeMX qW scales 2 with + | .error _ => false + | .ok result => result.toFloat32Tree! == .root [-4.0, -8.0] + +-- case 4: scale byte 0 = 2^-127 (smallest E8M0, fp32 subnormal) +-- 1.0 * 2^-127 = fp32 subnormal 0x00400000 +#guard + let qW := Tensor.ofFloat32List! [1.0] + let scales := { dtype := .float8_e8m0, shape := TensorLib.Shape.mk [1], data := ByteArray.mk #[0] : Tensor } + match Tensor.dequantizeMX qW scales 1 with + | .error _ => false + | .ok result => result.toFloat32Tree! == .root [Float32.ofBits 0x00400000] + +-- case 5: scale byte 254 = 2^127 (largest E8M0 value) +-- 1.0 * 2^127 = Float32.ofBits 0x7F000000 +#guard + let qW := Tensor.ofFloat32List! [1.0] + let scales := { dtype := .float8_e8m0, shape := TensorLib.Shape.mk [1], data := ByteArray.mk #[254] : Tensor } + match Tensor.dequantizeMX qW scales 1 with + | .error _ => false + | .ok result => result.toFloat32Tree! == .root [Float32.ofBits 0x7F000000] + +-- case 6: mixed NaN and non-NaN groups +-- group 1: byte 127 = 1.0, so [2.0, 4.0] -> [2.0, 4.0] +-- group 2: byte 255 = NaN, so [6.0, 8.0] -> [NaN, NaN] per OCP 5.1 +#guard + let qW := Tensor.ofFloat32List! [2.0, 4.0, 6.0, 8.0] + let scales := { dtype := .float8_e8m0, shape := TensorLib.Shape.mk [2], data := ByteArray.mk #[127, 255] : Tensor } + match Tensor.dequantizeMX qW scales 2 with + | .error _ => false + | .ok result => match result.toFloat32Tree! with + | .root [a, b, c, d] => a == 2.0 && b == 4.0 && c.isNaN && d.isNaN + | _ => false + +-- Quantize a fp32 tensor to MX format using NVIDIA's scale computation: +-- m = floor_pw2(fp8Max / amax), scale byte = floor(log2(1/m)) + 127 +-- Reference: NVIDIA TensorEngine +-- Returns (qW, scales) where qW is fp32 scaled values and scales is e8m0 byte tensor +def quantizeMX (x : Tensor) (groupSize : Nat) (computeDtype : Dtype) : Err (Tensor × Tensor) := do + -- lookup fp8Max for the compute dtype -- returns none for non fp8 dtypes + let fp8Max <- match Dtype.fp8Max computeDtype with + | none => .error s!"quantizeMX: unsupported compute dtype {computeDtype}" + | some v => .ok v + -- x must be fp32 + if x.dtype != .float32 then .error "quantizeMX: inpute tensor must be float32" + else + -- last dim must divide evenly by groupsize + let lastDim <- match x.shape.val.getLast? with + | none => .error "quantizeMX: input tensor must have >= 1 dimension" + | some d => .ok d + if lastDim % groupSize != 0 then .error "quantizeMX: groupSize must divide the last dimension of x" else + -- flatten x to a list of raw bytes, one ByteArray per element + let xElems <- x.toList + -- split into consecutive groups of groupSize along the last dim + let groups := List.toChunks groupSize xElems + -- for each group, compute the E8M0 scale byte using NVIDIA's formula: + -- m = floor_pow2(fp8Max / amax), scale byte = floor(log2(1/m)) + 127 + let results <- groups.mapM fun group => do + -- decode each element to Float32 + let vals <- group.mapM (Dtype.byteArrayToFloat32 .float32) + -- amax = max absolute value in the group + let amax := vals.foldl (fun acc v => + let absV := if v < 0.0 then -v else v + if absV > acc then absV else acc) 0.0 + -- compute scale byte + let scaleByte : UInt8 := + if amax == 0.0 then + -- zero group: scale = 1.0, no-op + 127 + else if amax.isInf || amax.isNaN then + -- inf or NaN group: encode as NaN scale + 255 + else + -- m = floor_pow2(fp8Max / amax) + -- floor(log2(x)) via Float32.log2 and Float32.floor + let logM := (fp8Max / amax).log2.floor + -- scale byte encodes 1/m: s = -floor(log2(m)) + 127 + let s := (-logM + 127.0) + s.toUInt8 + -- scale qW elements: qW_i = x_i * m = x_i * 2^(-logM) + -- scale qW elements: qW_i = x_i * m = x_i * 2^(logM) + -- use m=1.0 for zero/inf/NaN groups to avoid dividing by zero + let scaledVals <- group.mapM fun elemBytes => do + let v <- Dtype.byteArrayToFloat32 .float32 elemBytes + let m := if amax == 0.0 || amax.isInf || amax.isNaN then 1.0 + else Float32.pow 2.0 (fp8Max / amax).log2.floor + Dtype.byteArrayOfFloat32 .float32 (v * m) + return (scaledVals, scaleByte) + -- separate scaled values and bytes from results + let scaledGroups := results.map Prod.fst + let scaleBytes := results.map Prod.snd + -- flatten scaled groups into a single ByteArray for qW + let flatScaled := scaledGroups.flatten + let qwData := flatScaled.foldl (fun acc bytes => acc.append bytes) + (ByteArray.emptyWithCapacity (flatScaled.length * Dtype.float32.itemsize)) + -- pack scale bytes into a ByteArray for scales tensor + let scaleData := ByteArray.mk (scaleBytes.toArray) + -- qW has same shape as x, scales has last dim divided by groupSize + let scalesShape := TensorLib.Shape.mk (x.shape.val.dropLast ++ [x.shape.val.getLast?.getD 0 / groupSize]) + return ( + { dtype := .float32, shape := x.shape, data := qwData }, + { dtype := .float8_e8m0, shape := scalesShape, data := scaleData } + ) + section Test open TensorLib.Tensor.Format.Tree @@ -777,6 +950,8 @@ open TensorLib.Tensor.Format.Tree #guard match (Tensor.zeros .float8_e4m3 (Shape.mk [2])).toNpy with | .ok _ => true | .error _ => false #guard match (Tensor.zeros .float8_e8m0 (Shape.mk [2])).toNpy with | .error _ => true | .ok _ => false + + end Test end Tensor diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 313f224..6753c56 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -737,6 +737,76 @@ private def testFloat8E8M0EdgeCases : IO Bool := do return checks.all id +-- dequantizeMX: OCP MX spec 5.1 v_i = X * P_i +-- Tests three cases: identity scale, doubling scale, NaN scale +private def testDequantizeMX : IO Bool := do + let mut checks : List Bool := [] + + -- scale byte 127 = 2^0 = 1.0: output should equal input + let qW <- IO.ofExcept (Tensor.ofFloat32List [2.0, 4.0, 6.0, 8.0]) + let scaleData := ByteArray.mk #[127, 127] -- two groups of 2, scale = 1.0 + let scales := { dtype := .float8_e8m0, shape := TensorLib.Shape.mk [2], data := scaleData : Tensor } + let result <- IO.ofExcept (Tensor.dequantizeMX qW scales 2) + let tree <- IO.ofExcept result.toFloat32Tree + let pass := tree == .root [2.0, 4.0, 6.0, 8.0] + IO.println s!"dequantizeMX identity scale (byte 127 = 1.0): {pass}" + checks := pass :: checks + + -- scale byte 128 = 2^1 = 2.0: output should be input * 2.0 + let scaleData2 := ByteArray.mk #[128, 128] -- two groups of 2, scale = 2.0 + let scales2 := { dtype := .float8_e8m0, shape := TensorLib.Shape.mk [2], data := scaleData2 : Tensor } + let result2 <- IO.ofExcept (Tensor.dequantizeMX qW scales2 2) + let tree2 <- IO.ofExcept result2.toFloat32Tree + let pass2 := tree2 == .root [4.0, 8.0, 12.0, 16.0] + IO.println s!"dequantizeMX doubling scale (byte 128 = 2.0): {pass2}" + checks := pass2 :: checks + + -- scale byte 255 = NaN: output should be all NaN + let scaleData3 := ByteArray.mk #[255, 255] + let scales3 := { dtype := .float8_e8m0, shape := TensorLib.Shape.mk [2], data := scaleData3 : Tensor } + let result3 <- IO.ofExcept (Tensor.dequantizeMX qW scales3 2) + let tree3 <- IO.ofExcept result3.toFloat32Tree + let pass3 := match tree3 with + | .root vs => vs.all Float32.isNaN + | _ => false + IO.println s!"dequantizeMX NaN scale (byte 255): {pass3}" + checks := pass3 :: checks + + return checks.all id + +-- quantizeMX: NVIDIA scale computation +-- Tests scale byte computation and value scaling per group +private def testQuantizeMX : IO Bool := do + let mut checks : List Bool := [] + + -- two groups with different scales + -- group 1: [100, 200], amax=200, m=floor_pow2(448/200)=2.0, scale byte=126, scaled=[200, 400] + -- group 2: [300, 800], amax=800, m=floor_pow2(448/800)=0.5, scale byte=128, scaled=[150, 400] + let x <- IO.ofExcept (Tensor.ofFloat32List [100.0, 200.0, 300.0, 800.0]) + let (qW, scales) <- IO.ofExcept (Tensor.quantizeMX x 2 .float8_e4m3) + let scalesOk := scales.data == ByteArray.mk #[126, 128] + let qwOk := qW.toFloat32Tree! == .root [200.0, 400.0, 150.0, 400.0] + let pass := scalesOk && qwOk + IO.println s!"quantizeMX two groups different scales: {pass}" + checks := pass :: checks + + -- zero group: amax=0, scale byte=127 (no-op), scaled values unchanged + let x2 <- IO.ofExcept (Tensor.ofFloat32List [0.0, 0.0]) + let (qW2, scales2) <- IO.ofExcept (Tensor.quantizeMX x2 2 .float8_e4m3) + let pass2 := scales2.data == ByteArray.mk #[127] && qW2.toFloat32Tree! == .root [0.0, 0.0] + IO.println s!"quantizeMX zero group (scale byte 127): {pass2}" + checks := pass2 :: checks + + -- negative values: sign preserved after scaling + -- group: [-100, -200], amax=200, m=2.0, scale byte=126, scaled=[-200, -400] + let x3 <- IO.ofExcept (Tensor.ofFloat32List [-100.0, -200.0]) + let (qW3, scales3) <- IO.ofExcept (Tensor.quantizeMX x3 2 .float8_e4m3) + let pass3 := scales3.data == ByteArray.mk #[126] && qW3.toFloat32Tree! == .root [-200.0, -400.0] + IO.println s!"quantizeMX negative values sign preserved: {pass3}" + checks := pass3 :: checks + + return checks.all id + def runAllTests : IO Bool := do return (<- testTensorElementBV Dtype.uint16) && (<- testTensorElementBV Dtype.uint32) && @@ -745,7 +815,9 @@ def runAllTests : IO Bool := do (<- testFloat8E4M3EdgeCases) && (<- testFloat8E5M2EdgeCases) && (<- testFloat8E3M4EdgeCases) && - (<- testFloat8E8M0EdgeCases) + (<- testFloat8E8M0EdgeCases) && + (<- testDequantizeMX) && + (<- testQuantizeMX) end Test end TensorLib From 4d627961304e128bd0d2b97a7e896ea0f4bb5cbc Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Thu, 6 Aug 2026 16:33:55 -0700 Subject: [PATCH 3/9] PR review: deleted encoder comment, byte 0 comment, and added NaN guard --- TensorLib/Float.lean | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index f497238..f029c27 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -654,7 +654,7 @@ def _root_.Float32.toFloat8E3M4Bits (f : Float32) : UInt8 := -- Negative overflow #guard (Float32.ofBits 0xC1800000).toFloat8E3M4Bits == (240 : UInt8) -- -16.0 -> -inf --- Encoder for fp8_e8m0 (scale type) + -- 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 @@ -679,7 +679,8 @@ def _root_.UInt8.toFloat32FromFloat8E8M0 (bits: UInt8) : Float32 := #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: fp32 exp=0, mant=0 = +0 (not 2^-127) +#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 From b37b9f41d3d9b2132802fd70791023b981eff629 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Fri, 7 Aug 2026 11:17:15 -0700 Subject: [PATCH 4/9] Add roundToComputeDtype + quantizeMX/dequantizeMX with tests and comments --- TensorLib/Dtype.lean | 20 +++++++++++++++++ TensorLib/Test.lean | 51 +++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 9964958..c1cb702 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -1087,6 +1087,25 @@ 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]) + -- 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 @@ -1225,6 +1244,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 diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 6753c56..b7e7b4c 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -807,6 +807,54 @@ private def testQuantizeMX : IO Bool := do return checks.all id +-- roundToComputeDtype: encode to compute dtype then decode back to fp32 +-- Exactly representable values: no rounding error (these happen to land on the format's grid) +-- Non-representable values: rounding error introduced by quantization +private def testRoundToComputeDtype : IO Bool := do + let mut checks : List Bool := [] + + -- exactly representable value: no rounding + let v <- IO.ofExcept (Dtype.roundToComputeDtype 3.0 .float8_e4m3) + let pass := v == 3.0 + IO.println s!"roundToComputeDtype 3.0 e4m3 (exact): {pass}" + checks := pass :: checks + + -- 3.1 rounds down to 3.0 (midpoint is 3.125) + let v <- IO.ofExcept (Dtype.roundToComputeDtype 3.1 .float8_e4m3) + let pass := v == 3.0 + IO.println s!"roundToComputeDtype 3.1 e4m3 (rounds to 3.0): {pass}" + checks := pass :: checks + + -- 500.0 overflows to NaN in e4m3 (e4m3fn has no infinity, overflow = NaN) + let v <- IO.ofExcept (Dtype.roundToComputeDtype 500.0 .float8_e4m3) + let pass := v != v -- IEEE NaN property: NaN != NaN + IO.println s!"roundToComputeDtype 500.0 e4m3 (overflows to NaN): {pass}" + + -- 2.3 rounds up to 2.5 in E5M2 (representable values: 2.0, 2.5; midpoint=2.25, 2.3 > 2.25) + let v <- IO.ofExcept (Dtype.roundToComputeDtype 2.3 .float8_e5m2) + let pass := v == 2.5 + IO.println s!"roundToComputeDtype 2.3 e5m2 (rounds to 2.5): {pass}" + + -- zero is always exactly representable + let v <- IO.ofExcept (Dtype.roundToComputeDtype 0.0 .float8_e4m3) + let pass := v == 0.0 + IO.println s!"roundToComputeDtype 0.0 e4m3 (exact zero): {pass}" + checks := pass :: checks + + -- sign preserved + let v <- IO.ofExcept (Dtype.roundToComputeDtype (-3.0) .float8_e4m3) + let pass := v == -3.0 + IO.println s!"roundToComputeDtype -3.0 e4m3 (sign preserved): {pass}" + checks := pass :: checks + + -- fp16 round-trip for exactly representable value + let v <- IO.ofExcept (Dtype.roundToComputeDtype 1.5 .float16) + let pass := v == 1.5 + IO.println s!"roundToComputeDtype 1.5 float16 (exact): {pass}" + checks := pass :: checks + + return checks.all id + def runAllTests : IO Bool := do return (<- testTensorElementBV Dtype.uint16) && (<- testTensorElementBV Dtype.uint32) && @@ -817,7 +865,8 @@ def runAllTests : IO Bool := do (<- testFloat8E3M4EdgeCases) && (<- testFloat8E8M0EdgeCases) && (<- testDequantizeMX) && - (<- testQuantizeMX) + (<- testQuantizeMX) && + (<- testRoundToComputeDtype) end Test end TensorLib From 15ace7e3defc646974bc120588bf3a014cfbe8c9 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Fri, 7 Aug 2026 12:27:46 -0700 Subject: [PATCH 5/9] PR fix: qW dtype check, groupSize=0 guard, revert decodeFloat8 to private, and fix comments --- TensorLib/Dtype.lean | 14 ++++++++------ TensorLib/Tensor.lean | 29 ++++++++++++++++++----------- TensorLib/Test.lean | 2 ++ 3 files changed, 28 insertions(+), 17 deletions(-) diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index c1cb702..410b31b 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -247,6 +247,7 @@ def join (x y : Dtype) : Option Dtype := -- 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 @@ -451,8 +452,7 @@ def decodeFloat8E8M0 (arr : ByteArray) : Err Float32 := else .ok (arr.data[0]!.toFloat32FromFloat8E8M0) -- Dispatch fp8 decode by dtype --- change from private since I need to call this in anotehr file for quantizing -def decodeFloat8 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with +private def decodeFloat8 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with | .float8_e4m3 => decodeFloat8E4M3 arr | .float8_e5m2 => decodeFloat8E5M2 arr | .float8_e3m4 => decodeFloat8E3M4 arr @@ -1096,6 +1096,7 @@ def roundToComputeDtype (v : Float32) (dtype : Dtype) : Err Float32 := match dty | .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 @@ -1367,7 +1368,7 @@ private def canCastLosslessRoundTrip (fromDtype : Dtype) (data : ByteArray) (toD | .error _ => false private def canCastLosslessIntRoundTrip (fromDtype : Dtype) (n : Int) (toDtype : Dtype) : Bool := - if fromDtype == .float8_e8m0 || toDtype == .float8_e8m0 then true + if fromDtype == .float8_e8m0 || toDtype == .float8_e8m0 then false else let res := do let n <- fromDtype.byteArrayOfInt n @@ -1438,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 /-- @@ -1451,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 diff --git a/TensorLib/Tensor.lean b/TensorLib/Tensor.lean index e16beb8..e919795 100644 --- a/TensorLib/Tensor.lean +++ b/TensorLib/Tensor.lean @@ -699,14 +699,16 @@ def toNpy (arr : Tensor) : Err Npy.Ndarray := -- reconstructs the original fp32 values from a block scaled quantized tensor by multiplying each elemt by its group's decoded E8M0 scale def dequantizeMX (qW : Tensor) (scales : Tensor) (groupSize: Nat) : Err Tensor := do -- scales tensor must be E8M0 (the only MX scale format supported by tensorlib) - if scales.dtype != .float8_e8m0 then .error "dequantizeMX: scales must have dtype float8_e8m0" + if qW.dtype != .float32 then .error "dequantizeMX: qW must have dtype float32" + else if scales.dtype != .float8_e8m0 then .error "dequantizeMX: scales must have dtype float8_e8m0" else -- dequantization is defined along the last dimension so we need atleast 1 let lastDim <- match qW.shape.val.getLast? with | none => .error "dequantizeMX: qW must have atleast one dimension" | some d => .ok d -- each group of groupSize elements shares one scale byte - if lastDim % groupSize != 0 then + if groupSize == 0 then .error "dequantizeMX: groupSize must be positive" + else if lastDim % groupSize != 0 then .error "dequantizeMX: groupSize must divide the last dimension of qW" else -- expected scales shape same as qW but last dimension is divided by groupSize @@ -814,7 +816,8 @@ def quantizeMX (x : Tensor) (groupSize : Nat) (computeDtype : Dtype) : Err (Tens let lastDim <- match x.shape.val.getLast? with | none => .error "quantizeMX: input tensor must have >= 1 dimension" | some d => .ok d - if lastDim % groupSize != 0 then .error "quantizeMX: groupSize must divide the last dimension of x" else + if groupSize == 0 then .error "quantizeMX: groupSize must be positive" + else if lastDim % groupSize != 0 then .error "quantizeMX: groupSize must divide the last dimension of x" else -- flatten x to a list of raw bytes, one ByteArray per element let xElems <- x.toList -- split into consecutive groups of groupSize along the last dim @@ -837,19 +840,23 @@ def quantizeMX (x : Tensor) (groupSize : Nat) (computeDtype : Dtype) : Err (Tens -- inf or NaN group: encode as NaN scale 255 else - -- m = floor_pow2(fp8Max / amax) - -- floor(log2(x)) via Float32.log2 and Float32.floor - let logM := (fp8Max / amax).log2.floor - -- scale byte encodes 1/m: s = -floor(log2(m)) + 127 - let s := (-logM + 127.0) - s.toUInt8 - -- scale qW elements: qW_i = x_i * m = x_i * 2^(-logM) + let ratio := fp8Max / amax + if ratio.isInf then 254 -- amax too small so we use max scale byte + else + let logM := ratio.log2.floor + let s := (-logM + 127.0) + if s < 0.0 then 0 + else if s > 254.0 then 254 + else s.toUInt8 -- scale qW elements: qW_i = x_i * m = x_i * 2^(logM) -- use m=1.0 for zero/inf/NaN groups to avoid dividing by zero let scaledVals <- group.mapM fun elemBytes => do let v <- Dtype.byteArrayToFloat32 .float32 elemBytes let m := if amax == 0.0 || amax.isInf || amax.isNaN then 1.0 - else Float32.pow 2.0 (fp8Max / amax).log2.floor + else + let ratio := fp8Max / amax + if ratio.isInf then Float32.ofBits 0x7F000000 -- 2 ^ 127 + else Float32.pow 2.0 ratio.log2.floor Dtype.byteArrayOfFloat32 .float32 (v * m) return (scaledVals, scaleByte) -- separate scaled values and bytes from results diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index b7e7b4c..8fe8ac2 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -829,11 +829,13 @@ private def testRoundToComputeDtype : IO Bool := do let v <- IO.ofExcept (Dtype.roundToComputeDtype 500.0 .float8_e4m3) let pass := v != v -- IEEE NaN property: NaN != NaN IO.println s!"roundToComputeDtype 500.0 e4m3 (overflows to NaN): {pass}" + checks := pass :: checks -- 2.3 rounds up to 2.5 in E5M2 (representable values: 2.0, 2.5; midpoint=2.25, 2.3 > 2.25) let v <- IO.ofExcept (Dtype.roundToComputeDtype 2.3 .float8_e5m2) let pass := v == 2.5 IO.println s!"roundToComputeDtype 2.3 e5m2 (rounds to 2.5): {pass}" + checks := pass :: checks -- zero is always exactly representable let v <- IO.ofExcept (Dtype.roundToComputeDtype 0.0 .float8_e4m3) From 85a592e7503e121f24e9ce52007c06b57b51e3c1 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Fri, 7 Aug 2026 13:21:14 -0700 Subject: [PATCH 6/9] PR fix: use Fp32.abs, remove double decode/ratio computation, use lastDim directly --- TensorLib/Tensor.lean | 45 +++++++++++++++---------------------------- 1 file changed, 15 insertions(+), 30 deletions(-) diff --git a/TensorLib/Tensor.lean b/TensorLib/Tensor.lean index e919795..e5b52d1 100644 --- a/TensorLib/Tensor.lean +++ b/TensorLib/Tensor.lean @@ -828,36 +828,21 @@ def quantizeMX (x : Tensor) (groupSize : Nat) (computeDtype : Dtype) : Err (Tens -- decode each element to Float32 let vals <- group.mapM (Dtype.byteArrayToFloat32 .float32) -- amax = max absolute value in the group - let amax := vals.foldl (fun acc v => - let absV := if v < 0.0 then -v else v - if absV > acc then absV else acc) 0.0 - -- compute scale byte - let scaleByte : UInt8 := - if amax == 0.0 then - -- zero group: scale = 1.0, no-op - 127 - else if amax.isInf || amax.isNaN then - -- inf or NaN group: encode as NaN scale - 255 + let amax := vals.foldl (fun acc v => if v.abs > acc then v.abs else acc) 0.0 + -- compute scale byte and multiplier together (avoids recomputing ratio/logM) + let ratio := fp8Max / amax + let (scaleByte, m) : UInt8 × Float32 := + if amax == 0.0 then (127, 1.0) + else if amax.isInf || amax.isNaN then (255, 1.0) + else if ratio.isInf then (254, Float32.ofBits 0x7F000000) else - let ratio := fp8Max / amax - if ratio.isInf then 254 -- amax too small so we use max scale byte - else - let logM := ratio.log2.floor - let s := (-logM + 127.0) - if s < 0.0 then 0 - else if s > 254.0 then 254 - else s.toUInt8 - -- scale qW elements: qW_i = x_i * m = x_i * 2^(logM) - -- use m=1.0 for zero/inf/NaN groups to avoid dividing by zero - let scaledVals <- group.mapM fun elemBytes => do - let v <- Dtype.byteArrayToFloat32 .float32 elemBytes - let m := if amax == 0.0 || amax.isInf || amax.isNaN then 1.0 - else - let ratio := fp8Max / amax - if ratio.isInf then Float32.ofBits 0x7F000000 -- 2 ^ 127 - else Float32.pow 2.0 ratio.log2.floor - Dtype.byteArrayOfFloat32 .float32 (v * m) + let logM := ratio.log2.floor + let s := (-logM + 127.0) + let byte := if s < 0.0 then 0 + else if s > 254.0 then 254 + else s.toUInt8 + (byte, Float32.pow 2.0 logM) + let scaledVals <- vals.mapM fun v => Dtype.byteArrayOfFloat32 .float32 (v * m) return (scaledVals, scaleByte) -- separate scaled values and bytes from results let scaledGroups := results.map Prod.fst @@ -869,7 +854,7 @@ def quantizeMX (x : Tensor) (groupSize : Nat) (computeDtype : Dtype) : Err (Tens -- pack scale bytes into a ByteArray for scales tensor let scaleData := ByteArray.mk (scaleBytes.toArray) -- qW has same shape as x, scales has last dim divided by groupSize - let scalesShape := TensorLib.Shape.mk (x.shape.val.dropLast ++ [x.shape.val.getLast?.getD 0 / groupSize]) + let scalesShape := TensorLib.Shape.mk (x.shape.val.dropLast ++ [lastDim / groupSize]) return ( { dtype := .float32, shape := x.shape, data := qwData }, { dtype := .float8_e8m0, shape := scalesShape, data := scaleData } From 249cd9726d066187ff0ba311803e7e8ba09f547b Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Fri, 7 Aug 2026 14:55:26 -0700 Subject: [PATCH 7/9] PR fix: scale byte fix for ratio.isInf case, add subnormal amax test --- TensorLib/Tensor.lean | 2 +- TensorLib/Test.lean | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/TensorLib/Tensor.lean b/TensorLib/Tensor.lean index e5b52d1..a7e9d91 100644 --- a/TensorLib/Tensor.lean +++ b/TensorLib/Tensor.lean @@ -834,7 +834,7 @@ def quantizeMX (x : Tensor) (groupSize : Nat) (computeDtype : Dtype) : Err (Tens let (scaleByte, m) : UInt8 × Float32 := if amax == 0.0 then (127, 1.0) else if amax.isInf || amax.isNaN then (255, 1.0) - else if ratio.isInf then (254, Float32.ofBits 0x7F000000) + else if ratio.isInf then (0, Float32.ofBits 0x7F000000) else let logM := ratio.log2.floor let s := (-logM + 127.0) diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 8fe8ac2..c806853 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -805,6 +805,13 @@ private def testQuantizeMX : IO Bool := do IO.println s!"quantizeMX negative values sign preserved: {pass3}" checks := pass3 :: checks + -- subnormal amax: ratio overflows to inf, should use max multiplier with scale byte 0 (= 2^-127) + let x4 <- IO.ofExcept (Tensor.ofFloat32List [Float32.ofBits 0x00000001]) -- smallest fp32 subnormal + let (_, scales4) <- IO.ofExcept (Tensor.quantizeMX x4 1 .float8_e4m3) + let pass := scales4.data == ByteArray.mk #[0] + IO.println s!"quantizeMX subnormal amax (scale byte 0 = 2^-127): {pass}" + checks := pass :: checks + return checks.all id -- roundToComputeDtype: encode to compute dtype then decode back to fp32 From 047063feeccf6d5f4d9e1f4cb3232d95cd972256 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Fri, 7 Aug 2026 16:13:09 -0700 Subject: [PATCH 8/9] PR fix: add NaN detection in quantizeMX (IEEE fold skips NaN, now pre-checked) and fix spelling error --- TensorLib/Tensor.lean | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/TensorLib/Tensor.lean b/TensorLib/Tensor.lean index a7e9d91..ba2a725 100644 --- a/TensorLib/Tensor.lean +++ b/TensorLib/Tensor.lean @@ -810,7 +810,7 @@ def quantizeMX (x : Tensor) (groupSize : Nat) (computeDtype : Dtype) : Err (Tens | none => .error s!"quantizeMX: unsupported compute dtype {computeDtype}" | some v => .ok v -- x must be fp32 - if x.dtype != .float32 then .error "quantizeMX: inpute tensor must be float32" + if x.dtype != .float32 then .error "quantizeMX: input tensor must be float32" else -- last dim must divide evenly by groupsize let lastDim <- match x.shape.val.getLast? with @@ -828,7 +828,9 @@ def quantizeMX (x : Tensor) (groupSize : Nat) (computeDtype : Dtype) : Err (Tens -- decode each element to Float32 let vals <- group.mapM (Dtype.byteArrayToFloat32 .float32) -- amax = max absolute value in the group - let amax := vals.foldl (fun acc v => if v.abs > acc then v.abs else acc) 0.0 + let hasNaN := vals.any (fun v => v != v) -- IEEE: NaN != NaN + let amax := if hasNaN then Float32.ofBits 0x7FC00000 -- NaN + else vals.foldl (fun acc v => if v.abs > acc then v.abs else acc) 0.0 -- compute scale byte and multiplier together (avoids recomputing ratio/logM) let ratio := fp8Max / amax let (scaleByte, m) : UInt8 × Float32 := From 75858a3b284324dd313ced4b3b7c08a7b039f451 Mon Sep 17 00:00:00 2001 From: SmoothThunk Date: Fri, 7 Aug 2026 16:51:03 -0700 Subject: [PATCH 9/9] PR fix: NaN group detection and test for NaN input producing scale byte 255 --- TensorLib/Test.lean | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index c806853..7f95213 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -812,6 +812,13 @@ private def testQuantizeMX : IO Bool := do IO.println s!"quantizeMX subnormal amax (scale byte 0 = 2^-127): {pass}" checks := pass :: checks + -- NaN input: group with NaN should emit scale byte 255 (OCP NaN sentinel) + let xNaN <- IO.ofExcept (Tensor.ofFloat32List [Float32.ofBits 0x7FC00000, 1.0]) + let (_, scalesNaN) <- IO.ofExcept (Tensor.quantizeMX xNaN 2 .float8_e4m3) + let passNaN := scalesNaN.data == ByteArray.mk #[255] + IO.println s!"quantizeMX NaN input group (scale byte 255): {passNaN}" + checks := passNaN :: checks + return checks.all id -- roundToComputeDtype: encode to compute dtype then decode back to fp32