scripts/dec/build_next.py

"""v3.1b: next-token decoder functions n16/n32/n64/n96 -> mf_next/dec_mf_next_fp16.mlpackage (multifunction, shared weights).
Input  ids int32 [1,T]  = BOS + left text, RIGHT-padded with anything (causal model: pads never reach position `last`)
       last int32 [1]   = index of the last real token (= len(BOS + left ids) - 1)
Output lp fp32 [1,V]    = log-softmax over the 16k vocab of the NEXT token after ids[0, last] (max-shifted logsumexp).
Same QDec body as the shipped t{T} functions (fp16, ane_kwargs, t16 single-block GDN, t>=32 chunked c16).
usage: build_next.py [--merge-test]   (--merge-test: also merge with the shipped mf/parts/*fp16 into a throwaway package
to measure whether weights de-duplicate; prints sizes, then deletes it)"""
import sys, time, shutil, torch, numpy as np
import torch.nn.functional as F
from pathlib import Path
from qdec import QDec, ane_kwargs, FF

TS = (16, 32, 64, 96)


class QNext(QDec):
    def __init__(self, T, **kw):
        super().__init__(T, B=1, **kw)
        self.register_buffer("pos", torch.arange(T, dtype=torch.int32)[None, :])     # [1,T]

    def forward(self, ids, last):
        x = F.embedding(ids, self.emb)
        for m in self.layers:
            h = m.n1(x)
            x = x + (self.gdn(m, h) if m.kind == "linear_attention" else self.attn(m, h))
            gu = m.gate_up(m.n2(x))
            x = x + m.down(F.silu(gu[..., :FF]) * gu[..., FF:])
        # pick row `last` BEFORE the final RMSNorm (row-wise, so exact): one-hot mask * x, sum over T
        sel = (self.pos == last.to(torch.int32)[:, None]).to(x.dtype)                   # [1,T]
        h = self.norm((sel[:, :, None] * x).sum(1))                                      # [1,640]
        logits = h @ self.emb.t()                                                        # tied head [1,V]
        mx = logits.max(-1, keepdim=True).values                                         # max-shifted lse (ANE fp16, qdec.py)
        return logits - (mx + torch.log(torch.exp(logits - mx).sum(-1, keepdim=True)))



if __name__ == "__main__":
    import coremltools as ct
    from coremltools.models.utils import MultiFunctionDescriptor, save_multifunction
    out = Path("mf_next"); (out / "parts").mkdir(parents=True, exist_ok=True)
    size = lambda p: sum(f.stat().st_size for f in Path(p).rglob("*") if f.is_file()) / 1e6
    torch.set_num_threads(1)
    parts = {}
    for T in TS:
        m = QNext(T, chunk=16 if T > 16 else None, **ane_kwargs()).eval()
        ex = (torch.zeros((1, T), dtype=torch.long), torch.tensor([T - 1], dtype=torch.int32))
        with torch.no_grad():
            tr = torch.jit.trace(m, ex)
        t0 = time.time()
        mm = ct.convert(tr, inputs=[ct.TensorType("ids", shape=(1, T), dtype=np.int32), ct.TensorType("last", shape=(1,), dtype=np.int32)],
                        outputs=[ct.TensorType("lp", dtype=np.float32)], compute_precision=ct.precision.FLOAT16,
                        minimum_deployment_target=ct.target.macOS15, convert_to="mlprogram")
        p = out / "parts" / f"dec_b1_n{T}_fp16.mlpackage"; shutil.rmtree(p, ignore_errors=True); mm.save(str(p)); parts[T] = p
        print(f"OK part n{T} {time.time()-t0:.1f}s {size(p):.1f} MB", flush=True)

    def merge(fns, dst, default):
        d = MultiFunctionDescriptor()
        for name, p in fns.items():
            d.add_function(str(p), src_function_name="main", target_function_name=name)
        d.default_function_name = default
        shutil.rmtree(dst, ignore_errors=True); save_multifunction(d, str(dst))
        return size(dst)

    dst = out / "dec_mf_next_fp16.mlpackage"
    print(f"OK multifunction n16..n96: {merge({f'n{T}': p for T, p in parts.items()}, dst, 'n16'):.1f} MB "
          f"(parts total {sum(size(p) for p in parts.values()):.1f} MB) -> {dst}")
    if "--merge-test" in sys.argv:   # shipped t* parts + new n* parts in one package: do the weights de-duplicate?
        fns = {f"t{T}": Path(f"mf/parts/dec_b3_t{T}_fp16.mlpackage") for T in TS} | {f"n{T}": p for T, p in parts.items()}
        mt = out / "merge_test_fp16.mlpackage"
        print(f"MERGE-TEST t16..t96 + n16..n96: {merge(fns, mt, 't16'):.1f} MB (shipped mf/dec_mf_fp16 {size('mf/dec_mf_fp16.mlpackage'):.1f} MB)")
        for w in mt.rglob("*.bin"):
            print("  weight file", w.relative_to(mt), f"{w.stat().st_size/1e6:.1f} MB")
        shutil.rmtree(mt)