From 2b2f6d3a681c3ddf3539157bf2d99908c2dbdcf3 Mon Sep 17 00:00:00 2001
From: Scott Friedman <3011922+scttfrdmn@users.noreply.github.com>
Date: Tue, 21 Apr 2026 13:30:21 -0700
Subject: [PATCH 1/6] docs: update Known Gaps to reflect v0.5.4 NKI coverage +
 add run_bench.sh
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

architecture.md Known Gaps was stale since v0.4.0:
- syrk and trsm (left-side) now have NKI kernels; only symm and trmm are
  PyTorch-only, and neither is in the DF-MP2 hot path
- nki_mp2_energy dispatch overhead bullet superseded by batched-pair energy
  (#43/#46): warm 3.6×/5.2× faster than torch at small/medium shape
- FP64/double-double bullet updated: decision made 2026-04-18, #10 closed
  "not needed", #22 deferred indefinitely

scripts/run_bench.sh: runs df_mp2.py --bench via SSM on trnblas-ci-trn1.
Supports --shape large/medium (default: both). Follows base64-SSM pattern
from run_pyscf_tests.sh; polls up to 120 min for cold NEFF compile.
---
 docs/architecture.md         |  30 +++---
 infra/terraform-trn2/main.tf |   1 +
 scripts/run_bench.sh         | 187 +++++++++++++++++++++++++++++++++++
 3 files changed, 204 insertions(+), 14 deletions(-)
 create mode 100755 scripts/run_bench.sh

diff --git a/docs/architecture.md b/docs/architecture.md
index 3a29767..e25e399 100644
--- a/docs/architecture.md
+++ b/docs/architecture.md
@@ -44,20 +44,22 @@ since it's reused across all auxiliary basis indices `P`.
 
 ## Known gaps
 
-- **Level 3 NKI coverage is partial.** As of v0.4.0, `gemm`, `batched_gemm`,
-  and the custom `nki_mp2_energy` reduction have NKI kernels. `symm`, `syrk`,
-  `trsm`, `trmm` still dispatch straight to PyTorch — these are the next
-  targets (tracked for v0.5.0). `syrk` and `trsm` appear in the DF-MP2 hot
-  path (metric construction, Cholesky-based metric inversion).
-- **`nki_mp2_energy` matches torch at medium, doesn't beat it.** Kernel is
-  correct; perf is gated by per-(i, j) dispatch/load overhead. Phase 2
-  restructuring (batch multiple (i, j) per dispatch) is open under
-  [#15](https://github.com/trnsci/trnblas/issues/15).
-- **No FP64.** Trainium's Tensor Engine maxes out at FP32. See
-  [Precision envelope](#precision-envelope) below for the measured FP32 vs
-  PySCF-FP64 picture. Double-double emulation is gated on whether cc-pVTZ or
-  larger basis sets exceed µHartree
-  ([#10](https://github.com/trnsci/trnblas/issues/10)).
+- **Level 3 NKI coverage is partial.** `gemm`, `batched_gemm`, `syrk`, and
+  `trsm` (left-side blocked) have NKI kernels. `symm` and `trmm` still
+  dispatch straight to PyTorch. Neither is in the DF-MP2 hot path.
+- **Batched-pair energy (v0.5.2–v0.5.4) solved the dispatch overhead.**
+  `nki_batched_pair_energy`
+  ([#43](https://github.com/trnsci/trnblas/issues/43),
+  [#46](https://github.com/trnsci/trnblas/issues/46)) replaces the nocc²-loop
+  dispatch with a single `@nki.jit` call (small shape) or chunked i-loop
+  (medium/large). Warm: **3.6× faster than torch at small shape, 5.2× at
+  medium shape.**
+- **No FP64.** Trainium's Tensor Engine maxes out at FP32.
+  **Decision (2026-04-18):** FP32 is sufficient — both gate cases are well
+  below 1 µHartree.
+  [#10](https://github.com/trnsci/trnblas/issues/10) closed "not needed";
+  [#22](https://github.com/trnsci/trnblas/issues/22) (double-double) deferred
+  indefinitely. See [Precision envelope](#precision-envelope) below.
 - **Level 1/2 are PyTorch-only.** The Tensor Engine is wasted on vector ops;
   Level 3 is where NKI acceleration pays off. Not planned to change.
 
diff --git a/infra/terraform-trn2/main.tf b/infra/terraform-trn2/main.tf
index 09e3b09..b761063 100644
--- a/infra/terraform-trn2/main.tf
+++ b/infra/terraform-trn2/main.tf
@@ -169,6 +169,7 @@ resource "aws_instance" "ci" {
   tags = {
     Name = var.instance_tag
   }
+
 }
 
 # ---------------------------------------------------------------------------
diff --git a/scripts/run_bench.sh b/scripts/run_bench.sh
new file mode 100755
index 0000000..b3e7d91
--- /dev/null
+++ b/scripts/run_bench.sh
@@ -0,0 +1,187 @@
+#!/usr/bin/env bash
+#
+# Run the df_mp2.py bench on the trnblas CI trn1 instance.
+#
+# Usage:
+#   AWS_PROFILE=aws ./scripts/run_bench.sh                    # medium + large shapes
+#   AWS_PROFILE=aws ./scripts/run_bench.sh --medium-only      # medium shape only
+#   AWS_PROFILE=aws ./scripts/run_bench.sh --shape large      # one shape
+#
+# Runs `python examples/df_mp2.py --bench --batched-pair-energy` with cold
+# and warm timing. Cold compile can take 30–90 min for large shapes; the
+# script polls up to 120 min.
+#
+# The trnblas CI instance (trnblas-ci-trn1) must be provisioned via
+# infra/terraform/ before running. See docs/aws_setup.md.
+
+set -euo pipefail
+
+SHAPES=("medium" "large")
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --medium-only)
+      SHAPES=("medium")
+      shift
+      ;;
+    --shape)
+      SHAPES=("$2")
+      shift 2
+      ;;
+    *)
+      shift
+      ;;
+  esac
+done
+
+INSTANCE_TYPE="${INSTANCE_TYPE:-trn1}"
+TAG="trnblas-ci-${INSTANCE_TYPE}"
+case "$INSTANCE_TYPE" in
+  trn2*) REGION="${AWS_REGION:-sa-east-1}" ;;
+  *)     REGION="${AWS_REGION:-us-east-1}" ;;
+esac
+SHA="$(git rev-parse HEAD)"
+
+: "${AWS_PROFILE:?Set AWS_PROFILE, e.g. AWS_PROFILE=aws ./scripts/run_bench.sh}"
+
+echo "Looking up instance with Name=$TAG in $REGION..."
+INSTANCE_ID=$(aws ec2 describe-instances \
+  --filters "Name=tag:Name,Values=$TAG" \
+            "Name=instance-state-name,Values=stopped,stopping,running,pending" \
+  --query 'Reservations[0].Instances[0].InstanceId' \
+  --output text \
+  --region "$REGION")
+
+if [[ -z "$INSTANCE_ID" || "$INSTANCE_ID" == "None" ]]; then
+  echo "ERROR: No instance found with Name=$TAG" >&2
+  echo "Provision with: cd infra/terraform && terraform apply" >&2
+  exit 1
+fi
+
+echo "Instance: $INSTANCE_ID"
+
+cleanup() {
+  local exit_code=$?
+  echo ""
+  echo "Stopping $INSTANCE_ID..."
+  aws ec2 stop-instances --instance-ids "$INSTANCE_ID" --region "$REGION" >/dev/null
+  exit "$exit_code"
+}
+trap cleanup EXIT
+
+STATE=$(aws ec2 describe-instances --instance-ids "$INSTANCE_ID" --region "$REGION" \
+  --query 'Reservations[0].Instances[0].State.Name' --output text)
+
+if [[ "$STATE" == "stopping" ]]; then
+  echo "Instance is stopping — waiting for stopped..."
+  aws ec2 wait instance-stopped --instance-ids "$INSTANCE_ID" --region "$REGION"
+  STATE=stopped
+fi
+if [[ "$STATE" == "stopped" ]]; then
+  echo "Starting instance..."
+  aws ec2 start-instances --instance-ids "$INSTANCE_ID" --region "$REGION" >/dev/null
+fi
+
+echo "Waiting for instance-running..."
+aws ec2 wait instance-running --instance-ids "$INSTANCE_ID" --region "$REGION"
+echo "Waiting for SSM agent..."
+for _ in $(seq 1 60); do
+  PING=$(aws ssm describe-instance-information \
+    --filters "Key=InstanceIds,Values=$INSTANCE_ID" \
+    --region "$REGION" \
+    --query 'InstanceInformationList[0].PingStatus' --output text 2>/dev/null || true)
+  [[ "$PING" == "Online" ]] && break
+  sleep 5
+done
+if [[ "$PING" != "Online" ]]; then
+  echo "ERROR: SSM agent not Online after 5 minutes (last PingStatus=$PING)" >&2
+  exit 1
+fi
+
+# Build the remote bench command: run df_mp2.py --bench for each shape.
+# Large shape cold compile can take 30-90 min (chunked NEFF compilations).
+SHAPES_ARG="${SHAPES[*]}"
+echo "Sending bench command (SHA=$SHA, shapes=${SHAPES_ARG})..."
+
+# Build SSM parameters via Python with a base64-encoded bash script.
+# Same pattern as run_pyscf_tests.sh — see that file for rationale.
+PARAMS_FILE=$(mktemp /tmp/trnblas-bench-XXXXXX.json)
+SHA_VAL="$SHA" SHAPES_VAL="$SHAPES_ARG" python3 - <<'PYEOF' > "$PARAMS_FILE"
+import base64, json, os
+
+sha    = os.environ["SHA_VAL"]
+shapes = os.environ["SHAPES_VAL"].split()
+
+# Build the bench commands: one python invocation per shape.
+bench_cmds = ""
+for shape in shapes:
+    bench_cmds += (
+        f"echo '--- shape={shape} ---'\n"
+        f"sudo -u ubuntu env PATH=$NEURON_VENV/bin:/usr/bin:/bin TMPDIR=/var/tmp"
+        f" TRNBLAS_REQUIRE_NKI=1"
+        f" $NEURON_VENV/bin/python /home/ubuntu/trnblas/examples/df_mp2.py"
+        f" --bench --shape {shape} --batched-pair-energy\n"
+    )
+
+script = (
+    "#!/bin/bash\n"
+    "set -euo pipefail\n"
+    "cd /home/ubuntu/trnblas\n"
+    "sudo -u ubuntu git fetch --all\n"
+    f"sudo -u ubuntu git checkout {sha}\n"
+    "NEURON_VENV=$(ls -d /opt/aws_neuronx_venv_pytorch_* | head -1)\n"
+    "sudo -u ubuntu $NEURON_VENV/bin/pip install -e '/home/ubuntu/trnblas[dev]' --quiet\n"
+    + bench_cmds
+)
+encoded = base64.b64encode(script.encode()).decode()
+print(json.dumps({"commands": [f"echo '{encoded}' | base64 -d | bash"]}))
+PYEOF
+
+CMD_ID=$(aws ssm send-command \
+  --instance-ids "$INSTANCE_ID" \
+  --document-name "AWS-RunShellScript" \
+  --comment "trnblas bench @ $SHA shapes=${SHAPES_ARG}" \
+  --parameters "file://$PARAMS_FILE" \
+  --timeout-seconds 7200 \
+  --region "$REGION" \
+  --output text --query 'Command.CommandId')
+rm -f "$PARAMS_FILE"
+
+echo "Command ID: $CMD_ID"
+echo "Waiting for bench to complete (cold compile for large shape: 30-90 min)..."
+
+# Poll every 30s, up to 120 min.
+STATUS=InProgress
+for _ in $(seq 1 240); do
+  STATUS=$(aws ssm get-command-invocation \
+    --command-id "$CMD_ID" \
+    --instance-id "$INSTANCE_ID" \
+    --region "$REGION" \
+    --query 'Status' --output text 2>/dev/null || echo Pending)
+  case "$STATUS" in
+    Success|Failed|Cancelled|TimedOut|DeliveryTimedOut|ExecutionTimedOut)
+      break ;;
+  esac
+  echo "  Status: $STATUS — waiting..."
+  sleep 30
+done
+
+echo ""
+echo "=== STDOUT ==="
+aws ssm get-command-invocation \
+  --command-id "$CMD_ID" \
+  --instance-id "$INSTANCE_ID" \
+  --region "$REGION" \
+  --query 'StandardOutputContent' --output text
+
+echo ""
+echo "=== STDERR ==="
+aws ssm get-command-invocation \
+  --command-id "$CMD_ID" \
+  --instance-id "$INSTANCE_ID" \
+  --region "$REGION" \
+  --query 'StandardErrorContent' --output text
+
+echo ""
+echo "=== Status: $STATUS ==="
+
+[[ "$STATUS" == "Success" ]]

From 5fcbf822c78524e056c33de3e5784fd6aa684f21 Mon Sep 17 00:00:00 2001
From: Scott Friedman <3011922+scttfrdmn@users.noreply.github.com>
Date: Tue, 21 Apr 2026 16:18:27 -0700
Subject: [PATCH 2/6] feat(bench): separate cold/warm passes to avoid HBM OOM +
 record medium cold numbers

---
 CHANGELOG.md         | 34 ++++++++++++++++++++++++++++++++++
 docs/benchmarks.md   | 26 ++++++++++++++++++--------
 examples/df_mp2.py   | 23 ++++++++++++++++++++++-
 scripts/run_bench.sh | 22 +++++++++++++++++-----
 4 files changed, 91 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8137c64..b0c9643 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,40 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- **`--passes` flag for `df_mp2.py --bench`** (`cold`/`warm`/`both`, default `both`).
+  On Trainium at medium/large shapes, running both passes in the same process OOMs:
+  after the cold pass, all loaded NEFFs remain resident in HBM (64 × 244 MB = 15.6 GB
+  DMA spill at medium shape), leaving no headroom for tensor allocation in the warm pass.
+  The fix is two separate process invocations; `run_bench.sh` now does this automatically.
+
+- **`scripts/run_bench.sh`** — runs `df_mp2.py --bench --batched-pair-energy` on the
+  trn1 CI instance via SSM, cold and warm as separate processes. Supports
+  `--shape medium|large` (default: both). Follows the base64-SSM pattern from
+  `run_pyscf_tests.sh`.
+
+### Hardware (2026-04-21, trn1.2xlarge, neuronxcc 2.24.5133)
+
+**Medium-shape cold timing** (`nbasis=512, nocc=64, nvir=448, naux=1536`):
+
+| Step | Cold |
+|---|---:|
+| Cholesky | 29.7 s |
+| Half-transform | 103.5 s |
+| Metric contraction | 4.0 s |
+| Energy (64 i-dispatches) | 101.3 s |
+| **Total** | **238.5 s** |
+
+Measured from a partially-warm EBS cache (GEMM/SYRK/TRSM NEFFs cached from prior
+test-suite runs; energy kernel compiled fresh). E = −2.487218×10⁰ Ha.
+
+**HBM constraint confirmed:** after the medium cold pass, 64 energy NEFFs + GEMM/SYRK/TRSM
+NEFFs fill 15.9 GB of the 16 GB device. A subsequent in-process warm pass fails with
+`Failed to allocate 1.500GB (usage: tensors)`. Warm timing from prior benchmarks (1.536 s
+energy, 4.784 s total) was measured via a separate process loading from EBS NEFF cache —
+that remains the correct production number. Large-shape warm pending `run_bench.sh` rerun.
+
 ## [0.5.4] — 2026-04-17
 
 ### Added
diff --git a/docs/benchmarks.md b/docs/benchmarks.md
index 8566ecc..c34189d 100644
--- a/docs/benchmarks.md
+++ b/docs/benchmarks.md
@@ -146,9 +146,12 @@ cached-failed-NEFF path → torch.matmul fallback on CPU.
 `@nki.jit` call per i-row processes all `nocc` j-pairs. 64 i-dispatches × ~24 ms
 each = 1.536 s warm energy (XLA dispatch overhead dominates; Tensor Engine executes
 each kernel in ~1 ms). Cold energy = 34 min (77 NEFF compilations at ~27 s each;
-paid once per instance lifetime). Device HBM note: 64 loaded energy NEFFs ×
-244 MB DMA spill ≈ 15.6 GB fills the 16 GB device; a `Failed to allocate 1.5 GB`
-warning is logged during warm setup but computation succeeds.
+paid once per instance lifetime). **Device HBM note (confirmed 2026-04-21):** at
+medium shape, all 64 loaded energy NEFFs remain resident after the cold pass —
+12.6 GB DMA spill + 900 MB model code = 15.9 GB total. A warm pass in the same
+process fails with `Failed to allocate 1.500GB (usage: tensors)` — no headroom
+remains. Warm timing must be measured in a separate process that loads from the
+EBS NEFF cache; `run_bench.sh` does this via `--passes cold` then `--passes warm`.
 
 Energies agree to FP32 noise: -2.487220e+00 (torch), -2.487219e+00 (fused-gemm),
 -2.487221e+00 (batched-pair fallback), -2.487218e+00 (chunked NKI).
@@ -169,11 +172,18 @@ Energy matches bit-for-bit within fp32 reduction-order noise.
 (GA102 Ampere) launched Apr 2021 — closest single-GPU match on AWS.
 A10G via `g5.xlarge` (~$1/hr), trn1 via `trn1.2xlarge` (~$1.34/hr).
 
-| Shape                | Flops   | trn1 NKI warm | A10G warm | A10G vs trn1 |
-|----------------------|--------:|--------------:|----------:|-------------:|
-| small (128/16/384)   | 3.4 G   | 0.091 s       | 0.001 s   | 91×          |
-| medium (512/64/1536) | 2 757 G | **4.784 s** (v0.5.4†) | 0.266 s | **18×** |
-| large (768/96/2304)  | 20 352 G | (not re-run) | 2.018 s   | —            |
+| Shape                | Flops   | trn1 NKI cold | trn1 NKI warm | A10G warm | A10G vs trn1 |
+|----------------------|--------:|--------------:|--------------:|----------:|-------------:|
+| small (128/16/384)   | 3.4 G   | —             | 0.091 s       | 0.001 s   | 91×          |
+| medium (512/64/1536) | 2 757 G | **238.5 s**†† | **4.784 s** (v0.5.4†) | 0.266 s | **18×** |
+| large (768/96/2304)  | 20 352 G | —            | —             | 2.018 s   | —            |
+
+†† Medium cold (2026-04-21, `run_bench.sh --shape medium`): chol 29.7 s, half 103.5 s,
+metric 4.0 s, energy 101.3 s = 238.5 s total. Measured from a partially-warm EBS cache
+(GEMM/SYRK/TRSM NEFFs hit cache; energy kernel compiled fresh). A fully cold start
+(empty cache) would take longer. Cold is paid once per instance lifetime; warm numbers
+are the production-relevant metric. Large shape: in-process warm OOMs at medium (HBM
+saturated); separate-process cold pending.
 
 † v0.5.4 chunked dispatch (warm total 4.784 s). Prior v0.5.3 used CPU fallback
 (9.910 s). The 18× gap vs A10G is down from 37× in v0.5.3.
diff --git a/examples/df_mp2.py b/examples/df_mp2.py
index 27912ec..6b9631f 100644
--- a/examples/df_mp2.py
+++ b/examples/df_mp2.py
@@ -262,7 +262,16 @@ def bench(
     use_fused: bool = False,
     use_fused_gemm: bool = False,
     use_batched_pair: bool = False,
+    passes: str = "both",
 ):
+    """Run cold and/or warm timing for a single bench shape.
+
+    On Trainium, all loaded NEFFs stay resident in HBM after the cold pass.
+    At medium/large shapes (nocc≥64), this saturates the 16 GB device and
+    leaves no room for tensor allocations in the warm pass.  The correct
+    way to measure warm timing is to run this script twice in separate
+    processes (run_bench.sh does this automatically via --passes cold/warm).
+    """
     nbasis, nocc, naux = _BENCH_SHAPES[shape_name]
     nvir = nbasis - nocc
     flops = _flops(nbasis, nocc, naux)
@@ -282,7 +291,8 @@ def bench(
         f"device: {device}  energy_mode: {energy_mode}"
     )
 
-    for label in ("cold", "warm"):
+    labels = {"cold": ["cold"], "warm": ["warm"], "both": ["cold", "warm"]}[passes]
+    for label in labels:
         t = {}
         t0 = time.perf_counter()
         e = df_mp2_energy(
@@ -341,6 +351,16 @@ def main():
         help="Route the energy step through nki_batched_pair_energy (single dispatch "
         "for all nocc² pairs, #43 v0.5.2 — eliminates ~100ms × nocc² overhead).",
     )
+    parser.add_argument(
+        "--passes",
+        choices=["cold", "warm", "both"],
+        default="both",
+        help="Which timing pass(es) to run (default: both). On Trainium at medium/large "
+        "shapes, running 'both' in-process OOMs because all loaded NEFFs stay resident "
+        "in HBM after the cold pass. Use 'cold' for the first invocation and 'warm' for "
+        "a second invocation after the EBS NEFF cache is populated. run_bench.sh does "
+        "this automatically.",
+    )
     args = parser.parse_args()
 
     if args.bench:
@@ -352,6 +372,7 @@ def main():
                 use_fused=args.fused_energy,
                 use_fused_gemm=args.fused_gemm_energy,
                 use_batched_pair=args.batched_pair_energy,
+                passes=args.passes,
             )
         return
 
diff --git a/scripts/run_bench.sh b/scripts/run_bench.sh
index b3e7d91..bbaa04c 100755
--- a/scripts/run_bench.sh
+++ b/scripts/run_bench.sh
@@ -98,7 +98,13 @@ if [[ "$PING" != "Online" ]]; then
 fi
 
 # Build the remote bench command: run df_mp2.py --bench for each shape.
-# Large shape cold compile can take 30-90 min (chunked NEFF compilations).
+# Cold and warm are run as SEPARATE process invocations to avoid HBM OOM.
+#
+# At medium/large shapes, the Neuron runtime keeps all loaded NEFFs resident
+# in HBM after the cold pass (~15.9 GB at medium, nocc=64). A second in-process
+# warm pass then fails to allocate tensor workspace (needs 1.5 GB, HBM full).
+# Solution: cold pass populates the EBS NEFF cache; a fresh process for the
+# warm pass loads NEFFs from EBS (fast) without carrying the cold residue.
 SHAPES_ARG="${SHAPES[*]}"
 echo "Sending bench command (SHA=$SHA, shapes=${SHAPES_ARG})..."
 
@@ -111,15 +117,21 @@ import base64, json, os
 sha    = os.environ["SHA_VAL"]
 shapes = os.environ["SHAPES_VAL"].split()
 
-# Build the bench commands: one python invocation per shape.
+# Each shape: cold pass first (compiles + caches NEFFs), then warm pass
+# in a fresh process (loads from EBS cache).
 bench_cmds = ""
 for shape in shapes:
-    bench_cmds += (
-        f"echo '--- shape={shape} ---'\n"
+    run = (
         f"sudo -u ubuntu env PATH=$NEURON_VENV/bin:/usr/bin:/bin TMPDIR=/var/tmp"
         f" TRNBLAS_REQUIRE_NKI=1"
         f" $NEURON_VENV/bin/python /home/ubuntu/trnblas/examples/df_mp2.py"
-        f" --bench --shape {shape} --batched-pair-energy\n"
+        f" --bench --shape {shape} --batched-pair-energy"
+    )
+    bench_cmds += (
+        f"echo '--- shape={shape} cold ---'\n"
+        f"{run} --passes cold\n"
+        f"echo '--- shape={shape} warm ---'\n"
+        f"{run} --passes warm\n"
     )
 
 script = (

From 35a2e00d8ca070ee4ea5e7af50c16638143caad7 Mon Sep 17 00:00:00 2001
From: Scott Friedman <3011922+scttfrdmn@users.noreply.github.com>
Date: Tue, 21 Apr 2026 16:42:37 -0700
Subject: [PATCH 3/6] bench: record EBS-warm timing + fix log truncation +
 document disk-full for large
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three findings from the 2026-04-21 medium bench run:

1. EBS-warm cold (all NEFFs loaded from disk, not compiled): 137.2 s total
   - Half-transform 5.1 s (was 103.5 s compile-cold; 20× faster from EBS)
   - Energy 101.0 s (64 energy NEFFs still load serially at ~1.3 s/NEFF ≈ 83 s)
   - Adds a distinct "EBS-warm" column to the end-to-end table

2. Corrects CHANGELOG: prior 4.784 s warm figure is in-process HBM-warm,
   not separate-process EBS-warm (which gives ~137 s as now confirmed).
   HBM-warm cannot be reproduced (OOM after cold pass at medium shape).

3. Large-shape cold failed: LLVM ERROR IO failure (No space left on device)
   during neuronxcc compilation. Documented; needs disk investigation.

Fix: add NEURON_RT_LOG_LEVEL=WARNING to run_bench.sh to suppress per-NEFF
INFO messages that ate the SSM 24 KB stdout budget and truncated warm output.
Also add df -h /var/tmp before each shape for disk diagnostics.
---
 CHANGELOG.md         | 38 +++++++++++++++++++++++---------------
 docs/benchmarks.md   | 39 ++++++++++++++++++++++++++++-----------
 scripts/run_bench.sh |  7 +++++++
 3 files changed, 58 insertions(+), 26 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b0c9643..8efed4f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,24 +22,32 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Hardware (2026-04-21, trn1.2xlarge, neuronxcc 2.24.5133)
 
-**Medium-shape cold timing** (`nbasis=512, nocc=64, nvir=448, naux=1536`):
-
-| Step | Cold |
-|---|---:|
-| Cholesky | 29.7 s |
-| Half-transform | 103.5 s |
-| Metric contraction | 4.0 s |
-| Energy (64 i-dispatches) | 101.3 s |
-| **Total** | **238.5 s** |
-
-Measured from a partially-warm EBS cache (GEMM/SYRK/TRSM NEFFs cached from prior
-test-suite runs; energy kernel compiled fresh). E = −2.487218×10⁰ Ha.
+**Medium-shape timing** (`nbasis=512, nocc=64, nvir=448, naux=1536`):
+
+| Step | Compile-cold | EBS-warm |
+|---|---:|---:|
+| Cholesky | 29.7 s | 30.5 s |
+| Half-transform | 103.5 s | 5.1 s |
+| Metric contraction | 4.0 s | 0.6 s |
+| Energy (64 i-dispatches) | 101.3 s | 101.0 s |
+| **Total** | **238.5 s** | **137.2 s** |
+
+Compile-cold: energy kernel compiled fresh; GEMM/SYRK/TRSM NEFFs hit EBS cache from
+prior test-suite runs.
+EBS-warm: all NEFFs loaded from EBS cache (no compilation), but not yet in device HBM.
+Half-transform NEFF load drops 20× (5.1 s vs 103.5 s); energy remains ~101 s because
+64 energy NEFFs still load serially at ~1.3 s/NEFF ≈ 83 s DMA + kernel time.
+E = −2.487218×10⁰ Ha (both passes).
 
 **HBM constraint confirmed:** after the medium cold pass, 64 energy NEFFs + GEMM/SYRK/TRSM
 NEFFs fill 15.9 GB of the 16 GB device. A subsequent in-process warm pass fails with
-`Failed to allocate 1.500GB (usage: tensors)`. Warm timing from prior benchmarks (1.536 s
-energy, 4.784 s total) was measured via a separate process loading from EBS NEFF cache —
-that remains the correct production number. Large-shape warm pending `run_bench.sh` rerun.
+`Failed to allocate 1.500GB (usage: tensors)`. The prior 4.784 s warm figure (1.536 s
+energy) is in-process HBM-warm; it cannot be reproduced via separate-process EBS loading
+(that takes ~137 s as shown above).
+
+**Large-shape cold:** failed with `LLVM ERROR: IO failure on output stream: No space left
+on device` during neuronxcc compilation of large-shape kernels. EBS disk was full after
+medium NEFF cache + compilation artifacts. Needs disk investigation before large can run.
 
 ## [0.5.4] — 2026-04-17
 
diff --git a/docs/benchmarks.md b/docs/benchmarks.md
index c34189d..01abaf7 100644
--- a/docs/benchmarks.md
+++ b/docs/benchmarks.md
@@ -172,21 +172,38 @@ Energy matches bit-for-bit within fp32 reduction-order noise.
 (GA102 Ampere) launched Apr 2021 — closest single-GPU match on AWS.
 A10G via `g5.xlarge` (~$1/hr), trn1 via `trn1.2xlarge` (~$1.34/hr).
 
-| Shape                | Flops   | trn1 NKI cold | trn1 NKI warm | A10G warm | A10G vs trn1 |
-|----------------------|--------:|--------------:|--------------:|----------:|-------------:|
-| small (128/16/384)   | 3.4 G   | —             | 0.091 s       | 0.001 s   | 91×          |
-| medium (512/64/1536) | 2 757 G | **238.5 s**†† | **4.784 s** (v0.5.4†) | 0.266 s | **18×** |
-| large (768/96/2304)  | 20 352 G | —            | —             | 2.018 s   | —            |
+| Shape                | Flops   | trn1 compile-cold | trn1 EBS-warm‡ | trn1 HBM-warm | A10G warm | A10G vs trn1 |
+|----------------------|--------:|------------------:|---------------:|--------------:|----------:|-------------:|
+| small (128/16/384)   | 3.4 G   | —                 | —              | 0.091 s       | 0.001 s   | 91×          |
+| medium (512/64/1536) | 2 757 G | **238.5 s**††     | **137.2 s**‡‡  | **4.784 s**†  | 0.266 s   | **18×**      |
+| large (768/96/2304)  | 20 352 G | —                | —              | —             | 2.018 s   | —            |
 
-†† Medium cold (2026-04-21, `run_bench.sh --shape medium`): chol 29.7 s, half 103.5 s,
+†† Medium compile-cold (2026-04-21, `run_bench.sh --shape medium`): chol 29.7 s, half 103.5 s,
 metric 4.0 s, energy 101.3 s = 238.5 s total. Measured from a partially-warm EBS cache
 (GEMM/SYRK/TRSM NEFFs hit cache; energy kernel compiled fresh). A fully cold start
-(empty cache) would take longer. Cold is paid once per instance lifetime; warm numbers
-are the production-relevant metric. Large shape: in-process warm OOMs at medium (HBM
-saturated); separate-process cold pending.
+(empty cache) would take longer.
 
-† v0.5.4 chunked dispatch (warm total 4.784 s). Prior v0.5.3 used CPU fallback
-(9.910 s). The 18× gap vs A10G is down from 37× in v0.5.3.
+‡ **EBS-warm** = fresh process, all NEFFs loaded from EBS NEFF cache (no compilation),
+but not yet resident in device HBM. This is the timing experienced by any fresh process
+after the instance has been used at least once at this shape.
+
+‡‡ Medium EBS-warm (2026-04-21, second `run_bench.sh --shape medium`): chol 30.5 s,
+half 5.1 s, metric 0.6 s, energy 101.0 s = 137.2 s total. Half-transform NEFF load from
+EBS is now 5.1 s (was 103.5 s when compiled; 20× faster). Energy remains ~101 s because
+the 64 energy NEFFs still load serially from EBS at ~1.3 s/NEFF ≈ 83 s DMA + 17 s execution.
+
+† HBM-warm = NEFFs already resident in device HBM (in-process second pass). Energy step
+costs only the kernel dispatch: 64 i-dispatches × ~24 ms = 1.536 s. v0.5.4 chunked
+dispatch. **HBM-warm is not reproducible at medium via a separate process:** after the
+cold pass, 64 energy NEFFs + GEMM/SYRK/TRSM NEFFs fill 15.9 GB of the 16 GB HBM, leaving
+no headroom for tensor allocation in a second pass. The 4.784 s warm figure is from an
+earlier tracing run; current architecture cannot re-measure it without HBM OOM.
+
+Large cold: failed with `No space left on device` during LLVM compilation (EBS disk full
+after medium NEFF cache + compilation artifacts). Needs investigation before large can run.
+
+† v0.5.4 chunked dispatch. Prior v0.5.3 used CPU fallback (9.910 s). The 18× gap vs A10G
+is down from 37× in v0.5.3.
 
 **Energy bit-exact across platforms:** E_MP2 matches to fp32 noise for
 small (-1.619250e-04) and medium (-2.487218) under real NKI dispatch.
diff --git a/scripts/run_bench.sh b/scripts/run_bench.sh
index bbaa04c..47db885 100755
--- a/scripts/run_bench.sh
+++ b/scripts/run_bench.sh
@@ -119,15 +119,22 @@ shapes = os.environ["SHAPES_VAL"].split()
 
 # Each shape: cold pass first (compiles + caches NEFFs), then warm pass
 # in a fresh process (loads from EBS cache).
+#
+# NEURON_RT_LOG_LEVEL=WARNING suppresses the per-NEFF "[INFO]: Using a
+# cached neff" messages that would otherwise flood SSM's 24 KB stdout
+# limit and truncate the timing output.
 bench_cmds = ""
 for shape in shapes:
     run = (
         f"sudo -u ubuntu env PATH=$NEURON_VENV/bin:/usr/bin:/bin TMPDIR=/var/tmp"
         f" TRNBLAS_REQUIRE_NKI=1"
+        f" NEURON_RT_LOG_LEVEL=WARNING"
         f" $NEURON_VENV/bin/python /home/ubuntu/trnblas/examples/df_mp2.py"
         f" --bench --shape {shape} --batched-pair-energy"
     )
     bench_cmds += (
+        f"echo '--- shape={shape} cold (df -h /var/tmp) ---'\n"
+        f"df -h /var/tmp\n"
         f"echo '--- shape={shape} cold ---'\n"
         f"{run} --passes cold\n"
         f"echo '--- shape={shape} warm ---'\n"

From 2f403a8d8cfd5c1e29afe8ee0afcb6e209be5d00 Mon Sep 17 00:00:00 2001
From: Scott Friedman <3011922+scttfrdmn@users.noreply.github.com>
Date: Tue, 21 Apr 2026 16:53:17 -0700
Subject: [PATCH 4/6] =?UTF-8?q?infra+bench:=20increase=20EBS=20100?=
 =?UTF-8?q?=E2=86=92200G=20+=20fix=20SSM=20output=20truncation?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Root cause of large-shape disk-full: the 100G EBS fills to 99% (95G
used) from Neuron SDK packages + medium NEFF cache alone, leaving only
1.1G for large-shape compilation artifacts.

Changes:
- infra/terraform/main.tf: root_block_device volume_size 100 → 200
- scripts/run_bench.sh:
  - Redirect bench output to /tmp log file; grep timing line to stdout.
    NEURON_RT_LOG_LEVEL=WARNING doesn't suppress the [INFO] NEFF messages
    (they come from libnrt.so); file redirect + grep is the only reliable
    approach to staying under SSM's 24 KB stdout budget.
  - Add growpart/resize2fs at startup to expand filesystem after terraform
    resizes the EBS volume (idempotent: no-op if already full size).
  - df -h / before each pass for disk diagnostics.

Next step: run `terraform apply` in infra/terraform/ to resize the EBS,
then re-run `AWS_PROFILE=aws ./scripts/run_bench.sh`.
---
 infra/terraform/main.tf |  4 ++--
 scripts/run_bench.sh    | 43 +++++++++++++++++++++++++++++------------
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/infra/terraform/main.tf b/infra/terraform/main.tf
index 3bf058e..93eded6 100644
--- a/infra/terraform/main.tf
+++ b/infra/terraform/main.tf
@@ -111,7 +111,7 @@ resource "aws_instance" "ci" {
   associate_public_ip_address = true  # Needed for SSM agent to reach regional endpoint without VPC endpoints
 
   root_block_device {
-    volume_size = 100
+    volume_size = 200  # 100G filled up: Neuron SDK (~70G) + medium NEFF cache leaves <2G for large compilation
     volume_type = "gp3"
   }
 
@@ -126,7 +126,7 @@ resource "aws_instance" "ci" {
     sudo -u ubuntu $NEURON_VENV/bin/pip install -e '/home/ubuntu/trnblas[dev]'
     # neuronxcc compile workdirs can be >5 GB for large NKI kernels.  /tmp is
     # tmpfs (RAM-backed, ~16 GB on trn1.2xlarge) and runs out.  Redirect the
-    # compiler to /var/tmp (EBS-backed, 100 GB) for all ubuntu-user sessions.
+    # compiler to /var/tmp (EBS-backed, 200 GB) for all ubuntu-user sessions.
     echo 'export TMPDIR=/var/tmp' >> /home/ubuntu/.profile
   EOF
 
diff --git a/scripts/run_bench.sh b/scripts/run_bench.sh
index 47db885..9283fe4 100755
--- a/scripts/run_bench.sh
+++ b/scripts/run_bench.sh
@@ -120,30 +120,49 @@ shapes = os.environ["SHAPES_VAL"].split()
 # Each shape: cold pass first (compiles + caches NEFFs), then warm pass
 # in a fresh process (loads from EBS cache).
 #
-# NEURON_RT_LOG_LEVEL=WARNING suppresses the per-NEFF "[INFO]: Using a
-# cached neff" messages that would otherwise flood SSM's 24 KB stdout
-# limit and truncate the timing output.
+# Output filtering: bench output is redirected to a temp file; only the
+# timing line and errors are echoed to stdout.  This avoids the SSM 24 KB
+# stdout limit being consumed by per-NEFF "[INFO]: Using a cached neff"
+# lines that NEURON_RT_LOG_LEVEL cannot suppress (they come from libnrt.so
+# before the Python runtime log level is applied).
 bench_cmds = ""
 for shape in shapes:
     run = (
         f"sudo -u ubuntu env PATH=$NEURON_VENV/bin:/usr/bin:/bin TMPDIR=/var/tmp"
         f" TRNBLAS_REQUIRE_NKI=1"
-        f" NEURON_RT_LOG_LEVEL=WARNING"
         f" $NEURON_VENV/bin/python /home/ubuntu/trnblas/examples/df_mp2.py"
         f" --bench --shape {shape} --batched-pair-energy"
     )
-    bench_cmds += (
-        f"echo '--- shape={shape} cold (df -h /var/tmp) ---'\n"
-        f"df -h /var/tmp\n"
-        f"echo '--- shape={shape} cold ---'\n"
-        f"{run} --passes cold\n"
-        f"echo '--- shape={shape} warm ---'\n"
-        f"{run} --passes warm\n"
-    )
+    for pass_name in ("cold", "warm"):
+        log = f"/tmp/bench_{shape}_{pass_name}.log"
+        bench_cmds += (
+            f"echo '--- shape={shape} {pass_name} (disk) ---'\n"
+            f"df -h / | tail -1\n"
+            f"echo '--- shape={shape} {pass_name} ---'\n"
+            f"set +e\n"
+            f"{run} --passes {pass_name} > {log} 2>&1\n"
+            f"BENCH_EXIT=$?\n"
+            f"set -e\n"
+            f"grep -E '^  {pass_name}:' {log} || true\n"
+            f"if [[ $BENCH_EXIT -ne 0 ]]; then\n"
+            f"  echo 'BENCH FAILED (exit=$BENCH_EXIT):'\n"
+            f"  tail -10 {log}\n"
+            f"  false\n"
+            f"fi\n"
+        )
 
 script = (
     "#!/bin/bash\n"
     "set -euo pipefail\n"
+    # Expand filesystem if the EBS volume was resized via terraform.
+    # growpart/resize2fs are idempotent: they no-op if already at full size.
+    "ROOT_DEV=$(df / --output=source | tail -1)\n"
+    "PARENT=$(lsblk -no PKNAME \"$ROOT_DEV\" 2>/dev/null | head -1 || true)\n"
+    "if [[ -n \"$PARENT\" ]]; then\n"
+    "  sudo growpart \"/dev/$PARENT\" 1 2>/dev/null || true\n"
+    "fi\n"
+    "sudo resize2fs \"$ROOT_DEV\" 2>/dev/null || true\n"
+    "echo \"disk after grow: $(df -h / | tail -1)\"\n"
     "cd /home/ubuntu/trnblas\n"
     "sudo -u ubuntu git fetch --all\n"
     f"sudo -u ubuntu git checkout {sha}\n"

From b6daa1373f1c87e0c61e4bf3cd3ee6f02c3f3219 Mon Sep 17 00:00:00 2001
From: Scott Friedman <3011922+scttfrdmn@users.noreply.github.com>
Date: Tue, 21 Apr 2026 20:41:35 -0700
Subject: [PATCH 5/6] =?UTF-8?q?bench:=20fix=20SSM=20executionTimeout=20?=
 =?UTF-8?q?=E2=80=94=20set=20to=204hr,=20increase=20poll=20loop=20to=204hr?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 scripts/run_bench.sh | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/scripts/run_bench.sh b/scripts/run_bench.sh
index 9283fe4..625ebee 100755
--- a/scripts/run_bench.sh
+++ b/scripts/run_bench.sh
@@ -171,7 +171,13 @@ script = (
     + bench_cmds
 )
 encoded = base64.b64encode(script.encode()).decode()
-print(json.dumps({"commands": [f"echo '{encoded}' | base64 -d | bash"]}))
+# executionTimeout overrides the AWS-RunShellScript default of 3600s.
+# --timeout-seconds on send-command controls delivery only; this param
+# controls how long the script is allowed to run on the instance.
+print(json.dumps({
+    "commands": [f"echo '{encoded}' | base64 -d | bash"],
+    "executionTimeout": ["14400"],
+}))
 PYEOF
 
 CMD_ID=$(aws ssm send-command \
@@ -179,17 +185,18 @@ CMD_ID=$(aws ssm send-command \
   --document-name "AWS-RunShellScript" \
   --comment "trnblas bench @ $SHA shapes=${SHAPES_ARG}" \
   --parameters "file://$PARAMS_FILE" \
-  --timeout-seconds 7200 \
+  --timeout-seconds 14400 \
   --region "$REGION" \
   --output text --query 'Command.CommandId')
 rm -f "$PARAMS_FILE"
 
 echo "Command ID: $CMD_ID"
-echo "Waiting for bench to complete (cold compile for large shape: 30-90 min)..."
+echo "Waiting for bench to complete (cold compile for large shape: 2-4 hr)..."
 
-# Poll every 30s, up to 120 min.
+# Poll every 30s, up to 4 hr (480 polls).
+# Large-shape cold compile observed to take >2 hr (>120 NEFF compilations).
 STATUS=InProgress
-for _ in $(seq 1 240); do
+for _ in $(seq 1 480); do
   STATUS=$(aws ssm get-command-invocation \
     --command-id "$CMD_ID" \
     --instance-id "$INSTANCE_ID" \

From b690efbe9abc34f5950466a8a544047967c0c313 Mon Sep 17 00:00:00 2001
From: Scott Friedman <3011922+scttfrdmn@users.noreply.github.com>
Date: Tue, 21 Apr 2026 20:41:54 -0700
Subject: [PATCH 6/6] infra: ignore user_data/public_ip changes to prevent NEFF
 cache destruction

---
 infra/terraform/main.tf | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/infra/terraform/main.tf b/infra/terraform/main.tf
index 93eded6..83cacc4 100644
--- a/infra/terraform/main.tf
+++ b/infra/terraform/main.tf
@@ -133,6 +133,13 @@ resource "aws_instance" "ci" {
   tags = {
     Name = var.instance_tag
   }
+
+  lifecycle {
+    # Prevent instance replacement when only user_data comments change.
+    # The EBS NEFF cache (100s of GB of compiled kernels) is attached to
+    # this instance; replacement destroys it and forces a full recompile.
+    ignore_changes = [user_data, associate_public_ip_address]
+  }
 }
 
 # ---------------------------------------------------------------------------