trntensor 0.17.0__tar.gz → 0.19.0__tar.gz

{trntensor-0.17.0 → trntensor-0.19.0}/.github/workflows/ci.yml

@@ -30,8 +30,10 @@ jobs:
       - uses: astral-sh/setup-uv@v6
         with:
           python-version: ${{ matrix.python-version }}
-      - run: uv sync --extra dev
-      - run: uv run pytest tests/ -v -x --tb=short -m "not neuron and not nki_simulator" --cov=trntensor --cov-report=xml
+      # uv pip (not uv sync): sync resolves every extra, and [torchneuron]'s nki>=0.6
+      # lives only on the Neuron index.
+      - run: uv venv && uv pip install -e ".[dev]"
+      - run: uv run --no-sync pytest tests/ -v -x --tb=short -m "not neuron and not nki_simulator" --cov=trntensor --cov-report=xml
 
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v5
@@ -50,10 +52,10 @@ jobs:
       - uses: astral-sh/setup-uv@v6
         with:
           python-version: "3.12"
-      - run: uv sync --extra dev
+      - run: uv venv && uv pip install -e ".[dev]"
       - name: Install NKI simulator
         run: uv pip install --extra-index-url https://pip.repos.neuron.amazonaws.com "nki>=0.3.0"
       - name: Run simulator-backed kernel tests
         env:
           TRNTENSOR_USE_SIMULATOR: "1"
-        run: uv run pytest tests/ -v -m nki_simulator --tb=short
+        run: uv run --no-sync pytest tests/ -v -m nki_simulator --tb=short

{trntensor-0.17.0 → trntensor-0.19.0}/CHANGELOG.md

