gpu_latency.py

"""Prizma-Seq vs Transformer — the LONG-CONTEXT decode-latency & memory probe (CUDA/Colab-ready).

WHY THIS EXISTS
---------------
gpu_bench.py phase5 (`decode_latency`) measured per-step decode at sequence lengths n<=4096 and
found BOTH models were *overhead-bound*: the per-step wall-clock was flat (TF ~4.5ms, Prizma ~7.0ms),
so the asymptotic O(t)-per-step (Transformer KV-cache) vs O(1)-per-step (Prizma constant state)
WALL-CLOCK crossover did NOT appear at that scale. Prizma only won on MEMORY there (constant state vs
a linearly-growing KV-cache), and that was an *analytic* (floats) claim, not a measured one.

This probe pushes decode MUCH further — bigger n (up to 65k) AND a bigger model (heavier per-step
attention term) — to HONESTLY test whether the O(1) inference advantage is observable in wall-clock,
not just in memory. Two outcomes are both legitimate and reported as-is:
  * a crossover n* exists in the tested range  -> the O(1) latency advantage is real & measured; or
  * no crossover in range                       -> reported plainly; the MEMORY advantage still stands
                                                   (and is now MEASURED via torch.cuda.max_memory_allocated,
                                                   not merely analytic).

It also records the Transformer's practical OOM ceiling (the KV-cache is O(n) memory): hitting OOM is
ITSELF a result favoring Prizma's constant footprint, so OOM is caught and logged, never fatal.

FAIRNESS / HONESTY (no rigging)
-------------------------------
  * BOTH models decode through their `model.step()` streaming API — the fair O(1)-API path. The TF
    `step()` is genuinely KV-cached (O(t) compute & memory at step t); Prizma `step()` carries a fixed
    state + a length-`window` ring (verified O(1)). Same dtype, same device, same warmup, same reps.
  * Warmup steps before every timed measurement; median over reps; proper device sync
    (torch.cuda.synchronize / torch.mps.synchronize) bracketing each timed region.
  * We report BOTH per-step ms (the asymptotics) AND total decode time. The crossover is defined on
    per-step ms (the thing that actually grows for the TF).
  * Memory: analytic floats (KV vs constant state) AT EVERY n, PLUS measured peak GPU bytes for the TF
    decode at each n on CUDA (reset_peak_memory_stats -> decode -> max_memory_allocated), and a measured
    Prizma peak for contrast. On MPS the proper peak API is absent, so we sample current_allocated_memory
    as a best-effort (clearly labelled); the headline measured-memory result is a CUDA/Colab deliverable.

Crash-safe + resumable: every measured (model, n) cell streams to $PRIZMA_RESULTS/gpu_latency.json via an
atomic _save (mirrors gpu_bench.py). A Colab disconnect never loses progress, and re-running skips done
cells. We write to the SIBLING gpu_latency.json and NEVER touch gpu_bench.json.

Env:
  PRIZMA_RESULTS   dir for the JSON ledger (default ./results).
  PRIZMA_LAT_NS    comma-list overriding the n grid (e.g. "128,512,2048").
  PRIZMA_LAT_SMOKE =1 -> fast local smoke (tiny model, n in {128,512,2048}, reps=2) for MPS/CPU verify.
  PRIZMA_LAT_REPS  override reps (default 5; smoke uses 2).
Run:
  python3 gpu_latency.py                 # full campaign (both model sizes) — for Colab/A100/L4
  PRIZMA_LAT_SMOKE=1 python3 gpu_latency.py   # local machinery + step()-correctness smoke
  python3 gpu_latency.py --smoke         # same as PRIZMA_LAT_SMOKE=1
"""
from __future__ import annotations

import json
import math
import os
import sys
import time

import numpy as np
import torch

from seq.common import get_device
from seq.transformer import Transformer, TFConfig
from seq.prizma_seq import PrizmaSeqLM, PrizmaSeqConfig

# Prefer CUDA (the target), then MPS (local smoke), then CPU. Mirrors gpu_bench.py's selection.
DEV = torch.device("cuda" if torch.cuda.is_available()
                   else ("mps" if torch.backends.mps.is_available() else "cpu"))
