#!/usr/bin/env python3 """Decode the packed Mach-2 Additive Medium checkpoint into a bf16 Hugging Face checkpoint. python decode.py materialize --packed-dir . --out /path/ckpt [--workers 8] [--no-base] """ import argparse import hashlib import json import math import os import numpy as np HIDDEN, INTER, NEXP, NLAYERS = 2560, 640, 512, 48 BASE_REPO, BASE_REVISION = "Qwen/Qwen3.8-Flash-Next", "de4b8e4d43b917e7706784d8bb445c9af86a3540" FORMAT = "fn_packed_v2" def _hmat(q_elems, chi, sub): n = len(q_elems) Q = np.array([[chi(sub(q_elems[i], q_elems[j])) for j in range(n)] for i in range(n)], np.float64) C = np.zeros((n + 1, n + 1)) C[0, 1:] = 1.0 C[1:, 0] = 1.0 C[1:, 1:] = Q A = np.array([[1.0, 1.0], [1.0, -1.0]]) B = np.array([[1.0, -1.0], [-1.0, -1.0]]) H = np.kron(C, A) + np.kron(np.eye(n + 1), B) r = 2 * (n + 1) assert np.array_equal(H, H.T) and np.array_equal(H @ H, r * np.eye(r)) return H.astype(np.float32) def hadamard_12(): chi5 = {0: 0, 1: 1, 2: -1, 3: -1, 4: 1} return _hmat(list(range(5)), lambda e: chi5[e], lambda u, v: (u - v) % 5) def hadamard_20(): els = [(a, b) for a in range(3) for b in range(3)] def mul(u, v): (a, b), (c, d) = u, v return ((a * c - b * d) % 3, (a * d + b * c) % 3) squares = {mul(e, e) for e in els if e != (0, 0)} return _hmat(els, lambda e: 0 if e == (0, 0) else (1 if e in squares else -1), lambda u, v: ((u[0] - v[0]) % 3, (u[1] - v[1]) % 3)) _HS = { 12: ["+-----------", "++-+---+++-+", "+++-+---+++-", "+-++-+---+++", "++-++-+---++", "+++-++-+---+", "++++-++-+---", "+-+++-++-+--", "+--+++-++-+-", "+---+++-++-+", "++---+++-++-", "+-+---+++-++"], 20: ["+-------------------", "++-++----+-+-++++--+", "+++-++----+-+-++++--", "+-++-++----+-+-++++-", "+--++-++----+-+-++++", "++--++-++----+-+-+++", "+++--++-++----+-+-++", "++++--++-++----+-+-+", "+++++--++-++----+-+-", "+-++++--++-++----+-+", "++-++++--++-++----+-", "+-+-++++--++-++----+", "++-+-++++--++-++----", "+-+-+-++++--++-++---", "+--+-+-++++--++-++--", "+---+-+-++++--++-++-", "+----+-+-++++--++-++", "++----+-+-++++--++-+", "+++----+-+-++++--++-", "+-++----+-+-++++--++"], } _HR = {} def _hr(kind, radix): if (kind, radix) not in _HR: if kind == "spine": H = np.array([[1.0 if c == "+" else -1.0 for c in r] for r in _HS[radix]], np.float32) assert np.array_equal(H @ H.T, radix * np.eye(radix, dtype=np.float32)) else: H = hadamard_12() if radix == 12 else hadamard_20() _HR[(kind, radix)] = H return _HR[(kind, radix)] def hadamard(x, kind="expert"): x = np.asarray(x, np.float32) shape, N = x.shape, x.shape[-1] if N & (N - 1): radix = 12 if (N % 12 == 0 and (N // 12) & (N // 12 - 1) == 0) else 20 M = N // radix assert N % radix == 0 and M & (M - 1) == 0, f"Hadamard needs 2^k, 12*2^k or 20*2^k, got {N}" if kind == "spine": xb = _butterflies(x.reshape(-1, radix, M)) if M > 1 else x.reshape(-1, radix, 1) return (np.matmul(_hr(kind, radix), xb) / np.float32(math.sqrt(N))).reshape(shape) xb = hadamard(x.reshape(-1, radix, M)) if M > 1 else x.reshape(-1, radix, 1) return (np.matmul(_hr(kind, radix), xb) / np.float32(math.sqrt(radix))).reshape(shape) return (_butterflies(x) / np.float32(math.sqrt(N))).reshape(shape) def _butterflies(x): N = x.shape[-1] cur = np.ascontiguousarray(x, np.float32).reshape(-1, N) span = 1 while span < N: blk = cur.reshape(-1, N // (2 * span), 2, span) cur = np.stack([blk[:, :, 0] + blk[:, :, 1], blk[:, :, 0] - blk[:, :, 1]], 2).reshape(-1, N) span *= 2 return cur.reshape(x.shape) def bf16_round(x): u = np.ascontiguousarray(x, np.float32).view(np.uint32) return ((u + np.uint32(0x7FFF) + ((u >> np.uint32(16)) & np.uint32(1))) & np.uint32(0xFFFF0000)).view(np.float32) def _read(path): from safetensors import safe_open with safe_open(path, framework="np") as fh: return {k: fh.get_tensor(k) for k in fh.keys()}, dict(fh.metadata() or {}) EX_V, EX_L, EX_TD = 4, 16, 16 EX_SHAPES = {"gate": (INTER, HIDDEN), "up": (INTER, HIDDEN), "down": (HIDDEN, INTER)} def _ex_states(words, K4): words = np.asarray(words).view(np.uint16).astype(np.int64) rows, nstep, step = words.shape[0], EX_TD * EX_TD // EX_V, int(K4) bits = ((words[:, :, None] >> np.arange(15, -1, -1)) & 1).reshape(rows, -1)[:, :step * nstep] bits = np.concatenate([bits, bits[:, :EX_L - step]], axis=1) w = 1 << np.arange(EX_L - 1, -1, -1, dtype=np.int64) idx = np.arange(nstep)[:, None] * step + np.arange(EX_L)[None, :] return bits[:, idx] @ w def wave_map(Mb, Nb): starts = [(Mb - i - 1, Nb - 1) for i in range(Mb)] + [(0, Nb - i - 1) for i in range(Nb)] idx = np.zeros((Mb, Nb), dtype=np.int64) for w, (jm, jn) in enumerate(starts): while 0 <= jm < Mb and 0 <= jn < Nb: idx[jm, jn] = w jm, jn = jm + 1, jn - 1 return idx def sign_seed(layer, expert, proj, side, ns="canon"): return int.from_bytes(hashlib.sha256(f"{ns}|L{layer}|e{expert}|{proj}|{side}".encode()).digest()[:7], "big") _PM0, _PM1 = np.uint64(0xD2511F53), np.uint64(0xCD9E8D57) _PW0, _PW1, _PMASK = np.uint64(0x9E3779B9), np.uint64(0xBB67AE85), np.uint64(0xFFFFFFFF) def expert_signs(seeds, dim): E = len(seeds) c2 = np.tile(np.arange(dim, dtype=np.uint64), E) k0 = np.repeat(np.array([s & 0xFFFFFFFF for s in seeds], np.uint64), dim) k1 = np.repeat(np.array([s >> 32 for s in seeds], np.uint64), dim) c0 = c1 = c3 = np.zeros(E * dim, np.uint64) for r in range(10): if r: k0, k1 = (k0 + _PW0) & _PMASK, (k1 + _PW1) & _PMASK p0, p1 = _PM0 * c0, _PM1 * c2 c0, c1, c2, c3 = (p1 >> np.uint64(32)) ^ c1 ^ k0, p1 & _PMASK, (p0 >> np.uint64(32)) ^ c3 ^ k1, p0 & _PMASK inv2pi = np.float32(2 * np.pi / 2 ** 32) v = (c1.astype(np.float32) * inv2pi + inv2pi / np.float32(2)).astype(np.float32) u = (c0.astype(np.float32) * np.float32(2.0 ** -32) + np.float32(2.0 ** -33)).astype(np.float32) return np.where((v <= np.float32(np.pi)) | (u == np.float32(1.0)), 1.0, -1.0).astype(np.float32).reshape(E, dim) def _lut(cb, K4): return np.asarray(cb.get(f"lut.k{K4}", cb["lut"]), np.float32) def decode_expert(t, cb, proj, e): m, n = EX_SHAPES[proj] K4 = int(t[f"{proj}.rate_k4"][e]) ex = np.asarray(t[f"k{K4}.{proj}.experts"]) row = int(np.searchsorted(ex, e)) assert row < ex.size and ex[row] == e, (proj, e, K4) states = _ex_states(t[f"k{K4}.{proj}.trellis"][row], K4) Mb, Nb = m // EX_TD, n // EX_TD unit = _lut(cb, K4)[states].reshape(Mb, Nb, EX_TD, EX_TD).transpose(0, 2, 1, 3).reshape(m, n) g = np.asarray(t[f"{proj}.wave_gamma"][e], np.float32)[wave_map(Mb, Nb)] unit = (unit.reshape(Mb, EX_TD, Nb, EX_TD) * g[:, None, :, None]).reshape(m, n) unit = unit * np.float32(t[f"{proj}.wscale"][e]) rows = hadamard(unit) * t[f"{proj}.SU"][e] cols = hadamard(np.ascontiguousarray(rows.T)) * t[f"{proj}.SV"][e] return np.ascontiguousarray(cols.T) def load_expert_layer(packed_dir, layer): t, meta = _read(os.path.join(packed_dir, "experts", f"L{layer:02d}.safetensors")) assert meta.get("format") == FORMAT, f"L{layer}: not an {FORMAT} expert shard ({meta.get('format')})" cb = _read(os.path.join(packed_dir, "experts", "codebook.safetensors"))[0] return with_signs(t, layer), cb def with_signs(t, layer): for p, (m, n) in EX_SHAPES.items(): if f"{p}.SU" not in t: t[f"{p}.SU"] = expert_signs([sign_seed(layer, e, p, "SU") for e in range(NEXP)], n) t[f"{p}.SV"] = expert_signs([sign_seed(layer, e, p, "SV") for e in range(NEXP)], m) return t def decode_expert_layer(packed_dir, layer, experts=None): t, cb = load_expert_layer(packed_dir, layer) experts = list(range(NEXP)) if experts is None else list(experts) gu = np.empty((len(experts), 2 * INTER, HIDDEN), np.float32) dn = np.empty((len(experts), HIDDEN, INTER), np.float32) for i, e in enumerate(experts): gu[i, :INTER] = decode_expert(t, cb, "gate", e) gu[i, INTER:] = decode_expert(t, cb, "up", e) dn[i] = decode_expert(t, cb, "down", e) return {"gate_up_proj": gu, "down_proj": dn} NE_K, NE_L, NE_V, NE_TLUT_BITS, NE_TD = 4, 16, 2, 9, 16 _FULL_LUT = {} def _ne_full_lut(tlut): key = np.asarray(tlut).tobytes() if key not in _FULL_LUT: small = np.asarray(tlut, np.float32) s = np.arange(1 << NE_L, dtype=np.int64) p = s * (s + 1) row = (p >> (16 - NE_TLUT_BITS - 1)) & ((1 << NE_TLUT_BITS) - 1) table = small[row].copy() table[:, 0] *= (1 - ((p >> 15) & 1) * 2).astype(np.float32) _FULL_LUT[key] = table return _FULL_LUT[key] def _ne_states(words): words = np.asarray(words).view(np.uint16).astype(np.int64) rows, T = words.shape[0], NE_TD * NE_TD step, nstep = NE_K * NE_V, T // NE_V bits = ((words[:, :, None] >> np.arange(15, -1, -1)) & 1).reshape(rows, -1)[:, :T * NE_K] bits = np.concatenate([bits, bits[:, :NE_L - step]], axis=1) w = 1 << np.arange(NE_L - 1, -1, -1, dtype=np.int64) idx = np.arange(nstep)[:, None] * step + np.arange(NE_L)[None, :] return bits[:, idx] @ w def decode_ne_tensor(t, name, m, n, tlut): states = _ne_states(t[f"{name}|trellis"]) vals = _ne_full_lut(tlut)[states] unit = np.ascontiguousarray(vals.reshape(m // NE_TD, n // NE_TD, NE_TD, NE_TD).transpose(0, 2, 1, 3)).reshape(m, n) unit = unit * np.float32(np.asarray(t[f"{name}|Wscale"]).reshape(-1)[0]) su, sv = np.asarray(t[f"{name}|SU"]), np.asarray(t[f"{name}|SV"]) rows = hadamard(unit, "spine") * np.sign(su).astype(np.float32) cols = hadamard(np.ascontiguousarray(rows.T), "spine") * np.sign(sv).astype(np.float32) w = np.ascontiguousarray(cols.T) mx = np.asarray(t[f"{name}|rc_max"], np.float32) r, c = rc_grid(mx[0], np.abs(sv.astype(np.int32))), rc_grid(mx[1], np.abs(su.astype(np.int32))) return (r[:, None] * bf16_round(w)) * c[None, :] def rc_grid(mx, k): return (np.float32(mx) * k.astype(np.float32)).astype(np.float32) * np.float32(1.0 / 127.0) def decode_ne_shard(packed_dir, layer): t, meta = _read(os.path.join(packed_dir, "ne", f"L{layer:02d}.safetensors")) tlut = _read(os.path.join(packed_dir, "ne", "tlut.safetensors"))[0]["tlut"] dims = json.loads(meta["dims"]) return {name: decode_ne_tensor(t, name, d[2], d[3], tlut)[:d[0], :d[1]] for name, d in dims.items()} def unpack_int5(qp, n): m = qp.shape[0] by = np.asarray(qp).reshape(m, n // 8, 5) full = np.zeros((m, n // 8, 8), dtype=np.uint8) full[:, :, :5] = by word = full.reshape(m, n).view("> np.uint64(5 * i)) & np.uint64(31)).astype(np.int8) - 16 return out.reshape(m, n) def decode_head(packed_dir): d = os.path.join(packed_dir, "head") parts = [] for f in sorted(x for x in os.listdir(d) if x.startswith("head_c") and x.endswith(".safetensors")): t, meta = _read(os.path.join(d, f)) parts += [(int(name.split(":")[1]), m0, n0, int(meta.get("group", 64)), t, name) for name, (m0, n0) in json.loads(meta["dims"]).items()] parts.sort(key=lambda x: x[0]) out = np.empty((sum(p[1] for p in parts), parts[0][2]), np.float32) for r0, m0, n0, g, t, name in parts: q = unpack_int5(t[f"{name}|qp"], n0).astype(np.float32) out[r0:r0 + m0] = q * np.repeat(np.asarray(t[f"{name}|gscale"], np.float32), g, axis=1)[:, :n0] return out def decode_embed(packed_dir, bits=4, rows_per=16384): t, _ = _read(os.path.join(packed_dir, "ne", f"embed_int{bits}.safetensors")) rows, ng = t["mn"].shape hid = t["q_packed"].shape[1] * 8 // bits out = np.empty((rows, hid), np.float32) for r0 in range(0, rows, rows_per): sl = slice(r0, min(r0 + rows_per, rows)) b = np.unpackbits(t["q_packed"][sl], axis=1, count=hid * bits).reshape(-1, hid, bits) q = np.zeros(b.shape[:2], np.uint8) for j in range(bits): q = (q << 1) | b[..., j] mn = t["mn"][sl].astype(np.float32)[..., None] mx = t["mx"][sl].astype(np.float32)[..., None] step = np.maximum(mx - mn, np.float32(1e-8)) * np.float32(1.0 / (2 ** bits - 1)) out[sl] = (mn + q.reshape(-1, ng, hid // ng).astype(np.float32) * step).reshape(-1, hid) return out def decode_int8_rows(packed_dir): from safetensors import safe_open out = {} with safe_open(os.path.join(packed_dir, "ne", "int8_rows.safetensors"), framework="pt") as fh: names = sorted({k.split("|")[0] for k in fh.keys()}) for name in names: q = fh.get_tensor(f"{name}|q").numpy().astype(np.float32) s = fh.get_tensor(f"{name}|scale").float().numpy() out[name] = q * s return out def _bf16(a): import torch return torch.from_numpy(np.ascontiguousarray(a, np.float32)).to(torch.bfloat16) def _expert_layer_job(args): packed_dir, out, L = args import torch from safetensors.torch import save_file t, cb = load_expert_layer(packed_dir, L) gu = torch.empty((NEXP, 2 * INTER, HIDDEN), dtype=torch.bfloat16) dn = torch.empty((NEXP, HIDDEN, INTER), dtype=torch.bfloat16) for e in range(NEXP): gu[e, :INTER] = _bf16(decode_expert(t, cb, "gate", e)) gu[e, INTER:] = _bf16(decode_expert(t, cb, "up", e)) dn[e] = _bf16(decode_expert(t, cb, "down", e)) p = f"model.language_model.layers.{L}.mlp.experts." fn = f"experts-L{L:02d}.safetensors" save_file({p + "gate_up_proj": gu, p + "down_proj": dn}, os.path.join(out, fn), metadata={"format": "pt"}) return {p + "gate_up_proj": fn, p + "down_proj": fn} def _ne_layer_job(args): packed_dir, out, L = args from safetensors.torch import save_file dec = decode_ne_shard(packed_dir, L) fn = f"spine-L{L:02d}.safetensors" save_file({k: _bf16(v) for k, v in dec.items()}, os.path.join(out, fn), metadata={"format": "pt"}) return {k: fn for k in dec} def materialize(packed_dir, out, workers=4, base=True): import shutil from concurrent.futures import ProcessPoolExecutor from safetensors import safe_open from safetensors.torch import save_file os.makedirs(out, exist_ok=True) root = os.path.dirname(os.path.abspath(packed_dir.rstrip("/"))) if os.path.basename(packed_dir.rstrip("/")) == "packed" \ else packed_dir pk = os.path.join(root, "packed") wm = {} with ProcessPoolExecutor(workers) as ex: for r in ex.map(_ne_layer_job, [(pk, out, L) for L in range(NLAYERS)]): wm.update(r) for r in ex.map(_expert_layer_job, [(pk, out, L) for L in range(NLAYERS)]): wm.update(r) save_file({"lm_head.weight": _bf16(decode_head(pk))}, os.path.join(out, "head.safetensors"), metadata={"format": "pt"}) wm["lm_head.weight"] = "head.safetensors" emb = "model.language_model.embed_tokens.weight" save_file({emb: _bf16(decode_embed(pk))}, os.path.join(out, "embed.safetensors"), metadata={"format": "pt"}) wm[emb] = "embed.safetensors" i8 = decode_int8_rows(pk) save_file({k: _bf16(v) for k, v in i8.items()}, os.path.join(out, "int8rows.safetensors"), metadata={"format": "pt"}) wm.update({k: "int8rows.safetensors" for k in i8}) for f in ["extras.safetensors"] + sorted(os.path.join("packed", "table", x) for x in os.listdir(os.path.join(pk, "table")) if x.endswith(".safetensors")): dst = os.path.basename(f) shutil.copyfile(os.path.join(root, f), os.path.join(out, dst)) with safe_open(os.path.join(out, dst), framework="pt") as fh: wm.update({k: dst for k in fh.keys()}) if base: from huggingface_hub import hf_hub_download idx = json.load(open(hf_hub_download(BASE_REPO, "model.safetensors.index.json", revision=BASE_REVISION)))["weight_map"] want = {k: f for k, f in idx.items() if k.startswith("mtp.") or k.startswith("model.visual.")} for f in sorted(set(want.values())): src = hf_hub_download(BASE_REPO, f, revision=BASE_REVISION) with safe_open(src, framework="pt") as fh: ks = [k for k in fh.keys() if k in want] save_file({k: fh.get_tensor(k) for k in ks}, os.path.join(out, f"base-{f}"), metadata={"format": "pt"}) wm.update({k: f"base-{f}" for k in ks}) for f in os.listdir(root): if f.endswith((".json", ".jinja", ".txt")) and f not in ("MANIFEST.json", "model.safetensors.index.json"): shutil.copyfile(os.path.join(root, f), os.path.join(out, f)) json.dump({"metadata": {}, "weight_map": dict(sorted(wm.items()))}, open(os.path.join(out, "model.safetensors.index.json"), "w"), indent=1) print(f"MATERIALIZED {len(wm)} tensors -> {out}", flush=True) if __name__ == "__main__": ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) sub = ap.add_subparsers(dest="cmd", required=True) m = sub.add_parser("materialize") m.add_argument("--packed-dir", default=".") m.add_argument("--out", required=True) m.add_argument("--workers", type=int, default=4) m.add_argument("--no-base", action="store_true", help="text-only: skip MTP and vision weights") a = ap.parse_args() materialize(a.packed_dir, a.out, a.workers, not a.no_base)