tools: shared Metal bench harness for lab agents - #31
Conversation
Every lab keeps rewriting the same pyobjc boilerplate (and producing 10+ Pyright errors per file). Centralize it under SuperKittens/benchmark/harness/. bench_harness.py owns device + queue + library load, PSO cache, type-safe FunctionConstants (set_bool/short/ushort/int/uint/float), shared-storage buffer creation with explicit zeroing, GPU-timed run loop (warmup + reps x dispatches_per_cmd, GPUEndTime-GPUStartTime, min/mean/p10 + GB/s), and an aligned print_table. roofline.py exposes DEVICE_SPECS for M1-M4 and roofline_us(bytes, flops, device). numeric_ref.py implements matvec/rmsnorm/softmax_online/silu/geglu and explicitly NotImplementedError-s the q-quant dequants (with TODO notes to prevent a repeat of the q2_K reference-encoder bug).
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 341363fdda
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| from .bench_harness import ( | ||
| BenchHarness, | ||
| BenchResult, | ||
| Buf, | ||
| FunctionConstants, |
There was a problem hiding this comment.
Adding this package-level export changes what from benchmark.harness import BenchHarness resolves to, so existing MLX benchmark scripts that previously imported benchmark/harness.py now get the new Metal harness class instead. In this repo, benchmark/gemm/gemm_mlx.py and benchmark/gemm/gemm_bias_silu_mlx.py instantiate BenchHarness() with no arguments, which will fail once this module takes precedence because the new constructor requires metallib_path.
Useful? React with 👍 / 👎.
| key = (kernel_name, id(fc) if fc is not None else 0) | ||
| cached = self._pso_cache.get(key) |
There was a problem hiding this comment.
The PSO cache key uses id(fc) rather than the actual constant entries, so cached pipelines can be returned for the wrong constants when a FunctionConstants object is mutated and reused, or when Python reuses a freed object's id. In those cases pso() may silently run kernels with stale specialization constants, producing incorrect benchmark/correctness results.
Useful? React with 👍 / 👎.
1 participant
Summary
SuperKittens/benchmark/harness/withbench_harness.py,numeric_ref.py,roofline.py, plus an__init__.pyand short README.BenchHarnessowns the Metal device, command queue, and metallib load; caches PSOs; provides type-safeFunctionConstants; runs warmup + reps × dispatches/cmdbuf withGPUEndTime - GPUStartTimetiming and returns min/mean/p10 µs + GB/s.numeric_refprovides correct fp32 references for matvec/rmsnorm/softmax_online/silu/geglu; the quant-dequant entries are explicitNotImplementedErrorto prevent the q2_K reference-encoder bug from recurring.rooflineexposes M1-M4 device specs and a one-callroofline_us(bytes, flops, device)so labs can compare measurements to a ceiling.rmsnorm_bf16(D=1024 → 10.0 µs, D=4096 → 31.3 µs on local M2; rel ≈ 9e-3 bf16 noise floor).Test plan
from benchmark.harness import BenchHarness, FunctionConstantsresolves and a trivialh.run(...)returns aBenchResultend-to-end on lexie M4.