"""Lab 2: tensor-core fragments and the weight repack (_pack), emulated in numpy. Run: python3 lab2_fragments.py (parts 1 + 2) python3 lab2_fragments.py --break (part 2 with a deliberate layout mistake) Needs: Python 3 + numpy (pip install numpy) Part 1 Who holds what. mma.sync.m16n8k16 multiplies a 16x16 tile A by a 16x8 tile B and adds a 16x8 accumulator C. The operands live in registers spread over the 32 lanes of a warp. This prints which lane holds each element (the official PTX layout) and checks an emulated mma against numpy's matmul. Part 2 The repack trick. It implements _pack() from submission.py (lines 640-643) and emulates one CTA of gemv_kernel (lines 67-208): NW warps split K, each lane does ONE 16-byte load of weights per mma step and ONE 16-byte load of x, the mma runs, the warps' partial sums are reduced, and the C fragment is stored. The result must equal W @ x.T. Things to try - Run with --break: the x halves are fed to the wrong mma step. The answer becomes wrong. - Change NW (warps per CTA) or U (K blocks per loop trip) in main(); keep K/32 divisible by NW*U. - Change M (tokens) to 1, 5, 8, 13, 16, 32 and watch NB (8-token column blocks) change. """ import sys import numpy as np # ---------------------------------------------------------------------------------------------- # Part 1: the official register layouts of mma.sync.aligned.m16n8k16 (.bf16 inputs, .f32 accum) # lane = 4*g + t with g = lane >> 2 ("group", 0..7) and t = lane & 3 ("thread in group", 0..3) # ---------------------------------------------------------------------------------------------- def a_coords(lane): """(row, col) of a0..a7, the 8 values `lane` holds of the 16x16 A tile (4 registers x 2 bf16).""" g, t = lane >> 2, lane & 3 return [(g + 8 * ((i >> 1) & 1), 2 * t + (i & 1) + 8 * (i >> 2)) for i in range(8)] def b_coords(lane): """(k, n) of b0..b3, the 4 values `lane` holds of the 16x8 B tile (2 registers x 2 bf16).""" g, t = lane >> 2, lane & 3 return [(2 * t + (i & 1) + 8 * (i >> 1), g) for i in range(4)] def c_coords(lane): """(row, col) of c0..c3, the 4 fp32 values `lane` holds of the 16x8 accumulator.""" g, t = lane >> 2, lane & 3 return [(g + 8 * (i >> 1), 2 * t + (i & 1)) for i in range(4)] def mma_m16n8k16(a_regs, b_regs, c_regs): """Emulate one warp-wide mma: a_regs[lane] = 8 values, b_regs[lane] = 4, c_regs[lane] = 4. Gathers the tiles from the lanes, computes D = A @ B + C, hands D back in the C layout.""" A, B, C = np.zeros((16, 16)), np.zeros((16, 8)), np.zeros((16, 8)) for lane in range(32): for i, (r, c) in enumerate(a_coords(lane)): A[r, c] = a_regs[lane][i] for i, (k, n) in enumerate(b_coords(lane)): B[k, n] = b_regs[lane][i] for i, (r, c) in enumerate(c_coords(lane)): C[r, c] = c_regs[lane][i] D = A @ B + C return [np.array([D[r, c] for (r, c) in c_coords(lane)]) for lane in range(32)] def part1(): print("=== Part 1: which lane holds which element ===\n") grid = np.zeros((16, 16), dtype=int) for lane in range(32): for r, c in a_coords(lane): grid[r, c] = lane print("A tile (16 rows x 16 k), entry = lane that holds it in its registers:") for r in range(16): print(" " + " ".join(f"{v:2d}" for v in grid[r])) gb = np.zeros((16, 8), dtype=int) for lane in range(32): for k, n in b_coords(lane): gb[k, n] = lane print("\nB tile (16 k x 8 columns):") for k in range(16): print(" " + " ".join(f"{v:2d}" for v in gb[k])) gc = np.zeros((16, 8), dtype=int) for lane in range(32): for r, c in c_coords(lane): gc[r, c] = lane print("\nC / D accumulator (16 rows x 8 columns):") for r in range(16): print(" " + " ".join(f"{v:2d}" for v in gc[r])) rng = np.random.default_rng(0) A, B, C = rng.standard_normal((16, 16)), rng.standard_normal((16, 8)), rng.standard_normal((16, 8)) a_regs = [[A[r, c] for r, c in a_coords(l)] for l in range(32)] b_regs = [[B[k, n] for k, n in b_coords(l)] for l in range(32)] c_regs = [[C[r, c] for r, c in c_coords(l)] for l in range(32)] D = mma_m16n8k16(a_regs, b_regs, c_regs) Dm = np.zeros((16, 8)) for l in range(32): for i, (r, c) in enumerate(c_coords(l)): Dm[r, c] = D[l][i] print(f"\nemulated mma vs numpy A @ B + C: max abs difference = {np.abs(Dm - (A @ B + C)).max():.1e}\n") # ---------------------------------------------------------------------------------------------- # Part 2: _pack + one CTA of gemv_kernel # ---------------------------------------------------------------------------------------------- def pack(w): """submission.py line 643, in numpy (permute == transpose): [N, K] -> [RB, KB, step, lane(g, t), (j, h, e)], each lane's 8 values = one 16-byte A fragment.""" n, k = w.shape return w.reshape(n // 16, 2, 8, k // 32, 4, 2, 2, 2).transpose(0, 3, 5, 2, 4, 6, 1, 7).copy() def gemv_cta(x, P, rb, N, K, NW, U, broken=False): """Emulate CTA `rb` (rows 16*rb .. 16*rb+15) of gemv_kernel with S = 1. P is the packed weight array viewed as 'uint4' units: P[i] = the 8 bf16 values of 16 bytes.""" M = x.shape[0] NB = 1 if M <= 8 else 2 if M <= 16 else 4 # 8-token column blocks (nb_of) KB = K // 32 # 32-wide K blocks (line 74) KBW = KB // NW # K blocks per warp (line 74) red = np.zeros((NW, NB * 32, 4)) # __shared__ red[NW][NB*32] (line 164) for warp in range(NW): kb0 = warp * KBW # this warp's first K block (line 75) acc = [[np.zeros(4) for _ in range(NB)] for _ in range(32)] # acc[lane][nb] (line 85) for i in range(0, KBW, U): # main loop (line 90) for u in range(U): kb = kb0 + i + u woff = (rb * KB + kb) * 64 # (line 76): 64 x 16 B per (row block, K block) a0 = [P[woff + lane] for lane in range(32)] # step 0: ld_stream(wp + (i+u)*64) a1 = [P[woff + 32 + lane] for lane in range(32)] # step 1: ld_stream(wp + (i+u)*64 + 32) for nb in range(NB): xv = [] # __ldg 16 B of x per lane (line 105) for lane in range(32): g, t = lane >> 2, lane & 3 tok = min(nb * 8 + g, M - 1) # clamp: rows past M re-read the last token (line 82) xv.append(x[tok, kb * 32 + t * 8: kb * 32 + t * 8 + 8]) lo = [v[0:4] for v in xv] # xv.x, xv.y = x values 0..3 of the lane hi = [v[4:8] for v in xv] # xv.z, xv.w = x values 4..7 if broken: lo, hi = hi, lo d = mma_m16n8k16(a0, lo, [acc[l][nb] for l in range(32)]) # line 106 d = mma_m16n8k16(a1, hi, d) # line 107 for l in range(32): acc[l][nb] = d[l] for nb in range(NB): # each warp parks its partial (line 167) for lane in range(32): red[warp, nb * 32 + lane] = acc[lane][nb] out = np.full((M, N), np.nan) # __syncthreads(); then threads tid < NB*32: for tid in range(NB * 32): r = red[:, tid].sum(0) # sum over warps (lines 173-176) nb, lane = tid >> 5, tid & 31 g, t = lane >> 2, lane & 3 for q in range(4): # C fragment -> (token, row) (lines 199-205) tok = nb * 8 + 2 * t + (q & 1) row = rb * 16 + g + ((q >> 1) << 3) if tok < M: out[tok, row] = r[q] return out def part2(broken): print("=== Part 2: _pack + one CTA of gemv_kernel ===\n") N, K, M, NW, U = 32, 512, 5, 4, 2 # small sizes so pure Python is quick rng = np.random.default_rng(1) W = rng.standard_normal((N, K)).astype(np.float32) x = rng.standard_normal((M, K)).astype(np.float32) packed = pack(W) P = packed.reshape(-1, 8) # one row = one 16-byte (uint4) load print("what lanes 0..5 load in their first two 16-byte weight loads (row block 0, K block 0):") for lane in range(6): g, t = lane >> 2, lane & 3 for step in range(2): vals = P[step * 32 + lane] where = [tuple(int(v) for v in np.argwhere(W == val)[0]) for val in vals] print(f" lane {lane} (g={g}, t={t}) step {step}: (row, k) = {where}") print(" -> 8 values of rows g and g+8, at k = 8t+4*step+{0,1,2,3}: the same k the lane's 16-byte x load covers\n") out = np.zeros((M, N)) for rb in range(N // 16): # one CTA per 16-row block (grid.x = N/16) cta = gemv_cta(x, P, rb, N, K, NW, U, broken) out[:, rb * 16:(rb + 1) * 16] = cta[:, rb * 16:(rb + 1) * 16] ref = x.astype(np.float64) @ W.astype(np.float64).T err = np.abs(out - ref).max() / np.abs(ref).max() print(f"N={N} K={K} M={M} NW={NW} U={U}: emulated gemv vs x @ W.T -> max relative error {err:.1e}" f" {'OK' if err < 1e-6 else 'WRONG'}") if broken: print("(--break fed x values 4..7 to the step-0 fragment: the k labels of A and B no longer match)") if __name__ == "__main__": brk = "--break" in sys.argv if not brk: part1() part2(brk)