156 lines
6.0 KiB
Python
156 lines
6.0 KiB
Python
#!/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())
|