#!/usr/bin/env python3
"""stack-qwen38-official-fp8.py — reshape the official Qwen/Qwen3.8-Flash-Next-FP8
checkpoint into the form the shipped radiance qwen4exp plugin maps.

Run INSIDE the staging dir's tree (python3 + numpy on the host; data never
leaves the box). Two transforms, both VALUE-PRESERVING:

1. EXPERTS: the shipped plugin's stacked layout names ONE tensor per family
   (`mlp.experts.gate_up_proj` [E, 2F, H], `mlp.experts.down_proj` [E, H, F])
   plus a `<family>_scale_inv` plane at the same index (NO `.weight` in the
   stacked name — scale_name() appends _scale_inv to the tensor name verbatim),
   while the official FP8 repo ships per-expert unfused tensors
   (`experts.<e>.{gate,up,down}_proj.weight` + scale). Stacking is pure byte
   concatenation along dim 0 — gate rows
   first inside each expert's gate_up slab, which is the order the block
   wants (rad_block_moe_fp8.h: "a stacked gate_up is ALREADY FUSED in the
   order this block wants"). No requantisation: bytes are moved, not decoded.

2. PLE TABLE: the plugin's fp8 branch wants the table's codes beside a
   ONE-ELEMENT bf16 scale plane (`blk.%d.ple_ngram.scale`, view {1}); a
   per-shard-blocked scale can never fit that view, and `rtn` re-encoding the
   RAW scaleless codes would overflow (v / 1.99e-4 exceeds e4m3's 448 max).
   So staging decodes the scaleless shards with the checkpoint's OWN bf16
   scale into a bf16 table (exact: every e4m3 code value and the scale are
   exactly representable in bf16 to within the code's own 2^-3 grid), and the
   convert's rtn rule re-encodes at that same scale — the nearest-code
   distance is 2^-3 while the bf16 rounding perturbs by 2^-9, so every
   original code is restored exactly and the container gets the checkpoint's
   own codes plus the one-element scale plane.

The convert then runs with NO recipe except:
  blk.*.ple_ngram.weight=rtn:codes=fp8_e4m3,block=*x*,scale=bf16,rule=fixed,scale_value=<s>
with <s> read from the checkpoint and written to ple_scale.txt.

Idempotent: re-runs skip when the staged index already carries the stacked
names, the merged bf16 shard_0, and split_ngram_parts == 1. --force redoes.
"""

import argparse
import json
import mmap
import multiprocessing as mp
import os
import re
import struct
import sys
import time
import numpy as np
F8, BF16 = "F8_E4M3", "BF16"


def log(msg):
    print(msg, flush=True)


# ---------------------------------------------------------------- safetensors

def parse_header(path):
    """Return (header_len, {name: (dtype, shape, data_offsets)})."""
    with open(path, "rb") as fh:
        (hlen,) = struct.unpack("<Q", fh.read(8))
        hdr = json.loads(fh.read(hlen))
    return hlen, hdr


def write_safetensors(path, tensors, data_iter):
    """tensors: [(name, dtype, shape, nbytes)]; data_iter yields bytes per tensor."""
    hdr = {}
    off = 0
    for name, dtype, shape, nbytes in tensors:
        hdr[name] = {"dtype": dtype, "shape": list(shape),
                     "data_offsets": [off, off + nbytes]}
        off += nbytes
    blob = json.dumps(hdr, separators=(",", ":")).encode("utf-8")
    pad = (-len(blob)) % 8  # keep the data start 8-byte aligned
    blob += b" " * pad
    with open(path, "wb") as fh:
        fh.write(struct.pack("<Q", len(blob)))
        fh.write(blob)
        for _, _, _, nbytes in tensors:
            fh.write(next(data_iter).tobytes())
        fh.flush()
        os.fsync(fh.fileno())
    return 8 + len(blob) + sum(t[3] for t in tensors)


# ------------------------------------------------------------------- preflight

def build_meta(stage, index):
    """name -> (file, global_byte_offset, nbytes, dtype, shape) for every tensor."""
    files = sorted(set(index["weight_map"].values()))
    meta = {}
    for f in files:
        path = os.path.join(stage, f)
        hlen, hdr = parse_header(path)
        for name, e in hdr.items():
            if name == "__metadata__":
                continue
            b0, b1 = e["data_offsets"]
            meta[name] = (f, 8 + hlen + b0, b1 - b0, e["dtype"], e["shape"])
    return meta


