|
| 1 | +// SPDX-Licence-Identifier: EUPL-1.2 |
| 2 | + |
| 3 | +// Package mxfp4 writes OCP Microscaling FP4 (MXFP4) blockwise weights: E2M1 |
| 4 | +// elements in 32-element blocks that share one E8M0 power-of-two scale. This |
| 5 | +// is the AMD-native low-precision lane AMD Quark emits for its MX/FP4 output |
| 6 | +// (see model/quant/quantfmt.go's MethodQuark, and model.QuantConfig's |
| 7 | +// "mxfp4" mode in model/quant_config.go, which already validates |
| 8 | +// group_size==32/bits==4 for this format). |
| 9 | +// |
| 10 | +// Bit layout is the OCP Microscaling Formats (MX) Specification v1.0 |
| 11 | +// (https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf): |
| 12 | +// each element is E2M1 (1 sign + 2 exponent + 1 mantissa bit, exponent bias |
| 13 | +// 1), giving the 8 non-negative magnitudes {0, 0.5, 1, 1.5, 2, 3, 4, 6}; each |
| 14 | +// block of 32 elements shares one E8M0 scale (an unsigned 8-bit exponent, |
| 15 | +// bias 127, 0xFF reserved — https://sw23.github.io/fp-conv/formats/fp4-e2m1.html). |
| 16 | +// Both the block size and the magnitude range are cross-checked against a |
| 17 | +// real AMD Quark MXFP4 checkpoint |
| 18 | +// (huggingface.co/amd/Qwen3.5-397B-A17B-MXFP4/blob/main/config.json: |
| 19 | +// quantization_config.global_quant_config.weight = {dtype: "fp4", qscheme: |
| 20 | +// "per_group", group_size: 32, scale_format: "e8m0"}). Neither E2M1 nor E8M0 |
| 21 | +// represent Inf/NaN — the format is finite-only. |
| 22 | +// |
| 23 | +// q, _ := mxfp4.Quantize(weights, []int{64}) |
| 24 | +// back, _ := mxfp4.Dequantize(q) |
| 25 | +package mxfp4 |
| 26 | + |
| 27 | +import ( |
| 28 | + "math" |
| 29 | + |
| 30 | + core "dappco.re/go" |
| 31 | +) |
| 32 | + |
| 33 | +// BlockSize is the number of elements that share one E8M0 scale — fixed at |
| 34 | +// 32 by the OCP MX spec, and confirmed both in model.QuantConfig's mxfp4 |
| 35 | +// validation (model/quant_config.go) and in AMD Quark's emitted group_size. |
| 36 | +const BlockSize = 32 |
| 37 | + |
| 38 | +// MaxE2M1 is the largest finite magnitude an E2M1 element can represent. |
| 39 | +const MaxE2M1 = float32(6) |
| 40 | + |
| 41 | +// magnitudes is the E2M1 non-negative value table (OCP MX spec v1.0): index |
| 42 | +// is the 3-bit unsigned code (2 exponent bits + 1 mantissa bit, bias 1). |
| 43 | +var magnitudes = [8]float32{0, 0.5, 1, 1.5, 2, 3, 4, 6} |
| 44 | + |
| 45 | +// Tensor is a blockwise-quantised MXFP4 weight: one E2M1 nibble per element, |
| 46 | +// two packed per byte (like nf4.Tensor), and one decoded E8M0 scale per |
| 47 | +// BlockSize-element block. |
| 48 | +type Tensor struct { |
| 49 | + Data []byte |
| 50 | + Scale []float32 |
| 51 | + Shape []int |
| 52 | +} |
| 53 | + |
| 54 | +// Quantize packs values into blockwise MXFP4. Each BlockSize-run gets its |
| 55 | +// own E8M0 power-of-two scale — the smallest power of two that keeps the |
| 56 | +// block's peak magnitude within E2M1's finite range (MaxE2M1) — and each |
| 57 | +// element rounds to the nearest E2M1 code. |
| 58 | +// |
| 59 | +// q, _ := mxfp4.Quantize([]float32{0, 1, 2, 3}, []int{4}) |
| 60 | +func Quantize(values []float32, shape []int) (Tensor, error) { |
| 61 | + if len(values) == 0 || product(shape) != len(values) { |
| 62 | + return Tensor{}, core.NewError("mxfp4: values must match a non-empty shape") |
| 63 | + } |
| 64 | + out := Tensor{ |
| 65 | + Data: make([]byte, (len(values)+1)/2), |
| 66 | + Scale: make([]float32, (len(values)+BlockSize-1)/BlockSize), |
| 67 | + Shape: append([]int(nil), shape...), |
| 68 | + } |
| 69 | + for block := range out.Scale { |
| 70 | + start, end := block*BlockSize, min((block+1)*BlockSize, len(values)) |
| 71 | + var peak float32 |
| 72 | + for _, value := range values[start:end] { |
| 73 | + if math.IsNaN(float64(value)) || math.IsInf(float64(value), 0) { |
| 74 | + return Tensor{}, core.NewError("mxfp4: values must be finite") |
| 75 | + } |
| 76 | + peak = max(peak, float32(math.Abs(float64(value)))) |
| 77 | + } |
| 78 | + scale := blockScale(peak) |
| 79 | + out.Scale[block] = scale |
| 80 | + for i := start; i < end; i++ { |
| 81 | + code := encodeE2M1(values[i] / scale) |
| 82 | + if i&1 == 0 { |
| 83 | + out.Data[i/2] = code << 4 |
| 84 | + } else { |
| 85 | + out.Data[i/2] |= code |
| 86 | + } |
| 87 | + } |
| 88 | + } |
| 89 | + return out, nil |
| 90 | +} |
| 91 | + |
| 92 | +// Dequantize expands a blockwise MXFP4 Tensor back to float32. |
| 93 | +// |
| 94 | +// values, _ := mxfp4.Dequantize(q) |
| 95 | +func Dequantize(tensor Tensor) ([]float32, error) { |
| 96 | + count := product(tensor.Shape) |
| 97 | + if count == 0 || len(tensor.Data) != (count+1)/2 || len(tensor.Scale) != (count+BlockSize-1)/BlockSize { |
| 98 | + return nil, core.NewError("mxfp4: invalid tensor") |
| 99 | + } |
| 100 | + out := make([]float32, count) |
| 101 | + for i := range out { |
| 102 | + code := tensor.Data[i/2] & 15 |
| 103 | + if i&1 == 0 { |
| 104 | + code = tensor.Data[i/2] >> 4 |
| 105 | + } |
| 106 | + out[i] = decodeE2M1(code) * tensor.Scale[i/BlockSize] |
| 107 | + } |
| 108 | + return out, nil |
| 109 | +} |
| 110 | + |
| 111 | +// blockScale picks the E8M0 power-of-two scale for a block given its peak |
| 112 | +// (largest absolute) value: the smallest power of two large enough that |
| 113 | +// peak/scale never exceeds MaxE2M1, so the block's largest element always |
| 114 | +// stays representable. peak==0 (an all-zero block) uses scale 1 rather than |
| 115 | +// dividing by zero. |
| 116 | +// |
| 117 | +// blockScale(6) == 1 |
| 118 | +func blockScale(peak float32) float32 { |
| 119 | + if peak == 0 { |
| 120 | + return 1 |
| 121 | + } |
| 122 | + exp := clampExponent(int(math.Ceil(math.Log2(float64(peak) / float64(MaxE2M1))))) |
| 123 | + return float32(math.Ldexp(1, exp)) |
| 124 | +} |
| 125 | + |
| 126 | +// clampExponent bounds an E8M0 exponent to its storable range. E8M0 is an |
| 127 | +// unsigned 8-bit field biased by 127 with 0xFF reserved (NaN) per the OCP MX |
| 128 | +// spec, so the representable signed exponent range is [-127, 127]. |
| 129 | +// |
| 130 | +// clampExponent(500) == 127 |
| 131 | +func clampExponent(exp int) int { |
| 132 | + switch { |
| 133 | + case exp < -127: |
| 134 | + return -127 |
| 135 | + case exp > 127: |
| 136 | + return 127 |
| 137 | + default: |
| 138 | + return exp |
| 139 | + } |
| 140 | +} |
| 141 | + |
| 142 | +// encodeE2M1 maps a value already divided by its block scale to the nearest |
| 143 | +// E2M1 code: bit 3 is the sign, bits 2:0 index magnitudes. Values outside |
| 144 | +// magnitudes' range clamp to the nearest (largest) representable magnitude |
| 145 | +// rather than overflowing. |
| 146 | +// |
| 147 | +// encodeE2M1(2) == 4 |
| 148 | +func encodeE2M1(value float32) byte { |
| 149 | + sign := byte(0) |
| 150 | + if value < 0 { |
| 151 | + sign, value = 0x8, -value |
| 152 | + } |
| 153 | + best, distance := byte(0), float32(math.MaxFloat32) |
| 154 | + for i, candidate := range magnitudes { |
| 155 | + if d := float32(math.Abs(float64(value - candidate))); d < distance { |
| 156 | + best, distance = byte(i), d |
| 157 | + } |
| 158 | + } |
| 159 | + return sign | best |
| 160 | +} |
| 161 | + |
| 162 | +// decodeE2M1 reverses encodeE2M1: bit 3 is the sign, bits 2:0 index |
| 163 | +// magnitudes. |
| 164 | +// |
| 165 | +// decodeE2M1(0xC) == -2 |
| 166 | +func decodeE2M1(code byte) float32 { |
| 167 | + value := magnitudes[code&0x7] |
| 168 | + if code&0x8 != 0 { |
| 169 | + return -value |
| 170 | + } |
| 171 | + return value |
| 172 | +} |
| 173 | + |
| 174 | +// product returns the element count a shape describes, 0 for an empty shape. |
| 175 | +// |
| 176 | +// product([]int{5, 13}) == 65 |
| 177 | +func product(shape []int) int { |
| 178 | + if len(shape) == 0 { |
| 179 | + return 0 |
| 180 | + } |
| 181 | + out := 1 |
| 182 | + for _, value := range shape { |
| 183 | + out *= value |
| 184 | + } |
| 185 | + return out |
| 186 | +} |
0 commit comments