RES = os.environ.get("PRIZMA_RESULTS", os.path.join(os.path.dirname(__file__), "results"))
os.makedirs(RES, exist_ok=True)
OUT = os.path.join(RES, "gpu_latency.json")   # SIBLING of gpu_bench.json — never clobbered.

V = 512          # vocab (matches gpu_bench.py phase5); decode tokens are random/argmax, content-agnostic.
FEAT_N2 = 224    # Prizma quad2 monomials (matches gpu_bench.py phase2/phase5 Prizma-quad2 headline arm).


# --------------------------------- crash-safe ledger ------------------------------------- #
def _load():
    return json.load(open(OUT)) if os.path.exists(OUT) else {}


def _save(d):
    """Atomic write (mirrors gpu_bench.py _save): tmp file -> os.replace, so a crash mid-write
    never corrupts the ledger."""
    tmp = OUT + ".tmp"
    json.dump(d, open(tmp, "w"), indent=2)
    os.replace(tmp, OUT)


# --------------------------------- device sync helpers ----------------------------------- #
def _sync():
    if DEV.type == "cuda":
        torch.cuda.synchronize()
    elif DEV.type == "mps":
        torch.mps.synchronize()


def _is_oom(err: Exception) -> bool:
    """Detect an out-of-memory condition across CUDA / MPS / generic backends."""
    if isinstance(err, torch.cuda.OutOfMemoryError) if hasattr(torch.cuda, "OutOfMemoryError") else False:
        return True
    s = str(err).lower()
    return ("out of memory" in s) or ("cuda oom" in s) or ("alloc" in s and "fail" in s)


def _empty_cache():
    if DEV.type == "cuda":
        torch.cuda.empty_cache()
    elif DEV.type == "mps":
        try:
            torch.mps.empty_cache()
        except Exception:
            pass


# --------------------------------- memory measurement ------------------------------------ #
def _reset_peak():
    if DEV.type == "cuda":
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.synchronize()


def _peak_bytes():
    """Peak allocated bytes since the last _reset_peak(). CUDA: true peak via max_memory_allocated.
    MPS: no peak API exists, so sample current_allocated_memory (best-effort, labelled in output).
    CPU: not measurable -> None."""
    if DEV.type == "cuda":
        torch.cuda.synchronize()
        return int(torch.cuda.max_memory_allocated())
    if DEV.type == "mps":
        try:
            return int(torch.mps.current_allocated_memory())
        except Exception:
            return None
    return None


# --------------------------------- the core timed probe ---------------------------------- #
@torch.no_grad()
def decode_latency(model, n, reps, warmup, measure_mem=False):
    """Decode `n` tokens through model.step() (the fair O(1)-API path for BOTH models) and return
    the MEDIAN total wall-clock over `reps` timed runs (after `warmup` untimed runs).

    measure_mem=True additionally measures the peak allocated memory over the decode (reset before,
    read after — CUDA gives a true peak; MPS samples current alloc). Returns:
        {"total_s": median_total_seconds, "per_step_ms": median_total*1000/n, "peak_bytes": int|None}
    """
    lat = []
    for r in range(reps + warmup):
        st = model.init_state(1, DEV)
        tok = torch.randint(0, V, (1, 1), device=DEV)
        if r == reps + warmup - 1 and measure_mem:
            _reset_peak()
        _sync()
        t0 = time.time()
        for _ in range(n):
            lg, st = model.step(tok, st)
            tok = lg[:, -1:].argmax(-1)
        _sync()
        dt = time.time() - t0
        if r >= warmup:
            lat.append(dt)
    peak = _peak_bytes() if measure_mem else None
    total = float(np.median(lat))
    return {"total_s": total, "per_step_ms": total * 1000.0 / n, "peak_bytes": peak}


# --------------------------------- analytic memory --------------------------------------- #
def kv_floats(L, H, dh, n):
    """Transformer KV-cache size in floats at decode length n: K and V, per layer, per head, n slots,
    dh each. Grows LINEARLY in n. (Identical formula to gpu_bench.py phase5 tf_kv_floats.)"""
    return 2 * L * H * dh * n


