scripts/cpu_bench.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "torch",
#   "numpy",
#   "datasets>=2.19",
#   "huggingface_hub>=0.24",
#   "safetensors>=0.4",
# ]
# ///
"""CPU inference benchmark: Trident (O(1) state) vs the matched Transformer.

The product claim is CPU-only inference with bounded memory. This measures, on
actual CPU (no GPU), for both models:
  * throughput (kB/s) scoring bytes, in fp32 and int8-dynamic-quantized form,
  * quality retention under int8 (BPB fp32 vs int8 - a speedup only counts if
    quality holds),
  * how wall-time scales with context length,
  * the memory each needs to condition on a long context (Trident's constant
    state vs the transformer's O(context) KV cache).

int8 dynamic quantization (weights int8, activations quantized per-op at
runtime) is the honest "fortify the CPU moat" lever: it is exactly the kind of
optimization that matters for CPU-only deployment and is applied identically to
both models so the comparison stays fair.

Configurable via TRIDENT_CKPT / XF_CKPT so the same bench serves the 44M and
150M checkpoints. Pushes results/cpu_bench.json (or RESULT_NAME).
"""
import json
import math
import os
import sys
import time
from pathlib import Path

import torch
import torch.nn as nn


def env(k, d=None):
    v = os.environ.get(k)
    return v if v not in (None, "") else d


def log(m):
    print(f"[cpu] {m}", flush=True)


def median_time(fn, reps=3):
    ts = []
    for _ in range(reps):
        t0 = time.perf_counter()
        fn()
        ts.append(time.perf_counter() - t0)
    ts.sort()
    return ts[len(ts) // 2]


def quantize(model):
    """int8 dynamic quantization of all Linear layers (weights int8, dynamic
    activation quant). Leaves everything else in fp32."""
    return torch.ao.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)


def build_stream(dataset, name, split, field, total):
    from datasets import load_dataset
    ds = load_dataset(dataset, name=name, split=split, streaming=True)
    buf = bytearray()
    for row in ds:
        t = row.get(field)
        if not t:
            continue
        buf.extend(t.encode("utf-8"))
        buf.extend(b"\n\n")
        if len(buf) >= total:
            break
    return bytes(buf[:total])


def int8_bytes(model):
    """Approximate weight memory after int8 quantization: quantized Linear
    weights are 1 byte/elt; everything else stays fp32 (4 bytes/elt)."""
    lin = 0
    other = 0
    for m in model.modules():
        if isinstance(m, nn.Linear):
            lin += m.weight.numel()
            if m.bias is not None:
                other += m.bias.numel() * 4
        elif isinstance(m, nn.Embedding):
            other += m.weight.numel() * 4
    return lin * 1 + other


