fix(bench): bench_mlxlm.py hardcodes a 128 GB memory budget #1858

Problem / Background

scripts/bench_mlxlm.py computes the sweep's admission budget from a fixed 128 GB rather than from the host it runs on (scripts/bench_mlxlm.py:61-68):

# Memory budget: 85% of 128 GB. Override with PYLM_BENCH_MAX_GB env var to
# enforce a tighter cap (e.g. PYLM_BENCH_MAX_GB=65 to skip very large MoE
# models on a 128 GB host where 85% × 128 GB ≈ 108 GB still admits them).
_env_max_gb = os.environ.get("PYLM_BENCH_MAX_GB")
if _env_max_gb:
    MEMORY_LIMIT_BYTES = int(float(_env_max_gb) * 1024 * 1024 * 1024)
else:
    MEMORY_LIMIT_BYTES = int(128 * 1024 * 1024 * 1024 * 0.85)

The comment is written as though 128 GB were the only host size.

The failure this causes is not a skipped row, it is a missing comparison. bench_mlxlm.py and bench_decode.sh produce the paired halves of one measurement, so when one harness skips a checkpoint the other is left with nothing to compare against. On a 512 GB machine all three checkpoints of a recent sweep came back SKIP:oom_estimate (141.1 GB > 108.8 GB) while bench_decode.sh had just measured every one of them on that same host. PYLM_BENCH_MAX_GB can paper over this per run, but only if the operator already suspects the budget is wrong, and the budget is not printed anywhere except inside the skip line itself.

On hosts below 128 GB it errs the other way, and that direction is not cosmetic either. Measured on the GB10 box this repository's CUDA work runs on, free -g reports 121 GB total and os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') returns 121.69 GiB, so the true 85 percent budget is 103.4 GiB against the hardcoded 108.8 GiB. The sweep therefore admits about 5.4 GiB more than the host can hold. This host has been lost twice to GPU memory exhaustion during current work, once as a hard reboot on 2026-09-10 at about 18:29 with two large models resident and once after the NVIDIA driver accumulated 586 NVRM: NV_ERR_NO_MEMORY failures, so an over-generous admission default is a plausible contributor rather than a purely cosmetic defect.

Current Behavior

The guard at scripts/bench_mlxlm.py:552-557 is commented "OOM guard (same as bench_decode.sh)", and it is not the same. The shell harness already derives its budget from the host: detect_memory_bytes() at scripts/bench_decode.sh:315-321 uses sysctl -n hw.memsize on Darwin and free -b 2>/dev/null | awk '/^Mem:/{print $2}' elsewhere, and scripts/bench_decode.sh:323-325 takes 85 percent of that. Only the Python half is hardcoded, which is exactly why the two harnesses disagree about which checkpoints exist.

The resolved budget is printed only inside the skip message (bench_mlxlm.py:554), so a run that skips nothing never records the limit it applied, and a CSV cannot afterwards be audited for which checkpoints were even eligible.

detect_hardware() (bench_mlxlm.py:79-102) shells out to sysctl -n machdep.cpu.brand_string and returns ("unknown", ...) on anything else, and all 50 committed benchmarks/pylm_*.csv files come from Apple hosts (m1ultra, m5max). Nothing in the file gates it to macOS, so any detection added here has to work on Linux as well.

Status: a fix already exists, unmerged

This issue is a record of the bug and its evidence, not a work request. Another benchmarking session on the fleet, on a Mac Studio M3 Ultra, found the same bug from the large-host direction and has an uncommitted local fix to this file, alongside an unrelated --text-via-vlm mode. That change has not been approved for pushing by its owner. Do not implement this independently: it would duplicate work that is already written and would collide in the same file. Ownership is being resolved separately, and this issue should not be picked up until that is settled.

Proposed Solution (for whoever lands the existing fix)

Derive MEMORY_LIMIT_BYTES from the host's physical memory, keeping the 85 percent fraction and keeping PYLM_BENCH_MAX_GB as the override that wins.

• os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') covers macOS and Linux without a subprocess, and was verified on this host to return 121.69 GiB against free -g's 121.
• Fall back to sysctl -n hw.memsize on Darwin and MemTotal from /proc/meminfo on Linux when sysconf raises, mirroring how bench_decode.sh:315-321 branches.
• If every probe fails, disable the guard rather than substituting a number. estimate_model_size already returns 0 for an undeterminable size and bench_decode.sh:349 treats that as "try anyway".
• Rewrite the comment so it states the policy (85 percent of detected physical memory) instead of a single host size, and drop the PYLM_BENCH_MAX_GB example phrased around 128 GB.
• Print the resolved budget and its source once at startup, beside the >>> [store] and >>> [candidates] lines added at bench_mlxlm.py:475-477 and bench_mlxlm.py:499, so a run's admission decisions are auditable from its stderr afterwards.

Out of scope: BENCH_MEM_OVERHEAD_FACTOR (bench_decode.sh:211-216, applied at bench_decode.sh:345-356), which the Python harness also lacks. That is a second divergence between the two guards and belongs in its own issue; folding it in here would change admission for an unrelated reason.

Acceptance Criteria

• With PYLM_BENCH_MAX_GB unset on a host of N GiB physical memory, the resolved budget is 0.85 * N: 103.4 GiB on this GB10 box rather than 108.8.
• PYLM_BENCH_MAX_GB=65 still resolves to exactly 65 GiB on any host.
• The resolved budget and whether it came from detection or from the override are printed once at startup, on a run that skips nothing.
• No literal 128 remains in the budget computation or in its comment.
• A checkpoint that bench_decode.sh measures successfully on a given host is not returned as SKIP:oom_estimate by bench_mlxlm.py on that same host.

Verification

# Budget resolution is module level today, so importing the file is enough.
python3 -c "import importlib.util as u; s=u.spec_from_file_location('b','scripts/bench_mlxlm.py'); m=u.module_from_spec(s); s.loader.exec_module(m); print(m.MEMORY_LIMIT_BYTES/1024**3)"
# Before this fix: 108.79999999981374 on every host, measured.
# After: about 103.4 on this 121.69 GiB GB10 box.

python3 -c "import os; print(os.sysconf('SC_PAGE_SIZE')*os.sysconf('SC_PHYS_PAGES')/1024**3)"
free -g                # Linux cross-check
sysctl -n hw.memsize   # macOS cross-check

PYLM_BENCH_MAX_GB=65 python3 -c "import importlib.util as u; s=u.spec_from_file_location('b','scripts/bench_mlxlm.py'); m=u.module_from_spec(s); s.loader.exec_module(m); print(m.MEMORY_LIMIT_BYTES/1024**3)"
# Expect exactly 65.0 regardless of host.

Verified against 60341873.