"""Lab 3: the lossless 12-bit weight format (pack_z + decz), emulated bit for bit in numpy. Run: python3 lab3_compress.py (a 1024 x 8192 matrix, about 10 s) python3 lab3_compress.py --full (pooler size, 4096 x 8192) Needs: Python 3 + numpy. Part 1 Why it works: the histogram of the top 7 exponent bits of BF16 weights drawn like the judge's (normal, std 1/sqrt(K)), the 8-value window pack_z picks, and how many weights fall outside it. Part 2 pack_z (submission.py lines 650-708) ported to numpy, then the kernel's decoder decz (lines 257-273) emulated with exact prmt/lop3 semantics. Every decoded fragment is compared with the expected bits. Part 3 One CTA of gemvz: decoded weights times x, plus the exception warp's corrections. Checks that decoded + corrections reproduces W exactly, including a one-hot "probe" row like the judge's. Things to try - NWIN_SWEEP: how the outside-the-window fraction changes with the window size (2, 4, 8, 16 values). - Replace the normal weights with heavier tails: W = rng.standard_t(3, (N, K)) / sqrt(K). """ import sys import numpy as np # ---------- BF16 helpers: a BF16 value is the top 16 bits of an FP32 value ---------- def f32_to_bf16_bits(f): """Round-to-nearest-even FP32 -> BF16 bit pattern (uint16), like torch's .to(torch.bfloat16).""" u = np.asarray(f, dtype=np.float32).view(np.uint32).astype(np.uint64) u = (u + 0x7FFF + ((u >> 16) & 1)) >> 16 return u.astype(np.uint16) def bf16_bits_to_f32(b): return (np.asarray(b).astype(np.uint32) << 16).view(np.float32) # ---------- PTX instruction emulation ---------- def prmt(a, b, sel): """prmt.b32 d, a, b, sel: byte i of d = byte number (sel >> 4i) & 7 of the 8-byte value {b, a} (a = bytes 0-3, b = bytes 4-7). If bit 3 of the selector nibble is set, the byte's sign is replicated.""" src = [(a >> (8 * i)) & 0xFF for i in range(4)] + [(b >> (8 * i)) & 0xFF for i in range(4)] d = 0 for i in range(4): nib = (sel >> (4 * i)) & 0xF byte = src[nib & 7] if nib & 8: byte = 0xFF if byte & 0x80 else 0x00 d |= byte << (8 * i) return d def lop3(a, b, c, lut): """lop3.b32 d, a, b, c, lut: any 3-input bitwise function. lut is the function's truth table, built by applying it to the constants a=0xF0, b=0xCC, c=0xAA. 0xF8 = 0xF0 | (0xCC & 0xAA) means a | (b & c).""" d = 0 for bit in range(32): idx = (((a >> bit) & 1) << 2) | (((b >> bit) & 1) << 1) | ((c >> bit) & 1) d |= ((lut >> idx) & 1) << bit return d def decz(S0, S1, Q, TL, TH): """submission.py lines 264-273: 8 compressed weights -> one 16-byte mma A fragment (4 words).""" H0 = lop3(prmt(TL, TH, Q & 0x7777), Q, 0x80808080, 0xF8) # high bytes, weights 0-3 H1 = lop3(prmt(TL, TH, (Q >> 16) & 0x7777), (Q << 4) & 0xFFFFFFFF, 0x80808080, 0xF8) # weights 4-7 return (prmt(S0, H0, 0x5140), prmt(S0, H0, 0x7362), # a.x = weights 0,1 a.y = weights 2,3 prmt(S1, H1, 0x5140), prmt(S1, H1, 0x7362)) # a.z = weights 4,5 a.w = weights 6,7 # ---------- pack_z, ported line by line ---------- def perm(a, rb, kb): """[N, K] -> [RB, G, f(4), lane(32), i(8)]: mma A-fragment order, 4 fragments per 64-wide group.""" p = a.reshape(rb, 2, 8, kb, 4, 2, 2, 2).transpose(0, 3, 5, 2, 4, 6, 1, 7) return p.reshape(rb, kb // 2, 4, 32, 8) def pack_z(wbits, nwin=8): n, k = wbits.shape rb, kb = n // 16, k // 32 bits = wbits.astype(np.int64) # line 658: the 16-bit patterns e7 = (bits >> 8) & 0x7F # line 659: top 7 exponent bits h = np.bincount(e7.ravel(), minlength=128) # line 660: histogram c = np.concatenate([[0], np.cumsum(h)]) # line 661 cover = c[nwin:] - c[:-nwin] # line 662: weights inside [b, b+nwin) for each b base = int(np.argmax(cover)) # line 663: best window start code = np.clip(e7 - base, 0, nwin - 1) # line 664: 3-bit code (clamped outside the window) sign = bits >> 15 # line 665 low = bits & 0xFF # line 666: raw low byte dbits = (sign << 15) | ((base + code) << 8) | low # line 667: what the kernel will decode dec = bf16_bits_to_f32(dbits.astype(np.uint16)) w = bf16_bits_to_f32(wbits) d64 = w.astype(np.float64) - dec.astype(np.float64) # lines 673-675: exact difference, d32 = d64.astype(np.float32) # its fp32 value, res = (d64 - d32.astype(np.float64)).astype(np.float32) # and a residual if fp32 was not exact i1, i2 = np.nonzero(d32), np.nonzero(res) rows = np.concatenate([i1[0], i2[0]]) ks = np.concatenate([i1[1], i2[1]]) vals = np.concatenate([d32[i1], res[i2]]) order = np.argsort(rows * k + ks, kind="stable") # lines 681-689: per-row exception lists rows, ks, vals = rows[order], ks[order], vals[order] cnt = np.bincount(rows, minlength=n) capr = max(1, int(cnt.max()) if len(rows) else 1) start = np.cumsum(cnt) - cnt pos = np.arange(len(rows)) - start[rows] esc = np.zeros((n, capr, 2), dtype=np.int32) esc[rows, pos, 0] = ks esc[rows, pos, 1] = vals.astype(np.float32).view(np.int32) lowp = perm(low.astype(np.uint8), rb, kb) # line 695 A = lowp.transpose(0, 1, 3, 2, 4).reshape(rb, kb // 2, 32, 2, 16).transpose(0, 1, 3, 2, 4) # line 696 codp, sgnp = perm(code, rb, kb), perm(sign, rb, kb) # lines 697-698 csh = np.array([0, 4, 8, 12, 16, 20, 24, 28]) # line 699 ssh = np.array([7, 15, 23, 31, 3, 11, 19, 27]) # line 700 q = (codp << csh).sum(-1) + (sgnp << ssh).sum(-1) # line 701 B = q.astype(np.uint32).transpose(0, 1, 3, 2).copy().view(np.uint8).reshape(rb, kb // 2, 1, 32, 16) # 703 packed = np.concatenate([A, B], 2) # line 704 lut = [base + j for j in range(8)] # lines 705-707 TL = sum(lut[j] << (8 * j) for j in range(4)) TH = sum(lut[4 + j] << (8 * j) for j in range(4)) return dict(w=packed, esc=esc, capr=capr, TL=TL, TH=TH, base=base, nesc=len(rows), n2=len(i2[0]), dbits=dbits, h=h, cover=cover) def words(b16): """16 bytes -> the 4 little-endian 32-bit words a uint4 load puts in .x .y .z .w""" return [int(v) for v in np.frombuffer(b16.tobytes(), dtype="> 15, (b >> 8) & 0x7F, b & 0xFF code = min(max(e7 - base, 0), 7) rebuilt = (sign << 15) | ((base + code) << 8) | low print(f" weight {value:+.6g} -> BF16 bits {b:016b} (0x{b:04x}) = {float(bf16_bits_to_f32(b)):+.6g}") print(f" sign {sign} e7 (bits 14-8) {e7:3d} = {e7:07b} low byte (bits 7-0) 0x{low:02x} = {low:08b}") print(f" stored: low byte 0x{low:02x} + nibble (sign {sign}, code e7-base = {e7}-{base} = {e7 - base}" f"{' -> clamped to ' + str(code) if code != e7 - base else ''}) = 12 bits") print(f" rebuilt: sign | (base+code)<<8 | low = 0x{rebuilt:04x} = {float(bf16_bits_to_f32(rebuilt)):+.6g}" f" {'exact' if rebuilt == b else '-> needs an exception: ' + format(float(bf16_bits_to_f32(b)) - float(bf16_bits_to_f32(rebuilt)), '+.6g')}") def main(): full = "--full" in sys.argv N, K = (4096, 8192) if full else (1024, 8192) rng = np.random.default_rng(5151) W = (rng.standard_normal((N, K), dtype=np.float32) / np.sqrt(K)).astype(np.float32) wbits = f32_to_bf16_bits(W) # the BF16 weights the kernel must reproduce print(f"=== Part 1: exponent histogram ({N} x {K} BF16 weights, normal with std 1/sqrt(K)) ===") z = pack_z(wbits) h, base = z["h"], z["base"] total = h.sum() for e in range(128): if h[e] and h[e] / total > 1e-6: inside = "<- window" if base <= e < base + 8 else "" print(f" e7={e:3d} |w| in [2^{2 * e - 127}, 2^{2 * e - 125}) {100 * h[e] / total:8.4f}% " f"{'#' * int(60 * h[e] / h.max())} {inside}") print(f" window: e7 in [{base}, {base + 7}] covers {100 * z['cover'][base] / total:.4f}% of weights") print(f" exceptions: {z['nesc']} entries ({100 * z['nesc'] / total:.4f}%), {z['n2']} of them fp32 residuals;" f" longest per-row list (capr) = {z['capr']}") print(" outside the window (pooler at full size has about 0.014%):" " tiny or zero weights, stored clamped + an exact fp32 correction") NWIN_SWEEP = (2, 4, 8, 16) c = np.concatenate([[0], np.cumsum(h)]) for nw in NWIN_SWEEP: cov = (c[nw:] - c[:-nw]).max() / total bits = 8 + 1 + int(np.log2(nw)) print(f" window of {nw:2d} values ({bits} bits per weight): {100 * (1 - cov):.4f}% outside") mb16 = N * K * 2 / 1e6 print(f" bytes: {mb16:.1f} MB as BF16 -> {mb16 * 0.75:.1f} MB packed + {z['esc'].nbytes / 1e6:.2f} MB exception table") print("\n=== Part 1b: one weight, step by step ===") for v in (0.0123, -0.000457, 3.1e-7, 0.0): explain(v, base) print("\n=== Part 2: decode every fragment of the first 16-row blocks with emulated prmt/lop3 ===") rb, kb = N // 16, K // 32 expect = perm(z["dbits"], rb, kb) # [RB, G, f, lane, 8] expected 16-bit patterns bad = 0 for r in range(2): # 2 row blocks x 128 groups x 32 lanes x 4 fragments for g in range(kb // 2): for lane in range(32): a0 = words(z["w"][r, g, 0, lane]) # A0: low bytes of fragments 0, 1 a1 = words(z["w"][r, g, 1, lane]) # A1: low bytes of fragments 2, 3 bq = words(z["w"][r, g, 2, lane]) # B: one 32-bit code/sign word per fragment frags = [decz(a0[0], a0[1], bq[0], z["TL"], z["TH"]), decz(a0[2], a0[3], bq[1], z["TL"], z["TH"]), decz(a1[0], a1[1], bq[2], z["TL"], z["TH"]), decz(a1[2], a1[3], bq[3], z["TL"], z["TH"])] for f, fr in enumerate(frags): got = [(w >> (16 * h_)) & 0xFFFF for w in fr for h_ in (0, 1)] bad += int(np.any(np.array(got) != expect[r, g, f, lane])) print(f" decoded {2 * (kb // 2) * 32 * 4} fragments: {bad} mismatches {'OK' if bad == 0 else 'WRONG'}") print("\n=== Part 3: one CTA of gemvz (decoded weights x x + the exception warp) ===") M = 4 x = f32_to_bf16_bits(rng.standard_normal((M, K), dtype=np.float32)) xf = bf16_bits_to_f32(x).astype(np.float64) rows = slice(0, 16) # CTA rb = 0 dec = bf16_bits_to_f32(z["dbits"][rows].astype(np.uint16)).astype(np.float64) main_loop = xf @ dec.T # what the streaming warps compute (decoded weights) ecorr = np.zeros((16, M)) # what the exception warp computes (lines 291-314) for r in range(16): for j in range(z["capr"]): kk, bitsv = z["esc"][r, j] if bitsv != 0: d = float(np.array([bitsv], dtype=np.int32).view(np.float32)[0]) ecorr[r] += d * xf[:, kk] out = main_loop + ecorr.T ref = xf @ bf16_bits_to_f32(wbits[rows]).astype(np.float64).T print(f" rows 0-15, {M} tokens: |out - ref| max = {np.abs(out - ref).max():.2e} (decoded only: " f"{np.abs(main_loop - ref).max():.2e})") e_rows = np.nonzero(z["esc"][:, 0, 1])[0] if len(e_rows): r = int(e_rows[0]) kk = int(z["esc"][r, 0, 0]) probe = np.zeros(K) probe[kk] = 8.0 # a judge-style probe row: one nonzero, a power of two decw = bf16_bits_to_f32(z["dbits"][r, kk].astype(np.uint16)) corr = sum(float(np.array([z["esc"][r, j, 1]], dtype=np.int32).view(np.float32)[0]) for j in range(z["capr"]) if z["esc"][r, j, 0] == kk and z["esc"][r, j, 1] != 0) exact = 8.0 * float(bf16_bits_to_f32(wbits[r, kk])) print(f" probe on exception (row {r}, k {kk}): exact {exact:.6e}, decoded only {8.0 * float(decw):.6e}," f" decoded + correction {8.0 * (float(decw) + corr):.6e}") if __name__ == "__main__": main()