| |
| """Recompute a whole paged MoE layer in NumPy and compare with sdq-chat's output. |
| |
| `sdq-chat` with SDQ_MOE_DUMP=<prefix> SDQ_MOE_DUMP_LAYER=<il> writes, from inside |
| the custom MoE op: |
| int32 meta[6] = { n_embd, n_tok, n_expert_used, n_expert, 1, 2 } |
| float cur [n_tok][n_embd] (FFN input) |
| int32 ids [n_tok][n_expert_used] (routed experts) |
| float wts [n_tok][n_expert_used] (routed weights) |
| float out [n_tok][n_embd] (SDQ MoE output) |
| |
| This script reads the same expert weights straight out of the GGUF, runs the |
| SiLU-gated expert forward in fp32 and diffs the result, so the whole |
| pager + custom kernel + weight-scaling path is checked at once. |
| """ |
| import importlib.util |
| import struct |
| import sys |
|
|
| import numpy as np |
|
|
| spec = importlib.util.spec_from_file_location('vl', str(__import__('pathlib').Path(__file__).with_name('verify_ling.py'))) |
| vl = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(vl) |
|
|
| QK_K = 256 |
|
|
|
|
| def dequant_q4_K_vec(raw, n_elem): |
| """raw: uint8 [..., 144] -> float32 [..., 256]""" |
| raw = raw.reshape(-1, 144) |
| nb = raw.shape[0] |
| out = np.empty((nb, QK_K), dtype=np.float32) |
| d = raw[:, 0:2].copy().view(np.float16).astype(np.float32)[:, 0] |
| dmin = raw[:, 2:4].copy().view(np.float16).astype(np.float32)[:, 0] |
| b = np.ascontiguousarray(raw[:, 4:16]) |
| sc = np.empty((nb, 8), dtype=np.int32) |
| mn = np.empty((nb, 8), dtype=np.int32) |
| for j in range(4): |
| sc[:, j] = b[:, j] & 63 |
| mn[:, j] = b[:, j + 4] & 63 |
| for j in range(4, 8): |
| sc[:, j] = (b[:, j + 4] & 0xF) | ((b[:, j - 4] >> 6) << 4) |
| mn[:, j] = (b[:, j + 4] >> 4) | ((b[:, j] >> 6) << 4) |
| qs = raw[:, 16:144].reshape(nb, 4, 32).astype(np.int32) |
| q = np.empty((nb, QK_K), dtype=np.int32) |
| for j in range(4): |
| q[:, j * 64:j * 64 + 32] = qs[:, j, :] & 0xF |
| q[:, j * 64 + 32:j * 64 + 64] = qs[:, j, :] >> 4 |
| gidx = np.repeat(np.arange(8), 32) |
| out[:] = d[:, None] * sc[:, gidx] * q - dmin[:, None] * mn[:, gidx] |
| return out |
|
|
|
|
| def dequant_q6_K_vec(raw, n_elem): |
| raw = raw.reshape(-1, 210) |
| nb = raw.shape[0] |
| out = np.empty((nb, QK_K), dtype=np.float32) |
| d = raw[:, 208:210].copy().view(np.float16).astype(np.float32)[:, 0] |
| ql = raw[:, 0:128].reshape(nb, 4, 32).astype(np.int32) |
| qh = raw[:, 128:192].reshape(nb, 2, 32).astype(np.int32) |
| sc = raw[:, 192:208].copy().view(np.int8).astype(np.int32) |
| q = np.empty((nb, QK_K), dtype=np.int32) |
| sidx = np.empty(QK_K, dtype=np.int32) |
| for g in range(2): |
| for l in range(32): |
| is_ = l // 16 |
| o = g * 128 |
| a = ql[:, 2 * g, l] |
| b_ = ql[:, 2 * g + 1, l] |
| h = qh[:, g, l] |
| q[:, o + l] = ((a & 0xF) | ((h & 3) << 4)) - 32 |
| q[:, o + l + 32] = ((b_ & 0xF) | (((h >> 2) & 3) << 4)) - 32 |
| q[:, o + l + 64] = ((a >> 4) | (((h >> 4) & 3) << 4)) - 32 |
| q[:, o + l + 96] = ((b_ >> 4) | (((h >> 6) & 3) << 4)) - 32 |
| sidx[o + l] = g * 8 + is_ |
| sidx[o + l + 32] = g * 8 + is_ + 2 |
| sidx[o + l + 64] = g * 8 + is_ + 4 |
| sidx[o + l + 96] = g * 8 + is_ + 6 |
| out[:] = d[:, None] * sc[:, sidx] * q |
| return out |
|
|
|
|
| def dequant_mat(raw, n_rows, ne0, tt): |
| nrow = vl.row_bytes(ne0, tt) |
| nblk = ne0 // QK_K |
| arr = np.frombuffer(raw[:n_rows * nrow], dtype=np.uint8).reshape(n_rows, nblk, nrow // nblk) |
| if tt == 12: |
| return dequant_q4_K_vec(arr, ne0).reshape(n_rows, ne0) |
| if tt == 14: |
| return dequant_q6_K_vec(arr, ne0).reshape(n_rows, ne0) |
| raise ValueError(tt) |
|
|
|
|
| def main(): |
| dump = sys.argv[1] |
| il = int(sys.argv[2]) |
| kv, ts = vl.parse_headers() |
| n_embd = int(kv['bailingmoe3.embedding_length']) |
| n_ff = int(kv['bailingmoe3.expert_feed_forward_length']) |
| n_expert = int(kv['bailingmoe3.expert_count']) |
|
|
| raw = open(dump, 'rb').read() |
| meta = struct.unpack('<6i', raw[:24]) |
| n_tok, n_used = meta[1], meta[2] |
| off = 24 |
| cur = np.frombuffer(raw, dtype=np.float32, count=n_tok * n_embd, offset=off).reshape(n_tok, n_embd) |
| off += 4 * n_tok * n_embd |
| ids = np.frombuffer(raw, dtype=np.int32, count=n_tok * n_used, offset=off).reshape(n_tok, n_used) |
| off += 4 * n_tok * n_used |
| wts = np.frombuffer(raw, dtype=np.float32, count=n_tok * n_used, offset=off).reshape(n_tok, n_used) |
| off += 4 * n_tok * n_used |
| got = np.frombuffer(raw, dtype=np.float32, count=n_tok * n_embd, offset=off).reshape(n_tok, n_embd) |
| print(f'layer {il}: n_tok={n_tok} n_expert_used={n_used} n_expert={n_expert} ' |
| f'n_embd={n_embd} n_ff={n_ff}') |
|
|
| def seg(nm, ie): |
| si, off_t, dims, tt = ts[f'blk.{il}.{nm}'] |
| nrow = vl.row_bytes(dims[0], tt) |
| nb = nrow * dims[1] |
| with open(vl.SHARDS[si], 'rb') as f: |
| f.seek(off_t + ie * nb) |
| return f.read(nb), nrow, tt, dims[0] |
|
|
| ref = np.zeros((n_tok, n_embd), dtype=np.float32) |
| experts = sorted(set(int(e) for e in ids.reshape(-1))) |
| print(f'distinct experts: {len(experts)}') |
| for ie in experts: |
| g_raw, g_nrow, g_t, g_ne0 = seg('ffn_gate_exps.weight', ie) |
| u_raw, _, u_t, u_ne0 = seg('ffn_up_exps.weight', ie) |
| d_raw, _, d_t, d_ne0 = seg('ffn_down_exps.weight', ie) |
| Wg = dequant_mat(g_raw, n_ff, g_ne0, g_t) |
| Wu = dequant_mat(u_raw, n_ff, u_ne0, u_t) |
| Wd = dequant_mat(d_raw, n_embd, d_ne0, d_t) |
| rows = [t for t in range(n_tok) if ie in ids[t]] |
| x = cur[rows].T |
| w = np.array([wts[t][list(ids[t]).index(ie)] for t in rows], dtype=np.float32) |
| h = ((Wg @ x) / (1.0 + np.exp(-(Wg @ x)))) * (Wu @ x) |
| y = Wd @ h |
| ref[rows] += (y * w).T |
| del Wg, Wu, Wd |
|
|
| denom = max(float(np.abs(ref).max()), 1e-9) |
| diff = np.abs(got - ref) |
| corr = float(np.corrcoef(got.reshape(-1), ref.reshape(-1))[0, 1]) |
| print(f'max|got-ref| = {diff.max():.6g} max|ref| = {denom:.6g} ' |
| f'max_rel = {diff.max() / denom:.6g} corr = {corr:.8f}') |
| ok = corr > 0.99 and diff.max() / denom < 0.05 |
| print('PASS' if ok else 'FAIL') |
| return 0 if ok else 1 |
|
|
|
|
| if __name__ == '__main__': |
| sys.exit(main()) |
|
|
