docs: tell the fused mp2_energy story + honest perf note

- docs/architecture.md: new section explaining the fused
  contract+elementwise+reduce pattern as trntensor's architectural
  core, with the mp2_energy kernel as the canonical demonstration
- docs/api/quantum.md (new): mp2_energy reference
- docs/api/index.md + mkdocs.yml: Quantum primitives nav entry
- docs/tutorial_df_mp2.md: 'Fusing the whole calculation' section
- docs/benchmarks.md: add fused vs Python-loop row with honest note
  that dispatch+compile overhead still dominates at current sizes

Author: scttfrdmn

docs/api/index.md

@@ -5,4 +5,5 @@ trntensor's public API is organized into four modules:
 - **[einsum](einsum.md)** — `einsum`, `multi_einsum`
 - **[plan](plan.md)** — `plan_contraction`, `ContractionPlan`, `estimate_flops`
 - **[decompose](decompose.md)** — `cp_decompose`, `tucker_decompose`, and reconstructors
+- **[quantum](quantum.md)** — fused chemistry primitives (`mp2_energy`)
 - **[nki](nki.md)** — backend dispatch (`set_backend`, `get_backend`, `HAS_NKI`)

docs/api/quantum.md

@@ -0,0 +1,56 @@
+# Quantum-chemistry primitives
+
+Domain-specific fused kernels where the whole computation — contraction, elementwise shaping, reduction — lives in a single NKI program. The architectural core of trntensor (see [Architecture](../architecture.md)).
+
+## `trntensor.mp2_energy(B, eps_occ, eps_vir) -> scalar`
+
+Density-fitted second-order Møller–Plesset correlation energy.
+
+```
+E_MP2 = Σ_{i,j,a,b} T_{i,j,a,b} (2 T_{i,j,a,b} - T_{i,j,b,a}) / Δ_{i,j,a,b}
+
+T_{i,j,a,b}  = Σ_P  B[i, a, P] B[j, b, P]
+Δ_{i,j,a,b}  = ε_i + ε_j - ε_a - ε_b
+```
+
+### Arguments
+
+- `B: (nocc, nvir, naux) tensor` — density-fitted ERI coefficients
+- `eps_occ: (nocc,) tensor` — occupied orbital energies
+- `eps_vir: (nvir,) tensor` — virtual orbital energies
+
+### Returns
+
+A 0-D tensor containing the correlation energy.
+
+### Example
+
+```python
+import torch
+import trntensor
+
+nocc, nvir, naux = 5, 19, 72
+B = torch.randn(nocc, nvir, naux) * 0.1
+eps_occ = -torch.sort(torch.rand(nocc))[0] - 0.5
+eps_vir =  torch.sort(torch.rand(nvir))[0] + 0.1
+
+E = trntensor.mp2_energy(B, eps_occ, eps_vir)
+print(f"E_MP2 = {E.item():.6f}")
+```
+
+### Backend behaviour
+
+- **CPU**: falls back to a Python loop over `(i, j)` pairs composing `torch.einsum`
+  and element-wise ops. Same as `examples/df_mp2_einsum.py`.
+- **Trainium (NKI)**: dispatches a single `@nki.jit` program that
+  - accumulates `T` in PSUM via `nisa.nc_matmul`,
+  - builds `Δ` on the Vector Engine from SBUF-resident `ε` tiles,
+  - folds the energy into a scalar accumulator in SBUF,
+  - writes one partial per `(i, j)` pair to HBM.
+
+The host sums the `(nocc, nocc)` partial matrix into the final scalar. No intermediate four-index `T` tensor is ever materialized.
+
+### Current limitations
+
+- Single-tile path only: `nvir ≤ 128` and `naux ≤ 128`. Larger systems raise `NotImplementedError` — K/M tiling is a follow-up.
+- Dispatch overhead still dominates at small sizes; see [Benchmarks](../benchmarks.md) for the honest comparison against the Python loop.

docs/architecture.md

@@ -26,6 +26,31 @@ The planner analyzes einsum subscripts and selects a dispatch target:
 
 `ContractionPlan` carries the dispatch decision, the FLOPs estimate, and any reshape/transpose preamble so the same plan can be re-executed cheaply.
 