@@ -7,12 +7,139 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.19.0] — 2026-10-31
+
+Validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0, native TorchNeuron) — see trnsci/trnsci#35.
+
+Headline: new native MP2-energy and AO→MO kernels beat aten by 1.25–1.9× and 1.3–1.6× in the `auto` region; native NKI matmul/bmm kernels deprecated (#52).
+
+Resolves #51: on the native TorchNeuron path (Trn2/Trn3), the AO→MO and MP2-energy kernels are rewritten for Trn2 and now beat aten, and the NKI matmul/bmm kernels are retired in favour of aten.
+
+### Added
+
+- **`mp2_energy_row_kernel`, the native MP2-energy kernel (#51).**
+  - One launch per occupied row, using the pair symmetry `e_ij = e_ji`.
+  - Per pair, one GEMM fills `T` in SBUF; the `Tᵀ` blocks are Tensor Engine transposes. `T` never reaches device memory.
+  - Gridless launch: at LNC=2 each physical core takes half of the `j` window (`num_programs()` / `program_id`), and it also runs at LNC=1.
+  - The design follows trnblas `pair_energy_sym` (trnblas#55, #58).
+  - Native `"auto"` uses it for device-resident `B` once each row launch carries at least 3 GFLOP. On Trn2 it measured 1.25–1.9× faster than aten: for example 0.69 s against 1.30 s at `nocc=96, nvir=672, naux=2304`.
+  - Near a full device the kernel's input staging (about 1.5× the padded `B`) needs more memory than aten's 64 MiB chunks. `"auto"` keeps aten when live allocations plus that staging would leave less than 3 GiB of the core free, and remembers shapes whose launch failed.
+  - `TRNTENSOR_MP2_ENERGY=auto|rows|aten|legacy` overrides the choice.
+- **`ao_to_mo_trn2_kernel`, the native AO→MO kernel (#51).**
+  - Two plain GEMM stages with the contraction index on the partition axis, so there are no DMA transposes.
+  - The auxiliary index is split across both physical cores at LNC=2.
+  - Native `"auto"` uses it for device-resident `eri` with `nocc ≤ 128` and stage-1 work of at least 5e8 multiply-adds (about nbasis ≥ 256). There it measured 1.3–1.6× faster than the two aten einsums on Trn2, at nbasis 256–768. Smaller problems stay on aten.
+  - `TRNTENSOR_AO_TO_MO=auto|trn2|aten|legacy` overrides the choice.
+- `scripts/bench_native_kernels.py`: a kernel-level A/B harness (aten against the original and the Trn2 kernels), with `--hold-gb` for the memory-pressure regime of aws-neuron-sdk#1447.
+
+### Deprecated
+
+- **On the native path, the NKI `matmul` / `batched_matmul` kernels and the original `mp2_energy` / `ao_to_mo_transform` kernels (#51).** `set_backend("nki")` still runs the matmul/bmm kernels there, with a one-time `FutureWarning`; a later release will route them to aten.
+  - A Trn2 rewrite of matmul used both cores, Tensor Engine transposes and B-tile reuse. It was 1.2–4.3× faster than the original kernel, but still slower than aten at every size (4096³ fp32: 4.10 ms against 3.80 ms; bf16: 1.61 ms against 1.43 ms). It was not kept.
+  - Native `"auto"` already used aten for these.
+  - The XLA path (Trn1) keeps all four original kernels unchanged and does not warn.
+
+### Changed
+
+- Under `set_backend("nki")` on the native path, `mp2_energy` and `ao_to_mo_transform` run the Trn2 kernels instead of the original ones.
+
+## [0.18.0] — 2026-10-30
+
+Native TorchNeuron dispatch (Trn2/Trn3) alongside the PyTorch/XLA path (Trn1, maintenance); validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0) — see trnsci/trnsci#35.
+
+Headline performance: native `"auto"` routes to aten on the device, and the planner contraction-order fix takes the DF-MP2 AO→MO transform (`mi,na,mnP->iaP`, nb=256) from 70 to 1.16 ms (#49).
+
+### Changed
+
+- **Native dispatch thresholds are now measured (#41).** They were measured
+  on trn2.3xlarge (SDK 2.33, NKI 0.7.0) with `scripts/bench_native_einsum.py`
+  and `scripts/autotune_dispatch.py`.
+  - On device-resident operands, aten `torch.matmul` / `torch.bmm` beat the
+    NKI matmul and batched-matmul kernels at every size (128³ to 4096³, fp32
+    and bf16). The native pinned threshold now defaults to `never`, so
+    device work runs as aten on `neuron`.
+  - Host operands at or above `TRNTENSOR_MIN_NKI_FLOPS_NATIVE` (now
+    500 M multiply-adds, previously a provisional 1 G) are offloaded to
+    aten on the device instead of to the NKI kernel.
+  - Thresholds accept `never` / `inf`.
+  - The XLA path is unchanged.
+- **Native `"auto"` runs the fused quantum primitives as aten.** This covers
+  `mp2_energy` and `ao_to_mo_transform`. On Trn2 the AO→MO kernel measured
+  6–20× slower than aten, and the MP2 kernel 1.0–1.5× slower.
+  `set_backend("nki")` still runs both kernels, and the XLA path is
+  unchanged.
+- **`mp2_energy` reference path is vectorized.** Without NKI it now uses
+  aten ops chunked over `i` instead of a nocc² Python loop. That is 10–40×
+  faster on device. On an Apple M-series CPU it is 3× faster at small sizes
+  and about even (±10%) at `nocc=32, nvir=128, naux=512`.
+- **`ao_to_mo_transform` reference path** is now two explicit binary
+  contractions.
+
+### Fixed
+
+- **Greedy contraction path no longer starts with an outer product.** For
+  example, `einsum("mi,na,mnP->iaP")` used to contract `mi,na` first,
+  building a 4-index intermediate. Pairs that share an index are now
+  preferred.
+- **`scripts/autotune_dispatch.py` now reports M·K·N**, the quantity dispatch
+  compares, where it used to report 2·M·K·N. On the native path it now
+  derives both thresholds: host vs device offload, and NKI vs aten on device.
+
 ## [0.17.0] — 2026-10-28
 
 Validated on Trn2 (trn2.3xlarge, Neuron SDK 2.33.0, native TorchNeuron device) — see trnsci/trnsci#35.
 
 ### Added
 
+- **Native PyTorch dispatch on TorchNeuron (#41, trnsci/trnsci#35).** On
+  TorchNeuron (Neuron SDK 2.33+, `torch-neuronx` 2.14, PyTorch 2.14,
+  Trn2/Trn3) every host-wrapper NKI kernel is a PyTorch custom op registered
+  with `torch_neuronx.nki_op` and launched on `device="neuron"`, with no
+  `torch_xla` import: `trntensor::matmul`, `batched_matmul`, `mp2_energy`,
+  `ao_to_mo_transform`. Kernel bodies are unchanged. The PyTorch/XLA path
+  (Trn1) is unchanged.
+  - The dispatch path is detected on the first kernel launch, not at
+    import: native when `torch_neuronx` exports `nki_op` and
+    `torch_neuronx.device_count() > 0`, XLA when `torch_xla` is installed,
+    otherwise none. `TRNSCI_NKI_DISPATCH=auto|native|xla` (suite-wide)
+    forces a path; forcing an unavailable one raises.
+  - The ops are registered lazily (`trntensor/nki/_native.py` is imported
+    only on the native path). Every launch goes through one helper,
+    `dispatch._launch`. That covers `einsum` (matmul/bmm/path steps),
+    `multi_einsum` (batched and per-contraction), `ao_to_mo_transform`,
+    `mp2_energy`, and `precision="sr"` (which passes `use_sr=True` to the op).
+    `precision="kahan"` stays a host fp64 contraction, and `"dd"` still
+    raises on NKI hosts (trnblas#22). The reduce-parallel `nki.collectives`
+    all-reduce from Phase 4 is unchanged.
+  - New `to_device` / `from_device` place tensors on the active path's
+    device (`neuron` or XLA). `to_xla` / `from_xla` stay public. On the
+    native path `to_xla` places on `neuron`, so existing residency code keeps
+    working. `multi_einsum` pre-pins shared operands on either device.
+  - FLOP dispatch thresholds are path-aware. The XLA values are unchanged
+    (2 GFLOPs unpinned, 1 GFLOPs pinned). The native path uses provisional
+    1 GFLOPs / 128 MFLOPs, because eager native dispatch costs about 0.2 ms
+    per launch. Override them with `TRNTENSOR_MIN_NKI_FLOPS_NATIVE[_PINNED]`
+    or the `trntensor_min_nki_flops_native[_pinned]` autotune cache keys.
+    `scripts/autotune_dispatch.py` measures and writes the native keys on a
+    TorchNeuron host.
+  - Dormant deprecation hook: with `TRNSCI_XLA_DEPRECATION=1` (off by
+    default), selecting the XLA path and calling `to_xla` / `from_xla` each
+    emit a one-time `DeprecationWarning`. Migration: replace `to_xla` with
+    `to_device` and `from_xla` with `from_device` (docs/api/nki.md).
+  - The `mp2_energy` CPU fallback copies `neuron`-resident operands to the
+    host before its scalar loop, then moves the result back.
+  - New extras: `[torchneuron]` (`torch>=2.14`, `torch-neuronx>=2.14`,
+    `nki>=0.6`) and `[neuron]` for the XLA lineage (`torch-neuronx>=2.9`,
+    `nki>=0.3.0`).
+  - Tests: `tests/conftest.py` sets `NKI_ENABLE_TRACE_CACHE=0` by default.
+    New CPU tests in `tests/test_native_dispatch.py` cover path selection,
+    launch routing for every entry point, lazy registration, op schemas,
+    fake shapes, thresholds, and the deprecation hook. New
+    `@pytest.mark.neuron` tests in `tests/test_nki_native.py` have not yet
+    run on hardware. `TestXlaResidency::test_matmul_stays_on_xla` is now
+    path-aware: it expects `neuron` on TorchNeuron and `xla` on Trn1, and
+    still checks residency.
+
 - **Reduce-parallel sharding on real multi-rank all-reduce (Phase 4 prototype).**
   Before this change, reduce-parallel and mixed output/reduce-parallel einsums
   raised `NotImplementedError` whenever `HAS_NKI` was true.

{trntensor-0.17.0 → trntensor-0.19.0}/CLAUDE.md

@@ -24,7 +24,10 @@ trntensor/
 │   ├── decompose.py     # cp_decompose, tucker_decompose (HOSVD)
 │   └── nki/
 │       ├── __init__.py
-│       └── dispatch.py  # Fused contraction kernels
+│       ├── _kernels.py  # @nki.jit kernels
+│       ├── _native.py   # TorchNeuron custom ops (nki_op), loaded lazily
+│       ├── collectives.py  # nki.collectives all-reduce (reduce-parallel)
+│       └── dispatch.py  # host wrappers, dispatch-path selection, _launch
 ├── tests/
 │   ├── test_einsum.py   # Matmul, bmm, 4-index transform, DF-MP2 pattern
 │   └── test_decompose.py # CP rank-1/low-rank, Tucker orthogonality

{trntensor-0.17.0/trntensor.egg-info → trntensor-0.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trntensor
-Version: 0.17.0
+Version: 0.19.0
 Summary: Tensor contractions for AWS Trainium via NKI
 Author-email: Scott Friedman <scttfrdmn@gmail.com>
 License-Expression: Apache-2.0
@@ -18,8 +18,15 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.1
 Requires-Dist: numpy>=1.24
+Provides-Extra: neuron
+Requires-Dist: nki>=0.3.0; extra == "neuron"
+Requires-Dist: torch-neuronx>=2.9; extra == "neuron"
+Provides-Extra: torchneuron
+Requires-Dist: torch>=2.14; extra == "torchneuron"
+Requires-Dist: torch-neuronx>=2.14; extra == "torchneuron"
+Requires-Dist: nki>=0.6; extra == "torchneuron"
 Provides-Extra: dev
-Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest<9,>=7.0; extra == "dev"
 Requires-Dist: pytest-benchmark>=4.0; extra == "dev"
 Requires-Dist: pytest-cov>=4.0; extra == "dev"
 Requires-Dist: ruff>=0.6; extra == "dev"
@@ -59,8 +66,32 @@ Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1)
 
 ```bash
 pip install trntensor
+
+# Trn2/Trn3 with TorchNeuron (native PyTorch, Neuron SDK 2.33+)
+pip install "trntensor[torchneuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com
+
+# Trn1 / PyTorch-XLA lineage
+pip install "trntensor[neuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com
 ```
 
+Kernels dispatch on whichever Neuron PyTorch lineage is installed:
+
+- **TorchNeuron (native)**: Trn2/Trn3, Neuron SDK 2.33+, Python ≥ 3.12,
+  `torch` 2.14, `torch-neuronx` 2.14.*. Each kernel is a PyTorch custom op
+  (`torch.ops.trntensor.*`) on the `neuron` device; no `torch_xla`.
+- **PyTorch/XLA**: Trn1 (and any host on the `torch-neuronx` 2.9 lineage).
+  Kernels dispatch through `torch_xla`, unchanged.
+
+trntensor picks the native path when the native `torch_neuronx` and a
+`neuron` device are present, and the XLA path when `torch_xla` is
+installed; otherwise it runs on PyTorch. Set `TRNSCI_NKI_DISPATCH=native`
+or `xla` (shared by all trnsci libraries) to force a path. Pin operands on
+the device with `to_device` / `from_device`; the older `to_xla` /
+`from_xla` still work on both lineages (see the
+[migration note](docs/api/nki.md#migrating-from-to_xla-from_xla)). With an
+editable install (`pip install -e`), set `NKI_ENABLE_TRACE_CACHE=0`: the NKI
+kernel cache only versions installed wheels and can serve stale kernels.
+
 ## Usage
 
 ```python
@@ -103,7 +134,8 @@ core, factors = trntensor.tucker_decompose(tensor, ranks=(5, 5, 5))
 - [x] Tucker decomposition (HOSVD)
 - [x] DF-MP2 einsum example
 - [x] NKI fused contraction kernels (`mp2_energy`, `ao_to_mo_transform`)
-- [x] XLA operand residency (`to_xla` / `from_xla`)
+- [x] Device operand residency (`to_device` / `from_device`; `to_xla` / `from_xla` kept)
+- [x] Native TorchNeuron dispatch (Trn2/Trn3, `torch.ops.trntensor.*`) alongside PyTorch/XLA (Trn1)
 - [x] NKI CPU simulator + `nki-simulator` CI gate
 - [x] Optimal contraction ordering — greedy path search for 3+ operands
 - [x] Multi-contraction shared-operand XLA pre-pinning

trntensor-0.17.0/PKG-INFO → trntensor-0.19.0/README.md

@@ -1,34 +1,3 @@
-Metadata-Version: 2.4
-Name: trntensor
-Version: 0.17.0
-Summary: Tensor contractions for AWS Trainium via NKI
-Author-email: Scott Friedman <scttfrdmn@gmail.com>
-License-Expression: Apache-2.0
-Project-URL: Homepage, https://github.com/trnsci/trntensor
-Project-URL: Documentation, https://trnsci.dev/trntensor/
-Project-URL: Repository, https://github.com/trnsci/trntensor
-Project-URL: Issues, https://github.com/trnsci/trntensor/issues
-Classifier: Development Status :: 3 - Alpha
-Classifier: Intended Audience :: Science/Research
-Classifier: Programming Language :: Python :: 3
-Classifier: Topic :: Scientific/Engineering
-Classifier: Topic :: Scientific/Engineering :: Mathematics
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: torch>=2.1
-Requires-Dist: numpy>=1.24
-Provides-Extra: dev
-Requires-Dist: pytest>=7.0; extra == "dev"
-Requires-Dist: pytest-benchmark>=4.0; extra == "dev"
-Requires-Dist: pytest-cov>=4.0; extra == "dev"
-Requires-Dist: ruff>=0.6; extra == "dev"
-Requires-Dist: pre-commit>=3.7; extra == "dev"
-Requires-Dist: scipy>=1.11; extra == "dev"
-Requires-Dist: mkdocs>=1.5; extra == "dev"
-Requires-Dist: mkdocs-material>=9.0; extra == "dev"
-Dynamic: license-file
-
 # trntensor
 
 [![CI](https://github.com/trnsci/trntensor/actions/workflows/ci.yml/badge.svg)](https://github.com/trnsci/trntensor/actions/workflows/ci.yml)
@@ -59,8 +28,32 @@ Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1)
 
 ```bash
 pip install trntensor
+
+# Trn2/Trn3 with TorchNeuron (native PyTorch, Neuron SDK 2.33+)
+pip install "trntensor[torchneuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com
+
+# Trn1 / PyTorch-XLA lineage
+pip install "trntensor[neuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com
 ```
 
+Kernels dispatch on whichever Neuron PyTorch lineage is installed:
+
+- **TorchNeuron (native)**: Trn2/Trn3, Neuron SDK 2.33+, Python ≥ 3.12,
+  `torch` 2.14, `torch-neuronx` 2.14.*. Each kernel is a PyTorch custom op
+  (`torch.ops.trntensor.*`) on the `neuron` device; no `torch_xla`.
+- **PyTorch/XLA**: Trn1 (and any host on the `torch-neuronx` 2.9 lineage).
+  Kernels dispatch through `torch_xla`, unchanged.
+
+trntensor picks the native path when the native `torch_neuronx` and a
+`neuron` device are present, and the XLA path when `torch_xla` is
+installed; otherwise it runs on PyTorch. Set `TRNSCI_NKI_DISPATCH=native`
+or `xla` (shared by all trnsci libraries) to force a path. Pin operands on
+the device with `to_device` / `from_device`; the older `to_xla` /
+`from_xla` still work on both lineages (see the
+[migration note](docs/api/nki.md#migrating-from-to_xla-from_xla)). With an
+editable install (`pip install -e`), set `NKI_ENABLE_TRACE_CACHE=0`: the NKI
+kernel cache only versions installed wheels and can serve stale kernels.
+
 ## Usage
 
 ```python
@@ -103,7 +96,8 @@ core, factors = trntensor.tucker_decompose(tensor, ranks=(5, 5, 5))
 - [x] Tucker decomposition (HOSVD)
 - [x] DF-MP2 einsum example
 - [x] NKI fused contraction kernels (`mp2_energy`, `ao_to_mo_transform`)
-- [x] XLA operand residency (`to_xla` / `from_xla`)
+- [x] Device operand residency (`to_device` / `from_device`; `to_xla` / `from_xla` kept)
+- [x] Native TorchNeuron dispatch (Trn2/Trn3, `torch.ops.trntensor.*`) alongside PyTorch/XLA (Trn1)
 - [x] NKI CPU simulator + `nki-simulator` CI gate
 - [x] Optimal contraction ordering — greedy path search for 3+ operands
 - [x] Multi-contraction shared-operand XLA pre-pinning

trntensor-0.19.0/docs/api/nki.md

@@ -0,0 +1,187 @@
+# NKI dispatch
+
+Backend selection between the PyTorch fallback and the NKI-accelerated
+path on Trainium. Kernels live in `trntensor.nki._kernels`; the
+public-facing dispatch is in `trntensor.nki.dispatch`.
+
+## `set_backend(backend: str)`
+
+- `"auto"` — use NKI if `nki` is importable, else PyTorch
+- `"pytorch"` — always use `torch.matmul` / `torch.einsum`
+- `"nki"` — require NKI; raises `RuntimeError` on non-Neuron hosts
+
+## `get_backend() -> str`
+
+Return the current backend selection.
+
+## `HAS_NKI: bool`
+
+Module-level flag set at import time. `True` when `nki` (0.3.0+) is
+importable, otherwise `False`.
+
+## Dispatch path
+
+On Neuron hardware, kernels run on one of two `torch-neuronx` lineages,
+detected on the first kernel launch (never at import):
+
+- **Native (TorchNeuron, Trn2/Trn3, Neuron SDK 2.33+)** — chosen when
+  `torch_neuronx` exports `nki_op` and a `neuron` device is available
+  (`torch_neuronx.device_count() > 0`). Each kernel is a PyTorch custom op
+  on `device="neuron"`, also callable directly and inside
+  `torch.compile(backend="neuron")`: `torch.ops.trntensor.matmul`,
+  `batched_matmul`, `mp2_energy`, `ao_to_mo_transform`. The ops are
+  registered only once this path is selected, and this path never imports
+  `torch_xla`.
+- **XLA (Trn1)** — chosen when `torch_xla` is installed and the native
+  path is unavailable. Kernels run on the `torch_xla` device, unchanged.
+- Otherwise there is no device path: the host wrappers fall back to
+  `torch` (or raise under `TRNTENSOR_REQUIRE_NKI=1`), as before.
+
+`set_backend` is unaffected: it still chooses NKI vs PyTorch. Set the
+suite-wide `TRNSCI_NKI_DISPATCH=native` or `xla` to force a path; forcing an
+unavailable path raises `RuntimeError`.
+
+The reduce-parallel all-reduce kernel (`trntensor.nki.collectives`) is
+native-only already; it keeps its `torch_neuronx.wrap_nki` launch because
+its `ReplicaGroup` argument is a compile-time constant that a custom-op
+schema cannot carry.
+
+Precision modes follow the same launch: `precision="sr"` passes
+`use_sr=True` to the `matmul` / `batched_matmul` op (stochastic rounding
+still needs `NEURON_RT_STOCHASTIC_ROUNDING_EN=1` at process start and warns
+once otherwise); `precision="kahan"` contracts in true fp64 on the host
+(Trainium has no fp64 datapath) on both paths; `precision="dd"` still
+raises on NKI hosts pending trnblas#22.
+
+## `to_device(tensor) -> Tensor` / `from_device(tensor) -> Tensor`
+
+Pin an operand on the Trainium device of the active dispatch path —
+`neuron` on TorchNeuron, the XLA device on Trn1 — so repeated trntensor
+calls skip the host↔device transfer that otherwise dominates dispatch
+overhead.
+
+```python
+import trntensor
+
+# One-time transfer onto the accelerator.
+eri_d = trntensor.to_device(eri)
+C_occ_d = trntensor.to_device(C_occ)
+C_vir_d = trntensor.to_device(C_vir)
+eps_occ_d = trntensor.to_device(eps_occ)
+eps_vir_d = trntensor.to_device(eps_vir)
+
+# Full DF-MP2 pipeline — B_d never leaves the device.
+B_d = trntensor.ao_to_mo_transform(eri_d, C_occ_d, C_vir_d)
+E_d = trntensor.mp2_energy(B_d, eps_occ_d, eps_vir_d)
+
+# Pull the scalar back when we actually need it in Python.
+E = trntensor.from_device(E_d)
+```
+
+- `to_device`: no-op when the tensor is already on a Neuron device.
+  Raises `RuntimeError` on hosts without the NKI runtime.
+- `from_device`: no-op when the tensor is already on CPU.
+
+When every operand is already on the device, dispatch skips the per-call
+transfer, uses the lower "pinned" FLOP threshold, and returns the result on
+the device — the caller controls when to pull back.
+
+## `to_xla(tensor) -> Tensor` / `from_xla(tensor) -> Tensor`
+
+The original names, kept for compatibility. On Trn1 they behave exactly as
+before. On TorchNeuron there is no XLA device, so `to_xla` places the
+tensor on `neuron` instead (residency works the same way) and `from_xla`
+moves it back to CPU.
+
+### Migrating from `to_xla` / `from_xla`
+
+Replace `to_xla` with `to_device` and `from_xla` with `from_device`. Same
+behaviour on Trn1; on Trn2/Trn3 the new names say what they do. Checks such
+as `t.device.type == "xla"` become `t.device.type in ("xla", "neuron")`.
+
+`to_xla` / `from_xla` stay available. When the suite-wide
+`TRNSCI_XLA_DEPRECATION=1` is set (off by default; trnsci turns it on at
+TorchNeuron GA), each emits a one-time `DeprecationWarning` pointing at the
+new names, and selecting the XLA dispatch path emits: "the PyTorch/XLA path
+is deprecated on Trn2/Trn3; use TorchNeuron (native). Trn1/Inf2 remain
+supported in maintenance mode."
+
+## Environment variables
+
+| Variable | Effect |
+|---|---|
+| `TRNTENSOR_REQUIRE_NKI=1` | Re-raise kernel exceptions instead of falling back to PyTorch. Useful in the validation loop to surface silent kernel breakage. |
+| `TRNTENSOR_MIN_NKI_FLOPS=<int>` | XLA path: override the FLOP threshold (M·K·N) below which dispatch skips NKI and uses `torch.matmul` / `torch.bmm` directly. Default: `2_000_000_000`. Set to `0` to force NKI for kernel validation. |
+| `TRNTENSOR_MIN_NKI_FLOPS_PINNED=<int>` | XLA path, operands already on the device. Default: `1_000_000_000`. |
+| `TRNTENSOR_MIN_NKI_FLOPS_NATIVE=<int>` | Native (TorchNeuron) path, operands on CPU: M·K·N at which the work leaves the host for the `neuron` device. Below it, `torch.matmul` / `torch.bmm` run on the host. Above it, the operands are offloaded. The offload runs aten unless the pinned threshold below selects NKI. Default: `500_000_000` (measured on Trn2, see below). |
+| `TRNTENSOR_MIN_NKI_FLOPS_NATIVE_PINNED=<int>\|never` | Native path: M·K·N at which the NKI kernel replaces aten on the device. Below it, device-resident operands run `torch.matmul` / `torch.bmm` (aten) on `neuron`. Default: `never`, because on Trn2 aten beat the NKI kernels at every measured size. |
+| `TRNSCI_NKI_DISPATCH=auto\|native\|xla` | Suite-wide. Force the native or XLA dispatch path; `auto` (default) detects. Forcing an unavailable path raises. |
+| `TRNSCI_XLA_DEPRECATION=1` | Suite-wide, off by default. Emit the one-time XLA-path and `to_xla` / `from_xla` deprecation warnings. |
+| `NKI_ENABLE_TRACE_CACHE=0` | TorchNeuron setting (not trntensor's). Turns off the cross-process NKI kernel cache, which cannot version editable installs and may serve a stale kernel after a source edit. Set it with `pip install -e`; the test suite sets it by default. |
+| `TRNTENSOR_MP2_ENERGY=auto\|rows\|aten\|legacy` | Native path: force the `mp2_energy` executor (`rows` = the Trn2 row kernel, `legacy` = the original kernel, deprecated there). |
+| `TRNTENSOR_AO_TO_MO=auto\|trn2\|aten\|legacy` | Native path: force the `ao_to_mo_transform` executor (`trn2` = the Trn2 kernel, `legacy` = the original kernel, deprecated there). |
+| `TRNTENSOR_FORCE_BACKEND=pytorch\|nki\|auto` | Override the backend selection at runtime without calling `set_backend()`. Used by benchmarks to sweep both paths on the same machine. |
+| `TRNTENSOR_USE_SIMULATOR=1` | Route kernel dispatch through `nki.simulate(kernel)(numpy_args)` on CPU instead of XLA → NEFF → hardware. Catches Python-trace-level errors without AWS round-trips. MLIR verifier errors remain hardware-only. |
+
+### Native thresholds: measured on Trn2
+
+These were measured on a trn2.3xlarge (Neuron SDK 2.33, NKI 0.7.0), using one
+NeuronCore and two host threads, with `scripts/bench_native_einsum.py` and
+`scripts/autotune_dispatch.py` (#41):
+
+- **Pinned operands: NKI never beats aten.** On device-resident operands,
+  aten `torch.matmul` / `torch.bmm` beat the trntensor NKI matmul and
+  batched-matmul kernels at every size from 128³ to 4096³, in both fp32 and
+  bf16. Examples:
+  - 4096³ bf16: 1.38 ms with aten, 6.82 ms with NKI.
+  - 1024³ fp32: 0.20 ms with aten, 0.55 ms with NKI.
+
+  The eager launch floor is about 0.13 ms for an aten op and about 0.24 ms for
+  an NKI custom op.
+- **Host operands: offload to aten from about 0.5 G multiply-adds.**
+  Offloading to aten (upload, run on device, download) beats offloading
+  through the NKI kernel at every size. It also beats host torch:
+  - fp32: from about 134 M multiply-adds (512³) when the host has two
+    threads.
+  - bf16: from about 0.5 to 1 G multiply-adds.
+
+  The default of 500 M sits between these. A host with more threads moves the
+  crossover up. Re-measure on your host with `scripts/autotune_dispatch.py`.
+
+Kernel status on the native path after #51:
+
+- **matmul / bmm: retired in favour of aten.** A Trn2 rewrite used both
+  physical cores through the `program_id` split, Tensor Engine transposes of
+  A and B-tile reuse across four stationaries. It was 1.2–4.3× faster than
+  the original kernel and still slower than aten at every size: 4096³ fp32
+  took 4.10 ms against 3.80 ms, and 4096³ bf16 1.61 ms against 1.43 ms.
+  It was not kept. `set_backend("nki")` still runs the original kernels on
+  the native path, with a one-time `FutureWarning`. A later release will
+  route them to aten.
+- **AO→MO: kept, rewritten for Trn2** (`ao_to_mo_trn2_kernel`). It is
+  1.3–1.6× faster than the two aten einsums at nbasis 256–768, and native
+  `"auto"` uses it there for device-resident `eri`. Below nbasis ≈ 256
+  `"auto"` keeps aten.
+- **MP2 energy: kept, rewritten for Trn2** (`mp2_energy_row_kernel`). It is
+  1.25–1.9× faster than aten once each row launch carries at least 3 GFLOP,
+  and native `"auto"` uses it there.
+
+The original AO→MO and MP2 kernels remain the XLA-path (Trn1) kernels. See
+[Quantum-chemistry primitives](quantum.md) for the tables.
+
+## Kernels (internal)
+
+- `matmul_kernel(a, b)` — 2-index matmul with stationary-A tile reuse.
+  Tile constants: `TILE_M=128`, `TILE_K=128`, `TILE_N=512`.
+- `batched_matmul_kernel(a, b)` — per-batch-slice matmul. Batch dim
+  iterated via `nl.affine_range`; each slice reuses the stationary-A
+  tile layout.
+
+The DF-MP2 pair pattern `einsum("ap,bp->ab")` is served by
+`matmul_kernel` through the planner's `transB=True` route — no
+dedicated kernel needed. Fused energy-denominator kernel is tracked
+in [#13][i13].
+
+[i13]: https://github.com/trnsci/trntensor/issues/13
+
+[m2]: https://github.com/trnsci/trntensor/milestone/2

trntensor-0.19.0/docs/api/quantum.md

@@ -0,0 +1,152 @@
+# 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, and small problems on native TorchNeuron**: aten ops on the
+  operands' device. `T[i, j, a, b]` is contracted with `torch.einsum` in
+  chunks over `i`, which bounds peak memory to about 64 MiB of `T` at a time.
+  The `T`-weighted sum is vectorized.
+- **Native TorchNeuron (Trn2/Trn3): the row kernel.** `mp2_energy_row_kernel`
+  is launched once per occupied row `i`. It covers `j = i + d (mod nocc)` for
+  `d = 0..nocc/2` and uses the pair symmetry `e_ij = e_ji`. Per pair, one
+  GEMM fills `T = B[i] B[j]ᵀ` in SBUF, and the `Tᵀ` blocks come from Tensor
+  Engine transposes. `T` never reaches device memory. At LNC=2 each physical
+  core takes half of the `j` window. The design follows trnblas
+  `pair_energy_sym` (trnblas#55, #58).
+  - `"auto"` uses the row kernel when `B` is on the device and each row
+    launch carries at least 3 GFLOP, counted on the padded sizes as
+    `2·(nocc//2+1)·nvir²·naux`. Below that, the ~nocc custom-op launch
+    floors lose to aten.
+  - `set_backend("nki")` uses it whenever its SBUF working set fits.
+  - `TRNTENSOR_MP2_ENERGY=auto|rows|aten|legacy` overrides the choice.
+    `legacy` is the original single-launch kernel, which is deprecated on
+    the native path.
+
+  Measured on Trn2 (fp32, one LNC=2 core, device-resident, warm, #51):
+
+  | nocc, nvir, naux | aten | row kernel |
+  |---|---|---|
+  | 16, 128, 512 | 1.5 ms | 5.1 ms |
+  | 32, 256, 768 | 12.1 ms | 11.6 ms |
+  | 32, 320, 1024 | 24.9 ms | 19.8 ms |
+  | 48, 384, 1152 | 75.7 ms | 39.2 ms |
+  | 64, 448, 1536 | 240 ms | 120 ms |
+  | 96, 672, 2304 | 1.30 s | 0.69 s |
+
+  Holding 12 or 18 GB of the core's 24 GB on the device did not slow either
+  path. The aten path's 64 MiB `T` chunks stay clear of the slowdown in
+  aws-neuron-sdk#1447. The row kernel's input staging is about 1.5× the
+  padded `B`, so near a full device it needs more memory than aten does.
+  `"auto"` therefore keeps aten when this process's live device allocations
+  plus that staging would leave less than 3 GiB free. A shape whose kernel
+  launch fails anyway falls back to aten and is remembered for the rest of
+  the process.
+- **PyTorch/XLA path (Trn1)**: the original single-launch `mp2_energy_kernel`.
+  It accumulates `T` in PSUM, builds `Δ` on the Vector Engine, and writes one
+  partial per `(i, j)` pair.
+
+## `trntensor.ao_to_mo_transform(eri, C_occ, C_vir) -> Tensor`
+
+Fused 4-index AO→MO integral transform. The canonical quantum-chemistry primitive.
+
+```
+B[i, a, P] = Σ_{μ, ν} C_occ[μ, i] · C_vir[ν, a] · eri[μ, ν, P]
+```
+
+### Arguments
+
+- `eri: (nbasis, nbasis, naux) tensor` — density-fitted two-electron integrals in the AO basis
+- `C_occ: (nbasis, nocc) tensor` — occupied molecular orbital coefficients
+- `C_vir: (nbasis, nvir) tensor` — virtual molecular orbital coefficients
+
+### Returns
+
+A `(nocc, nvir, naux)` tensor of transformed DF coefficients B_ia^P.
+
+### Example — full DF-MP2 pipeline
+
+```python
+import trntensor
+
+B = trntensor.ao_to_mo_transform(eri, C_occ, C_vir)
+E = trntensor.mp2_energy(B, eps_occ, eps_vir)
+```
+
+### Backend behaviour
+
+- **CPU**: two explicit aten steps on the operands' device:
+  `einsum("mi,mnP->inP")` then `einsum("na,inP->iaP")`.
+- **Native TorchNeuron (Trn2/Trn3): `ao_to_mo_trn2_kernel`.** Both stages are
+  plain GEMMs with the contraction index on the partition axis, so there are
+  no DMA transposes.
+  - Stage 1: `X[i, ν, P] = Σ_μ C_occ[μ, i] eri[μ, ν, P]`.
+  - Stage 2: `B[i, a, P] = Σ_ν C_vir[ν, a] X[i, ν, P]`.
+  - The auxiliary index is split across both physical cores at LNC=2.
+  - `"auto"` uses the kernel for device-resident `eri` with `nocc ≤ 128` when
+    its SBUF working set fits and the stage-1 work `nocc·nbasis²·naux` is at
+    least 5e8 multiply-adds (about nbasis ≥ 256). Smaller problems lose to
+    the launch floor and the padding copies; otherwise it runs the two aten
+    steps.
+    `TRNTENSOR_AO_TO_MO=auto|trn2|aten|legacy` overrides the choice.
+
+  Measured on Trn2 (fp32, one LNC=2 core, device-resident, warm, #51):
+
+  | nbasis, nocc, naux | two aten einsums | Trn2 kernel (op) | `ao_to_mo_transform`, `"auto"` | original kernel |
+  |---|---|---|---|---|
+  | 64, 8, 128 | 0.41 ms | — | 0.41 ms (aten) | 6.0 ms (#49) |
+  | 128, 16, 256 | 0.49 ms | 0.43 ms | 0.49 ms (aten) | 16.1 ms |
+  | 256, 32, 512 | 1.23 ms | 1.11 ms | 0.92 ms | 112 ms (#49) |
+  | 512, 32, 1024 | 7.30 ms | 4.56 ms | 4.72 ms | — |
+  | 768, 48, 1536 | 18.3 ms | 11.8 ms | 12.0 ms | — |
+
+  The op column takes tile-multiple inputs (the raw kernel). The `"auto"`
+  column includes the padding of `nvir` to 128 and is the shipped route. At
+  nbasis 64 and 128 it measured 0.79 and 0.57 ms with the kernel forced,
+  which is why `"auto"` keeps aten there.
+
+- **PyTorch/XLA path (Trn1)**: the original kernel. It loops over P, with the
+  `(i, ν)` intermediate round-tripping through kernel-scratch HBM. On the
+  native path it is deprecated.
+
+### Current limitations
+
+- Both NKI kernels need `nocc ≤ 128`; larger `nocc` runs the aten steps (native) or raises `NotImplementedError` (XLA).
+- The `(nocc, nbasis, naux)` intermediate goes through kernel-scratch HBM between the two stages, as it does between the two aten einsums. Stage 1 leaves `i` on the partition axis, and stage 2 needs `ν` there.