1 # BendTorch's dtypes (plan M11): BF16 and F16 stored two per U32, int8 with a scale per row stored four per U32, a 2 # two-F32 ("double-word") accumulator, and the layers that give a model BF16 storage around F32 arithmetic. 3 # 4 # Every number is computed in F32. A 16-bit value lives in the low 16 bits of a U32 while it is handled alone, and two 5 # of them share a word in storage (element 2k in the low half of word k, element 2k + 1 in the high half: the bytes of 6 # a little-endian torch.bfloat16 / torch.float16 tensor). The conversions from F32 round to nearest, ties to even, and 7 # treat NaN as PyTorch 2.14.1's CPU kernels do on this machine (tested bit for bit by reference/test_dtype.py): 8 # - F32 -> BF16: every NaN becomes 0xFFFF (the vectorized kernel's answer; c10's scalar code would give 0x7FC0). 9 # - F32 -> F16: a NaN keeps its sign and the top 9 bits of its payload, with the quiet bit set (F16C's vcvtps2ph). 10 # - F16 -> F32: a NaN gets the quiet bit set; subnormals become normal F32 numbers, exactly. 11 # No @unsafe and no foreign code: plain Bend over U32 and F32 bit operations. 12 13 import Base 14 import ./main.bend as T 15 import bend-ml-tensor-array@0.1.14.0/main.bend as TA 16 17 # ===================================================================================================== 18 # Bits 19 # ===================================================================================================== 20 21 # U32 and F32 hold the same Word(32n) 22 def u2f(u: U32) -> F32: 23 match u: 24 case U32{w}: 25 F32{w} 26 27 def f2u(x: F32) -> U32: 28 match x: 29 case F32{w}: 30 U32{w} 31 32 def pick_u(b: Bool, +x: U32, +y: U32) -> U32: 33 match b: 34 case True{}: 35 x 36 case False{}: 37 y 38 39 # ===================================================================================================== 40 # BF16: the top 16 bits of an F32 41 # ===================================================================================================== 42 43 # F32 -> BF16, round to nearest even: add 0x7FFF plus the lowest kept bit, keep the top half. The carry rounds up into 44 # the exponent when the mantissa overflows, and past the largest finite value to infinity, as IEEE asks. 45 def bf16_rne(+u: U32) -> U32: 46 U32.shrn(((u + 32767 : U32) + (U32.shrn(u, 16n) .&. 1 : U32) : U32), 16n) 47 48 def bf16_of_bits(+u: U32) -> U32: 49 pick_u(U32.is_gt((u .&. 2147483647 : U32), 2139095040), 65535, bf16_rne(u)) 50 51 def BF16.of_f32(x: F32) -> U32: 52 bf16_of_bits(f2u(x)) 53 54 # BF16 -> F32 is exact: the 16 bits become the top half 55 def BF16.to_f32(+h: U32) -> F32: 56 u2f(U32.shln((h .&. 65535 : U32), 16n)) 57 58 # x rounded to the nearest BF16 value, as an F32 (x.to(torch.bfloat16).to(torch.float32)) 59 def BF16.round(+x: F32) -> F32: 60 BF16.to_f32(BF16.of_f32(x)) 61 62 # ===================================================================================================== 63 # F16: 1 sign bit, 5 exponent bits (bias 15), 10 mantissa bits 64 # ===================================================================================================== 65 66 # a normal F16 (|x| >= 2^-14): rebias the exponent (127 - 15 = 112), round the 23-bit mantissa to 10 bits to nearest 67 # even; past 65504 the carry reaches the infinity pattern 0x7C00 and anything above it is infinity too 68 def f16_normal(+a: U32) -> U32: 69 +m = (a - 939524096 : U32) 70 +r = U32.shrn(((m + 4095 : U32) + (U32.shrn(m, 13n) .&. 1 : U32) : U32), 13n) 71 pick_u(U32.is_gt(r, 31744), 31744, r) 72 73 # a subnormal F16 (2^-25 < |x| < 2^-14): the value in units of 2^-24 is mant / 2^sh, sh = 126 - e in 14..24, rounded to 74 # nearest even; a result of 0x400 is the smallest normal, which is the right answer 75 def f16_sub2(+q: U32, +rem: U32, +half: U32) -> U32: 76 +up = Bool.or(U32.is_gt(rem, half), Bool.and(U32.is_eq(rem, half), U32.is_eq((q .&. 1 : U32), 1))) 77 (q + Bool.to_u32(up) : U32) 78 79 def f16_sub(+a: U32) -> U32: 80 +mant = ((a .&. 8388607 : U32) .|. 8388608 : U32) 81 +sh = U32.to_nat((126 - U32.shrn(a, 23n) : U32)) 82 f16_sub2(U32.shrn(mant, sh), (mant .&. (U32.shln(1, sh) - 1 : U32) : U32), U32.shln(1, Nat.sub(sh, 1n))) 83 84 # a NaN keeps the top 9 bits of its payload and gets the quiet bit 85 def f16_nan(+a: U32) -> U32: 86 (32256 .|. U32.shrn((a .&. 8388607 : U32), 13n) : U32) 87 88 def f16_mag(nan: Bool, normal: Bool, tiny: Bool, +a: U32) -> U32: 89 match nan normal tiny: 90 case True{} _ _: 91 f16_nan(a) 92 case False{} True{} _: 93 f16_normal(a) 94 case False{} False{} True{}: 95 0 96 case False{} False{} False{}: 97 f16_sub(a) 98 99 # F32 -> F16, round to nearest even, with the sign kept (so -0.0 and the negative tiny numbers give 0x8000) 100 def f16_of_bits(+u: U32) -> U32: 101 +a = (u .&. 2147483647 : U32) 102 ((U32.shrn(u, 16n) .&. 32768 : U32) .|. f16_mag(U32.is_gt(a, 2139095040), U32.is_ge(a, 947912704), U32.is_le(a, 855638016), a) : U32) 103 104 def F16.of_f32(x: F32) -> U32: 105 f16_of_bits(f2u(x)) 106 107 # F16 -> F32, exact. A subnormal is m · 2^-24: the product of two F32 numbers that is exact and normal. 108 def f16_up(zero: Bool, top: Bool, +m: U32, +s: U32, +eu: U32) -> F32: 109 match zero top: 110 case True{} _: 111 u2f((s .|. f2u((U32.to_f32(m) * 0.000000059604644775390625 : F32)) : U32)) 112 case False{} True{}: 113 u2f((((s .|. 2139095040 : U32) .|. U32.shln(m, 13n) : U32) .|. pick_u(U32.is_eq(m, 0), 0, 4194304) : U32)) 114 case False{} False{}: 115 u2f(((s .|. U32.shln((eu + 112 : U32), 23n) : U32) .|. U32.shln(m, 13n) : U32)) 116 117 def F16.to_f32(+h: U32) -> F32: 118 +eu = (U32.shrn(h, 10n) .&. 31 : U32) 119 f16_up(U32.is_eq(eu, 0), U32.is_eq(eu, 31), (h .&. 1023 : U32), U32.shln((h .&. 32768 : U32), 16n), eu) 120 121 def F16.round(+x: F32) -> F32: 122 F16.to_f32(F16.of_f32(x)) 123 124 # ===================================================================================================== 125 # Two halves in a word 126 # ===================================================================================================== 127 128 def half_join(+lo: U32, +hi: U32) -> U32: 129 ((lo .&. 65535 : U32) .|. U32.shln(hi, 16n) : U32) 130 131 def half_lo(+w: U32) -> U32: 132 (w .&. 65535 : U32) 133 134 def half_hi(+w: U32) -> U32: 135 U32.shrn(w, 16n) 136 137 # a pair of 16-bit values, the unit of storage 138 type HW is Data: 139 HW{lo: U32, hi: U32} 140 141 # the values two by two; an odd last one is paired with 0 142 def pairs(xs: List<&2, U32>) -> List<&2, HW>: 143 match xs: 144 case Nil{}: 145 Nil{} 146 case Con{a, Nil{}}: 147 [HW{a, 0}] 148 case Con{a, Con{b, t}}: 149 HW{a, b} <> pairs(t) 150 151 # every value of a list of pairs, in order 152 def unpairs(ws: List<&2, HW>) -> List<&2, U32>: 153 match ws: 154 case Nil{}: 155 Nil{} 156 case Con{HW{a, b}, t}: 157 a <> (b <> unpairs(t)) 158 159 # the first n of them: the padding of an odd count is left out 160 def unpairs_n(+n: Nat, ws: List<&2, HW>) -> List<&2, U32>: 161 List.take(&2, U32, unpairs(ws), n) 162 163 # ===================================================================================================== 164 # Packed 16-bit matrices: r·c values in ceil(r·c / 2) words, element 2k in the low half of word k 165 # ===================================================================================================== 166 167 type HMat<-r: Nat, -c: Nat> is Type: 168 HMat{d: Array} 169 170 type UL is Type: 171 UL{a: Array, l: List<&2, U32>} 172 173 def fill_u(xs: List<&2, U32>, +i: U32, a: Array) -> Array: 174 match xs: 175 case Nil{}: 176 a 177 case Con{h, t}: 178 fill_u(t, (i + 1 : U32), Array.set(U32, a, i, h)) 179 180 # the k + 1 values from index i, in order (ra holds the array and the value at i, already read) 181 def read_u(k: Nat, +i: U32, acc: List<&2, U32>, ra: Array & U32) -> UL: 182 match k ra: 183 case 0n Tuple{a, v}: 184 UL{a, List.reverse(&2, U32, v <> acc)} 185 case 1n+p Tuple{a, +v}: 186 read_u(p, (i + 1 : U32), v <> acc, Array.get(U32, a, (i + 1 : U32))) 187 188 def ul_list(r: UL) -> List<&2, U32>: 189 match r: 190 case UL{a, l}: 191 l 192 193 def enc_l(~f: F32 -> U32, xs: List<&2, F32>) -> List<&2, U32>: 194 match xs: 195 case Nil{}: 196 Nil{} 197 case Con{h, t}: 198 f(h) <> enc_l(~f, t) 199 200 def dec_l(~f: @+h: U32 -> F32, xs: List<&2, U32>) -> List<&2, F32>: 201 match xs: 202 case Nil{}: 203 Nil{} 204 case Con{+h, t}: 205 f(h) <> dec_l(~f, t) 206 207 def join_l(ws: List<&2, HW>) -> List<&2, U32>: 208 match ws: 209 case Nil{}: 210 Nil{} 211 case Con{HW{+a, +b}, t}: 212 half_join(a, b) <> join_l(t) 213 214 def split_l(ws: List<&2, U32>) -> List<&2, HW>: 215 match ws: 216 case Nil{}: 217 Nil{} 218 case Con{+w, t}: 219 HW{half_lo(w), half_hi(w)} <> split_l(t) 220 221 # the number of words for n halves 222 def hwords(+n: Nat) -> Nat: 223 Nat.div(Nat.add(n, 1n), 2n) 224 225 def ml_l(-r: Nat, -c: Nat, m: TA.ML) -> List<&2, F32>: 226 match m: 227 case TA.ML{d, xs}: 228 xs 229 230 # the values of a 16-bit list packed two per word 231 def Half.words(+n: Nat, hs: List<&2, U32>) -> Array: 232 fill_u(join_l(pairs(hs)), 0, Array.new(U32, T.depth(hwords(n)), 0)) 233 234 # the first n 16-bit values of a packed array 235 def Half.values(+n: Nat, a: Array) -> List<&2, U32>: 236 unpairs_n(n, split_l(ul_list(read_u(Nat.sub(hwords(n), 1n), 0, Nil{}, Array.get(U32, a, 0))))) 237 238 # a matrix encoded by f (BF16.of_f32 or F16.of_f32) and packed 239 def Half.pack(~f: F32 -> U32, +r: Nat, +c: Nat, m: TA.Mat) -> HMat: 240 HMat{Half.words(Nat.mul(r, c), enc_l(~f, ml_l(r, c, TA.Mat.to_list(r, c, m))))} 241 242 # a packed matrix decoded by g (BF16.to_f32 or F16.to_f32) into F32 243 def Half.unpack(~g: @+h: U32 -> F32, +r: Nat, +c: Nat, h: HMat) -> TA.Mat: 244 match h: 245 case HMat{d}: 246 TA.Mat.from_list(r, c, dec_l(~g, Half.values(Nat.mul(r, c), d))) 247 248 def BF16.pack(+r: Nat, +c: Nat, m: TA.Mat) -> HMat: 249 Half.pack(~BF16.of_f32, r, c, m) 250 251 def BF16.unpack(+r: Nat, +c: Nat, h: HMat) -> TA.Mat: 252 Half.unpack(~BF16.to_f32, r, c, h) 253 254 def F16.pack(+r: Nat, +c: Nat, m: TA.Mat) -> HMat: 255 Half.pack(~F16.of_f32, r, c, m) 256 257 def F16.unpack(+r: Nat, +c: Nat, h: HMat) -> TA.Mat: 258 Half.unpack(~F16.to_f32, r, c, h) 259 260 # the words of a packed matrix, for printing and tests 261 def HMat.words(+r: Nat, +c: Nat, h: HMat) -> List<&2, U32>: 262 match h: 263 case HMat{d}: 264 ul_list(read_u(Nat.sub(hwords(Nat.mul(r, c)), 1n), 0, Nil{}, Array.get(U32, d, 0))) 265 266 # ===================================================================================================== 267 # Laws of the packing 268 # ===================================================================================================== 269 270 # LAW (unpair_pair): packing a list two by two and reading back as many values as it had gives the list back, in 271 # order, for any length (an odd last value is paired with a 0 that the read leaves out). Proof: by induction two 272 # values at a time; the pair HW{a, b} unpacks to a then b, and the rest is the induction hypothesis. 273 def unpair_pair(xs: List<&2, U32>) -> {List.take(&2, U32, unpairs(pairs(xs)), List.length(&2, U32, xs)) == xs : List<&2, U32>}: 274 match xs: 275 case Nil{}: 276 {==} 277 case Con{a, Nil{}}: 278 {==} 279 case Con{+a, Con{+b, +t}}: 280 Equal.cong(List<&2, U32>, List<&2, U32>, l => a <> (b <> l), List.take(&2, U32, unpairs(pairs(t)), List.length(&2, U32, t)), t, unpair_pair(t)) 281 282 # Not proved (debt, tested instead by reference/test_dtype.py): 283 # - half_lo(half_join(a, b)) == a and half_hi(half_join(a, b)) == b for 16-bit a and b: a fact about Word(32n)'s bit 284 # operations for every value, which the checker does not reach by computation (`dtype_test joins` checks all 65536 285 # values of each half against a second one walking every value); 286 # - that fill_u and read_u give back what was written (Array.set then Array.get), and the int8 words' byte order; 287 # - the rounding itself: F32 is opaque to the checker, so BF16, F16 and int8 are tested bit for bit against PyTorch. 288 289 # ===================================================================================================== 290 # int8 with a scale per row, for inference 291 # ===================================================================================================== 292 # Row i of an r x c weight matrix W gets the scale s_i = max_j |W[i, j]| / 127 and the integers 293 # q[i, j] = round(W[i, j] / s_i), to nearest even, in -127..127 (0 for a row of zeros). q is stored four per word, 294 # element 4k + b in byte b of word k, as two's complement bytes: the bytes of a torch.int8 tensor. A product with F32 295 # activations reads q straight from the words: y[n, i] = s_i · Σ_j x[n, j] · q[i, j], summed in order of j. 296 297 type Q8<-r: Nat, -c: Nat> is Type: 298 Q8{q: Array, s: Array} 299 300 def pick_f(b: Bool, +x: F32, +y: F32) -> F32: 301 match b: 302 case True{}: 303 x 304 case False{}: 305 y 306 307 # v rounded to the nearest integer, ties to even (torch.round), for |v| < 2^23 (here |v| <= 127 and a bit) 308 def rne_f(+v: F32) -> F32: 309 +f = F32.floor(v) 310 +d = (v - f : F32) 311 +h = (f * 0.5 : F32) 312 +up = Bool.or(F32.is_gt(d, 0.5), Bool.and(F32.is_eq(d, 0.5), F32.is_ne(F32.floor(h), h))) 313 pick_f(up, (f + 1.0 : F32), f) 314 315 def q8_of(+w: F32, +s: F32) -> F32: 316 pick_f(F32.is_eq(s, 0.0), 0.0, F32.clamp(rne_f((w / s : F32)), F32.neg(127.0), 127.0)) 317 318 # an integer -127..127 (as an F32) to its two's complement byte, and back 319 def q8_byte(+q: F32) -> U32: 320 ((F32.to_u32((q + 128.0 : F32)) + 128 : U32) .&. 255 : U32) 321 322 def q8_val(+b: U32) -> F32: 323 (U32.to_f32(((b + 128 : U32) .&. 255 : U32)) - 128.0 : F32) 324 325 def amax_l(xs: List<&2, F32>, +acc: F32) -> F32: 326 match xs: 327 case Nil{}: 328 acc 329 case Con{+h, t}: 330 amax_l(t, F32.max(acc, F32.abs(h))) 331 332 def q8_row(xs: List<&2, F32>, +s: F32) -> List<&2, U32>: 333 match xs: 334 case Nil{}: 335 Nil{} 336 case Con{+h, t}: 337 q8_byte(q8_of(h, s)) <> q8_row(t, s) 338 339 type QR is Type: 340 QR{q: List<&2, U32>, s: List<&2, F32>} 341 342 def qr_add(+row: List<&2, F32>, r: QR) -> QR: 343 match r: 344 case QR{qs, ss}: 345 +s = (amax_l(row, 0.0) / 127.0 : F32) 346 QR{List.append(&2, U32, q8_row(row, s), qs), s <> ss} 347 348 # the rows of a flat list of k rows of c, from the last one: each row's bytes and scale 349 def q8_rows(k: Nat, +c: Nat, +xs: List<&2, F32>) -> QR: 350 match k: 351 case 0n: 352 QR{Nil{}, Nil{}} 353 case 1n+p: 354 qr_add(List.take(&2, F32, xs, c), q8_rows(p, c, List.drop(&2, F32, xs, c))) 355 356 # four bytes to a word, little-endian; a short last group is padded with zeros 357 def quads(xs: List<&2, U32>) -> List<&2, U32>: 358 match xs: 359 case Nil{}: 360 Nil{} 361 case Con{+a, Nil{}}: 362 [a] 363 case Con{+a, Con{+b, Nil{}}}: 364 [(a .|. U32.shln(b, 8n) : U32)] 365 case Con{+a, Con{+b, Con{+c, Nil{}}}}: 366 [((a .|. U32.shln(b, 8n) : U32) .|. U32.shln(c, 16n) : U32)] 367 case Con{+a, Con{+b, Con{+c, Con{+d, t}}}}: 368 (((a .|. U32.shln(b, 8n) : U32) .|. U32.shln(c, 16n) : U32) .|. U32.shln(d, 24n) : U32) <> quads(t) 369 370 def fill_f(xs: List<&2, F32>, +i: U32, a: Array) -> Array: 371 match xs: 372 case Nil{}: 373 a 374 case Con{h, t}: 375 fill_f(t, (i + 1 : U32), Array.set(F32, a, i, h)) 376 377 def q8_fin(-r: Nat, -c: Nat, +n: Nat, +rr: Nat, qr: QR) -> Q8: 378 match qr: 379 case QR{qs, ss}: 380 Q8{fill_u(quads(qs), 0, Array.new(U32, T.depth(Nat.div(Nat.add(n, 3n), 4n)), 0)), fill_f(ss, 0, Array.new(F32, T.depth(rr), 0.0))} 381 382 # W quantized by rows (the matrix is consumed) 383 def Q8.quantize(+r: Nat, +c: Nat, m: TA.Mat) -> Q8: 384 q8_fin(r, c, Nat.mul(r, c), r, q8_rows(r, c, ml_l(r, c, TA.Mat.to_list(r, c, m)))) 385 386 type FL is Type: 387 FL{a: Array, l: List<&2, F32>} 388 389 def read_f(k: Nat, +i: U32, acc: List<&2, F32>, ra: Array & F32) -> FL: 390 match k ra: 391 case 0n Tuple{a, v}: 392 FL{a, List.reverse(&2, F32, v <> acc)} 393 case 1n+p Tuple{a, +v}: 394 read_f(p, (i + 1 : U32), v <> acc, Array.get(F32, a, (i + 1 : U32))) 395 396 def fl_list(r: FL) -> List<&2, F32>: 397 match r: 398 case FL{a, l}: 399 l 400 401 # byte e of the packed words, as its integer 402 def q8_at(+e: U32, +w: U32) -> F32: 403 q8_val((U32.shrn(w, U32.to_nat((8 * (e .&. 3 : U32) : U32))) .&. 255 : U32)) 404 405 # the words and scales, for printing and tests 406 type QL is Type: 407 QL{q: List<&2, U32>, s: List<&2, F32>} 408 409 def Q8.lists(+r: Nat, +c: Nat, w: Q8) -> QL: 410 match w: 411 case Q8{q, s}: 412 QL{ul_list(read_u(Nat.sub(Nat.div(Nat.add(Nat.mul(r, c), 3n), 4n), 1n), 0, Nil{}, Array.get(U32, q, 0))), fl_list(read_f(Nat.sub(r, 1n), 0, Nil{}, Array.get(F32, s, 0)))} 413 414 # W back in F32: q[i, j] · s_i, one F32 product each 415 def dq_row(k: Nat, +e: U32, +s: F32, acc: List<&2, F32>, rq: Array & U32) -> UF: 416 match k rq: 417 case 0n Tuple{q, w}: 418 UF{q, acc} 419 case 1n+p Tuple{q, +w}: 420 dq_row(p, (e + 1 : U32), s, (q8_at(e, w) * s : F32) <> acc, Array.get(U32, q, U32.shrn((e + 1 : U32), 2n))) 421 422 type UF is Type: 423 UF{q: Array, l: List<&2, F32>} 424 425 def dq_rows(k: Nat, +c: Nat, +cu: U32, +i: U32, ss: List<&2, F32>, r: UF) -> List<&2, F32>: 426 match k ss r: 427 case 0n _ UF{q, acc}: 428 List.reverse(&2, F32, acc) 429 case 1n+p Nil{} UF{q, acc}: 430 List.reverse(&2, F32, acc) 431 case 1n+p Con{+s, st} UF{q, acc}: 432 dq_rows(p, c, cu, (i + 1 : U32), st, dq_row(c, (i * cu : U32), s, acc, Array.get(U32, q, U32.shrn((i * cu : U32), 2n)))) 433 434 def dq_fin(+r: Nat, +c: Nat, q: Array, ss: List<&2, F32>) -> TA.Mat: 435 TA.Mat.from_list(r, c, dq_rows(r, c, U32.from_nat(c), 0, ss, UF{q, Nil{}})) 436 437 def dq_go(+r: Nat, +c: Nat, q: Array, sl: FL) -> TA.Mat: 438 match sl: 439 case FL{s, ss}: 440 dq_fin(r, c, q, ss) 441 442 def Q8.dequantize(+r: Nat, +c: Nat, w: Q8) -> TA.Mat: 443 match w: 444 case Q8{q, s}: 445 dq_go(r, c, q, read_f(Nat.sub(r, 1n), 0, Nil{}, Array.get(F32, s, 0))) 446 447 # ---- the product y = x · Wᵀ (n x m) for x (n x k) in F32 and W (m x k) in int8, reading the bytes from the words ---- 448 449 type QD is Type: 450 QD{x: Array, q: Array, acc: F32} 451 452 # Σ_l x[xi + l] · q[qe + l] for l < k, in order, from acc 453 def q8_dot(k: Nat, +xi: U32, +qe: U32, +acc: F32, rx: Array & F32, rq: Array & U32) -> QD: 454 match k rx rq: 455 case 0n Tuple{x, v} Tuple{q, w}: 456 QD{x, q, acc} 457 case 1n+p Tuple{x, +v} Tuple{q, +w}: 458 q8_dot(p, (xi + 1 : U32), (qe + 1 : U32), (acc + (v * q8_at(qe, w) : F32) : F32), Array.get(F32, x, (xi + 1 : U32)), Array.get(U32, q, U32.shrn((qe + 1 : U32), 2n))) 459 460 type QS is Type: 461 QS{x: Array, q: Array, s: Array, y: Array} 462 463 def q8_put2(+idx: U32, +acc: F32, x: Array, q: Array, rs: Array & F32, y: Array) -> QS: 464 match rs: 465 case Tuple{s, +sc}: 466 QS{x, q, s, Array.set(F32, y, idx, (sc * acc : F32))} 467 468 def q8_put(+idx: U32, +j: U32, d: QD, s: Array, y: Array) -> QS: 469 match d: 470 case QD{x, q, acc}: 471 q8_put2(idx, acc, x, q, Array.get(F32, s, j), y) 472 473 # the m outputs of row i 474 def q8_cols(ml: Nat, +i: U32, +j: U32, +k: Nat, +ku: U32, +mu: U32, st: QS) -> QS: 475 match ml st: 476 case 0n QS{x, q, s, y}: 477 QS{x, q, s, y} 478 case 1n+p QS{x, q, s, y}: 479 q8_cols(p, i, (j + 1 : U32), k, ku, mu, q8_put(((i * mu : U32) + j : U32), j, q8_dot(k, (i * ku : U32), (j * ku : U32), 0.0, Array.get(F32, x, (i * ku : U32)), Array.get(U32, q, U32.shrn((j * ku : U32), 2n))), s, y)) 480 481 def q8_rowsm(nl: Nat, +i: U32, +m: Nat, +k: Nat, +ku: U32, +mu: U32, st: QS) -> QS: 482 match nl: 483 case 0n: 484 st 485 case 1n+p: 486 q8_rowsm(p, (i + 1 : U32), m, k, ku, mu, q8_cols(m, i, 0, k, ku, mu, st)) 487 488 type Q8Mul<-n: Nat, -k: Nat, -m: Nat> is Type: 489 Q8Mul{x: TA.Mat, w: Q8, y: TA.Mat} 490 491 def q8_mm_fin(-n: Nat, -k: Nat, -m: Nat, st: QS) -> Q8Mul: 492 match st: 493 case QS{x, q, s, y}: 494 Q8Mul{TA.Mat{x}, Q8{q, s}, TA.Mat{y}} 495 496 # y = x · Wᵀ; x and W come back 497 def Q8.matmul_nt(+n: Nat, +k: Nat, +m: Nat, x: TA.Mat, w: Q8) -> Q8Mul: 498 match x w: 499 case TA.Mat{xa} Q8{q, s}: 500 q8_mm_fin(n, k, m, q8_rowsm(n, 0, m, k, U32.from_nat(k), U32.from_nat(m), QS{xa, q, s, T.zeros_a(Nat.mul(n, m))})) 501 502 # ===================================================================================================== 503 # A two-F32 ("double-word") accumulator, for long sums (an experiment: Bend 2.0.37 has no F64) 504 # ===================================================================================================== 505 # The running sum is hi + lo, with |lo| <= ulp(hi) / 2. Each addition is Knuth's TwoSum (the exact error of hi + x), 506 # folded into lo, then renormalized; a product is first split exactly by Dekker's TwoProduct (Veltkamp's split with 507 # 4097 = 2^12 + 1), which needs every operation rounded on its own: Bend's emitted C turns off FMA contraction. 508 509 type DW is Data: 510 DW{hi: F32, lo: F32} 511 512 def DW.zero() -> DW: 513 DW{0.0, 0.0} 514 515 def DW.value(d: DW) -> F32: 516 match d: 517 case DW{+hi, +lo}: 518 (hi + lo : F32) 519 520 # hi + lo when |hi| >= |lo| (Fast2Sum), as a renormalized pair 521 def dw_norm(+s: F32, +e: F32) -> DW: 522 +h = (s + e : F32) 523 DW{h, (e - (h - s : F32) : F32)} 524 525 # hi + x: TwoSum's s and its exact error, plus the old lo and an extra error term c 526 def dw_fin(+s: F32, +bb: F32, +a: F32, +x: F32, +lo: F32, +c: F32) -> DW: 527 +e = ((a - (s - bb : F32) : F32) + (x - bb : F32) : F32) 528 dw_norm(s, ((e + lo : F32) + c : F32)) 529 530 def dw_add3(d: DW, +x: F32, +c: F32) -> DW: 531 match d: 532 case DW{+hi, +lo}: 533 +s = (hi + x : F32) 534 dw_fin(s, (s - hi : F32), hi, x, lo, c) 535 536 def DW.add(d: DW, +x: F32) -> DW: 537 dw_add3(d, x, 0.0) 538 539 def split_hi(+a: F32) -> F32: 540 +t = (4097.0 * a : F32) 541 (t - (t - a : F32) : F32) 542 543 # a · b exactly as p + e 544 def two_prod_e(+a: F32, +b: F32, +p: F32) -> F32: 545 +ah = split_hi(a) 546 +al = (a - ah : F32) 547 +bh = split_hi(b) 548 +bl = (b - bh : F32) 549 ((((((ah * bh : F32) - p : F32) + (ah * bl : F32) : F32) + (al * bh : F32) : F32)) + (al * bl : F32) : F32) 550 551 # acc + a · b, with the product's rounding error kept 552 def DW.fma(d: DW, +a: F32, +b: F32) -> DW: 553 +p = (a * b : F32) 554 dw_add3(d, p, two_prod_e(a, b, p)) 555 556 # ---- sums over arrays: a running F32 sum, and the same in double-word ---- 557 558 type AF is Type: 559 AF{a: Array, v: F32} 560 561 def sum_f(k: Nat, +i: U32, +acc: F32, ra: Array & F32) -> AF: 562 match k ra: 563 case 0n Tuple{a, v}: 564 AF{a, acc} 565 case 1n+p Tuple{a, +v}: 566 sum_f(p, (i + 1 : U32), (acc + v : F32), Array.get(F32, a, (i + 1 : U32))) 567 568 def sum_dw(k: Nat, +i: U32, acc: DW, ra: Array & F32) -> AF: 569 match k ra: 570 case 0n Tuple{a, v}: 571 AF{a, DW.value(acc)} 572 case 1n+p Tuple{a, +v}: 573 sum_dw(p, (i + 1 : U32), DW.add(acc, v), Array.get(F32, a, (i + 1 : U32))) 574 575 type AF2 is Type: 576 AF2{a: Array, b: Array, v: F32} 577 578 def dot_f(k: Nat, +i: U32, +acc: F32, ra: Array & F32, rb: Array & F32) -> AF2: 579 match k ra rb: 580 case 0n Tuple{a, x} Tuple{b, y}: 581 AF2{a, b, acc} 582 case 1n+p Tuple{a, +x} Tuple{b, +y}: 583 dot_f(p, (i + 1 : U32), (acc + (x * y : F32) : F32), Array.get(F32, a, (i + 1 : U32)), Array.get(F32, b, (i + 1 : U32))) 584 585 def dot_dw(k: Nat, +i: U32, acc: DW, ra: Array & F32, rb: Array & F32) -> AF2: 586 match k ra rb: 587 case 0n Tuple{a, x} Tuple{b, y}: 588 AF2{a, b, DW.value(acc)} 589 case 1n+p Tuple{a, +x} Tuple{b, +y}: 590 dot_dw(p, (i + 1 : U32), DW.fma(acc, x, y), Array.get(F32, a, (i + 1 : U32)), Array.get(F32, b, (i + 1 : U32))) 591 592 # the sums of the first n numbers of a (the array comes back) 593 def Sum.f32(+n: Nat, a: Array) -> AF: 594 sum_f(n, 0, 0.0, Array.get(F32, a, 0)) 595 596 def Sum.dw(+n: Nat, a: Array) -> AF: 597 sum_dw(n, 0, DW.zero(), Array.get(F32, a, 0)) 598 599 def Dot.f32(+n: Nat, a: Array, b: Array) -> AF2: 600 dot_f(n, 0, 0.0, Array.get(F32, a, 0), Array.get(F32, b, 0)) 601 602 def Dot.dw(+n: Nat, a: Array, b: Array) -> AF2: 603 dot_dw(n, 0, DW.zero(), Array.get(F32, a, 0), Array.get(F32, b, 0)) 604 605 # ===================================================================================================== 606 # Mixed precision: BF16 storage around F32 arithmetic 607 # ===================================================================================================== 608 # Two wrappers keep a layer's types (P, X, Y, S), so a model gets BF16 storage by wrapping its layers, with the same 609 # parameters, checkpoints and optimizer state: 610 # - CastOut(l): l's output rounded to BF16 (as if stored in BF16), and the gradient arriving at it rounded too, which 611 # is what autograd does through x.to(torch.bfloat16).to(torch.float32). 612 # - CastP(l), for a layer whose parameter is one matrix: the parameter stays F32 (the master weights the optimizer 613 # updates); l's forward and backward passes see its BF16 rounding, and the gradient l accumulates is rounded to 614 # BF16 before the optimizer sees it, as autograd does for w.to(torch.bfloat16).to(torch.float32). The rounded copy 615 # is made again in the backward pass (no saved state), so S does not change. 616 # Inside the layers every operation is F32. The values are BF16 values held in F32 arrays: this gives BF16's numbers, 617 # not yet its memory (no kernel reads packed operands; that is the next step). 618 619 def bf16_mat(+r: Nat, +c: Nat, m: TA.Mat) -> TA.Mat: 620 T.Ew.map(~BF16.round, r, c, m) 621 622 def cast_f(-P: Type, -S: Type, +r: Nat, +c: Nat, f: T.Fwd, S>) -> T.Fwd, S>: 623 match f: 624 case T.Fwd{p, y, s}: 625 T.Fwd{p, bf16_mat(r, c, y), s} 626 627 def cast_fwd(~P: Type, ~X: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer, S>, p: P, x: X) -> T.Fwd, S>: 628 cast_f(P, S, r, c, T.Layer.fwd(P, X, TA.Mat, S, l, p, x)) 629 630 def cast_bwd(~P: Type, ~X: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer, S>, p: P, s: S, dy: TA.Mat, g: P) -> T.Bwd: 631 T.Layer.bwd(P, X, TA.Mat, S, l, p, s, bf16_mat(r, c, dy), g) 632 633 def CastOut(~P: Type, ~X: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer, S>) -> T.Layer, S>: 634 T.Layer{p => x => cast_fwd(~P, ~X, ~S, ~r, ~c, ~l, p, x), p => s => dy => g => cast_bwd(~P, ~X, ~S, ~r, ~c, ~l, p, s, dy, g), op => p => g => st => T.Layer.step(P, X, TA.Mat, S, l, op, p, g, st), u => T.Layer.zeros(P, X, TA.Mat, S, l), p => T.Layer.dump(P, X, TA.Mat, S, l, p), p => xs => off => T.Layer.load(P, X, TA.Mat, S, l, p, xs, off), p => xs => off => T.Layer.store(P, X, TA.Mat, S, l, p, xs, off), u => T.Layer.count(P, X, TA.Mat, S, l)} 635 636 def castp_f2(-X: Type, -Y: Type, -S: Type, -r: Nat, -c: Nat, w: TA.Mat, f: T.Fwd, Y, S>) -> T.Fwd, Y, S>: 637 match f: 638 case T.Fwd{wc, y, s}: 639 T.Fwd{w, y, s} 640 641 def castp_f(~X: Type, ~Y: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer, X, Y, S>, d: T.Two, TA.Mat>, x: X) -> T.Fwd, Y, S>: 642 match d: 643 case T.Two{w, wc}: 644 castp_f2(X, Y, S, r, c, w, T.Layer.fwd(TA.Mat, X, Y, S, l, bf16_mat(r, c, wc), x)) 645 646 def castp_b2(-X: Type, +r: Nat, +c: Nat, w: TA.Mat, b: T.Bwd, X>) -> T.Bwd, X>: 647 match b: 648 case T.Bwd{wc, dx, g}: 649 T.Bwd{w, dx, bf16_mat(r, c, g)} 650 651 def castp_b(~X: Type, ~Y: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer, X, Y, S>, d: T.Two, TA.Mat>, s: S, dy: Y, g: TA.Mat) -> T.Bwd, X>: 652 match d: 653 case T.Two{w, wc}: 654 castp_b2(X, r, c, w, T.Layer.bwd(TA.Mat, X, Y, S, l, bf16_mat(r, c, wc), s, dy, g)) 655 656 def CastP(~X: Type, ~Y: Type, ~S: Type, ~r: Nat, ~c: Nat, ~l: T.Layer, X, Y, S>) -> T.Layer, X, Y, S>: 657 T.Layer{p => x => castp_f(~X, ~Y, ~S, ~r, ~c, ~l, T.Mat.dup(r, c, p), x), p => s => dy => g => castp_b(~X, ~Y, ~S, ~r, ~c, ~l, T.Mat.dup(r, c, p), s, dy, g), op => p => g => st => T.Layer.step(TA.Mat, X, Y, S, l, op, p, g, st), u => T.Layer.zeros(TA.Mat, X, Y, S, l), p => T.Layer.dump(TA.Mat, X, Y, S, l, p), p => xs => off => T.Layer.load(TA.Mat, X, Y, S, l, p, xs, off), p => xs => off => T.Layer.store(TA.Mat, X, Y, S, l, p, xs, off), u => T.Layer.count(TA.Mat, X, Y, S, l)}