def e4m3_lut():
    lut = np.zeros(256, np.float32)
    for b in range(256):
        sign = -1.0 if b & 0x80 else 1.0
        e = (b >> 3) & 0xF
        m = b & 0x7
        if e == 15 and m == 7:
            v = np.float32(np.nan)  # OCP e4m3fn: only 0x7F/0xFF are NaN
        elif e == 0:
            v = sign * (m / 8.0) * 2.0 ** -6
        else:
            v = sign * (1.0 + m / 8.0) * 2.0 ** (e - 7)
        lut[b] = v
    assert lut[0x38] == 1.0 and lut[0x7E] == 448.0
    assert np.isnan(lut[0x7F]) and lut[0x00] == 0.0
    assert lut[0x08] == 2.0 ** -6 and lut[0x80] == 0.0
    return lut


def f32_to_bf16_rne(v):
    """float32 array -> uint16 bf16 bits, round-to-nearest-even."""
    q = np.ascontiguousarray(v, dtype=np.float32)
    q[~np.isfinite(q)] = 0.0
    bits = q.view(np.uint32)
    r = bits + 0x7FFF + ((bits >> 16) & 1)
    return (r >> 16).astype(np.uint16)


# ------------------------------------------------------------------- workers

def stack_prefix(job):
    """Byte-stack one layer's 512 experts into a 4-tensor safetensors file."""
    (stage, pre, idx, n_e, shapes, meta) = job
    (f2, h) = shapes["ffn2_h"]          # 2*ffn, n_embd
    (f,) = shapes["ffn"]                # ffn
    gs, ds = shapes["gs"], shapes["ds"] # scale shapes [5,20] and [20,5]

    # byte widths: fp8 is 1 B/element, bf16 2 B; the fused gate_up scale is
    # the gate and up scale planes concatenated (gate rows first)
    g_bytes = (f2 // 2) * h
    gs_bytes = gs[0] * gs[1] * 2          # ONE projection's scale plane
    gus_bytes = 2 * gs_bytes              # fused scale per expert
    dn_bytes = h * f
    ds_bytes = ds[0] * ds[1] * 2

    gu = np.empty(n_e * f2 * h, np.uint8)
    gus = np.empty(n_e * gus_bytes, np.uint8)
    dn = np.empty(n_e * dn_bytes, np.uint8)
    dns = np.empty(n_e * ds_bytes, np.uint8)

    mmaps = {}

    def view(name):
        f_, off, ln, dt, shp = meta[name]
        if f_ not in mmaps:
            fh = open(os.path.join(stage, f_), "rb")
            mmaps[f_] = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
        return memoryview(mmaps[f_])[off:off + ln], (dt, shp)

    def need(name, dt, shp):
        mv, (got_dt, got_shp) = view(name)
        assert got_dt == dt and list(got_shp) == list(shp), f"{name}: {got_dt}{got_shp}"
        return mv

    for e in range(n_e):
        base = f"{pre}.mlp.experts.{e}"
        slab = gu[e * f2 * h:(e + 1) * f2 * h]
        slab[:g_bytes] = need(f"{base}.gate_proj.weight", F8, [f2 // 2, h])
        slab[g_bytes:] = need(f"{base}.up_proj.weight", F8, [f2 // 2, h])
        sslab = gus[e * gus_bytes:(e + 1) * gus_bytes]
        sslab[:gs_bytes] = need(f"{base}.gate_proj.weight_scale_inv", BF16, gs)
        sslab[gs_bytes:] = need(f"{base}.up_proj.weight_scale_inv", BF16, gs)
        dn[e * dn_bytes:(e + 1) * dn_bytes] = need(f"{base}.down_proj.weight", F8, [h, f])
        dns[e * ds_bytes:(e + 1) * ds_bytes] = need(f"{base}.down_proj.weight_scale_inv", BF16, ds)

    for mm in mmaps.values():
        mm.close()

    # byte-identity check on two sampled experts (first and last), all four
    # tensors — the stack IS the checkpoint bytes, so they must compare equal
    for e in (0, n_e - 1):
        base = f"{pre}.mlp.experts.{e}"
        s = e * f2 * h
        assert bytes(gu[s:s + g_bytes]) == bytes(stack_src(f"{base}.gate_proj.weight", meta, stage))
        assert bytes(gu[s + g_bytes:s + f2 * h]) == bytes(stack_src(f"{base}.up_proj.weight", meta, stage))
        t = e * gus_bytes
        assert bytes(gus[t:t + gs_bytes]) == bytes(stack_src(f"{base}.gate_proj.weight_scale_inv", meta, stage))
        assert bytes(gus[t + gs_bytes:t + gus_bytes]) == bytes(stack_src(f"{base}.up_proj.weight_scale_inv", meta, stage))
        d = e * dn_bytes
        assert bytes(dn[d:d + dn_bytes]) == bytes(stack_src(f"{base}.down_proj.weight", meta, stage))
        u = e * ds_bytes
        assert bytes(dns[u:u + ds_bytes]) == bytes(stack_src(f"{base}.down_proj.weight_scale_inv", meta, stage))

    out = os.path.join(stage, f"stacked.{idx:03d}.safetensors")
    tensors = [
        (f"{pre}.mlp.experts.gate_up_proj", F8, [n_e, f2, h], n_e * f2 * h),
        (f"{pre}.mlp.experts.gate_up_proj_scale_inv", BF16,
         [n_e, gs[0] * 2, gs[1]], n_e * gus_bytes),
        (f"{pre}.mlp.experts.down_proj", F8, [n_e, h, f], n_e * h * f),
        (f"{pre}.mlp.experts.down_proj_scale_inv", BF16,
         [n_e, ds[0], ds[1]], n_e * ds_bytes),
    ]
    bufs = iter([gu, gus, dn, dns])
    written = write_safetensors(out, tensors, bufs)
    return (pre, written)


def stack_src(name, meta, stage):
    """Re-open a source tensor for the sampled-expert check."""
    f_, off, ln, dt, shp = meta[name]
    fh = open(os.path.join(stage, f_), "rb")
    mm = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
    fh.close()  # mmap keeps the mapping alive
    return mm[off:off + ln]


def fill_ple_shard(job):
    """Decode one fp8 shard with the table scale, write bf16 rows into the mmap."""
    (src_path, src_off, src_rows, dst_path, dst_off, scale_f32,
     cols, chunk_rows) = job
    lut = e4m3_lut()
    with open(src_path, "rb") as fh:
        sm = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
        with open(dst_path, "r+b") as dfh:
            dm = mmap.mmap(dfh.fileno(), 0, access=mmap.ACCESS_WRITE)
            for r0 in range(0, src_rows, chunk_rows):
                r1 = min(r0 + chunk_rows, src_rows)
                codes = np.frombuffer(sm, dtype=np.uint8,
                                      count=(r1 - r0) * cols,
                                      offset=src_off + r0 * cols)
                q = lut.take(codes)
                del codes  # release the mmap export before sm.close()
                q[~np.isfinite(q)] = 0.0
                q *= np.float32(scale_f32)
                packed = f32_to_bf16_rne(q)
                dm[dst_off + r0 * cols * 2: dst_off + r1 * cols * 2] = packed.tobytes()
            dm.flush()
            os.fsync(dfh.fileno())
            dm.close()
        sm.close()


# ---------------------------------------------------------------------- main

def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--src", default=None,
                    help="raw checkpoint dir (documentation only; unused)")
    ap.add_argument("--stage", required=True, help="staging dir to reshape in place")
    ap.add_argument("--force", action="store_true", help="redo even if already staged")
    ap.add_argument("--workers", type=int, default=min(8, os.cpu_count() or 1))
    args = ap.parse_args()
    stage = args.stage
    t0 = time.time()

    index_path = os.path.join(stage, "model.safetensors.index.json")
    cfg_path = os.path.join(stage, "config.json")
    if not (os.path.isfile(index_path) and os.path.isfile(cfg_path)):
        log(f"ERROR: {stage} is not a staged checkpoint (missing index/config)")
        return 2
    with open(index_path) as fh:
        index = json.load(fh)
    with open(cfg_path) as fh:
        cfg = json.load(fh)
    wm = index["weight_map"]

    meta = build_meta(stage, index)

    # ---- collect prefixes -------------------------------------------------
    pre_re = re.compile(r"^(?P<pre>.+)\.mlp\.experts\.(?P<e>\d+)\.gate_proj\.weight$")
    pres, n_e = {}, 0
    for k in wm:
        m = pre_re.match(k)
        if m:
            pres.setdefault(m.group("pre"), set()).add(int(m.group("e")))
            n_e = max(n_e, int(m.group("e")) + 1)
    prefixes = sorted(pres)
    if not prefixes:
        log("ERROR: no per-expert tensors found in the staged index")
        return 2
    mtp = [p for p in prefixes if p.startswith("mtp.")]
    trunk = [p for p in prefixes if p.startswith("model.language_model.layers.")]
    log(f"expert prefixes: {len(prefixes)} (trunk {len(trunk)}, mtp {len(mtp)}), "
        f"{n_e} experts each")

    # shapes from expert 0 of the first prefix
    pre0 = prefixes[0]
    shapes = {}
    gdt, gsh = meta[f"{pre0}.mlp.experts.0.gate_proj.weight"][3:5]
    ddt, dsh = meta[f"{pre0}.mlp.experts.0.down_proj.weight"][3:5]
    gsd, gssh = meta[f"{pre0}.mlp.experts.0.gate_proj.weight_scale_inv"][3:5]
    dsd, dssh = meta[f"{pre0}.mlp.experts.0.down_proj.weight_scale_inv"][3:5]
    scale_law = (gdt == F8 and ddt == F8 and gsd == BF16 and dsd == BF16
                 and list(gssh) == [(gsh[0] + 127) // 128, (gsh[1] + 127) // 128]
                 and list(dssh) == [(dsh[0] + 127) // 128, (dsh[1] + 127) // 128])
    if not scale_law:
        log(f"ERROR: unexpected expert shapes: gate {gdt}{gsh} scale {gsd}{gssh} "
            f"down {ddt}{dsh} scale {dsd}{dssh} (scale law 128x128)")
        return 2
    shapes = {"ffn2_h": (gsh[0] * 2, gsh[1]), "ffn": (dsh[1],),
              "gs": tuple(gssh), "ds": tuple(dssh)}
    f2, h = shapes["ffn2_h"]
    f = shapes["ffn"][0]
    log(f"shapes: gate_up [{n_e},{f2},{h}] F8, down [{n_e},{h},{f}] F8, "
        f"scales {gssh}x2 / {dssh} bf16")

    # ---- PLE table --------------------------------------------------------
    ple_re = re.compile(r"^(?P<pre>.+)\.ple\.ple_embedding\.ngram_embedding\.shard_(?P<i>\d+)\.weight$")
    ple_pre, shard_ids = None, {}
    for k in wm:
        m = ple_re.match(k)
        if m:
            ple_pre = m.group("pre")
            shard_ids[int(m.group("i"))] = k
    if ple_pre is None:
        log("ERROR: no PLE table shards found")
        return 2
    n_shards = len(shard_ids)
    if sorted(shard_ids) != list(range(n_shards)):
        log(f"ERROR: PLE shard ids not contiguous 0..{n_shards - 1}")
        return 2
    sc_name = f"{ple_pre}.ple.ple_embedding.ngram_embedding.weight_scale"
    sdt, ssh, sln = meta[sc_name][3], meta[sc_name][4], meta[sc_name][2]
    if sdt != BF16 or list(ssh) != [1]:
        log(f"ERROR: PLE scale is {sdt}{ssh}, expected BF16 [1]")
        return 2
    s_bits = struct.unpack("<H", bytes(view_src_tuple(meta, sc_name, stage)))[0]
    scale_f32 = struct.unpack("<f", struct.pack("<I", s_bits << 16))[0]
    log(f"PLE: {ple_pre}, {n_shards} shards, scale bf16 = {scale_f32!r}")

    rows = [meta[shard_ids[i]][4][0] for i in range(n_shards)]
    total_rows = sum(rows)
    hd = meta[shard_ids[0]][4][1]
    if any(meta[shard_ids[i]][4][1] != hd for i in range(n_shards)):
        log("ERROR: PLE shards disagree on HD")
        return 2
    if hd % 128 != 0 and hd != 160:
        log(f"WARNING: PLE HD={hd} is neither 160 nor a multiple of 128")

    # ---- idempotence gate -------------------------------------------------
    stacked_names = [f"{p}.mlp.experts.gate_up_proj" for p in prefixes]
    merged_name = f"{ple_pre}.ple.ple_embedding.ngram_embedding.shard_0.weight"
    already = (all(n in wm for n in stacked_names)
               and merged_name in wm
               and find_split(cfg) is not None and read_split(cfg) == 1)
    if already and not args.force:
        log("already staged (stacked names, merged shard_0 and split_ngram_parts=1 "
            "all present) — nothing to do")
        return 0

    # ---- disk gate --------------------------------------------------------
    stacked_bytes = len(prefixes) * n_e * (f2 * h + h * f)
    ple_bytes = total_rows * hd * 2
    need = stacked_bytes + ple_bytes
    st = os.statvfs(stage)
    avail = st.f_bavail * st.f_frsize
    if avail < need * 1.05:
        log(f"ERROR: {avail / 2**30:.1f} GiB free < {need / 2**30:.1f} GiB needed (+5%)")
        return 2
    log(f"disk: {avail / 2**30:.1f} GiB free, need {need / 2**30:.1f} GiB "
        f"(experts {stacked_bytes / 2**30:.1f}, ple {ple_bytes / 2**30:.1f})")

    # ---- stage experts ----------------------------------------------------
    jobs = []
    for idx, pre in enumerate(prefixes):
        jmeta = {n: meta[n] for n in meta if n.startswith(f"{pre}.mlp.experts.")}
        jobs.append((stage, pre, idx, n_e, shapes, jmeta))
    ctx = mp.get_context("fork")
    written_total = 0
    with ctx.Pool(args.workers) as pool:
        for pre, nbytes in pool.imap_unordered(stack_prefix, jobs):
            written_total += nbytes
            log(f"  stacked {pre} ({nbytes / 2**30:.2f} GiB)")
    log(f"experts stacked: {written_total / 2**30:.2f} GiB in {len(prefixes)} files")

    # ---- stage PLE --------------------------------------------------------
    # The file name sorts BEFORE the model-000NN shards ON PURPOSE: the reader
    # opens every index-named file in sorted order and emplaces EVERY tensor
    # it physically holds, first-wins (checkpoint.cpp add_shard). The raw
    # shards carry tensors of the SAME name as the merged one
    # (...ple.ple_embedding.ngram_embedding.shard_0.weight, [2500012,160]
    # fp8), so ple.merged.safetensors lost the race and the plan read one raw
    # shard as the codes. "00-" wins the sort and the name.
    out_name = "00-ple.merged.safetensors"
    out_path = os.path.join(stage, out_name)
    data_bytes = total_rows * hd * 2
    hdr = {merged_name: {"dtype": BF16, "shape": [total_rows, hd],
                         "data_offsets": [0, data_bytes]}}
    blob = json.dumps(hdr, separators=(",", ":")).encode("utf-8")
    pad = (-len(blob)) % 8
    blob += b" " * pad
    hlen = len(blob)
    with open(out_path, "wb") as fh:
        fh.write(struct.pack("<Q", hlen))
        fh.write(blob)
        fh.truncate(8 + hlen + data_bytes)
    jobs = []
    row0 = 0
    for i in range(n_shards):
        f_, off, ln, _, shp = meta[shard_ids[i]]
        jobs.append((os.path.join(stage, f_), off, shp[0], out_path,
                     8 + hlen + row0 * hd * 2, scale_f32, hd, 131072))
        row0 += shp[0]
    assert row0 == total_rows
    with ctx.Pool(args.workers) as pool:
        pool.map(fill_ple_shard, jobs)
    log(f"PLE merged: {data_bytes / 2**30:.2f} GiB bf16 "
        f"({total_rows} rows x {hd} cols)")

    # ---- post-verify ------------------------------------------------------
    verify(stage, prefixes, n_e, shapes, meta, shard_ids, rows, hd,
           scale_f32, out_path, hlen)

    # ---- rewrite index + config -------------------------------------------
    new_wm = {k: v for k, v in wm.items() if k not in set(shard_ids.values())}
    for idx, pre in enumerate(prefixes):
        fname = f"stacked.{idx:03d}.safetensors"
        for t in (f"{pre}.mlp.experts.gate_up_proj",
                  f"{pre}.mlp.experts.gate_up_proj_scale_inv",
                  f"{pre}.mlp.experts.down_proj",
                  f"{pre}.mlp.experts.down_proj_scale_inv"):
            new_wm[t] = fname
    new_wm[merged_name] = out_name
    total = sum(meta_size(meta, k) for k in new_wm
                if k in meta)  # originals
    # new tensors' sizes
    total += written_total + data_bytes
    index["weight_map"] = new_wm
    if "metadata" in index and isinstance(index["metadata"], dict):
        index["metadata"]["total_size"] = total
    # ATOMIC REPLACE, never open(path, "w"): the stage's index is a symlink
    # to the RAW checkpoint (staging symlinks everything but config.json) —
    # open("w") would rewrite the raw file through the link. os.replace
    # swaps the link itself for a real file; the raw tree stays pristine.
    for path, payload in ((index_path, json.dumps(index, indent=2)),
                          (cfg_path, json.dumps(cfg, indent=2))):
        tmp = path + ".tmp"
        with open(tmp, "w") as fh:
            fh.write(payload)
        os.replace(tmp, path)
    n = set_split(cfg, 1)
    log(f"config: split_ngram_parts=1 written at {n} location(s)")

    with open(os.path.join(stage, "ple_scale.txt"), "w") as fh:
        fh.write(repr(scale_f32) + "\n")

    dt = time.time() - t0
    log(f"staged in {dt / 60:.1f} min: {written_total / 2**30:.1f} GiB experts + "
        f"{data_bytes / 2**30:.1f} GiB PLE; scale {scale_f32!r} -> ple_scale.txt")
    return 0


def meta_size(meta, name):
    _, _, ln, _, _ = meta[name]
    return ln


def view_src_tuple(meta, name, stage):
    f_, off, ln, _, _ = meta[name]
    with open(os.path.join(stage, f_), "rb") as fh:
        fh.seek(off)
        return fh.read(ln)


def find_split(cfg):
    """Locate split_ngram_parts keys (recursive); return list of paths."""
    hits = []

    def walk(d, path):
        if isinstance(d, dict):
            for k, v in d.items():
                if k == "split_ngram_parts":
                    hits.append(path + [k])
                walk(v, path + [k])
    walk(cfg, [])
    return hits


def read_split(cfg):
    hits = find_split(cfg)
    if not hits:
        return None
    d = cfg
    for k in hits[0][:-1]:
        d = d[k]
    return d[hits[0][-1]]


def set_split(cfg, value):
    hits = find_split(cfg)
    for path in hits:
        d = cfg
        for k in path[:-1]:
            d = d[k]
        d[path[-1]] = value
    if not hits:  # plugin's meta_int2 falls back to the top-level key
        cfg["split_ngram_parts"] = value
        hits.append(["split_ngram_parts"])
    return hits


def verify(stage, prefixes, n_e, shapes, meta, shard_ids, rows, hd,
           scale_f32, out_path, hlen):
    """Cheap strong checks: headers parse, sizes match, sampled bytes identical."""
    f2, h = shapes["ffn2_h"]
    f = shapes["ffn"][0]
    gs, ds = shapes["gs"], shapes["ds"]
    # headers of every new file
    for idx, pre in enumerate(prefixes):
        path = os.path.join(stage, f"stacked.{idx:03d}.safetensors")
        hlen2, hdr = parse_header(path)
        want = {
            f"{pre}.mlp.experts.gate_up_proj": (F8, [n_e, f2, h]),
            f"{pre}.mlp.experts.gate_up_proj_scale_inv":
                (BF16, [n_e, gs[0] * 2, gs[1]]),
            f"{pre}.mlp.experts.down_proj": (F8, [n_e, h, f]),
            f"{pre}.mlp.experts.down_proj_scale_inv": (BF16, [n_e, ds[0], ds[1]]),
        }
        for name, (dt, shp) in want.items():
            e = hdr[name]
            assert e["dtype"] == dt and list(e["shape"]) == shp, name
            assert e["data_offsets"][1] - e["data_offsets"][0] == \
                int(np.prod(shp)) * (1 if dt == F8 else 2), name
    _, phdr = parse_header(out_path)
    e = phdr[next(iter(phdr))]
    assert e["dtype"] == BF16 and list(e["shape"]) == [sum(rows), hd]

    # sampled PLE rows: recompute and compare bytes exactly
    lut = e4m3_lut()
    out_off = 8 + hlen
    for i in (0, len(shard_ids) // 2, len(shard_ids) - 1):
        name = shard_ids[i]
        f_, off, ln, _, shp = meta[name]
        with open(os.path.join(stage, f_), "rb") as fh:
            mm = mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ)
            for r in (0, shp[0] - 1):
                codes = np.frombuffer(mm, dtype=np.uint8, count=hd,
                                      offset=off + r * hd)
                q = lut.take(codes) * np.float32(scale_f32)
                del codes  # release the mmap export before mm.close()
                want = f32_to_bf16_rne(q).tobytes()
                got_row = sum(rows[:i]) + r
                with open(out_path, "rb") as ofh:
                    ofh.seek(out_off + got_row * hd * 2)
                    got = ofh.read(hd * 2)
                assert got == want, f"PLE row {got_row} mismatch"
            mm.close()
    log("verify: new-file headers, sizes and sampled PLE rows all exact; "
        "expert slabs were byte-checked in-stack")


if __name__ == "__main__":
    sys.exit(main())