feat(mlip): ALCHEMI Polaris benchmark harness + batched MD driver

alcf/polaris/mlip/benchmark/README.md

@@ -0,0 +1,104 @@
+# MLIP throughput benchmarks on Polaris
+
+These scripts measure **performance** of MatKit's MLIP backends — how many
+structures per second a backend sustains, how timing scales with structure size,
+and how NVIDIA ALCHEMI native batching compares to sequential ASE MACE.
+
+This is distinct from `../smoke.py`, which is a correctness/integration test and
+explicitly **not** a performance benchmark. Run the smoke test first to confirm
+the backend works on the node, then run these to characterize throughput.
+
+All numbers are hardware-, checkpoint-, and version-specific. Every run writes a
+`bench_meta.json` capturing the GPU model, package versions, and arguments —
+keep it with any results you report.
+
+## Prerequisites
+
+Install the MLIP environment and pass the smoke test as described in
+[`../README.md`](../README.md):
+
+```bash
+export MATKIT_MLIP_ENV=/lus/eagle/projects/<PROJECT>/<USER>/envs/matkit-mlip
+bash alcf/polaris/mlip/install.sh
+qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" alcf/polaris/mlip/smoke.pbs
+```
+
+Edit the `#PBS -A PROJECT` line in the `.pbs` scripts before submitting.
+
+## Single-GPU sweeps
+
+```bash
+qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" \
+  alcf/polaris/mlip/benchmark/bench.pbs
+```
+
+Runs two sweeps for both `nvalchemi-mace` and `ase-mace`:
+
+- **batch_size** — a fixed set of `--n-structures` copies, ALCHEMI batch size
+  swept over `--batch-sizes`. ASE MACE is sequential (batch size 1) and serves as
+  the baseline. Throughput should rise with batch size, then plateau when the GPU
+  saturates. Peak GPU memory is sampled from `nvidia-smi`.
+- **structure_size** — one structure per supercell factor (`--size-factors`),
+  measuring per-structure time vs atom count.
+
+Override defaults at submit time:
+
+```bash
+qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV",\
+MATKIT_BENCH_INPUT=/path/to/structure.cif,\
+MACE_CHECKPOINT=medium,\
+MATKIT_BENCH_DRIVER=opt,\
+MATKIT_BENCH_NSTRUCT=128 \
+  alcf/polaris/mlip/benchmark/bench.pbs
+```
+
+Run `bench.py --help` for the full flag list (checkpoint, dtype, driver, steps,
+batch sizes, size factors).
+
+## Multi-GPU node throughput
+
+Polaris nodes have 4x A100. The runner uses one GPU per process, so node-level
+throughput is measured by launching one process per GPU over a round-robin shard
+of the inputs:
+
+```bash
+qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" \
+  alcf/polaris/mlip/benchmark/bench_multigpu.pbs
+```
+
+`node_result.json` reports aggregate `node_throughput_structs_per_s`. Compare it
+to the single-GPU throughput at the same batch size — ideal scaling is ~4x.
+
+## Reading results
+
+`bench.py` writes `bench_results.jsonl`, one JSON object per configuration:
+
+| field | meaning |
+|-------|---------|
+| `sweep` | `batch_size` or `structure_size` |
+| `backend` | `nvalchemi-mace` or `ase-mace` |
+| `batch_size`, `n_structures`, `n_atoms` | configuration |
+| `setup_time_s` | model load time (from the manifest) |
+| `wall_time_s` | batch execution wall time (from the manifest) |
+| `elapsed_s` | end-to-end `run_mlip_batch` time |
+| `throughput_structs_per_s` | `succeeded / elapsed_s` |
+| `calc_time_s_mean` / `_max` | per-item calculation time |
+| `peak_gpu_mem_mib` | max used GPU memory sampled during the run |
+
+Quick summary with `jq`:
+
+```bash
+jq -r 'select(.sweep=="batch_size") |
+  [.backend, .batch_size, .throughput_structs_per_s, .peak_gpu_mem_mib] | @tsv' \
+  bench_results.jsonl
+```
+
+## Caveats
+
+- Let model weights finish downloading before trusting timings. The PBS scripts
+  do a throwaway warm-up run first; the ALCF HTTP proxy is exported for
+  compute-node downloads.
+- `setup_time_s` (model load) is reported separately from per-item time; when
+  comparing backends, look at steady-state throughput, not the first item.
+- These measure MatKit's execution path, not raw kernel performance, and do not
+  establish model parity or scientific accuracy.

