trntensor 0.17.0__tar.gz → 0.18.0__tar.gz

{trntensor-0.17.0 → trntensor-0.18.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.18.0}/CHANGELOG.md

@@ -7,12 +7,103 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [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.18.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.18.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: trntensor
-Version: 0.17.0
+Version: 0.18.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.18.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.18.0/docs/api/nki.md

@@ -0,0 +1,174 @@
+# 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_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`.
+
+Under `"auto"` on the native path, the fused quantum kernels (`mp2_energy`,
+`ao_to_mo_transform`) also run as aten, because both measured slower than
+aten on the device:
+
+- AO→MO was 6–20× slower.
+- MP2 energy was 1.0–1.5× slower, reaching parity at best.
+
+`set_backend("nki")` still runs the NKI kernels. The kernels are unchanged,
+and the XLA path keeps them under `"auto"`.
+
+## 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.17.0 → trntensor-0.18.0}/docs/api/quantum.md

@@ -40,9 +40,15 @@ 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
+- **CPU, and native TorchNeuron under `"auto"`**: 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. This replaces the former Python loop over `(i, j)` pairs, which
+  ran 10–40× slower on a Trainium device.
+  On Trn2 the aten form matched or beat the fused kernel. For example, with
+  `nocc=32, nvir=128, naux=512` it took 3.1 ms, against 4.5 ms for the
+  kernel (#41).
+- **Trainium (NKI): the PyTorch/XLA path, or `set_backend("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,
@@ -85,8 +91,12 @@ E = trntensor.mp2_energy(B, eps_occ, eps_vir)
 
 ### Backend behaviour
 
-- **CPU**: `torch.einsum("mi,na,mnP->iaP", C_occ, C_vir, eri)` — a single three-operand einsum.
-- **Trainium (NKI)**: dispatches a single `@nki.jit` program that
+- **CPU, and native TorchNeuron under `"auto"`**: two explicit aten steps,
+  `einsum("mi,mnP->inP")` then `einsum("na,inP->iaP")`, on the operands'
+  device. On Trn2 this was 6–20× faster than the fused kernel. For example,
+  with `nbasis=128, nocc=16, naux=256` it took 0.82 ms, against 16.4 ms for
+  the kernel (#41).
+- **Trainium (NKI): the PyTorch/XLA path, or `set_backend("nki")`**: dispatches a single `@nki.jit` program that
   - loads `C_occ` and `C_vir` once, SBUF-resident across all P iterations,
   - for each auxiliary index P: two sequential `nisa.nc_matmul` calls with the intermediate `(i, ν)` tile round-tripping through kernel-scratch HBM to handle the partition-dim change,
   - writes one `(nocc, nvir)` slice per P to the output HBM tensor.

{trntensor-0.17.0 → trntensor-0.18.0}/docs/contraction_patterns.md

@@ -15,7 +15,7 @@ as the primary design surface.
 | Multiple `Plan` objects with HBM intermediates | `multi_einsum()` with shared-operand pre-pinning |
 | `cutensorDecompositionCreate` (Tucker HOSVD) | `tucker_decompose()` |
 | Custom THC / Cholesky-factored ERIs | `cp_decompose()` (tensor hypercontraction) |
-| CUDA stream + `cudaMemcpy` to pin device memory | `to_xla(tensor)` / `from_xla(tensor)` |
+| CUDA stream + `cudaMemcpy` to pin device memory | `to_device(tensor)` / `from_device(tensor)` (formerly `to_xla` / `from_xla`) |
 
 ## 2-Operand contractions — the generic path
 
@@ -119,12 +119,14 @@ XLA-resident copy is used for all subsequent contractions in the batch.
 
 On CPU (no NKI), `multi_einsum` behaves exactly like a loop over `einsum()`.
 
-## XLA residency: `to_xla` / `from_xla`
+## Device residency: `to_device` / `from_device`
 
-The NKI dispatch path always transfers operands to the Trainium XLA device
-before kernel launch. For a single large matmul this cost is negligible; for a
-tight loop of small contractions it dominates. `to_xla` / `from_xla` expose the
-transfer explicitly so callers can control when it happens:
+The NKI dispatch path always transfers operands to the Trainium device
+(`neuron` on TorchNeuron, the XLA device on Trn1) before kernel launch. For a
+single large matmul this cost is negligible; for a tight loop of small
+contractions it dominates. `to_device` / `from_device` expose the transfer
+explicitly so callers can control when it happens. `to_xla` / `from_xla` are
+the older names for the same thing and still work on both lineages:
 
 ```python
 # Cold path: B transferred on every call (3 transfers for 3 calls)
@@ -132,17 +134,17 @@ for _ in range(3):
     result = trntensor.einsum("ij,jk->ik", A, B)
 
 # Residency path: B transferred once, result stays on device until needed
-A_x = trntensor.to_xla(A)
-B_x = trntensor.to_xla(B)
+A_x = trntensor.to_device(A)
+B_x = trntensor.to_device(B)
 for _ in range(3):
     result_x = trntensor.einsum("ij,jk->ik", A_x, B_x)  # zero transfer cost
-result = trntensor.from_xla(result_x)
+result = trntensor.from_device(result_x)
 ```
 
 The residency API maps to the CUDA concept of a device-pinned buffer, but with
 explicit Python-level transfer rather than stream-implicit staging.
 
-`to_xla` raises `RuntimeError` on hosts without `nki` installed.
+`to_device` (like `to_xla`) raises `RuntimeError` on hosts without `nki` installed.
 
 ## 3+ operand einsums: contraction-path search
 

{trntensor-0.17.0 → trntensor-0.18.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "trntensor"
-version = "0.17.0"
+version = "0.18.0"
 description = "Tensor contractions for AWS Trainium via NKI"
 readme = "README.md"
 license = "Apache-2.0"
@@ -25,8 +25,24 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+neuron = [
+    # PyTorch/XLA lineage (Trn1, and Trn2 on PyTorch 2.9). NKI 0.3.0 Stable
+    # (Neuron SDK 2.29) promotes `nki.*` as the canonical namespace. Ships with
+    # the Deep Learning AMI Neuron or via https://pip.repos.neuron.amazonaws.com.
+    "nki>=0.3.0",
+    "torch-neuronx>=2.9",
+]
+torchneuron = [
+    # TorchNeuron / native PyTorch (Neuron SDK 2.33+, Trn2/Trn3; torch-neuronx
+    # 2.14 itself requires Python >= 3.12). Install with
+    # --extra-index-url https://pip.repos.neuron.amazonaws.com. The nki floor
+    # also keeps pip off the unrelated `nki` 0.0.1 on public PyPI.
+    "torch>=2.14",
+    "torch-neuronx>=2.14",
+    "nki>=0.6",
+]
 dev = [
-    "pytest>=7.0",
+    "pytest>=7.0,<9",  # pytest-benchmark 5.3 breaks on pytest 9 (Function._nodeid)
     "pytest-benchmark>=4.0",
     "pytest-cov>=4.0",
     "ruff>=0.6",
@@ -47,11 +63,13 @@ include = ["trntensor*"]
 
 # Neuron SDK note: on-hardware runs use the AWS Deep Learning AMI's
 # pre-built venv at /opt/aws_neuronx_venv_pytorch_*, which bundles
-# neuronxcc and the NKI runtime. There is no pip-install step for Neuron —
-# the DLAMI's terraform `most_recent = true` pulls the current image and
-# its venv tracks the current SDK (2.29 as of 2026-04-09, with NKI 0.3.0
-# Stable and the CPU simulator). trntensor does not declare neuronxcc in
-# its dep set; CPU-only installs resolve cleanly from PyPI.
+# the NKI runtime. The DLAMI's terraform `most_recent = true` pulls the
+# current image and its venv tracks the current SDK. Outside the DLAMI, use
+# the `[torchneuron]` (Trn2/Trn3, native) or `[neuron]` (Trn1, PyTorch/XLA)
+# extra with the Neuron pip repo. Both lineages publish a package named
+# `torch-neuronx`, so pip markers cannot pick between them; choose the extra
+# that matches the instance. The base dependency set stays CPU-only and
+# resolves cleanly from PyPI.
 
 [tool.ruff]
 # Match the rest of the trnsci suite. Line length 100 is a compromise