def prizma_state_floats(L, H, dh, d_phi, window):
    """Prizma carried-state size in floats — CONSTANT in n: the d_h x d_phi associative state per head
    per layer, plus the length-`window` k/v ring (2*window*dh per head per layer). Independent of n.
    (Generalizes gpu_bench.py phase5 prizma_state_floats, which hard-coded window=16 and d_phi=dh+feat_n2.)"""
    return L * H * dh * d_phi + 2 * L * H * window * dh


# --------------------------------- model builders ---------------------------------------- #
def build_tf(d, L, H, max_len):
    return Transformer(TFConfig(vocab=V, d_model=d, n_layers=L, n_heads=H, max_len=max_len + 8, rope=True)).to(DEV)


def build_prizma(d, L, H, max_len):
    # max_len does NOT bound Prizma's O(1) decode (no learned pos here); kept generous for safety.
    return PrizmaSeqLM(PrizmaSeqConfig(vocab=V, d_model=d, n_layers=L, n_heads=H, max_len=max_len + 8,
                                     feat_map="quad2", feat_n2=FEAT_N2)).to(DEV)


# --------------------------------- per-model-size sweep ---------------------------------- #
def run_size(res, size_key, d, L, H, ns, reps, warmup):
    """Run the TF-vs-Prizma decode sweep over `ns` for one model size; stream every cell to the
    ledger; build the per-size summary (per-step curves, crossover, OOM ceiling, measured peak mem)."""
    print(f"\n==== MODEL SIZE {size_key}: d{d} L{L} H{H}  (V={V}, feat_n2={FEAT_N2}) ====", flush=True)
    cells = res.setdefault("cells", {})
    dh = d // H

    # Build models once per size; the TF KV-cache is freed by reallocating init_state each rep.
    tf = build_tf(d, L, H, max(ns)); tf.train(False)
    ps = build_prizma(d, L, H, max(ns)); ps.train(False)
    d_phi = ps.cfg.d_phi
    window = ps.cfg.window

    tf_oom_ceiling = None   # first n at which the TF decode OOMs (its practical memory ceiling)
    for n in ns:
        base = f"{size_key}.n{n}"

        # ---- Transformer (KV-cache, O(t)/step, O(n) memory) ---- #
        tkey = f"TF.{base}"
        if tkey not in cells:
            if tf_oom_ceiling is not None:
                # already OOM'd at a smaller n -> everything larger also OOMs; record without retrying.
                cells[tkey] = {"oom": True, "note": f"skipped (OOM at n={tf_oom_ceiling})"}
            else:
                try:
                    rec = decode_latency(tf, n, reps, warmup, measure_mem=True)
                    rec["kv_floats"] = kv_floats(L, H, dh, n)
                    cells[tkey] = rec
                    pm = rec["peak_bytes"]
                    print(f"  n={n:<6} TF(KV)   per-step={rec['per_step_ms']:.3f}ms  "
                          f"total={rec['total_s']:.3f}s  peak={_fmt_mb(pm)}", flush=True)
                except Exception as e:   # noqa: BLE001 — OOM (or any decode failure) is a DATA POINT
                    if _is_oom(e):
                        tf_oom_ceiling = n
                        cells[tkey] = {"oom": True, "note": str(e)[:200]}
                        _empty_cache()
                        print(f"  n={n:<6} TF(KV)   OOM -> practical ceiling (KV-cache too large)", flush=True)
                    else:
                        cells[tkey] = {"error": str(e)[:200]}
                        print(f"  n={n:<6} TF(KV)   ERROR: {str(e)[:120]}", flush=True)
            _save(res)
        elif cells[tkey].get("oom") and tf_oom_ceiling is None:
            tf_oom_ceiling = n   # resume: re-learn the ceiling from the ledger

        # ---- Prizma-quad2 (constant state, O(1)/step, O(1) memory) ---- #
        pkey = f"Prizma-quad2.{base}"
        if pkey not in cells:
            try:
                rec = decode_latency(ps, n, reps, warmup, measure_mem=True)
                rec["state_floats"] = prizma_state_floats(L, H, dh, d_phi, window)
                cells[pkey] = rec
                pm = rec["peak_bytes"]
                print(f"  n={n:<6} Prizma    per-step={rec['per_step_ms']:.3f}ms  "
                      f"total={rec['total_s']:.3f}s  peak={_fmt_mb(pm)}", flush=True)
            except Exception as e:   # noqa: BLE001
                cells[pkey] = ({"oom": True, "note": str(e)[:200]} if _is_oom(e)
                               else {"error": str(e)[:200]})
                print(f"  n={n:<6} Prizma    {'OOM' if _is_oom(e) else 'ERROR'}: {str(e)[:120]}", flush=True)
            _save(res)

    # free this size's models before the next (bigger) size
    del tf, ps
    _empty_cache()

    summary = _summarize_size(cells, size_key, d, L, H, dh, d_phi, window, ns, tf_oom_ceiling)
    res.setdefault("size_summaries", {})[size_key] = summary
    _save(res)
    return summary


