File size: 2,289 Bytes
76bbe95 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """Decode-path comparison: tokens/sec and memory vs context length.
HOARD carries constant state (GDN fast weights + a 256-token attention window)
while the Transformer's KV cache grows with context. Measures batch-1 decode
after prefills of increasing length.
python bench_decode.py --runs runs/fw_hoard_small runs/fw_transformer_small \
--contexts 1024 4096 16384 32768 --tokens 96
"""
import argparse, json, os, time
import numpy as np
import mlx.core as mx
from model import HOARD, HoardConfig
def load(run_dir):
meta = json.load(open(os.path.join(run_dir, "config.json")))
cfg = HoardConfig.from_dict(meta["config"])
model = HOARD(cfg)
model.load_weights(os.path.join(run_dir, "model.safetensors"))
model.eval()
return model, meta["args"]["config"]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--runs", nargs="+", required=True)
ap.add_argument("--contexts", type=int, nargs="+", default=[1024, 4096, 16384, 32768])
ap.add_argument("--tokens", type=int, default=96)
ap.add_argument("--data", default="data/fineweb/val.bin")
a = ap.parse_args()
toks = np.memmap(a.data, dtype=np.uint16, mode="r")
print(f"{'model':<22}{'context':>9}{'prefill s':>11}{'decode tok/s':>14}{'peak GB':>9}")
for rd in a.runs:
model, name = load(rd)
for L in a.contexts:
mx.reset_peak_memory()
prompt = mx.array(np.asarray(toks[:L]).astype(np.int32))[None]
cache = model.new_cache()
t0 = time.time()
h = model.forward_hidden(prompt, None, cache)
logits = model.logits(h[:, -1:])
mx.eval(logits)
t_prefill = time.time() - t0
ids = mx.argmax(logits[:, -1], axis=-1)[:, None]
t0 = time.time()
for _ in range(a.tokens):
h = model.forward_hidden(ids, None, cache)
logits = model.logits(h)
ids = mx.argmax(logits[:, -1], axis=-1)[:, None]
mx.eval(ids)
dt = time.time() - t0
print(f"{name:<22}{L:>9}{t_prefill:>11.2f}{a.tokens / dt:>14.1f}"
f"{mx.get_peak_memory() / 2**30:>9.2f}")
del model
mx.clear_cache()
if __name__ == "__main__":
main()
|