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:
Shay 2026-06-24 13:53:50 -07:00 committed by GitHub
parent e77b0299a6
commit b30716a19c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 513 additions and 0 deletions

View 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())

View file

@ -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.

View 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.

View 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.

View 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