alcf/polaris/mlip/benchmark/bench.pbs

@@ -0,0 +1,59 @@
+#!/bin/bash -l
+#PBS -N matkit-mlip-bench
+#PBS -l select=1:system=polaris
+#PBS -l place=scatter
+#PBS -l walltime=01:00:00
+#PBS -l filesystems=home:eagle
+#PBS -q debug
+#PBS -A PROJECT
+
+# Single-GPU throughput benchmark for MatKit MLIP backends.
+# Submit from the MatKit checkout:
+#   qsub -v MATKIT_MLIP_ENV="$MATKIT_MLIP_ENV" alcf/polaris/mlip/benchmark/bench.pbs
+# Override inputs/checkpoint/output as needed, e.g.:
+#   qsub -v MATKIT_MLIP_ENV=...,MATKIT_BENCH_INPUT=/path/to.cif,\
+#   MACE_CHECKPOINT=medium alcf/polaris/mlip/benchmark/bench.pbs
+
+set -euo pipefail
+
+: "${MATKIT_MLIP_ENV:?Submit with -v MATKIT_MLIP_ENV=/path/to/env}"
+
+MATKIT_REPO="${MATKIT_REPO:-${PBS_O_WORKDIR}}"
+INPUT_FILE="${MATKIT_BENCH_INPUT:-${MATKIT_REPO}/tests/data/test_structure.cif}"
+OUTPUT_DIR="${MATKIT_BENCH_OUTPUT:-${MATKIT_REPO}/projects/mlip_bench_${PBS_JOBID}}"
+MACE_CHECKPOINT="${MACE_CHECKPOINT:-medium}"
+DRIVER="${MATKIT_BENCH_DRIVER:-energy}"
+
+module use /soft/modulefiles
+module load conda/2025-09-25
+source "${MATKIT_MLIP_ENV}/bin/activate"
+
+export HTTP_PROXY="http://proxy.alcf.anl.gov:3128"
+export HTTPS_PROXY="http://proxy.alcf.anl.gov:3128"
+export http_proxy="${HTTP_PROXY}"
+export https_proxy="${HTTPS_PROXY}"
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
+
+cd "${MATKIT_REPO}"
+
+# Pre-warm model weights so downloads do not pollute timed runs. A throwaway
+# single-structure batch forces the checkpoint to resolve before the sweeps.
+python -m matkit.cli mlip run \
+    --backend nvalchemi-mace \
+    --checkpoint "${MACE_CHECKPOINT}" \
+    --device cuda \
+    --dtype float32 \
+    --driver energy \
+    --input "${INPUT_FILE}" \
+    --output "${OUTPUT_DIR}_warmup/result.json" || true
+
+python alcf/polaris/mlip/benchmark/bench.py \
+    --input "${INPUT_FILE}" \
+    --output-dir "${OUTPUT_DIR}" \
+    --checkpoint "${MACE_CHECKPOINT}" \
+    --driver "${DRIVER}" \
+    --backends nvalchemi-mace ase-mace \
+    --n-structures "${MATKIT_BENCH_NSTRUCT:-64}" \
+    --batch-sizes ${MATKIT_BENCH_BATCH_SIZES:-1 2 4 8 16 32 64} \
+    --size-factors ${MATKIT_BENCH_SIZE_FACTORS:-1 2 3}