def main():
    torch.set_grad_enabled(False)
    repo = env("CODE_REPO", "farguney/trident")
    threads = int(env("THREADS", "4"))
    torch.set_num_threads(threads)
    # pick an available int8 backend (fbgemm on x86 runners, qnnpack on ARM)
    for eng in ("fbgemm", "qnnpack"):
        if eng in torch.backends.quantized.supported_engines:
            torch.backends.quantized.engine = eng
            log(f"int8 engine={eng}")
            break
    result_name = env("RESULT_NAME", "cpu_bench.json")
    tri_ckpt = env("TRIDENT_CKPT", "checkpoints_tbptt/final")
    xf_ckpt = env("XF_CKPT", "checkpoints_baseline/final")

    from huggingface_hub import HfApi, snapshot_download
    from safetensors.torch import load_file

    work = Path("/tmp/code")
    snapshot_download(repo_id=repo, repo_type="model", allow_patterns=["src/**"], local_dir=str(work))
    sys.path.insert(0, str(work / "src"))
    from trident import Trident, TridentConfig
    from trident.baseline import ByteTransformer

    ck = Path("/tmp/ck")
    snapshot_download(repo_id=repo, repo_type="model",
                      allow_patterns=[f"{tri_ckpt}/*", f"{xf_ckpt}/*"], local_dir=str(ck))
    tdir = ck / tri_ckpt
    tcfg = TridentConfig(**json.loads((tdir / "config.json").read_text()))
    trident = Trident(tcfg).eval()
    trident.load_state_dict(load_file(str(tdir / "model.safetensors")), strict=False)

    bdir = ck / xf_ckpt
    bcfg = json.loads((bdir / "config.json").read_text())
    xf = ByteTransformer(**bcfg).eval()
    xf.load_state_dict(load_file(str(bdir / "model.safetensors")), strict=False)

    tri_params = sum(p.numel() for p in trident.parameters())
    log(f"threads={threads} torch={torch.__version__}")
    log(f"Trident {tri_params/1e6:.1f}M ({tri_ckpt}) | Transformer {xf.num_params()/1e6:.1f}M ({xf_ckpt})")

    # int8-quantized copies
    trident_q = quantize(trident)
    xf_q = quantize(xf)
    log(f"int8 weight memory: Trident {int8_bytes(trident)/1e6:.1f} MB (fp32 {tri_params*4/1e6:.1f} MB) | "
        f"Transformer {int8_bytes(xf)/1e6:.1f} MB (fp32 {xf.num_params()*4/1e6:.1f} MB)")

    W = 2048
    trident_state_bytes = tcfg.n_blocks * tcfg.n_heads * tcfg.d_v * tcfg.d_k * 4

    def trident_stream(model, L):
        data = torch.randint(0, 256, (1, L), dtype=torch.long)
        def run():
            state = None
            for w in range(0, L, W):
                out = model(data[:, w:w + W], state0=state, return_state=True)
                state = out["state"]
        return run

    def xf_forward(model, L):
        data = torch.randint(0, 256, (1, L), dtype=torch.long)
        return lambda: model(data)

    results = {"threads": threads, "trident_state_bytes": trident_state_bytes, "window": W,
               "trident_params_M": tri_params / 1e6, "transformer_params_M": xf.num_params() / 1e6,
               "int8_weight_bytes": {"trident": int8_bytes(trident), "transformer": int8_bytes(xf)},
               "throughput": {}, "quality_retention": {}, "trident": {}, "transformer": {}}

    # ---- throughput: fp32 vs int8 (single 2 KB window) ----
    log("== throughput (single 2KB window) ==")
    for name, tri_m, xf_m in [("fp32", trident, xf), ("int8", trident_q, xf_q)]:
        t_tri = median_time(trident_stream(tri_m, W))
        t_xf = median_time(xf_forward(xf_m, W))
        results["throughput"][name] = {"trident_kBps": W / t_tri / 1e3, "transformer_kBps": W / t_xf / 1e3,
                                       "trident_ms": t_tri * 1e3, "transformer_ms": t_xf * 1e3}
        log(f"  [{name}] Trident {W/t_tri/1e3:6.1f} kB/s ({t_tri*1e3:5.0f} ms) | "
            f"Transformer {W/t_xf/1e3:6.1f} kB/s ({t_xf*1e3:5.0f} ms)")

    # ---- quality retention: does int8 hold BPB? (real held-out bytes) ----
    try:
        nqw = int(env("QUAL_WINDOWS", "16"))
        raw = build_stream(env("VAL_DATASET", "Salesforce/wikitext"),
                           env("VAL_NAME", "wikitext-103-raw-v1"),
                           env("VAL_SPLIT", "validation"), env("VAL_TEXT_FIELD", "text"), nqw * W)
        qd = torch.tensor(list(raw), dtype=torch.long)[: nqw * W].view(nqw, W)

        def tri_bpb(model):
            nll, tot = 0.0, 0
            for i in range(nqw):
                win = qd[i:i + 1]
                out = model(win, return_logits=True)
                lg = out["logits"].float()
                nll += torch.nn.functional.cross_entropy(lg.reshape(-1, lg.shape[-1]),
                                                          win.reshape(-1), reduction="sum").item()
                tot += win.numel()
            return nll / tot / math.log(2)

        def xf_bpb(model):
            nll, tot = 0.0, 0
            for i in range(nqw):
                win = qd[i:i + 1]
                out = model(win, return_logits=True)
                lg = out["logits"].float()
                nll += torch.nn.functional.cross_entropy(lg.reshape(-1, lg.shape[-1]),
                                                          win.reshape(-1), reduction="sum").item()
                tot += win.numel()
            return nll / tot / math.log(2)

        tri_fp32, tri_int8 = tri_bpb(trident), tri_bpb(trident_q)
        xf_fp32, xf_int8 = xf_bpb(xf), xf_bpb(xf_q)
        results["quality_retention"] = {
            "trident_bpb_fp32": tri_fp32, "trident_bpb_int8": tri_int8,
            "transformer_bpb_fp32": xf_fp32, "transformer_bpb_int8": xf_int8,
        }
        log("== int8 quality retention (BPB, lower=better) ==")
        log(f"  Trident     fp32 {tri_fp32:.4f} -> int8 {tri_int8:.4f}  (Δ{tri_int8-tri_fp32:+.4f})")
        log(f"  Transformer fp32 {xf_fp32:.4f} -> int8 {xf_int8:.4f}  (Δ{xf_int8-xf_fp32:+.4f})")
    except Exception as e:  # noqa
        log(f"quality retention skipped ({e})")

    # ---- wall-time vs context length (fp32 + int8 for Trident stream) ----
    log("== wall-time vs context length ==")
    for L in [2048, 8192, 32768, 98304]:
        t = median_time(trident_stream(trident, L), reps=2)
        tq = median_time(trident_stream(trident_q, L), reps=2)
        results["trident"][L] = {"sec": t, "kBps": L / t / 1e3, "sec_int8": tq,
                                 "kBps_int8": L / tq / 1e3, "mem_bytes": trident_state_bytes}
        log(f"  Trident     ctx={L:6d}  fp32 {t*1e3:7.0f} ms ({L/t/1e3:5.1f} kB/s) | "
            f"int8 {tq*1e3:7.0f} ms ({L/tq/1e3:5.1f} kB/s)  state={trident_state_bytes/1e6:.2f} MB (const)")
    for L in [1024, 2048, 4096, 8192]:
        t = median_time(xf_forward(xf, L), reps=2)
        tq = median_time(xf_forward(xf_q, L), reps=2)
        kv = 2 * bcfg["n_layers"] * bcfg["d_model"] * L * 4  # fp32 KV bytes
        results["transformer"][L] = {"sec": t, "kBps": L / t / 1e3, "sec_int8": tq,
                                     "kBps_int8": L / tq / 1e3, "kv_bytes": kv}
        log(f"  Transformer ctx={L:6d}  fp32 {t*1e3:7.0f} ms ({L/t/1e3:5.1f} kB/s) | "
            f"int8 {tq*1e3:7.0f} ms ({L/tq/1e3:5.1f} kB/s)  KV={kv/1e6:.2f} MB")

    # ---- dominance: Trident/transformer speed ratio + transformer memory ceiling ----
    mem_budget_gb = float(env("MEM_BUDGET_GB", "8"))
    log("== DOMINANCE (Trident vs Transformer) ==")
    dominance = {"speed_ratio": {}, "doc_scale": {}, "mem_budget_gb": mem_budget_gb}
    for L in [2048, 8192]:
        if L in results["trident"] and L in results["transformer"]:
            r = results["trident"][L]["kBps"] / results["transformer"][L]["kBps"]
            rq = results["trident"][L]["kBps_int8"] / results["transformer"][L]["kBps_int8"]
            dominance["speed_ratio"][L] = {"fp32": r, "int8": rq}
            log(f"  ctx={L:5d}: Trident {r:.2f}x faster (fp32), {rq:.2f}x (int8)")
    for L in [65536, 262144, 1048576]:
        kv_gb = 2 * bcfg["n_layers"] * bcfg["d_model"] * L * 4 / 1e9  # transformer fp32 KV
        feasible = kv_gb <= mem_budget_gb
        dominance["doc_scale"][L] = {"transformer_kv_GB": kv_gb, "feasible": feasible,
                                     "trident_state_MB": trident_state_bytes / 1e6}
        log(f"  doc L={L:7d}: transformer KV={kv_gb:6.2f} GB "
            f"({'ok' if feasible else 'INFEASIBLE'} @ {mem_budget_gb:.0f}GB) | "
            f"Trident state={trident_state_bytes/1e6:.2f} MB (const)")
    results["dominance"] = dominance

    out = Path("/tmp/results"); out.mkdir(exist_ok=True)
    (out / result_name).write_text(json.dumps(results, indent=2))
    try:
        HfApi().upload_folder(folder_path=str(out), repo_id=repo, repo_type="model",
                              path_in_repo="results", commit_message=f"cpu inference benchmark ({result_name})")
        log(f"pushed results/{result_name}")
    except Exception as e:  # noqa
        log(f"upload failed: {e}")


if __name__ == "__main__":
    main()