#!/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 fcntl
import json
import os
import struct
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "vendor" / "strata" / "metal" / "tools"))
import metal_repack as R
ALIGN = 16384
CHUNK = 64 << 20
SYNC_EVERY = 1 << 30
NL = 48
NE = 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())