feat: add native MLX inference benchmarks

benchmarks/inference/direct_mlx_benchmark.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+"""Measure direct MLX-LM generation without an HTTP serving layer."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import platform
+import statistics
+import time
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+from benchmark import PROMPTS
+from mlx_lm import load, stream_generate
+from mlx_lm.sample_utils import make_sampler
+
+
+def percentile(values: list[float], fraction: float) -> float:
+    ordered = sorted(values)
+    position = (len(ordered) - 1) * fraction
+    lower = math.floor(position)
+    upper = math.ceil(position)
+    if lower == upper:
+        return ordered[lower]
+    return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower)
+
+
+def summarize(rows: list[dict[str, Any]]) -> dict[str, Any]:
+    measured = [row for row in rows if row["phase"] == "measured"]
+    result: dict[str, Any] = {
+        "runtime": "mlx-lm-direct",
+        "model": measured[0]["model"],
+        "requests": len(measured),
+        "successes": sum(row["error"] is None for row in measured),
+    }
+    for field in ("ttft_s", "total_s", "prompt_tps", "decode_tps", "request_tps"):
+        values = [float(row[field]) for row in measured if row[field] is not None]
+        result[f"{field}_mean"] = statistics.fmean(values)
+        result[f"{field}_p50"] = percentile(values, 0.5)
+        result[f"{field}_p95"] = percentile(values, 0.95)
+    result["peak_memory_gb"] = max(float(row["peak_memory_gb"]) for row in measured)
+    return result
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--label", required=True)
+    parser.add_argument("--warmups", type=int, default=2)
+    parser.add_argument("--rounds", type=int, default=5)
+    parser.add_argument("--max-tokens", type=int, default=128)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--output-dir", type=Path, default=Path(__file__).parent / "direct-results"
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    if args.warmups < 0 or args.rounds < 1 or args.max_tokens < 1:
+        raise SystemExit("warmups must be >= 0; rounds and max-tokens must be >= 1")
+
+    run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
+    output_dir = args.output_dir / f"{run_id}-{args.label}"
+    output_dir.mkdir(parents=True, exist_ok=False)
+
+    load_started = time.perf_counter()
+    model, tokenizer = load(args.model, trust_remote_code=True)
+    load_s = time.perf_counter() - load_started
+    sampler = make_sampler(temp=0.0)
+    prompt_items = list(PROMPTS.items())
+    rows: list[dict[str, Any]] = []
+
+    for phase, count in (("warmup", args.warmups), ("measured", args.rounds)):
+        for round_number in range(count):
+            prompt_name, prompt = prompt_items[round_number % len(prompt_items)]
+            mx.random.seed(args.seed)
+            started = time.perf_counter()
+            first_token_at: float | None = None
+            final = None
+            error = None
+            try:
+                for response in stream_generate(
+                    model,
+                    tokenizer,
+                    prompt,
+                    max_tokens=args.max_tokens,
+                    sampler=sampler,
+                ):
+                    if first_token_at is None:
+                        first_token_at = time.perf_counter()
+                    final = response
+            except Exception as exception:  # Preserve evidence before failing the run.
+                error = f"{type(exception).__name__}: {exception}"
+            total_s = time.perf_counter() - started
+            tokens = final.generation_tokens if final is not None else 0
+            row = {
+                "runtime": "mlx-lm-direct",
+                "model": args.label,
+                "source": args.model,
+                "phase": phase,
+                "round": round_number + 1,
+                "prompt_name": prompt_name,
+                "prompt_tokens": final.prompt_tokens if final is not None else None,
+                "output_tokens": tokens,
+                "ttft_s": first_token_at - started if first_token_at is not None else None,
+                "total_s": total_s,
+                "prompt_tps": final.prompt_tps if final is not None else None,
+                "decode_tps": final.generation_tps if final is not None else None,
+                "request_tps": tokens / total_s if tokens else None,
+                "peak_memory_gb": final.peak_memory if final is not None else None,
+                "finish_reason": final.finish_reason if final is not None else None,
+                "error": error,
+            }
+            rows.append(row)
+            print(
+                f"{phase} {round_number + 1}/{count} prompt={prompt_name} "
+                f"total={total_s:.3f}s tokens={tokens} "
+                f"decode={row['decode_tps'] or 0:.2f}tok/s error={error or 'none'}",
+                flush=True,
+            )
+
+    summary = summarize(rows)
+    metadata = {
+        "run_id": run_id,
+        "created_at": datetime.now(UTC).isoformat(),
+        "host": {
+            "platform": platform.platform(),
+            "machine": platform.machine(),
+            "python": platform.python_version(),
+        },
+        "arguments": vars(args) | {"output_dir": str(args.output_dir)},
+        "load_s": load_s,
+        "mlx_version": getattr(mx, "__version__", None),
+        "device": str(mx.default_device()),
+        "prompts": PROMPTS,
+        "summary": summary,
+    }
+    (output_dir / "raw.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
+    (output_dir / "metadata.json").write_text(
+        json.dumps(metadata, indent=2, default=str), encoding="utf-8"
+    )
+    (output_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    print(json.dumps(summary, indent=2))
+    print(f"Results: {output_dir}")
+    return 1 if any(row["error"] for row in rows if row["phase"] == "measured") else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())