Repository navigation
feat(core): shape-bucketed kernel autotuner and cold-L2 benchmarks - #965
Merged
Merged
Conversation
Kernel configuration in this tree was manual and shape-blind: the Blackwell qmm CTA tile was tuned once by hand and set by env, the multirow-qmv decode window was a hardcoded 8-row ceiling, and the v1 paged-decode kernel picked its `NumSplits` launch shape from a threadgroup-memory budget alone. Those are per-shape decisions, and a per-shape decision made once by hand is wrong for most shapes. `mlxcel-core::autotune` adds the durable part of a serving-runtime autotuner: the cache key and its invalidation rules. Keys are `(op, kernel/runner identity, power-of-two shape bucket, dtype)` so nearby shapes share one entry and the tuning matrix stays bounded. Entries carry the mlxcel version and pinned MLX commit they were measured under, and a mismatch in either discards the entry instead of applying a stale config. The store mirrors `PolicyStore` in `src/server/batch/mtp_policy.rs`: one JSON file per key hash under the mlxcel cache root, atomic write-plus-rename, and a load path that validates the version and every key field. A corrupt or unreadable file warns and is ignored, never fatal. An out-of-bucket lookup warns once and returns the op's default. Selection is min-latency over a median-of-5 timing per candidate, with each repetition individually bracketed by an eval and a backend sync. A candidate must beat the op's own default by more than 2% to be selected, so a noisy host converges back to today's behavior rather than to a coin flip; candidates that fail to launch are dropped from the sweep rather than aborting it. First consumer, validated on Apple Silicon: the v1 paged-attention decode `NumSplits` launch shape. `paged_attention_decode` gains a `num_splits_override` parameter through the cxx bridge, and `paged_attention_num_splits_cap` exposes the budget ceiling so the tuner enumerates candidates from the launcher's own arithmetic instead of duplicating it. An override of `0` reproduces the pre-change launch exactly, and out-of-range values clamp back to the ceiling, so a stale cached tactic can never produce an infeasible launch. `NumSplits` is a template argument, so every candidate is a distinct JIT specialization. The two CUDA consumers (qmm CTA `tile_m`, multirow-qmv row window) are wired but UNVALIDATED: no CUDA host was available, so their candidate sets and defaults are derived from the kernel sources and the existing GB10 reports rather than from a run. Both are process-wide env knobs read by patched MLX kernels, so they are tuned only by the offline `mlxcel tune` path (`lazy_tunable() == false`) and applied at process start by `apply_tuned_cuda_kernel_env`, following the `apply_cuda_graph_cache_default` precedent. `qmv.cu` gains `MLXCEL_QMV_MULTIROW_MAX_ROWS`, which narrows the dispatch window without touching the compile-time accumulator width. Issue #898's paged-decode v2 `kv_chunk_size` is left as a documented seam with a reserved op name; its feasible chunk sizes depend on v2 plan accounting that does not exist yet. `mlxcel tune` profiles the matrix offline and writes the winners. It always profiles, so it works with `MLXCEL_AUTOTUNE` unset, and `--model` reads only `config.json` for head geometry (no weights are loaded). `MLXCEL_AUTOTUNE` defaults to off and is then fully inert: no cache reads, no profiling, no filesystem access on the decode path. `cache` consumes tuned tactics without profiling; `1` additionally tunes an unseen bucket on first use. An explicitly-set env var for any tuned knob always wins and suppresses tuning for it. `mlxcel-core::bench_rotation` adds cold-last-level-cache benchmark support: rotating input buffers sized from the detected cache so each timed iteration reads evicted memory, with the Apple SLC estimated by device family and `MLXCEL_BENCH_LLC_BYTES` as the override (the CUDA L2 path is a clearly-marked stub needing an FFI helper). `examples/page_gather_microbench.rs` gains `--cold-l2`, prints the mode and rotation count, and appends both as CSV columns; warm remains the default so historical numbers stay reproducible. Validation: 42 autotuner unit tests, 15 rotation tests, 10 `mlxcel tune` tests, 4 microbench tests, and a GPU-backed integration test (`tests/autotuner_outcome.rs`) asserting on a reduced matrix that the selected tactic is never slower than the default. The existing fused and pooled paged-decode parity tests pass unchanged. The CUDA paths are not exercised anywhere here. Closes #906
inureyes
added
status:review
Issue #906 carries a mandatory determinism guard: repeated tuning runs on the same machine must converge to tactics whose measured times are within noise of each other, with a flaky-tactic guard on top of median-of-N. The harness satisfied the median-of-N half (warmup=2, reps=5) but not the guard, and the constants were far too weak to make the first half mean anything. Two back-to-back `mlxcel tune --op paged-decode-splits` runs on an idle M1 Ultra disagreed in 3 of 12 cells, and the same default tactic re-measured 554.4us then 417.2us at batch 4 / ctx 1024, a 33% swing, while every ctx 16384 cell re-measured within a few percent. The instability was concentrated where per-launch overhead dominates and 5 repetitions after 2 warmup iterations do not average it out. Four changes to the measurement methodology, all in the profiling harness: Warmup is now budgeted in wall-clock time (`warmup_budget_us`, 40ms) with the iteration count as a floor. Two iterations of a sub-millisecond launch is under a millisecond of work, nowhere near enough to bring an Apple GPU to a steady clock, so the old warmup left the ramp inside the measurement. Timed repetitions scale with measured launch cost. The warmup doubles as a cost probe and each candidate then targets `sample_budget_us` (250ms) of samples, clamped to `[reps, max_reps]`. Cheap cells are both the noisy ones and the cheap ones to sample harder, so a ~500us launch now gets ~500 repetitions instead of 5. `reps` stays the documented floor of 5 and `mlxcel tune` refuses a lower value, so the existing gate in `tests/autotuner_outcome.rs` still holds. Candidates are sampled round-robin on a stride rather than one at a time to completion. Measuring A fully and then B charges A for whatever thermal and clock state the sweep started in, which is a systematic bias in candidate order rather than noise; interleaving spreads every candidate's repetitions across the same window so drift is common-mode. The switch threshold is now variance-aware, which is the flaky-tactic guard the issue asked for and the part that was missing. Each measurement carries a scaled median absolute deviation over its repetitions, and switching away from the default requires clearing both `min_improvement` and the default's and candidate's combined relative spread. Candidates that are statistically indistinguishable from the leader collapse to whichever sits nearest the default, so two equally good tactics stop alternating between runs. A tactic that cannot be shown better than the default falls back to the default, which is the safe direction. The observed spread and the threshold actually applied are recorded in the cache entry (`spread`, `default_spread`, `required_improvement`) and printed as `spread` / `reps` / `guard` columns by `mlxcel tune`, so a human can see which rows were decided by the kernel and which by the noise floor. No cell is silently dropped. `TACTIC_VERSION` goes to 2 because a v1 entry was chosen under a rule that ignored dispersion and is not comparable; the new fields are `serde(default)` so an old file is a clean miss rather than a parse warning. Behavior gated on `MLXCEL_AUTOTUNE` is untouched: unset stays fully inert, and an explicit env override still wins over any tuned or cached value. Validated on an idle M1 Ultra: four consecutive `mlxcel tune --op paged-decode-splits` runs now select an identical tactic in all 12 of 12 cells (short context uniformly keeps the default; every ctx 16384 cell selects num_splits=16 at 1.15x to 1.23x). Per-sample spread at short context dropped from 13-17% to 3-5%, and the marginal batch 1 / ctx 16384 cell that flipped now clears its guard by a factor of four. A run takes ~26s. Refs #906
Adds the performance-validation report issue #906 requires, covering the autotuner outcome gate and the cold last-level-cache methodology on Apple M1 Ultra with the Metal backend. The headline result is that the shipped static split count is the wrong choice at long context. `num_splits` defaults to the largest value the threadgroup-memory budget allows (32 at head_dim 128), chosen without reference to how many tokens there are to stripe. At context 16384 the tuner selects 16 instead and gains 1.15x to 1.23x, consistently across every batch size measured and across seven independent sweeps. Short-context cells show no measurable win and keep the default. Records two limits honestly rather than burying them. The host carried a load average of 12.8 to 18.0 from unrelated concurrent work, so short-context cells spread 5 to 19 percent per candidate and cannot resolve differences smaller than that; the long-context cells spread 1 to 5 percent under the same load, which is why their result is the one worth quoting. Separately, no CUDA hardware was available, so the qmm CTA tile and multirow-qmv tuners ship wired but unvalidated and the `qmv.cu` change has not been compiled by any toolchain. Reconciles the cold-L2 section of `docs/benchmarks.md` with what was measured. The warm-to-cold median delta on this part is only -1.0 to +5.0 percent, not the large correction the section implied; the reproducible benefit is dispersion, which falls from 27 percent to 4.0 percent. Also corrects the stated no-rotation crossover, which lands at batch 4 / context 16384 on a 96 MiB SLC rather than at context 32768.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
mlxcel-core::autotune, a shape-bucketed kernel autotuner whose durable contribution is the cache key and its invalidation rules rather than the profiling loop, plusmlxcel-core::bench_rotationfor cold-last-level-cache microbenchmarking. The v1 paged-attention decodeNumSplitslaunch shape becomes the first consumer, validated on Apple Silicon; the two CUDA kernel knobs named in the issue are wired but explicitly unvalidated.Design
Cache keys are
(op, kernel/runner identity, power-of-two shape bucket, dtype), so nearby shapes share one entry and the tuning matrix stays bounded. Every entry carries the mlxcel version and the pinned MLX commit it was measured under, and a mismatch in either discards it instead of applying a stale config. The store deliberately mirrorsPolicyStoreinsrc/server/batch/mtp_policy.rs: directory frommlxcel_core::cache_root(), one JSON file per key hash, atomic write-plus-rename, a load path that validates the version field and every key field, a silent no-op when no cache root resolves, and a#[cfg(test)]explicit-dir constructor for injection. A corrupt or unreadable file warns and is ignored, never panics or propagates. An out-of-bucket lookup (saturated bucket, or an op with no candidates at that shape) warns once and returns the op's default.Selection is min-latency over a median-of-5 timing per candidate, each repetition individually bracketed by an eval and a backend sync. A candidate must beat the op's own default by more than 2% to be selected, so a noisy host converges back to today's behavior rather than to a coin flip. Candidates that fail to launch are dropped from the sweep rather than aborting it, because a launch configuration can legitimately be infeasible on a given device.
Precedence, highest first: an explicitly-set env var for that knob, a valid cached tactic, a tactic profiled now (only under
MLXCEL_AUTOTUNE=1), then the op's default, which is exactly the pre-change behavior.Issue task items
1. Autotuner module.
TunableOp(candidates per bucket, default tactic, env override, run, sync), the median-of-N profiling harness with real MLX synchronization, and the persistent JSON cache with metadata-mismatch invalidation and out-of-bucket default fallback. Files:src/lib/mlxcel-core/src/autotune/{mod,tactic,bucket,profile,store}.rs.2. First consumers. The v1 paged-decode
NumSplitsop is the Metal-tunable one and is wired end to end (src/lib/mlxcel-core/src/autotune/ops/paged_decode_splits.rs, consumed inPagedBlockPool::paged_decode_fused).paged_attention_decodegains anum_splits_overrideparameter through the cxx bridge andpaged_attention_num_splits_capexposes the budget ceiling, so the tuner enumerates candidates from the launcher's own arithmetic instead of duplicating the28672 / (dim * 4)budget. An override of0reproduces the pre-change launch exactly and out-of-range values clamp back to the ceiling, so a stale cached tactic can never produce an infeasible launch. The qmm CTAtile_mand multirow-qmv row-window ops are wired insrc/lib/mlxcel-core/src/autotune/ops/cuda_kernel_knobs.rsand are UNVALIDATED, see the section below. Issue #898's paged-decode v2kv_chunk_sizeis left as a documented seam with a reserved op name (OP_PAGED_DECODE_V2_KV_CHUNK) and a three-step plug-in description; no v2 plan is invented here because the feasible chunk sizes depend on v2 memory accounting that does not exist yet.3. Tuning entry points.
mlxcel tuneprofiles the matrix offline and writes the winners, always profiling so it works withMLXCEL_AUTOTUNEunset.--modelreads onlyconfig.jsonfor head geometry (head_dim,num_attention_heads,num_key_value_heads, with ahidden_size / num_attention_headsfallback); no weights are loaded, because the split-count optimum is a launch-shape property.--dry-run,--op, and per-sweep knobs (--batch-sizes,--context-lengths,--reps,--min-improvement) are available. Lazy first-use tuning sits behindMLXCEL_AUTOTUNE=1, default off;MLXCEL_AUTOTUNE=cacheconsumes an earliermlxcel tuneresult without ever profiling at runtime.4. Cold-L2 benchmark support.
src/lib/mlxcel-core/src/bench_rotation.rscomputes a rotation count asceil(2 * last_level_cache / per_iteration_read_bytes)clamped to 64, so the rotation set exceeds the cache and a buffer is evicted before the rotation returns to it. The Apple SLC is estimated by device family from the silicon generation and performance-core count (macOS exposes no SLC size throughsysctl), biased high because under-estimating silently reintroduces the bias the mode exists to remove.cuda_l2_bytes()is a clearly-marked stub that returnsNoneand documents thecudaDeviceProp::l2CacheSizeFFI helper it needs;MLXCEL_BENCH_LLC_BYTESis the override for both platforms.examples/page_gather_microbench.rsgains--cold-l2and--rotation, prints the mode and rotation count in its header and per config, and appendsmodeandrotationas the last two CSV columns so column-indexed readers of the historical schema keep working. Warm stays the default, so existing numbers remain reproducible. Methodology is documented indocs/benchmarks.md, including why large-context configs need no rotation at all and why a kernel whose warm and cold numbers agree is not memory-limited.Default-off guarantee
With
MLXCEL_AUTOTUNEunset and no cache present,resolve_num_splitsreturns0after oneOnceLockread and one mode read: no cache read, no memo lock, no filesystem access, and the C++ launcher takes exactly the pre-change path.apply_tuned_cuda_kernel_envreturns immediately. Explicit env vars are still honored verbatim.What is NOT validated
The qmm CTA
tile_mand multirow-qmv row-window ops were written on an Apple-Silicon-only host with no CUDA hardware, so their candidate sets, defaults, and workload shapes are derived from the kernel sources and the existing GB10 reports rather than from a run. Nothing in this PR measures them. The module carries a prominent UNVALIDATED banner,docs/environment-variables.mdrepeats it forMLXCEL_QMV_MULTIROW_MAX_ROWS, and both ops arelazy_tunable() == falseso a serving process never profiles them.src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/qmm/qmv.cuwas edited but could not be compiled here, since a Metal-only build never compiles the CUDA patch tree.Test plan
cargo test --release -p mlxcel-core --lib --features metal,accelerate autotune::(42 passed)cargo test --release -p mlxcel-core --lib --features metal,accelerate bench_rotation::(15 passed)cargo test --bin mlxcel --features metal,accelerate commands::tune(10 passed)cargo test --example page_gather_microbench --features metal,accelerate(15 passed)cargo test --release --test autotuner_outcome --features metal,accelerate(2 passed, GPU-backed reduced matrix asserting the tuned tactic is never slower than the default)cargo test --release -p mlxcel-core --lib --features metal,accelerate ffi_tests::test_fused_paged_decode(3 passed, unchanged)cargo test --release -p mlxcel-core --lib --features metal,accelerate ffi_tests::test_pooled_paged_decode(3 passed, unchanged)cargo clippy -p mlxcel-core --lib --tests --features metal,accelerate -- -D warningscargo clippy --lib --bins --test autotuner_outcome --example page_gather_microbench --features metal,accelerate -- -D warningscargo fmt --checkpython3 scripts/ci/check_cross_repo_refs.pymlxcel tune --dry-runand a single-cellmlxcel tunerun confirming the cache file is written with the expected key, tactic, and metadata fieldsFull benchmark measurement is deliberately left to a separate serialized pass so GPU contention does not pollute the numbers.
Closes #906