Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286#!/usr/bin/env python3"""Prepare pruned K=256 Qwen3.8-Flash-Next for Strata Metal engine:1. Generates K256-expert-profile.bin (STRP v1, 48 layers x 256 experts) ranked by multi-domain calibration frequency & REAP saliency (results/reap_saliency_v2.json).2. Writes K256-metal.gguf (4.22 GiB = 4.53 GB dense tensors: 2.90 GB active GPU dense + 1.62 GB skipped BF16 HC / CPU token_embd, 16 KiB aligned, F_NOCACHE) so Strata wraps dense weights zero-copy as clean file-backed pages (0 GiB anonymous RAM).3. Writes K256-metal.gguf.experts (23.42 GiB STRX v1 contiguous [gate | up | down] per-expert slices, F_NOCACHE) so each missed expert is read in 1 contiguous pread."""
import fcntlimport jsonimport osimport structimport sysimport timefrom pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "vendor" / "strata" / "metal" / "tools"))import metal_repack as R
ALIGN = 16384CHUNK = 64 << 20SYNC_EVERY = 1 << 30NL = 48NE = 256
def nocache(f) -> None: try: fcntl.fcntl(f.fileno(), getattr(fcntl, "F_NOCACHE", 48), 1) except Exception: pass
def build_expert_profile(saliency_path: Path, keep_map_path: Path, out_path: Path) -> None: sal = json.loads(saliency_path.read_text()) km = json.loads(keep_map_path.read_text()) scored = [] for l in range(NL): l_str = str(l) kept_orig = km[l_str] assert len(kept_orig) == NE, f"Layer {l} has {len(kept_orig)} experts != {NE}" counts = sal["layers"][l_str]["count"] probs = sal["layers"][l_str]["prob_sum"] for new_e, orig_e in enumerate(kept_orig): # Rank by empirical token selection count + prob_sum tie-breaker score = float(counts[orig_e]) + 1e-3 * float(probs[orig_e]) scored.append((score, l, new_e)) scored.sort(key=lambda x: (-x[0], x[1], x[2])) assert len(scored) == NL * NE
buf = bytearray() buf += struct.pack("<4s5I", b"STRP", 1, NL, NE, len(scored), len(scored)) for _, l, new_e in scored: buf += struct.pack("<2H", l, new_e) out_path.write_bytes(bytes(buf)) print(f"[1/3] Wrote STRP profile {out_path} ({len(buf)} bytes, top pair=L{scored[0][1]}:E{scored[0][2]} hits={scored[0][0]:.0f})", flush=True)
def build_metal_gguf_and_experts(shard1: Path, out_gguf: Path, out_exp: Path) -> None: t0 = time.monotonic() shards = [R.Reader(p) for p in R.shard_paths(shard1)] every = [t for s in shards for t in s.tensors] keep = [t for t in every if t["name"] != R.PLE]
def layer_of(t): import re m = re.match(r"blk\.(\d+)\.", t["name"]) return int(m.group(1)) if m else -1
dense = [t for t in keep if "_exps." not in t["name"]] role = {"ffn_gate_exps": 0, "ffn_up_exps": 1, "ffn_down_exps": 2} exps = sorted( (t for t in keep if "_exps." in t["name"]), key=lambda t: (layer_of(t), role.get(t["name"].split(".")[2], 9)), ) order = dense + exps
kv = [raw for key, raw in shards[0].kv if not key.startswith("split.") and key != "general.alignment"] kv.append(R.gguf_str("general.alignment") + struct.pack("<II", 4, ALIGN))
off = 0 infos = [] dense_end_rel = 0 for t in order: t["new"] = off infos.append( R.gguf_str(t["name"]) + struct.pack("<I", len(t["dims"])) + struct.pack(f"<{len(t['dims'])}Q", *t["dims"]) + struct.pack("<IQ", t["type"], off) ) off = (off + t["bytes"] + ALIGN - 1) // ALIGN * ALIGN if "_exps." not in t["name"]: dense_end_rel = off
head = struct.pack("<4sIQQ", b"GGUF", 3, len(order), len(kv)) + b"".join(kv) + b"".join(infos) data0 = (len(head) + ALIGN - 1) // ALIGN * ALIGN dense_total = data0 + dense_end_rel
if out_gguf.exists() and out_gguf.stat().st_size == dense_total: print(f"[2/3] {out_gguf} already exists ({dense_total / 2**30:.2f} GiB)", flush=True) else: tmp_g = out_gguf.with_name(out_gguf.name + ".part") with open(tmp_g, "wb", buffering=0) as o: nocache(o) R.write_all(o, head) R.write_all(o, b"\0" * (data0 - len(head))) done = 0 since = 0 for t in dense: o.seek(data0 + t["new"]) src = t["src"].f src.seek(t["abs"]) left = t["bytes"] while left: buf = src.read(min(left, CHUNK)) if not buf: raise RuntimeError(f"Short read in {t['name']}") R.write_all(o, buf) left -= len(buf) since += len(buf) if since >= SYNC_EVERY: os.fsync(o.fileno()) since = 0 done += t["bytes"] o.truncate(dense_total) os.fsync(o.fileno()) os.replace(tmp_g, out_gguf) print(f"[2/3] Wrote dense-only {out_gguf} ({dense_total / 2**30:.2f} GiB in {time.monotonic() - t0:.1f}s)", flush=True)
# 3. Write STRX .experts file (contiguous [gate_e | up_e | down_e] per expert) t1 = time.monotonic() tab = {t["name"]: t for t in exps} lay = [] for l in range(NL): g, u, d = (tab[f"blk.{l}.ffn_{k}_exps.weight"] for k in ("gate", "up", "down")) gs, us, ds = g["bytes"] // NE, u["bytes"] // NE, d["bytes"] // NE lay.append((g, u, d, gs, us, ds))
exp_head_len = 64 + NL * 32 exp_data0 = (exp_head_len + ALIGN - 1) // ALIGN * ALIGN offs, at = [], exp_data0 for g, u, d, gs, us, ds in lay: offs.append(at) at += NE * (gs + us + ds) at = (at + ALIGN - 1) // ALIGN * ALIGN
if out_exp.exists() and out_exp.stat().st_size == at: print(f"[3/3] {out_exp} already exists ({at / 2**30:.2f} GiB)", flush=True) else: exp_head = struct.pack("<4sIIIQ", b"STRX", 1, NL, NE, dense_total).ljust(64, b"\0") for (g, u, d, gs, us, ds), o_off in zip(lay, offs): exp_head += struct.pack("<QQQQ", o_off, gs, us, ds)
tmp_e = out_exp.with_name(out_exp.name + ".part") with open(tmp_e, "wb", buffering=0) as f: nocache(f) R.write_all(f, exp_head.ljust(exp_data0, b"\0")) for l, ((g, u, d, gs, us, ds), o_off) in enumerate(zip(lay, offs)): g["src"].f.seek(g["abs"]); G = g["src"].f.read(g["bytes"]) u["src"].f.seek(u["abs"]); U = u["src"].f.read(u["bytes"]) d["src"].f.seek(d["abs"]); D = d["src"].f.read(d["bytes"]) f.seek(o_off) layer_blob = b"".join( G[e * gs : (e + 1) * gs] + U[e * us : (e + 1) * us] + D[e * ds : (e + 1) * ds] for e in range(NE) ) R.write_all(f, layer_blob) if (l + 1) % 8 == 0 or l == NL - 1: os.fsync(f.fileno()) print(f" STRX layer {l + 1}/{NL} ({f.tell() / 2**30:.2f} GiB)", flush=True) f.truncate(at) os.fsync(f.fileno()) os.replace(tmp_e, out_exp) print(f"[3/3] Wrote STRX {out_exp} ({at / 2**30:.2f} GiB in {time.monotonic() - t1:.1f}s)", flush=True)
for s in shards: s.f.close()
def main() -> int: root = Path(__file__).resolve().parent.parent model_dir = root / "models" / "K256_reap_l2_sum" shard1 = model_dir / "Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf" out_prof = model_dir / "K256-expert-profile.bin" out_gguf = model_dir / "K256-metal.gguf" out_exp = model_dir / "K256-metal.gguf.experts"
build_expert_profile( root / "results" / "reap_saliency_v2.json", root / "results" / "keep_maps" / "keep_reap_l2_sum_k256.json", out_prof, ) build_metal_gguf_and_experts(shard1, out_gguf, out_exp)
tok_dir = model_dir / "tokenizer" if not (tok_dir / "vocab.json").exists(): sys.path.insert(0, str(root / "vendor" / "strata" / "tools")) import strata_tokenizer as ST ST.extract(str(out_gguf), str(model_dir)) jinja_p = tok_dir / "chat_template.jinja" if jinja_p.exists(): jtxt = jinja_p.read_text(encoding="utf-8") if "reasoning_effort|default('xhigh')" in jtxt: jinja_p.write_text(jtxt.replace("reasoning_effort|default('xhigh')", "reasoning_effort|default('low')"), encoding="utf-8")
healed_gguf = model_dir / "K256-q2_0-healed-metal.gguf" active_gguf = healed_gguf if healed_gguf.exists() else out_gguf vocab_out = model_dir / "vocab.gguf" if not vocab_out.exists(): rd_v = R.Reader(active_gguf) rd_v.f.close() vocab_out.write_bytes(struct.pack("<4sIQQ", b"GGUF", rd_v.version, 0, len(rd_v.kv)) + b"".join(raw for _, raw in rd_v.kv)) del rd_v vexe = root / "vendor" / "strata" / "build-vision" / "bin" / "strata-vision" mmproj_p = root / "models" / "K512_iq2_xs" / "mmproj-Swift-Qwen3.8-Flash-Next-BF16.gguf" if not mmproj_p.exists(): mmproj_p = model_dir / "mmproj-Swift-Qwen3.8-Flash-Next-BF16.gguf" has_vision = vexe.exists() and mmproj_p.exists() cfg = { "exe": str(root / "vendor" / "strata" / "build-metal" / "strata_window_cli"), "args": [ "--metal-gguf", str(active_gguf), "--ple-gguf", str(shard1), "--n-expert", "256", "--max-context", "131072", "--cache-disk", str(model_dir / "cache_disk"), "--cache-disk-size", "16384", ] + (["--vision"] if has_vision else []), "cwd": str(root / "vendor" / "strata"), "tokenizer": str(tok_dir), "model_name": "swift-1.5-k256-healed" if healed_gguf.exists() else "swift-1.5-k256", "log": "/tmp/strata-k256-serve.log", "lib_dirs": [], "port": 8080, "backend": "metal", "fit_max_tokens": True, "env": { "STRATA_METAL_ARENA_GB": "6.0", "STRATA_METAL_MLOCK_DENSE": "1", "STRATA_METAL_WIRE_ARENA": "1", "STRATA_METAL_EXPERT_PROFILE": str(out_prof), "STRATA_METAL_MOCE_BONUS": "0.85", "STRATA_METAL_EXPERT_SPLIT": "1", "STRATA_METAL_SHEXP_OVERLAP": "1", "STRATA_METAL_SPLIT_TURN": "1", "STRATA_METAL_PRED_Q8": "1", "STRATA_METAL_ROUTER_BF16": "1", "STRATA_METAL_MOE_DN_TILE": "8,2,4", "STRATA_METAL_PREFETCH": "6", "STRATA_METAL_PREFETCH2": "2", "STRATA_METAL_IO_THREADS": "8", "STRATA_METAL_IO_URGENT": "2", "STRATA_ANTI_SLOP_PENALTY": "6.0", }, "sampling": { "temperature": 0.6, "top_p": 0.95, "top_k": 40, "min_p": 0.05, "presence_penalty": 0.4, "repeat_penalty": 1.08, }, "reasoning_effort": "low", } if has_vision: cfg["vram_elastic"] = True cfg["vision"] = { "exe": str(vexe), "mmproj": str(mmproj_p), "model": str(vocab_out), "gpu": True, "max_tokens": 1024, "elastic": True, "elastic_idle_s": 60, } (model_dir / "strata-k256.json").write_text(json.dumps(cfg, indent=2) + "\n", encoding="utf-8") return 0
if __name__ == "__main__": raise SystemExit(main())