def _summarize_size(cells, size_key, d, L, H, dh, d_phi, window, ns, tf_oom_ceiling):
    """Build the per-step ms curves for both models, find the crossover n (first n where TF per-step
    ms strictly exceeds Prizma per-step ms), and assemble the measured + analytic memory curves."""
    tf_ps_ms, ps_ps_ms = {}, {}
    tf_total, ps_total = {}, {}
    tf_peak, ps_peak = {}, {}
    kv_fl, st_fl = {}, {}
    crossover = "no crossover in tested range"
    for n in ns:
        tf_rec = cells.get(f"TF.{size_key}.n{n}", {})
        ps_rec = cells.get(f"Prizma-quad2.{size_key}.n{n}", {})
        kv_fl[n] = kv_floats(L, H, dh, n)
        st_fl[n] = prizma_state_floats(L, H, dh, d_phi, window)
        if "per_step_ms" in tf_rec:
            tf_ps_ms[n] = round(tf_rec["per_step_ms"], 4)
            tf_total[n] = round(tf_rec["total_s"], 4)
            if tf_rec.get("peak_bytes") is not None:
                tf_peak[n] = int(tf_rec["peak_bytes"])
        if "per_step_ms" in ps_rec:
            ps_ps_ms[n] = round(ps_rec["per_step_ms"], 4)
            ps_total[n] = round(ps_rec["total_s"], 4)
            if ps_rec.get("peak_bytes") is not None:
                ps_peak[n] = int(ps_rec["peak_bytes"])
        # crossover: first n where BOTH measured and TF per-step > Prizma per-step.
        if (crossover == "no crossover in tested range"
                and n in tf_ps_ms and n in ps_ps_ms and tf_ps_ms[n] > ps_ps_ms[n]):
            crossover = n

    # measured memory crossover (first n where measured TF peak > Prizma peak), CUDA-meaningful.
    mem_crossover = "no measured-memory crossover in tested range"
    for n in ns:
        if n in tf_peak and n in ps_peak and tf_peak[n] > ps_peak[n]:
            mem_crossover = n
            break

    return {
        "config": {"d": d, "L": L, "H": H, "d_h": dh, "d_phi": d_phi, "window": window,
                   "vocab": V, "feat_n2": FEAT_N2},
        "ns_tested": ns,
        "tf_per_step_ms": tf_ps_ms,
        "prizma_per_step_ms": ps_ps_ms,
        "tf_total_s": tf_total,
        "prizma_total_s": ps_total,
        "latency_crossover_n": crossover,
        "tf_oom_ceiling_n": tf_oom_ceiling,
        "analytic_kv_floats": kv_fl,
        "analytic_prizma_state_floats": st_fl,
        "analytic_kv_over_state_ratio": {n: round(kv_fl[n] / st_fl[n], 3) for n in ns},
        "measured_tf_peak_bytes": tf_peak,
        "measured_prizma_peak_bytes": ps_peak,
        "measured_mem_crossover_n": mem_crossover,
        "mem_measure_kind": ("cuda_max_memory_allocated" if DEV.type == "cuda"
                             else ("mps_current_allocated_memory_bestEffort" if DEV.type == "mps"
                                   else "unavailable_cpu")),
    }


def _fmt_mb(b):
    return "n/a" if b is None else f"{b / 1e6:.1f}MB"


