feat(bench): monotonic OOM backstop and persistent OOM record for decode benchmark sweeps #807

Summary

On the GB10 (DGX Spark, CUDA unified memory) machine, benchmark cells that OOM are worse than wasted time: near-OOM CUDA runs have repeatedly left the driver in a degraded state (NVRM NV_ERR_NO_MEMORY accumulating for hours, eventually wedging the kernel and requiring a reboot), so cells that are certain to OOM should never be launched at all. An earlier WIP gate paired a bash-formula skip with an analytic single-pass memory estimator (MLXCEL_ESTIMATE_SINGLE_PASS_PREFILL extending src/execution/memory_estimate.rs), but the estimator was measured to under-predict allocator-driven OOMs by more than 2.5x (estimated ~42 GiB vs actual >115 GiB for gemma-4-31b at 32768 prompt tokens, before PR #676), so it was judged unreliable as a gate and that WIP was discarded (the maintainer keeps the patch locally for reference). This issue asks only for the two empirical layers that remained sound, implemented fresh as shell-level bookkeeping in the bench harness.

Reference context: #648 and #672 (the July long-prompt failures), PR #676 (the chunked-prefill fix that removed the original motivating gemma-4 32k OOM).

Design

Two layers, both purely empirical (they react to observed OOMs, they never predict):

1. Monotonic backstop (within one sweep run). When a (model, prompt-length) cell fails with a true OOM signature, skip every remaining cell for that same model whose prompt length is greater than or equal to the OOM'd length. Each skipped cell is logged as an explicit CSV row with a new trailing classification token (SKIP:oom_backstop), never silently.
2. Persistent OOM record (across runs). Append each true-OOM cell as a tuple to a gitignored record file at benchmarks/.oom-record (one line per event: model,prompt_target_len,date,commit,peak_bytes,source_csv, with peak_bytes blank when unknown since the harness does not currently capture peak memory, and commit from git rev-parse --short HEAD which is cheap). On later runs, pre-skip any cell whose model has a recorded OOM at a length less than or equal to the cell's length, logging the skip as SKIP:oom_record. Setting MLXCEL_BENCH_IGNORE_OOM_RECORD=1 disables the pre-skip (new OOMs are still appended) so cells can be re-tested after a memory fix lands.

Staleness. Records go stale whenever memory behavior improves: the gemma-4-31b 32768-token OOM that motivated this design was fixed by chunked prefill in PR #676, so a record of it taken before that merge must be prunable. The per-line date and commit exist precisely to make manual pruning easy (grep -v the model, or delete lines older than the fix), and the override env var covers one-off re-tests without editing the file.

Non-goal. No analytic memory estimator, in the harness or in Rust. The discarded estimator under-predicted allocator-driven OOMs by more than 2.5x; a gate that under-predicts is worse than no gate because it launches exactly the runs that wedge the driver. The existing weight-size check in bench_decode.sh (model_fits_in_memory / SKIP:oom_estimate) stays as-is; it is length-independent and out of scope here.

Scope

First pass: scripts/bench_longprompt.sh only. It is the only script with a (model x prompt-length) grid (model-major, length-minor nested loop) and it is where the GB10 driver-wedge incidents happened. scripts/bench_decode.sh runs one prompt length per invocation, so the monotonic backstop has no meaning there; it needs no changes beyond an optional hardening item (below). Other bench scripts (bench_kv_cache.sh, bench_m5_*.sh, bench_na_attention.sh) are Metal/feature-specific harnesses without the length ladder; extending the persistent-record pre-skip to bench_decode.sh --prompt-tokens N invocations can be a follow-up if single-cell ad-hoc runs turn out to re-trigger known OOMs in practice.

Current harness facts (from inspection)

• scripts/bench_decode.sh runs one cell per invocation via bench_one(), which classifies failures using is_oom_failure(rc, err_text): exit 124 (timeout(1) expiry) is explicitly never OOM, exit 137 (SIGKILL, OS OOM-killer) is OOM on the exit code alone, and exit 134 (SIGABRT) or 1 (cxx-propagated exception) is OOM only when the captured combined stdout+stderr (raw) matches allocator text (out of memory, failed to allocate, bad_alloc, metal::malloc, etc.). True OOM emits a row ending in SKIP:oom; non-OOM failures emit FAIL:bench. This already distinguishes OOM from the exit-134 non-OOM aborts seen before (e.g. the CUDA: long-prompt prefill aborts (exit 134, invalid launch config) when a kernel grid dim exceeds 65535 #648 grid.z > 65535 launch-config crash was correctly FAIL:bench).
• scripts/bench_longprompt.sh drives the grid: for model_name in $MODELS; do for len in $LADDER; do ... (default ladder 512 2048 8192 32768), invoking "$BENCH_DECODE" "$model_path" --prompt-tokens "$len" ... --output "$TMP_CELL" per cell and appending tail -n +2 "$TMP_CELL" to the aggregate CSV. The cell's classification is therefore already recoverable in the sweep script from the trailing CSV field of the data row in $TMP_CELL.
• The trailing classification token is CSV field 16 (after prompt_target_len); successful rows have no trailing token, so the check must be "last field matches ^(SKIP|FAIL):", not "field 16 is non-empty" (the quoted prompt field contains a comma, so parse from the end, e.g. awk -F, '{print $NF}').
• benchmarks/ CSVs are git-tracked (135 files), so benchmarks/.oom-record needs an explicit .gitignore entry.
• Benches read ./models, a gitignored symlink to a local model store; if the symlink is missing, bench_longprompt.sh prints [miss] per model and produces zero rows. Infra failures (missing symlink, missing model dir, watchdog timeout exit 124) must never be recorded as OOM; only cells classified SKIP:oom by is_oom_failure feed the backstop and the record. SKIP:oom_estimate is also excluded: it is a length-independent weight-size decision that re-derives instantly on every run, so persisting it adds nothing.

Implementation plan

All changes in scripts/bench_longprompt.sh unless noted.

1. Add a cell_classification helper that extracts the trailing token of the data row in $TMP_CELL (tail -n +2 "$TMP_CELL" | head -n 1 | awk -F, '{print $NF}') and returns it only when it matches ^(SKIP|FAIL):, empty otherwise.
2. Add an emit_skip_row helper that composes a schema-compatible 16-field row for a skipped cell (model,model_path filled, profiling fields empty, date and max_tokens from the sweep's own variables, prompt_target_len set to the skipped length, trailing token SKIP:oom_backstop or SKIP:oom_record) and appends it to $OUTPUT. The sweep script does not compute hardware/mlx_version today; either reuse the values from the most recent real cell row or leave those two fields blank, and document the choice in a comment (downstream consumers filter on the SKIP: prefix, so blanks are acceptable).
3. Monotonic backstop: declare -A OOM_MIN_LEN keyed by model_name. In the inner loop, before invoking $BENCH_DECODE, if OOM_MIN_LEN[$model_name] is set and len >= OOM_MIN_LEN[$model_name], emit the SKIP:oom_backstop row, log to stderr, and continue. After each cell, if cell_classification returned SKIP:oom, update OOM_MIN_LEN[$model_name] to min(current, len) (use >= and min so an unsorted --lengths list still behaves).
4. Persistent record: define OOM_RECORD_FILE="${BENCHMARKS_DIR}/.oom-record". Before invoking a cell, unless MLXCEL_BENCH_IGNORE_OOM_RECORD=1, scan the file for any line whose model equals model_name and whose recorded length is <= len; on match, emit the SKIP:oom_record row (stderr log includes the matching record line so the operator can see why). After a SKIP:oom cell, append model_name,len,$(date '+%Y-%m-%d'),$(git rev-parse --short HEAD 2>/dev/null || echo unknown),,$OUTPUT to the record file. Tolerate a missing or empty record file and malformed lines (skip them with a stderr warning).
5. Add benchmarks/.oom-record to .gitignore and document the file format, the pre-skip semantics, the MLXCEL_BENCH_IGNORE_OOM_RECORD=1 override, and the staleness/pruning guidance (PR fix(cuda/attention): bound long-prompt prefill memory for flash-ineligible models #676 precedent) in the header comment of bench_longprompt.sh and in its usage() text.
6. Optional hardening in scripts/bench_decode.sh (small, same PR is fine): extend the is_oom_failure pattern list with CUDA-specific markers (CUDA_ERROR_OUT_OF_MEMORY, cudaErrorMemoryAllocation, cuMemAlloc failed) alongside the existing generic and Metal patterns, keeping the same "specific and grounded" rule stated in that function's comment.

Tests / validation

No 100 GB run is needed to exercise the logic:

1. Injectable failure hook. Add MLXCEL_BENCH_FAKE_OOM="model:len[,model:len...]" to bench_longprompt.sh: when the current cell matches an entry, skip invoking the runner and synthesize a $TMP_CELL containing a SKIP:oom row. Then run the sweep against any tiny local model (e.g. qwen2.5-0.5b) with a short ladder and assert: the faked cell row says SKIP:oom, every longer cell for that model says SKIP:oom_backstop, other models are unaffected, and benchmarks/.oom-record gained exactly one line with the expected fields.
2. Hand-written record file. Write a record line for (tiny-model, 512) by hand, run the sweep, and assert all cells for that model at lengths >= 512 emit SKIP:oom_record without launching the runner; re-run with MLXCEL_BENCH_IGNORE_OOM_RECORD=1 and assert the cells execute normally.
3. Non-OOM isolation. Fake a FAIL:bench cell (or point one model at a missing directory) and assert neither the backstop nor the record file is touched.
4. Real smoke (optional, cheap). One real sweep on GB10 with the default ladder over a small model set to confirm the aggregate CSV schema is unchanged for successful rows.

Acceptance criteria

• Within one bench_longprompt.sh run, a SKIP:oom cell causes all subsequent cells for the same model with len >= the OOM'd length to be skipped before launch, each logged as an explicit SKIP:oom_backstop CSV row and a stderr message.
• True OOM cells are appended to gitignored benchmarks/.oom-record with model, length, date, short commit, peak-if-known (blank allowed), and source CSV path.
• On later runs, recorded (model, length) cells and all longer lengths for that model are pre-skipped with explicit SKIP:oom_record rows.
• MLXCEL_BENCH_IGNORE_OOM_RECORD=1 bypasses the pre-skip and is documented in the script header and usage().
• Only SKIP:oom classifications feed the backstop and the record; FAIL:bench, timeout (exit 124), SKIP:oom_estimate, missing ./models symlink, and missing model dirs never do.
• Successful-row CSV schema is byte-compatible with the current 16-column output; skip rows use the same column layout with a SKIP: trailing token.
• The fake-OOM and hand-written-record validation scenarios above pass.
• Staleness pruning guidance (with the PR fix(cuda/attention): bound long-prompt prefill memory for flash-ineligible models #676 precedent) is documented in the script header.

Effort: low. Shell-only changes, one script plus a .gitignore line and an optional pattern-list extension.