Shape-bucketed kernel autotuner and cold-L2 benchmark methodology

[inureyes]

Background

mlxcel's kernel configuration is manual and shape-blind:

• Blackwell qmm tile shape is set by env (MLXCEL_QMM_TILE_M/N/K, patch patches/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu), tuned once by hand per report qmm-sm121-tile-tuning-gb10-2026-07-10.md.
• The multirow-qmv decode window is a hardcoded range (rows 2-7, env MLXCEL_QMV_MULTIROW), and CUDA batched decode is documented to regress past 7 rows because the small-M qmm shape takes over (docs/CONTINUOUS_BATCHING.md backend note, qmv-multirow-gb10-2026-07-11.md).
• Chunk sizes (prefill chunk, paged decode split) are static defaults.

These are per-shape decisions a tuner should make. The durable design ideas for a serving-runtime autotuner are the cache-key structure and invalidation rules, independent of any particular profiling loop:

• Key = (op, kernel/runner identity, nearest shape bucket, extras such as dtype). Shapes are rounded to buckets (power-of-two style rounding) so nearby shapes share one entry and the tuning matrix stays bounded.
• Cached tactics are validated against environment metadata (runtime version, kernel-library pin, device) collected at tuning time; mismatch invalidates the entry instead of silently applying a stale config.
• Out-of-bucket lookups fall back to a default with a warning instead of failing.

Separately, benchmark timing should defeat last-level-cache warming: rotate input buffers so each timed iteration reads cold memory, with the rotation count derived from the device's L2/SLC size. mlxcel's microbenches currently reuse the same buffers every iteration, which flatters bandwidth-bound kernels (exactly the kernels most of this epic touches).

Task

1. Autotuner module in mlxcel-core (Rust): a TunableOp trait exposing candidates(bucket) -> Vec<Tactic> and run(&Tactic); a profiling harness that times candidates with proper MLX synchronization (eval + backend sync) and picks the min-latency tactic; a persistent JSON cache under the mlxcel cache dir keyed by (op, device model string, shape bucket, dtype, metadata {mlxcel version, MLX pin commit}), with metadata-mismatch invalidation and out-of-bucket default fallback.
2. First consumers:
   • qmm tile selection (replaces static MLXCEL_QMM_TILE_*; env still wins when set explicitly),
   • multirow-qmv row-window boundary (tunes the crossover instead of the hardcoded 7),
   • paged decode v2 kv_chunk_size heuristic (coordinate with Fused paged-attention decode v2: CSR page table, cross-CTA split-KV, and variable-length merge kernels #898; the plan consults the tuner when available).
3. Tuning entry points: an explicit mlxcel tune CLI subcommand that profiles the matrix for a loaded model offline, plus optional first-use lazy tuning behind MLXCEL_AUTOTUNE=1 (default off so cold-start latency is unchanged).
4. Cold-L2 benchmark support: add rotating-buffer input support to the microbench harnesses (benchmarks/, examples/page_gather_microbench.rs, the sampling and paged-decode benches from sibling issues): compute rotation count from the device's last-level-cache size (Apple SLC estimate by device family; CUDA L2 via device props through a small FFI helper) and document the methodology in docs/benchmarks.md.

Performance validation (mandatory)

• Autotuner outcome check on GB10 (and Apple Silicon for any Metal-tunable op): for every tuned op, the autotuned tactic must be >= the current manual/default configuration on the tuning matrix (assert in an integration test that runs a reduced matrix; full numbers recorded in docs/benchmark_results/autotuner-<hw>-<date>.md).
• Batched decode sweep 1-16 rows on GB10 before/after qmv/qmm tuning: the documented 8+ row regression cliff must improve or, at minimum, the tuner must reproduce the current best-known configuration (record which).
• Cold-L2 methodology: demonstrate on one bandwidth-bound kernel (paged decode or rmsnorm) the warm-vs-cold measurement delta and note it in docs/benchmarks.md so future reports state which mode they used.

Regression guard (mandatory)

• With MLXCEL_AUTOTUNE unset and no cache present, behavior and performance are byte-identical to today (defaults untouched; explicit envs still honored verbatim).
• Cache correctness: metadata mismatch (different mlxcel version or MLX pin) triggers re-tune, covered by a unit test with a doctored cache file; corrupted cache files are ignored with a warning, never fatal.
• Determinism: repeated tuning runs on the same machine converge to tactics whose measured times are within noise of each other (flaky-tactic guard: median-of-N timing, N >= 5).
• End-to-end: decode and prefill on two models with tuning enabled show no regression vs defaults (within 3%); full cargo test -p mlxcel-core passes.

References

• patches/mlx/backend/cuda/quantized/qmm/{qmm_sm80.cu,qmv.cu}
• docs/CONTINUOUS_BATCHING.md
• docs/benchmark_results/qmv-multirow-gb10-2026-07-11.md
• docs/benchmarks.md

Line numbers are indicative as of current main; search for the named symbols.

---

Part of epic #909.