feat(bench): add MLX exact CGA recall experiment (#909)
* feat(bench): add MLX exact CGA recall experiment * test(bench): cover MLX exact recall experiment contracts * docs(bench): document MLX exact recall experiment * docs(bench): record MLX experiment integration status * docs(bench): add MLX local validation handoff
This commit is contained in:
parent
e77b0299a6
commit
b30716a19c
5 changed files with 513 additions and 0 deletions
282
benchmarks/apple_uma_mlx_exact_recall.py
Normal file
282
benchmarks/apple_uma_mlx_exact_recall.py
Normal file
|
|
@ -0,0 +1,282 @@
|
|||
"""Benchmark-only MLX exact CGA recall experiment.
|
||||
|
||||
ADR-0235 Lane 3: optional MLX score-vector experiment for CORE's exact
|
||||
Cl(4,1) CGA recall workload. This module does not serve answers, does not
|
||||
replace Python/Rust as semantic source of truth, does not use ANN, and does not
|
||||
claim MLX as a runtime backend.
|
||||
|
||||
The MLX path computes the exact diagonal CGA score vector over deterministic
|
||||
(N, 32) float32 fixtures. Scores are copied back to NumPy for the same stable
|
||||
canonical top-k ordering used by the Python/Rust exact-recall oracle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable
|
||||
|
||||
import numpy as np
|
||||
|
||||
from benchmarks.apple_uma_mechanical_sympathy import (
|
||||
DEFAULT_MEASURED,
|
||||
DEFAULT_WARMUP,
|
||||
N_COMPONENTS,
|
||||
RECALL_N_VALUES,
|
||||
RECALL_TOP_K,
|
||||
synthetic_matrix,
|
||||
synthetic_mv,
|
||||
)
|
||||
|
||||
BENCHMARK_NAME = "CORE Apple Silicon MLX Exact CGA Recall Experiment"
|
||||
BENCHMARK_VERSION = "0.1.0"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TimingStats:
|
||||
warmup_iterations: int
|
||||
measured_iterations: int
|
||||
min_ms: float
|
||||
p50_ms: float
|
||||
p95_ms: float
|
||||
max_ms: float
|
||||
mean_ms: float
|
||||
ops_per_sec: float
|
||||
|
||||
def as_dict(self) -> dict[str, float | int]:
|
||||
return {
|
||||
"warmup_iterations": self.warmup_iterations,
|
||||
"measured_iterations": self.measured_iterations,
|
||||
"min_ms": round(self.min_ms, 6),
|
||||
"p50_ms": round(self.p50_ms, 6),
|
||||
"p95_ms": round(self.p95_ms, 6),
|
||||
"max_ms": round(self.max_ms, 6),
|
||||
"mean_ms": round(self.mean_ms, 6),
|
||||
"ops_per_sec": round(self.ops_per_sec, 3),
|
||||
}
|
||||
|
||||
|
||||
def _measure_timing(
|
||||
fn: Callable[[], Any],
|
||||
*,
|
||||
warmup: int = DEFAULT_WARMUP,
|
||||
measured: int = DEFAULT_MEASURED,
|
||||
) -> TimingStats:
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
samples_ms: list[float] = []
|
||||
for _ in range(measured):
|
||||
t0 = time.perf_counter()
|
||||
fn()
|
||||
samples_ms.append((time.perf_counter() - t0) * 1000.0)
|
||||
samples_ms.sort()
|
||||
p95_index = max(0, int(round(0.95 * (len(samples_ms) - 1))))
|
||||
mean_ms = float(np.mean(samples_ms))
|
||||
return TimingStats(
|
||||
warmup_iterations=warmup,
|
||||
measured_iterations=measured,
|
||||
min_ms=samples_ms[0],
|
||||
p50_ms=float(np.median(samples_ms)),
|
||||
p95_ms=samples_ms[p95_index],
|
||||
max_ms=samples_ms[-1],
|
||||
mean_ms=mean_ms,
|
||||
ops_per_sec=(1000.0 / mean_ms) if mean_ms > 0 else 0.0,
|
||||
)
|
||||
|
||||
|
||||
def mlx_import_status() -> dict[str, Any]:
|
||||
"""Return optional MLX availability without making it a dependency."""
|
||||
try:
|
||||
import mlx # type: ignore[import-not-found]
|
||||
import mlx.core as mx # type: ignore[import-not-found]
|
||||
except ImportError as exc:
|
||||
return {"import_succeeded": False, "reason": str(exc)}
|
||||
except Exception as exc:
|
||||
return {"import_succeeded": False, "reason": f"MLX import failed: {exc}"}
|
||||
|
||||
status: dict[str, Any] = {
|
||||
"import_succeeded": True,
|
||||
"module": "mlx.core",
|
||||
"version": getattr(mlx, "__version__", None),
|
||||
"benchmark_only": True,
|
||||
"serving_authorized": False,
|
||||
}
|
||||
try:
|
||||
status["default_device"] = str(mx.default_device())
|
||||
except Exception as exc:
|
||||
status["default_device_error"] = str(exc)
|
||||
return status
|
||||
|
||||
|
||||
def _stable_top_k_from_scores(scores: np.ndarray, top_k: int) -> list[tuple[int, float]]:
|
||||
scores = np.asarray(scores, dtype=np.float32)
|
||||
k = min(top_k, scores.shape[0])
|
||||
if k <= 0:
|
||||
return []
|
||||
if k < scores.shape[0]:
|
||||
cand = np.argpartition(-scores, k - 1)[:k]
|
||||
else:
|
||||
cand = np.arange(scores.shape[0])
|
||||
order = np.lexsort((cand, -scores[cand]))
|
||||
cand = cand[order]
|
||||
return [(int(i), float(scores[i])) for i in cand]
|
||||
|
||||
|
||||
def _cga_inner_metric() -> np.ndarray:
|
||||
from algebra import backend as alg_backend
|
||||
|
||||
metric = getattr(alg_backend, "_CGA_INNER_METRIC")
|
||||
return np.asarray(metric, dtype=np.float32)
|
||||
|
||||
|
||||
def mlx_exact_score_vector(matrix: np.ndarray, query: np.ndarray) -> np.ndarray:
|
||||
"""Compute exact CGA recall scores with MLX, then copy scores to NumPy.
|
||||
|
||||
This intentionally performs only the score-vector workload in MLX. The
|
||||
stable top-k ordering remains canonical NumPy/Python to avoid depending on
|
||||
MLX top-k API details and to preserve CORE's deterministic ordering rule.
|
||||
"""
|
||||
import mlx.core as mx # type: ignore[import-not-found]
|
||||
|
||||
matrix_f32 = np.ascontiguousarray(matrix, dtype=np.float32)
|
||||
query_f32 = np.ascontiguousarray(query, dtype=np.float32)
|
||||
metric_f32 = np.ascontiguousarray(_cga_inner_metric(), dtype=np.float32)
|
||||
|
||||
mx_matrix = mx.array(matrix_f32)
|
||||
mx_query = mx.array(query_f32)
|
||||
mx_metric = mx.array(metric_f32)
|
||||
scores = mx.zeros((matrix_f32.shape[0],), dtype=mx.float32)
|
||||
for i in range(N_COMPONENTS):
|
||||
scores = scores + (mx_metric[i] * mx_matrix[:, i]) * mx_query[i]
|
||||
eval_fn = getattr(mx, "eval", None)
|
||||
if callable(eval_fn):
|
||||
eval_fn(scores)
|
||||
return np.asarray(scores, dtype=np.float32)
|
||||
|
||||
|
||||
def _parity_report(
|
||||
*,
|
||||
canonical: list[tuple[int, float]],
|
||||
candidate: list[tuple[int, float]],
|
||||
) -> dict[str, Any]:
|
||||
canonical_indices = [i for i, _ in canonical]
|
||||
candidate_indices = [i for i, _ in candidate]
|
||||
deltas = [abs(float(a[1]) - float(b[1])) for a, b in zip(canonical, candidate)]
|
||||
max_abs_score_delta = max(deltas) if deltas else 0.0
|
||||
return {
|
||||
"top_k_indices_match": canonical_indices == candidate_indices,
|
||||
"max_abs_score_delta": round(float(max_abs_score_delta), 8),
|
||||
"scores_close": bool(max_abs_score_delta <= 1e-4),
|
||||
"parity_pass": canonical_indices == candidate_indices and max_abs_score_delta <= 1e-4,
|
||||
}
|
||||
|
||||
|
||||
def run_mlx_exact_recall_experiment(
|
||||
*,
|
||||
warmup: int = DEFAULT_WARMUP,
|
||||
measured: int = DEFAULT_MEASURED,
|
||||
mlx_status: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
from algebra import backend as alg_backend
|
||||
|
||||
status = mlx_status or mlx_import_status()
|
||||
if not status.get("import_succeeded"):
|
||||
return {
|
||||
"benchmark_name": BENCHMARK_NAME,
|
||||
"benchmark_version": BENCHMARK_VERSION,
|
||||
"track": "mlx_exact_cga_recall",
|
||||
"skipped": True,
|
||||
"reason": f"MLX unavailable: {status.get('reason', 'mlx.core import failed')}",
|
||||
"mlx_status": status,
|
||||
"benchmark_only": True,
|
||||
"serving_authorized": False,
|
||||
"semantic_backend": "python/rust canonical exact recall",
|
||||
"non_claims": [
|
||||
"No MLX semantic-backend claim.",
|
||||
"No serving integration.",
|
||||
"No ANN or approximate recall.",
|
||||
"No CoreML or Neural Engine claim.",
|
||||
],
|
||||
}
|
||||
|
||||
cases: list[dict[str, Any]] = []
|
||||
for n in RECALL_N_VALUES:
|
||||
matrix = synthetic_matrix(n, seed=n % 17)
|
||||
query = synthetic_mv(seed=5)
|
||||
canonical = alg_backend.vault_recall(
|
||||
[],
|
||||
query,
|
||||
top_k=RECALL_TOP_K,
|
||||
prebuilt_matrix=matrix,
|
||||
)
|
||||
|
||||
def _run_scores() -> np.ndarray:
|
||||
return mlx_exact_score_vector(matrix, query)
|
||||
|
||||
timing = _measure_timing(_run_scores, warmup=warmup, measured=measured)
|
||||
scores = _run_scores()
|
||||
candidate = _stable_top_k_from_scores(scores, RECALL_TOP_K)
|
||||
parity = _parity_report(canonical=canonical, candidate=candidate)
|
||||
rows_per_sec = (n / (timing.mean_ms / 1000.0)) if timing.mean_ms > 0 else 0.0
|
||||
cases.append(
|
||||
{
|
||||
"N": n,
|
||||
"top_k": RECALL_TOP_K,
|
||||
"dtype": "float32",
|
||||
"contiguous": bool(matrix.flags["C_CONTIGUOUS"]),
|
||||
"backend_used": "mlx",
|
||||
"semantic_backend": "canonical exact recall via algebra.backend.vault_recall",
|
||||
"copy_in_boundary": "NumPy contiguous float32 matrix/query copied into MLX arrays",
|
||||
"copy_out_boundary": "MLX score vector copied to NumPy for canonical stable top-k ordering",
|
||||
"timing": timing.as_dict(),
|
||||
"rows_per_sec": round(rows_per_sec, 3),
|
||||
"parity": parity,
|
||||
"top_result_preview": candidate[:3],
|
||||
"canonical_preview": canonical[:3],
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"benchmark_name": BENCHMARK_NAME,
|
||||
"benchmark_version": BENCHMARK_VERSION,
|
||||
"track": "mlx_exact_cga_recall",
|
||||
"skipped": False,
|
||||
"mlx_status": status,
|
||||
"benchmark_only": True,
|
||||
"serving_authorized": False,
|
||||
"semantic_backend": "python/rust canonical exact recall",
|
||||
"score_computation": "MLX exact diagonal CGA score vector; no ANN or approximate search",
|
||||
"top_k_ordering": "canonical NumPy stable ordering after score copy-out",
|
||||
"copy_boundary": {
|
||||
"input": "NumPy -> MLX array copy at benchmark boundary",
|
||||
"output": "MLX score vector -> NumPy copy for stable top-k",
|
||||
"zero_copy_input": "no",
|
||||
},
|
||||
"non_claims": [
|
||||
"No MLX semantic-backend claim.",
|
||||
"No serving integration.",
|
||||
"No ANN or approximate recall.",
|
||||
"No CoreML or Neural Engine claim.",
|
||||
],
|
||||
"cases": cases,
|
||||
}
|
||||
|
||||
|
||||
def _cli_main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description=BENCHMARK_NAME)
|
||||
parser.add_argument("--json", action="store_true", help="emit machine-readable JSON")
|
||||
parser.add_argument("--warmup", type=int, default=DEFAULT_WARMUP)
|
||||
parser.add_argument("--measured", type=int, default=DEFAULT_MEASURED)
|
||||
args = parser.parse_args(argv)
|
||||
report = run_mlx_exact_recall_experiment(warmup=args.warmup, measured=args.measured)
|
||||
if args.json:
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True))
|
||||
else:
|
||||
print(f"{BENCHMARK_NAME} — use --json")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(_cli_main())
|
||||
|
|
@ -0,0 +1,45 @@
|
|||
# Gemini Local Validation Handoff — MLX Exact Recall
|
||||
|
||||
Branch: `feat/apple-mlx-exact-cga-recall-experiment`
|
||||
|
||||
This branch is deliberately narrow. It adds the MLX exact CGA recall experiment as an isolated benchmark module before wiring it into the main Apple UMA report.
|
||||
|
||||
## Commands
|
||||
|
||||
```bash
|
||||
cd /Users/kaizenpro/Projects/core
|
||||
git fetch origin --prune
|
||||
git switch feat/apple-mlx-exact-cga-recall-experiment
|
||||
git pull --ff-only origin feat/apple-mlx-exact-cga-recall-experiment
|
||||
|
||||
uv run python -m pytest -q tests/test_apple_uma_mlx_exact_recall.py
|
||||
uv run python -m benchmarks.apple_uma_mlx_exact_recall --json
|
||||
```
|
||||
|
||||
If MLX is available or installable locally:
|
||||
|
||||
```bash
|
||||
CORE_BACKEND=rust uv run python -m benchmarks.apple_uma_mlx_exact_recall --json
|
||||
```
|
||||
|
||||
## Touch-up limits
|
||||
|
||||
Only patch:
|
||||
|
||||
- MLX import/API mismatches
|
||||
- formatting or lint issues
|
||||
- test failure fixes directly tied to this module
|
||||
|
||||
Do not add serving integration.
|
||||
Do not add Metal.
|
||||
Do not add CoreML/ANE.
|
||||
Do not change semantic source of truth.
|
||||
Do not use ANN or approximate recall.
|
||||
|
||||
## Success criteria
|
||||
|
||||
- Skip path passes when MLX is unavailable.
|
||||
- MLX-present path emits cases.
|
||||
- Every emitted case has `parity.parity_pass: true`.
|
||||
- Copy-in/copy-out fields remain explicit.
|
||||
- Non-claims remain visible.
|
||||
32
docs/benchmarks/apple-uma-mlx-exact-recall-status.md
Normal file
32
docs/benchmarks/apple-uma-mlx-exact-recall-status.md
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
# Apple UMA MLX Exact Recall Status
|
||||
|
||||
This branch intentionally lands the MLX exact CGA recall lane as an isolated benchmark module first.
|
||||
|
||||
## Implemented
|
||||
|
||||
- `benchmarks/apple_uma_mlx_exact_recall.py`
|
||||
- `tests/test_apple_uma_mlx_exact_recall.py`
|
||||
- `docs/benchmarks/apple-uma-mlx-exact-recall.md`
|
||||
|
||||
## Not yet integrated
|
||||
|
||||
The main `benchmarks/apple_uma_mechanical_sympathy.py` report has not yet been modified on this branch.
|
||||
|
||||
Reason: the isolated module is the safest first landing surface for local Apple Silicon validation. After MLX is verified locally, the module can be integrated into the main Apple UMA report as the `mlx_exact_cga_recall` track.
|
||||
|
||||
## Local validation request
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
uv run python -m pytest -q tests/test_apple_uma_mlx_exact_recall.py
|
||||
uv run python -m benchmarks.apple_uma_mlx_exact_recall --json
|
||||
```
|
||||
|
||||
When MLX is installed:
|
||||
|
||||
```bash
|
||||
CORE_BACKEND=rust uv run python -m benchmarks.apple_uma_mlx_exact_recall --json
|
||||
```
|
||||
|
||||
All emitted cases must show `parity.parity_pass: true` before the results are used in Apple-facing material.
|
||||
65
docs/benchmarks/apple-uma-mlx-exact-recall.md
Normal file
65
docs/benchmarks/apple-uma-mlx-exact-recall.md
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
# Apple UMA MLX Exact CGA Recall Experiment
|
||||
|
||||
ADR-0235 Lane 3 introduces an optional, benchmark-only MLX exact-recall experiment for CORE's Cl(4,1) CGA recall workload.
|
||||
|
||||
This is **not** a serving backend. It does not replace Python or Rust as the semantic source of truth. It does not use ANN, HNSW, approximate recall, sampling, CoreML, or Neural Engine acceleration.
|
||||
|
||||
## What it measures
|
||||
|
||||
`benchmarks/apple_uma_mlx_exact_recall.py` measures one narrow workload:
|
||||
|
||||
```text
|
||||
Deterministic (N, 32) float32 fixture matrix
|
||||
+ deterministic length-32 query
|
||||
→ MLX exact diagonal CGA score vector
|
||||
→ score vector copied back to NumPy
|
||||
→ canonical stable top-k ordering
|
||||
→ parity check against algebra.backend.vault_recall
|
||||
```
|
||||
|
||||
The MLX path computes exact scores only. The final top-k ordering is intentionally kept in NumPy/Python so the experiment does not depend on MLX sorting/top-k API behavior and can reuse CORE's canonical descending-score / ascending-index tie break.
|
||||
|
||||
## Run
|
||||
|
||||
```bash
|
||||
uv run python -m benchmarks.apple_uma_mlx_exact_recall --json
|
||||
```
|
||||
|
||||
With MLX unavailable, the report must skip cleanly with an explicit reason.
|
||||
|
||||
With MLX available, the report emits cases for the standard Apple UMA recall sizes and includes:
|
||||
|
||||
- MLX import status and default device when observable
|
||||
- `N`, `top_k`, dtype, and contiguity
|
||||
- p50/p95/mean timing and rows/sec
|
||||
- copy-in boundary: NumPy fixture to MLX array
|
||||
- copy-out boundary: MLX score vector to NumPy
|
||||
- parity against `algebra.backend.vault_recall`
|
||||
- top result preview and canonical preview
|
||||
|
||||
## Non-claims
|
||||
|
||||
This experiment does **not** claim:
|
||||
|
||||
- MLX is a semantic backend
|
||||
- MLX is serving-authorized
|
||||
- CoreML or Neural Engine acceleration
|
||||
- zero-copy everywhere
|
||||
- ANN or approximate recall
|
||||
- token-generation throughput
|
||||
- Apple endorsement, sponsorship, or product integration
|
||||
|
||||
## Validation
|
||||
|
||||
```bash
|
||||
uv run python -m pytest -q tests/test_apple_uma_mlx_exact_recall.py
|
||||
uv run python -m benchmarks.apple_uma_mlx_exact_recall --json
|
||||
```
|
||||
|
||||
When MLX is installed on Apple Silicon, also run:
|
||||
|
||||
```bash
|
||||
CORE_BACKEND=rust uv run python -m benchmarks.apple_uma_mlx_exact_recall --json
|
||||
```
|
||||
|
||||
The `parity.parity_pass` field must be true for every emitted case before using results in any Apple-facing material.
|
||||
89
tests/test_apple_uma_mlx_exact_recall.py
Normal file
89
tests/test_apple_uma_mlx_exact_recall.py
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
"""Tests for the optional Apple UMA MLX exact recall experiment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import socket
|
||||
from unittest import mock
|
||||
|
||||
import numpy as np
|
||||
|
||||
from benchmarks.apple_uma_mlx_exact_recall import (
|
||||
BENCHMARK_NAME,
|
||||
_parity_report,
|
||||
_stable_top_k_from_scores,
|
||||
run_mlx_exact_recall_experiment,
|
||||
)
|
||||
|
||||
|
||||
def test_mlx_unavailable_skips_with_explicit_reason() -> None:
|
||||
report = run_mlx_exact_recall_experiment(
|
||||
warmup=1,
|
||||
measured=1,
|
||||
mlx_status={"import_succeeded": False, "reason": "No module named 'mlx'"},
|
||||
)
|
||||
|
||||
assert report["benchmark_name"] == BENCHMARK_NAME
|
||||
assert report["track"] == "mlx_exact_cga_recall"
|
||||
assert report["skipped"] is True
|
||||
assert "MLX unavailable" in report["reason"]
|
||||
assert report["benchmark_only"] is True
|
||||
assert report["serving_authorized"] is False
|
||||
assert report["semantic_backend"] == "python/rust canonical exact recall"
|
||||
assert any("No ANN" in item for item in report["non_claims"])
|
||||
assert any("No serving" in item for item in report["non_claims"])
|
||||
|
||||
|
||||
def test_stable_top_k_orders_by_score_then_index() -> None:
|
||||
scores = np.array([0.5, 2.0, 2.0, -1.0, 1.5], dtype=np.float32)
|
||||
|
||||
assert _stable_top_k_from_scores(scores, 3) == [
|
||||
(1, float(scores[1])),
|
||||
(2, float(scores[2])),
|
||||
(4, float(scores[4])),
|
||||
]
|
||||
|
||||
|
||||
def test_parity_report_requires_indices_and_scores_close() -> None:
|
||||
canonical = [(2, 1.25), (5, 0.75)]
|
||||
candidate = [(2, 1.25001), (5, 0.75001)]
|
||||
|
||||
report = _parity_report(canonical=canonical, candidate=candidate)
|
||||
|
||||
assert report["top_k_indices_match"] is True
|
||||
assert report["scores_close"] is True
|
||||
assert report["parity_pass"] is True
|
||||
|
||||
|
||||
def test_parity_report_rejects_index_mismatch() -> None:
|
||||
canonical = [(2, 1.25), (5, 0.75)]
|
||||
candidate = [(5, 0.75), (2, 1.25)]
|
||||
|
||||
report = _parity_report(canonical=canonical, candidate=candidate)
|
||||
|
||||
assert report["top_k_indices_match"] is False
|
||||
assert report["parity_pass"] is False
|
||||
|
||||
|
||||
def test_experiment_makes_no_network_calls(monkeypatch) -> None:
|
||||
def _blocked(*_a: object, **_k: object) -> None:
|
||||
raise AssertionError("network access attempted during MLX experiment")
|
||||
|
||||
monkeypatch.setattr(socket.socket, "connect", _blocked)
|
||||
run_mlx_exact_recall_experiment(
|
||||
warmup=1,
|
||||
measured=1,
|
||||
mlx_status={"import_succeeded": False, "reason": "No module named 'mlx'"},
|
||||
)
|
||||
|
||||
|
||||
def test_mlx_score_vector_is_not_called_when_mlx_unavailable() -> None:
|
||||
with mock.patch(
|
||||
"benchmarks.apple_uma_mlx_exact_recall.mlx_exact_score_vector",
|
||||
side_effect=AssertionError("MLX score path should not run when unavailable"),
|
||||
):
|
||||
report = run_mlx_exact_recall_experiment(
|
||||
warmup=1,
|
||||
measured=1,
|
||||
mlx_status={"import_succeeded": False, "reason": "No module named 'mlx'"},
|
||||
)
|
||||
assert report["skipped"] is True
|
||||
Loading…
Reference in a new issue