| """Measure 512/1024/2048-token prefill on frozen original acceptance documents.""" |
|
|
| import argparse |
| import json |
| import time |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from solomon_mlx import Solomon |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument("--model", default="models/quality") |
| p.add_argument("--output", default="evaluations/chunk-benchmark.json") |
| a = p.parse_args() |
| out = Path(a.output) |
| if out.exists(): |
| raise FileExistsError("Use a new immutable benchmark output") |
| model = Solomon.load(a.model) |
| fixtures = json.loads(Path("evaluations/cuda-acceptance/input/documents.json").read_text()) |
| results = [] |
| for document, parts in fixtures["documents"].items(): |
| for target in (1242, 2048): |
| text = "".join(p["text"] for p in parts) |
| while True: |
| rendered = model.engine.render([{"text": text}], "X") |
| end = rendered.rfind("\n\nX") |
| length = len(model.engine.t.encode(rendered[:end], add_special_tokens=False)) - 1 |
| if length >= target: |
| break |
| text += fixtures["filler"] |
| for chunk in (512, 1024, 2048): |
| model.engine.chunk_size = chunk |
| mx.reset_peak_memory() |
| started = time.perf_counter() |
| with model.prefill(text) as state: |
| results.append( |
| { |
| "document": document, |
| "chunk": chunk, |
| "tokens": state.prefix_tokens, |
| "seconds": time.perf_counter() - started, |
| "peak_metal_bytes": mx.get_peak_memory(), |
| "cache_bytes": sum(c.nbytes for c in state._data["cache"]), |
| } |
| ) |
| print(results[-1], flush=True) |
| out.write_text( |
| json.dumps({"runtime": model.identity, "measurements": results, "hardware_simulation": False}, indent=2) |
| ) |
|
|
