docs/bench: architecture gaps cleanup + HBM OOM finding + cold bench numbers

CHANGELOG.md

@@ -7,6 +7,48 @@ 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 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)`. 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
 
 ### Added

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.
 

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,14 +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 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   | —            |
-
-† 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.
+| 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 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.
+
+‡ **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.

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
 

infra/terraform-trn2/main.tf

@@ -169,6 +169,7 @@ resource "aws_instance" "ci" {
   tags = {
     Name = var.instance_tag
   }
+
 }
 
 # ---------------------------------------------------------------------------

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,13 +126,20 @@ 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
 
   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]
+  }
 }
 
 # ---------------------------------------------------------------------------