| |
| """SDQ MoE kernel 的 NumPy 對照實作（單元測試）。 |
| |
| 用途：拿 C++ 端 `epq-chat --selftest IL:IE:SEED:OUT.bin` 算出來的結果， |
| 用純 NumPy 從 GGUF 的同一段位元組重算一次，確認分頁 + 量化點積 + SiLU + 加權 |
| 這條路徑的數學完全正確。 |
| |
| 實作完全照抄 ggml 的定義： |
| * dequantize_row_q4_K / dequantize_row_q5_K / dequantize_row_q6_K |
| * quantize_row_q8_K_ref（Q8_K activation） |
| * SiLU(x) = x * sigmoid(x) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import struct |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| QK_K = 256 |
|
|
|
|
| def get_scale_min_k4(j: int, q: np.ndarray) -> tuple[int, int]: |
| if j < 4: |
| return int(q[j] & 63), int(q[j + 4] & 63) |
| return ((int(q[j + 4] & 0xF) | ((int(q[j - 4]) >> 6) << 4)), |
| ((int(q[j + 4]) >> 4) | ((int(q[j]) >> 6) << 4))) |
|
|
|
|
| def fp16_to_fp32(raw: bytes) -> float: |
| return struct.unpack("<e", raw)[0] |
|
|
|
|
| def dequant_q4_K(raw: bytes, n: int) -> np.ndarray: |
| nb = n // QK_K |
| out = np.empty(n, dtype=np.float32) |
| for i in range(nb): |
| d = fp16_to_fp32(raw[i * 144:i * 144 + 2]) |
| dmin = fp16_to_fp32(raw[i * 144 + 2:i * 144 + 4]) |
| scales = np.frombuffer(raw, dtype=np.uint8, count=12, offset=i * 144 + 4) |
| qs = np.frombuffer(raw, dtype=np.uint8, count=128, offset=i * 144 + 16) |
| base = i * QK_K |
| is_ = 0 |
| for j in range(0, QK_K, 64): |
| sc1, m1 = get_scale_min_k4(is_ + 0, scales) |
| sc2, m2 = get_scale_min_k4(is_ + 1, scales) |
| d1, mm1 = d * sc1, dmin * m1 |
| d2, mm2 = d * sc2, dmin * m2 |
| q = qs[j // 2: j // 2 + 32] |
| out[base + j: base + j + 32] = d1 * (q & 0xF).astype(np.float32) - mm1 |
| out[base + j + 32: base + j + 64] = d2 * (q >> 4).astype(np.float32) - mm2 |
| is_ += 2 |
| return out |
|
|
|
|
| def dequant_q5_K(raw: bytes, n: int) -> np.ndarray: |
| nb = n // QK_K |
| out = np.empty(n, dtype=np.float32) |
| for i in range(nb): |
| o = i * 176 |
| d = fp16_to_fp32(raw[o:o + 2]) |
| dmin = fp16_to_fp32(raw[o + 2:o + 4]) |
| scales = np.frombuffer(raw, dtype=np.uint8, count=12, offset=o + 4) |
| qh = np.frombuffer(raw, dtype=np.uint8, count=32, offset=o + 16) |
| qs = np.frombuffer(raw, dtype=np.uint8, count=128, offset=o + 48) |
| base = i * QK_K |
| is_ = 0 |
| u1, u2 = 1, 2 |
| for j in range(0, QK_K, 64): |
| sc1, m1 = get_scale_min_k4(is_ + 0, scales) |
| sc2, m2 = get_scale_min_k4(is_ + 1, scales) |
| d1, mm1 = d * sc1, dmin * m1 |
| d2, mm2 = d * sc2, dmin * m2 |
| q = qs[j // 2: j // 2 + 32] |
| |
| h1 = ((qh[0:32] & u1) != 0).astype(np.float32) * 16 |
| h2 = ((qh[0:32] & u2) != 0).astype(np.float32) * 16 |
| out[base + j: base + j + 32] = d1 * ((q & 0xF).astype(np.float32) + h1) - mm1 |
| out[base + j + 32: base + j + 64] = d2 * ((q >> 4).astype(np.float32) + h2) - mm2 |
| is_ += 2 |
| u1 <<= 2 |
| u2 <<= 2 |
| return out |
|
|
|
|
| def dequant_q6_K(raw: bytes, n: int) -> np.ndarray: |
| """逐位元組對照 ggml/src/ggml-quants.c 的 dequantize_row_q6_K。 |
| |
| block_q6_K（210 bytes）：ql[128] qh[64] scales[16(int8)] d(fp16) |
| 對每個 128 元素半區塊（n 迴圈，ql += 64 / qh += 32 / sc += 8）： |
| q1 = (ql[l] & 0xF) | ((qh[l] >> 0 & 3) << 4) → y[l] |
| q2 = (ql[l+32] & 0xF) | ((qh[l] >> 2 & 3) << 4) → y[l+32] |
| q3 = (ql[l] >> 4) | ((qh[l] >> 4 & 3) << 4) → y[l+64] |
| q4 = (ql[l+32] >> 4) | ((qh[l] >> 6 & 3) << 4) → y[l+96] |
| 縮放：y[l] 用 sc[is+0]、y[l+32] 用 sc[is+2]、y[l+64] 用 sc[is+4]、y[l+96] 用 sc[is+6]， |
| 其中 is = l / 16（l<16 用偶數、l>=16 用緊接的奇數）→ 等價於「每 16 個元素一個 scale」。 |
| """ |
| nb = n // QK_K |
| out = np.empty(n, dtype=np.float32) |
| for i in range(nb): |
| o = i * 210 |
| ql = np.frombuffer(raw, dtype=np.uint8, count=128, offset=o) |
| qh = np.frombuffer(raw, dtype=np.uint8, count=64, offset=o + 128) |
| sc = np.frombuffer(raw, dtype=np.int8, count=16, offset=o + 192).astype(np.float32) |
| d = fp16_to_fp32(raw[o + 208:o + 210]) |
| base = i * QK_K |
| for half in range(2): |
| qa = ql[half * 64: half * 64 + 64].astype(np.int32) |
| ha = qh[half * 32: half * 32 + 32].astype(np.int32) |
| sa = sc[half * 8: half * 8 + 8] |
| scale = np.repeat(sa, 16) |
|
|
| lo = qa[:32] & 0xF |
| hi = qa[32:] & 0xF |
| blo = qa[:32] >> 4 |
| bhi = qa[32:] >> 4 |
|
|
| q1 = lo | ((ha & 0x3) << 4) |
| q2 = hi | (((ha >> 2) & 0x3) << 4) |
| q3 = blo | (((ha >> 4) & 0x3) << 4) |
| q4 = bhi | (((ha >> 6) & 0x3) << 4) |
|
|
| p = base + half * 128 |
| out[p + 0: p + 32] = d * (q1 - 32) * scale[0:32] |
| out[p + 32: p + 64] = d * (q2 - 32) * scale[32:64] |
| out[p + 64: p + 96] = d * (q3 - 32) * scale[64:96] |
| out[p + 96: p + 128] = d * (q4 - 32) * scale[96:128] |
| return out |
|
|
|
|
| def quant_q8_K(x: np.ndarray) -> np.ndarray: |
| """照抄 quantize_row_q8_K_ref：回傳反量化後的 activation（float32）。""" |
| n = x.size |
| out = np.empty(n, dtype=np.float32) |
| for i in range(n // QK_K): |
| blk = x[i * QK_K:(i + 1) * QK_K] |
| ax = np.abs(blk) |
| amax = ax.max() |
| if amax == 0: |
| out[i * QK_K:(i + 1) * QK_K] = 0 |
| continue |
| |
| idx = int(np.argmax(ax)) |
| maxv = float(blk[idx]) |
| iscale = -127.0 / maxv |
| q = np.clip(np.rint(iscale * blk), -127, 127).astype(np.int8) |
| out[i * QK_K:(i + 1) * QK_K] = q.astype(np.float32) / iscale |
| return out |
|
|
|
|
| def read_expert(fh, off: int, n: int, typ: str) -> np.ndarray: |
| fh.seek(off) |
| raw = fh.read(n) |
| if typ == "Q4_K": |
| return dequant_q4_K(raw, n) |
| if typ == "Q5_K": |
| return dequant_q5_K(raw, n) |
| if typ == "Q6_K": |
| return dequant_q6_K(raw, n) |
| raise SystemExit(f"selftest 暫不支援 {typ}") |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description="SDQ MoE kernel 的 NumPy 對照") |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--index", required=True, help="tools/gguf_index.py 產生的 JSON") |
| ap.add_argument("--bin", required=True, help="epq-chat --selftest 的輸出") |
| ap.add_argument("--layer", type=int, required=True) |
| ap.add_argument("--expert", type=int, required=True) |
| ap.add_argument("--n-tok", type=int, default=1, help="幾個 token 共用同一個 expert") |
| ap.add_argument("--tol", type=float, default=2e-3, help="相對誤差上限") |
| args = ap.parse_args() |
|
|
| idx = json.loads(Path(args.index).read_text()) |
| layer = idx["per_layer"][args.layer]["experts"] |
| n_embd = idx["n_embd"] |
| n_ff = idx["n_ff_exp"] |
| e = args.expert |
|
|
| n_tok = args.n_tok |
| raw = np.fromfile(args.bin, dtype=np.float32) |
| assert raw.size == 2 * n_embd * n_tok, f"selftest 輸出大小 {raw.size} 應為 {2 * n_embd * n_tok}" |
| cpp_out = raw[:n_embd * n_tok].reshape(n_tok, n_embd) |
| x = raw[n_embd * n_tok:].reshape(n_tok, n_embd) |
|
|
| with open(args.model, "rb") as fh: |
| g = layer["ffn_gate_exps"] |
| u = layer["ffn_up_exps"] |
| d = layer["ffn_down_exps"] |
| gate = read_expert(fh, g["file_offset"] + e * g["bytes_per_expert"], |
| n_ff * n_embd, g["type"]).reshape(n_ff, n_embd) |
| up = read_expert(fh, u["file_offset"] + e * u["bytes_per_expert"], |
| n_ff * n_embd, u["type"]).reshape(n_ff, n_embd) |
| down = read_expert(fh, d["file_offset"] + e * d["bytes_per_expert"], |
| n_embd * n_ff, d["type"]).reshape(n_embd, n_ff) |
|
|
| |
| outs = [] |
| for t in range(n_tok): |
| xq = quant_q8_K(x[t]) |
| gate_out = gate @ xq |
| up_out = up @ xq |
| h = (gate_out / (1.0 + np.exp(-gate_out))) * up_out |
| hq = quant_q8_K(h) |
| outs.append(down @ hq) |
| out = np.stack(outs) |
|
|
| diff = np.abs(out - cpp_out) |
| scale = np.maximum(np.abs(out), 1e-3) |
| rel = diff / scale |
| print(f"layer={args.layer} expert={e} n_embd={n_embd} n_ff={n_ff}") |
| print(f" NumPy out[0:5] = {out[:5]}") |
| print(f" C++ out[0:5] = {cpp_out[:5]}") |
| print(f" max abs diff = {diff.max():.6g} max rel = {rel.max():.6g}") |
| print(f" out norm = {np.linalg.norm(out):.4f}") |
| ok = rel.max() <= args.tol |
| print(" 結果：", "PASS ✅" if ok else "FAIL ❌") |
| raise SystemExit(0 if ok else 1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