+## Fused contraction-reduction: the architectural core
+
+What makes trntensor different from `torch.einsum` or `trnblas.gemm` is **fusion across the boundary between contraction and reduction**. On Trainium:
+
+- the Tensor Engine does the contraction, accumulating into PSUM
+- the Vector Engine does the elementwise reshaping, reading from SBUF
+- a small SBUF scalar accumulates the reduction
+- nothing lands in HBM until the whole program finishes
+
+A cuTENSOR clone would compose these three phases as separate ops with HBM round-trips between them. NKI lets us write them as one `@nki.jit` program with the same data living in PSUM → SBUF → accumulator across phases. That's the architectural lever.
+
+The canonical demonstration is `trntensor.mp2_energy` — the DF-MP2 correlation energy is
+
+```
+T_{i,j,a,b} = Σ_P B_{i,a,P} B_{j,b,P}              (contraction)
+term        = T * (2T - T^T) / Δ_{i,j,a,b}         (elementwise)
+E           = Σ_{i,j,a,b} term                      (reduction)
+```
+
+The fused kernel in `trntensor/nki/_kernels.py::mp2_energy_kernel` does all three in one program, iterating over `(i, j)` pairs in `affine_range`. No intermediate `T` tensor is ever materialized to HBM.
+
+This pattern — fused contract + elementwise + reduce with SBUF-resident intermediates — is what trntensor extends across the suite. Generic `einsum` dispatch to `matmul` / `bmm` kernels is the foundation; primitives like `mp2_energy` are the architectural flagships.
+
+**Performance honest note.** The fused kernel is architecturally correct but currently carries ~15–40 ms of XLA dispatch + compile overhead per call, which means at small chemistry sizes the CPU fallback path still wins. The fusion work pays off once per-call overhead is amortized (tracked in #33/#34). See [Benchmarks](benchmarks.md) for the current numbers.
+
 ## Decompositions for quantum chemistry
 
 - **CP (CANDECOMP/PARAFAC)** — Tensor hypercontraction (THC) of two-electron integrals. Reduces $O(N^4)$ storage to $O(N^2 R)$.

docs/benchmarks.md

@@ -33,6 +33,19 @@ Both columns ran on the same trn1.2xlarge instance — CPU (Intel Xeon 8375C) vs
 
 **NKI wins**: 2048×2048 matmul (1.6× faster than CPU). All other sizes still favor CPU on this hardware.
 
+## DF-MP2 energy — fused vs Python-loop
+
+Fused `trntensor.mp2_energy` (single NKI program) vs the reference loop from `examples/df_mp2_einsum.py` (25–256 einsum calls), both on trn1.2xlarge:
+
+| Workload | Python loop | Fused NKI | Winner |
+|---|---:|---:|---|
+| `(nocc, nvir, naux) = (5, 19, 72)` | **1.5 ms** | 15.6 ms | Loop 10× |
+| `(nocc, nvir, naux) = (16, 128, 128)` | **25.5 ms** | 41.3 ms | Loop 1.6× |
+
+The fused kernel is **architecturally** what we want — one program, no intermediate HBM materialization of the four-index `T` tensor, PSUM/SBUF-resident across contract → elementwise → reduce. But NKI's per-call dispatch + compile overhead (~15–40 ms floor) still eats the win at current sizes. The gap closes with scale (10× → 1.6× as workload grows), consistent with the overhead being a fixed cost the compute amortizes. See [#33][i33] and [#34][i34].
+
+
+
 ## Size-based dispatch threshold
 
 Because per-call NKI dispatch currently carries ~1 ms of XLA launch overhead, `nki_matmul` and `nki_batched_matmul` short-circuit to the PyTorch path when the contraction is below `TRNTENSOR_MIN_NKI_FLOPS` (default **2 GFLOPs**, calibrated at ≈ half the smallest NKI-winning size). The `plan.backend` field reflects this: it reports `"nki"` only when the dispatch will actually invoke a kernel, and `"pytorch"` otherwise.
@@ -54,3 +67,4 @@ Recommendation for users:
 - Until #33 lands, tight loops of small contractions (DF-MP2 pair-energy style) see no NKI benefit and are served by the PyTorch path.
 
 [i33]: https://github.com/trnsci/trntensor/issues/33
+[i34]: https://github.com/trnsci/trntensor/issues/34

docs/tutorial_df_mp2.md

@@ -92,6 +92,23 @@ for i in range(nocc):
 
 will reveal exactly where NKI is or isn't being used.
 
+## Fusing the whole calculation
+
+The Python loop above does 25 einsum calls, 25 elementwise passes, and 25 host reductions — each landing in HBM between steps. On Trainium, trntensor can compile the entire DF-MP2 energy into **one NKI program**:
+
+```python
+E = trntensor.mp2_energy(B, eps_occ, eps_vir)
+```
+
+Inside that single call:
+
+1. `T_ab = Σ_P B[i,a,P] B[j,b,P]` accumulates in PSUM via `nc_matmul`
+2. The spin-adapted numerator `(2T − T^T)` and the denominator `Δ_ab = ε_i + ε_j − ε_a − ε_b` are built on the Vector Engine from SBUF-resident ε tiles
+3. `(T·(2T − T^T) / Δ).sum()` folds into a scalar accumulator in SBUF
+4. One HBM partial per `(i, j)` pair; host sums to the final scalar
+
+No intermediate `T` tensor is ever materialized to HBM. See [API: quantum](api/quantum.md) and [Architecture](architecture.md) for the full description.
+
 ## Current limitations
 
 - The energy denominator is a separate element-wise op. A fused kernel that folds the division into the PSUM accumulation — `E = Σ T²/Δ` in one kernel — is tracked in [#13][i13]. Landing it collapses the pair-energy loop into a single tensor-engine invocation per `(i,j)`.

mkdocs.yml

@@ -25,6 +25,7 @@ nav:
       - einsum: api/einsum.md
       - Planning: api/plan.md
       - Decompositions: api/decompose.md
+      - Quantum primitives: api/quantum.md
       - NKI dispatch: api/nki.md
   - Architecture: architecture.md
   - Benchmarks: benchmarks.md