# --------------------------------- top-level driver -------------------------------------- #
def default_grid(smoke: bool):
    """The n-grid and (model-size list, reps, warmup) for full vs smoke runs."""
    if smoke:
        ns = [128, 512, 2048]
        sizes = [("small", 128, 4, 4)]          # tiny: just prove the machinery + step() correctness
        reps, warmup = 2, 1
    else:
        ns = [4096, 8192, 16384, 32768, 65536]  # push WAY past phase5's 4096 to surface a crossover
        sizes = [("small", 128, 4, 4),          # the headline small model (overhead-bound at <=4k)
                 ("big",   512, 8, 8)]          # heavier per-step attention term -> crossover more likely
        reps, warmup = 5, 2
    env_ns = os.environ.get("PRIZMA_LAT_NS")
    if env_ns:
        ns = [int(x) for x in env_ns.split(",") if x.strip()]
    env_reps = os.environ.get("PRIZMA_LAT_REPS")
    if env_reps:
        reps = int(env_reps)
    return ns, sizes, reps, warmup


def main():
    smoke = ("--smoke" in sys.argv) or (os.environ.get("PRIZMA_LAT_SMOKE", "0") == "1")
    ns, sizes, reps, warmup = default_grid(smoke)
    print(f"device={DEV} torch={torch.__version__} results={OUT}", flush=True)
    print(f"mode={'SMOKE' if smoke else 'FULL'} ns={ns} sizes={[s[0] for s in sizes]} "
          f"reps={reps} warmup={warmup}", flush=True)
    if DEV.type == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}", flush=True)

    res = _load()
    res["meta"] = {"device": DEV.type, "torch": torch.__version__, "vocab": V, "feat_n2": FEAT_N2,
                   "mode": "smoke" if smoke else "full", "ns": ns, "reps": reps, "warmup": warmup,
                   "fairness": "both models decode via model.step() (KV-cache for TF, O(1) state for Prizma)"}
    _save(res)

    for (size_key, d, L, H) in sizes:
        run_size(res, size_key, d, L, H, ns, reps, warmup)

    # ---- top-level latency_summary across sizes ---- #
    latency_summary = {}
    for (size_key, d, L, H) in sizes:
        s = res["size_summaries"][size_key]
        latency_summary[size_key] = {
            "config": s["config"],
            "tf_per_step_ms": s["tf_per_step_ms"],
            "prizma_per_step_ms": s["prizma_per_step_ms"],
            "latency_crossover_n": s["latency_crossover_n"],
            "tf_oom_ceiling_n": s["tf_oom_ceiling_n"],
            "measured_tf_peak_bytes": s["measured_tf_peak_bytes"],
            "measured_prizma_peak_bytes": s["measured_prizma_peak_bytes"],
            "measured_mem_crossover_n": s["measured_mem_crossover_n"],
            "analytic_kv_over_state_ratio": s["analytic_kv_over_state_ratio"],
        }
    res["latency_summary"] = latency_summary
    _save(res)

    # ---- honest verdict line ---- #
    verdicts = []
    for size_key, s in latency_summary.items():
        c = s["latency_crossover_n"]
        if isinstance(c, int):
            verdicts.append(f"{size_key}: latency crossover at n={c} (Prizma O(1) wins wall-clock beyond it)")
        else:
            verdicts.append(f"{size_key}: NO latency crossover in {ns} (overhead-bound; memory advantage stands)")
        if s["tf_oom_ceiling_n"] is not None:
            verdicts[-1] += f"; TF OOM ceiling n={s['tf_oom_ceiling_n']}"
    res["verdict"] = verdicts
    _save(res)

    print("\n==== LATENCY VERDICT ====", flush=True)
    for v in verdicts:
        print("  " + v, flush=True)
    print(f"\nsaved -> {OUT}", flush=True)

    # Machine-readable block for downstream parsing (mirrors the ===...=== convention).
    print("\n===LATENCY_RESULTS===", flush=True)
    print(json.dumps({"meta": res["meta"], "latency_summary": latency_summary, "verdict": verdicts},
                     indent=2), flush=True)
    print("===END_LATENCY_RESULTS===", flush=True)
    return res


if __name__ == "__main__":
    main()