diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..6d3a392 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,30 @@ +name: ci + +# Host-only checks. TT hardware is not available on GitHub runners, so device / accuracy / +# OOM / perf tests are NOT run here — they gate releases on real cards before a tag is cut +# (see RELEASING.md). This job proves the package builds, is publishable, and imports +# without a card. + +on: + pull_request: + push: + branches: [master] + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + - name: Build sdist + wheel + run: | + python -m pip install --upgrade build twine + python -m build + - name: Check package metadata + run: python -m twine check dist/* + - name: Import check (no card, no heavy deps) + run: | + python -m pip install --no-deps dist/*.whl + python -c "import tt_atom; print('tt_atom', tt_atom.__version__)" diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml new file mode 100644 index 0000000..3da6885 --- /dev/null +++ b/.github/workflows/release.yaml @@ -0,0 +1,66 @@ +name: release + +# Cut a release by pushing a version tag AFTER the on-hardware release gate is green +# (accuracy parity, no OOM across the supported size range, no perf regression — see +# RELEASING.md): +# git tag v0.1.0 && git push origin master --tags +# This builds the artifacts and publishes a GitHub Release with notes + wheel. No PyPI +# job: tt-atom requires a source tt-metal/ttnn build, so a pip wheel can't run standalone +# (see the NOTE below and RELEASING.md). + +on: + push: + tags: ["v*"] + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + - name: Tag must match pyproject version + run: | + tag="${GITHUB_REF_NAME#v}" + ver="$(grep -m1 -E '^version *=' pyproject.toml | sed -E 's/.*"([^"]+)".*/\1/')" + if [ "$tag" != "$ver" ]; then + echo "::error::tag v$tag != pyproject version $ver"; exit 1 + fi + - name: Build + run: | + python -m pip install --upgrade build twine + python -m build + python -m twine check dist/* + - uses: actions/upload-artifact@v4 + with: + name: dist + path: dist/* + + github-release: + needs: build + runs-on: ubuntu-latest + permissions: + contents: write + steps: + - uses: actions/checkout@v4 + - uses: actions/download-artifact@v4 + with: + name: dist + path: dist + - name: Create GitHub Release (notes from CHANGELOG section) + env: + GH_TOKEN: ${{ github.token }} + run: | + ver="${GITHUB_REF_NAME#v}" + awk -v v="$ver" '/^## /{p=($0 ~ "\\["v"\\]")} p' CHANGELOG.md > NOTES.md || true + if [ -s NOTES.md ]; then + gh release create "$GITHUB_REF_NAME" dist/* --title "$GITHUB_REF_NAME" --notes-file NOTES.md + else + gh release create "$GITHUB_REF_NAME" dist/* --title "$GITHUB_REF_NAME" --generate-notes + fi + + # NOTE: tt-atom is distributed via GitHub Releases only, NOT PyPI. It is the + # custom-kernel-only build and requires a source tt-metal/ttnn build (see RELEASING.md + # / README Install), so a `pip install tt-atom` wheel can't run standalone — publishing + # to PyPI would be misleading. (tt-bio, which is genuinely pip-installable, keeps PyPI.) diff --git a/.gitignore b/.gitignore index 3ad5228..a302338 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,26 @@ __pycache__/ checkpoints/ assets/*.npz .DS_Store + +# Generated benchmark artifacts (committed charts live in assets/; raw data does not) +benchmarks/results/ + +# Golden parity fixtures ARE committed (small, our own random weights, not gated). Allow them: +!tests/data/ +!tests/data/*.npz + +# Vendored fairchem Wigner-D coefficient table is a needed source asset (not a checkpoint). Allow it: +!tt_atom/assets/ +!tt_atom/assets/*.pt + +# ttnn runtime artifacts +generated/ + +# packaging +*.egg-info/ + +# benchmark outputs (machine-specific, regenerated) +benchmarks/results/ +dist/ +build/ +*.egg-info/ diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..1e49de0 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,29 @@ +# Changelog + +All notable changes to TT-Atom are recorded here. Versioning is [SemVer](https://semver.org); +releases are cut only from a commit that has passed the on-hardware release gate — accuracy +parity, no OOM across the supported size range, and no perf regression (see `RELEASING.md`). + +## [0.1.0] - 2026-07-08 + +Initial release. The **custom-kernel-only, highest-performance build for `uma-s`** — the per-edge +Wigner rotation runs as a custom tt-metal kernel, so `ttnn` comes from a source tt-metal build +that includes the op (see README "Install"); there is no slow fallback path. + +### Added +- Tenstorrent inference for Meta **UMA** (eSEN / eSCN-MD) equivariant ML interatomic potentials: + energy, conservative analytic forces, and stress for molecules and periodic materials, behind an + **ASE** calculator that mirrors fairchem's (moving off fairchem is a one-line change). Validated + against the released `uma-s-1`. +- Device-resident trace loop for MD / relaxation; multi-card data-parallel throughput path. +- `tt-atom verify` device round-trip check and a one-command checkpoint converter. + +### Performance (uma-s-1, Blackhole p150a) +- Fused-rotation kernel: **4.3×** vs the addcmul MAC in isolation (7.01 → 1.62 ms, PCC 0.999995); + **1.4–1.68× faster end-to-end** traced MD/relax across N=54–2662, no regression at any size. +- Accuracy (vs fairchem reference): energy rel-error ≤ 5.4e-4, force PCC ≥ 0.9996 across + molecular / periodic / slab; traced == eager (PCC 1.0). pytest 51 passed / 1 skipped. + +### Scope +- `uma-s` (lmax=mmax=2) is the supported target. Other checkpoints (e.g. `uma-m`) raise a clear + error rather than silently falling back. `ttnn` is not a pip dependency (source build required). diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..2eb10bd --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Moritz Thüning + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index 0fafe09..9b2d612 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,178 @@ # TT-Atom -High-performance Tenstorrent inference for eSEN/eSCN-MD (UMA-family) ML interatomic potentials. +![Caffeine molecular dynamics, uma-s-1 on a Tenstorrent Blackhole card](assets/caffeine_md.gif) -_Under construction._ +Run Meta's [UMA](https://huggingface.co/facebook/UMA) interatomic potential on [Tenstorrent](https://tenstorrent.com). Energy, forces and stress for molecules and periodic materials, behind an [ASE](https://wiki.fysik.dtu.dk/ase/) calculator. Bring your own UMA checkpoint. + +## Install + +TT-Atom is the custom-kernel-only, highest-performance build for `uma-s`. Its per-edge Wigner rotation runs as a custom tt-metal kernel that the pip `ttnn` wheel does not carry, so `ttnn` comes from a **source tt-metal build**. The op is pre-integrated on the [`moritztng/tt-atom`](https://github.com/tenstorrent/tt-metal/tree/moritztng/tt-atom) branch of tt-metal, so the build is a plain clone-and-build — no patching. You need a Tenstorrent card and its driver. + +**1. Build and install tt-metal with the op** (branch `moritztng/tt-atom`): + +```bash +git clone --recursive -b moritztng/tt-atom https://github.com/tenstorrent/tt-metal.git +cd tt-metal +export TT_METAL_HOME=$PWD +./build_metal.sh --build-type Release # full build (tens of minutes) +pip install -e . # tt-metal's own dev-install path +``` + +The branch is the validated base `b5522097b39` plus the `fused_rotate` op library — nothing else. Its source and contract are mirrored in [`custom_kernels/README.md`](custom_kernels/README.md) as the authoritative backup (and for re-integrating onto a newer tt-metal commit). + +`TT_METAL_HOME` must stay exported at **runtime** too — the JIT-compiled kernels load from `$TT_METAL_HOME/build_Release`, so don't delete that directory after installing. + +On some boards/firmware this base commit's UMD misreads the board ID as a dual-chip P300 (`Board ... has 1 chips, but expected 2 chips for board type p300` -> `TT_FATAL: Custom fabric mesh graph descriptor path must be specified for CUSTOM cluster type`), which blocks opening *any* device, single-card included. If you hit that, export `TT_MESH_GRAPH_DESC_PATH=$TT_METAL_HOME/tt_metal/fabric/mesh_graph_descriptors/p150_mesh_graph_descriptor.textproto` before opening a device — this also needs to be set in the parent process before constructing `tt_atom.batch.MultiCard`, since its per-card worker processes inherit it. + +**2. Install TT-Atom into the same venv:** + +```bash +git clone https://github.com/moritztng/tt-atom.git +pip install -e ./tt-atom # numpy<2, torch (CPU), ase — NOT ttnn +``` + +**3. Verify the op is loaded:** + +```bash +python -c "import ttnn; e=ttnn._ttnn.operations.experimental; print(hasattr(e,'fused_rotate'), hasattr(e,'fused_rotate_gc'))" # -> True True +``` + +`uma-s` (lmax=mmax=2) is the validated target; other checkpoints (e.g. uma-m) raise a clear error. `import tt_atom` never imports ttnn, so it imports fine on a machine without a card. + +## Quickstart + +```bash +tt-atom run structure.xyz +``` + +```python +from ase.io import read +from tt_atom import UMA + +atoms = read("structure.xyz") +atoms.calc = UMA(atoms) +atoms.get_potential_energy() +atoms.get_forces() +``` + +`UMA(atoms)` uses `uma-s-1`, infers the task (`omat` if the cell is periodic, else `omol`), and builds a device-resident model for that composition on first use. Later calls load it from cache. Everything downstream is plain ASE. + +## Relax and MD + +```bash +tt-atom run structure.xyz --relax --out relaxed.xyz +tt-atom run structure.xyz --md --steps 200 --temp 300 +``` + +Add `--trace` (or `UMA(atoms, trace=True)`) to replay the captured device graph over the loop. About 2x on relax/MD, forces stay bit-identical. + +## What it supports + +- Models: `uma-s-1` (default), `uma-s-1.2`. See [Model coverage](#model-coverage) for what else + exists upstream and why this build doesn't run it. +- Tasks: `omol`, `omat`, `oc20`, `odac`, `omc`. +- Systems: isolated molecules and periodic cells. Charge and spin via `UMA(atoms, charge=-1, spin=2)`. +- Properties: energy, conservative analytic forces, and stress, so variable-cell relaxation works (see [`examples/relax_cell.py`](examples/relax_cell.py)). + +## Model coverage + +Meta has released two UMA sizes: `uma-s-1` (`.1`/`.2`) and `uma-m-1p1` — there is no `uma-l`. The +[paper](https://arxiv.org/abs/2506.23971) scales capacity via mixture-of-linear-experts on the +small and medium models rather than shipping a third, larger dense tier, and +[facebook/UMA](https://huggingface.co/facebook/UMA) carries checkpoints for only those two. + +Of the two that exist, only `uma-s` runs on this build (both `uma-s-1` and `uma-s-1.2`). `uma-s` is +square (lmax=mmax=2), so its per-edge Wigner rotation is a 9x9 tile that fits the fused kernel's L1 CB +budget. `uma-m-1p1` uses mmax19, W=256) — that overflows the kernel's L1 budget, and this build has no MAC fallback, so it +raises a clear `RuntimeError` naming the shape rather than silently running slow or wrong +(`tests/test_umam.py` anchors this contract). A hypothetical `uma-l`, sized above `uma-m`, would +need L1 headroom `uma-m` already overflows, so it isn't a new question, just a bigger version of +the one above — and moot, since the checkpoint doesn't exist to test it against. + +### uma-s-1.2 + +`uma-s-1.2` adds fairchem's charge-balanced channels: the `l=0` charge channels are re-balanced to the system charge after every block. TT-Atom applies this automatically — point `UMA` at the checkpoint (gated; bring your own): + +```python +atoms.calc = UMA(atoms, checkpoint="uma-s-1p2.pt") +``` + +Parity with fairchem — forces, energy, and stress across 757 molecular and periodic systems, plus a CPU throughput comparison — is written up in [`docs/uma-s-1p2-validation.md`](docs/uma-s-1p2-validation.md). + +## Accuracy + +Every task is checked on-device against the released `uma-s-1` checkpoint run through fairchem on the same structure. + +| task | system | energy rel. err | force PCC | stress PCC | +|------|--------|----------------:|----------:|-----------:| +| omol | ethanol | 2e-7 | 0.9996 | | +| omat | bulk Si | 3e-4 | 0.99999 | 0.99999 | +| oc20 | Cu(100) + H slab| 9e-5 | 1.0000 | | +| odac | MgO framework | 2e-4 | 0.99999 | | +| omc | solid CO2 | 8e-5 | 1.0000 | | + +Dynamics are stable: NVE energy drift is about 1 meV/atom/ps. These numbers are from `ttnn` 0.68.0. Op numerics can shift slightly between `ttnn` versions, so confirm parity on the version you actually run: + +Reproduce it yourself. Every bundle embeds the fairchem reference energy and forces from build time, so: + +```bash +tt-atom verify model.npz # device output vs the embedded fairchem reference +pytest tests/ # full parity suite against fairchem goldens +``` + +## Throughput + +Batch independent systems into a single device pass: + +```python +out = calc.evaluate_batch(list_of_atoms) # out["energy"], out["forces"] +``` + +For many small molecules this is roughly 13x over looping on one card. To use several cards, fan systems across them with `tt_atom.batch` (one process per card). + +For a **batched MD ensemble / relaxation** — K fixed-composition replicas evolving with a stable neighbour list — add `trace=True` to capture the batched device graph once and replay it (forces stay bit-identical; it re-captures whenever the neighbour list changes): + +```python +out = calc.evaluate_batch(replicas, trace=True) # per-step in the ensemble loop +``` + +At small per-system sizes the eager batched forward is host-dispatch-bound below saturation, so the trace lets a *modest* ensemble reach near-peak throughput: measured on one p150 (uma-s-1, 9-atom molecules) K=4 gives 4.2x (59→246 systems/s), K=16 2.6x (207→528 sys/s) — approaching the K≥128 eager device-bound plateau (~700 sys/s) at a fraction of the batch size. Leave it `False` for one-shot screening, where a fresh batch each call would re-capture every time. + +## Compared to fairchem + +TT-Atom is an inference runtime, not a rewrite of fairchem. It reuses the released weights and matches them. + +| | fairchem | TT-Atom | +|--|:--------:|:-------:| +| Hardware | GPU, CPU | Tenstorrent | +| Energy, forces, stress | ✅ | ✅ | +| Molecules, periodic (PBC) | ✅ | ✅ | +| Tasks (omol/omat/oc20/odac/omc) | ✅ | ✅ | +| Models | uma-s, uma-m | uma-s-1, uma-s-1.2 | +| ASE relax and MD | ✅ | ✅ (plus a traced loop) | +| Batched inference | ✅ | ✅ (one composition per batch) | +| LAMMPS interface | ✅ | ❌ | +| Training, fine-tuning | ✅ | ❌ (inference only) | + +## Bundles and the reference environment + +The model is a "bundle": UMA weights merged for one composition. `UMA(atoms)` builds and caches bundles for you, so most users never touch this. To build one yourself: + +```bash +refenv/bin/python tools/export_weights.py --uma-s-1 --xyz structure.xyz --task omol --out model.npz +``` + +then `TTAtomCalculator("model.npz")`. + +Building a bundle needs `fairchem` to read the checkpoint and merge the experts. `fairchem` wants `numpy>=2`, which cannot share a process with `ttnn`'s `numpy<2`, so keep it in its own venv: + +```bash +python -m venv refenv && refenv/bin/pip install "fairchem-core>=2.10" +``` + +`UMA(atoms)` and `tt-atom run` call it automatically the first time they see a new composition, then cache the result. Set `TT_ATOM_REFENV` to its python if it is not found automatically. Cached runs never need it. + +## License + +MIT for this code, which reimplements the UMA / eSCN-MD architecture from [fairchem](https://github.com/facebookresearch/fairchem) (also MIT). It depends on `ttnn` (Apache-2.0) and `ase` (LGPL-2.1+). The UMA weights are separately licensed under the [FAIR Chemistry License](https://huggingface.co/facebook/UMA), are gated, and are not included. Bring your own. diff --git a/RELEASE.md b/RELEASE.md new file mode 100644 index 0000000..aaae2e4 --- /dev/null +++ b/RELEASE.md @@ -0,0 +1,106 @@ +# TT-Atom — release notes & announcement draft + +*Draft for Moritz. Every number below is measured on this machine and reproducible with the +scripts in `benchmarks/`. The public release (GitHub repo, social post) is your call — this +file is the prepared material, nothing has been pushed or posted.* + +--- + +## What it is + +**TT-Atom** is a clean, minimal, high-performance port of Meta's **UMA** (eSEN / eSCN-MD) +equivariant ML interatomic potential to **Tenstorrent** via `ttnn` — energy and **conservative +analytic forces**, fully device-resident, behind an ASE calculator that **mirrors fairchem's**. +Moving off fairchem is a one-line change. Molecules **and** periodic materials, validated against +the released `uma-s-1` checkpoint. + +It does one thing well: fast, accurate inference for this architecture, with a device-resident +trace loop for MD/relaxation and a multi-card throughput path. No framework sprawl, no dead code. + +## Headline numbers (Blackhole p150, real & reproducible) + +- **Drop-in for fairchem:** `FAIRChemCalculator(...)` → `TTAtomCalculator(bundle, task_name=...)`, + same ASE surface. One-command checkpoint converter + `tt-atom verify` device roundtrip. +- **Validated vs released uma-s-1 across three tasks / all graph regimes** (energy rel < 1e-3, + force PCC > 0.99): omol ethanol (1.8e-7 / 0.99958), omat bulk Si (3.0e-4 / 0.99999), oc20 + Cu(100)+H slab (8.6e-5 / 1.00000). The periodic neighbour list reproduces fairchem's + `radius_graph_pbc` edge-for-edge. +- **Trace-captured MD/relaxation loop: 2.33× wall-clock** (FIRE, real uma-s-1) — identical + trajectory, bit-for-bit the same analytic forces. +- **Device compute up to 5.3× faster than 16-thread PyTorch CPU** at 250 atoms, growing with + system size (device latency is nearly flat). +- **3.95× near-linear throughput scaling across 4 cards** (validated on qb1). + +![device vs CPU](assets/device_vs_cpu.png) +![multi-card scaling](assets/multicard_scaling.png) + +## Why it's interesting (engineering story) + +The eSCN SO(2) trick turns the SO(3) tensor product into per-`m` dense GEMMs, so ~85–90 % of +the model is matmul — a great fit for Tenstorrent. The port got its speed from three structural +moves, each measured: + +1. **SO(2) per-`m` conv as flat 2-D GEMMs** — removing a `[E,2,K]` reshape that tile-padded a + length-2 axis to 32 (a 16× data blowup + a per-edge batched matmul). ~12× on that module. +2. **Wigner rotation as a sparse multiply-accumulate** over its fixed nonzero pattern, in a + flat `[E, 9·C]` layout — replacing a launch-bound batched `[E,9,9]×[E,9,C]` matmul (~2.9 µs + *per edge*). ~3.8×, and it composes into a fully tile-aligned pipeline. +3. **Analytic on-device reverse pass** for forces (matmul backward = transpose-matmul on + device; the cheap geometric Jacobian finishes on host). + +Net: the full forward went **251 ms → 33 ms** at ~4800 edges, and the device latency stopped +scaling with the system — which is exactly why the CPU gap widens as systems grow. + +An honest negative result worth keeping: a `bfloat8_b` "fast" mode gives **no speedup** here — +the forward is data-movement bound, not flop bound — so `bf16` (with `HiFi4` + `fp32` +accumulation) is the recommended default. + +## What's validated, and what isn't + +Validated with **real `facebook/UMA` uma-s-1 weights** against the official fairchem reference on +a single p150 (reproducible via `tests/test_realweight.py` + `tests/test_periodic.py`): + +| task | system | graph | energy rel err | force PCC | MAE (eV/Å) | +|---|---|---|---:|---:|---:| +| omol | ethanol | aperiodic | 1.8e-7 | 0.99958 | 3.4e-3 | +| omat | bulk Si | periodic [T,T,T] | 3.0e-4 | 0.99999 | 6.5e-3 | +| oc20 | Cu(100)+H | mixed [T,T,F] | 8.6e-5 | 1.00000 | 9.7e-4 | + +- MoLE experts host-merged to a plain `eSCNMDBackbone` (fairchem's own `merge_mole` path): + merged vs unmerged-MoE oracle **E rel 1.3e-12, force PCC 1.0** — the merge is exact. +- Periodic neighbour list reproduces fairchem's `radius_graph_pbc` edge-for-edge (edges + image + offsets), so materials tasks work. `odac`/`omc` use the identical data-driven path (export with + `--task`). +- Real-weight ASE FIRE relaxation of ethanol **converges on device** (fmax 9.16 → 0.049 eV/Å). +- **Trace path** (`trace=True`): device fwd+bwd is ~96% of a step and dispatch-bound, so capturing + and replaying the op-stream gives **2.14× per step / 2.33× a full FIRE relaxation**, with + bit-for-bit identical forces. + +Also validated with **random weights** against a bit-exact PyTorch reference (the always-available, +ungated CI path): per-module PCC ≥ 0.99, end-to-end energy/forces, module VJPs. 22 tests pass. + +**Model coverage / honest ceiling.** `uma-s-1` (lmax=mmax=2) is the validated default. `uma-m-1p1` +exports cleanly but is **not supported**: it uses lmax=4/mmax=2 spherical-harmonic coefficient +subselection, a code path TT-Atom does not implement (the calculator raises a clear error), and its +oracle+merge parity harness OOMs on a 30 GB host. **Single-card only on pc**; the 4-card 3.95× was +validated on qb1 and is not re-measured here. + +No weights are shipped or redistributed; the `facebook/UMA` checkpoint is gated under the FAIR +Chemistry License and the real-weight tests auto-skip when absent. Conversion is one command +(`tt-atom convert-checkpoint` / `tools/export_weights.py`), and `tt-atom verify` closes the +device roundtrip against a fairchem reference embedded in the bundle at convert time. + +## Suggested social post (draft) + +> Got Meta's UMA interatomic potential running on Tenstorrent as a **drop-in for fairchem** — +> swap one calculator class and your ASE relaxations/MD run on the card. Energy **and analytic +> forces**, device-resident. Matches the **released uma-s-1** across molecules *and* periodic +> materials (omol / omat / oc20): energy to ≤3e-4, force PCC ≥ 0.9996 vs the fairchem oracle. +> The SO(2)-convolution trick makes it ~90 % dense GEMM, so it maps beautifully: a trace-captured +> MD loop runs **2.3× faster**, device compute up to **5.3× over CPU** (gap grows with size), and +> **~4× linear scaling across 4 Blackhole cards**. MIT, bring your own checkpoint. 🧪⚡ + +## Status + +MIT (our code). Tests pass (`pytest tests/ -q`). Benchmarks reproducible. Nothing pushed +or posted — ready for your review. diff --git a/RELEASING.md b/RELEASING.md new file mode 100644 index 0000000..7f69d5c --- /dev/null +++ b/RELEASING.md @@ -0,0 +1,43 @@ +# Releasing TT-Atom + +`master` is the development branch — it may contain experimental, not-yet-validated work. +**A tagged release is a promise to customers that it works.** So a release is cut only from a +commit that has passed the full on-hardware gate below. + +## Release gate — MUST pass on real Tenstorrent hardware before tagging + +GitHub CI only builds and imports the package (no card). Everything that matters is verified +on-device, on the exact commit to be tagged: + +1. **Accuracy / correctness** — full test suite green **and** numerical parity vs the reference + (fairchem `uma-s-1`) within tolerance (energy rel-error and force/stress PCC) across every + supported task and graph regime (molecular, bulk/periodic, slab). **No accuracy regression** + vs the previous release. +2. **No OOM** — run the full supported size range (small molecules → large systems) on the + target card(s), single- and multi-card, to completion. No out-of-memory. Any hard size limit + is documented in the release notes, not discovered by a customer. +3. **No perf regression** — benchmark the release commit against the previous release; latency + and throughput must not regress beyond noise. Record the numbers in the release notes. + +If any of the three fails, it does not ship — fix it or hold the release. + +## Cut a release + +1. Run the gate above on hardware; capture the accuracy table + benchmark numbers. +2. Bump the version in `pyproject.toml` (SemVer) and add a dated section to `CHANGELOG.md` + (include the measured accuracy + perf numbers). +3. Tag and push: + ```bash + git tag v0.1.0 + git push origin master --tags + ``` +4. CI (`.github/workflows/release.yaml`) builds the sdist + wheel, checks the tag matches the + `pyproject` version, and publishes a **GitHub Release** with the changelog notes + wheel. + +## Distribution: GitHub Releases only (NOT PyPI) + +tt-atom is the **custom-kernel-only** build and **requires a source tt-metal/ttnn build** with +the `fused_rotate` op (see the README "Install"). A `pip install tt-atom` wheel therefore can't +run standalone, so publishing to PyPI would be misleading — tt-atom ships via **GitHub Releases** +(source + build instructions + tagged versions). There is intentionally no `pypi-publish` job. +(tt-bio, which *is* pip-installable, does publish to PyPI.) diff --git a/assets/caffeine_md.gif b/assets/caffeine_md.gif new file mode 100644 index 0000000..0ea5d67 Binary files /dev/null and b/assets/caffeine_md.gif differ diff --git a/assets/device_vs_cpu.png b/assets/device_vs_cpu.png new file mode 100644 index 0000000..6915fb3 Binary files /dev/null and b/assets/device_vs_cpu.png differ diff --git a/assets/multicard_scaling.png b/assets/multicard_scaling.png new file mode 100644 index 0000000..80b75af Binary files /dev/null and b/assets/multicard_scaling.png differ diff --git a/benchmarks/bench_batch.py b/benchmarks/bench_batch.py new file mode 100644 index 0000000..9d3b14a --- /dev/null +++ b/benchmarks/bench_batch.py @@ -0,0 +1,123 @@ +"""Disjoint-union batching throughput on ONE card: K small systems batched vs one-at-a-time. + +Batching's win is in the dispatch-bound regime — MANY SMALL systems, where per-call host +overhead (build geometry, upload, launch, read back) dominates device compute. Concatenating K +systems into one block-diagonal graph pays that overhead once instead of K times. + +This is a strict apples-to-apples: both paths call the SAME code (``energy_and_forces_batch``, +energy-only). "one-at-a-time" runs it K times with a 1-system batch; "batched" runs it once with +a K-system batch. So the only variable is the disjoint union. We report systems/s for each, the +batched speedup, the crossover K (where batched first overtakes), and the batch-size ceiling +(largest K before device OOM). + + ~/.ttatom_run/env/bin/python benchmarks/bench_batch.py --weights ~/.ttatom_run/uma_s_ethanol.npz +""" +from __future__ import annotations + +import argparse +import json +import pathlib +import time + +from ase.build import molecule + +from tt_atom import device as D +from tt_atom.model import Backbone +from tt_atom.geometry import HostGeometry +from tt_atom.weights import WeightBundle +from tt_atom import forces, disjoint + +RESULTS = pathlib.Path(__file__).parent / "results" + + +def conformers(k, mol, seed0=10): + out = [] + for i in range(k): + a = molecule(mol) + a.rattle(stdev=0.08, seed=seed0 + i) + a.info.update(charge=0, spin=1) + out.append(a) + return out + + +def time_it(fn, iters): + fn() # warm (program-cache fill for this shape) + t0 = time.perf_counter() + for _ in range(iters): + fn() + return (time.perf_counter() - t0) / iters + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", required=True) + ap.add_argument("--mol", default="CH3CH2OH", help="ASE molecule name (small system)") + ap.add_argument("--ks", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32, 64, 128]) + ap.add_argument("--iters", type=int, default=10) + ap.add_argument("--device-id", type=int, default=0) + args = ap.parse_args() + + b = WeightBundle.load(args.weights) + cfg, w = b.config, b.weights + C = cfg["sphere_channels"] + dev = D.open_device(args.device_id) + bb = Backbone(w, dev, cfg, b.to_grid_mat, b.from_grid_mat) + geo = HostGeometry(w, cfg, b.to_m, b.gauss_offset, b.gauss_coeff, gamma=0.0) + + natoms = len(molecule(args.mol)) + print(f"molecule {args.mol}: {natoms} atoms/system") + + rows = [] + ceiling = None + for k in args.ks: + systems = conformers(k, args.mol) + try: + # one-at-a-time: same code, K single-system batches + def seq(): + for a in systems: + bg1 = disjoint.assemble([a], cfg["cutoff"], w, C, task=b.task) + forces.energy_and_forces_batch(bb, geo, bg1, compute_forces=False) + + # batched: one K-system disjoint-union forward + bgK = disjoint.assemble(systems, cfg["cutoff"], w, C, task=b.task) + + def bat(): + forces.energy_and_forces_batch(bb, geo, bgK, compute_forces=False) + + seq_s = time_it(seq, max(2, args.iters // 2)) + bat_s = time_it(bat, args.iters) + except RuntimeError as e: # device OOM etc. -> record the ceiling and stop + print(f"K={k}: FAILED ({str(e).splitlines()[0][:80]}) -> batch-size ceiling below {k}") + ceiling = k + break + + Etot = bgK.edge_index.shape[1] + seq_thru = k / seq_s + bat_thru = k / bat_s + rows.append(dict(K=k, natoms_total=int(bgK.pos.shape[0]), nedges_total=Etot, + seq_ms=seq_s * 1e3, batched_ms=bat_s * 1e3, + seq_sys_per_s=seq_thru, batched_sys_per_s=bat_thru, + speedup=seq_s / bat_s)) + print(f"K={k:4d} Ntot={bgK.pos.shape[0]:5d} Etot={Etot:6d} " + f"seq={seq_s*1e3:8.2f}ms ({seq_thru:7.1f} sys/s) " + f"batched={bat_s*1e3:8.2f}ms ({bat_thru:7.1f} sys/s) x{seq_s/bat_s:.2f}") + + import ttnn + ttnn.close_device(dev) + + speedups = [r for r in rows if r["speedup"] > 1.0] + crossover = speedups[0]["K"] if speedups else None + best = max(rows, key=lambda r: r["speedup"]) if rows else None + summary = dict(molecule=args.mol, natoms_per_system=natoms, crossover_K=crossover, + batch_ceiling=ceiling, + best_speedup=best["speedup"] if best else None, + best_K=best["K"] if best else None) + print("\nSUMMARY:", json.dumps(summary)) + RESULTS.mkdir(exist_ok=True) + out = RESULTS / "batch_throughput.json" + out.write_text(json.dumps(dict(config=cfg, summary=summary, rows=rows), indent=2)) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_multicard.py b/benchmarks/bench_multicard.py new file mode 100644 index 0000000..c25e346 --- /dev/null +++ b/benchmarks/bench_multicard.py @@ -0,0 +1,75 @@ +"""Multi-card throughput scaling for TT-Atom. + +Evaluates a fixed pool of independent systems on 1..N cards (one worker process per card, +weights resident) and reports aggregate Medges/s and the scaling factor at each card count. +All numbers are measured here; nothing is hardcoded. Results -> benchmarks/results/multicard.json. + +Run: ~/.ttatom_run/env/bin/python benchmarks/bench_multicard.py --weights /tmp/tt_full.npz +""" +from __future__ import annotations + +import argparse +import json +import pathlib +import time + +import numpy as np +from ase.build import bulk + +from tt_atom.batch import MultiCard + +RESULTS = pathlib.Path(__file__).parent / "results" + + +def make_systems(n_systems, cells): + # Same-size systems (constant edge count) so the device program cache stays warm: this + # measures steady-state throughput capacity. Production screening of differently-sized + # systems would bucket/pad edges to a few fixed sizes for the same effect. + systems = [] + for i in range(n_systems): + a = bulk("Si", "diamond", a=5.43) * (cells, cells, cells) + a.rattle(stdev=0.1, seed=1) # identical geometry -> constant E + systems.append((a.get_positions().astype(np.float32), a.get_atomic_numbers())) + return systems + + +def measure(weights, device_ids, systems, fast): + with MultiCard(weights, device_ids=device_ids, fast=fast) as mc: + mc.energies(systems[: len(device_ids)]) # warmup (compile per worker) + t0 = time.perf_counter() + _, total_edges = mc.energies(systems) + dt = time.perf_counter() - t0 + return total_edges / dt / 1e6, dt + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", required=True) + ap.add_argument("--cells", type=int, default=4) # ~128 atoms / ~2200 edges each + ap.add_argument("--systems", type=int, default=64) + ap.add_argument("--fast", action="store_true") + ap.add_argument("--max-cards", type=int, default=4) + args = ap.parse_args() + + systems = make_systems(args.systems, args.cells) + natoms, nedges = len(systems[0][1]), None + rows = [] + for n in range(1, args.max_cards + 1): + ids = list(range(n)) + medges, dt = measure(args.weights, ids, systems, args.fast) + rows.append(dict(cards=len(ids), device_ids=ids, medges_per_s=medges, wall_s=dt)) + print(f"{len(ids)} card(s): {medges:6.3f} Medges/s ({args.systems} systems x ~{natoms} atoms in {dt:.2f}s)") + + base = rows[0]["medges_per_s"] + for r in rows: + r["scaling_vs_1card"] = r["medges_per_s"] / base + print(f"{args.max_cards}-card scaling: {rows[-1]['scaling_vs_1card']:.2f}x") + RESULTS.mkdir(exist_ok=True) + out = RESULTS / ("multicard_fast.json" if args.fast else "multicard.json") + out.write_text(json.dumps(dict(systems=args.systems, cells=args.cells, natoms=natoms, + fast=args.fast, rows=rows), indent=2)) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_throughput.py b/benchmarks/bench_throughput.py new file mode 100644 index 0000000..5c35113 --- /dev/null +++ b/benchmarks/bench_throughput.py @@ -0,0 +1,141 @@ +"""Single-card throughput + end-to-end CPU-vs-TT benchmark for TT-Atom. + +Measures, over a system-size sweep, on real hardware: + * warm device-resident forward (energy) ms/eval and Medges/s -- pure device compute + * end-to-end ms/eval (host geometry + upload + forward + readback) -- what a user pays + * a PyTorch-CPU reference ms/eval (tests/mirror.py, a bit-exact transcription of the same + function) for an honest CPU-vs-TT speedup + +All numbers are real and measured here; nothing is hardcoded. Results -> benchmarks/results/. +Run: ~/.ttatom_run/env/bin/python benchmarks/bench_throughput.py --weights /tmp/tt_full.npz +""" +from __future__ import annotations + +import argparse +import json +import pathlib +import sys +import time + +import torch +from ase.build import bulk + +sys.path.insert(0, str(pathlib.Path(__file__).parent.parent / "tests")) +import mirror # noqa: E402 + +from tt_atom import device as D # noqa: E402 +from tt_atom.model import Backbone, GraphContext # noqa: E402 +from tt_atom.geometry import HostGeometry, csd_embedding, radius_graph # noqa: E402 +from tt_atom.weights import WeightBundle # noqa: E402 + +RESULTS = pathlib.Path(__file__).parent / "results" + + +def make_system(n_cells): + a = bulk("Si", "diamond", a=5.43) * (n_cells, n_cells, n_cells) + a.rattle(stdev=0.1, seed=1) + return a + + +def time_it(fn, iters): + fn() # warmup (compile / program-cache fill) + t0 = time.perf_counter() + for _ in range(iters): + fn() + return (time.perf_counter() - t0) / iters + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", required=True) + ap.add_argument("--cells", type=int, nargs="+", default=[1, 2, 3]) + ap.add_argument("--iters", type=int, default=20) + ap.add_argument("--device-id", type=int, default=0) + ap.add_argument("--fast", action="store_true", help="bf8 weights + rotation coefficients") + args = ap.parse_args() + + b = WeightBundle.load(args.weights) + cfg, w = b.config, b.weights + C = cfg["sphere_channels"] + dev = D.open_device(args.device_id) + bb = Backbone(w, dev, cfg, b.to_grid_mat, b.from_grid_mat, fast=args.fast) + geo = HostGeometry(w, cfg, b.to_m, b.gauss_offset, b.gauss_coeff, gamma=0.0) + import ttnn + + rows = [] + for nc in args.cells: + atoms = make_system(nc) + pos = torch.tensor(atoms.get_positions(), dtype=torch.float32) + Z = torch.tensor(atoms.get_atomic_numbers()) + ei, _ = radius_graph(pos, cfg["cutoff"]) + N, E = Z.shape[0], ei.shape[1] + se = csd_embedding(w, torch.tensor([0.0]), torch.tensor([0.0]), C)[torch.zeros(N, dtype=torch.long)] + + def host_geom(): + return geo(pos, Z, ei, se) + + t = host_geom() + + def upload(t): + graph = GraphContext(dev, edge_index=ei, wigner=t["wigner"].detach(), + wigner_inv=t["wigner_inv"].detach(), x_edge=t["x_edge"].detach(), + edge_envelope=t["edge_envelope"].detach(), num_nodes=N, fast=args.fast) + se3 = ttnn.from_torch(se.reshape(N, 1, C), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=dev) + xi = ttnn.from_torch(t["x_init"].detach(), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=dev) + return graph, se3, xi + + graph, se3, xi = upload(t) + + # warm device-resident forward (energy) + def dev_fwd(): + bb(xi, graph, se3) + ttnn.synchronize_device(dev) + + dev_ms = time_it(dev_fwd, args.iters) * 1e3 + + # end-to-end (host geom + upload + forward + readback) + def e2e(): + tt = host_geom() + g, s, x = upload(tt) + _, en = bb(x, g, s) + float(ttnn.to_torch(en).reshape(-1)[0]) + + e2e_ms = time_it(e2e, max(3, args.iters // 4)) * 1e3 + + # CPU reference (bit-exact mirror) + def cpu_fwd(): + ne = mirror.backbone(w, cfg, t["x_init"], t["wigner"], t["wigner_inv"], + t["x_edge"], t["edge_envelope"], se, ei, b.to_grid_mat, b.from_grid_mat) + float(mirror.energy(ne, w)) + + with torch.no_grad(): + cpu_ms = time_it(cpu_fwd, max(3, args.iters // 4)) * 1e3 + + # honest accuracy: TT energy vs the CPU reference (same fp32 mirror) + _, en_tt = bb(xi, graph, se3) + e_tt = float(ttnn.to_torch(en_tt).reshape(-1)[0]) + with torch.no_grad(): + e_cpu = float(mirror.energy(mirror.backbone( + w, cfg, t["x_init"], t["wigner"], t["wigner_inv"], t["x_edge"], + t["edge_envelope"], se, ei, b.to_grid_mat, b.from_grid_mat), w)) + rel_err = abs(e_tt - e_cpu) / (abs(e_cpu) + 1e-6) + + medges = E / (dev_ms * 1e-3) / 1e6 + rows.append(dict(cells=nc, natoms=N, nedges=E, dev_ms=dev_ms, e2e_ms=e2e_ms, + cpu_ms=cpu_ms, medges_per_s=medges, speedup_dev=cpu_ms / dev_ms, + speedup_e2e=cpu_ms / e2e_ms, energy_rel_err=rel_err)) + print(f"N={N:4d} E={E:5d} dev={dev_ms:7.2f}ms e2e={e2e_ms:7.2f}ms cpu={cpu_ms:7.2f}ms " + f"{medges:5.3f} Medges/s dev x{cpu_ms/dev_ms:.2f} e2e x{cpu_ms/e2e_ms:.2f} " + f"Erel={rel_err:.1e}") + + ttnn.close_device(dev) + RESULTS.mkdir(exist_ok=True) + out = RESULTS / ("throughput_fast.json" if args.fast else "throughput.json") + out.write_text(json.dumps(dict(config=cfg, fast=args.fast, rows=rows), indent=2)) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_trace.py b/benchmarks/bench_trace.py new file mode 100644 index 0000000..ce407e0 --- /dev/null +++ b/benchmarks/bench_trace.py @@ -0,0 +1,71 @@ +"""Trace vs eager benchmark for the MD / relaxation loop (single card). + +Measures the real end-to-end per-step cost (host geometry + device forward + device backward + +host force finish) of the eager path vs the trace-captured device-resident path, plus a phase +breakdown that shows why tracing is the right lever (the device forward+backward dominates and is +host-dispatch-bound for these graph sizes). Forces are bit-for-bit identical between the two. + + PYTHONPATH=~/TT-Atom ~/.ttatom_run/env/bin/python benchmarks/bench_trace.py \ + --weights ~/.ttatom_run/uma_s_ethanol.npz + +With no --weights it falls back to the committed random-weight demo bundle (architecture-only). +""" +from __future__ import annotations + +import argparse +import pathlib +import time + +import numpy as np +from ase.build import molecule + +HERE = pathlib.Path(__file__).parent + + +def _median_ms(fn, n=15, warm=3): + for _ in range(warm): + fn() + ts = [] + for _ in range(n): + t = time.perf_counter() + fn() + ts.append((time.perf_counter() - t) * 1000) + return float(np.median(ts)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", default=str(HERE.parent / "examples" / "model_tiny_demo.npz")) + ap.add_argument("--device-id", type=int, default=0) + args = ap.parse_args() + + from tt_atom.calculator import TTAtomCalculator + + atoms = molecule("CH3CH2OH") + atoms.info.update(charge=0, spin=0) + atoms.rattle(stdev=0.03, seed=5) + + eager = TTAtomCalculator(args.weights, device_id=args.device_id) + atoms.calc = eager + eager_ms = _median_ms(lambda: eager.calculate(atoms)) + Ee, Fe = eager.results["energy"], eager.results["forces"] + eager.close() + + traced = TTAtomCalculator(args.weights, device_id=args.device_id, trace=True) + atoms.calc = traced + traced.calculate(atoms) # capture + traced_ms = _median_ms(lambda: traced.calculate(atoms)) + Et, Ft = traced.results["energy"], traced.results["forces"] + traced.close() + + print(f"system: ethanol (9 atoms), weights={pathlib.Path(args.weights).name}") + print(f" eager E+F per step : {eager_ms:6.2f} ms") + print(f" traced E+F per step : {traced_ms:6.2f} ms") + print(f" speedup : {eager_ms / traced_ms:.2f}x") + print(f" energy diff : {abs(Et - Ee):.2e} eV") + print(f" force PCC / maxdiff : {np.corrcoef(Ft.ravel(), Fe.ravel())[0, 1]:.6f} / " + f"{np.abs(Ft - Fe).max():.2e} eV/A (trace only removes host dispatch)") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/plot.py b/benchmarks/plot.py new file mode 100644 index 0000000..3218732 --- /dev/null +++ b/benchmarks/plot.py @@ -0,0 +1,79 @@ +"""Generate the release perf charts from measured benchmark JSON (no hardcoded numbers). + +Run after the benchmarks: + ~/.ttatom_run/env/bin/python benchmarks/plot.py +Reads benchmarks/results/{throughput,multicard}.json -> assets/*.png +""" +from __future__ import annotations + +import json +import pathlib + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +ROOT = pathlib.Path(__file__).parent.parent +RESULTS = ROOT / "benchmarks" / "results" +ASSETS = ROOT / "assets" +TT = "#7C3AED" +CPU = "#9CA3AF" +GRN = "#10B981" + + +def plot_device_vs_cpu(): + d = json.loads((RESULTS / "throughput.json").read_text()) + rows = d["rows"] + N = [r["natoms"] for r in rows] + dev = [r["dev_ms"] for r in rows] + cpu = [r["cpu_ms"] for r in rows] + spd = [r["speedup_dev"] for r in rows] + + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.2)) + ax1.plot(N, cpu, "o-", color=CPU, label="PyTorch CPU", lw=2) + ax1.plot(N, dev, "o-", color=TT, label="TT-Atom (1x Blackhole p150)", lw=2) + ax1.set_xlabel("atoms"); ax1.set_ylabel("ms / energy eval (warm)") + ax1.set_title("Device compute vs CPU"); ax1.legend(); ax1.grid(alpha=0.3) + + ax2.plot(N, spd, "o-", color=GRN, lw=2) + ax2.axhline(1.0, color="k", ls="--", lw=1, alpha=0.5) + ax2.set_xlabel("atoms"); ax2.set_ylabel("device speedup vs CPU (x)") + ax2.set_title("Speedup grows with system size") + for x, y in zip(N, spd): + ax2.annotate(f"{y:.1f}x", (x, y), textcoords="offset points", xytext=(0, 8), ha="center") + ax2.grid(alpha=0.3) + fig.suptitle("TT-Atom — eSEN / eSCN-MD inference on Tenstorrent (full config, random weights)", + fontweight="bold") + fig.tight_layout() + ASSETS.mkdir(exist_ok=True) + fig.savefig(ASSETS / "device_vs_cpu.png", dpi=130) + print("wrote", ASSETS / "device_vs_cpu.png") + + +def plot_multicard(): + p = RESULTS / "multicard.json" + if not p.exists(): + return + d = json.loads(p.read_text()) + rows = d["rows"] + cards = [r["cards"] for r in rows] + medges = [r["medges_per_s"] for r in rows] + ideal = [medges[0] * c for c in cards] + + fig, ax = plt.subplots(figsize=(5.2, 4.2)) + ax.plot(cards, ideal, "--", color=CPU, label="ideal linear", lw=1.5) + ax.plot(cards, medges, "o-", color=TT, label="measured", lw=2) + for c, m in zip(cards, medges): + ax.annotate(f"{m:.3f}", (c, m), textcoords="offset points", xytext=(6, -4)) + ax.set_xlabel("cards"); ax.set_ylabel("aggregate Medges/s") + ax.set_title(f"Multi-card throughput ({d['natoms']}-atom systems)\n" + f"{rows[-1]['scaling_vs_1card']:.2f}x on {cards[-1]} cards") + ax.set_xticks(cards); ax.legend(); ax.grid(alpha=0.3) + fig.tight_layout() + fig.savefig(ASSETS / "multicard_scaling.png", dpi=130) + print("wrote", ASSETS / "multicard_scaling.png") + + +if __name__ == "__main__": + plot_device_vs_cpu() + plot_multicard() diff --git a/custom_kernels/README.md b/custom_kernels/README.md new file mode 100644 index 0000000..e6474fd --- /dev/null +++ b/custom_kernels/README.md @@ -0,0 +1,76 @@ +# Custom tt-metal kernels for TT-Atom + +`tt_atom/rotation.py` routes the per-edge Wigner rotation through custom tt-metal compute kernels +that the pip `ttnn` wheel does not carry, so TT-Atom needs a **source tt-metal build** that +includes this op. The op is pre-integrated on the +[`moritztng/tt-atom`](https://github.com/tenstorrent/tt-metal/tree/moritztng/tt-atom) +branch of tt-metal, so the normal install just clones and builds that branch (see the top-level +README). This directory is the authoritative backup of the op source and the recipe for +re-integrating it onto a newer tt-metal commit. + +## What the op provides +One tt-metal experimental op library (`TTNN::Ops::Experimental::FusedRotate`) exposing four kernels +under `ttnn._ttnn.operations.experimental`: + +- **`fused_rotate`** — the per-edge sparse Wigner rotation. Replaces the ~35 `ttnn.addcmul` + dispatches of `rotation.rotate` with ONE compute-kernel launch: reads x once, keeps all `nnz` + multiply-accumulates in the dest registers (fp32 accumulate, HiFi4), writes out once. **Measured + 4.3x faster** than the split+addcmul MAC it replaces on the uma-s forward-rotate shape (E=46016, + n=9, W=256, nnz=35): 7.01 ms -> 1.62 ms, PCC=0.999995. This is the ALWAYS-ON path for uma-s. +- **`fused_rotate_gc`** — the coefficient-adjoint (`dE/dcoef`) mul-reduce used in the backward at + large graphs (products L1-resident, one accumulating matmul reduces + places into `gc[E,nnz]`). +- `fused_gate`, `fused_ln_bw` — the SO(3) gate activation and the LayerNorm-backward reduction + (optional; used by other modules). + +### `fused_rotate` contract +`ttnn._ttnn.operations.experimental.fused_rotate(x_flat, coef_exp, n_in, n_out, W, deg, ks, js)` +- `x_flat` `[E, n_in*W]` bf16/bf8_b TILE +- `coef_exp` `[E, nnz*32]` — each nonzero k's coefficient broadcast across a 32-col tile + (`ttnn.repeat_interleave(coef[E,nnz], 32, dim=1)`) +- `deg[i]` nonzeros feeding output block i; `ks`/`js` (len nnz) the coef-tile index and input + block for each nonzero, grouped by output block i in order. +- returns `[E, n_out*W]`. + +The kernel fans-in all `d` per-block products into `dst[0..d-1]` and sums them there, so `max(deg)` +must fit the fp32 DST register file (`dst_full_sync_en` -> 8 slots) and the per-core CBs must fit +L1 (~1.5 MB). uma-s (square 9x9, W=128/256) fits; uma-m (rectangular 19x25, W=256) overflows and is +unsupported — `rotation.rotate` raises rather than falling back. + +## Re-integrating onto a newer tt-metal commit + +The `moritztng/tt-atom` branch already carries this op on top of validated commit +**`b5522097b39`** (`Migrate experimental/ssm leftovers to ProgramDescriptor`, #44403), so you only +need the steps below to rebase the op onto a *different* tt-metal commit. They are exactly how that +branch was produced. + +1. Copy this op into the tt-metal tree: + ``` + cp -r custom_kernels/fused_rotate \ + $TT_METAL_HOME/ttnn/cpp/ttnn/operations/experimental/fused_rotate + ``` +2. Apply the 3 registration edits (the op's own `CMakeLists.txt`/`sources.cmake` build all four + kernels into one library; these hook that library into ttnn): + - `ttnn/CMakeLists.txt`: add `add_subdirectory(cpp/ttnn/operations/experimental/fused_rotate)` + and add `TTNN::Ops::Experimental::FusedRotate` to the `target_link_libraries(ttnn ...)` list. + - `ttnn/cpp/ttnn/operations/experimental/experimental_nanobind.cpp`: add + `#include "ttnn/operations/experimental/fused_rotate/fused_rotate_nanobind.hpp"` and, inside + `py_module`, `fr_detail::bind_fused_rotate(mod);` (binds all four kernels). + - `ttnn/sources.cmake`: add + `cpp/ttnn/operations/experimental/fused_rotate/fused_rotate_nanobind.cpp` to `TTNN_SRC_PYBIND` + (omitting this links but leaves the Python symbols undefined — the classic trap). +3. Rebuild and re-stage the shared object (the `install` target both compiles and copies + `_ttnn.so` into `ttnn/ttnn/` — this is what `build_metal.sh` itself runs): + ``` + cmake --build $TT_METAL_HOME/build_Release --target install + ``` + `pip install -e .` (top-level README) only needs to run once per venv — it's a packaging step, + not a build step, so re-running it after this rebuild is unnecessary. +4. Verify: + ``` + TT_METAL_HOME=$TT_METAL_HOME python3 -c \ + "import ttnn; e=ttnn._ttnn.operations.experimental; \ + print(hasattr(e,'fused_rotate'), hasattr(e,'fused_rotate_gc'))" # -> True True + ``` + +Run TT-Atom (both `ttnn` and `tt-atom` are `pip install -e`'d into the active venv): +`TT_METAL_HOME=$TT_METAL_HOME python ...` diff --git a/custom_kernels/fused_rotate/CMakeLists.txt b/custom_kernels/fused_rotate/CMakeLists.txt new file mode 100644 index 0000000..dc72566 --- /dev/null +++ b/custom_kernels/fused_rotate/CMakeLists.txt @@ -0,0 +1,44 @@ +include(sources.cmake) + +add_library(ttnn_op_experimental_fused_rotate ${LIB_TYPE}) +add_library(TTNN::Ops::Experimental::FusedRotate ALIAS ttnn_op_experimental_fused_rotate) + +tt_reuse_precompile_headers(ttnn_op_experimental_fused_rotate TTNN::PCH) +TT_ENABLE_UNITY_BUILD(ttnn_op_experimental_fused_rotate) + +set_target_properties( + ttnn_op_experimental_fused_rotate + PROPERTIES + INTERFACE_HEADER_SETS_TO_VERIFY + api +) + +file(GLOB_RECURSE kernels device/kernels/*) + +target_sources( + ttnn_op_experimental_fused_rotate + PUBLIC + FILE_SET api + TYPE HEADERS + BASE_DIRS ${FixmeOpAPIDir} + FILES ${TTNN_OP_EXPERIMENTAL_FUSED_ROTATE_API_HEADERS} + FILE_SET kernels + TYPE HEADERS + BASE_DIRS ${CMAKE_CURRENT_SOURCE_DIR} + FILES ${kernels} + PRIVATE + ${TTNN_OP_EXPERIMENTAL_FUSED_ROTATE_SRCS} +) + +target_link_libraries(ttnn_op_experimental_fused_rotate PRIVATE TT::Metalium PUBLIC TTNN::Core) + +install( + TARGETS + ttnn_op_experimental_fused_rotate + FILE_SET + kernels + DESTINATION ${CMAKE_INSTALL_LIBEXECDIR}/tt-metalium/ttnn/cpp/ttnn/operations/experimental/fused_rotate + COMPONENT ttnn-runtime +) + +install(TARGETS ttnn_op_experimental_fused_rotate FILE_SET api COMPONENT ttnn-dev LIBRARY COMPONENT tar) diff --git a/custom_kernels/fused_rotate/device/fused_rotate_device_operation.cpp b/custom_kernels/fused_rotate/device/fused_rotate_device_operation.cpp new file mode 100644 index 0000000..822cc47 --- /dev/null +++ b/custom_kernels/fused_rotate/device/fused_rotate_device_operation.cpp @@ -0,0 +1,118 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "fused_rotate_device_operation.hpp" +#include "fused_rotate_program_factory.hpp" +#include "ttnn/device_operation.hpp" + +#include + +using namespace tt::constants; +using namespace tt::tt_metal; + +namespace ttnn::experimental::prim { + +void FusedRotateDeviceOperation::validate_on_program_cache_miss( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + const auto& x = inputs.x_flat; + const auto& coef = inputs.coef_exp; + TT_FATAL( + x.storage_type() == StorageType::DEVICE && coef.storage_type() == StorageType::DEVICE, + "fused_rotate operands must be on device"); + TT_FATAL(x.layout() == Layout::TILE && coef.layout() == Layout::TILE, "fused_rotate requires TILE layout"); + // bf16 or bf8_b; x and coef must share a format (the program factory derives one tile size + // from x.dtype() for both CBs). bf8_b coef is parity-safe (orthogonal basis change, O(1) coefs) + // and halves the [E,W] edge-activation DRAM traffic that dominates the bandwidth-bound replay. + TT_FATAL( + (x.dtype() == DataType::BFLOAT16 || x.dtype() == DataType::BFLOAT8_B) && coef.dtype() == x.dtype(), + "fused_rotate requires bf16 or bf8_b inputs (x and coef same dtype)"); + + const auto& xs = x.padded_shape(); + const auto& cs = coef.padded_shape(); + TT_FATAL(attrs.W % TILE_WIDTH == 0, "W ({}) must be a multiple of TILE_WIDTH", attrs.W); + TT_FATAL(xs[-1] == attrs.n_in * attrs.W, "x_flat last dim {} != n_in*W {}", xs[-1], attrs.n_in * attrs.W); + TT_FATAL(cs[-1] == attrs.nnz * TILE_WIDTH, "coef_exp last dim {} != nnz*32 {}", cs[-1], attrs.nnz * TILE_WIDTH); + TT_FATAL(xs[-2] == cs[-2], "x_flat and coef_exp must have the same number of rows (edges)"); + TT_FATAL(attrs.deg.size() == attrs.n_out, "deg size {} != n_out {}", attrs.deg.size(), attrs.n_out); + TT_FATAL(attrs.ks.size() == attrs.nnz && attrs.js.size() == attrs.nnz, "ks/js size must equal nnz"); + uint32_t sum = 0; + for (auto d : attrs.deg) { + sum += d; + } + TT_FATAL(sum == attrs.nnz, "sum(deg)={} != nnz={}", sum, attrs.nnz); +} + +FusedRotateDeviceOperation::spec_return_value_t FusedRotateDeviceOperation::compute_output_specs( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + const auto& x = inputs.x_flat; + ttnn::Shape out_shape(x.logical_shape()); + out_shape[-1] = attrs.n_out * attrs.W; + return TensorSpec(out_shape, TensorLayout(x.dtype(), PageConfig(Layout::TILE), x.memory_config())); +} + +FusedRotateDeviceOperation::tensor_return_value_t FusedRotateDeviceOperation::create_output_tensors( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + return create_device_tensor(compute_output_specs(attrs, inputs), inputs.x_flat.device()); +} + +ttsl::hash::hash_t FusedRotateDeviceOperation::compute_program_hash( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + // The sparsity pattern (deg/ks/js) is set as RUNTIME args in create() but NOT refreshed in + // override_runtime_arguments, so two calls that share shapes but differ in pattern (e.g. the + // forward rotation grouped by output i vs. the backward g_in grouped by input j -- identical + // [n_in,n_out,W,nnz] for a square rotation) MUST NOT share a cached program. Fold the pattern + // into the hash so each distinct pattern gets its own program. + uint64_t ph = 1469598103934665603ULL; // FNV-1a + auto mix = [&](uint32_t v) { + ph = (ph ^ v) * 1099511628211ULL; + }; + for (auto v : attrs.deg) { + mix(v); + } + for (auto v : attrs.ks) { + mix(v); + } + for (auto v : attrs.js) { + mix(v); + } + return tt::tt_metal::operation::hash_operation( + attrs.n_in, + attrs.n_out, + attrs.W, + attrs.nnz, + static_cast(ph), + static_cast(ph >> 32), + inputs.x_flat.dtype(), + inputs.x_flat.memory_config(), + inputs.x_flat.padded_shape(), + inputs.coef_exp.padded_shape()); +} + +} // namespace ttnn::experimental::prim + +namespace ttnn::prim { + +Tensor fused_rotate( + const Tensor& x_flat, + const Tensor& coef_exp, + uint32_t n_in, + uint32_t n_out, + uint32_t W, + const std::vector& deg, + const std::vector& ks, + const std::vector& js) { + using OperationType = ttnn::experimental::prim::FusedRotateDeviceOperation; + auto attrs = OperationType::operation_attributes_t{ + .n_in = n_in, + .n_out = n_out, + .W = W, + .nnz = static_cast(ks.size()), + .deg = deg, + .ks = ks, + .js = js}; + return ttnn::device_operation::launch( + attrs, OperationType::tensor_args_t{.x_flat = x_flat, .coef_exp = coef_exp}); +} + +} // namespace ttnn::prim diff --git a/custom_kernels/fused_rotate/device/fused_rotate_device_operation.hpp b/custom_kernels/fused_rotate/device/fused_rotate_device_operation.hpp new file mode 100644 index 0000000..789cc3d --- /dev/null +++ b/custom_kernels/fused_rotate/device/fused_rotate_device_operation.hpp @@ -0,0 +1,38 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "ttnn/tensor/tensor.hpp" +#include "fused_rotate_program_factory.hpp" +#include "fused_rotate_device_operation_types.hpp" + +namespace ttnn::experimental::prim { + +struct FusedRotateDeviceOperation { + using operation_attributes_t = FusedRotateParams; + using tensor_args_t = FusedRotateInputs; + using spec_return_value_t = TensorSpec; + using tensor_return_value_t = Tensor; + using program_factory_t = std::variant; + + static void validate_on_program_cache_miss(const operation_attributes_t&, const tensor_args_t&); + static spec_return_value_t compute_output_specs(const operation_attributes_t&, const tensor_args_t&); + static tensor_return_value_t create_output_tensors(const operation_attributes_t&, const tensor_args_t&); + static ttsl::hash::hash_t compute_program_hash(const operation_attributes_t&, const tensor_args_t&); +}; + +} // namespace ttnn::experimental::prim + +namespace ttnn::prim { +Tensor fused_rotate( + const Tensor& x_flat, + const Tensor& coef_exp, + uint32_t n_in, + uint32_t n_out, + uint32_t W, + const std::vector& deg, + const std::vector& ks, + const std::vector& js); +} // namespace ttnn::prim diff --git a/custom_kernels/fused_rotate/device/fused_rotate_device_operation_types.hpp b/custom_kernels/fused_rotate/device/fused_rotate_device_operation_types.hpp new file mode 100644 index 0000000..67e428f --- /dev/null +++ b/custom_kernels/fused_rotate/device/fused_rotate_device_operation_types.hpp @@ -0,0 +1,39 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include + +#include "ttnn/tensor/tensor.hpp" + +namespace ttnn::experimental::prim { + +// Fused per-edge sparse Wigner rotation: +// out[e, i*W + w] = sum_{(i,j,k)} coef_exp[e, k*32 + w'] * x[e, j*W + w] +// where the nonzero pattern (i,j) is fixed for the topology (block-diagonal in l). +// coef_exp is the packed [E, nnz] coefficients broadcast to [E, nnz*32] (each nonzero +// occupies one tile, its value replicated across the 32 tile columns). The rotation is a +// dense fan-in per output block; this op fuses all `nnz` multiply-accumulates into ONE +// kernel launch (1 DRAM read of x + 1 write of out) instead of `nnz` ttnn.addcmul dispatches. +struct FusedRotateParams { + uint32_t n_in; // number of input coordinate blocks + uint32_t n_out; // number of output coordinate blocks + uint32_t W; // channels per coordinate (must be a multiple of TILE_WIDTH) + uint32_t nnz; // number of structural nonzeros + // Grouped-by-output-block sparsity pattern. deg[i] = nonzeros feeding output block i. + // For each output block i (in order), the next deg[i] entries of (ks, js) give the + // coefficient index k (into coef tiles) and the input block j. + std::vector deg; // length n_out + std::vector ks; // length nnz + std::vector js; // length nnz +}; + +struct FusedRotateInputs { + Tensor x_flat; // [E, n_in*W] TILE bf16 + Tensor coef_exp; // [E, nnz*32] TILE bf16 (each nonzero broadcast across 32 cols) +}; + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/fused_rotate_program_factory.cpp b/custom_kernels/fused_rotate/device/fused_rotate_program_factory.cpp new file mode 100644 index 0000000..4782174 --- /dev/null +++ b/custom_kernels/fused_rotate/device/fused_rotate_program_factory.cpp @@ -0,0 +1,134 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "fused_rotate_program_factory.hpp" +#include "fused_rotate_device_operation_types.hpp" + +#include +#include +#include +#include + +using namespace tt::constants; +using namespace tt::tt_metal; + +namespace ttnn::experimental::prim { + +static const char* kKernelDir = "ttnn/cpp/ttnn/operations/experimental/fused_rotate/device/kernels/"; + +FusedRotateProgramFactory::cached_program_t FusedRotateProgramFactory::create( + const FusedRotateParams& attrs, const FusedRotateInputs& inputs, Tensor& output) { + Program program{}; + + const auto& x = inputs.x_flat; + const auto& coef = inputs.coef_exp; + + const uint32_t Wt = attrs.W / TILE_WIDTH; + const uint32_t n_in_tiles = attrs.n_in * Wt; + const uint32_t n_out_tiles = attrs.n_out * Wt; + const uint32_t coef_tiles = attrs.nnz; + const uint32_t Et = x.padded_shape()[-2] / TILE_HEIGHT; // number of tile-rows (edges/32) + + tt::DataFormat data_format = datatype_to_dataformat_converter(x.dtype()); + const uint32_t tile_bytes = tile_size(data_format); + + auto* device = x.device(); + CoreCoord grid = device->compute_with_storage_grid_size(); + auto [num_cores, all_cores, core_group_1, core_group_2, rows_per_core_1, rows_per_core_2] = + tt::tt_metal::split_work_to_cores(grid, Et); + + // Circular buffers. Keep a whole tile-row resident (compute indexes randomly into the + // n_in input blocks), double-buffered across rows to overlap DRAM with compute. + auto make_cb = [&](uint32_t cb_index, uint32_t num_tiles) { + CircularBufferConfig cfg = + CircularBufferConfig(num_tiles * tile_bytes, {{cb_index, data_format}}).set_page_size(cb_index, tile_bytes); + CreateCircularBuffer(program, all_cores, cfg); + }; + constexpr uint32_t cb_x = tt::CBIndex::c_0; + constexpr uint32_t cb_coef = tt::CBIndex::c_1; + constexpr uint32_t cb_out = tt::CBIndex::c_16; + make_cb(cb_x, 2 * n_in_tiles); + make_cb(cb_coef, 2 * coef_tiles); + make_cb(cb_out, 2 * n_out_tiles); + + // ---- reader ---- + std::vector reader_ct = {cb_x, cb_coef, n_in_tiles, coef_tiles, tile_bytes}; + TensorAccessorArgs(*x.buffer()).append_to(reader_ct); + TensorAccessorArgs(*coef.buffer()).append_to(reader_ct); + KernelHandle reader_id = CreateKernel( + program, + std::string(kKernelDir) + "reader.cpp", + all_cores, + ReaderDataMovementConfig(reader_ct)); + + // ---- writer ---- + std::vector writer_ct = {cb_out, n_out_tiles, tile_bytes}; + TensorAccessorArgs(*output.buffer()).append_to(writer_ct); + KernelHandle writer_id = CreateKernel( + program, + std::string(kKernelDir) + "writer.cpp", + all_cores, + WriterDataMovementConfig(writer_ct)); + + // ---- compute ---- + std::vector compute_ct = {cb_x, cb_coef, cb_out, n_in_tiles, coef_tiles, n_out_tiles, attrs.n_out, Wt}; + KernelHandle compute_id = CreateKernel( + program, + std::string(kKernelDir) + "compute.cpp", + all_cores, + ComputeConfig{ + .math_fidelity = MathFidelity::HiFi4, + .fp32_dest_acc_en = true, + .dst_full_sync_en = true, // 8 fp32 dest slots so the fan-in (d<=5) fits in dst[0..d-1] + .compile_args = compute_ct}); + + // Runtime args. The sparsity pattern (deg/ks/js) is identical on every core. + auto* x_buf = x.buffer(); + auto* coef_buf = coef.buffer(); + auto* out_buf = output.buffer(); + auto cores = corerange_to_cores(all_cores, num_cores, true); + + uint32_t row_offset = 0; + for (const auto& core : cores) { + uint32_t rows; + if (core_group_1.contains(core)) { + rows = rows_per_core_1; + } else { + rows = rows_per_core_2; + } + SetRuntimeArgs(program, reader_id, core, {x_buf->address(), coef_buf->address(), row_offset, rows}); + SetRuntimeArgs(program, writer_id, core, {out_buf->address(), row_offset, rows}); + + std::vector compute_rt = {rows}; + compute_rt.insert(compute_rt.end(), attrs.deg.begin(), attrs.deg.end()); + compute_rt.insert(compute_rt.end(), attrs.ks.begin(), attrs.ks.end()); + compute_rt.insert(compute_rt.end(), attrs.js.begin(), attrs.js.end()); + SetRuntimeArgs(program, compute_id, core, compute_rt); + + row_offset += rows; + } + + return cached_program_t{ + std::move(program), {reader_id, writer_id, compute_id, cores}}; +} + +void FusedRotateProgramFactory::override_runtime_arguments( + cached_program_t& cached_program, const FusedRotateParams&, const FusedRotateInputs& inputs, Tensor& output) { + auto& program = cached_program.program; + const auto& cores = cached_program.shared_variables.cores; + const auto reader_id = cached_program.shared_variables.reader_kernel_id; + const auto writer_id = cached_program.shared_variables.writer_kernel_id; + auto* x_buf = inputs.x_flat.buffer(); + auto* coef_buf = inputs.coef_exp.buffer(); + auto* out_buf = output.buffer(); + for (const auto& core : cores) { + auto& ra = GetRuntimeArgs(program, reader_id, core); + ra[0] = x_buf->address(); + ra[1] = coef_buf->address(); + auto& wa = GetRuntimeArgs(program, writer_id, core); + wa[0] = out_buf->address(); + } +} + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/fused_rotate_program_factory.hpp b/custom_kernels/fused_rotate/device/fused_rotate_program_factory.hpp new file mode 100644 index 0000000..18f31ec --- /dev/null +++ b/custom_kernels/fused_rotate/device/fused_rotate_program_factory.hpp @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "fused_rotate_device_operation_types.hpp" +#include "ttnn/device_operation.hpp" + +namespace ttnn::experimental::prim { + +struct FusedRotateSharedVariables { + tt::tt_metal::KernelHandle reader_kernel_id = 0; + tt::tt_metal::KernelHandle writer_kernel_id = 0; + tt::tt_metal::KernelHandle compute_kernel_id = 0; + std::vector cores; +}; + +struct FusedRotateProgramFactory { + using shared_variables_t = FusedRotateSharedVariables; + using cached_program_t = ttnn::device_operation::CachedProgram; + + static cached_program_t create( + const FusedRotateParams& operation_attributes, const FusedRotateInputs& inputs, Tensor& output); + + static void override_runtime_arguments( + cached_program_t& cached_program, + const FusedRotateParams& operation_attributes, + const FusedRotateInputs& inputs, + Tensor& output); +}; + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/gate_device_operation.cpp b/custom_kernels/fused_rotate/device/gate_device_operation.cpp new file mode 100644 index 0000000..d7920cc --- /dev/null +++ b/custom_kernels/fused_rotate/device/gate_device_operation.cpp @@ -0,0 +1,72 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "gate_device_operation.hpp" +#include "gate_program_factory.hpp" +#include "ttnn/device_operation.hpp" + +#include + +using namespace tt::constants; +using namespace tt::tt_metal; + +namespace ttnn::experimental::prim { + +void GateDeviceOperation::validate_on_program_cache_miss( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + const auto& a = inputs.a; + const auto& gate = inputs.gate; + TT_FATAL( + a.storage_type() == StorageType::DEVICE && gate.storage_type() == StorageType::DEVICE, + "fused_gate operands must be on device"); + TT_FATAL(a.layout() == Layout::TILE && gate.layout() == Layout::TILE, "fused_gate requires TILE layout"); + // bf16 or bf8_b (one tile size for all CBs -> a and gate share a dtype). bf8 halves the + // [E,W] edge-activation traffic on the bandwidth-bound replay. + TT_FATAL( + (a.dtype() == DataType::BFLOAT16 || a.dtype() == DataType::BFLOAT8_B) && gate.dtype() == a.dtype(), + "fused_gate requires bf16 or bf8_b (a and gate same dtype)"); + TT_FATAL(attrs.Wt == attrs.Ht + attrs.Gt, "Wt ({}) must equal Ht ({}) + Gt ({})", attrs.Wt, attrs.Ht, attrs.Gt); + const auto& as_ = a.padded_shape(); + const auto& gs = gate.padded_shape(); + TT_FATAL(as_[-1] == attrs.Wt * TILE_WIDTH, "a last dim {} != Wt*32 {}", as_[-1], attrs.Wt * TILE_WIDTH); + TT_FATAL(gs[-1] == attrs.Gt * TILE_WIDTH, "gate last dim {} != Gt*32 {}", gs[-1], attrs.Gt * TILE_WIDTH); + TT_FATAL(as_[-2] == gs[-2], "a and gate must have the same number of rows"); + if (attrs.mode == 1) { + const auto& b = inputs.b; + TT_FATAL(b.storage_type() == StorageType::DEVICE && b.layout() == Layout::TILE, "fused_gate b device/TILE"); + TT_FATAL(b.dtype() == a.dtype(), "fused_gate b must match a dtype"); + TT_FATAL(b.padded_shape()[-2] == as_[-2], "b and a must have the same number of rows"); + } +} + +GateDeviceOperation::spec_return_value_t GateDeviceOperation::compute_output_specs( + const operation_attributes_t&, const tensor_args_t& inputs) { + const auto& a = inputs.a; + return TensorSpec(a.logical_shape(), TensorLayout(a.dtype(), PageConfig(Layout::TILE), a.memory_config())); +} + +GateDeviceOperation::tensor_return_value_t GateDeviceOperation::create_output_tensors( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + return create_device_tensor(compute_output_specs(attrs, inputs), inputs.a.device()); +} + +ttsl::hash::hash_t GateDeviceOperation::compute_program_hash( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + return tt::tt_metal::operation::hash_operation( + attrs.Wt, attrs.Gt, attrs.Ht, attrs.mode, inputs.a.dtype(), inputs.a.memory_config(), inputs.a.padded_shape()); +} + +} // namespace ttnn::experimental::prim + +namespace ttnn::prim { + +Tensor fused_gate( + const Tensor& a, const Tensor& gate, const Tensor& b, uint32_t Wt, uint32_t Gt, uint32_t Ht, uint32_t mode) { + using OperationType = ttnn::experimental::prim::GateDeviceOperation; + auto attrs = OperationType::operation_attributes_t{.Wt = Wt, .Gt = Gt, .Ht = Ht, .mode = mode}; + return ttnn::device_operation::launch( + attrs, OperationType::tensor_args_t{.a = a, .gate = gate, .b = b}); +} + +} // namespace ttnn::prim diff --git a/custom_kernels/fused_rotate/device/gate_device_operation.hpp b/custom_kernels/fused_rotate/device/gate_device_operation.hpp new file mode 100644 index 0000000..9366456 --- /dev/null +++ b/custom_kernels/fused_rotate/device/gate_device_operation.hpp @@ -0,0 +1,31 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "ttnn/tensor/tensor.hpp" +#include "gate_program_factory.hpp" +#include "gate_device_operation_types.hpp" + +namespace ttnn::experimental::prim { + +struct GateDeviceOperation { + using operation_attributes_t = GateParams; + using tensor_args_t = GateInputs; + using spec_return_value_t = TensorSpec; + using tensor_return_value_t = Tensor; + using program_factory_t = std::variant; + + static void validate_on_program_cache_miss(const operation_attributes_t&, const tensor_args_t&); + static spec_return_value_t compute_output_specs(const operation_attributes_t&, const tensor_args_t&); + static tensor_return_value_t create_output_tensors(const operation_attributes_t&, const tensor_args_t&); + static ttsl::hash::hash_t compute_program_hash(const operation_attributes_t&, const tensor_args_t&); +}; + +} // namespace ttnn::experimental::prim + +namespace ttnn::prim { +Tensor fused_gate( + const Tensor& a, const Tensor& gate, const Tensor& b, uint32_t Wt, uint32_t Gt, uint32_t Ht, uint32_t mode); +} // namespace ttnn::prim diff --git a/custom_kernels/fused_rotate/device/gate_device_operation_types.hpp b/custom_kernels/fused_rotate/device/gate_device_operation_types.hpp new file mode 100644 index 0000000..7e87d4a --- /dev/null +++ b/custom_kernels/fused_rotate/device/gate_device_operation_types.hpp @@ -0,0 +1,31 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include + +#include "ttnn/tensor/tensor.hpp" + +namespace ttnn::experimental::prim { + +// Fused SO(3) gate activation (column-split elementwise, NO reduction). Per edge row of nsph*H cols: +// out[:, :H] = silu(a[:, :H]) (mode 0, forward) OR +// a[:, :H] * silu'(b[:, :H]) (mode 1, backward = silu_bw(a, b)) +// out[:, H:] = a[:, H:] * gate[:, :] (both modes; gate is [E,(nsph-1)*H]) +// Collapses the slice+silu+slice+multiply+concat chain into one kernel. b is only read in mode 1. +struct GateParams { + uint32_t Wt; // total tiles per row = nsph*H/32 + uint32_t Gt; // gate tiles per row = (nsph-1)*H/32 + uint32_t Ht; // scalar(l=0) tiles = H/32 + uint32_t mode; // 0 = forward (silu), 1 = backward (silu_bw) +}; + +struct GateInputs { + Tensor a; // [E, nsph*H] TILE bf16 (fwd: x; bw: g_out) + Tensor gate; // [E, (nsph-1)*H] TILE bf16 (expanded sigmoid gate) + Tensor b; // [E, nsph*H] TILE bf16 (bw: x for silu'; fwd: pass a, unused) +}; + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/gate_program_factory.cpp b/custom_kernels/fused_rotate/device/gate_program_factory.cpp new file mode 100644 index 0000000..0b34916 --- /dev/null +++ b/custom_kernels/fused_rotate/device/gate_program_factory.cpp @@ -0,0 +1,128 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "gate_program_factory.hpp" +#include "gate_device_operation_types.hpp" + +#include +#include +#include +#include + +using namespace tt::constants; +using namespace tt::tt_metal; + +namespace ttnn::experimental::prim { + +static const char* kGateKernelDir = "ttnn/cpp/ttnn/operations/experimental/fused_rotate/device/kernels/"; + +GateProgramFactory::cached_program_t GateProgramFactory::create( + const GateParams& attrs, const GateInputs& inputs, Tensor& output) { + Program program{}; + + const auto& a = inputs.a; + const auto& gate = inputs.gate; + const auto& b = inputs.b; + + const uint32_t Wt = attrs.Wt; + const uint32_t Gt = attrs.Gt; + const uint32_t Ht = attrs.Ht; + const uint32_t Et = a.padded_shape()[-2] / TILE_HEIGHT; + + tt::DataFormat data_format = datatype_to_dataformat_converter(a.dtype()); + const uint32_t tile_bytes = tile_size(data_format); + + auto* device = a.device(); + CoreCoord grid = device->compute_with_storage_grid_size(); + auto [num_cores, all_cores, core_group_1, core_group_2, rows_per_core_1, rows_per_core_2] = + tt::tt_metal::split_work_to_cores(grid, Et); + + auto make_cb = [&](uint32_t cb_index, uint32_t num_tiles) { + CircularBufferConfig cfg = + CircularBufferConfig(num_tiles * tile_bytes, {{cb_index, data_format}}).set_page_size(cb_index, tile_bytes); + CreateCircularBuffer(program, all_cores, cfg); + }; + constexpr uint32_t cb_a = tt::CBIndex::c_0; + constexpr uint32_t cb_gate = tt::CBIndex::c_1; + constexpr uint32_t cb_b = tt::CBIndex::c_2; + constexpr uint32_t cb_sp = tt::CBIndex::c_3; // silu'(b), Ht tiles (bw only) + constexpr uint32_t cb_s = tt::CBIndex::c_4; // scratch sigmoid(b) (bw only) + constexpr uint32_t cb_p = tt::CBIndex::c_5; // scratch silu(b)=b*s (bw only) + constexpr uint32_t cb_r = tt::CBIndex::c_6; // scratch p*s (bw only) + constexpr uint32_t cb_tmp = tt::CBIndex::c_7; // scratch s+p (bw only) + constexpr uint32_t cb_out = tt::CBIndex::c_16; + make_cb(cb_a, 2 * Wt); + make_cb(cb_gate, 2 * Gt); + make_cb(cb_b, 2 * Ht); + make_cb(cb_sp, 2 * Ht); + make_cb(cb_s, 2 * Ht); + make_cb(cb_p, 2 * Ht); + make_cb(cb_r, 2 * Ht); + make_cb(cb_tmp, 2 * Ht); + make_cb(cb_out, 2 * Wt); + + // ---- reader ---- + std::vector reader_ct = {cb_a, cb_gate, cb_b, Wt, Gt, Ht, tile_bytes, attrs.mode}; + TensorAccessorArgs(*a.buffer()).append_to(reader_ct); + TensorAccessorArgs(*gate.buffer()).append_to(reader_ct); + TensorAccessorArgs(*b.buffer()).append_to(reader_ct); + KernelHandle reader_id = CreateKernel( + program, std::string(kGateKernelDir) + "gate_reader.cpp", all_cores, ReaderDataMovementConfig(reader_ct)); + + // ---- writer ---- + std::vector writer_ct = {cb_out, Wt, tile_bytes}; + TensorAccessorArgs(*output.buffer()).append_to(writer_ct); + KernelHandle writer_id = CreateKernel( + program, std::string(kGateKernelDir) + "gate_writer.cpp", all_cores, WriterDataMovementConfig(writer_ct)); + + // ---- compute ---- + std::vector compute_ct = { + cb_a, cb_gate, cb_b, cb_sp, cb_out, Wt, Gt, Ht, attrs.mode, cb_s, cb_p, cb_r, cb_tmp}; + KernelHandle compute_id = CreateKernel( + program, + std::string(kGateKernelDir) + "gate_compute.cpp", + all_cores, + ComputeConfig{.math_fidelity = MathFidelity::HiFi4, .fp32_dest_acc_en = true, .compile_args = compute_ct}); + + auto* a_buf = a.buffer(); + auto* gate_buf = gate.buffer(); + auto* b_buf = b.buffer(); + auto* out_buf = output.buffer(); + auto cores = corerange_to_cores(all_cores, num_cores, true); + + uint32_t row_offset = 0; + for (const auto& core : cores) { + uint32_t rows = core_group_1.contains(core) ? rows_per_core_1 : rows_per_core_2; + SetRuntimeArgs( + program, reader_id, core, + {a_buf->address(), gate_buf->address(), b_buf->address(), row_offset, rows}); + SetRuntimeArgs(program, writer_id, core, {out_buf->address(), row_offset, rows}); + SetRuntimeArgs(program, compute_id, core, {rows}); + row_offset += rows; + } + + return cached_program_t{std::move(program), {reader_id, writer_id, compute_id, cores}}; +} + +void GateProgramFactory::override_runtime_arguments( + cached_program_t& cached_program, const GateParams&, const GateInputs& inputs, Tensor& output) { + auto& program = cached_program.program; + const auto& cores = cached_program.shared_variables.cores; + const auto reader_id = cached_program.shared_variables.reader_kernel_id; + const auto writer_id = cached_program.shared_variables.writer_kernel_id; + auto* a_buf = inputs.a.buffer(); + auto* gate_buf = inputs.gate.buffer(); + auto* b_buf = inputs.b.buffer(); + auto* out_buf = output.buffer(); + for (const auto& core : cores) { + auto& ra = GetRuntimeArgs(program, reader_id, core); + ra[0] = a_buf->address(); + ra[1] = gate_buf->address(); + ra[2] = b_buf->address(); + auto& wa = GetRuntimeArgs(program, writer_id, core); + wa[0] = out_buf->address(); + } +} + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/gate_program_factory.hpp b/custom_kernels/fused_rotate/device/gate_program_factory.hpp new file mode 100644 index 0000000..63a0586 --- /dev/null +++ b/custom_kernels/fused_rotate/device/gate_program_factory.hpp @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "gate_device_operation_types.hpp" +#include "ttnn/device_operation.hpp" + +namespace ttnn::experimental::prim { + +struct GateSharedVariables { + tt::tt_metal::KernelHandle reader_kernel_id = 0; + tt::tt_metal::KernelHandle writer_kernel_id = 0; + tt::tt_metal::KernelHandle compute_kernel_id = 0; + std::vector cores; +}; + +struct GateProgramFactory { + using shared_variables_t = GateSharedVariables; + using cached_program_t = ttnn::device_operation::CachedProgram; + + static cached_program_t create(const GateParams& operation_attributes, const GateInputs& inputs, Tensor& output); + + static void override_runtime_arguments( + cached_program_t& cached_program, + const GateParams& operation_attributes, + const GateInputs& inputs, + Tensor& output); +}; + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/gc_device_operation.cpp b/custom_kernels/fused_rotate/device/gc_device_operation.cpp new file mode 100644 index 0000000..d3402e8 --- /dev/null +++ b/custom_kernels/fused_rotate/device/gc_device_operation.cpp @@ -0,0 +1,105 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "gc_device_operation.hpp" +#include "gc_program_factory.hpp" +#include "ttnn/device_operation.hpp" + +#include + +using namespace tt::constants; +using namespace tt::tt_metal; + +namespace ttnn::experimental::prim { + +void FusedGcDeviceOperation::validate_on_program_cache_miss( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + const auto& gout = inputs.gout; + const auto& xin = inputs.xin; + const auto& sel = inputs.sel; + TT_FATAL( + gout.storage_type() == StorageType::DEVICE && xin.storage_type() == StorageType::DEVICE && + sel.storage_type() == StorageType::DEVICE, + "fused_rotate_gc operands must be on device"); + TT_FATAL( + gout.layout() == Layout::TILE && xin.layout() == Layout::TILE && sel.layout() == Layout::TILE, + "fused_rotate_gc requires TILE layout"); + // bf16 or bf8_b (one tile size for all CBs -> gout/xin/sel share a dtype). + TT_FATAL( + (gout.dtype() == DataType::BFLOAT16 || gout.dtype() == DataType::BFLOAT8_B) && + xin.dtype() == gout.dtype() && sel.dtype() == gout.dtype(), + "fused_rotate_gc requires bf16 or bf8_b inputs (gout/xin/sel same dtype)"); + const auto& gs = gout.padded_shape(); + const auto& xs = xin.padded_shape(); + const auto& ss = sel.padded_shape(); + TT_FATAL(attrs.W % TILE_WIDTH == 0, "W ({}) must be a multiple of TILE_WIDTH", attrs.W); + TT_FATAL(gs[-1] == attrs.n_out * attrs.W, "gout last dim {} != n_out*W {}", gs[-1], attrs.n_out * attrs.W); + TT_FATAL(xs[-1] == attrs.n_in * attrs.W, "xin last dim {} != n_in*W {}", xs[-1], attrs.n_in * attrs.W); + TT_FATAL(gs[-2] == xs[-2], "gout and xin must have the same number of rows (edges)"); + TT_FATAL(ss[-2] == TILE_HEIGHT && ss[-1] == 32 * TILE_WIDTH, "sel must be [32, 32*32]"); + TT_FATAL(attrs.is_.size() == attrs.nnz && attrs.js.size() == attrs.nnz, "is_/js size must equal nnz"); +} + +FusedGcDeviceOperation::spec_return_value_t FusedGcDeviceOperation::compute_output_specs( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + const auto& gout = inputs.gout; + ttnn::Shape out_shape(gout.logical_shape()); + out_shape[-1] = attrs.nnz; // TILE layout pads to ceil(nnz/32)*32 + return TensorSpec(out_shape, TensorLayout(gout.dtype(), PageConfig(Layout::TILE), gout.memory_config())); +} + +FusedGcDeviceOperation::tensor_return_value_t FusedGcDeviceOperation::create_output_tensors( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + return create_device_tensor(compute_output_specs(attrs, inputs), inputs.gout.device()); +} + +ttsl::hash::hash_t FusedGcDeviceOperation::compute_program_hash( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + uint64_t ph = 1469598103934665603ULL; // FNV-1a over the sparsity pattern (set as runtime args) + auto mix = [&](uint32_t v) { ph = (ph ^ v) * 1099511628211ULL; }; + for (auto v : attrs.is_) { + mix(v); + } + for (auto v : attrs.js) { + mix(v); + } + return tt::tt_metal::operation::hash_operation( + attrs.n_out, + attrs.n_in, + attrs.W, + attrs.nnz, + static_cast(ph), + static_cast(ph >> 32), + inputs.gout.dtype(), + inputs.gout.memory_config(), + inputs.gout.padded_shape(), + inputs.xin.padded_shape()); +} + +} // namespace ttnn::experimental::prim + +namespace ttnn::prim { + +Tensor fused_rotate_gc( + const Tensor& gout, + const Tensor& xin, + const Tensor& sel, + uint32_t n_out, + uint32_t n_in, + uint32_t W, + const std::vector& is_, + const std::vector& js) { + using OperationType = ttnn::experimental::prim::FusedGcDeviceOperation; + auto attrs = OperationType::operation_attributes_t{ + .n_out = n_out, + .n_in = n_in, + .W = W, + .nnz = static_cast(is_.size()), + .is_ = is_, + .js = js}; + return ttnn::device_operation::launch( + attrs, OperationType::tensor_args_t{.gout = gout, .xin = xin, .sel = sel}); +} + +} // namespace ttnn::prim diff --git a/custom_kernels/fused_rotate/device/gc_device_operation.hpp b/custom_kernels/fused_rotate/device/gc_device_operation.hpp new file mode 100644 index 0000000..072e898 --- /dev/null +++ b/custom_kernels/fused_rotate/device/gc_device_operation.hpp @@ -0,0 +1,38 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "ttnn/tensor/tensor.hpp" +#include "gc_program_factory.hpp" +#include "gc_device_operation_types.hpp" + +namespace ttnn::experimental::prim { + +struct FusedGcDeviceOperation { + using operation_attributes_t = FusedGcParams; + using tensor_args_t = FusedGcInputs; + using spec_return_value_t = TensorSpec; + using tensor_return_value_t = Tensor; + using program_factory_t = std::variant; + + static void validate_on_program_cache_miss(const operation_attributes_t&, const tensor_args_t&); + static spec_return_value_t compute_output_specs(const operation_attributes_t&, const tensor_args_t&); + static tensor_return_value_t create_output_tensors(const operation_attributes_t&, const tensor_args_t&); + static ttsl::hash::hash_t compute_program_hash(const operation_attributes_t&, const tensor_args_t&); +}; + +} // namespace ttnn::experimental::prim + +namespace ttnn::prim { +Tensor fused_rotate_gc( + const Tensor& gout, + const Tensor& xin, + const Tensor& sel, + uint32_t n_out, + uint32_t n_in, + uint32_t W, + const std::vector& is_, + const std::vector& js); +} // namespace ttnn::prim diff --git a/custom_kernels/fused_rotate/device/gc_device_operation_types.hpp b/custom_kernels/fused_rotate/device/gc_device_operation_types.hpp new file mode 100644 index 0000000..b88248c --- /dev/null +++ b/custom_kernels/fused_rotate/device/gc_device_operation_types.hpp @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include + +#include "ttnn/tensor/tensor.hpp" + +namespace ttnn::experimental::prim { + +// Fused per-edge coefficient adjoint (rotate_bw dE/dcoef): +// gc[e, k] = sum_w gout[e, is[k]*W + w] * xin[e, js[k]*W + w] for each structural nonzero k. +// Output is compact [E, ceil(nnz/32)*32]; column k holds the dot for nonzero k. +struct FusedGcParams { + uint32_t n_out; // number of gout coordinate blocks + uint32_t n_in; // number of xin coordinate blocks + uint32_t W; // channels per coordinate (multiple of TILE_WIDTH) + uint32_t nnz; // number of structural nonzeros + std::vector is_; // length nnz: gout block per nonzero + std::vector js; // length nnz: xin block per nonzero +}; + +struct FusedGcInputs { + Tensor gout; // [E, n_out*W] TILE bf16 + Tensor xin; // [E, n_in*W] TILE bf16 + Tensor sel; // [32, 32*32] TILE bf16 (tile c has column c all-ones) +}; + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/gc_program_factory.cpp b/custom_kernels/fused_rotate/device/gc_program_factory.cpp new file mode 100644 index 0000000..2c6ae63 --- /dev/null +++ b/custom_kernels/fused_rotate/device/gc_program_factory.cpp @@ -0,0 +1,127 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "gc_program_factory.hpp" +#include "gc_device_operation_types.hpp" + +#include +#include +#include +#include + +using namespace tt::constants; +using namespace tt::tt_metal; + +namespace ttnn::experimental::prim { + +static const char* kGcKernelDir = "ttnn/cpp/ttnn/operations/experimental/fused_rotate/device/kernels/"; + +FusedGcProgramFactory::cached_program_t FusedGcProgramFactory::create( + const FusedGcParams& attrs, const FusedGcInputs& inputs, Tensor& output) { + Program program{}; + + const auto& gout = inputs.gout; + const auto& xin = inputs.xin; + const auto& sel = inputs.sel; + + const uint32_t Wt = attrs.W / TILE_WIDTH; + const uint32_t n_out_tiles = attrs.n_out * Wt; + const uint32_t n_in_tiles = attrs.n_in * Wt; + const uint32_t out_tiles = (attrs.nnz + TILE_WIDTH - 1) / TILE_WIDTH; // ceil(nnz/32) + const uint32_t Et = gout.padded_shape()[-2] / TILE_HEIGHT; + + tt::DataFormat data_format = datatype_to_dataformat_converter(gout.dtype()); + const uint32_t tile_bytes = tile_size(data_format); + + auto* device = gout.device(); + CoreCoord grid = device->compute_with_storage_grid_size(); + auto [num_cores, all_cores, core_group_1, core_group_2, rows_per_core_1, rows_per_core_2] = + tt::tt_metal::split_work_to_cores(grid, Et); + + auto make_cb = [&](uint32_t cb_index, uint32_t num_tiles) { + CircularBufferConfig cfg = + CircularBufferConfig(num_tiles * tile_bytes, {{cb_index, data_format}}).set_page_size(cb_index, tile_bytes); + CreateCircularBuffer(program, all_cores, cfg); + }; + constexpr uint32_t cb_gout = tt::CBIndex::c_0; + constexpr uint32_t cb_xin = tt::CBIndex::c_1; + constexpr uint32_t cb_sel = tt::CBIndex::c_2; + constexpr uint32_t cb_prod = tt::CBIndex::c_24; + constexpr uint32_t cb_out = tt::CBIndex::c_16; + make_cb(cb_gout, 2 * n_out_tiles); + make_cb(cb_xin, 2 * n_in_tiles); + make_cb(cb_sel, 32); + make_cb(cb_prod, 32 * Wt); // one output-tile worth of products (d<=32 nonzeros x Wt) + make_cb(cb_out, 2 * out_tiles); + + // ---- reader ---- + std::vector reader_ct = {cb_gout, cb_xin, cb_sel, n_out_tiles, n_in_tiles, tile_bytes}; + TensorAccessorArgs(*gout.buffer()).append_to(reader_ct); + TensorAccessorArgs(*xin.buffer()).append_to(reader_ct); + TensorAccessorArgs(*sel.buffer()).append_to(reader_ct); + KernelHandle reader_id = + CreateKernel(program, std::string(kGcKernelDir) + "gc_reader.cpp", all_cores, ReaderDataMovementConfig(reader_ct)); + + // ---- writer (reuse the generic writer: cb_out, out_tiles, tile_bytes) ---- + std::vector writer_ct = {cb_out, out_tiles, tile_bytes}; + TensorAccessorArgs(*output.buffer()).append_to(writer_ct); + KernelHandle writer_id = + CreateKernel(program, std::string(kGcKernelDir) + "writer.cpp", all_cores, WriterDataMovementConfig(writer_ct)); + + // ---- compute ---- + std::vector compute_ct = {cb_gout, cb_xin, cb_sel, cb_prod, cb_out, + n_out_tiles, n_in_tiles, Wt, attrs.nnz, out_tiles}; + KernelHandle compute_id = CreateKernel( + program, + std::string(kGcKernelDir) + "gc_compute.cpp", + all_cores, + ComputeConfig{ + .math_fidelity = MathFidelity::HiFi4, .fp32_dest_acc_en = true, .compile_args = compute_ct}); + + auto* gout_buf = gout.buffer(); + auto* xin_buf = xin.buffer(); + auto* sel_buf = sel.buffer(); + auto* out_buf = output.buffer(); + auto cores = corerange_to_cores(all_cores, num_cores, true); + + uint32_t row_offset = 0; + for (const auto& core : cores) { + uint32_t rows = core_group_1.contains(core) ? rows_per_core_1 : rows_per_core_2; + SetRuntimeArgs( + program, reader_id, core, + {gout_buf->address(), xin_buf->address(), sel_buf->address(), row_offset, rows}); + SetRuntimeArgs(program, writer_id, core, {out_buf->address(), row_offset, rows}); + + std::vector compute_rt = {rows}; + compute_rt.insert(compute_rt.end(), attrs.is_.begin(), attrs.is_.end()); + compute_rt.insert(compute_rt.end(), attrs.js.begin(), attrs.js.end()); + SetRuntimeArgs(program, compute_id, core, compute_rt); + + row_offset += rows; + } + + return cached_program_t{std::move(program), {reader_id, writer_id, compute_id, cores}}; +} + +void FusedGcProgramFactory::override_runtime_arguments( + cached_program_t& cached_program, const FusedGcParams&, const FusedGcInputs& inputs, Tensor& output) { + auto& program = cached_program.program; + const auto& cores = cached_program.shared_variables.cores; + const auto reader_id = cached_program.shared_variables.reader_kernel_id; + const auto writer_id = cached_program.shared_variables.writer_kernel_id; + auto* gout_buf = inputs.gout.buffer(); + auto* xin_buf = inputs.xin.buffer(); + auto* sel_buf = inputs.sel.buffer(); + auto* out_buf = output.buffer(); + for (const auto& core : cores) { + auto& ra = GetRuntimeArgs(program, reader_id, core); + ra[0] = gout_buf->address(); + ra[1] = xin_buf->address(); + ra[2] = sel_buf->address(); + auto& wa = GetRuntimeArgs(program, writer_id, core); + wa[0] = out_buf->address(); + } +} + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/gc_program_factory.hpp b/custom_kernels/fused_rotate/device/gc_program_factory.hpp new file mode 100644 index 0000000..0d9013c --- /dev/null +++ b/custom_kernels/fused_rotate/device/gc_program_factory.hpp @@ -0,0 +1,29 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "gc_device_operation_types.hpp" +#include "ttnn/device_operation.hpp" + +namespace ttnn::experimental::prim { + +struct FusedGcSharedVariables { + tt::tt_metal::KernelHandle reader_kernel_id = 0; + tt::tt_metal::KernelHandle writer_kernel_id = 0; + tt::tt_metal::KernelHandle compute_kernel_id = 0; + std::vector cores; +}; + +struct FusedGcProgramFactory { + using shared_variables_t = FusedGcSharedVariables; + using cached_program_t = ttnn::device_operation::CachedProgram; + + static cached_program_t create(const FusedGcParams& attrs, const FusedGcInputs& inputs, Tensor& output); + + static void override_runtime_arguments( + cached_program_t& cached_program, const FusedGcParams& attrs, const FusedGcInputs& inputs, Tensor& output); +}; + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/kernels/compute.cpp b/custom_kernels/fused_rotate/device/kernels/compute.cpp new file mode 100644 index 0000000..eb33024 --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/compute.cpp @@ -0,0 +1,74 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Fused per-edge sparse Wigner rotation. For each tile-row (32 edges) and each output block i, +// accumulate the fan-in out[i] = sum_{(i,j,k)} coef_tile[k] * x_tile[j] entirely in the dest +// registers, then pack once. All `nnz` multiply-accumulates run in a single kernel launch with +// one DRAM read of x and one write of out (vs `nnz` separate ttnn.addcmul dispatches). + +#include +#include "api/compute/compute_kernel_api.h" +#include "api/compute/eltwise_binary.h" +#include "api/compute/eltwise_binary_sfpu.h" +#include "api/compute/tile_move_copy.h" + +void kernel_main() { + constexpr uint32_t cb_x = get_compile_time_arg_val(0); + constexpr uint32_t cb_coef = get_compile_time_arg_val(1); + constexpr uint32_t cb_out = get_compile_time_arg_val(2); + constexpr uint32_t n_in_tiles = get_compile_time_arg_val(3); + constexpr uint32_t coef_tiles = get_compile_time_arg_val(4); + constexpr uint32_t n_out_tiles = get_compile_time_arg_val(5); + constexpr uint32_t n_out = get_compile_time_arg_val(6); + constexpr uint32_t Wt = get_compile_time_arg_val(7); + + uint32_t arg = 0; + const uint32_t num_rows = get_arg_val(arg++); + // deg[0..n_out-1], then ks[0..nnz-1], then js[0..nnz-1] + const uint32_t deg_base = arg; + const uint32_t ks_base = deg_base + n_out; + const uint32_t nnz = coef_tiles; + const uint32_t js_base = ks_base + nnz; + + binary_op_init_common(cb_coef, cb_x, cb_out); + + for (uint32_t r = 0; r < num_rows; r++) { + cb_wait_front(cb_x, n_in_tiles); + cb_wait_front(cb_coef, coef_tiles); + cb_reserve_back(cb_out, n_out_tiles); + + uint32_t off = 0; // running offset into ks/js + for (uint32_t i = 0; i < n_out; i++) { + const uint32_t d = get_arg_val(deg_base + i); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + // Fan-in: compute all d products into dst[0..d-1] (FPU), then sum them into + // dst[0] with the SFPU dst-to-dst adder. Doing all muls first then all adds + // (rather than alternating FPU/SFPU) is required for correctness. Needs d dst + // slots -> the program uses dst_full_sync_en (8 fp32 slots; d<=5 for uma-s lmax=2). + mul_tiles_init(cb_coef, cb_x); + for (uint32_t m = 0; m < d; m++) { + const uint32_t k = get_arg_val(ks_base + off + m); + const uint32_t j = get_arg_val(js_base + off + m); + mul_tiles(cb_coef, cb_x, k, j * Wt + wt, m); + } + if (d > 1) { + add_binary_tile_init(); + for (uint32_t m = 1; m < d; m++) { + add_binary_tile(0, m, 0); + } + } + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_out, i * Wt + wt); + tile_regs_release(); + } + off += d; + } + + cb_push_back(cb_out, n_out_tiles); + cb_pop_front(cb_x, n_in_tiles); + cb_pop_front(cb_coef, coef_tiles); + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/gate_compute.cpp b/custom_kernels/fused_rotate/device/kernels/gate_compute.cpp new file mode 100644 index 0000000..0127da6 --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/gate_compute.cpp @@ -0,0 +1,159 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Fused SO(3) gate activation (column-split elementwise, NO reduction). Per edge tile-row: +// scalar tiles [0, Ht): mode 0 (fwd): out = silu(a) +// mode 1 (bw): out = a * silu'(b) (silu_bw, b = x) +// vector tiles [Ht, Wt): out = a * gate[t-Ht] (both modes) +// silu'(b) = s + p - p*s with s = sigmoid(b), p = silu(b) = b*s. One op per tile_regs_acquire. + +#include +#include "api/compute/compute_kernel_api.h" +#include "api/compute/eltwise_binary.h" +#include "api/compute/tile_move_copy.h" + +void kernel_main() { + constexpr uint32_t cb_a = get_compile_time_arg_val(0); + constexpr uint32_t cb_gate = get_compile_time_arg_val(1); + constexpr uint32_t cb_b = get_compile_time_arg_val(2); + constexpr uint32_t cb_sp = get_compile_time_arg_val(3); + constexpr uint32_t cb_out = get_compile_time_arg_val(4); + constexpr uint32_t Wt = get_compile_time_arg_val(5); + constexpr uint32_t Gt = get_compile_time_arg_val(6); + constexpr uint32_t Ht = get_compile_time_arg_val(7); + constexpr uint32_t mode = get_compile_time_arg_val(8); + constexpr uint32_t cb_s = get_compile_time_arg_val(9); + constexpr uint32_t cb_p = get_compile_time_arg_val(10); + constexpr uint32_t cb_r = get_compile_time_arg_val(11); + constexpr uint32_t cb_tmp = get_compile_time_arg_val(12); + + uint32_t arg = 0; + const uint32_t num_rows = get_arg_val(arg++); + + binary_op_init_common(cb_a, cb_gate, cb_out); + + for (uint32_t r = 0; r < num_rows; r++) { + cb_wait_front(cb_a, Wt); + cb_wait_front(cb_gate, Gt); + + cb_reserve_back(cb_out, Wt); + + // ---- scalar (l=0) tiles [0, Ht) ---- + if (mode == 0) { + // out = silu(a) + for (uint32_t t = 0; t < Ht; t++) { + tile_regs_acquire(); + copy_tile_to_dst_init_short(cb_a); + copy_tile(cb_a, t, 0); + silu_tile_init(); + silu_tile(0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_out, t); + tile_regs_release(); + } + } else { + // silu'(b) into cb_sp, then out = a * silu'(b) + cb_wait_front(cb_b, Ht); + // s = sigmoid(b) -> cb_s + cb_reserve_back(cb_s, Ht); + for (uint32_t t = 0; t < Ht; t++) { + tile_regs_acquire(); + copy_tile_to_dst_init_short(cb_b); + copy_tile(cb_b, t, 0); + sigmoid_tile_init(); + sigmoid_tile(0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_s, t); + tile_regs_release(); + } + cb_push_back(cb_s, Ht); + cb_wait_front(cb_s, Ht); + // p = b * s -> cb_p + cb_reserve_back(cb_p, Ht); + mul_tiles_init(cb_b, cb_s); + for (uint32_t t = 0; t < Ht; t++) { + tile_regs_acquire(); + mul_tiles(cb_b, cb_s, t, t, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_p, t); + tile_regs_release(); + } + cb_push_back(cb_p, Ht); + cb_wait_front(cb_p, Ht); + cb_pop_front(cb_b, Ht); + // r = p * s -> cb_r + cb_reserve_back(cb_r, Ht); + mul_tiles_init(cb_p, cb_s); + for (uint32_t t = 0; t < Ht; t++) { + tile_regs_acquire(); + mul_tiles(cb_p, cb_s, t, t, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_r, t); + tile_regs_release(); + } + cb_push_back(cb_r, Ht); + cb_wait_front(cb_r, Ht); + // tmp = s + p -> cb_tmp + cb_reserve_back(cb_tmp, Ht); + add_tiles_init(cb_s, cb_p); + for (uint32_t t = 0; t < Ht; t++) { + tile_regs_acquire(); + add_tiles(cb_s, cb_p, t, t, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_tmp, t); + tile_regs_release(); + } + cb_push_back(cb_tmp, Ht); + cb_wait_front(cb_tmp, Ht); + cb_pop_front(cb_s, Ht); + cb_pop_front(cb_p, Ht); + // silup = tmp - r -> cb_sp + cb_reserve_back(cb_sp, Ht); + sub_tiles_init(cb_tmp, cb_r); + for (uint32_t t = 0; t < Ht; t++) { + tile_regs_acquire(); + sub_tiles(cb_tmp, cb_r, t, t, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_sp, t); + tile_regs_release(); + } + cb_push_back(cb_sp, Ht); + cb_wait_front(cb_sp, Ht); + cb_pop_front(cb_tmp, Ht); + cb_pop_front(cb_r, Ht); + // out = a * silup -> cb_out[0:Ht) + mul_tiles_init(cb_a, cb_sp); + for (uint32_t t = 0; t < Ht; t++) { + tile_regs_acquire(); + mul_tiles(cb_a, cb_sp, t, t, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_out, t); + tile_regs_release(); + } + cb_pop_front(cb_sp, Ht); + } + + // ---- vector tiles [Ht, Wt): out = a * gate[t-Ht] ---- + mul_tiles_init(cb_a, cb_gate); + for (uint32_t t = Ht; t < Wt; t++) { + tile_regs_acquire(); + mul_tiles(cb_a, cb_gate, t, t - Ht, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_out, t); + tile_regs_release(); + } + + cb_push_back(cb_out, Wt); + cb_pop_front(cb_a, Wt); + cb_pop_front(cb_gate, Gt); + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/gate_reader.cpp b/custom_kernels/fused_rotate/device/kernels/gate_reader.cpp new file mode 100644 index 0000000..2ae2594 --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/gate_reader.cpp @@ -0,0 +1,72 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Reader for fused_gate: streams a (Wt tiles/row) and gate (Gt tiles/row) per edge tile-row; in +// backward mode (mode==1) also streams the first Ht tiles of b (= x, for silu'). a and b share the +// [E, nsph*H] row stride (Wt tiles); gate is [E, (nsph-1)*H] (Gt tiles). + +#include +#include "api/dataflow/dataflow_api.h" + +void kernel_main() { + constexpr uint32_t cb_a = get_compile_time_arg_val(0); + constexpr uint32_t cb_gate = get_compile_time_arg_val(1); + constexpr uint32_t cb_b = get_compile_time_arg_val(2); + constexpr uint32_t Wt = get_compile_time_arg_val(3); + constexpr uint32_t Gt = get_compile_time_arg_val(4); + constexpr uint32_t Ht = get_compile_time_arg_val(5); + constexpr uint32_t tile_bytes = get_compile_time_arg_val(6); + constexpr uint32_t mode = get_compile_time_arg_val(7); + + constexpr auto a_args = TensorAccessorArgs<8>(); + constexpr auto gate_args = TensorAccessorArgs(); + constexpr auto b_args = TensorAccessorArgs(); + + uint32_t arg = 0; + const uint32_t a_addr = get_arg_val(arg++); + const uint32_t gate_addr = get_arg_val(arg++); + const uint32_t b_addr = get_arg_val(arg++); + const uint32_t start_row = get_arg_val(arg++); + const uint32_t num_rows = get_arg_val(arg++); + + const auto a_gen = TensorAccessor(a_args, a_addr, tile_bytes); + const auto gate_gen = TensorAccessor(gate_args, gate_addr, tile_bytes); + const auto b_gen = TensorAccessor(b_args, b_addr, tile_bytes); + + for (uint32_t r = 0; r < num_rows; r++) { + const uint32_t row = start_row + r; + const uint32_t abase = row * Wt; + const uint32_t gbase = row * Gt; + + cb_reserve_back(cb_a, Wt); + uint32_t aw = get_write_ptr(cb_a); + for (uint32_t t = 0; t < Wt; t++) { + noc_async_read_tile(abase + t, a_gen, aw); + aw += tile_bytes; + } + + cb_reserve_back(cb_gate, Gt); + uint32_t gw = get_write_ptr(cb_gate); + for (uint32_t t = 0; t < Gt; t++) { + noc_async_read_tile(gbase + t, gate_gen, gw); + gw += tile_bytes; + } + + if (mode == 1) { + cb_reserve_back(cb_b, Ht); + uint32_t bw = get_write_ptr(cb_b); + for (uint32_t t = 0; t < Ht; t++) { + noc_async_read_tile(abase + t, b_gen, bw); + bw += tile_bytes; + } + } + + noc_async_read_barrier(); + cb_push_back(cb_a, Wt); + cb_push_back(cb_gate, Gt); + if (mode == 1) { + cb_push_back(cb_b, Ht); + } + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/gate_writer.cpp b/custom_kernels/fused_rotate/device/kernels/gate_writer.cpp new file mode 100644 index 0000000..c47e69b --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/gate_writer.cpp @@ -0,0 +1,36 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Writer for fused_gate: writes Wt output tiles per edge tile-row. + +#include +#include "api/dataflow/dataflow_api.h" + +void kernel_main() { + constexpr uint32_t cb_out = get_compile_time_arg_val(0); + constexpr uint32_t Wt = get_compile_time_arg_val(1); + constexpr uint32_t tile_bytes = get_compile_time_arg_val(2); + + constexpr auto out_args = TensorAccessorArgs<3>(); + + uint32_t arg = 0; + const uint32_t out_addr = get_arg_val(arg++); + const uint32_t start_row = get_arg_val(arg++); + const uint32_t num_rows = get_arg_val(arg++); + + const auto out_gen = TensorAccessor(out_args, out_addr, tile_bytes); + + for (uint32_t r = 0; r < num_rows; r++) { + const uint32_t row = start_row + r; + cb_wait_front(cb_out, Wt); + uint32_t rd = get_read_ptr(cb_out); + const uint32_t base = row * Wt; + for (uint32_t t = 0; t < Wt; t++) { + noc_async_write_tile(base + t, out_gen, rd); + rd += tile_bytes; + } + noc_async_write_barrier(); + cb_pop_front(cb_out, Wt); + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/gc_compute.cpp b/custom_kernels/fused_rotate/device/kernels/gc_compute.cpp new file mode 100644 index 0000000..4ff2280 --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/gc_compute.cpp @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Fused per-edge coefficient adjoint (the rotate_bw dE/dcoef). For each tile-row (32 edges) and +// each structural nonzero k=(i,j): +// gc[e, k] = sum_w gout[e, i*W + w] * xin[e, j*W + w] +// i.e. the per-edge dot product over the W channels of output block i of `gout` with input block j +// of `xin`. Done WITHOUT materialising the [E, nnz*W] product concat that the ttnn path builds in +// DRAM: the products stay L1-resident and one accumulating matmul against a column-selector tile +// (`sel[c]` has column c all-ones) does BOTH the W-reduction AND the placement into output column +// c in a single op. Reads gout+xin once, writes the compact gc[E, nnz] once. + +#include +#include "api/compute/compute_kernel_api.h" +#include "api/compute/eltwise_binary.h" +#include "api/compute/matmul.h" +#include "api/compute/tile_move_copy.h" + +void kernel_main() { + constexpr uint32_t cb_gout = get_compile_time_arg_val(0); + constexpr uint32_t cb_xin = get_compile_time_arg_val(1); + constexpr uint32_t cb_sel = get_compile_time_arg_val(2); + constexpr uint32_t cb_prod = get_compile_time_arg_val(3); + constexpr uint32_t cb_out = get_compile_time_arg_val(4); + constexpr uint32_t n_out_tiles = get_compile_time_arg_val(5); // n_out*Wt (gout blocks) + constexpr uint32_t n_in_tiles = get_compile_time_arg_val(6); // n_in*Wt (xin blocks) + constexpr uint32_t Wt = get_compile_time_arg_val(7); + constexpr uint32_t nnz = get_compile_time_arg_val(8); + constexpr uint32_t out_tiles = get_compile_time_arg_val(9); // ceil(nnz/32) + + uint32_t arg = 0; + const uint32_t num_rows = get_arg_val(arg++); + const uint32_t is_base = arg; // is[0..nnz-1]: gout block per nonzero + const uint32_t js_base = is_base + nnz; // js[0..nnz-1]: xin block per nonzero + + binary_op_init_common(cb_gout, cb_xin, cb_out); + + // the 32 column-selector tiles are pos-independent -> loaded once, resident for the kernel. + cb_wait_front(cb_sel, 32); + + for (uint32_t r = 0; r < num_rows; r++) { + cb_wait_front(cb_gout, n_out_tiles); + cb_wait_front(cb_xin, n_in_tiles); + + for (uint32_t ot = 0; ot < out_tiles; ot++) { + uint32_t d = nnz - ot * 32; + if (d > 32) { + d = 32; + } + + // Phase A: products gout_i * xin_j for the d nonzeros of this output tile -> cb_prod + // (d*Wt tiles, packed at explicit slots c*Wt+wt). + mul_tiles_init(cb_gout, cb_xin); + cb_reserve_back(cb_prod, d * Wt); + for (uint32_t c = 0; c < d; c++) { + const uint32_t k = ot * 32 + c; + const uint32_t i = get_arg_val(is_base + k); + const uint32_t j = get_arg_val(js_base + k); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles(cb_gout, cb_xin, i * Wt + wt, j * Wt + wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_prod, c * Wt + wt); + tile_regs_release(); + } + } + cb_push_back(cb_prod, d * Wt); + + // Phase B: one accumulating matmul per (c,wt). matmul_tiles does dst += prod @ sel[c]; + // sel[c] has column c all-ones so prod@sel[c] = rowsum(prod) placed in column c. Summing + // over wt gives the full W-reduction; over c fills the distinct output columns. + cb_wait_front(cb_prod, d * Wt); + mm_init(cb_prod, cb_sel, cb_out); + cb_reserve_back(cb_out, 1); + tile_regs_acquire(); + uint32_t pt = 0; + for (uint32_t c = 0; c < d; c++) { + for (uint32_t wt = 0; wt < Wt; wt++) { + matmul_tiles(cb_prod, cb_sel, pt, c, 0); + pt++; + } + } + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_out, ot); + tile_regs_release(); + cb_pop_front(cb_prod, d * Wt); + cb_push_back(cb_out, 1); + } + + cb_pop_front(cb_gout, n_out_tiles); + cb_pop_front(cb_xin, n_in_tiles); + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/gc_reader.cpp b/custom_kernels/fused_rotate/device/kernels/gc_reader.cpp new file mode 100644 index 0000000..9eca6b0 --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/gc_reader.cpp @@ -0,0 +1,67 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Reader for the coefficient-adjoint (gc) kernel: streams gout + xin tile-rows and loads the 32 +// pos-independent column-selector tiles once (resident for the whole kernel). + +#include +#include "api/dataflow/dataflow_api.h" + +void kernel_main() { + constexpr uint32_t cb_gout = get_compile_time_arg_val(0); + constexpr uint32_t cb_xin = get_compile_time_arg_val(1); + constexpr uint32_t cb_sel = get_compile_time_arg_val(2); + constexpr uint32_t n_out_tiles = get_compile_time_arg_val(3); + constexpr uint32_t n_in_tiles = get_compile_time_arg_val(4); + constexpr uint32_t tile_bytes = get_compile_time_arg_val(5); + + constexpr auto gout_args = TensorAccessorArgs<6>(); + constexpr auto xin_args = TensorAccessorArgs(); + constexpr auto sel_args = TensorAccessorArgs(); + + uint32_t arg = 0; + const uint32_t gout_addr = get_arg_val(arg++); + const uint32_t xin_addr = get_arg_val(arg++); + const uint32_t sel_addr = get_arg_val(arg++); + const uint32_t start_row = get_arg_val(arg++); + const uint32_t num_rows = get_arg_val(arg++); + + const auto gout_gen = TensorAccessor(gout_args, gout_addr, tile_bytes); + const auto xin_gen = TensorAccessor(xin_args, xin_addr, tile_bytes); + const auto sel_gen = TensorAccessor(sel_args, sel_addr, tile_bytes); + + // selector tiles 0..31 (one tile-row of the [32, 32*32] constant), loaded once. + cb_reserve_back(cb_sel, 32); + uint32_t sw = get_write_ptr(cb_sel); + for (uint32_t t = 0; t < 32; t++) { + noc_async_read_tile(t, sel_gen, sw); + sw += tile_bytes; + } + noc_async_read_barrier(); + cb_push_back(cb_sel, 32); + + for (uint32_t r = 0; r < num_rows; r++) { + const uint32_t row = start_row + r; + + cb_reserve_back(cb_gout, n_out_tiles); + uint32_t gw = get_write_ptr(cb_gout); + const uint32_t gout_base = row * n_out_tiles; + for (uint32_t t = 0; t < n_out_tiles; t++) { + noc_async_read_tile(gout_base + t, gout_gen, gw); + gw += tile_bytes; + } + + cb_reserve_back(cb_xin, n_in_tiles); + uint32_t xw = get_write_ptr(cb_xin); + const uint32_t xin_base = row * n_in_tiles; + for (uint32_t t = 0; t < n_in_tiles; t++) { + noc_async_read_tile(xin_base + t, xin_gen, xw); + xw += tile_bytes; + } + + noc_async_read_barrier(); + cb_push_back(cb_gout, n_out_tiles); + cb_push_back(cb_xin, n_in_tiles); + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/lnbw_compute.cpp b/custom_kernels/fused_rotate/device/kernels/lnbw_compute.cpp new file mode 100644 index 0000000..7774d81 --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/lnbw_compute.cpp @@ -0,0 +1,333 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Fused LayerNorm backward (grad wrt the LN input, affine scale folded into gy on host). +// Drop-in for tt_atom's hand-written _ln_bw. For each tile-row (32 edges) over W channels: +// mean_x = mean_w(x); xc = x - mean_x; rstd = rsqrt(mean_w(xc^2) + eps); xhat = xc*rstd +// m1 = mean_w(gy); m2 = mean_w(gy*xhat) +// dx = rstd * (gy - m1 - xhat*m2) +// where gy = g_out * gamma is precomputed on host. Reductions are one accumulating matmul against +// a [32,32] tile whose column 0 is 1/W (rowsum-to-col0, then broadcast back with bcast_cols) -- +// the same matmul-reduce trick as the gc kernel. All W stays L1-resident across the chain; one +// DRAM read of gy+x, one write of dx (vs ~15 ttnn ops). + +#include +#include "api/compute/compute_kernel_api.h" +#include "api/compute/eltwise_binary.h" +#include "api/compute/eltwise_binary_sfpu.h" +#include "api/compute/bcast.h" +#include "api/compute/matmul.h" +#include "api/compute/tile_move_copy.h" +#include "api/compute/eltwise_unary/rsqrt.h" +#include "api/compute/eltwise_unary/binop_with_scalar.h" + +void kernel_main() { + constexpr uint32_t cb_gout = get_compile_time_arg_val(0); // [E, Wt] g_out (matmul, pre-silu-bw) + constexpr uint32_t cb_x = get_compile_time_arg_val(1); // [E, Wt] cached LN input + constexpr uint32_t cb_red = get_compile_time_arg_val(2); // [32,32] col0 = 1/W (resident) + constexpr uint32_t cb_xc = get_compile_time_arg_val(3); // scratch: x - mean_x + constexpr uint32_t cb_xhat = get_compile_time_arg_val(4); // scratch: xc * rstd + constexpr uint32_t cb_prod = get_compile_time_arg_val(5); // scratch: 1-tile products for reduce + constexpr uint32_t cb_s = get_compile_time_arg_val(6); // scratch: 1-tile scalars (mean/rstd) + constexpr uint32_t cb_rstd = get_compile_time_arg_val(7); // rstd (col0) held across the row + constexpr uint32_t cb_dx = get_compile_time_arg_val(8); // [E, Wt] output + constexpr uint32_t Wt = get_compile_time_arg_val(9); + constexpr uint32_t eps_bits = get_compile_time_arg_val(10); + constexpr uint32_t cb_n = get_compile_time_arg_val(11); // [E, Wt] pre-silu activation (LN output) + constexpr uint32_t cb_gamma = get_compile_time_arg_val(12); // [1, Wt] LN affine scale (row bcast, resident) + constexpr uint32_t cb_gy = get_compile_time_arg_val(13); // internal: gy = g_out*silu'(n)*gamma + constexpr uint32_t cb_g1 = get_compile_time_arg_val(14); // internal preamble scratch + + uint32_t arg = 0; + const uint32_t num_rows = get_arg_val(arg++); + + binary_op_init_common(cb_x, cb_red, cb_dx); + cb_wait_front(cb_gamma, Wt); // resident LN scale, loaded once + + // Accumulating matmul reduce of `cb_in` (Wt tiles) into col0 of dst0 -> pack to `cb_scalar`. + // cb_red col0 = 1/W so this is the row-wise mean over W. Caller must have cb_in filled. + auto reduce_mean = [&](uint32_t cb_in, uint32_t cb_scalar) { + mm_init(cb_in, cb_red, cb_scalar); + tile_regs_acquire(); + for (uint32_t wt = 0; wt < Wt; wt++) { + matmul_tiles(cb_in, cb_red, wt, 0, 0); + } + tile_regs_commit(); + tile_regs_wait(); + cb_reserve_back(cb_scalar, 1); + pack_tile(0, cb_scalar); + cb_push_back(cb_scalar, 1); + tile_regs_release(); + }; + + for (uint32_t r = 0; r < num_rows; r++) { + // ===== silu+gamma fold: build gy = g_out * silu'(n) * gamma, one op per tile_regs_acquire. + // silu'(n) = s + p - p*s with s = sigmoid(n), p = n*s (= silu(n)). ===== + cb_wait_front(cb_gout, Wt); + cb_wait_front(cb_n, Wt); + + // s = sigmoid(n) -> cb_xc + cb_reserve_back(cb_xc, Wt); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + copy_tile_to_dst_init_short(cb_n); + copy_tile(cb_n, wt, 0); + sigmoid_tile_init(); + sigmoid_tile(0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_xc, wt); + tile_regs_release(); + } + cb_push_back(cb_xc, Wt); + cb_wait_front(cb_xc, Wt); // s + + // p = n * s -> cb_xhat + cb_reserve_back(cb_xhat, Wt); + mul_tiles_init(cb_n, cb_xc); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles(cb_n, cb_xc, wt, wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_xhat, wt); + tile_regs_release(); + } + cb_push_back(cb_xhat, Wt); + cb_wait_front(cb_xhat, Wt); // p = silu(n) + cb_pop_front(cb_n, Wt); + + // r = p * s -> cb_prod + cb_reserve_back(cb_prod, Wt); + mul_tiles_init(cb_xhat, cb_xc); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles(cb_xhat, cb_xc, wt, wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_prod, wt); + tile_regs_release(); + } + cb_push_back(cb_prod, Wt); + cb_wait_front(cb_prod, Wt); // r = p*s + + // tmp = s + p -> cb_g1 + cb_reserve_back(cb_g1, Wt); + add_tiles_init(cb_xc, cb_xhat); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + add_tiles(cb_xc, cb_xhat, wt, wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_g1, wt); + tile_regs_release(); + } + cb_push_back(cb_g1, Wt); + cb_wait_front(cb_g1, Wt); // tmp = s+p + cb_pop_front(cb_xc, Wt); + cb_pop_front(cb_xhat, Wt); + + // silup = tmp - r -> cb_xc + cb_reserve_back(cb_xc, Wt); + sub_tiles_init(cb_g1, cb_prod); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + sub_tiles(cb_g1, cb_prod, wt, wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_xc, wt); + tile_regs_release(); + } + cb_push_back(cb_xc, Wt); + cb_wait_front(cb_xc, Wt); // silup = silu'(n) + cb_pop_front(cb_g1, Wt); + cb_pop_front(cb_prod, Wt); + + // g1 = g_out * silup -> cb_xhat + cb_reserve_back(cb_xhat, Wt); + mul_tiles_init(cb_gout, cb_xc); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles(cb_gout, cb_xc, wt, wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_xhat, wt); + tile_regs_release(); + } + cb_push_back(cb_xhat, Wt); + cb_wait_front(cb_xhat, Wt); + cb_pop_front(cb_xc, Wt); + cb_pop_front(cb_gout, Wt); + + // gy = g1 * gamma (bcast gamma row across the 32 edges) -> cb_gy + cb_reserve_back(cb_gy, Wt); + mul_bcast_rows_init_short(cb_xhat, cb_gamma); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles_bcast_rows(cb_xhat, cb_gamma, wt, wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_gy, wt); + tile_regs_release(); + } + cb_push_back(cb_gy, Wt); + cb_pop_front(cb_xhat, Wt); + + cb_wait_front(cb_gy, Wt); + cb_wait_front(cb_x, Wt); + + // --- mean_x = mean_w(x) into cb_s --- + reduce_mean(cb_x, cb_s); + + // --- xc = x - mean_x (bcast col0 of cb_s across cols) -> cb_xc --- + cb_wait_front(cb_s, 1); + cb_reserve_back(cb_xc, Wt); + sub_bcast_cols_init_short(cb_x, cb_s); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + sub_tiles_bcast_cols(cb_x, cb_s, wt, 0, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_xc, wt); + tile_regs_release(); + } + cb_push_back(cb_xc, Wt); + cb_pop_front(cb_s, 1); + + // --- var = mean_w(xc^2): build xc^2 into cb_prod (Wt tiles), reduce -> cb_s ; rstd = rsqrt(var+eps) --- + cb_wait_front(cb_xc, Wt); + cb_reserve_back(cb_prod, Wt); + mul_tiles_init(cb_xc, cb_xc); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles(cb_xc, cb_xc, wt, wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_prod, wt); + tile_regs_release(); + } + cb_push_back(cb_prod, Wt); + cb_wait_front(cb_prod, Wt); + // reduce cb_prod -> var(col0) in dst0, then add eps + rsqrt in-place -> cb_rstd + mm_init(cb_prod, cb_red, cb_rstd); + tile_regs_acquire(); + for (uint32_t wt = 0; wt < Wt; wt++) { + matmul_tiles(cb_prod, cb_red, wt, 0, 0); + } + binop_with_scalar_tile_init(); + add_unary_tile(0, eps_bits); + rsqrt_tile_init(); + rsqrt_tile(0); + tile_regs_commit(); + tile_regs_wait(); + cb_reserve_back(cb_rstd, 1); + pack_tile(0, cb_rstd); + cb_push_back(cb_rstd, 1); + tile_regs_release(); + cb_pop_front(cb_prod, Wt); + + // --- xhat = xc * rstd (bcast) -> cb_xhat --- + cb_wait_front(cb_rstd, 1); + cb_reserve_back(cb_xhat, Wt); + mul_bcast_cols_init_short(cb_xc, cb_rstd); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles_bcast_cols(cb_xc, cb_rstd, wt, 0, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_xhat, wt); + tile_regs_release(); + } + cb_push_back(cb_xhat, Wt); + cb_pop_front(cb_xc, Wt); + + // The dx assembly keeps only ONE scalar live in cb_s at a time (sharing cb_s for m1 AND m2 + // simultaneously corrupts the 2nd entry on multi-row cores). Each tile_regs_acquire runs a + // SINGLE bcast/eltwise op then packs (multi-op-per-acquire with init switches also corrupts). + // Scratch: cb_prod (a / gy*xhat product), cb_xc (b), cb_xhat reused for e. + + // --- m2 = mean_w(gy*xhat) -> cb_s ; b = xhat * m2 -> cb_xc --- + cb_wait_front(cb_xhat, Wt); + cb_reserve_back(cb_prod, Wt); + mul_tiles_init(cb_gy, cb_xhat); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles(cb_gy, cb_xhat, wt, wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_prod, wt); + tile_regs_release(); + } + cb_push_back(cb_prod, Wt); + cb_wait_front(cb_prod, Wt); + reduce_mean(cb_prod, cb_s); // m2 -> cb_s (front, sole entry) + cb_pop_front(cb_prod, Wt); + cb_wait_front(cb_s, 1); + cb_reserve_back(cb_xc, Wt); + mul_bcast_cols_init_short(cb_xhat, cb_s); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles_bcast_cols(cb_xhat, cb_s, wt, 0, 0); // b = xhat * m2 + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_xc, wt); + tile_regs_release(); + } + cb_push_back(cb_xc, Wt); + cb_pop_front(cb_s, 1); + cb_pop_front(cb_xhat, Wt); // xhat done; reuse cb_xhat for e below + + // --- m1 = mean_w(gy) -> cb_s ; a = gy - m1 -> cb_prod --- + reduce_mean(cb_gy, cb_s); + cb_wait_front(cb_s, 1); + cb_reserve_back(cb_prod, Wt); + sub_bcast_cols_init_short(cb_gy, cb_s); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + sub_tiles_bcast_cols(cb_gy, cb_s, wt, 0, 0); // a = gy - m1 + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_prod, wt); + tile_regs_release(); + } + cb_push_back(cb_prod, Wt); + cb_pop_front(cb_s, 1); + + // --- e = a - b -> cb_xhat (reused) --- + cb_wait_front(cb_prod, Wt); + cb_wait_front(cb_xc, Wt); + cb_reserve_back(cb_xhat, Wt); + sub_tiles_init(cb_prod, cb_xc); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + sub_tiles(cb_prod, cb_xc, wt, wt, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_xhat, wt); + tile_regs_release(); + } + cb_push_back(cb_xhat, Wt); + cb_pop_front(cb_prod, Wt); + cb_pop_front(cb_xc, Wt); + + // --- dx = e * rstd -> cb_dx --- + cb_wait_front(cb_xhat, Wt); + cb_reserve_back(cb_dx, Wt); + mul_bcast_cols_init_short(cb_xhat, cb_rstd); + for (uint32_t wt = 0; wt < Wt; wt++) { + tile_regs_acquire(); + mul_tiles_bcast_cols(cb_xhat, cb_rstd, wt, 0, 0); + tile_regs_commit(); + tile_regs_wait(); + pack_tile(0, cb_dx, wt); + tile_regs_release(); + } + cb_push_back(cb_dx, Wt); + cb_pop_front(cb_xhat, Wt); + + cb_pop_front(cb_rstd, 1); + cb_pop_front(cb_gy, Wt); + cb_pop_front(cb_x, Wt); + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/lnbw_reader.cpp b/custom_kernels/fused_rotate/device/kernels/lnbw_reader.cpp new file mode 100644 index 0000000..cfe46d2 --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/lnbw_reader.cpp @@ -0,0 +1,87 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Reader for fused_ln_bw (silu+gamma folded): loads the [32,32] reduction tile (col0 = 1/W) and the +// [1,W] LN affine scale (gamma) once (both resident), then streams g_out, x and n tile-rows (Wt tiles +// each) per edge tile-row. The compute kernel builds gy = g_out*silu'(n)*gamma internally. + +#include +#include "api/dataflow/dataflow_api.h" + +void kernel_main() { + constexpr uint32_t cb_gout = get_compile_time_arg_val(0); + constexpr uint32_t cb_x = get_compile_time_arg_val(1); + constexpr uint32_t cb_red = get_compile_time_arg_val(2); + constexpr uint32_t Wt = get_compile_time_arg_val(3); + constexpr uint32_t tile_bytes = get_compile_time_arg_val(4); + constexpr uint32_t cb_n = get_compile_time_arg_val(5); + constexpr uint32_t cb_gamma = get_compile_time_arg_val(6); + + constexpr auto gout_args = TensorAccessorArgs<7>(); + constexpr auto x_args = TensorAccessorArgs(); + constexpr auto red_args = TensorAccessorArgs(); + constexpr auto n_args = TensorAccessorArgs(); + constexpr auto gamma_args = TensorAccessorArgs(); + + uint32_t arg = 0; + const uint32_t gout_addr = get_arg_val(arg++); + const uint32_t x_addr = get_arg_val(arg++); + const uint32_t red_addr = get_arg_val(arg++); + const uint32_t start_row = get_arg_val(arg++); + const uint32_t num_rows = get_arg_val(arg++); + const uint32_t n_addr = get_arg_val(arg++); + const uint32_t gamma_addr = get_arg_val(arg++); + + const auto gout_gen = TensorAccessor(gout_args, gout_addr, tile_bytes); + const auto x_gen = TensorAccessor(x_args, x_addr, tile_bytes); + const auto red_gen = TensorAccessor(red_args, red_addr, tile_bytes); + const auto n_gen = TensorAccessor(n_args, n_addr, tile_bytes); + const auto gamma_gen = TensorAccessor(gamma_args, gamma_addr, tile_bytes); + + // reduction tile 0, loaded once (resident) + cb_reserve_back(cb_red, 1); + uint32_t rw = get_write_ptr(cb_red); + noc_async_read_tile(0, red_gen, rw); + // gamma row (Wt tiles, tile-row 0), loaded once (resident) + cb_reserve_back(cb_gamma, Wt); + uint32_t gaw = get_write_ptr(cb_gamma); + for (uint32_t t = 0; t < Wt; t++) { + noc_async_read_tile(t, gamma_gen, gaw); + gaw += tile_bytes; + } + noc_async_read_barrier(); + cb_push_back(cb_red, 1); + cb_push_back(cb_gamma, Wt); + + for (uint32_t r = 0; r < num_rows; r++) { + const uint32_t row = start_row + r; + const uint32_t base = row * Wt; + + cb_reserve_back(cb_gout, Wt); + uint32_t gw = get_write_ptr(cb_gout); + for (uint32_t t = 0; t < Wt; t++) { + noc_async_read_tile(base + t, gout_gen, gw); + gw += tile_bytes; + } + + cb_reserve_back(cb_x, Wt); + uint32_t xw = get_write_ptr(cb_x); + for (uint32_t t = 0; t < Wt; t++) { + noc_async_read_tile(base + t, x_gen, xw); + xw += tile_bytes; + } + + cb_reserve_back(cb_n, Wt); + uint32_t nw = get_write_ptr(cb_n); + for (uint32_t t = 0; t < Wt; t++) { + noc_async_read_tile(base + t, n_gen, nw); + nw += tile_bytes; + } + + noc_async_read_barrier(); + cb_push_back(cb_gout, Wt); + cb_push_back(cb_x, Wt); + cb_push_back(cb_n, Wt); + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/lnbw_writer.cpp b/custom_kernels/fused_rotate/device/kernels/lnbw_writer.cpp new file mode 100644 index 0000000..82c2513 --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/lnbw_writer.cpp @@ -0,0 +1,36 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +// +// Writer for fused_ln_bw: writes Wt output (dx) tiles per edge tile-row. + +#include +#include "api/dataflow/dataflow_api.h" + +void kernel_main() { + constexpr uint32_t cb_dx = get_compile_time_arg_val(0); + constexpr uint32_t Wt = get_compile_time_arg_val(1); + constexpr uint32_t tile_bytes = get_compile_time_arg_val(2); + + constexpr auto out_args = TensorAccessorArgs<3>(); + + uint32_t arg = 0; + const uint32_t out_addr = get_arg_val(arg++); + const uint32_t start_row = get_arg_val(arg++); + const uint32_t num_rows = get_arg_val(arg++); + + const auto out_gen = TensorAccessor(out_args, out_addr, tile_bytes); + + for (uint32_t r = 0; r < num_rows; r++) { + const uint32_t row = start_row + r; + cb_wait_front(cb_dx, Wt); + uint32_t rd = get_read_ptr(cb_dx); + const uint32_t base = row * Wt; + for (uint32_t t = 0; t < Wt; t++) { + noc_async_write_tile(base + t, out_gen, rd); + rd += tile_bytes; + } + noc_async_write_barrier(); + cb_pop_front(cb_dx, Wt); + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/reader.cpp b/custom_kernels/fused_rotate/device/kernels/reader.cpp new file mode 100644 index 0000000..c7a9ab5 --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/reader.cpp @@ -0,0 +1,52 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include +#include "api/dataflow/dataflow_api.h" + +void kernel_main() { + constexpr uint32_t cb_x = get_compile_time_arg_val(0); + constexpr uint32_t cb_coef = get_compile_time_arg_val(1); + constexpr uint32_t n_in_tiles = get_compile_time_arg_val(2); + constexpr uint32_t coef_tiles = get_compile_time_arg_val(3); + constexpr uint32_t tile_bytes = get_compile_time_arg_val(4); + + constexpr auto x_args = TensorAccessorArgs<5>(); + const auto coef_args = TensorAccessorArgs(); + + uint32_t arg = 0; + const uint32_t x_addr = get_arg_val(arg++); + const uint32_t coef_addr = get_arg_val(arg++); + const uint32_t start_row = get_arg_val(arg++); + const uint32_t num_rows = get_arg_val(arg++); + + const auto x_gen = TensorAccessor(x_args, x_addr, tile_bytes); + const auto coef_gen = TensorAccessor(coef_args, coef_addr, tile_bytes); + + for (uint32_t r = 0; r < num_rows; r++) { + const uint32_t row = start_row + r; + + // input feature blocks: tiles [row*n_in_tiles .. +n_in_tiles) + cb_reserve_back(cb_x, n_in_tiles); + uint32_t xw = get_write_ptr(cb_x); + const uint32_t x_base = row * n_in_tiles; + for (uint32_t t = 0; t < n_in_tiles; t++) { + noc_async_read_tile(x_base + t, x_gen, xw); + xw += tile_bytes; + } + + // per-nonzero coefficient tiles: tiles [row*coef_tiles .. +coef_tiles) + cb_reserve_back(cb_coef, coef_tiles); + uint32_t cw = get_write_ptr(cb_coef); + const uint32_t coef_base = row * coef_tiles; + for (uint32_t t = 0; t < coef_tiles; t++) { + noc_async_read_tile(coef_base + t, coef_gen, cw); + cw += tile_bytes; + } + + noc_async_read_barrier(); + cb_push_back(cb_x, n_in_tiles); + cb_push_back(cb_coef, coef_tiles); + } +} diff --git a/custom_kernels/fused_rotate/device/kernels/writer.cpp b/custom_kernels/fused_rotate/device/kernels/writer.cpp new file mode 100644 index 0000000..c2ca64b --- /dev/null +++ b/custom_kernels/fused_rotate/device/kernels/writer.cpp @@ -0,0 +1,34 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include +#include "api/dataflow/dataflow_api.h" + +void kernel_main() { + constexpr uint32_t cb_out = get_compile_time_arg_val(0); + constexpr uint32_t n_out_tiles = get_compile_time_arg_val(1); + constexpr uint32_t tile_bytes = get_compile_time_arg_val(2); + + constexpr auto out_args = TensorAccessorArgs<3>(); + + uint32_t arg = 0; + const uint32_t out_addr = get_arg_val(arg++); + const uint32_t start_row = get_arg_val(arg++); + const uint32_t num_rows = get_arg_val(arg++); + + const auto out_gen = TensorAccessor(out_args, out_addr, tile_bytes); + + for (uint32_t r = 0; r < num_rows; r++) { + const uint32_t row = start_row + r; + cb_wait_front(cb_out, n_out_tiles); + uint32_t rd = get_read_ptr(cb_out); + const uint32_t out_base = row * n_out_tiles; + for (uint32_t t = 0; t < n_out_tiles; t++) { + noc_async_write_tile(out_base + t, out_gen, rd); + rd += tile_bytes; + } + noc_async_write_barrier(); + cb_pop_front(cb_out, n_out_tiles); + } +} diff --git a/custom_kernels/fused_rotate/device/lnbw_device_operation.cpp b/custom_kernels/fused_rotate/device/lnbw_device_operation.cpp new file mode 100644 index 0000000..89b43c0 --- /dev/null +++ b/custom_kernels/fused_rotate/device/lnbw_device_operation.cpp @@ -0,0 +1,92 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "lnbw_device_operation.hpp" +#include "lnbw_program_factory.hpp" +#include "ttnn/device_operation.hpp" + +#include + +using namespace tt::constants; +using namespace tt::tt_metal; + +namespace ttnn::experimental::prim { + +void LnBwDeviceOperation::validate_on_program_cache_miss( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + const auto& gy = inputs.gy; + const auto& x = inputs.x; + const auto& red = inputs.red; + TT_FATAL( + gy.storage_type() == StorageType::DEVICE && x.storage_type() == StorageType::DEVICE && + red.storage_type() == StorageType::DEVICE, + "fused_ln_bw operands must be on device"); + TT_FATAL( + gy.layout() == Layout::TILE && x.layout() == Layout::TILE && red.layout() == Layout::TILE, + "fused_ln_bw requires TILE layout"); + TT_FATAL( + gy.dtype() == DataType::BFLOAT16 && x.dtype() == DataType::BFLOAT16 && red.dtype() == DataType::BFLOAT16, + "fused_ln_bw requires BFLOAT16 inputs"); + TT_FATAL(attrs.W % TILE_WIDTH == 0, "W ({}) must be a multiple of TILE_WIDTH", attrs.W); + const auto& gs = gy.padded_shape(); + const auto& xsh = x.padded_shape(); + TT_FATAL(gs[-1] == attrs.W, "gy last dim {} != W {}", gs[-1], attrs.W); + TT_FATAL(xsh[-1] == attrs.W, "x last dim {} != W {}", xsh[-1], attrs.W); + TT_FATAL(gs[-2] == xsh[-2], "gy and x must have the same number of rows"); + const auto& n = inputs.n; + const auto& gamma = inputs.gamma; + TT_FATAL( + n.storage_type() == StorageType::DEVICE && gamma.storage_type() == StorageType::DEVICE, + "fused_ln_bw n/gamma must be on device"); + TT_FATAL( + n.layout() == Layout::TILE && gamma.layout() == Layout::TILE, "fused_ln_bw n/gamma require TILE layout"); + TT_FATAL( + n.dtype() == DataType::BFLOAT16 && gamma.dtype() == DataType::BFLOAT16, + "fused_ln_bw n/gamma require BFLOAT16"); + TT_FATAL(n.padded_shape()[-1] == attrs.W, "n last dim {} != W {}", n.padded_shape()[-1], attrs.W); + TT_FATAL(n.padded_shape()[-2] == gs[-2], "n and gy must have the same number of rows"); + TT_FATAL(gamma.padded_shape()[-1] == attrs.W, "gamma last dim {} != W {}", gamma.padded_shape()[-1], attrs.W); +} + +LnBwDeviceOperation::spec_return_value_t LnBwDeviceOperation::compute_output_specs( + const operation_attributes_t&, const tensor_args_t& inputs) { + const auto& gy = inputs.gy; + return TensorSpec( + gy.logical_shape(), TensorLayout(gy.dtype(), PageConfig(Layout::TILE), gy.memory_config())); +} + +LnBwDeviceOperation::tensor_return_value_t LnBwDeviceOperation::create_output_tensors( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + return create_device_tensor(compute_output_specs(attrs, inputs), inputs.gy.device()); +} + +ttsl::hash::hash_t LnBwDeviceOperation::compute_program_hash( + const operation_attributes_t& attrs, const tensor_args_t& inputs) { + return tt::tt_metal::operation::hash_operation( + attrs.W, + attrs.eps_bits, + inputs.gy.dtype(), + inputs.gy.memory_config(), + inputs.gy.padded_shape()); +} + +} // namespace ttnn::experimental::prim + +namespace ttnn::prim { + +Tensor fused_ln_bw( + const Tensor& gy, + const Tensor& x, + const Tensor& red, + const Tensor& n, + const Tensor& gamma, + uint32_t W, + uint32_t eps_bits) { + using OperationType = ttnn::experimental::prim::LnBwDeviceOperation; + auto attrs = OperationType::operation_attributes_t{.W = W, .eps_bits = eps_bits}; + return ttnn::device_operation::launch( + attrs, OperationType::tensor_args_t{.gy = gy, .x = x, .red = red, .n = n, .gamma = gamma}); +} + +} // namespace ttnn::prim diff --git a/custom_kernels/fused_rotate/device/lnbw_device_operation.hpp b/custom_kernels/fused_rotate/device/lnbw_device_operation.hpp new file mode 100644 index 0000000..344fc74 --- /dev/null +++ b/custom_kernels/fused_rotate/device/lnbw_device_operation.hpp @@ -0,0 +1,37 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "ttnn/tensor/tensor.hpp" +#include "lnbw_program_factory.hpp" +#include "lnbw_device_operation_types.hpp" + +namespace ttnn::experimental::prim { + +struct LnBwDeviceOperation { + using operation_attributes_t = LnBwParams; + using tensor_args_t = LnBwInputs; + using spec_return_value_t = TensorSpec; + using tensor_return_value_t = Tensor; + using program_factory_t = std::variant; + + static void validate_on_program_cache_miss(const operation_attributes_t&, const tensor_args_t&); + static spec_return_value_t compute_output_specs(const operation_attributes_t&, const tensor_args_t&); + static tensor_return_value_t create_output_tensors(const operation_attributes_t&, const tensor_args_t&); + static ttsl::hash::hash_t compute_program_hash(const operation_attributes_t&, const tensor_args_t&); +}; + +} // namespace ttnn::experimental::prim + +namespace ttnn::prim { +Tensor fused_ln_bw( + const Tensor& gy, + const Tensor& x, + const Tensor& red, + const Tensor& n, + const Tensor& gamma, + uint32_t W, + uint32_t eps_bits); +} // namespace ttnn::prim diff --git a/custom_kernels/fused_rotate/device/lnbw_device_operation_types.hpp b/custom_kernels/fused_rotate/device/lnbw_device_operation_types.hpp new file mode 100644 index 0000000..16dba6c --- /dev/null +++ b/custom_kernels/fused_rotate/device/lnbw_device_operation_types.hpp @@ -0,0 +1,29 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include + +#include "ttnn/tensor/tensor.hpp" + +namespace ttnn::experimental::prim { + +// Fused LayerNorm backward (grad wrt LN input). gy = g_out*gamma (affine scale folded on host). +// dx = rstd * (gy - mean_w(gy) - xhat*mean_w(gy*xhat)), xhat=(x-mean_w(x))*rstd, +// rstd = rsqrt(mean_w((x-mean_w(x))^2) + eps). W must be a multiple of TILE_WIDTH. +struct LnBwParams { + uint32_t W; // channels per row (multiple of 32) + uint32_t eps_bits; // fp32 bits of the LN epsilon +}; + +struct LnBwInputs { + Tensor gy; // [E, W] TILE bf16 (g_out, matmul result pre-silu-bw) + Tensor x; // [E, W] TILE bf16 (cached forward LN input) + Tensor red; // [32, 32] TILE bf16, column 0 = 1/W (reduction selector) + Tensor n; // [E, W] TILE bf16 (pre-silu activation = LN output; kernel applies silu'(n)) + Tensor gamma; // [1, W] TILE bf16 (LN affine scale; folded in via row-broadcast) +}; + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/lnbw_program_factory.cpp b/custom_kernels/fused_rotate/device/lnbw_program_factory.cpp new file mode 100644 index 0000000..b09e5af --- /dev/null +++ b/custom_kernels/fused_rotate/device/lnbw_program_factory.cpp @@ -0,0 +1,147 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "lnbw_program_factory.hpp" +#include "lnbw_device_operation_types.hpp" + +#include +#include +#include +#include + +using namespace tt::constants; +using namespace tt::tt_metal; + +namespace ttnn::experimental::prim { + +static const char* kLnBwKernelDir = "ttnn/cpp/ttnn/operations/experimental/fused_rotate/device/kernels/"; + +LnBwProgramFactory::cached_program_t LnBwProgramFactory::create( + const LnBwParams& attrs, const LnBwInputs& inputs, Tensor& output) { + Program program{}; + + const auto& gy = inputs.gy; // g_out (matmul, pre-silu-bw) + const auto& x = inputs.x; + const auto& red = inputs.red; + const auto& n = inputs.n; + const auto& gamma = inputs.gamma; + + const uint32_t Wt = attrs.W / TILE_WIDTH; + const uint32_t Et = gy.padded_shape()[-2] / TILE_HEIGHT; + + tt::DataFormat data_format = datatype_to_dataformat_converter(gy.dtype()); + const uint32_t tile_bytes = tile_size(data_format); + + auto* device = gy.device(); + CoreCoord grid = device->compute_with_storage_grid_size(); + auto [num_cores, all_cores, core_group_1, core_group_2, rows_per_core_1, rows_per_core_2] = + tt::tt_metal::split_work_to_cores(grid, Et); + + auto make_cb = [&](uint32_t cb_index, uint32_t num_tiles) { + CircularBufferConfig cfg = + CircularBufferConfig(num_tiles * tile_bytes, {{cb_index, data_format}}).set_page_size(cb_index, tile_bytes); + CreateCircularBuffer(program, all_cores, cfg); + }; + constexpr uint32_t cb_gout = tt::CBIndex::c_0; // g_out (matmul) + constexpr uint32_t cb_x = tt::CBIndex::c_1; + constexpr uint32_t cb_red = tt::CBIndex::c_2; + constexpr uint32_t cb_xc = tt::CBIndex::c_3; + constexpr uint32_t cb_xhat = tt::CBIndex::c_4; + constexpr uint32_t cb_prod = tt::CBIndex::c_5; + constexpr uint32_t cb_s = tt::CBIndex::c_6; + constexpr uint32_t cb_rstd = tt::CBIndex::c_7; + constexpr uint32_t cb_n = tt::CBIndex::c_8; // pre-silu activation + constexpr uint32_t cb_gamma = tt::CBIndex::c_9; // LN affine scale [1,W], resident + constexpr uint32_t cb_gy = tt::CBIndex::c_10; // internal: g_out*silu'(n)*gamma + constexpr uint32_t cb_g1 = tt::CBIndex::c_11; // internal preamble scratch + constexpr uint32_t cb_dx = tt::CBIndex::c_16; + make_cb(cb_gout, 2 * Wt); + make_cb(cb_x, 2 * Wt); + make_cb(cb_red, 2); + make_cb(cb_xc, 2 * Wt); + make_cb(cb_xhat, 2 * Wt); + make_cb(cb_prod, 2 * Wt); + make_cb(cb_s, 4); + make_cb(cb_rstd, 2); + make_cb(cb_n, 2 * Wt); + make_cb(cb_gamma, Wt); + make_cb(cb_gy, 2 * Wt); + make_cb(cb_g1, 2 * Wt); + make_cb(cb_dx, 2 * Wt); + + // ---- reader ---- + std::vector reader_ct = {cb_gout, cb_x, cb_red, Wt, tile_bytes, cb_n, cb_gamma}; + TensorAccessorArgs(*gy.buffer()).append_to(reader_ct); + TensorAccessorArgs(*x.buffer()).append_to(reader_ct); + TensorAccessorArgs(*red.buffer()).append_to(reader_ct); + TensorAccessorArgs(*n.buffer()).append_to(reader_ct); + TensorAccessorArgs(*gamma.buffer()).append_to(reader_ct); + KernelHandle reader_id = CreateKernel( + program, std::string(kLnBwKernelDir) + "lnbw_reader.cpp", all_cores, ReaderDataMovementConfig(reader_ct)); + + // ---- writer ---- + std::vector writer_ct = {cb_dx, Wt, tile_bytes}; + TensorAccessorArgs(*output.buffer()).append_to(writer_ct); + KernelHandle writer_id = CreateKernel( + program, std::string(kLnBwKernelDir) + "lnbw_writer.cpp", all_cores, WriterDataMovementConfig(writer_ct)); + + // ---- compute ---- + std::vector compute_ct = { + cb_gout, cb_x, cb_red, cb_xc, cb_xhat, cb_prod, cb_s, cb_rstd, cb_dx, Wt, attrs.eps_bits, + cb_n, cb_gamma, cb_gy, cb_g1}; + KernelHandle compute_id = CreateKernel( + program, + std::string(kLnBwKernelDir) + "lnbw_compute.cpp", + all_cores, + ComputeConfig{ + .math_fidelity = MathFidelity::HiFi4, .fp32_dest_acc_en = true, .compile_args = compute_ct}); + + auto* gy_buf = gy.buffer(); + auto* x_buf = x.buffer(); + auto* red_buf = red.buffer(); + auto* n_buf = n.buffer(); + auto* gamma_buf = gamma.buffer(); + auto* out_buf = output.buffer(); + auto cores = corerange_to_cores(all_cores, num_cores, true); + + uint32_t row_offset = 0; + for (const auto& core : cores) { + uint32_t rows = core_group_1.contains(core) ? rows_per_core_1 : rows_per_core_2; + SetRuntimeArgs( + program, reader_id, core, + {gy_buf->address(), x_buf->address(), red_buf->address(), row_offset, rows, + n_buf->address(), gamma_buf->address()}); + SetRuntimeArgs(program, writer_id, core, {out_buf->address(), row_offset, rows}); + SetRuntimeArgs(program, compute_id, core, {rows}); + row_offset += rows; + } + + return cached_program_t{std::move(program), {reader_id, writer_id, compute_id, cores}}; +} + +void LnBwProgramFactory::override_runtime_arguments( + cached_program_t& cached_program, const LnBwParams&, const LnBwInputs& inputs, Tensor& output) { + auto& program = cached_program.program; + const auto& cores = cached_program.shared_variables.cores; + const auto reader_id = cached_program.shared_variables.reader_kernel_id; + const auto writer_id = cached_program.shared_variables.writer_kernel_id; + auto* gy_buf = inputs.gy.buffer(); + auto* x_buf = inputs.x.buffer(); + auto* red_buf = inputs.red.buffer(); + auto* n_buf = inputs.n.buffer(); + auto* gamma_buf = inputs.gamma.buffer(); + auto* out_buf = output.buffer(); + for (const auto& core : cores) { + auto& ra = GetRuntimeArgs(program, reader_id, core); + ra[0] = gy_buf->address(); + ra[1] = x_buf->address(); + ra[2] = red_buf->address(); + ra[5] = n_buf->address(); + ra[6] = gamma_buf->address(); + auto& wa = GetRuntimeArgs(program, writer_id, core); + wa[0] = out_buf->address(); + } +} + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/device/lnbw_program_factory.hpp b/custom_kernels/fused_rotate/device/lnbw_program_factory.hpp new file mode 100644 index 0000000..46ba027 --- /dev/null +++ b/custom_kernels/fused_rotate/device/lnbw_program_factory.hpp @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "lnbw_device_operation_types.hpp" +#include "ttnn/device_operation.hpp" + +namespace ttnn::experimental::prim { + +struct LnBwSharedVariables { + tt::tt_metal::KernelHandle reader_kernel_id = 0; + tt::tt_metal::KernelHandle writer_kernel_id = 0; + tt::tt_metal::KernelHandle compute_kernel_id = 0; + std::vector cores; +}; + +struct LnBwProgramFactory { + using shared_variables_t = LnBwSharedVariables; + using cached_program_t = ttnn::device_operation::CachedProgram; + + static cached_program_t create(const LnBwParams& operation_attributes, const LnBwInputs& inputs, Tensor& output); + + static void override_runtime_arguments( + cached_program_t& cached_program, + const LnBwParams& operation_attributes, + const LnBwInputs& inputs, + Tensor& output); +}; + +} // namespace ttnn::experimental::prim diff --git a/custom_kernels/fused_rotate/fused_gate.cpp b/custom_kernels/fused_rotate/fused_gate.cpp new file mode 100644 index 0000000..b8aca15 --- /dev/null +++ b/custom_kernels/fused_rotate/fused_gate.cpp @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "device/gate_device_operation.hpp" +#include "ttnn/operations/experimental/fused_rotate/fused_gate.hpp" + +namespace ttnn::operations::experimental { + +ttnn::Tensor fused_gate( + const ttnn::Tensor& a, + const ttnn::Tensor& gate, + const ttnn::Tensor& b, + uint32_t Wt, + uint32_t Gt, + uint32_t Ht, + uint32_t mode) { + return ttnn::prim::fused_gate(a, gate, b, Wt, Gt, Ht, mode); +} + +} // namespace ttnn::operations::experimental diff --git a/custom_kernels/fused_rotate/fused_gate.hpp b/custom_kernels/fused_rotate/fused_gate.hpp new file mode 100644 index 0000000..95a32af --- /dev/null +++ b/custom_kernels/fused_rotate/fused_gate.hpp @@ -0,0 +1,23 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include + +#include "ttnn/tensor/tensor.hpp" +#include "ttnn/types.hpp" + +namespace ttnn::operations::experimental { + +ttnn::Tensor fused_gate( + const ttnn::Tensor& a, + const ttnn::Tensor& gate, + const ttnn::Tensor& b, + uint32_t Wt, + uint32_t Gt, + uint32_t Ht, + uint32_t mode); + +} // namespace ttnn::operations::experimental diff --git a/custom_kernels/fused_rotate/fused_ln_bw.cpp b/custom_kernels/fused_rotate/fused_ln_bw.cpp new file mode 100644 index 0000000..eda78f7 --- /dev/null +++ b/custom_kernels/fused_rotate/fused_ln_bw.cpp @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "device/lnbw_device_operation.hpp" +#include "ttnn/operations/experimental/fused_rotate/fused_ln_bw.hpp" + +namespace ttnn::operations::experimental { + +ttnn::Tensor fused_ln_bw( + const ttnn::Tensor& gy, + const ttnn::Tensor& x, + const ttnn::Tensor& red, + const ttnn::Tensor& n, + const ttnn::Tensor& gamma, + uint32_t W, + uint32_t eps_bits) { + return ttnn::prim::fused_ln_bw(gy, x, red, n, gamma, W, eps_bits); +} + +} // namespace ttnn::operations::experimental diff --git a/custom_kernels/fused_rotate/fused_ln_bw.hpp b/custom_kernels/fused_rotate/fused_ln_bw.hpp new file mode 100644 index 0000000..7b409c6 --- /dev/null +++ b/custom_kernels/fused_rotate/fused_ln_bw.hpp @@ -0,0 +1,23 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include + +#include "ttnn/tensor/tensor.hpp" +#include "ttnn/types.hpp" + +namespace ttnn::operations::experimental { + +ttnn::Tensor fused_ln_bw( + const ttnn::Tensor& gy, + const ttnn::Tensor& x, + const ttnn::Tensor& red, + const ttnn::Tensor& n, + const ttnn::Tensor& gamma, + uint32_t W, + uint32_t eps_bits); + +} // namespace ttnn::operations::experimental diff --git a/custom_kernels/fused_rotate/fused_rotate.cpp b/custom_kernels/fused_rotate/fused_rotate.cpp new file mode 100644 index 0000000..fb0ece6 --- /dev/null +++ b/custom_kernels/fused_rotate/fused_rotate.cpp @@ -0,0 +1,22 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "device/fused_rotate_device_operation.hpp" +#include "ttnn/operations/experimental/fused_rotate/fused_rotate.hpp" + +namespace ttnn::operations::experimental { + +ttnn::Tensor fused_rotate( + const ttnn::Tensor& x_flat, + const ttnn::Tensor& coef_exp, + uint32_t n_in, + uint32_t n_out, + uint32_t W, + const std::vector& deg, + const std::vector& ks, + const std::vector& js) { + return ttnn::prim::fused_rotate(x_flat, coef_exp, n_in, n_out, W, deg, ks, js); +} + +} // namespace ttnn::operations::experimental diff --git a/custom_kernels/fused_rotate/fused_rotate.hpp b/custom_kernels/fused_rotate/fused_rotate.hpp new file mode 100644 index 0000000..a6763a0 --- /dev/null +++ b/custom_kernels/fused_rotate/fused_rotate.hpp @@ -0,0 +1,25 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include + +#include "ttnn/tensor/tensor.hpp" +#include "ttnn/types.hpp" + +namespace ttnn::operations::experimental { + +ttnn::Tensor fused_rotate( + const ttnn::Tensor& x_flat, + const ttnn::Tensor& coef_exp, + uint32_t n_in, + uint32_t n_out, + uint32_t W, + const std::vector& deg, + const std::vector& ks, + const std::vector& js); + +} // namespace ttnn::operations::experimental diff --git a/custom_kernels/fused_rotate/fused_rotate_gc.cpp b/custom_kernels/fused_rotate/fused_rotate_gc.cpp new file mode 100644 index 0000000..c00219c --- /dev/null +++ b/custom_kernels/fused_rotate/fused_rotate_gc.cpp @@ -0,0 +1,22 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "device/gc_device_operation.hpp" +#include "ttnn/operations/experimental/fused_rotate/fused_rotate_gc.hpp" + +namespace ttnn::operations::experimental { + +ttnn::Tensor fused_rotate_gc( + const ttnn::Tensor& gout, + const ttnn::Tensor& xin, + const ttnn::Tensor& sel, + uint32_t n_out, + uint32_t n_in, + uint32_t W, + const std::vector& is_, + const std::vector& js) { + return ttnn::prim::fused_rotate_gc(gout, xin, sel, n_out, n_in, W, is_, js); +} + +} // namespace ttnn::operations::experimental diff --git a/custom_kernels/fused_rotate/fused_rotate_gc.hpp b/custom_kernels/fused_rotate/fused_rotate_gc.hpp new file mode 100644 index 0000000..83e261d --- /dev/null +++ b/custom_kernels/fused_rotate/fused_rotate_gc.hpp @@ -0,0 +1,25 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include + +#include "ttnn/tensor/tensor.hpp" +#include "ttnn/types.hpp" + +namespace ttnn::operations::experimental { + +ttnn::Tensor fused_rotate_gc( + const ttnn::Tensor& gout, + const ttnn::Tensor& xin, + const ttnn::Tensor& sel, + uint32_t n_out, + uint32_t n_in, + uint32_t W, + const std::vector& is_, + const std::vector& js); + +} // namespace ttnn::operations::experimental diff --git a/custom_kernels/fused_rotate/fused_rotate_nanobind.cpp b/custom_kernels/fused_rotate/fused_rotate_nanobind.cpp new file mode 100644 index 0000000..cd7c147 --- /dev/null +++ b/custom_kernels/fused_rotate/fused_rotate_nanobind.cpp @@ -0,0 +1,98 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "fused_rotate_nanobind.hpp" + +#include +#include + +#include "ttnn/operations/experimental/fused_rotate/fused_rotate.hpp" +#include "ttnn/operations/experimental/fused_rotate/fused_rotate_gc.hpp" +#include "ttnn/operations/experimental/fused_rotate/fused_ln_bw.hpp" +#include "ttnn/operations/experimental/fused_rotate/fused_gate.hpp" + +namespace ttnn::operations::experimental::fr_detail { + +void bind_fused_rotate(nb::module_& mod) { + mod.def( + "fused_ln_bw", + [](const ttnn::Tensor& gy, + const ttnn::Tensor& x, + const ttnn::Tensor& red, + const ttnn::Tensor& n, + const ttnn::Tensor& gamma, + uint32_t W, + uint32_t eps_bits) { + return ttnn::operations::experimental::fused_ln_bw(gy, x, red, n, gamma, W, eps_bits); + }, + nb::arg("gy").noconvert(), + nb::arg("x").noconvert(), + nb::arg("red").noconvert(), + nb::arg("n").noconvert(), + nb::arg("gamma").noconvert(), + nb::arg("W"), + nb::arg("eps_bits"), + "Fused SiLU-bw + LayerNorm-bw: gy=g_out*silu'(n)*gamma in-kernel, then dx."); + mod.def( + "fused_gate", + [](const ttnn::Tensor& a, + const ttnn::Tensor& gate, + const ttnn::Tensor& b, + uint32_t Wt, + uint32_t Gt, + uint32_t Ht, + uint32_t mode) { return ttnn::operations::experimental::fused_gate(a, gate, b, Wt, Gt, Ht, mode); }, + nb::arg("a").noconvert(), + nb::arg("gate").noconvert(), + nb::arg("b").noconvert(), + nb::arg("Wt"), + nb::arg("Gt"), + nb::arg("Ht"), + nb::arg("mode"), + "Fused SO(3) gate: out=[silu(a[:H]) | a[H:]*gate] (mode0 fwd) or [a[:H]*silu'(b) | a[H:]*gate] (mode1 bw)."); + mod.def( + "fused_rotate_gc", + [](const ttnn::Tensor& gout, + const ttnn::Tensor& xin, + const ttnn::Tensor& sel, + uint32_t n_out, + uint32_t n_in, + uint32_t W, + const std::vector& is_, + const std::vector& js) { + return ttnn::operations::experimental::fused_rotate_gc(gout, xin, sel, n_out, n_in, W, is_, js); + }, + nb::arg("gout").noconvert(), + nb::arg("xin").noconvert(), + nb::arg("sel").noconvert(), + nb::arg("n_out"), + nb::arg("n_in"), + nb::arg("W"), + nb::arg("is_"), + nb::arg("js"), + "Fused per-edge coefficient adjoint (rotate_bw dE/dcoef) mul+reduce+place in one kernel."); + mod.def( + "fused_rotate", + [](const ttnn::Tensor& x_flat, + const ttnn::Tensor& coef_exp, + uint32_t n_in, + uint32_t n_out, + uint32_t W, + const std::vector& deg, + const std::vector& ks, + const std::vector& js) { + return ttnn::operations::experimental::fused_rotate(x_flat, coef_exp, n_in, n_out, W, deg, ks, js); + }, + nb::arg("x_flat").noconvert(), + nb::arg("coef_exp").noconvert(), + nb::arg("n_in"), + nb::arg("n_out"), + nb::arg("W"), + nb::arg("deg"), + nb::arg("ks"), + nb::arg("js"), + "Fused per-edge sparse Wigner rotation (all nnz MACs in one kernel launch)."); +} + +} // namespace ttnn::operations::experimental::fr_detail diff --git a/custom_kernels/fused_rotate/fused_rotate_nanobind.hpp b/custom_kernels/fused_rotate/fused_rotate_nanobind.hpp new file mode 100644 index 0000000..c3a5ad1 --- /dev/null +++ b/custom_kernels/fused_rotate/fused_rotate_nanobind.hpp @@ -0,0 +1,15 @@ +// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include "ttnn-nanobind/nanobind_fwd.hpp" + +namespace ttnn::operations::experimental::fr_detail { + +namespace nb = nanobind; + +void bind_fused_rotate(nb::module_& mod); + +} // namespace ttnn::operations::experimental::fr_detail diff --git a/custom_kernels/fused_rotate/sources.cmake b/custom_kernels/fused_rotate/sources.cmake new file mode 100644 index 0000000..29aac6b --- /dev/null +++ b/custom_kernels/fused_rotate/sources.cmake @@ -0,0 +1,18 @@ +# Source files for ttnn_op_experimental_fused_rotate. + +set(TTNN_OP_EXPERIMENTAL_FUSED_ROTATE_API_HEADERS fused_rotate.hpp fused_rotate_gc.hpp fused_ln_bw.hpp fused_gate.hpp) + +set(TTNN_OP_EXPERIMENTAL_FUSED_ROTATE_SRCS + device/fused_rotate_device_operation.cpp + device/fused_rotate_program_factory.cpp + fused_rotate.cpp + device/gc_device_operation.cpp + device/gc_program_factory.cpp + fused_rotate_gc.cpp + device/lnbw_device_operation.cpp + device/lnbw_program_factory.cpp + fused_ln_bw.cpp + device/gate_device_operation.cpp + device/gate_program_factory.cpp + fused_gate.cpp +) diff --git a/docs/uma-s-1p2-validation.md b/docs/uma-s-1p2-validation.md new file mode 100644 index 0000000..91ea9c2 --- /dev/null +++ b/docs/uma-s-1p2-validation.md @@ -0,0 +1,130 @@ +# UMA-S-1.2 on Tenstorrent — validation & performance + +This documents the parity and performance of `uma-s-1.2` running on Tenstorrent through TT-Atom, +measured against Meta's fairchem on CPU using the released `uma-s-1p2.pt` checkpoint. It backs the +one-line claim in the README; see [`../README.md`](../README.md) for how to run the model. + +`uma-s-1.2` differs from `uma-s-1` by fairchem's **charge-balanced channels**: the `l=0` charge +channels are re-balanced to the system charge (a self-adjoint per-system mean-subtraction, plus a +`charge/natoms` target) after every block. TT-Atom ports this forward and backward through the +single-system, batched, and traced paths. Without it, the force PCC on the released checkpoint +collapses to ~0.83; with it, parity is restored. + +## Method + +- **Reference:** fairchem-core on CPU — the released `uma-s-1p2.pt`, MoLE-merged per composition, + through `FAIRChemCalculator`. +- **Device:** TT-Atom on a Tenstorrent Wormhole card, `ttnn` 0.68.0 (bf16 backbone with fp32 + accumulation and an fp32 energy head). +- **Per system:** a static set (equilibrium + seeded rattles) and a seeded NVE trajectory + (velocity-Verlet). The *identical* geometries are fed to both engines. +- **Metrics:** pooled force PCC (over all geometries), median energy discrepancy per atom + (meV/atom), stress PCC (periodic), and NVE total-energy drift for each engine. +- **Pass** (per the released `uma-s-1` tolerances) = static **and** trajectory force PCC > 0.99, + static **and** trajectory median energy < 5 meV/atom, and — when periodic — stress PCC > 0.99. + Chemical accuracy is ~43 meV/atom (1 kcal/mol), so the 5 meV/atom bar is deliberately strict. + +**Why pooled force PCC, not per-geometry.** At a symmetric equilibrium or a perfect crystal +lattice the forces are ~0 by symmetry, so a per-geometry PCC just correlates numerical noise and +reads low despite a tiny absolute error. Pooling over the rattled geometries (where `|F|` is +meaningful) is the honest measure; the artifact is called out explicitly below rather than hidden. + +## Coverage + +**757 systems**, 2–99 atoms, ~63 elements, molecular + **491 periodic**. Tasks `omol`, `omat`, +`oc20`, `odac`, `omc`. The screen spans 175 `oc20` catalysis surfaces (metal facets × CO/H/O/OH), +40 metal nanoparticles, the 22 S22 non-covalent complexes, open-shell radicals, charged +transition-metal complexes, 120 equation-of-state volume-strain curves, and a broad spread of +elemental / binary / perovskite crystals. + +## Forces + +Forces reproduce fairchem across the screen with **no genuine errors**: **741 / 757** systems keep +pooled force PCC **> 0.999** (754/757 > 0.99), and the three below 0.99 are all the **zero-force +perfect-lattice artifact** — PCC correlates noise where `|F|≈0` by symmetry, while the absolute error +is tiny (`el_Pb`, `el_Pb_exp`, `s_Pb100_clean`, all traj Fmae < 1.6 meV/Å). Wherever forces are +non-negligible, PCC is ≈ 1.0; the showcase systems are all exactly 1.0000. + +The re-screen surfaced one real force regression from `master`'s on-device-backward perf work — +`el_Sn_cmp` (2-atom compressed β-Sn, traj PCC 0.73 / 230 meV/Å): the radial-MLP backward computes its +LayerNorm VJP in **bf16** by default (unlike the fp32-accurate SH-norm/gate backwards), mis-directing +the gradient on that out-of-distribution geometry while the energy stays correct. **This PR fixes it** +by running the small radial backward in fp32 — verified against an fp64 `mirror.radial_mlp` oracle +(device `rad.bw` PCC 0.35 → 1.0); `el_Sn_cmp` is now **1.0000**, at ~4% eager cost. (`master`'s opt-in +`fused_ln_bw` kernel is bf16-only and off by default, so the default path carried the bug.) + +## Energy + +**691 / 757 pass** the strict bar (force PCC > 0.99, energy < 5 meV/atom, stress PCC > 0.99); passing +systems sit at a median of **1.33 meV/atom**. The misses are the **bf16 backbone energy floor**: the +error scales as `|E_raw| / N`, so it hits the highest-`|raw|`, smallest systems hardest — heavy-metal +diatomics / crystals (`el_Ta_exp` 11, `m_CH` 10, `el_Mo` 9, `el_Tc_exp` 8 meV/atom) and dense +ionic/oxide cells (`x_CaO_rock`). Their **forces are unaffected** (PCC ≈ 1.0) and **every system stays +under chemical accuracy** (≈ 43 meV/atom). + +This is a **fundamental bf16-forward precision limit, not a bug** (identical inputs give identical +outputs to within it): the reference-subtracted per-node energy rounds in the bf16 backbone. The fp32 +energy head removes the reduction-order term; the residual is the bf16-forward floor. Net, on the +current `master` (custom fused kernels + on-device backward) with these fixes, the pass rate is +**691/757 — above the 672/757 the earlier backbone reached** (the perf work's numerics plus the fp32 +head slightly *lower* the floor overall, even as it moves the exact set of borderline systems around). + +## Stress (periodic) + +Median stress PCC **0.99996**; **488 / 491 ≥ 0.99**. The few lowest are near-zero-stress soft metals / +centrosymmetric cells whose equilibrium stress is itself ~0 (PCC-of-noise). + +## Showcase systems + +Four end-to-end structures spanning catalysis, nanoparticles, and a charged solvated +transition-metal chelate — all reproduce fairchem forces exactly: + +| system | task | atoms | charge / spin | force PCC | energy | +|--------|------|------:|:-------------:|----------:|-------:| +| Pt(111) + CO + O (CO oxidation) | oc20 | 39 | 0 / 1 | **1.0000** | 0.31 meV/atom | +| Cu₅₅ nanoparticle + CO | omat | 57 | 0 / 1 | **1.0000** | 0.91 meV/atom | +| [Cu(EDTA)]²⁻ chelate | omol | 33 | −2 / 2 | **1.0000** | 1.18 meV/atom | +| [Cu(EDTA)]²⁻ + 22 H₂O | omol | 99 | −2 / 2 | **1.0000** | 0.88 meV/atom | + +## Performance — CPU vs Tenstorrent + +Energy + forces, milliseconds per call, MgO supercells, measured on the current `master` (custom fused +kernels + on-device backward) with this PR's fixes. TT (Wormhole, `ttnn` 0.68.0) eager and traced vs +fairchem on a 16-thread CPU; each point in its own process: + +| MgO cell | atoms | TT eager | TT trace | CPU (16-thread) | speedup (CPU / TT-best) | +|----------|------:|---------:|---------:|----------------:|------------------------:| +| 1×1×1 | 8 | 103 | 75 | 138 | 1.8× | +| 2×2×2 | 64 | 317 | 319 | 860 | 2.7× | +| 3×3×3 | 216 | 1059 | 973 | 5326 | **5.5×** | +| 4×4×4 | 512 | 2794 | OOM | — | — | + +Molecular water boxes (eager / trace ms): ethanol 90 / 47, H₂O×8 142 / 135, H₂O×27 (81 at) 358 / 347, +H₂O×64 (192 at) 807 / 764. + +TT overtakes the 16-core CPU by ~64 atoms and reaches **5.5× at 216 atoms** (trace) — the margin grows +with size. `master`'s kernel work makes this run **~2.2× faster than the pre-perf-pass backbone** +(MgO-216 eager 2307 → 1059 ms) and, notably, restores the trace path at scale: trace now captures +through 216-atom periodic / 192-atom molecular cells (it was capture-bound at large N on the earlier +on-device-backward), and eager reaches the 512-atom cell (2794 ms; only the 512 *trace* still exceeds +single-card DRAM). The trace loop (the MD fast path) is ~1.9× the eager call on small molecules (ethanol +90 → 47 ms), bit-identical forces. The fp32 radial-backward fix here adds ~4%. Below ~10 atoms the CPU +wins — TT's advantage is at real system sizes. + +## Reproduce + +- **A/B screen:** fairchem-CPU references vs TT-Atom on the card, over the system list above; + pooled force PCC, per-atom energy, stress PCC, and NVE drift. +- **Gated device parity test (in-repo):** [`../tests/test_realweight_uma_s_1p2.py`](../tests/test_realweight_uma_s_1p2.py) + checks device energy + analytic forces against the fairchem oracle for a real `uma-s-1.2` golden. + It auto-skips without the (gated) checkpoint. To run it: + + ```bash + HF_HUB_OFFLINE=1 /bin/python tests/gen_golden_real.py \ + --system molecule --task omol --ckpt uma-s-1p2 \ + --out ~/.ttatom_run/goldens_real/ethanol_omol_uma_s_1p2.npz + TT_VISIBLE_DEVICES=0 /bin/python -m pytest tests/test_realweight_uma_s_1p2.py -q + ``` + +Numbers are from `ttnn` 0.68.0; op numerics can shift slightly between `ttnn` versions, so confirm +parity on the version you actually run. diff --git a/examples/batch.py b/examples/batch.py new file mode 100644 index 0000000..20f013f --- /dev/null +++ b/examples/batch.py @@ -0,0 +1,49 @@ +"""Multi-card batch throughput on Tenstorrent via TT-Atom's fan-out. + +The evaluation of one system is independent of every other, so throughput scales by running one +worker process per Tenstorrent card (each pinned to its own device with the model resident) while +the parent streams systems to a shared queue. Near-linear scaling was validated on a 4-card +QuietBox (qb1): 3.95x on 4 cards. On a single-card host this still runs — it just uses one card. + + ~/.ttatom_run/env/bin/python examples/batch.py --device-ids 0 --n 32 +""" +from __future__ import annotations + +import argparse +import pathlib +import time + +import numpy as np +from ase.build import molecule + +from tt_atom.batch import MultiCard + +HERE = pathlib.Path(__file__).parent + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", default=str(HERE / "model_tiny_demo.npz")) + ap.add_argument("--device-ids", type=int, nargs="+", default=[0]) + ap.add_argument("--n", type=int, default=32, help="number of systems to evaluate") + args = ap.parse_args() + + base = molecule("CH3CH2OH") + rng = np.random.default_rng(0) + systems = [] + for _ in range(args.n): + pos = base.get_positions() + rng.normal(scale=0.05, size=base.get_positions().shape) + systems.append((pos, base.get_atomic_numbers())) + + with MultiCard(args.weights, device_ids=tuple(args.device_ids)) as pool: + t = time.perf_counter() + energies, total_edges = pool.energies(systems) + dt = time.perf_counter() - t + + print(f"evaluated {args.n} systems on {len(args.device_ids)} card(s) in {dt:.2f}s " + f"({args.n / dt:.1f} systems/s, {total_edges} edges total)") + print(f"first energies: {[round(float(e), 4) for e in energies[:4]]}") + + +if __name__ == "__main__": + main() diff --git a/examples/evaluate_batch.py b/examples/evaluate_batch.py new file mode 100644 index 0000000..381a3fe --- /dev/null +++ b/examples/evaluate_batch.py @@ -0,0 +1,60 @@ +"""Disjoint-union (block-diagonal) batched inference on ONE card. + +Evaluate K independent systems in a SINGLE device forward — the fairchem/PyG way: concatenate +them into one big block-diagonal graph, run once, recover per-system energies by segment-sum +(and per-system forces from the one shared analytic backward). This is the throughput lever for +many small systems, where one-at-a-time is host-dispatch-bound; see benchmarks/bench_batch.py. + +The single-system TTAtomCalculator API is unchanged — batching is just the extra +``evaluate_batch`` method. A merged uma-s-1 bundle bakes the MoLE routing for one reduced +composition, so the batch shares that composition (e.g. an MD ensemble / conformer set of one +molecule); pass a real merged bundle for physical energies. + + ~/.ttatom_run/env/bin/python examples/evaluate_batch.py --weights uma_s_ethanol.npz --n 32 +""" +from __future__ import annotations + +import argparse +import pathlib +import time + +from ase.build import molecule + +from tt_atom import TTAtomCalculator + +HERE = pathlib.Path(__file__).parent + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", default=str(HERE / "model_tiny_demo.npz")) + ap.add_argument("--mol", default="CH3CH2OH") + ap.add_argument("--n", type=int, default=32, help="number of systems in the batch") + ap.add_argument("--device-id", type=int, default=0) + args = ap.parse_args() + + systems = [] + for i in range(args.n): + a = molecule(args.mol) + a.rattle(stdev=0.05, seed=i) # a conformer set (one composition) + a.info.update(charge=0, spin=1) + systems.append(a) + + calc = TTAtomCalculator(args.weights, device_id=args.device_id) + try: + calc.evaluate_batch(systems) # warm the program cache for this shape + t = time.perf_counter() + out = calc.evaluate_batch(systems) + dt = time.perf_counter() - t + finally: + calc.close() + + E = out["energy"] + print(f"evaluated {args.n} systems in ONE device forward in {dt*1e3:.1f} ms " + f"({args.n / dt:.1f} systems/s)") + print(f"per-system energies (eV): {[round(float(e), 4) for e in E[:4]]} ...") + print(f"forces[0] shape: {out['forces'][0].shape}") + + +if __name__ == "__main__": + main() diff --git a/examples/md.py b/examples/md.py new file mode 100644 index 0000000..92040ad --- /dev/null +++ b/examples/md.py @@ -0,0 +1,60 @@ +"""Molecular dynamics on Tenstorrent via the TT-Atom ASE calculator. + +Runs Langevin MD with conservative analytic forces from the device. ``--trace`` captures the +device forward+backward once and replays it each step (fixed topology) for ~2x fewer host +dispatches — the forces are bit-for-bit identical to the eager path. The shipped demo bundle is +random-weight (arbitrary surface); point ``--weights`` at a bundle exported from a real UMA +checkpoint (``tools/export_weights.py`` / ``tt-atom convert-checkpoint``) for real dynamics. + + ~/.ttatom_run/env/bin/python examples/md.py --trace +""" +from __future__ import annotations + +import argparse +import pathlib + +from ase import units +from ase.build import molecule +from ase.md.langevin import Langevin +from ase.md.velocitydistribution import MaxwellBoltzmannDistribution + +from tt_atom.calculator import TTAtomCalculator + +HERE = pathlib.Path(__file__).parent + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", default=str(HERE / "model_tiny_demo.npz")) + ap.add_argument("--steps", type=int, default=100) + ap.add_argument("--dt", type=float, default=0.5, help="timestep (fs)") + ap.add_argument("--temp", type=float, default=300.0, help="temperature (K)") + ap.add_argument("--trace", action="store_true", help="trace-captured device loop (~2x)") + ap.add_argument("--device-id", type=int, default=0) + args = ap.parse_args() + + atoms = molecule("CH3CH2OH") + atoms.info.update(charge=0, spin=0) + calc = TTAtomCalculator(args.weights, device_id=args.device_id, trace=args.trace) + atoms.calc = calc + try: + MaxwellBoltzmannDistribution(atoms, temperature_K=args.temp) + dyn = Langevin(atoms, timestep=args.dt * units.fs, temperature_K=args.temp, + friction=0.01 / units.fs) + + def _log(): + ekin = atoms.get_kinetic_energy() + print(f"step {dyn.nsteps:4d} E={atoms.get_potential_energy():.5f} eV " + f"T={ekin / (1.5 * units.kB * len(atoms)):.1f} K") + + dyn.attach(_log, interval=max(1, args.steps // 10)) + e0 = atoms.get_potential_energy() + dyn.run(args.steps) + print(f"\nMD: {args.steps} steps x {args.dt} fs at {args.temp} K; " + f"E {e0:.5f} -> {atoms.get_potential_energy():.5f} eV (trace={args.trace})") + finally: + calc.close() + + +if __name__ == "__main__": + main() diff --git a/examples/model_tiny_demo.npz b/examples/model_tiny_demo.npz new file mode 100644 index 0000000..9860295 Binary files /dev/null and b/examples/model_tiny_demo.npz differ diff --git a/examples/periodic.py b/examples/periodic.py new file mode 100644 index 0000000..cb83320 --- /dev/null +++ b/examples/periodic.py @@ -0,0 +1,52 @@ +"""Periodic materials (PBC) on Tenstorrent via the TT-Atom ASE calculator. + +Energy + forces (and an optional cell-fixed relaxation) for a bulk crystal. TT-Atom's cell-aware +neighbour list (minimum-image, matching fairchem ``radius_graph_pbc``) makes materials tasks +work; validated against the fairchem uma-s-1 oracle on bulk Si (omat) to energy rel < 1e-3 and +force PCC > 0.99. Use a bundle exported with ``--task omat`` (or oc20/odac/omc) for real numbers; +the shipped demo bundle is random-weight but still exercises the full periodic path. + + ~/.ttatom_run/env/bin/python examples/periodic.py --weights si_omat.npz +""" +from __future__ import annotations + +import argparse +import pathlib + +import numpy as np +from ase.build import bulk +from ase.optimize import FIRE + +from tt_atom.calculator import TTAtomCalculator + +HERE = pathlib.Path(__file__).parent + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", default=str(HERE / "model_tiny_demo.npz")) + ap.add_argument("--relax", action="store_true", help="run a FIRE relaxation (atoms only)") + ap.add_argument("--device-id", type=int, default=0) + args = ap.parse_args() + + atoms = bulk("Si", "diamond", a=5.43) * (2, 1, 1) # 4-atom periodic cell + atoms.rattle(stdev=0.1, seed=1) + atoms.info.update(charge=0, spin=1) + print(f"periodic system: {atoms.get_chemical_formula()}, pbc={atoms.get_pbc().tolist()}, " + f"cell diag={np.round(atoms.cell.lengths(), 3).tolist()}") + + calc = TTAtomCalculator(args.weights, device_id=args.device_id) + atoms.calc = calc + try: + E = atoms.get_potential_energy() + fmax = float((atoms.get_forces() ** 2).sum(1).max() ** 0.5) + print(f"energy = {E:.6f} eV |F|max = {fmax:.4f} eV/A") + if args.relax: + FIRE(atoms, logfile="-").run(fmax=0.05, steps=100) + print(f"relaxed energy = {atoms.get_potential_energy():.6f} eV") + finally: + calc.close() + + +if __name__ == "__main__": + main() diff --git a/examples/relax.py b/examples/relax.py new file mode 100644 index 0000000..9a7bc30 --- /dev/null +++ b/examples/relax.py @@ -0,0 +1,52 @@ +"""Geometry relaxation on Tenstorrent via the TT-Atom ASE calculator. + +Runs a real FIRE optimization to convergence on device. The shipped demo bundle holds +*random* weights (the eSCN-MD architecture, no UMA checkpoint), so the energy surface is +arbitrary — but the forces are the exact analytic gradient of that energy, so the relaxation +genuinely converges. Point ``--weights`` at a bundle exported from a fairchem checkpoint +(``tools/export_weights.py``) to relax on the real potential. + + ~/.ttatom_run/env/bin/python examples/relax.py +""" +from __future__ import annotations + +import argparse +import pathlib + +from ase.build import molecule +from ase.optimize import FIRE + +from tt_atom.calculator import TTAtomCalculator + +HERE = pathlib.Path(__file__).parent + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", default=str(HERE / "model_tiny_demo.npz")) + ap.add_argument("--fmax", type=float, default=0.05) + ap.add_argument("--steps", type=int, default=200) + ap.add_argument("--device-id", type=int, default=0) + args = ap.parse_args() + + atoms = molecule("CH3CH2OH") # ethanol + atoms.info["charge"] = 0 + atoms.info["spin"] = 0 + atoms.rattle(stdev=0.05, seed=0) # perturb so there is something to relax + + calc = TTAtomCalculator(args.weights, device_id=args.device_id) + atoms.calc = calc + try: + e0 = atoms.get_potential_energy() + opt = FIRE(atoms, logfile="-") + opt.run(fmax=args.fmax, steps=args.steps) + e1 = atoms.get_potential_energy() + fmax = float((atoms.get_forces() ** 2).sum(1).max() ** 0.5) + print(f"\nrelaxation: E {e0:.6f} -> {e1:.6f} eV over {opt.nsteps} steps; " + f"fmax={fmax:.4f} (target {args.fmax}); converged={fmax <= args.fmax}") + finally: + calc.close() + + +if __name__ == "__main__": + main() diff --git a/examples/relax_cell.py b/examples/relax_cell.py new file mode 100644 index 0000000..3ee4e3a --- /dev/null +++ b/examples/relax_cell.py @@ -0,0 +1,62 @@ +"""Variable-cell relaxation on Tenstorrent — the end-to-end proof of the stress tensor. + +An ``ase.filters.FrechetCellFilter`` lets FIRE relax the atomic positions *and* the unit cell +together; the cell degrees of freedom are driven by the stress the TT-Atom calculator now +exposes (virial = symmetrized ``dE/dstrain``, divided by volume — fairchem's convention). This +is UMA's flagship materials use case (NPT / variable-cell geometry optimization). + +Point ``--weights`` at a bundle exported from the real uma-s-1 ``omat`` checkpoint to relax on +the physical potential (default: the periodic golden bundle if present, else the random-weight +demo bundle — which still converges, on its own arbitrary surface). + + ~/.ttatom_run/env/bin/python examples/relax_cell.py +""" +from __future__ import annotations + +import argparse +import pathlib + +from ase.build import bulk +from ase.filters import FrechetCellFilter +from ase.optimize import FIRE + +from tt_atom.calculator import TTAtomCalculator + +HERE = pathlib.Path(__file__).parent +_OMAT = pathlib.Path.home() / ".ttatom_run/goldens_real/si_omat.npz" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--weights", default=str(_OMAT if _OMAT.exists() + else HERE / "model_tiny_demo.npz")) + ap.add_argument("--fmax", type=float, default=0.05) + ap.add_argument("--steps", type=int, default=200) + ap.add_argument("--device-id", type=int, default=0) + args = ap.parse_args() + + atoms = bulk("Si", "diamond", a=5.35) * (2, 1, 1) # slightly compressed -> nonzero stress + atoms.info.update(charge=0, spin=1) + + calc = TTAtomCalculator(args.weights, device_id=args.device_id) + atoms.calc = calc + try: + e0 = atoms.get_potential_energy() + v0 = atoms.get_volume() + s0 = float((atoms.get_stress() ** 2).sum() ** 0.5) + opt = FIRE(FrechetCellFilter(atoms), logfile="-") + opt.run(fmax=args.fmax, steps=args.steps) + e1 = atoms.get_potential_energy() + v1 = atoms.get_volume() + fmax = float((atoms.get_forces() ** 2).sum(1).max() ** 0.5) + smax = float((atoms.get_stress() ** 2).sum() ** 0.5) + print(f"\nvariable-cell relax: E {e0:.5f} -> {e1:.5f} eV, V {v0:.3f} -> {v1:.3f} A^3 " + f"over {opt.nsteps} steps") + print(f" |stress| {s0:.4f} -> {smax:.4f} eV/A^3; fmax={fmax:.4f} " + f"(target {args.fmax}); converged={fmax <= args.fmax}") + finally: + calc.close() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..a81bb93 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,49 @@ +[build-system] +requires = ["setuptools>=64"] +build-backend = "setuptools.build_meta" + +[project] +name = "tt-atom" +version = "0.1.0" +description = "High-performance Tenstorrent inference for eSEN / eSCN-MD (UMA-family) ML interatomic potentials" +readme = "README.md" +requires-python = ">=3.10" +license = { text = "MIT" } +authors = [{ name = "Moritz Thüning" }] +keywords = ["tenstorrent", "ttnn", "machine-learning-potential", "eSEN", "eSCN", "UMA", "equivariant", "MLIP"] + +# Core runtime deps. ttnn is NOT a pip dependency: this build routes the rotation through custom +# tt-metal kernels (ttnn.experimental.fused_rotate) that the pip wheel does not carry, so ttnn must +# come from a source tt-metal build that includes the op (see README "Install"). numpy is pinned <2 +# to match the source-ttnn numpy 1.x C ABI; torch is CPU-only (the model runs on the card). +# `import tt_atom` never imports ttnn, so it imports fine on a machine without a card. +dependencies = [ + # ttnn is provided by the source tt-metal build with the custom op -- see README "Install". + "numpy<2", + "torch>=2.0,<3", + "ase>=3.26", +] + +[project.scripts] +tt-atom = "tt_atom.cli:main" + +[project.optional-dependencies] +# fairchem is needed ONLY to (a) load a real UMA/eSEN checkpoint and (b) generate the +# golden reference tensors used by the parity tests. It pulls numpy>=2 and therefore +# must live in a SEPARATE environment from ttnn (see README "Reference environment"). +reference = ["fairchem-core>=2.10"] +dev = ["pytest", "matplotlib"] + +[project.urls] +Homepage = "https://github.com/moritztng/tt-atom" + +[tool.setuptools.packages.find] +include = ["tt_atom*"] + +[tool.setuptools.package-data] +# the quaternion edge frame (default) loads this vendored fairchem Wigner-D table at runtime, so it +# must ship in the wheel (a non-editable `pip install` otherwise omits it -> FileNotFoundError). +tt_atom = ["assets/*.pt"] + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..7be995a --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,26 @@ +"""Shared parity-test fixtures: a single TT device for the session and the golden loader.""" +import os + +import pytest + +os.environ.setdefault("TT_METAL_LOGGER_LEVEL", "FATAL") + +from util import Golden # noqa: E402 (pytest puts tests/ on sys.path) + + +@pytest.fixture(scope="session") +def golden(): + return Golden("golden_tiny.npz") + + +@pytest.fixture(scope="session") +def device(): + from tt_atom import device as D + + # reserve a trace region so the trace-path test can capture on this shared device; harmless + # (only reserves DRAM) for the eager parity tests. + dev = D.open_device(0, trace_region_size=400_000_000) + yield dev + import ttnn + + ttnn.close_device(dev) diff --git a/tests/data/golden_bulk_tiny.npz b/tests/data/golden_bulk_tiny.npz new file mode 100644 index 0000000..2dd27b1 Binary files /dev/null and b/tests/data/golden_bulk_tiny.npz differ diff --git a/tests/data/golden_tiny.npz b/tests/data/golden_tiny.npz new file mode 100644 index 0000000..503051a Binary files /dev/null and b/tests/data/golden_tiny.npz differ diff --git a/tests/gen_golden.py b/tests/gen_golden.py new file mode 100644 index 0000000..e06ac98 --- /dev/null +++ b/tests/gen_golden.py @@ -0,0 +1,217 @@ +"""Generate golden reference tensors for TT-Atom parity tests. + +Run with the *reference* environment (fairchem-core, numpy>=2), NOT the ttnn env: + + ~/.ttatom_run/refenv/bin/python tests/gen_golden.py --tiny --out tests/data/golden_tiny.npz + +This instantiates the fairchem ``eSCNMDBackbone`` (the eSEN / eSCN-MD / UMA backbone) with +deterministic random weights, runs a forward + autograd-force pass on a small system, and +saves to a single ``.npz``: + + * inputs (atomic_numbers, pos, edge_index, ...) + * host geometric (wigner_and_M_mapping[_inv], edge_envelope, x_edge, sys_node_embedding, + x_message_init) -- the terms TT-Atom precomputes on host + * weights (the full state_dict, key -> array, under the ``w@`` prefix) + * activations (per-module inputs/outputs under the ``a@`` prefix) + * outputs (node_embedding, energy, forces) + * config (JSON string of the backbone config) + +The ttnn port (different env, numpy<2) loads this npz and checks PCC per-module and +end-to-end. .npz arrays are numpy-version agnostic, which is exactly why we decouple the +two environments through disk rather than importing fairchem next to ttnn. +""" +from __future__ import annotations + +import argparse +import json + +import numpy as np +import torch + +from ase.build import molecule, bulk +from fairchem.core.datasets.atomic_data import AtomicData +from fairchem.core.models.uma.escn_md import eSCNMDBackbone + + +TINY = dict( + sphere_channels=32, lmax=2, mmax=2, num_layers=2, hidden_channels=32, + edge_channels=16, num_distance_basis=32, +) +# Representative "uma-s-like" config (used for perf / full-size goldens). +FULL = dict( + sphere_channels=128, lmax=2, mmax=2, num_layers=2, hidden_channels=128, + edge_channels=128, num_distance_basis=512, +) +COMMON = dict( + max_num_elements=100, cutoff=5.0, max_neighbors=300, otf_graph=False, + direct_forces=False, regress_forces=True, regress_stress=False, + norm_type="rms_norm_sh", act_type="gate", ff_type="grid", + use_dataset_embedding=True, dataset_list=["omat"], distance_function="gaussian", +) + + +def build_system(kind: str): + if kind == "molecule": + atoms = molecule("CH3CH2OH") # ethanol, 9 atoms, aperiodic + atoms.info["charge"] = 0 + atoms.info["spin"] = 0 + return atoms + if kind == "bulk": + atoms = bulk("Si", "diamond", a=5.43) * (2, 1, 1) # small periodic cell + atoms.rattle(stdev=0.1, seed=1) # break symmetry -> nonzero forces + atoms.info["charge"] = 0 + atoms.info["spin"] = 0 + return atoms + raise ValueError(kind) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--tiny", action="store_true", help="use the tiny config (committed golden)") + ap.add_argument("--system", default="molecule", choices=["molecule", "bulk"]) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--out", required=True) + args = ap.parse_args() + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + + cfg = dict(COMMON) + cfg.update(TINY if args.tiny else FULL) + + backbone = eSCNMDBackbone(**cfg).eval() + # energy head: matches MLP_Energy_Head (sum of per-node MLP on l=0 channel) + sc, hc = cfg["sphere_channels"], cfg["hidden_channels"] + energy_block = torch.nn.Sequential( + torch.nn.Linear(sc, hc), torch.nn.SiLU(), + torch.nn.Linear(hc, hc), torch.nn.SiLU(), + torch.nn.Linear(hc, 1), + ) + + atoms = build_system(args.system) + # aperiodic molecules need a vacuum box for the periodic neighbour-list builder + mol_box = 12.0 if args.system == "molecule" else None + data = AtomicData.from_ase( + atoms, r_edges=True, radius=cfg["cutoff"], max_neigh=cfg["max_neighbors"], + molecule_cell_size=mol_box, task_name="omat", target_dtype=torch.float32, + ) + + saved: dict[str, np.ndarray] = {} + + def npy(t): + return t.detach().to(torch.float32).cpu().numpy() if t.dtype.is_floating_point \ + else t.detach().cpu().numpy() + + # ---- capture per-module activations via hooks ---------------------------------- + acts: dict[str, np.ndarray] = {} + + def save_io(name): + def hook(mod, inp, out): + for i, t in enumerate(inp): + if torch.is_tensor(t): + acts[f"{name}.in{i}"] = npy(t) + outs = out if isinstance(out, tuple) else (out,) + for i, t in enumerate(outs): + if torch.is_tensor(t): + acts[f"{name}.out{i}"] = npy(t) + return hook + + handles = [] + handles.append(backbone.edge_degree_embedding.register_forward_hook(save_io("edge_degree"))) + handles.append(backbone.norm.register_forward_hook(save_io("final_norm"))) + for li, blk in enumerate(backbone.blocks): + handles.append(blk.register_forward_hook(save_io(f"block{li}"))) + handles.append(blk.norm_1.register_forward_hook(save_io(f"block{li}.norm_1"))) + handles.append(blk.edge_wise.register_forward_hook(save_io(f"block{li}.edgewise"))) + handles.append(blk.edge_wise.so2_conv_1.register_forward_hook(save_io(f"block{li}.so2_1"))) + handles.append(blk.edge_wise.so2_conv_2.register_forward_hook(save_io(f"block{li}.so2_2"))) + handles.append(blk.norm_2.register_forward_hook(save_io(f"block{li}.norm_2"))) + handles.append(blk.atom_wise.register_forward_hook(save_io(f"block{li}.atomwise"))) + + # capture the host geometric terms produced inside forward by wrapping the method + orig_wigner = backbone._get_rotmat_and_wigner + captured = {} + + def wrap_wigner(edge_distance_vecs): + w, winv = orig_wigner(edge_distance_vecs) + captured["wigner"] = npy(w) + captured["wigner_inv"] = npy(winv) + return w, winv + + backbone._get_rotmat_and_wigner = wrap_wigner + + orig_graph = backbone._generate_graph + + def wrap_graph(dd): + gd = orig_graph(dd) + captured["edge_distance_vec"] = npy(gd["edge_distance_vec"]) + captured["edge_distance"] = npy(gd["edge_distance"]) + return gd + + backbone._generate_graph = wrap_graph + + # ---- forward + autograd forces ------------------------------------------------- + out = backbone(data) + node_emb = out["node_embedding"] + node_energy = energy_block(node_emb.narrow(1, 0, 1).squeeze(1)).view(-1) + energy = torch.zeros(len(data["natoms"])) + energy.index_add_(0, data["batch"], node_energy) + forces = -torch.autograd.grad(energy.sum(), data["pos"])[0] + + for h in handles: + h.remove() + + # ---- assemble npz -------------------------------------------------------------- + saved["config"] = np.frombuffer(json.dumps(cfg).encode(), dtype=np.uint8) + # inputs + saved["in@atomic_numbers"] = npy(data["atomic_numbers"]) + saved["in@pos"] = npy(data["pos"]) + saved["in@edge_index"] = npy(data["edge_index"]) + saved["in@cell"] = npy(data["cell"]) + saved["in@batch"] = npy(data["batch"]) + saved["in@natoms"] = npy(data["natoms"]) + saved["in@charge"] = npy(data["charge"]) + saved["in@spin"] = npy(data["spin"]) + # host geometric terms TT-Atom precomputes + saved["host@wigner"] = captured["wigner"] + saved["host@wigner_inv"] = captured["wigner_inv"] + saved["host@edge_distance_vec"] = captured["edge_distance_vec"] + saved["host@edge_distance"] = captured["edge_distance"] + # fixed geometry buffers needed to rebuild the host pos->geometric Jacobian (forces) + saved["host@to_m"] = npy(backbone.mappingReduced.to_m) + saved["host@gauss_offset"] = npy(backbone.distance_expansion.offset) + saved["host@gauss_coeff"] = np.array([backbone.distance_expansion.coeff], dtype=np.float32) + saved["host@x_edge"] = acts["block0.edgewise.in1"] # x_edge fed to edgewise + saved["host@edge_envelope"] = acts["block0.edgewise.in6"] # edge_envelope arg + saved["host@x_message_init"] = acts["edge_degree.out0"] # node feats after edge-degree emb + # SO3 grid transform matrices (fixed) for GridAtomwise; and the per-node system embedding + sg = backbone.SO3_grid["lmax_lmax"] + saved["host@to_grid_mat"] = npy(sg.to_grid_mat) + saved["host@from_grid_mat"] = npy(sg.from_grid_mat) + csd = backbone.csd_embedding(data["charge"], data["spin"], data.get("dataset", default=None)) + saved["host@sys_node_embedding"] = npy(csd[data["batch"]]) + # outputs + saved["out@node_embedding"] = npy(node_emb) + saved["out@energy"] = npy(energy) + saved["out@forces"] = npy(forces) + # weights + for k, v in backbone.state_dict().items(): + saved[f"w@{k}"] = npy(v) + for k, v in energy_block.state_dict().items(): + saved[f"w@energy_block.{k}"] = npy(v) + # activations + for k, v in acts.items(): + saved[f"a@{k}"] = v + + np.savez(args.out, **saved) + print(f"wrote {args.out}") + print(f" config: {cfg}") + print(f" natoms={int(data['natoms'].sum())} nedges={data['edge_index'].shape[1]} " + f"node_emb={tuple(node_emb.shape)}") + print(f" energy={energy.tolist()} |F|max={forces.abs().max().item():.4f}") + print(f" n weight tensors={sum(1 for k in saved if k.startswith('w@'))} " + f"n activations={len(acts)}") + + +if __name__ == "__main__": + main() diff --git a/tests/gen_golden_batch.py b/tests/gen_golden_batch.py new file mode 100644 index 0000000..a4e3682 --- /dev/null +++ b/tests/gen_golden_batch.py @@ -0,0 +1,111 @@ +"""Generate a REAL-weight BATCHED golden for disjoint-union batching parity (refenv only). + +Run in the reference environment (fairchem-core, numpy>=2), NOT the ttnn env: + + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tests/gen_golden_batch.py \ + --k 8 --task omol --out ~/.ttatom_run/goldens_real/batch_ethanol_omol.npz + +Builds K same-reduced-composition conformers (ethanol, rattled) — the regime a merged uma-s-1 +bundle is valid for (fairchem's ``merge_MOLE_model`` asserts one reduced composition per batch). +It then runs fairchem's OWN batched merged inference (``data_list_collater`` of K ``AtomicData`` +-> one ``predict`` call, per-system energies split by ``data.batch``) and cross-checks that +against the per-conformer ``FAIRChemCalculator`` — proving fairchem itself batches block- +diagonally — before dumping the per-system energies + forces as the golden. + +The ttnn env (test_realweight.py) then loads this, assembles the same conformers into one +block-diagonal graph, runs ``energy_and_forces_batch``, and asserts E rel<1e-3, F PCC>0.99. +""" +from __future__ import annotations + +import argparse +import os + +os.environ.setdefault("HF_HUB_OFFLINE", "1") + +import numpy as np + +from ase.build import molecule +from huggingface_hub import hf_hub_download +from fairchem.core import FAIRChemCalculator +from fairchem.core.datasets import data_list_collater +from fairchem.core.units.mlip_unit import load_predict_unit +from fairchem.core.units.mlip_unit.api.inference import InferenceSettings + + +def conformers(k, seed0=10): + """K rattled ethanol conformers — identical composition, different geometry.""" + out = [] + for i in range(k): + a = molecule("CH3CH2OH") + a.rattle(stdev=0.08, seed=seed0 + i) + a.info.update(charge=0, spin=1) + out.append(a) + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--k", type=int, default=8) + ap.add_argument("--task", default="omol") + ap.add_argument("--ckpt", default="uma-s-1") + ap.add_argument("--out", required=True) + args = ap.parse_args() + + ckpt = hf_hub_download("facebook/UMA", f"checkpoints/{args.ckpt}.pt") + settings = InferenceSettings( + tf32=False, activation_checkpointing=True, merge_mole=True, + compile=False, external_graph_gen=False, internal_graph_gen_version=2, + ) + pu = load_predict_unit(ckpt, inference_settings=settings, device="cpu") + calc = FAIRChemCalculator(pu, task_name=args.task) + + systems = conformers(args.k) + + # ---- per-conformer reference (fairchem's validated single-system inference) ----------- + E_single, F_single = [], [] + for a in systems: + a.calc = calc + E_single.append(float(a.get_potential_energy())) + F_single.append(a.get_forces().astype(np.float32)) + E_single = np.array(E_single, dtype=np.float64) + + # ---- fairchem's OWN batched merged inference (Batch.from_data_list equivalent) -------- + # build each AtomicData exactly as the calculator does (task_name + r_data_keys carry the + # charge/spin the merged model checks), then collate K of them into one batch. + data = [] + for a in systems: + calc.predictor.validate_atoms_data(a, args.task) + data.append(calc.a2g(a)) + batch = data_list_collater(data, otf_graph=True) + pred = pu.predict(batch) + Ebt = pred["energy"].detach().cpu().numpy().astype(np.float64).reshape(-1) + Fbt_all = pred["forces"].detach().cpu().numpy().astype(np.float32) + bidx = batch.batch.detach().cpu().numpy() + F_batched = [Fbt_all[bidx == i] for i in range(args.k)] + + # cross-check: fairchem batched == fairchem per-system (block-diagonal, no cross terms) + e_rel = np.abs(Ebt - E_single).max() / (np.abs(E_single).max() + 1e-6) + f_max = max(np.abs(F_batched[i] - F_single[i]).max() for i in range(args.k)) + print(f"fairchem batched vs single: E rel={e_rel:.2e} F maxdiff={f_max:.2e}") + assert e_rel < 1e-4 and f_max < 1e-3, "fairchem batched != fairchem single — investigate" + + # ---- dump the golden (positions + fairchem batched E/F per system) -------------------- + saved = { + "k": np.array([args.k]), + "task": np.frombuffer(args.task.encode(), dtype=np.uint8), + "charge": np.array([0.0]), + "spin": np.array([1.0]), + "natoms": np.array([len(a) for a in systems], dtype=np.int64), + "Z": np.concatenate([a.get_atomic_numbers() for a in systems]).astype(np.int64), + "pos": np.concatenate([a.get_positions() for a in systems]).astype(np.float64), + "energy": Ebt, # [K] fairchem batched energies + "forces": np.concatenate(F_batched).astype(np.float32), # [Ntot,3] + } + os.makedirs(os.path.dirname(args.out), exist_ok=True) + np.savez(args.out, **saved) + print(f"wrote {args.out}: K={args.k} |E|={np.abs(Ebt).mean():.2f} eV " + f"|F|max={np.abs(Fbt_all).max():.3f}") + + +if __name__ == "__main__": + main() diff --git a/tests/gen_golden_real.py b/tests/gen_golden_real.py new file mode 100644 index 0000000..91b3e70 --- /dev/null +++ b/tests/gen_golden_real.py @@ -0,0 +1,318 @@ +"""Generate REAL-weight golden tensors for TT-Atom uma-s-1 parity (refenv only). + +Run with the *reference* environment (fairchem-core, numpy>=2), NOT the ttnn env: + + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tests/gen_golden_real.py \ + --system molecule --task omol --out ~/.ttatom_run/goldens_real/ethanol_omol.npz + +Unlike ``gen_golden.py`` (random weights, self-chosen config), this loads the gated +``facebook/UMA`` ``uma-s-1`` checkpoint and reproduces the *released* model: + + * EMA weights (the predict unit uses ``use_ema=True``); + * MoLE merge on host -> a plain ``eSCNMDBackbone`` (fairchem's own + ``merge_MOLE_model``, the exact inference path with ``merge_mole=True``); + * ``ff_type=spectral`` atomwise, ``num_layers=4``, ``num_distance_basis=64``, ``cutoff=6``; + * the per-task energy normalizer (``E = rmsd*E_raw + sum_i refs[Z_i]``) and force + scale (``F = rmsd * F_raw``). + +It also records the *unmerged* MoE oracle energy+forces (a fresh predict unit with the +default ``merge_mole=False`` settings) so the merge can be validated to PCC>0.999 (the host +correctness anchor). Goldens are dumped to disk (NOT committed) and consumed by the ttnn env. + +The bundle layout mirrors ``gen_golden.py`` so the ttnn loader / tests share one code path; +spectral-atomwise activations replace the grid ones, and ``scale@*`` carries the normalizer. +""" +from __future__ import annotations + +import argparse +import json +import os + +os.environ.setdefault("HF_HUB_OFFLINE", "1") + +import numpy as np +import torch + +from ase.build import molecule, bulk +from huggingface_hub import hf_hub_download +from fairchem.core import FAIRChemCalculator +from fairchem.core.units.mlip_unit import load_predict_unit +from fairchem.core.units.mlip_unit.api.inference import InferenceSettings + + +def npy(t): + return t.detach().to(torch.float32).cpu().numpy() if t.dtype.is_floating_point \ + else t.detach().cpu().numpy() + + +def build_system(kind: str, task: str): + if kind == "molecule": + atoms = molecule("CH3CH2OH") # ethanol, 9 atoms, aperiodic + atoms.info.update(charge=0, spin=1) # omol default (closed-shell singlet) + return atoms + if kind == "bulk": + atoms = bulk("Si", "diamond", a=5.43) * (2, 1, 1) + atoms.rattle(stdev=0.1, seed=1) + atoms.info.update(charge=0, spin=1) + return atoms + if kind == "slab": + from ase.build import fcc100, add_adsorbate # Cu(100) slab + H adsorbate (oc20) + atoms = fcc100("Cu", (2, 2, 2), vacuum=8.0) + add_adsorbate(atoms, "H", height=1.5, position="hollow") + atoms.rattle(stdev=0.05, seed=2) + atoms.info.update(charge=0, spin=1) # pbc = [True, True, False] (mixed) + return atoms + if kind == "mof": + # odac (DAC / MOFs): a metal-oxide framework fragment (MgO), fully periodic. A minimal + # inorganic stand-in that exercises the odac dataset token + normalizer on the periodic path. + atoms = bulk("MgO", "rocksalt", a=4.21) * (2, 1, 1) + atoms.rattle(stdev=0.08, seed=3) + atoms.info.update(charge=0, spin=1) + return atoms + if kind == "molcrystal": + # omc (molecular crystals): solid CO2 (dry ice) as a periodic cubic cell dense enough that + # periodic images fall within the 6 A cutoff (tests the cell-aware neighbour list). + co2 = molecule("CO2") + co2.set_cell([5.0, 5.0, 5.0]) + co2.set_pbc(True) + co2.center() + co2.rattle(stdev=0.05, seed=4) + co2.info.update(charge=0, spin=1) + return co2 + raise ValueError(kind) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--system", default="molecule", + choices=["molecule", "bulk", "slab", "mof", "molcrystal"]) + ap.add_argument("--task", default="omol") + ap.add_argument("--ckpt", default="uma-s-1", help="UMA checkpoint name (e.g. uma-s-1, uma-s-1p2, uma-m-1p1)") + ap.add_argument("--ckpt-path", default=None, + help="local checkpoint .pt path; overrides --ckpt (skips the HF download)") + ap.add_argument("--out", required=True) + ap.add_argument("--merged-only", action="store_true", + help="skip the separate unmerged-MoE oracle unit and use the merged inference " + "path as the oracle. Needed for big checkpoints (uma-m-1p1 is 11GB — two " + "predict units OOM a 30GB host); the merge is already the released " + "inference path and was validated against the unmerged MoE on uma-s-1.") + args = ap.parse_args() + + ckpt = args.ckpt_path or hf_hub_download("facebook/UMA", f"checkpoints/{args.ckpt}.pt") + atoms = build_system(args.system, args.task) + + # ---- ground-truth oracle: unmerged MoE (the released inference default) -------------- + if args.merged_only: + E_oracle = F_oracle = S_oracle = None # filled from the merged path below + else: + pu_oracle = load_predict_unit(ckpt, inference_settings="default", device="cpu") + calc_oracle = FAIRChemCalculator(pu_oracle, task_name=args.task) + atoms.calc = calc_oracle + E_oracle = float(atoms.get_potential_energy()) + F_oracle = atoms.get_forces().astype(np.float32) + # stress (ASE Voigt-6) only for a fully-periodic cell; zeros as a sentinel otherwise + S_oracle = (atoms.get_stress().astype(np.float32) + if bool(np.all(atoms.get_pbc())) else np.zeros(6, dtype=np.float32)) + print(f"oracle (unmerged MoE): E={E_oracle:.6f} eV |F|max={np.abs(F_oracle).max():.4f} " + f"stress={S_oracle}") + del pu_oracle, calc_oracle # free ~11GB before loading the merged unit + + # ---- merged plain backbone: host MoLE merge (the correctness anchor) ------------------ + settings = InferenceSettings( + tf32=False, activation_checkpointing=True, merge_mole=True, + compile=False, external_graph_gen=False, internal_graph_gen_version=2, + ) + pu = load_predict_unit(ckpt, inference_settings=settings, device="cpu") + calc = FAIRChemCalculator(pu, task_name=args.task) + atoms2 = build_system(args.system, args.task) + atoms2.calc = calc + E_merged = float(atoms2.get_potential_energy()) # triggers the merge in _lazy_init + F_merged = atoms2.get_forces().astype(np.float32) + print(f"merged (host MoLE): E={E_merged:.6f} eV |F|max={np.abs(F_merged).max():.4f}") + + if args.merged_only: # merged path is the released inference oracle + E_oracle, F_oracle = E_merged, F_merged + S_oracle = (atoms2.get_stress().astype(np.float32) + if bool(np.all(atoms2.get_pbc())) else np.zeros(6, dtype=np.float32)) + + hydra = pu.model.module + backbone = hydra.backbone # plain eSCNMDBackbone after merge + head = hydra.output_heads["energyandforcehead"].head + energy_block = head.energy_block + assert type(backbone).__name__ == "eSCNMDBackbone", \ + f"backbone was not merged to a plain backbone (got {type(backbone).__name__})" + + # per-task normalizer (energy scale + element references) ------------------------------- + etask = pu.model.module.tasks[f"{args.task}_energy"] + rmsd = float(etask.normalizer.rmsd) + mean = float(etask.normalizer.mean) + elem_refs = etask.element_references.element_references.detach().cpu().numpy().astype(np.float64) + + # ---- hooks to capture per-module activations on the MERGED backbone ------------------- + acts: dict[str, np.ndarray] = {} + + def save_io(name): + def hook(mod, inp, out): + for i, t in enumerate(inp): + if torch.is_tensor(t): + acts[f"{name}.in{i}"] = npy(t) + outs = out if isinstance(out, tuple) else (out,) + for i, t in enumerate(outs): + if torch.is_tensor(t): + acts[f"{name}.out{i}"] = npy(t) + return hook + + handles = [] + # the real backbone pre-fuses the polynomial envelope into wigner_inv (no separate edgewise + # arg), so capture it straight off the PolynomialEnvelope module (TT-Atom applies it separately) + handles.append(backbone.envelope.register_forward_hook(save_io("envelope"))) + handles.append(backbone.edge_degree_embedding.register_forward_hook(save_io("edge_degree"))) + handles.append(backbone.norm.register_forward_hook(save_io("final_norm"))) + for li, blk in enumerate(backbone.blocks): + handles.append(blk.register_forward_hook(save_io(f"block{li}"))) + handles.append(blk.norm_1.register_forward_hook(save_io(f"block{li}.norm_1"))) + handles.append(blk.edge_wise.register_forward_hook(save_io(f"block{li}.edgewise"))) + handles.append(blk.edge_wise.so2_conv_1.register_forward_hook(save_io(f"block{li}.so2_1"))) + handles.append(blk.edge_wise.so2_conv_2.register_forward_hook(save_io(f"block{li}.so2_2"))) + handles.append(blk.norm_2.register_forward_hook(save_io(f"block{li}.norm_2"))) + aw = blk.atom_wise # SpectralAtomwise + handles.append(aw.register_forward_hook(save_io(f"block{li}.atomwise"))) + handles.append(aw.scalar_mlp.register_forward_hook(save_io(f"block{li}.aw_scalar"))) + handles.append(aw.so3_linear_1.register_forward_hook(save_io(f"block{li}.aw_so3lin1"))) + handles.append(aw.act.register_forward_hook(save_io(f"block{li}.aw_gate"))) + handles.append(aw.so3_linear_2.register_forward_hook(save_io(f"block{li}.aw_so3lin2"))) + + captured = {} + orig_wigner = backbone._get_rotmat_and_wigner + + def wrap_wigner(edge_distance_vecs): + w, winv = orig_wigner(edge_distance_vecs) + captured["wigner"] = npy(w) + captured["wigner_inv"] = npy(winv) + return w, winv + + backbone._get_rotmat_and_wigner = wrap_wigner + + orig_graph = backbone._generate_graph + + def wrap_graph(dd): + gd = orig_graph(dd) + captured["edge_distance_vec"] = npy(gd["edge_distance_vec"]) + captured["edge_distance"] = npy(gd["edge_distance"]) + captured["edge_index"] = npy(gd["edge_index"]) + return gd + + backbone._generate_graph = wrap_graph + + # re-run via the calculator to fire hooks on the merged path; capture the data object too + data = calc.a2g(atoms2) + captured["data"] = data + atoms2.calc.calculate(atoms2, ["energy", "forces"], ["positions"]) # re-run -> fills acts + + for h in handles: + h.remove() + backbone._get_rotmat_and_wigner = orig_wigner + backbone._generate_graph = orig_graph + + # ---- node embedding + raw/denorm energy from the merged backbone directly ------------- + data = data.clone() + data.pos.requires_grad_(True) + emb = backbone(data) + node_emb = emb["node_embedding"] + scalar = node_emb.narrow(1, 0, 1).squeeze(1) # l=0 channel [N, C] + node_energy = energy_block(scalar).view(-1) + n_sys = len(data["natoms"]) + E_raw = torch.zeros(n_sys, dtype=node_energy.dtype).index_add(0, data["batch"], node_energy) + refs_t = torch.from_numpy(elem_refs).to(node_energy.dtype) + ref_per_atom = refs_t[data["atomic_numbers"].long()] + E_final = E_raw * rmsd + mean + E_final = E_final.index_add(0, data["batch"], ref_per_atom) + F_final = -torch.autograd.grad(E_final.sum(), data.pos)[0] + print(f"recompute (merged bb): E={E_final.sum().item():.6f} eV " + f"|F|max={F_final.abs().max().item():.4f} E_raw={E_raw.sum().item():.6f}") + + # ---- assemble npz ---------------------------------------------------------------------- + saved: dict[str, np.ndarray] = {} + # config the ttnn port needs (real uma-s-1 values) + bb = backbone + out_cfg = dict( + sphere_channels=bb.sphere_channels, lmax=bb.lmax, mmax=bb.mmax, + num_layers=len(bb.blocks), hidden_channels=bb.hidden_channels, + num_distance_basis=int(bb.distance_expansion.offset.numel()), + cutoff=float(bb.cutoff), ff_type="spectral", act_type="gate", + norm_type="rms_norm_sh", chg_spin_emb_type=bb.chg_spin_emb_type, task=args.task, + # charge_balanced_channels (uma-s-1.2): l=0 charge channels re-balanced to the system + # charge after each block. cs==ce (uma-s-1) => disabled. Mirrors tools/export_weights.py. + charge_channel_start=int(getattr(bb, "charge_channel_start", 0)), + charge_channel_end=int(getattr(bb, "charge_channel_end", 0)), + ) + saved["config"] = np.frombuffer(json.dumps(out_cfg).encode(), dtype=np.uint8) + + di = captured["data"] + saved["in@atomic_numbers"] = npy(di["atomic_numbers"]) + saved["in@pos"] = npy(data.pos) + saved["in@edge_index"] = captured["edge_index"] + saved["in@cell"] = npy(di["cell"]) + saved["in@pbc"] = np.asarray(atoms2.get_pbc(), dtype=bool) + saved["in@batch"] = npy(di["batch"]) + saved["in@natoms"] = npy(di["natoms"]) + saved["in@charge"] = npy(di["charge"]) + saved["in@spin"] = npy(di["spin"]) + + saved["host@wigner"] = captured["wigner"] + saved["host@wigner_inv"] = captured["wigner_inv"] + saved["host@edge_distance_vec"] = captured["edge_distance_vec"] + saved["host@edge_distance"] = captured["edge_distance"] + saved["host@to_m"] = npy(backbone.mappingReduced.to_m) + # coefficient subselection for mmax 25 SH coeffs + # reduced to 19 m-space); None/absent when mmax==lmax (uma-s) + if backbone.mmax != backbone.lmax: + saved["host@coefficient_index"] = npy(backbone.coefficient_index) + saved["host@gauss_offset"] = npy(backbone.distance_expansion.offset) + saved["host@gauss_coeff"] = np.array([backbone.distance_expansion.coeff], dtype=np.float32) + saved["host@x_edge"] = acts["block0.edgewise.in1"] + saved["host@edge_envelope"] = acts["envelope.out0"].reshape(-1, 1, 1) + saved["host@x_message_init"] = acts["edge_degree.out0"] + sg = backbone.SO3_grid["lmax_lmax"] + saved["host@to_grid_mat"] = npy(sg.to_grid_mat) + saved["host@from_grid_mat"] = npy(sg.from_grid_mat) + csd = backbone.csd_embedding(di["charge"], di["spin"], di.get("dataset", default=None)) + saved["host@sys_node_embedding"] = npy(csd[di["batch"]]) + + # energy normalizer / element references + saved["scale@rmsd"] = np.array([rmsd], dtype=np.float64) + saved["scale@mean"] = np.array([mean], dtype=np.float64) + saved["scale@elem_refs"] = elem_refs + + # outputs + saved["out@node_embedding"] = npy(node_emb) + saved["out@energy_raw"] = npy(E_raw) + saved["out@energy"] = np.array([E_final.sum().item()], dtype=np.float64) + saved["out@forces"] = npy(F_final) + saved["out@energy_oracle"] = np.array([E_oracle], dtype=np.float64) + saved["out@forces_oracle"] = F_oracle + saved["out@stress_oracle"] = S_oracle + saved["out@energy_merged_oracle"] = np.array([E_merged], dtype=np.float64) + saved["out@forces_merged_oracle"] = F_merged + + # weights (merged plain backbone + energy head) + for k, v in backbone.state_dict().items(): + saved[f"w@{k}"] = npy(v) + for k, v in energy_block.state_dict().items(): + saved[f"w@energy_block.{k}"] = npy(v) + + for k, v in acts.items(): + saved[f"a@{k}"] = v + + os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) + np.savez(args.out, **saved) + print(f"wrote {args.out}") + print(f" config: {out_cfg}") + n_edges = captured["edge_index"].shape[1] + print(f" natoms={int(di['natoms'].sum())} nedges={n_edges} node_emb={tuple(node_emb.shape)}") + print(f" n weight tensors={sum(1 for k in saved if k.startswith('w@'))} " + f"n activations={len(acts)}") + + +if __name__ == "__main__": + main() diff --git a/tests/mirror.py b/tests/mirror.py new file mode 100644 index 0000000..8f6751b --- /dev/null +++ b/tests/mirror.py @@ -0,0 +1,134 @@ +"""A faithful PyTorch transcription of the TT-Atom device forward. + +This is the differentiable *oracle* for the analytic-force work: it computes exactly the same +mathematical function as the ttnn backbone (same m-primed SO(2) split, same one-hot scatter, +same RMS-norm-SH), so ``torch.autograd`` on it yields the ground-truth adjoints that the +hand-written on-device VJP (``tt_atom/forces.py``) must reproduce. It also serves as the host +forward that composes with ``tt_atom/geometry.py`` for the full ``-dE/dpos`` force. +""" +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tt_atom.activation import _expand_index_m_prime + + +def _l_of_coeff(lmax): + return [l for l in range(lmax + 1) for _ in range(2 * l + 1)] + + +def radial_mlp(x_edge, w, p): + x = F.linear(x_edge, w[f"{p}.net.0.weight"], w[f"{p}.net.0.bias"]) + x = F.layer_norm(x, (x.shape[-1],), w[f"{p}.net.1.weight"], w[f"{p}.net.1.bias"], 1e-5) + x = F.silu(x) + x = F.linear(x, w[f"{p}.net.3.weight"], w[f"{p}.net.3.bias"]) + x = F.layer_norm(x, (x.shape[-1],), w[f"{p}.net.4.weight"], w[f"{p}.net.4.bias"], 1e-5) + x = F.silu(x) + return F.linear(x, w[f"{p}.net.6.weight"], w[f"{p}.net.6.bias"]) + + +def so2(x, x_edge, w, p, lmax, mmax, Cin, H, extra): + E, nsph = x.shape[0], (lmax + 1) ** 2 + xf = x.reshape(E, nsph * Cin) + num_coef = [lmax - m + 1 for m in range(mmax + 1)] + off = [0, (lmax + 1) * Cin] + for m in range(1, mmax + 1): + off.append(off[-1] + 2 * num_coef[m] * Cin) + if f"{p}.rad_func.net.0.weight" in w: + rad = radial_mlp(x_edge, w, f"{p}.rad_func") + sizes = [num_coef[m] * Cin for m in range(mmax + 1)] + o, rms = 0, [] + for m in range(mmax + 1): + rms.append(rad[:, o:o + sizes[m]]); o += sizes[m] + mult = torch.cat([rms[0]] + sum(([rms[m], rms[m]] for m in range(1, mmax + 1)), []), 1) + xf = xf * mult + blocks = [] + x0 = F.linear(xf[:, off[0]:off[1]], w[f"{p}.fc_m0.weight"], w[f"{p}.fc_m0.bias"]) + extra_t = None + if extra: + extra_t, x0 = x0[:, :extra], x0[:, extra:] + blocks.append(x0) + for m in range(1, mmax + 1): + nc = num_coef[m] + blk = xf[:, off[m]:off[m + 1]].reshape(E, 2, nc * Cin) + blk = blk @ w[f"{p}.so2_m_conv.{m-1}.fc.weight"].T + Hh = blk.shape[-1] // 2 + blk = blk.reshape(E, 4 * Hh) + a, b, c, d = blk[:, :Hh], blk[:, Hh:2 * Hh], blk[:, 2 * Hh:3 * Hh], blk[:, 3 * Hh:] + blocks.append(a - d); blocks.append(c + b) + out = torch.cat(blocks, 1).reshape(E, nsph, H) + return (out, extra_t) if extra else out + + +def rms_norm_sh(x, w, p, lmax, C, eps=1e-5): + nsph = (lmax + 1) ** 2 + lc = _l_of_coeff(lmax) + bdw = torch.tensor([1.0 / (2 * l + 1) / (lmax + 1) for l in lc]).view(1, nsph, 1) + l0 = x[:, 0:1, :] + x = torch.cat([l0 - l0.mean(2, keepdim=True), x[:, 1:, :]], 1) + fn = (x * x * bdw).sum(1, keepdim=True).mean(2, keepdim=True) + fn = torch.rsqrt(fn + eps) + aw = w[f"{p}.affine_weight"][torch.tensor(lc)].view(1, nsph, C) + out = x * (fn * aw) + ab = w[f"{p}.affine_bias"].view(1, 1, C) + return torch.cat([out[:, 0:1, :] + ab, out[:, 1:, :]], 1) + + +def gate(gating, x, lmax, mmax, H): + N = x.shape[0] + ei = _expand_index_m_prime(lmax, mmax) + g = torch.sigmoid(gating).view(N, lmax, H)[:, ei, :] + return torch.cat([F.silu(x[:, 0:1, :]), x[:, 1:, :] * g], 1) + + +def grid_atomwise(x, w, p, to_grid, from_grid): + b, a, nsph = to_grid.shape + tg, fg = to_grid.reshape(b * a, nsph), from_grid.reshape(b * a, nsph) + g = torch.einsum("pi,nic->npc", tg, x) + g = F.silu(g @ w[f"{p}.grid_mlp.0.weight"].T) + g = F.silu(g @ w[f"{p}.grid_mlp.2.weight"].T) + g = g @ w[f"{p}.grid_mlp.4.weight"].T + return torch.einsum("pi,npc->nic", fg, g) + + +def edgewise(x, w, p, cfg, wigner, winv, x_edge, envelope, edge_index, N): + C, H = cfg["sphere_channels"], cfg["hidden_channels"] + lmax, mmax = cfg["lmax"], cfg["mmax"] + nsph = (lmax + 1) ** 2 + src, tgt = edge_index[0], edge_index[1] + m = torch.cat([x[src], x[tgt]], dim=2) + m = torch.bmm(wigner, m) + m, gating = so2(m, x_edge, w, f"{p}.so2_conv_1", lmax, mmax, 2 * C, H, lmax * H) + m = gate(gating, m, lmax, mmax, H) + m = so2(m, x_edge, w, f"{p}.so2_conv_2", lmax, mmax, H, C, 0) + m = m * envelope + m = torch.bmm(winv, m) + out = torch.zeros(N, nsph, C, dtype=m.dtype) + out.index_add_(0, tgt, m) + return out + + +def backbone(w, cfg, x_init, wigner, winv, x_edge, envelope, sys_emb, edge_index, + to_grid, from_grid): + C, lmax = cfg["sphere_channels"], cfg["lmax"] + N = x_init.shape[0] + x = x_init + for i in range(cfg["num_layers"]): + p = f"blocks.{i}" + x_res = x + n = rms_norm_sh(x, w, f"{p}.norm_1", lmax, C) + n = torch.cat([n[:, 0:1, :] + sys_emb.view(N, 1, C), n[:, 1:, :]], 1) + x = edgewise(n, w, f"{p}.edge_wise", cfg, wigner, winv, x_edge, envelope, edge_index, N) + x_res + x_res = x + n = rms_norm_sh(x, w, f"{p}.norm_2", lmax, C) + x = grid_atomwise(n, w, f"{p}.atom_wise", to_grid, from_grid) + x_res + return rms_norm_sh(x, w, "norm", lmax, C) + + +def energy(node_emb, w): + h = node_emb[:, 0, :] + h = F.silu(F.linear(h, w["energy_block.0.weight"], w["energy_block.0.bias"])) + h = F.silu(F.linear(h, w["energy_block.2.weight"], w["energy_block.2.bias"])) + h = F.linear(h, w["energy_block.4.weight"], w["energy_block.4.bias"]) + return h.sum() diff --git a/tests/test_batch.py b/tests/test_batch.py new file mode 100644 index 0000000..dd5d72e --- /dev/null +++ b/tests/test_batch.py @@ -0,0 +1,227 @@ +"""Disjoint-union (block-diagonal) batching parity. + +The whole correctness claim of batching is: evaluating K systems as one concatenated +block-diagonal graph gives the *same* per-system energies and forces as evaluating each system +separately. Block-diagonality means every backbone op stays within-system, so the only new piece +is the segment-sum energy readout (and forces need no change). The mechanism is pinned on the +tiny random-weight golden (device, no fairchem needed); ``test_batched_vs_fairchem`` closes the +loop against fairchem's OWN batched merged inference on the real uma-s-1 checkpoint. +""" +import os +import pathlib + +import numpy as np +import pytest +import torch + +from tt_atom.model import Backbone +from tt_atom.geometry import HostGeometry, radius_graph +from tt_atom import forces, disjoint +from util import pcc + +# real-weight batched parity (skipped unless the merged bundle + fairchem batched golden exist) +BUNDLE = os.environ.get("TTATOM_BUNDLE", str(pathlib.Path.home() / ".ttatom_run/uma_s_ethanol.npz")) +BATCH_GOLDEN = os.environ.get( + "TTATOM_BATCH_GOLDEN", + str(pathlib.Path.home() / ".ttatom_run/goldens_real/batch_ethanol_omol.npz")) + + +def _build(golden, device): + cfg = dict(golden.config) + w = golden.w() + bb = Backbone(w, device, cfg, golden.host("to_grid_mat"), golden.host("from_grid_mat")) + geo = HostGeometry(w, cfg, golden.host("to_m"), golden.host("gauss_offset"), + golden.host("gauss_coeff"), gamma=0.0) + return bb, geo, w, cfg + + +def _systems(golden, k, jitter): + """K variants of the tiny system: same atoms, deterministically jittered positions.""" + pos0 = golden.inp("pos").float() + Z = golden.inp("atomic_numbers").long() + g = torch.Generator().manual_seed(0) + out = [] + for i in range(k): + dp = (0.0 if i == 0 else jitter) * torch.randn(pos0.shape, generator=g) + out.append(dict(pos=pos0 + dp, Z=Z, charge=0.0, spin=0.0)) + return out + + +def test_assemble_block_diagonal(golden): + """Assembly concatenates atoms/edges with per-system offsets and a correct batch index.""" + w = golden.w() + cfg = dict(golden.config) + systems = _systems(golden, 3, jitter=0.02) + bg = disjoint.assemble(systems, cfg["cutoff"], w, cfg["sphere_channels"], task=cfg.get("task", "omat")) + n = systems[0]["Z"].shape[0] + assert bg.natoms == [n, n, n] + assert bg.pos.shape[0] == 3 * n + assert torch.equal(bg.batch, torch.arange(3).repeat_interleave(n)) + # edges of system k must reference only system k's node block + for k in range(3): + m = bg.batch[bg.edge_index[0]] == k + assert torch.equal(bg.batch[bg.edge_index[0][m]], bg.batch[bg.edge_index[1][m]]) + seg = bg.segment_matrix() + assert torch.equal(seg.sum(0), torch.ones(3 * n)) # each atom in exactly one system + assert torch.equal(seg.sum(1), torch.tensor([float(n)] * 3)) + + +def test_batched_equals_separate(golden, device): + """Batched per-system energies and forces == evaluating each system separately.""" + bb, geo, w, cfg = _build(golden, device) + cutoff, C = cfg["cutoff"], cfg["sphere_channels"] + task = cfg.get("task", "omat") + systems = _systems(golden, 4, jitter=0.03) + bg = disjoint.assemble(systems, cutoff, w, C, task=task) + + E_batch, F_batch = forces.energy_and_forces_batch(bb, geo, bg) + + # separate baseline: identical inputs, one system at a time + off = 0 + E_sep, F_sep = [], [] + for k, s in enumerate(systems): + pos, Z = s["pos"], s["Z"] + n = Z.shape[0] + ei, shift = radius_graph(pos, cutoff) + se = bg.sys_emb[off:off + n] + E, F = forces.energy_and_forces(bb, geo, pos, Z, ei, se, edge_cell_shift=shift) + E_sep.append(E) + F_sep.append(F) + off += n + E_sep = torch.tensor(E_sep) + F_sep = torch.cat(F_sep, dim=0) + + # Block-diagonal batching leaves every backbone op within-system, so batched == separate up + # to bf16 rounding only: the batched forward uses the linear O(E) gather+reduce scatter (a + # K>1 batch is block-diagonal, so the dense one-hot S[N,E] would be O(K^2) off-diagonal zeros) + # while the separate baseline keeps the dense matmul, and the energy readout accumulates over a + # larger E — so the tile/accumulation order, and thus the last bf16 bit, differs. On the tiny + # random-weight model energies are O(1), where + # one bf16 ULP is ~8e-3; agreement at that level is exact-to-precision. (PCC>0.999 / rel<1e-3 + # is asserted on the real, large-magnitude energies vs fairchem in test_realweight.py.) + assert (E_batch - E_sep).abs().max() < 2e-2, f"energy maxdiff {(E_batch - E_sep).abs().max()}" + assert pcc(F_batch, F_sep) > 0.99, f"force PCC {pcc(F_batch, F_sep)}" + f_err = (F_batch - F_sep).abs().max() + assert f_err < 2e-2, f"force max abs diff {f_err}" + + +def test_batch_forces_linear_scatter(golden, device): + """A K>1 disjoint-union batch must use the linear O(E) scatter regardless of node count (the + block-diagonal dense one-hot is O(K^2) off-diagonal zeros); a 1-system batch keeps the + single-system node-count threshold. Guards the batched-throughput optimization.""" + from tt_atom.model import GraphContext, SCATTER_LINEAR_THRESHOLD + + bb, geo, w, cfg = _build(golden, device) + cutoff, C = cfg["cutoff"], cfg["sphere_channels"] + task = cfg.get("task", "omat") + + captured = {} + orig_init = GraphContext.__init__ + + def spy(self, *a, **kw): + orig_init(self, *a, **kw) + captured["linear"] = self.linear_scatter + captured["N"] = self.N + + GraphContext.__init__ = spy + try: + bg2 = disjoint.assemble(_systems(golden, 2, jitter=0.03), cutoff, w, C, task=task) + forces.energy_and_forces_batch(bb, geo, bg2, compute_forces=False) + assert captured["N"] <= SCATTER_LINEAR_THRESHOLD, "test premise: batch below threshold" + assert captured["linear"] is True, "K>1 batch must force the linear scatter" + + bg1 = disjoint.assemble(_systems(golden, 1, jitter=0.0), cutoff, w, C, task=task) + forces.energy_and_forces_batch(bb, geo, bg1, compute_forces=False) + assert captured["linear"] is False, "K=1 batch (below threshold) keeps the dense scatter" + finally: + GraphContext.__init__ = orig_init + + +def test_batched_identical_copies(golden, device): + """A batch of identical copies must yield identical per-system energies (segment-sum sanity).""" + bb, geo, w, cfg = _build(golden, device) + systems = _systems(golden, 5, jitter=0.0) # all copies identical + bg = disjoint.assemble(systems, cfg["cutoff"], w, cfg["sphere_channels"], + task=cfg.get("task", "omat")) + E_batch, _ = forces.energy_and_forces_batch(bb, geo, bg, compute_forces=False) + assert (E_batch - E_batch[0]).abs().max() < 1e-3, f"copies disagree: {E_batch}" + + +def test_traced_batch_matches_eager(golden, device): + """The trace-replayed batched forward (``evaluate_batch(trace=True)`` engine) must be BIT-EXACT + vs the eager batched forward: a trace only removes host dispatch, it is the same device op + stream (and both use the K>1 linear scatter + segment-sum readout). Also exercises replay: a + second call on the same topology must match too. Guards the batched-MD throughput path.""" + from tt_atom.trace import TracedEngine + + bb, geo, w, cfg = _build(golden, device) + systems = _systems(golden, 4, jitter=0.03) + bg = disjoint.assemble(systems, cfg["cutoff"], w, cfg["sphere_channels"], + task=cfg.get("task", "omat")) + E_eager, F_eager = forces.energy_and_forces_batch(bb, geo, bg) + + eng = TracedEngine(bb, geo, bg.Z, bg.edge_index, bg.sys_emb, edge_cell_shift=bg.cell_shift, + seg=bg.segment_matrix(), linear_scatter=True) + E_cap, F_cap = eng(bg.pos) # capture step (records + replays) + E_replay, F_replay = eng(bg.pos) # pure replay + eng.close() + + assert (E_cap - E_eager).abs().max() == 0, f"traced E != eager: {E_cap} vs {E_eager}" + assert (F_cap - F_eager).abs().max() == 0, "traced forces != eager" + assert torch.equal(E_cap, E_replay) and torch.equal(F_cap, F_replay), "replay not deterministic" + + +@pytest.mark.skipif(not (pathlib.Path(BUNDLE).exists() and pathlib.Path(BATCH_GOLDEN).exists()), + reason="real merged bundle or fairchem batched golden not present") +def test_batched_vs_fairchem(device): + """Real uma-s-1: TT-Atom disjoint-union batched E+F vs fairchem's OWN batched merged inference + (Batch.from_data_list) on a same-composition conformer batch — E rel<1e-3, F PCC>0.99.""" + from ase import Atoms + from tt_atom import TTAtomCalculator + + d = np.load(BATCH_GOLDEN) + charge, spin = float(d["charge"][0]), float(d["spin"][0]) + natoms = d["natoms"].tolist() + Z, pos = d["Z"], d["pos"] + E_ref, F_ref = d["energy"].astype(np.float64), d["forces"].astype(np.float64) + + systems, off = [], 0 + for n in natoms: + a = Atoms(numbers=Z[off:off + n], positions=pos[off:off + n]) + a.info.update(charge=charge, spin=spin) + systems.append(a) + off += n + + calc = TTAtomCalculator(BUNDLE, device=device) + res = calc.evaluate_batch(systems) + E = res["energy"] + F = np.concatenate(res["forces"], axis=0) + + e_rel = np.abs(E - E_ref).max() / (np.abs(E_ref).max() + 1e-6) + fp = pcc(F, F_ref) + assert e_rel < 1e-3, f"batched energy rel err {e_rel:.2e} (E={E[:3]} vs {E_ref[:3]})" + assert fp > 0.99, f"batched force PCC {fp:.4f}" + + +real_bundle = pytest.mark.skipif( + not pathlib.Path(BUNDLE).exists(), + reason=f"real uma-s-1 bundle not found at {BUNDLE}") + + +@real_bundle +def test_evaluate_batch_rejects_composition_charge_spin_mismatch(device): + """A merged bundle bakes the MoLE routing for one (composition, charge, spin); evaluate_batch + must reject a batch that mixes them rather than silently returning wrong energies. The bundle + is ethanol (C2H6O), merged at charge=0, spin=1.""" + from ase.build import molecule + + from tt_atom import TTAtomCalculator + + calc = TTAtomCalculator(BUNDLE, device=device) + good = molecule("CH3CH2OH"); good.info.update(charge=0, spin=1) # matches the bundle + water = molecule("H2O"); water.info.update(charge=0, spin=1) # wrong composition + with pytest.raises(ValueError, match="reduced composition"): + calc.evaluate_batch([good, water]) + bad_cs = molecule("CH3CH2OH"); bad_cs.info.update(charge=0, spin=0) # wrong spin + with pytest.raises(ValueError, match="merged for"): + calc.evaluate_batch([good, bad_cs]) diff --git a/tests/test_calculator.py b/tests/test_calculator.py new file mode 100644 index 0000000..b891327 --- /dev/null +++ b/tests/test_calculator.py @@ -0,0 +1,34 @@ +"""The ASE calculator runs, gives finite energy/forces, and a relaxation converges on device.""" +import pathlib + +import numpy as np +from ase.build import molecule +from ase.optimize import FIRE + +from tt_atom.calculator import TTAtomCalculator +from tt_atom.weights import WeightBundle + +BUNDLE = pathlib.Path(__file__).parent.parent / "examples" / "model_tiny_demo.npz" + + +def test_coverage(): + ok, missing, n = WeightBundle.load(str(BUNDLE)).verify_coverage() + assert ok, f"missing weight keys: {missing}" + assert n > 0 + + +def test_calculator_relaxation(device): + calc = TTAtomCalculator(str(BUNDLE), device=device) + atoms = molecule("CH3CH2OH") + atoms.info["charge"] = 0 + atoms.info["spin"] = 0 + atoms.rattle(stdev=0.05, seed=0) + atoms.calc = calc + + e0 = atoms.get_potential_energy() + assert np.isfinite(e0) + FIRE(atoms, logfile=None).run(fmax=0.05, steps=200) + fmax = float((atoms.get_forces() ** 2).sum(1).max() ** 0.5) + e1 = atoms.get_potential_energy() + assert fmax <= 0.05, f"did not converge, fmax={fmax}" + assert e1 <= e0 + 1e-3, f"energy increased: {e0} -> {e1}" diff --git a/tests/test_edgewise.py b/tests/test_edgewise.py new file mode 100644 index 0000000..c363c81 --- /dev/null +++ b/tests/test_edgewise.py @@ -0,0 +1,53 @@ +"""Per-module parity: Edgewise message block on TT vs the fairchem golden (PCC >= 0.98).""" +import ttnn + +from tt_atom.edgewise import Edgewise +from tt_atom.model import GraphContext +from util import pcc + + +def _graph(golden, device): + return GraphContext( + device, + edge_index=golden.inp("edge_index"), + wigner=golden.host("wigner"), wigner_inv=golden.host("wigner_inv"), + x_edge=golden.host("x_edge"), edge_envelope=golden.host("edge_envelope"), + num_nodes=golden.act("block0.edgewise.in0").shape[0], + ) + + +def test_edgewise(golden, device): + cfg = golden.config + ew = Edgewise(golden.w(), "blocks.0.edge_wise", device, + sphere_channels=cfg["sphere_channels"], hidden_channels=cfg["hidden_channels"], + lmax=cfg["lmax"], mmax=cfg["mmax"]) + graph = _graph(golden, device) + x = ttnn.from_torch(golden.act("block0.edgewise.in0"), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + o = ttnn.to_torch(ew(x, graph)).float() + p = pcc(o, golden.act("block0.edgewise.out0")) + assert p >= 0.98, f"edgewise PCC {p}" + + +def test_edgewise_linear_scatter(golden, device): + """Fix B: the linear O(E) gather+reduce scatter (large-N path) matches the dense one-hot + matmul (small-N path) on the same system. Force it on by dropping the size threshold.""" + import tt_atom.model as M + + cfg = golden.config + ew = Edgewise(golden.w(), "blocks.0.edge_wise", device, + sphere_channels=cfg["sphere_channels"], hidden_channels=cfg["hidden_channels"], + lmax=cfg["lmax"], mmax=cfg["mmax"]) + x = ttnn.from_torch(golden.act("block0.edgewise.in0"), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + dense = ttnn.to_torch(ew(x, _graph(golden, device))).float() # default threshold -> dense + old = M.SCATTER_LINEAR_THRESHOLD + M.SCATTER_LINEAR_THRESHOLD = 0 # force linear + try: + g_lin = _graph(golden, device) + assert g_lin.linear_scatter + lin = ttnn.to_torch(ew(x, g_lin)).float() + finally: + M.SCATTER_LINEAR_THRESHOLD = old + p = pcc(lin, dense) + assert p >= 0.999, f"linear-vs-dense scatter PCC {p}" diff --git a/tests/test_end2end.py b/tests/test_end2end.py new file mode 100644 index 0000000..8109c90 --- /dev/null +++ b/tests/test_end2end.py @@ -0,0 +1,43 @@ +"""End-to-end parity from raw positions: host geometry + device forward + analytic forces. + +Validates the full production path (``tt_atom.forces.energy_and_forces``) against the fairchem +golden — energy and the conservative analytic force ``F = -dE/dpos`` (NOT finite differences).""" +import torch + +from tt_atom.model import Backbone +from tt_atom.geometry import HostGeometry +from tt_atom import forces +from util import pcc + + +def _build(golden, device): + cfg = dict(golden.config) + w = golden.w() + bb = Backbone(w, device, cfg, golden.host("to_grid_mat"), golden.host("from_grid_mat")) + geo = HostGeometry(w, cfg, golden.host("to_m"), golden.host("gauss_offset"), + golden.host("gauss_coeff"), gamma=0.0) + return bb, geo + + +def test_geometry_forward(golden, device): + # gamma-independent geometric terms must match the fairchem golden exactly + _, geo = _build(golden, device) + t = geo(golden.inp("pos").float(), golden.inp("atomic_numbers").long(), + golden.inp("edge_index").long(), golden.host("sys_node_embedding")) + assert pcc(t["edge_distance"], golden.host("edge_distance")) >= 0.999 + assert pcc(t["x_edge"], golden.host("x_edge")) >= 0.999 + assert pcc(t["edge_envelope"], golden.host("edge_envelope")) >= 0.999 + + +def test_energy_and_forces(golden, device): + bb, geo = _build(golden, device) + E, F = forces.energy_and_forces( + bb, geo, golden.inp("pos").float(), golden.inp("atomic_numbers").long(), + golden.inp("edge_index").long(), golden.host("sys_node_embedding")) + Eref = float(golden.out("energy").reshape(-1)[0]) + Fref = golden.out("forces") + assert abs(E - Eref) / (abs(Eref) + 1e-6) < 0.05, f"energy {E} vs {Eref}" + p = pcc(F, Fref) + cos = float(torch.nn.functional.cosine_similarity(F.reshape(1, -1), Fref.reshape(1, -1))) + assert p >= 0.98, f"force PCC {p}" + assert cos >= 0.98, f"force cosine {cos}" diff --git a/tests/test_forces.py b/tests/test_forces.py new file mode 100644 index 0000000..4a1e4d9 --- /dev/null +++ b/tests/test_forces.py @@ -0,0 +1,82 @@ +"""Analytic-force VJP parity: the on-device reverse pass must reproduce, to PCC >= 0.98, the +adjoints that ``torch.autograd`` produces on the bit-exact PyTorch mirror of the forward. + +This isolates the *device backward* (the hard, hand-written part). The remaining host +geometric Jacobian ``d(geometric)/dpos`` is exercised by the end-to-end force test once +``tt_atom/geometry.py`` lands. +""" +import torch +import ttnn + +import mirror +from tt_atom.model import Backbone, GraphContext +from tt_atom import forces +from util import pcc + + +def _leaves(golden): + xi = golden.host("x_message_init") + return dict(xi=xi, wig=golden.host("wigner"), winv=golden.host("wigner_inv"), + xe=golden.host("x_edge"), env=golden.host("edge_envelope")) + + +def test_backbone_vjp(golden, device): + cfg = golden.config + w = golden.w() + se = golden.host("sys_node_embedding") + ei = golden.inp("edge_index").long() + tg, fg = golden.host("to_grid_mat"), golden.host("from_grid_mat") + N = se.shape[0] + + # oracle: autograd through the PyTorch mirror + lv = {k: v.clone().requires_grad_() for k, v in _leaves(golden).items()} + ne = mirror.backbone(w, cfg, lv["xi"], lv["wig"], lv["winv"], lv["xe"], lv["env"], se, ei, tg, fg) + mirror.energy(ne, w).backward() + + # device forward + reverse VJP + bb = Backbone(w, device, cfg, tg, fg) + graph = GraphContext(device, edge_index=golden.inp("edge_index"), wigner=lv["wig"].detach(), + wigner_inv=lv["winv"].detach(), x_edge=lv["xe"].detach(), + edge_envelope=lv["env"].detach(), num_nodes=N) + se3 = ttnn.from_torch(se.reshape(N, 1, se.shape[1]), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + xi = ttnn.from_torch(lv["xi"].detach(), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + node_emb = bb.node_embedding(xi, graph, se3) + acc = forces.backbone_bw(bb, graph, node_emb) + + # Rotation adjoints are now per-edge sparse coefficients (rotation.py). Compare only the + # structural-nonzero pattern: off-pattern dE/dW entries are nonzero but irrelevant to the + # force (W is structurally zero there for *all* directions, so dW/dpos == 0). + nsph = graph.nsph + + def on_pattern(g_coef_key, ij, ref_grad): + g_coef = ttnn.to_torch(acc[g_coef_key]).float() + ref = torch.stack([ref_grad[:, i, j] for (i, j) in ij], dim=1) # [E, nnz] + return pcc(g_coef, ref) + + assert pcc(ttnn.to_torch(acc["x_init"]).float(), lv["xi"].grad) >= 0.98 + assert on_pattern("rot_fwd", graph.rot_fwd_ij, lv["wig"].grad) >= 0.98 + assert on_pattern("rot_inv", graph.rot_inv_ij, lv["winv"].grad) >= 0.98 + assert pcc(ttnn.to_torch(acc["envelope"]).float().reshape(-1, 1, 1), lv["env"].grad) >= 0.98 + + # host radial finish: g_rad (radial-MLP output adjoint) -> g_x_edge + xel = _leaves(golden)["xe"].clone().requires_grad_() + for conv, grad in acc["g_rad"]: + g = ttnn.to_torch(grad).float() + # the fused SO2 radial emits a duplicated (real|imag) mult [E, nsph*Cin] via a duplicated + # net.6 weight; its adjoint g_rad is at that duplicated output. The mirror radial_mlp emits + # the compact [E, sum rad_sizes] output, so collapse the duplicated rows (== what the fused + # rad.bw's dup weight does inside its matmul) before comparing. + di = getattr(conv.rad, "_dup_index", None) + if di is not None: + gc = torch.zeros(g.shape[0], max(di) + 1, dtype=g.dtype) + gc.index_add_(1, torch.as_tensor(di, dtype=torch.long), g) + g = gc + mirror.radial_mlp(xel, w, conv.rad_prefix).backward(g) + assert pcc(xel.grad, lv["xe"].grad) >= 0.98 + + # device radial finish: acc["x_edge"] (the on-device RadialMLP.bw path) must match the fp64 + # oracle. This is the device backward the perf pass introduced; a bf16 radial VJP mis-directs + # forces on OOD geometries (regressed el_Sn_cmp: 230 meV/A, PCC 0.35), so hold it tight — + # the default fp32 backward matches the oracle, a bf16 one would fail this. + assert pcc(ttnn.to_torch(acc["x_edge"]).float(), lv["xe"].grad) >= 0.999 diff --git a/tests/test_from_uma.py b/tests/test_from_uma.py new file mode 100644 index 0000000..aa4b177 --- /dev/null +++ b/tests/test_from_uma.py @@ -0,0 +1,246 @@ +"""Tests for the ``from_uma`` auto-bundle factory + composition cache (tt_atom.bundle_cache). + +Split into two tiers: + * pure cache logic (composition hashing, cache-path shape, refenv resolution/error) — always run, + no device / no fairchem; + * device tiers gated on availability: the cached fast path + factory-vs-direct parity use a real + uma-s-1 golden bundle (skip cleanly if absent), and the auto-build-vs-manual-export parity uses + the reference (fairchem) env + the gated UMA checkpoint (skip cleanly if either is absent). + +Nothing here commits or requires new weights; it reuses the out-of-repo real goldens. +""" +from __future__ import annotations + +import glob +import json +import os +import pathlib +import shutil +import subprocess + +import numpy as np +import pytest + +from tt_atom import bundle_cache as BC + +REAL_GOLDEN = pathlib.Path( + os.environ.get("TTATOM_REAL_GOLDEN", pathlib.Path.home() / ".ttatom_run/goldens_real/ethanol_omol.npz") +) +HF_CKPT = next(iter(glob.glob(str(pathlib.Path.home() / ".cache/huggingface/**/uma-s-1.pt"), + recursive=True)), None) + + +def _default_refenv(): + p = pathlib.Path.home() / ".ttatom_run/refenv/bin/python" + return str(p) if p.exists() else None + + +# ------------------------------------------------------------------ pure cache logic (no device) + +def test_reduced_composition_is_scale_invariant(): + from ase import Atoms + + h2o = Atoms("H2O", positions=[[0, 0, 0], [1, 0, 0], [0, 1, 0]]) + h4o2 = Atoms("H4O2", positions=[[i, 0, 0] for i in range(6)]) + assert BC.reduced_composition(h2o.numbers) == ((1, 2), (8, 1)) + assert BC.reduced_composition(h2o.numbers) == BC.reduced_composition(h4o2.numbers) + assert BC.composition_hash(h2o.numbers) == BC.composition_hash(h4o2.numbers) + assert BC.formula(h2o.numbers) == "H2O" + + +def test_composition_hash_is_order_independent(): + a = [6, 1, 1, 1, 1] # CH4 + b = [1, 1, 6, 1, 1] # same atoms, permuted + assert BC.composition_hash(a) == BC.composition_hash(b) + # a genuinely different composition hashes differently + assert BC.composition_hash([6, 1, 1, 1, 1]) != BC.composition_hash([6, 1, 1, 1]) + + +def test_bundle_path_shape(tmp_path): + p = BC.bundle_path("uma-s-1", "omol", [1, 1, 8], charge=0, spin=1, cache_dir=tmp_path) + assert p.parent == tmp_path + assert p.name.startswith("uma-s-1_omol_") and p.name.endswith("_c0_s1.npz") + # charge/spin land in the name so distinct charge states never collide + p2 = BC.bundle_path("uma-s-1", "omol", [1, 1, 8], charge=-1, spin=2, cache_dir=tmp_path) + assert p2.name.endswith("_c-1_s2.npz") and p2 != p + + +def test_resolve_refenv_positive_when_default_present(): + if _default_refenv() is None and not os.environ.get("TT_ATOM_REFENV"): + pytest.skip("no reference env installed on this machine") + assert pathlib.Path(BC.resolve_refenv()).exists() + + +def test_resolve_refenv_errors_clearly(monkeypatch, tmp_path): + monkeypatch.delenv("TT_ATOM_REFENV", raising=False) + monkeypatch.setattr(BC.pathlib.Path, "home", staticmethod(lambda: tmp_path)) + with pytest.raises(RuntimeError) as ei: + BC.resolve_refenv() + msg = str(ei.value) + assert "TT_ATOM_REFENV" in msg and "refenv" in msg # actionable, names the knobs + + +def test_from_uma_requires_atoms(): + from tt_atom import TTAtomCalculator + + with pytest.raises(ValueError, match="needs `atoms`"): + TTAtomCalculator.from_uma(atoms=None) + + +def test_infer_task_from_periodicity(): + from ase import Atoms + from ase.build import molecule + + assert BC.infer_task(molecule("H2O")) == "omol" # aperiodic -> molecules + bulk = Atoms("Si2", positions=[[0, 0, 0], [1.4, 1.4, 1.4]], cell=[5.4] * 3, pbc=True) + assert BC.infer_task(bulk) == "omat" # fully periodic -> materials + + +def test_uma_is_exported_and_delegates(): + import tt_atom + + assert "UMA" in tt_atom.__all__ + assert callable(tt_atom.UMA) + + +# ------------------------------------------------------------------ device: cached fast path + parity + +real_golden = pytest.mark.skipif( + not REAL_GOLDEN.exists(), + reason=f"real uma-s-1 golden not found at {REAL_GOLDEN} (UMA checkpoint not available)", +) + + +def _atoms_from_golden(d): + from ase import Atoms + + numbers = d["in@atomic_numbers"] + pos = d["in@pos"] + atoms = Atoms(numbers=numbers, positions=pos) + atoms.info["charge"] = float(d["in@charge"][0]) + atoms.info["spin"] = float(d["in@spin"][0]) + return atoms + + +@real_golden +def test_cached_fast_path_needs_no_refenv_and_matches_direct(tmp_path, device, monkeypatch): + """A cache hit loads without fairchem and yields exactly the same energy as constructing the + calculator directly from the same bundle — the factory-vs-direct parity claim.""" + from tt_atom import TTAtomCalculator + + d = np.load(REAL_GOLDEN) + task = json.loads(bytes(d["config"]).decode())["task"] + atoms = _atoms_from_golden(d) + charge, spin = atoms.info["charge"], atoms.info["spin"] + + # seed the cache: copy the golden to exactly the path from_uma will compute for this system + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + target = BC.bundle_path("uma-s-1", task, atoms.numbers, charge, spin, cache_dir=cache_dir) + shutil.copyfile(REAL_GOLDEN, target) + + # point the refenv at nothing: if the fast path tried to build, this would raise + monkeypatch.setenv("TT_ATOM_REFENV", "/nonexistent/python") + + calc = TTAtomCalculator.from_uma(task_name=task, atoms=atoms, charge=charge, spin=spin, + refenv="/nonexistent/python", cache_dir=str(cache_dir), + device=device) + atoms.calc = calc + e_factory = atoms.get_potential_energy() + + direct = TTAtomCalculator(str(target), device=device) + a2 = _atoms_from_golden(d) + a2.calc = direct + e_direct = a2.get_potential_energy() + + assert e_factory == pytest.approx(e_direct, abs=1e-6), f"{e_factory} vs {e_direct}" + + # the zero-config UMA(atoms) face must reach the identical result (task inferred = omol here) + if task == "omol": + from tt_atom import UMA + + a3 = _atoms_from_golden(d) + a3.calc = UMA(a3, refenv="/nonexistent/python", cache_dir=str(cache_dir), device=device) + assert a3.get_potential_energy() == pytest.approx(e_direct, abs=1e-6) + + +@real_golden +def test_uma_evaluates_with_the_bundles_charge_spin(tmp_path, device, monkeypatch): + """Regression (silent-wrong-answer): the flagship ``UMA(atoms)`` / ``from_uma`` path must + evaluate with the SAME charge/spin the bundle was merged for. The ethanol golden is merged at + charge=0, spin=1 (the UMA omol default); before the fix ``calculate`` fell back to spin=0, + silently disagreeing with the baked MoLE routing. Assert (a) the resolved values are stamped + onto ``atoms.info`` and (b) the energy matches an explicit spin=1 run and *differs* from a + spin=0 run (so the test has teeth — spin genuinely moves this bundle's answer).""" + from ase import Atoms + + from tt_atom import UMA, TTAtomCalculator + + d = np.load(REAL_GOLDEN) + assert int(d["in@charge"][0]) == 0 and int(d["in@spin"][0]) == 1 # golden's merge (charge, spin) + numbers, pos = d["in@atomic_numbers"], d["in@pos"] + + cache_dir = tmp_path / "cache" + cache_dir.mkdir() + target = BC.bundle_path("uma-s-1", "omol", numbers, 0, 1, cache_dir=cache_dir) + shutil.copyfile(REAL_GOLDEN, target) + monkeypatch.setenv("TT_ATOM_REFENV", "/nonexistent/python") # a build attempt would raise + + atoms = Atoms(numbers=numbers, positions=pos) # NO info: exercises the defaults + atoms.calc = UMA(atoms, refenv="/nonexistent/python", cache_dir=str(cache_dir), device=device) + e_uma = atoms.get_potential_energy() + assert atoms.info["charge"] == 0 and atoms.info["spin"] == 1 # resolved values stamped back + + a1 = Atoms(numbers=numbers, positions=pos) + a1.info.update(charge=0, spin=1) + a1.calc = TTAtomCalculator(str(target), device=device) + e_spin1 = a1.get_potential_energy() + + a0 = Atoms(numbers=numbers, positions=pos) + a0.info.update(charge=0, spin=0) + a0.calc = TTAtomCalculator(str(target), device=device) + e_spin0 = a0.get_potential_energy() + + assert e_uma == pytest.approx(e_spin1, abs=1e-6), f"UMA(atoms) used the wrong spin: {e_uma} vs {e_spin1}" + assert abs(e_spin1 - e_spin0) > 1e-4, "spin has no effect on this bundle — test would not catch the bug" + + +refenv_and_ckpt = pytest.mark.skipif( + _default_refenv() is None or HF_CKPT is None, + reason="reference (fairchem) env and/or UMA checkpoint not available for a live merge", +) + + +@refenv_and_ckpt +def test_autobuild_matches_manual_export(tmp_path): + """from_uma's transparent subprocess build produces a bundle numerically identical to a + hand-run tools/export_weights.py merge on the same structure (no device needed).""" + from ase.build import molecule + from ase.io import write + + atoms = molecule("H2O") + atoms.info.update(charge=0, spin=1) + + # (a) auto path: the factory's own build machinery on a fresh cache + cache_dir = tmp_path / "cache" + auto = BC.get_or_build(atoms, model="uma-s-1", task="omol", charge=0, spin=1, + cache_dir=cache_dir, log=False) + + # (b) manual path: invoke the exporter exactly as a user would in the reference env + xyz = tmp_path / "h2o.xyz" + write(str(xyz), atoms) + manual = tmp_path / "manual.npz" + env = dict(os.environ) + env.setdefault("HF_HUB_OFFLINE", "1") + tools = pathlib.Path(BC.__file__).resolve().parent.parent / "tools" / "export_weights.py" + subprocess.run([_default_refenv(), str(tools), "--uma-s-1", "--xyz", str(xyz), + "--task", "omol", "--charge", "0", "--spin", "1", "--out", str(manual)], + check=True, env=env) + + da, dm = np.load(auto), np.load(manual) + assert json.loads(bytes(da["config"]).decode()) == json.loads(bytes(dm["config"]).decode()) + payload = [k for k in da.files if k.startswith(("w@", "scale@", "host@"))] + assert payload, "no weight/scale/buffer arrays in the built bundle" + for k in payload: + assert k in dm.files, f"auto bundle has {k}, manual does not" + assert np.allclose(da[k], dm[k], atol=0, rtol=0), f"array {k} differs auto vs manual" diff --git a/tests/test_grid.py b/tests/test_grid.py new file mode 100644 index 0000000..48eef85 --- /dev/null +++ b/tests/test_grid.py @@ -0,0 +1,18 @@ +"""Per-module parity: GridAtomwise on TT vs the fairchem golden (PCC >= 0.98).""" +import ttnn + +from tt_atom.grid import GridAtomwise +from util import pcc + + +def test_grid_atomwise(golden, device): + w = golden.w() + aw = GridAtomwise( + w, "blocks.0.atom_wise", device, + golden.host("to_grid_mat"), golden.host("from_grid_mat"), + ) + x = ttnn.from_torch(golden.act("block0.atomwise.in0"), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + o = ttnn.to_torch(aw(x)).float() + p = pcc(o, golden.act("block0.atomwise.out0")) + assert p >= 0.98, f"grid atomwise PCC {p}" diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..110f0c5 --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,40 @@ +"""End-to-end backbone parity: node embedding + energy on TT vs the fairchem golden.""" +import ttnn + +from tt_atom.model import Backbone, GraphContext +from util import pcc + + +def _build(golden, device): + cfg = golden.config + bb = Backbone(golden.w(), device, cfg, + golden.host("to_grid_mat"), golden.host("from_grid_mat")) + graph = GraphContext( + device, + edge_index=golden.inp("edge_index"), + wigner=golden.host("wigner"), wigner_inv=golden.host("wigner_inv"), + x_edge=golden.host("x_edge"), edge_envelope=golden.host("edge_envelope"), + num_nodes=golden.host("x_message_init").shape[0], + ) + x_init = ttnn.from_torch(golden.host("x_message_init"), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + sys_emb = golden.host("sys_node_embedding") + sys_emb = ttnn.from_torch(sys_emb.reshape(sys_emb.shape[0], 1, sys_emb.shape[1]), + dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + return bb, graph, x_init, sys_emb + + +def test_node_embedding(golden, device): + bb, graph, x_init, sys_emb = _build(golden, device) + node_emb = bb.node_embedding(x_init, graph, sys_emb) + p = pcc(ttnn.to_torch(node_emb).float(), golden.out("node_embedding")) + assert p >= 0.98, f"node_embedding PCC {p}" + + +def test_energy(golden, device): + bb, graph, x_init, sys_emb = _build(golden, device) + _, energy = bb(x_init, graph, sys_emb) + e = ttnn.to_torch(energy).float().reshape(-1) + ref = golden.out("energy").reshape(-1) + rel = abs(float(e[0]) - float(ref[0])) / (abs(float(ref[0])) + 1e-6) + assert rel < 0.05, f"energy {float(e[0])} vs {float(ref[0])} rel {rel}" diff --git a/tests/test_multicard.py b/tests/test_multicard.py new file mode 100644 index 0000000..bce41a8 --- /dev/null +++ b/tests/test_multicard.py @@ -0,0 +1,56 @@ +"""Multi-card fan-out parity: MultiCard sharding independent systems across N cards must give the +SAME per-system energies, in the SAME order, as running every system sequentially on one card. + +Each system is evaluated by exactly one worker in exactly one code path (the ``_worker`` loop in +``tt_atom/batch.py``) regardless of which card it lands on or how many cards are in the pool, so +sharding is bit-exact by construction: there is no cross-system batching/regrouping (unlike +disjoint-union batching's bf16 accumulation-order sensitivity, see test_batch.py) that could make +the result depend on the shard layout. This test exercises the real queue/dispatch/gather path +(``MultiCard.energies``), not just the per-system compute. +""" +import glob +import pathlib + +import numpy as np +import pytest +from ase.build import molecule + +from tt_atom.batch import MultiCard + +HERE = pathlib.Path(__file__).parent +WEIGHTS = HERE.parent / "examples" / "model_tiny_demo.npz" + + +def _num_devices(): + return len(glob.glob("/dev/tenstorrent/[0-9]*")) + + +pytestmark = [ + pytest.mark.skipif(not WEIGHTS.exists(), reason="examples/model_tiny_demo.npz not present"), + pytest.mark.skipif(_num_devices() < 2, reason="needs >=2 Tenstorrent cards"), +] + + +def _systems(n, seed=0): + base = molecule("CH3CH2OH") + rng = np.random.default_rng(seed) + out = [] + for _ in range(n): + pos = base.get_positions() + rng.normal(scale=0.05, size=base.get_positions().shape) + out.append((pos.astype(np.float32), base.get_atomic_numbers())) + return out + + +@pytest.mark.parametrize("n_systems", [8, 10]) # 10 is uneven across e.g. 3 or 4 cards +def test_sharded_matches_sequential(n_systems): + n_dev = min(4, _num_devices()) + systems = _systems(n_systems) + + with MultiCard(str(WEIGHTS), device_ids=(0,)) as pool: + e_ref, edges_ref = pool.energies(systems) + + with MultiCard(str(WEIGHTS), device_ids=tuple(range(n_dev))) as pool: + e_sharded, edges_sharded = pool.energies(systems) + + assert edges_sharded == edges_ref + assert e_sharded == e_ref, f"sharded != sequential: {e_sharded} vs {e_ref}" diff --git a/tests/test_norm.py b/tests/test_norm.py new file mode 100644 index 0000000..a36228e --- /dev/null +++ b/tests/test_norm.py @@ -0,0 +1,16 @@ +"""Per-module parity: RMS norm SH on TT vs the fairchem golden (PCC >= 0.98).""" +import ttnn + +from tt_atom.norm import RMSNormSH +from util import pcc + + +def test_rms_norm_sh(golden, device): + cfg = golden.config + norm = RMSNormSH(golden.w(), "blocks.0.norm_1", device, + lmax=cfg["lmax"], num_channels=cfg["sphere_channels"]) + x = ttnn.from_torch(golden.act("block0.norm_1.in0"), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + out = ttnn.to_torch(norm(x)).float() + p = pcc(out, golden.act("block0.norm_1.out0")) + assert p >= 0.98, f"rms_norm_sh PCC {p}" diff --git a/tests/test_periodic.py b/tests/test_periodic.py new file mode 100644 index 0000000..7bae752 --- /dev/null +++ b/tests/test_periodic.py @@ -0,0 +1,161 @@ +"""Periodic (PBC) parity tests against the released uma-s-1 checkpoint on periodic systems. + +Materials are UMA's flagship domain, so this is the parity anchor for the periodic path: the +host cell-aware neighbour list (``geometry.radius_graph`` with a cell + pbc) plus the shared +device backbone must reproduce the fairchem oracle. Two periodic tasks are covered when their +(gated, uncommitted) golden bundles are present, else each auto-skips: + + * omat — bulk Si diamond, fully periodic pbc=[T,T,T]; + * oc20 — Cu(100) slab + H adsorbate, mixed pbc=[T,T,F] (catalysis); + * odac — MgO framework fragment, fully periodic (DAC / MOFs); + * omc — solid CO2 (dry ice) molecular crystal, fully periodic. + + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tests/gen_golden_real.py \ + --system bulk --task omat --out ~/.ttatom_run/goldens_real/si_omat.npz + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tests/gen_golden_real.py \ + --system slab --task oc20 --out ~/.ttatom_run/goldens_real/cuh_oc20.npz + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tests/gen_golden_real.py \ + --system mof --task odac --out ~/.ttatom_run/goldens_real/mgo_odac.npz + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tests/gen_golden_real.py \ + --system molcrystal --task omc --out ~/.ttatom_run/goldens_real/co2_omc.npz + PYTHONPATH=~/TT-Atom ~/.ttatom_run/env/bin/python -m pytest tests/test_periodic.py -q + +What is checked per task: + * the periodic neighbour list reproduces fairchem's ``edge_index`` + ``edge_distance_vec`` + exactly (same edge set, matching image offsets) — the graph-construction anchor; + * end-to-end device energy + analytic forces match the fairchem oracle + (energy rel err < 1e-3, force PCC > 0.99). +""" +from __future__ import annotations + +import json +import os +import pathlib + +import numpy as np +import pytest +import torch + +GOLDEN_DIR = pathlib.Path(os.environ.get( + "TTATOM_GOLDEN_DIR", str(pathlib.Path.home() / ".ttatom_run/goldens_real"))) + +# (task label, bundle filename) — parametrized; each case skips if its bundle is absent. +# omat/omc/odac are fully periodic (stress validated too); oc20 is mixed-pbc (stress skips). +PERIODIC_CASES = [("omat", "si_omat.npz"), ("oc20", "cuh_oc20.npz"), + ("odac", "mgo_odac.npz"), ("omc", "co2_omc.npz")] + + +def _pcc(a, b): + a = np.asarray(a, dtype=np.float64).ravel() + b = np.asarray(b, dtype=np.float64).ravel() + if a.std() == 0 and b.std() == 0: + return 1.0 + return float(np.corrcoef(a, b)[0, 1]) + + +def _load(fname): + path = GOLDEN_DIR / fname + if not path.exists(): + pytest.skip(f"periodic golden {path} not found (UMA checkpoint not available)") + return np.load(path), str(path) + + +def _pbc(rg): + return rg["in@pbc"].tolist() if "in@pbc" in rg.files else [True, True, True] + + +@pytest.mark.parametrize("task,fname", PERIODIC_CASES) +def test_neighbour_list_matches_fairchem(task, fname): + """The host cell-aware graph reproduces fairchem's edge set + image offsets exactly.""" + from tt_atom.geometry import radius_graph + + rg, _ = _load(fname) + cfg = json.loads(bytes(rg["config"]).decode()) + assert cfg["task"] == task + pos = torch.from_numpy(rg["in@pos"].copy()).float() + cell = torch.from_numpy(rg["in@cell"].reshape(3, 3).copy()).float() + ei, shift = radius_graph(pos, cfg["cutoff"], cell=cell, pbc=_pbc(rg)) + edge_vec = (pos[ei[0]] - pos[ei[1]] + shift).numpy() + + ei_fc, vec_fc = rg["in@edge_index"], rg["host@edge_distance_vec"] + assert ei.shape[1] == ei_fc.shape[1], f"edge count {ei.shape[1]} vs fairchem {ei_fc.shape[1]}" + + def keyset(ei_, vec_): + return {(int(ei_[0][k]), int(ei_[1][k]), tuple(np.round(vec_[k], 3))) + for k in range(ei_.shape[1])} + + assert keyset(ei.numpy(), edge_vec) == keyset(ei_fc, vec_fc) + + +@pytest.mark.parametrize("task,fname", PERIODIC_CASES) +def test_end_to_end_periodic_energy_forces(task, fname, device): + """Full periodic path — our own neighbour list + device forward + analytic forces — vs the + fairchem oracle (real uma-s-1).""" + from tt_atom import forces as Fmod + from tt_atom.geometry import HostGeometry, csd_embedding, radius_graph + from tt_atom.model import Backbone + from tt_atom.weights import WeightBundle + + rg, path = _load(fname) + rcfg = json.loads(bytes(rg["config"]).decode()) + b = WeightBundle.load(path) + w = b.weights + pos = torch.from_numpy(rg["in@pos"].copy()).float() + Z = torch.from_numpy(rg["in@atomic_numbers"].copy()).long() + cell = torch.from_numpy(rg["in@cell"].reshape(3, 3).copy()).float() + charge = torch.from_numpy(rg["in@charge"].copy()).float() + spin = torch.from_numpy(rg["in@spin"].copy()).float() + + edge_index, edge_cell_shift = radius_graph(pos, rcfg["cutoff"], cell=cell, pbc=_pbc(rg)) + geo = HostGeometry(w, rcfg, b.to_m, b.gauss_offset, b.gauss_coeff) + sys_emb = csd_embedding(w, charge, spin, rcfg["sphere_channels"], + dataset=b.task)[torch.zeros(Z.shape[0], dtype=torch.long)] + bb = Backbone(w, device, rcfg, b.to_grid_mat, b.from_grid_mat) + + E_raw, F_raw = Fmod.energy_and_forces(bb, geo, pos, Z, edge_index, sys_emb, + edge_cell_shift=edge_cell_shift) + E = b.scale_rmsd * E_raw + b.scale_mean + float(b.elem_refs[Z].sum()) + F = b.scale_rmsd * F_raw + + E_oracle = float(rg["out@energy_oracle"][0]) + F_oracle = torch.from_numpy(rg["out@forces_oracle"].copy()).float() + rel = abs(E - E_oracle) / abs(E_oracle) + fpcc = _pcc(F, F_oracle) + assert rel < 1e-3, f"[{task}] energy rel err {rel} (E={E}, oracle={E_oracle})" + assert fpcc > 0.99, f"[{task}] force PCC {fpcc}" + + +@pytest.mark.parametrize("task,fname", PERIODIC_CASES) +def test_stress_matches_fairchem(task, fname, device): + """Stress (virial = symmetrized dE/dstrain, / volume) vs the fairchem oracle on a fully + periodic cell — the anchor for variable-cell relaxation / NPT. Runs through the ASE + ``TTAtomCalculator`` so the Voigt output + normalizer scaling + volume are all exercised. + Skips a mixed-pbc case (stress ill-defined; oracle stored zeros).""" + from ase import Atoms + + from tt_atom.calculator import TTAtomCalculator + from tt_atom.weights import WeightBundle + + rg, path = _load(fname) + if "out@stress_oracle" not in rg.files or not np.any(rg["out@stress_oracle"]): + pytest.skip(f"[{task}] no oracle stress (mixed-pbc or pre-stress golden)") + pbc = _pbc(rg) + if not all(pbc): + pytest.skip(f"[{task}] stress only validated for a fully periodic cell") + + atoms = Atoms( + numbers=rg["in@atomic_numbers"].copy(), + positions=rg["in@pos"].copy(), + cell=rg["in@cell"].reshape(3, 3).copy(), + pbc=pbc, + ) + atoms.info.update(charge=int(rg["in@charge"][0]), spin=int(rg["in@spin"][0])) + calc = TTAtomCalculator(WeightBundle.load(path), device=device) + atoms.calc = calc + stress = atoms.get_stress() # ASE Voigt-6 + S_oracle = rg["out@stress_oracle"] + spcc = _pcc(stress, S_oracle) + maxrel = float(np.max(np.abs(stress - S_oracle) / (np.abs(S_oracle) + 1e-9))) + # stress is noisier than forces (bf16 device); accept fairchem parity by PCC or rel err + assert spcc > 0.99 or maxrel < 1e-2, ( + f"[{task}] stress PCC {spcc}, maxrel {maxrel}\n mine={stress}\n oracle={S_oracle}") diff --git a/tests/test_realweight.py b/tests/test_realweight.py new file mode 100644 index 0000000..8b906a4 --- /dev/null +++ b/tests/test_realweight.py @@ -0,0 +1,139 @@ +"""Real-weight parity tests against the released uma-s-1 checkpoint. + +These run only when a real-weight golden bundle is present (generated in the fairchem refenv by +``tests/gen_golden_real.py`` and stored OUTSIDE the repo, since the UMA checkpoint is gated and +must not be committed). Absent the bundle the whole module auto-skips, so the suite stays green +for anyone without UMA access. + + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tests/gen_golden_real.py \ + --system molecule --task omol --out ~/.ttatom_run/goldens_real/ethanol_omol.npz + PYTHONPATH=~/TT-Atom ~/.ttatom_run/env/bin/python -m pytest tests/test_realweight.py -q + +Point ``TTATOM_REAL_GOLDEN`` at a different bundle to override the default path. + +What is checked (all numbers measured on the p150, real uma-s-1, ethanol/omol): + * MoLE host-merge anchor: the merged plain backbone reproduces the unmerged-MoE fairchem + oracle E+F (the golden records both) to PCC>0.999; + * the device spectral atomwise matches the golden per-module (PCC>=0.98); + * WeightBundle.verify_coverage passes on the real merged bundle; + * end-to-end device energy + analytic forces match the fairchem oracle + (energy rel err < 1e-2, force PCC > 0.99). +""" +from __future__ import annotations + +import json +import os +import pathlib + +import numpy as np +import pytest +import torch + +REAL_GOLDEN = os.environ.get( + "TTATOM_REAL_GOLDEN", str(pathlib.Path.home() / ".ttatom_run/goldens_real/ethanol_omol.npz") +) + +pytestmark = pytest.mark.skipif( + not pathlib.Path(REAL_GOLDEN).exists(), + reason=f"real-weight golden bundle not found at {REAL_GOLDEN} (UMA checkpoint not available)", +) + + +def _pcc(a, b): + a = np.asarray(a, dtype=np.float64).ravel() + b = np.asarray(b, dtype=np.float64).ravel() + if a.std() == 0 and b.std() == 0: + return 1.0 + return float(np.corrcoef(a, b)[0, 1]) + + +@pytest.fixture(scope="module") +def rg(): + return np.load(REAL_GOLDEN) + + +@pytest.fixture(scope="module") +def rcfg(rg): + return json.loads(bytes(rg["config"]).decode()) + + +def _w(rg): + return {k[2:]: torch.from_numpy(rg[k].copy()).float() for k in rg.files if k.startswith("w@")} + + +def _act(rg, name): + return torch.from_numpy(rg[f"a@{name}"].copy()).float() + + +def test_config_is_real_uma_s(rcfg): + assert rcfg["ff_type"] == "spectral" + assert rcfg["num_layers"] == 4 + assert rcfg["lmax"] == 2 and rcfg["mmax"] == 2 + assert rcfg["chg_spin_emb_type"] == "rand_emb" + assert rcfg["task"] == "omol" + + +def test_merge_anchor(rg): + """Host MoLE merge reproduces the unmerged-MoE fairchem oracle to PCC>0.999.""" + Eo = float(rg["out@energy_oracle"][0]) + Em = float(rg["out@energy_merged_oracle"][0]) + Fo = rg["out@forces_oracle"] + Fm = rg["out@forces_merged_oracle"] + assert abs(Em - Eo) / abs(Eo) < 1e-6, f"merge energy rel err {abs(Em - Eo) / abs(Eo)}" + assert _pcc(Fm, Fo) > 0.999, f"merge force PCC {_pcc(Fm, Fo)}" + + +def test_verify_coverage(rg): + from tt_atom.weights import WeightBundle + + b = WeightBundle.load(REAL_GOLDEN) + ok, missing, present = b.verify_coverage() + assert ok, f"missing weight keys: {missing}" + assert b.scale_rmsd > 0 and b.elem_refs is not None and b.task == "omol" + + +def test_spectral_atomwise_module(rg, rcfg, device): + import ttnn + + from tt_atom.spectral import SpectralAtomwise + + w = _w(rg) + sp = SpectralAtomwise(w, "blocks.0.atom_wise", device, + sphere_channels=rcfg["sphere_channels"], hidden_channels=rcfg["hidden_channels"], + lmax=rcfg["lmax"], mmax=rcfg["mmax"]) + x = ttnn.from_torch(_act(rg, "block0.atomwise.in0"), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + o = ttnn.to_torch(sp(x)).float() + assert _pcc(o, _act(rg, "block0.atomwise.out0")) >= 0.98 + + +def test_end_to_end_energy_forces(rg, rcfg, device): + """Full device forward + analytic forces vs the fairchem oracle (real uma-s-1).""" + from tt_atom import forces as Fmod + from tt_atom.geometry import HostGeometry, csd_embedding + from tt_atom.model import Backbone + from tt_atom.weights import WeightBundle + + b = WeightBundle.load(REAL_GOLDEN) + w = b.weights + pos = torch.from_numpy(rg["in@pos"].copy()).float() + Z = torch.from_numpy(rg["in@atomic_numbers"].copy()).long() + edge_index = torch.from_numpy(rg["in@edge_index"].copy()).long() + charge = torch.from_numpy(rg["in@charge"].copy()).float() + spin = torch.from_numpy(rg["in@spin"].copy()).float() + + geo = HostGeometry(w, rcfg, b.to_m, b.gauss_offset, b.gauss_coeff) + sys_emb = csd_embedding(w, charge, spin, rcfg["sphere_channels"], + dataset=b.task)[torch.zeros(Z.shape[0], dtype=torch.long)] + bb = Backbone(w, device, rcfg, b.to_grid_mat, b.from_grid_mat) + + E_raw, F_raw = Fmod.energy_and_forces(bb, geo, pos, Z, edge_index, sys_emb) + E = b.scale_rmsd * E_raw + b.scale_mean + float(b.elem_refs[Z].sum()) + F = b.scale_rmsd * F_raw + + E_oracle = float(rg["out@energy_oracle"][0]) + F_oracle = torch.from_numpy(rg["out@forces_oracle"].copy()).float() + rel = abs(E - E_oracle) / abs(E_oracle) + fpcc = _pcc(F, F_oracle) + assert rel < 1e-2, f"energy rel err {rel} (E={E}, oracle={E_oracle})" + assert fpcc > 0.99, f"force PCC {fpcc}" diff --git a/tests/test_realweight_uma_s_1p2.py b/tests/test_realweight_uma_s_1p2.py new file mode 100644 index 0000000..b1e2008 --- /dev/null +++ b/tests/test_realweight_uma_s_1p2.py @@ -0,0 +1,121 @@ +"""Real-weight parity tests for uma-s-1.2 (charge-balanced channels) against fairchem. + +uma-s-1.2 differs from uma-s-1 by fairchem's ``charge_balanced_channels``: the l=0 charge +channels are re-balanced (a self-adjoint per-system mean-subtraction, plus the charge/natoms +target) after every block. Without it the force PCC on the released checkpoint collapses to +~0.83; with it, parity is restored (>0.99). This module is the on-device regression for that +path, and complements the 757-system CPU-vs-TT screen written up in +``docs/uma-s-1p2-validation.md``. + +Like ``test_realweight.py`` it runs only when a real-weight golden is present (the UMA checkpoint +is gated and must not be committed), so it auto-skips for anyone without access. Generate the +golden in the fairchem refenv, then run it on a card: + + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tests/gen_golden_real.py \ + --system molecule --task omol --ckpt uma-s-1p2 \ + --out ~/.ttatom_run/goldens_real/ethanol_omol_uma_s_1p2.npz + # (or add --ckpt-path /path/to/uma-s-1p2.pt to use a local checkpoint file) + TT_VISIBLE_DEVICES=0 ~/.ttatom_run/venv/bin/python -m pytest tests/test_realweight_uma_s_1p2.py -q + +Point ``TTATOM_REAL_GOLDEN_S1P2`` at a different bundle to override the default path. + +Neutral ethanol/omol still exercises the balancing: the per-system l=0 mean-subtraction runs on +every block (the charge/natoms target is 0 for a neutral system). Charged-system parity (e.g. +[Cu(EDTA)]2-) is covered by the A/B screen in the validation doc. +""" +from __future__ import annotations + +import json +import os +import pathlib + +import numpy as np +import pytest +import torch + +REAL_GOLDEN = os.environ.get( + "TTATOM_REAL_GOLDEN_S1P2", + str(pathlib.Path.home() / ".ttatom_run/goldens_real/ethanol_omol_uma_s_1p2.npz"), +) + +pytestmark = pytest.mark.skipif( + not pathlib.Path(REAL_GOLDEN).exists(), + reason=f"uma-s-1.2 real-weight golden not found at {REAL_GOLDEN} (UMA checkpoint not available)", +) + + +def _pcc(a, b): + a = np.asarray(a, dtype=np.float64).ravel() + b = np.asarray(b, dtype=np.float64).ravel() + if a.std() == 0 and b.std() == 0: + return 1.0 + return float(np.corrcoef(a, b)[0, 1]) + + +@pytest.fixture(scope="module") +def rg(): + return np.load(REAL_GOLDEN) + + +@pytest.fixture(scope="module") +def rcfg(rg): + return json.loads(bytes(rg["config"]).decode()) + + +def test_config_enables_charge_balancing(rcfg): + """The golden must be uma-s-1.2: an s-size spectral model with charge-balanced l=0 channels.""" + assert rcfg["ff_type"] == "spectral" + assert rcfg["num_layers"] == 4 + assert rcfg["lmax"] == 2 and rcfg["mmax"] == 2 + assert rcfg["task"] == "omol" + cs = int(rcfg.get("charge_channel_start", 0)) + ce = int(rcfg.get("charge_channel_end", 0)) + assert cs < ce, f"charge balancing inactive (cs={cs}, ce={ce}); this is not a uma-s-1.2 golden" + + +def test_verify_coverage(rg): + from tt_atom.weights import WeightBundle + + b = WeightBundle.load(REAL_GOLDEN) + ok, missing, present = b.verify_coverage() + assert ok, f"missing weight keys: {missing}" + assert b.scale_rmsd > 0 and b.elem_refs is not None and b.task == "omol" + + +def test_end_to_end_energy_forces(rg, rcfg, device): + """Full device forward + analytic forces vs the fairchem oracle (real uma-s-1.2). + + Exercises the charge-balanced channels end to end: with balancing the force PCC is >0.99; + without it (the pre-1.2 port) it collapses to ~0.83, so this is a genuine regression guard. + """ + from tt_atom import forces as Fmod + from tt_atom.geometry import HostGeometry, csd_embedding + from tt_atom.model import Backbone + from tt_atom.weights import WeightBundle + + b = WeightBundle.load(REAL_GOLDEN) + w = b.weights + pos = torch.from_numpy(rg["in@pos"].copy()).float() + Z = torch.from_numpy(rg["in@atomic_numbers"].copy()).long() + edge_index = torch.from_numpy(rg["in@edge_index"].copy()).long() + charge = torch.from_numpy(rg["in@charge"].copy()).float() + spin = torch.from_numpy(rg["in@spin"].copy()).float() + + geo = HostGeometry(w, rcfg, b.to_m, b.gauss_offset, b.gauss_coeff) + sys_emb = csd_embedding(w, charge, spin, rcfg["sphere_channels"], + dataset=b.task)[torch.zeros(Z.shape[0], dtype=torch.long)] + bb = Backbone(w, device, rcfg, b.to_grid_mat, b.from_grid_mat) + + # pass the golden's charge so the additive charge/natoms target branch of channel balancing is + # exercised (the fairchem oracle balanced to this charge); 0 for the neutral ethanol golden. + E_raw, F_raw = Fmod.energy_and_forces(bb, geo, pos, Z, edge_index, sys_emb, + charge=float(charge.reshape(-1)[0])) + E = b.scale_rmsd * E_raw + b.scale_mean + float(b.elem_refs[Z].sum()) + F = b.scale_rmsd * F_raw + + E_oracle = float(rg["out@energy_oracle"][0]) + F_oracle = torch.from_numpy(rg["out@forces_oracle"].copy()).float() + rel = abs(E - E_oracle) / abs(E_oracle) + fpcc = _pcc(F, F_oracle) + assert rel < 1e-2, f"energy rel err {rel} (E={E}, oracle={E_oracle})" + assert fpcc > 0.99, f"force PCC {fpcc}" diff --git a/tests/test_so2.py b/tests/test_so2.py new file mode 100644 index 0000000..2033834 --- /dev/null +++ b/tests/test_so2.py @@ -0,0 +1,45 @@ +"""Per-module parity: SO2Convolution on TT vs the fairchem golden (PCC >= 0.98).""" +import ttnn + +from tt_atom.so2 import SO2Convolution +from util import pcc + + +def _run(golden, device, prefix, act_prefix, Cin, H, extra): + cfg = golden.config + w = golden.w() + conv = SO2Convolution( + w, prefix, device, + sphere_channels_in=Cin, m_output_channels=H, + lmax=cfg["lmax"], mmax=cfg["mmax"], extra_m0_output_channels=extra, + ) + x = ttnn.from_torch(golden.act(f"{act_prefix}.in0"), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + xe = None + if conv.has_radial: + xe = ttnn.from_torch(golden.act(f"{act_prefix}.in1"), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + out = conv(x, xe) + if extra: + out, extra_t = out + e = ttnn.to_torch(extra_t).float() + assert pcc(e, golden.act(f"{act_prefix}.out1")) >= 0.98 + o = ttnn.to_torch(out).float() + return pcc(o, golden.act(f"{act_prefix}.out0")) + + +def test_so2_conv_1(golden, device): + # Edgewise feeds 2*sphere_channels in; extra_m0 = lmax*hidden gating channels. + cfg = golden.config + Cin = 2 * cfg["sphere_channels"] + extra = cfg["lmax"] * cfg["hidden_channels"] + p = pcc_val = _run(golden, device, "blocks.0.edge_wise.so2_conv_1", + "block0.so2_1", Cin, cfg["hidden_channels"], extra) + assert p >= 0.98, f"so2_conv_1 PCC {p}" + + +def test_so2_conv_2(golden, device): + cfg = golden.config + p = _run(golden, device, "blocks.0.edge_wise.so2_conv_2", + "block0.so2_2", cfg["hidden_channels"], cfg["sphere_channels"], 0) + assert p >= 0.98, f"so2_conv_2 PCC {p}" diff --git a/tests/test_symmetry.py b/tests/test_symmetry.py new file mode 100644 index 0000000..3fbb7a0 --- /dev/null +++ b/tests/test_symmetry.py @@ -0,0 +1,48 @@ +"""Regression for the exact-symmetry wrong-force bug (host-only, no device). + +The legacy ZYZ-Euler edge frame (``geometry._euler_angles``) has a coordinate singularity on the ++-Y axis: ``alpha = atan2(x, z)`` with ``x, z -> 0`` there. ``_Safeatan2``'s ``clamp(min=EPS)`` +backward then *annihilates* the azimuth's position-gradient (denominator clamped, numerator ->0), +so ``d(alpha)/dpos -> 0`` — a degenerate frame derivative. At an exactly-symmetric geometry every +edge sits on that set, so the analytic force (which needs ``d(wigner)/dpos``) was wrong while the +energy (roll-gauge invariant) stayed fine. Fix 1 swaps the frame for fairchem's smooth two-chart +quaternion (``quaternion.wigner_from_edge``), which is finite, orthogonal and non-degenerate on the +axes. + +These are host-only (torch + the coefficient asset — no card, no fairchem, no weight bundle), so the +symmetry gap that the ethanol-only goldens miss is covered in fast CI. Force-level correctness at +symmetry (vs fairchem) is validated end-to-end by the A/B harness (its ``symF`` metric evaluates the +exact-equilibrium geometry) and by the passing device parity suite. +""" +import torch + +from tt_atom import quaternion +from tt_atom.geometry import _euler_angles + +AXES = torch.tensor([[1.0, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, 1], [0, 0, -1]], + dtype=torch.float64) + + +def test_quaternion_wigner_finite_orthogonal_at_axes(): + """Quaternion Wigner-D and its gradient are finite, and D is orthogonal (a valid rotation), for + edges exactly on the coordinate axes — the singular set of the old Euler frame. lmax 2 and 4.""" + for lmax in (2, 4): + kern = quaternion.WignerKernels(lmax) + e = AXES.clone().requires_grad_(True) + W = quaternion.wigner_from_edge(e, lmax, kern, gamma=0.0) + assert torch.isfinite(W).all(), f"non-finite wigner at axes (lmax={lmax})" + g, = torch.autograd.grad(W.sum(), e) + assert torch.isfinite(g).all(), f"non-finite d(wigner)/d(edge) at axes (lmax={lmax})" + eye = torch.eye((lmax + 1) ** 2, dtype=torch.float64) + assert (W @ W.transpose(1, 2) - eye).abs().max() < 1e-10, "wigner not orthogonal on axes" + + +def test_euler_frame_degenerate_at_pole(): + """Guard the ROOT CAUSE: on the +-Y axis the old Euler azimuth gradient is ANNIHILATED + (``clamp`` denominator, vanishing numerator) -> ``d(alpha)/dpos ~= 0``. A degenerate frame + derivative is what corrupted forces at exact symmetry; the quaternion frame (above) is + non-degenerate there instead.""" + pole = torch.tensor([[0.0, 1.0, 0.0], [0.0, -1.0, 0.0]], dtype=torch.float64, requires_grad=True) + _, _, alpha = _euler_angles(pole, 0.0) + g, = torch.autograd.grad(alpha.sum(), pole) + assert g.abs().max() < 1e-9, "expected the Euler azimuth gradient to be annihilated at the pole" diff --git a/tests/test_trace.py b/tests/test_trace.py new file mode 100644 index 0000000..982da63 --- /dev/null +++ b/tests/test_trace.py @@ -0,0 +1,107 @@ +"""Trace path parity: the device-resident, trace-captured engine must return exactly the eager +energy+forces (it only removes host dispatch, never changes the math). Uses the committed +random-weight demo bundle so it runs without a UMA checkpoint. +""" +from __future__ import annotations + +import pathlib + +import numpy as np +from ase.build import molecule + +from tt_atom.calculator import TTAtomCalculator + +DEMO = str(pathlib.Path(__file__).parent.parent / "examples" / "model_tiny_demo.npz") + + +def _pcc(a, b): + return float(np.corrcoef(np.asarray(a).ravel(), np.asarray(b).ravel())[0, 1]) + + +def test_traced_matches_eager(device): + """Traced calculator == eager calculator on energy and forces, over several geometries + (including a re-capture triggered by moving atoms).""" + eager = TTAtomCalculator(DEMO, device=device) + traced = TTAtomCalculator(DEMO, device=device, trace=True) + try: + rng = np.random.default_rng(0) + for k in range(3): + atoms = molecule("CH3CH2OH") + atoms.info.update(charge=0, spin=0) + atoms.positions += rng.normal(scale=0.05, size=atoms.positions.shape) + atoms.calc = eager + Ee, Fe = atoms.get_potential_energy(), atoms.get_forces() + atoms.calc = traced + Et, Ft = atoms.get_potential_energy(), atoms.get_forces() + assert abs(Et - Ee) < 1e-4, f"step {k}: energy {Et} vs {Ee}" + assert _pcc(Ft, Fe) > 0.9999, f"step {k}: force PCC {_pcc(Ft, Fe)}" + assert np.abs(Ft - Fe).max() < 1e-3, f"step {k}: max force diff {np.abs(Ft - Fe).max()}" + finally: + traced.close() + + +def test_traced_recaptures_on_cell_shift_change(device): + """Regression: the traced engine bakes in ``edge_cell_shift``, so a changing cell at *fixed* + edge topology (the NPT / variable-cell step that doesn't cross the cutoff) must trigger a + re-capture — else the replay reuses the stale shift and returns wrong forces. Two cells that + yield the SAME edge_index but different shifts must both match eager.""" + import torch + from ase import Atoms + + from tt_atom.geometry import radius_graph + + eager = TTAtomCalculator(DEMO, device=device) + traced = TTAtomCalculator(DEMO, device=device, trace=True) + cut = eager.cfg["cutoff"] + + def _atoms(cx): + a = Atoms(numbers=[6, 8], positions=[[0, 0, 0], [2.0, 0, 0]], + cell=[cx, 40.0, 40.0], pbc=[True, False, False]) + a.info.update(charge=0, spin=0) + return a + + try: + # the two cells must share edge topology but differ in cell shift (else the test would be + # trivially satisfied by the edge-index re-capture that already existed) + pbc = [True, False, False] + ei1, sh1 = radius_graph(torch.tensor(_atoms(5.1).get_positions(), dtype=torch.float32), cut, + cell=torch.tensor(np.asarray(_atoms(5.1).get_cell()), dtype=torch.float32), + pbc=pbc) + ei2, sh2 = radius_graph(torch.tensor(_atoms(5.2).get_positions(), dtype=torch.float32), cut, + cell=torch.tensor(np.asarray(_atoms(5.2).get_cell()), dtype=torch.float32), + pbc=pbc) + assert torch.equal(ei1, ei2) and not torch.equal(sh1, sh2), "setup: need same edges, diff shift" + + for cx in (5.1, 5.2): # capture at 5.1, then shift-only change to 5.2 + a = _atoms(cx) + a.calc = eager + Ee, Fe = a.get_potential_energy(), a.get_forces() + a.calc = traced + Et, Ft = a.get_potential_energy(), a.get_forces() + assert abs(Et - Ee) < 1e-4, f"cx={cx}: energy {Et} vs {Ee}" + assert _pcc(Ft, Fe) > 0.9999, f"cx={cx}: force PCC {_pcc(Ft, Fe)}" + assert np.abs(Ft - Fe).max() < 1e-3, f"cx={cx}: max force diff {np.abs(Ft - Fe).max()}" + finally: + traced.close() + + +def test_traced_stress_falls_back_to_eager(device): + """trace=True must still deliver stress (via the eager fallback) instead of silently dropping + it — else an ASE variable-cell relaxation with trace=True would hit PropertyNotImplementedError.""" + from ase import Atoms + + a = Atoms(numbers=[6, 8], positions=[[0, 0, 0], [2.0, 0, 0]], + cell=[5.1, 40.0, 40.0], pbc=[True, False, False]) + a.info.update(charge=0, spin=0) + eager = TTAtomCalculator(DEMO, device=device) + traced = TTAtomCalculator(DEMO, device=device, trace=True) + try: + a.calc = eager + s_e, F_e = a.get_stress(), a.get_forces() + a.calc = traced + s_t, F_t = a.get_stress(), a.get_forces() + assert s_t is not None and s_t.shape == (6,) + assert np.abs(s_t - s_e).max() < 1e-4, f"stress {s_t} vs {s_e}" + assert _pcc(F_t, F_e) > 0.9999, f"force PCC {_pcc(F_t, F_e)}" + finally: + traced.close() diff --git a/tests/test_umam.py b/tests/test_umam.py new file mode 100644 index 0000000..cf4ca00 --- /dev/null +++ b/tests/test_umam.py @@ -0,0 +1,69 @@ +"""uma-m is UNSUPPORTED in this build: it must raise a clear error, not silently fall back. + +tt-atom is the custom-kernel-only, highest-performance build for uma-s. uma-m uses spherical- +harmonic coefficient subselection: the 25 SH coefficients of the node representation are reduced +to a 19-dim ``|m|<=mmax`` m-space inside the edgewise SO(2) block, so its Wigner rotation is +RECTANGULAR (node SH 25 <-> reduced m-space 19, W=256), unlike uma-s (square 9<->9). That shape +overflows the fused_rotate kernel's L1 CB budget, and this build has no slow MAC fallback -- so +the rotation raises a clear ``RuntimeError`` naming the unsupported shape. This test anchors that +contract (uma-s is the validated target; uma-m is explicitly out of scope here). + +Golden (gated, uncommitted; the checkpoint is 11 GB so the generator loads a single merged unit): + + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tests/gen_golden_real.py \ + --system molecule --task omol --ckpt uma-m-1p1 --merged-only \ + --out ~/.ttatom_run/goldens_real/ethanol_omol_umam.npz + PYTHONPATH=~/TT-Atom ~/.ttatom_run/env/bin/python -m pytest tests/test_umam.py -q +""" +from __future__ import annotations + +import json +import os +import pathlib + +import numpy as np +import pytest +import torch + +GOLDEN_DIR = pathlib.Path(os.environ.get( + "TTATOM_GOLDEN_DIR", str(pathlib.Path.home() / ".ttatom_run/goldens_real"))) +GOLDEN = GOLDEN_DIR / "ethanol_omol_umam.npz" + + +def _pcc(a, b): + a = np.asarray(a, dtype=np.float64).ravel() + b = np.asarray(b, dtype=np.float64).ravel() + return float(np.corrcoef(a, b)[0, 1]) + + +def test_umam_unsupported_raises(device): + """uma-m-1p1 (lmax=4/mmax=2) must raise a clear RuntimeError: its rectangular reduced-m Wigner + rotation (25<->19, W=256) overflows the fused kernel's L1 budget and this build has no fallback.""" + if not GOLDEN.exists(): + pytest.skip(f"uma-m golden {GOLDEN} not found (checkpoint not available)") + from tt_atom import forces as Fmod + from tt_atom.geometry import HostGeometry, csd_embedding, radius_graph + from tt_atom.model import Backbone + from tt_atom.weights import WeightBundle + + rg = np.load(GOLDEN) + rcfg = json.loads(bytes(rg["config"]).decode()) + assert rcfg["lmax"] == 4 and rcfg["mmax"] == 2, "expected the uma-m lmax=4/mmax=2 config" + b = WeightBundle.load(str(GOLDEN)) + assert b.coefficient_index is not None, "uma-m golden must carry coefficient_index" + w = b.weights + pos = torch.from_numpy(rg["in@pos"].copy()).float() + Z = torch.from_numpy(rg["in@atomic_numbers"].copy()).long() + charge = torch.from_numpy(rg["in@charge"].copy()).float() + spin = torch.from_numpy(rg["in@spin"].copy()).float() + + edge_index, edge_cell_shift = radius_graph(pos, rcfg["cutoff"]) + geo = HostGeometry(w, rcfg, b.to_m, b.gauss_offset, b.gauss_coeff, + coefficient_index=b.coefficient_index) + sys_emb = csd_embedding(w, charge, spin, rcfg["sphere_channels"], + dataset=b.task)[torch.zeros(Z.shape[0], dtype=torch.long)] + bb = Backbone(w, device, rcfg, b.to_grid_mat, b.from_grid_mat) + + with pytest.raises(RuntimeError, match="unsupported in this build"): + Fmod.energy_and_forces(bb, geo, pos, Z, edge_index, sys_emb, + edge_cell_shift=edge_cell_shift) diff --git a/tests/util.py b/tests/util.py new file mode 100644 index 0000000..a6d6216 --- /dev/null +++ b/tests/util.py @@ -0,0 +1,44 @@ +"""Golden-fixture helpers shared by the parity tests.""" +import json +import pathlib + +import numpy as np +import torch + +DATA = pathlib.Path(__file__).parent / "data" + + +def pcc(a, b): + """Pearson correlation of two tensors/arrays, flattened to fp64.""" + a = np.asarray(a, dtype=np.float64).ravel() + b = np.asarray(b, dtype=np.float64).ravel() + if a.std() == 0 and b.std() == 0: + return 1.0 + return float(np.corrcoef(a, b)[0, 1]) + + +class Golden: + """Accessor over a golden npz: weights (``w@``), activations (``a@``), + host terms (``host@``), inputs (``in@``), outputs (``out@``).""" + + def __init__(self, name): + self.d = np.load(DATA / name) + self.config = json.loads(bytes(self.d["config"]).decode()) + + def _t(self, key): + return torch.from_numpy(self.d[key].copy()) + + def w(self): + return {k[2:]: self._t(k).float() for k in self.d.files if k.startswith("w@")} + + def act(self, name): + return self._t(f"a@{name}").float() + + def host(self, name): + return self._t(f"host@{name}").float() + + def inp(self, name): + return self._t(f"in@{name}") + + def out(self, name): + return self._t(f"out@{name}").float() diff --git a/tools/export_weights.py b/tools/export_weights.py new file mode 100644 index 0000000..74b6c37 --- /dev/null +++ b/tools/export_weights.py @@ -0,0 +1,170 @@ +"""Export a TT-Atom weight bundle from a fairchem checkpoint (or random init). + +Run in the *fairchem* environment (fairchem-core, numpy>=2) — NOT the ttnn env — because it +instantiates the reference ``eSCNMDBackbone`` to obtain both the learned ``state_dict`` and the +fixed geometric buffers (Jd, to_m, SO3 grid matrices, gaussian basis) that a bare checkpoint +does not contain. The resulting ``.npz`` is loaded by ``tt_atom.weights.WeightBundle`` in the +ttnn env. This two-step export is what lets the two incompatible numpy worlds coexist. + + # random-weight demo bundle (architecture only, no checkpoint): + ~/.ttatom_run/refenv/bin/python tools/export_weights.py --out model.npz + + # plain backbone checkpoint (no MoLE), fresh energy head: + ~/.ttatom_run/refenv/bin/python tools/export_weights.py --checkpoint plain.pt --out model.npz + + # real released uma-s-1 (gated facebook/UMA; MoLE-merged on host for a fixed composition): + HF_HUB_OFFLINE=1 ~/.ttatom_run/refenv/bin/python tools/export_weights.py \ + --uma-s-1 --molecule CH3CH2OH --task omol --charge 0 --spin 1 --out uma_s_ethanol.npz + +For ``--uma-s-1`` the bundle is composition-specific (MoLE routing is fixed at merge time), so it +is valid for systems with the same reduced composition / charge / spin / dataset — exactly the +constant-composition regime of a relaxation or MD run. The exported bundle carries the per-task +energy normalizer (``scale@*``) and the real energy head. No weights are committed/redistributed. +""" +from __future__ import annotations + +import argparse +import json +import os + +import numpy as np +import torch + +from fairchem.core.models.uma.escn_md import eSCNMDBackbone + +TINY = dict(sphere_channels=32, lmax=2, mmax=2, num_layers=2, hidden_channels=32, + edge_channels=16, num_distance_basis=32) +FULL = dict(sphere_channels=128, lmax=2, mmax=2, num_layers=2, hidden_channels=128, + edge_channels=128, num_distance_basis=512) +COMMON = dict(max_num_elements=100, cutoff=5.0, max_neighbors=300, otf_graph=False, + direct_forces=False, regress_forces=True, regress_stress=False, + norm_type="rms_norm_sh", act_type="gate", ff_type="grid", + use_dataset_embedding=True, dataset_list=["omat"], distance_function="gaussian") + + +def npy(t): + return t.detach().to(torch.float32).cpu().numpy() + + +def export_uma_s_1(args): + """Export the released uma-s-1 checkpoint: host MoLE-merge to a plain backbone for the given + composition, then write a clean WeightBundle (weights + fixed buffers + energy normalizer).""" + os.environ.setdefault("HF_HUB_OFFLINE", "1") + from ase.build import molecule + from huggingface_hub import hf_hub_download + from fairchem.core import FAIRChemCalculator + from fairchem.core.units.mlip_unit import load_predict_unit + from fairchem.core.units.mlip_unit.api.inference import InferenceSettings + + ckpt = args.checkpoint or hf_hub_download("facebook/UMA", "checkpoints/uma-s-1.pt") + settings = InferenceSettings(tf32=False, activation_checkpointing=True, merge_mole=True, + compile=False, external_graph_gen=False, internal_graph_gen_version=2) + pu = load_predict_unit(ckpt, inference_settings=settings, device="cpu") + calc = FAIRChemCalculator(pu, task_name=args.task) + if args.xyz: + from ase.io import read as _read + atoms = _read(args.xyz) + else: + atoms = molecule(args.molecule) + atoms.info.update(charge=args.charge, spin=args.spin) + atoms.calc = calc + E_ref = float(atoms.get_potential_energy()) # triggers the host MoLE merge + F_ref = atoms.get_forces().astype(np.float32) + + bb = pu.model.module.backbone # plain eSCNMDBackbone after merge + assert type(bb).__name__ == "eSCNMDBackbone", "merge did not produce a plain backbone" + energy_block = pu.model.module.output_heads["energyandforcehead"].head.energy_block + etask = pu.model.module.tasks[f"{args.task}_energy"] + + cfg = dict(sphere_channels=bb.sphere_channels, lmax=bb.lmax, mmax=bb.mmax, + num_layers=len(bb.blocks), hidden_channels=bb.hidden_channels, + num_distance_basis=int(bb.distance_expansion.offset.numel()), + cutoff=float(bb.cutoff), ff_type="spectral", act_type="gate", + norm_type="rms_norm_sh", chg_spin_emb_type=bb.chg_spin_emb_type, task=args.task, + # charge_balanced_channels (uma-s-1.2): l=0 scalar channels re-balanced to the + # system charge after each block. cs==ce (default) => disabled (uma-s-1). + charge_channel_start=int(getattr(bb, "charge_channel_start", 0)), + charge_channel_end=int(getattr(bb, "charge_channel_end", 0))) + + saved = {"config": np.frombuffer(json.dumps(cfg).encode(), dtype=np.uint8)} + for k, v in bb.state_dict().items(): + saved[f"w@{k}"] = npy(v) + for k, v in energy_block.state_dict().items(): + saved[f"w@energy_block.{k}"] = npy(v) + sg = bb.SO3_grid["lmax_lmax"] + saved["host@to_m"] = npy(bb.mappingReduced.to_m) + saved["host@to_grid_mat"] = npy(sg.to_grid_mat) + saved["host@from_grid_mat"] = npy(sg.from_grid_mat) + saved["host@gauss_offset"] = npy(bb.distance_expansion.offset) + saved["host@gauss_coeff"] = np.array([bb.distance_expansion.coeff], dtype=np.float32) + saved["scale@rmsd"] = np.array([float(etask.normalizer.rmsd)], dtype=np.float64) + saved["scale@mean"] = np.array([float(etask.normalizer.mean)], dtype=np.float64) + saved["scale@elem_refs"] = etask.element_references.element_references.detach().cpu().numpy().astype(np.float64) + + # embed the fairchem reference E/F for this composition so `tt-atom verify` can close the + # roundtrip on device (the two numpy worlds cannot share a process, so we carry the numbers). + saved["ref@energy"] = np.array([E_ref], dtype=np.float64) + saved["ref@forces"] = F_ref + saved["ref@pos"] = npy(torch.as_tensor(atoms.get_positions())) + saved["ref@atomic_numbers"] = np.asarray(atoms.get_atomic_numbers(), dtype=np.int64) + saved["ref@charge"] = np.array([float(args.charge)], dtype=np.float64) + saved["ref@spin"] = np.array([float(args.spin)], dtype=np.float64) + saved["ref@cell"] = npy(torch.as_tensor(atoms.get_cell().array)) + saved["ref@pbc"] = np.asarray(atoms.get_pbc(), dtype=bool) + + np.savez(args.out, **saved) + print(f"wrote {args.out} ({sum(1 for k in saved if k.startswith('w@'))} weight tensors, " + f"uma-s-1 merged for {args.xyz or args.molecule} charge={args.charge} spin={args.spin} task={args.task})") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--checkpoint", default=None, help="fairchem state_dict .pt (optional)") + ap.add_argument("--uma-s-1", action="store_true", help="export the released uma-s-1 (MoLE-merged)") + ap.add_argument("--xyz", default=None, help="structure file (overrides --molecule; for compositions not in ASE g2)") + ap.add_argument("--molecule", default="CH3CH2OH", help="ASE molecule name for uma-s-1 routing") + ap.add_argument("--task", default="omol") + ap.add_argument("--charge", type=int, default=0) + ap.add_argument("--spin", type=int, default=1) + ap.add_argument("--tiny", action="store_true") + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--out", required=True) + args = ap.parse_args() + + if args.uma_s_1: + export_uma_s_1(args) + return + + torch.manual_seed(args.seed) + cfg = dict(COMMON) + cfg.update(TINY if args.tiny else FULL) + bb = eSCNMDBackbone(**cfg).eval() + if args.checkpoint: + sd = torch.load(args.checkpoint, map_location="cpu") + sd = sd.get("state_dict", sd) + sd = {k.replace("backbone.", "", 1): v for k, v in sd.items()} + missing, unexpected = bb.load_state_dict(sd, strict=False) + print(f"loaded checkpoint: {len(missing)} missing, {len(unexpected)} unexpected keys") + + sc, hc = cfg["sphere_channels"], cfg["hidden_channels"] + energy = torch.nn.Sequential(torch.nn.Linear(sc, hc), torch.nn.SiLU(), + torch.nn.Linear(hc, hc), torch.nn.SiLU(), torch.nn.Linear(hc, 1)) + + saved = {"config": np.frombuffer(json.dumps(cfg).encode(), dtype=np.uint8)} + for k, v in bb.state_dict().items(): + saved[f"w@{k}"] = npy(v) + for k, v in energy.state_dict().items(): + saved[f"w@energy_block.{k}"] = npy(v) + sg = bb.SO3_grid["lmax_lmax"] + saved["host@to_m"] = npy(bb.mappingReduced.to_m) + saved["host@to_grid_mat"] = npy(sg.to_grid_mat) + saved["host@from_grid_mat"] = npy(sg.from_grid_mat) + saved["host@gauss_offset"] = npy(bb.distance_expansion.offset) + saved["host@gauss_coeff"] = np.array([bb.distance_expansion.coeff], dtype=np.float32) + + np.savez(args.out, **saved) + print(f"wrote {args.out} ({sum(1 for k in saved if k.startswith('w@'))} weight tensors)") + + +if __name__ == "__main__": + main() diff --git a/tt_atom/__init__.py b/tt_atom/__init__.py new file mode 100644 index 0000000..90eeff9 --- /dev/null +++ b/tt_atom/__init__.py @@ -0,0 +1,44 @@ +"""TT-Atom — high-performance Tenstorrent inference for eSEN / eSCN-MD (UMA-family) +equivariant ML interatomic potentials. + +Public API is populated as modules land (model, calculator, weights, ...). Submodules +import ttnn lazily so that ``import tt_atom`` is cheap and never opens a device. +""" + +from importlib.metadata import PackageNotFoundError, version as _version + +try: + __version__ = _version("tt-atom") +except PackageNotFoundError: # running from a source tree, not an installed dist + __version__ = "0+unknown" + +__all__ = ["UMA", "TTAtomCalculator", "WeightBundle", "Backbone", "HostGeometry", "MultiCard"] + + +def __getattr__(name): + # lazy so that ``import tt_atom`` stays cheap and never imports ttnn/torch eagerly + if name == "UMA": + from .calculator import UMA + + return UMA + if name == "TTAtomCalculator": + from .calculator import TTAtomCalculator + + return TTAtomCalculator + if name == "WeightBundle": + from .weights import WeightBundle + + return WeightBundle + if name == "Backbone": + from .model import Backbone + + return Backbone + if name == "HostGeometry": + from .geometry import HostGeometry + + return HostGeometry + if name == "MultiCard": + from .batch import MultiCard + + return MultiCard + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tt_atom/activation.py b/tt_atom/activation.py new file mode 100644 index 0000000..fc23aca --- /dev/null +++ b/tt_atom/activation.py @@ -0,0 +1,70 @@ +"""Gated nonlinearity (``GateActivation``) used between the two SO(2) convolutions. + +The l=0 (scalar) coefficient gets a plain SiLU; every higher-degree (vector) coefficient is +multiplied by a sigmoid gate broadcast from a per-degree scalar. We operate in the m-primed +coefficient order produced by the host Wigner map, so the gate-expansion is a fixed gather +that we realise as a slice + concat (only ``lmax`` distinct gate rows exist). + +Reference: ``fairchem ... nn/activation.py:GateActivation`` (``m_prime=True``). +""" +from __future__ import annotations + +import os + +# Route the gate fwd/bw column-split glue (slice+silu+slice+multiply+concat) through the custom +# ttnn.experimental.fused_gate kernel (one launch, no reduction). Needs source ttnn. +_FUSED_GATE = os.environ.get("TT_ATOM_FUSED_GATE") == "1" + + +def _expand_index_m_prime(lmax, mmax): + """The per-vector-coefficient gate index in m-primed order (see reference).""" + idx = [] + idx += list(range(lmax)) # m=0 block: l1m0, l2m0, ... + for mval in range(1, mmax + 1): + r = list(range(mval - 1, lmax)) + idx += r + r # real + imag halves + return idx + + +class GateActivation: + def __init__(self, device, *, lmax, mmax, num_channels): + import ttnn + import torch + + self.ttnn = ttnn + self.device = device + self.lmax = lmax + self.H = num_channels + self.expand_index = _expand_index_m_prime(lmax, mmax) + # gate expand as ONE matmul: gate_exp[E,(nsph-1)*H] = sigmoid(gating)[E,lmax*H] @ Expand. + # Expand [lmax*H, (nsph-1)*H] is a 0/1 selector (column-block c = I_H at row-block + # expand_index[c]) -> replaces the lmax*(nsph-1) slice+concat gather (fwd) and its + # segment-sum transpose (bw) with a single (transpose-)matmul. Bit-identical (0/1, fp32 acc). + from .device import compute_kernel_config + self.kcfg = compute_kernel_config() + H = num_channels + ncol = len(self.expand_index) + Ex = torch.zeros(lmax * H, ncol * H) + for c, row in enumerate(self.expand_index): + Ex[row * H:(row + 1) * H, c * H:(c + 1) * H] = torch.eye(H) + self.expand_w = ttnn.from_torch(Ex.contiguous(), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + + def __call__(self, gating_scalars, x): + """gating_scalars: ttnn ``[E, lmax*H]``; x: flat ttnn ``[E, nsph*H]`` (m-primed). + Returns flat ``[E, nsph*H]``.""" + ttnn = self.ttnn + E, H = x.shape[0], self.H + edt = x.dtype if x.dtype == ttnn.bfloat8_b else None # keep the bf8 edge flow bf8 + self._cache_gating, self._cache_x = gating_scalars, x # for the analytic-force VJP + g = ttnn.sigmoid(gating_scalars) # [E, lmax*H], H-block per degree + # expand the gate rows per vector coefficient as ONE matmul (0/1 selector); see __init__ + gate = ttnn.matmul(g, self.expand_w, dtype=edt, compute_kernel_config=self.kcfg) # [E,(nsph-1)*H] + self._cache_gate = gate # expanded gate for the VJP (fewer bw ops) + if _FUSED_GATE and x.shape[1] % 32 == 0 and gate.shape[1] % 32 == 0: + # one kernel: out = [silu(x[:, :H]) | x[:, H:] * gate] + op = ttnn._ttnn.operations.experimental.fused_gate + return op(x, gate, x, x.shape[1] // 32, gate.shape[1] // 32, H // 32, 0) + scalar = ttnn.silu(ttnn.slice(x, [0, 0], [E, H])) # l=0 coeff + vector = ttnn.multiply(ttnn.slice(x, [0, H], [E, x.shape[1]]), gate) + return ttnn.concat([scalar, vector], dim=1) diff --git a/tt_atom/assets/wigner_d_coefficients.pt b/tt_atom/assets/wigner_d_coefficients.pt new file mode 100644 index 0000000..c4d9976 Binary files /dev/null and b/tt_atom/assets/wigner_d_coefficients.pt differ diff --git a/tt_atom/batch.py b/tt_atom/batch.py new file mode 100644 index 0000000..a881bd4 --- /dev/null +++ b/tt_atom/batch.py @@ -0,0 +1,104 @@ +"""Multi-card throughput: fan independent systems across all cards, one process per card. + +The eSEN/eSCN-MD evaluation of one system is independent of every other, so throughput scales +by running one worker process per Tenstorrent card (each pinned with ``TT_VISIBLE_DEVICES`` so it +owns exactly one device) with the model + weights resident on that card. The parent streams +systems to a shared queue and the workers pull, evaluate, and return energies — embarrassingly +parallel, so aggregate throughput is the sum across cards. + +``ttnn`` is imported only *inside* the worker (after the device is pinned); the parent never +touches a device, which is what keeps the fan-out deadlock-free. +""" +from __future__ import annotations + +import multiprocessing as mp + + +def _worker(device_id, weights_path, fast, in_q, out_q): + import os + + os.environ["TT_VISIBLE_DEVICES"] = str(device_id) # pin one card -> it is device 0 + # one host thread per worker: the host geometry (torch) otherwise grabs every core, so N + # workers oversubscribe the CPU and throttle each other (4-card went *slower* than 1). + os.environ.setdefault("OMP_NUM_THREADS", "1") + import torch + + torch.set_num_threads(1) + from .device import open_device + from .model import Backbone, GraphContext + from .geometry import HostGeometry, csd_embedding, radius_graph + from .weights import WeightBundle + import ttnn + + b = WeightBundle.load(weights_path) + cfg, w = b.config, b.weights + C = cfg["sphere_channels"] + dev = open_device(0) + bb = Backbone(w, dev, cfg, b.to_grid_mat, b.from_grid_mat, fast=fast) + geo = HostGeometry(w, cfg, b.to_m, b.gauss_offset, b.gauss_coeff, gamma=0.0) + out_q.put(("ready", device_id)) + + while True: + job = in_q.get() + if job is None: + break + idx, pos_np, Z_np = job + pos = torch.tensor(pos_np, dtype=torch.float32) + Z = torch.tensor(Z_np) + ei, _ = radius_graph(pos, cfg["cutoff"]) + N, E = Z.shape[0], ei.shape[1] + se = csd_embedding(w, torch.tensor([0.0]), torch.tensor([0.0]), C)[torch.zeros(N, dtype=torch.long)] + t = geo(pos, Z, ei, se) + # build the per-system mean operator when the bundle balances charge channels (uma-s-1.2); + # else node_embedding would deref a None node_meanM. MultiCard evaluates one neutral system + # per worker, so the (1/N) mean + balance_add=0 default is exact. + graph = GraphContext(dev, edge_index=ei, wigner=t["wigner"].detach(), + wigner_inv=t["wigner_inv"].detach(), x_edge=t["x_edge"].detach(), + edge_envelope=t["edge_envelope"].detach(), num_nodes=N, fast=fast, + build_mean_op=(bb.ce > bb.cs)) + se3 = ttnn.from_torch(se.reshape(N, 1, C), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) + xi = ttnn.from_torch(t["x_init"].detach(), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) + _, energy = bb(xi, graph, se3) + out_q.put((idx, float(ttnn.to_torch(energy).reshape(-1)[0]), E)) + + ttnn.close_device(dev) + + +class MultiCard: + """A persistent pool of one worker per device. Use as a context manager.""" + + def __init__(self, weights_path, device_ids=(0, 1, 2, 3), *, fast=False): + self.ctx = mp.get_context("spawn") + self.in_q = self.ctx.Queue() + self.out_q = self.ctx.Queue() + self.procs = [self.ctx.Process(target=_worker, args=(d, weights_path, fast, self.in_q, self.out_q), + daemon=True) for d in device_ids] + for p in self.procs: + p.start() + for _ in self.procs: # wait until every card is ready + self.out_q.get() + + def energies(self, systems): + """``systems``: list of (positions[N,3], atomic_numbers[N]) numpy arrays. + Returns (energies list in input order, total edges processed).""" + for i, (pos, Z) in enumerate(systems): + self.in_q.put((i, pos, Z)) + out = [None] * len(systems) + total_edges = 0 + for _ in systems: + idx, en, E = self.out_q.get() + out[idx] = en + total_edges += E + return out, total_edges + + def close(self): + for _ in self.procs: + self.in_q.put(None) + for p in self.procs: + p.join(timeout=10) + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.close() diff --git a/tt_atom/bundle_cache.py b/tt_atom/bundle_cache.py new file mode 100644 index 0000000..eb8c7b7 --- /dev/null +++ b/tt_atom/bundle_cache.py @@ -0,0 +1,156 @@ +"""Composition-cached uma bundle factory — the machinery behind ``TTAtomCalculator.from_uma``. + +UMA's MoLE routing is baked at merge time, so one merged bundle is valid for exactly one +*(reduced composition, charge, spin, task)*. This module turns those into a stable cache key and, +on a miss, transparently builds the bundle by invoking the reference (fairchem, numpy>=2) +environment as a subprocess — the ttnn-side user never has to touch the two-env split for the +common path. A cache *hit* needs no fairchem at all: it is a plain ``np.load``. + +Two inherent frictions this hides (see README): the per-composition merge, and the fact that ttnn +(numpy<2) and fairchem (numpy>=2) cannot live in one interpreter. Both are automated + cached here +so a scientist sees them at most once per composition. +""" +from __future__ import annotations + +import hashlib +import math +import os +import pathlib +import subprocess +import sys +import tempfile +from collections import Counter +from functools import reduce + +CACHE_DIR = pathlib.Path( + os.environ.get("TT_ATOM_CACHE", pathlib.Path.home() / ".cache" / "tt_atom" / "bundles") +) + + +def infer_task(atoms): + """Zero-config task default: a fully periodic cell -> ``'omat'`` (bulk materials), otherwise + ``'omol'`` (molecules). Slabs / MOFs / molecular crystals (oc20/odac/omc) should pass the task + explicitly — this only picks the right *common* default so the bare entry point Just Works.""" + import numpy as np + + return "omat" if np.asarray(atoms.get_pbc()).all() else "omol" + + +def reduced_composition(numbers): + """Return sorted ``((Z, reduced_count), ...)``. + + MoLE routes on the *fractional* composition, so H2O and H4O2 share a bundle — we divide the + counts by their GCD to make the key scale-invariant.""" + counts = Counter(int(z) for z in numbers) + g = reduce(math.gcd, counts.values()) + return tuple(sorted((z, c // g) for z, c in counts.items())) + + +def formula(numbers): + """Human-readable reduced formula (e.g. ``C2H6O``) for logging.""" + from ase.data import chemical_symbols + + return "".join(f"{chemical_symbols[z]}{c if c > 1 else ''}" + for z, c in reduced_composition(numbers)) + + +def composition_hash(numbers): + comp = reduced_composition(numbers) + s = ";".join(f"{z}:{c}" for z, c in comp) + return hashlib.sha1(s.encode()).hexdigest()[:12] + + +def bundle_path(model, task, numbers, charge, spin, cache_dir=None): + """Deterministic cache path for a merged bundle. Same composition/charge/spin/task/model -> + same file, regardless of atom ordering or an integer scaling of the counts.""" + h = composition_hash(numbers) + name = f"{model}_{task}_{h}_c{int(charge)}_s{int(spin)}.npz" + return pathlib.Path(cache_dir or CACHE_DIR) / name + + +def resolve_refenv(refenv=None): + """Locate the reference (fairchem) python: explicit arg > ``$TT_ATOM_REFENV`` > default + ``~/.ttatom_run/refenv/bin/python``. Raises a clear, actionable error if none is found — + only ever called on a cache *miss* (a hit needs no refenv).""" + candidates = [ + refenv, + os.environ.get("TT_ATOM_REFENV"), + str(pathlib.Path.home() / ".ttatom_run" / "refenv" / "bin" / "python"), + ] + for c in candidates: + if c and pathlib.Path(c).exists(): + return c + raise RuntimeError( + "No reference (fairchem) environment found to build the uma bundle.\n" + "The one-time MoLE merge needs fairchem (numpy>=2), which cannot share the ttnn env\n" + "(numpy<2). Create it once with this one command:\n\n" + " python -m venv ~/.ttatom_run/refenv && " + "~/.ttatom_run/refenv/bin/pip install 'fairchem-core>=2.10'\n\n" + "or point TT-Atom at an existing fairchem env via TT_ATOM_REFENV=/path/to/bin/python\n" + "(or the refenv= argument). A cached bundle needs no refenv — this is only the\n" + "first-use build per composition." + ) + + +def build_bundle(atoms, out_path, *, model="uma-s-1", task="omol", charge=0, spin=1, + refenv=None, checkpoint=None): + """Merge + export a bundle for ``atoms`` by running ``tools/export_weights.py`` in the + reference env. Writes atomically (build to a sidecar, then ``os.replace``) so an interrupted + build can never leave a half-written file that later looks like a cache hit.""" + if model != "uma-s-1": + raise ValueError( + f"auto-build supports model='uma-s-1'; got {model!r}. Export other checkpoints " + "manually with tools/export_weights.py and load the .npz directly." + ) + py = resolve_refenv(refenv) + tools = pathlib.Path(__file__).resolve().parent.parent / "tools" / "export_weights.py" + out_path = pathlib.Path(out_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + # sidecar ends in .npz so np.savez does not append a second extension + tmp_out = out_path.with_name(out_path.name + ".building.npz") + with tempfile.TemporaryDirectory() as td: + from ase.io import write + + xyz = pathlib.Path(td) / "structure.xyz" + write(str(xyz), atoms) + cmd = [py, str(tools), "--uma-s-1", "--xyz", str(xyz), "--task", task, + "--charge", str(int(charge)), "--spin", str(int(spin)), "--out", str(tmp_out)] + if checkpoint: + cmd += ["--checkpoint", str(checkpoint)] + env = dict(os.environ) + env.setdefault("HF_HUB_OFFLINE", "1") + try: + subprocess.run(cmd, check=True, env=env) + except subprocess.CalledProcessError as e: + tmp_out.unlink(missing_ok=True) + raise RuntimeError( + f"reference-env bundle build failed (exit {e.returncode}). Command:\n " + + " ".join(cmd) + + "\n\nIf this is a checkpoint/access error: the UMA weights are gated — accept the " + "license at\n https://huggingface.co/facebook/UMA\nand log in once with " + "`huggingface-cli login` (or set HF_TOKEN) in the reference env." + ) from e + os.replace(tmp_out, out_path) + return out_path + + +def get_or_build(atoms, *, model="uma-s-1", task="omol", charge=0, spin=1, refenv=None, + checkpoint=None, cache_dir=None, log=True): + """Return the cache path for this system's bundle, building it on a miss. Pure I/O + subprocess + — no ttnn, no device. ``TTAtomCalculator.from_uma`` wraps this and returns a calculator.""" + numbers = atoms.get_atomic_numbers() + path = bundle_path(model, task, numbers, charge, spin, cache_dir=cache_dir) + if path.exists(): + return path + if log: + print( + f"[tt-atom] building {model} bundle for composition {formula(numbers)} " + f"(task={task}, charge={int(charge)}, spin={int(spin)}) — one-time per composition, " + f"~30s via the reference env...", + file=sys.stderr, flush=True, + ) + build_bundle(atoms, path, model=model, task=task, charge=charge, spin=spin, + refenv=refenv, checkpoint=checkpoint) + if log: + print(f"[tt-atom] cached bundle -> {path}", file=sys.stderr, flush=True) + return path diff --git a/tt_atom/calculator.py b/tt_atom/calculator.py new file mode 100644 index 0000000..c5f6928 --- /dev/null +++ b/tt_atom/calculator.py @@ -0,0 +1,343 @@ +"""``TTAtomCalculator`` — an ASE calculator backed by the device-resident eSCN-MD engine. + +Wraps the host geometry + device backbone + analytic-force VJP behind ASE's interface so the +model is usable for real geometry relaxations and MD. Energy and conservative forces come from +``tt_atom.forces.energy_and_forces`` (forces are ``-dE/dpos`` via the on-device reverse pass, +not finite differences).""" +from __future__ import annotations + +import numpy as np +import torch +from ase.calculators.calculator import Calculator, all_changes + +from . import device as D +from . import forces as Fmod +from .geometry import HostGeometry, csd_embedding, radius_graph +from .model import Backbone +from .weights import WeightBundle + + +def _eval_rotation(): + """A fixed, generic 3x3 rotation used to evaluate molecules OFF the coordinate axes. + + The device backbone runs bf16. At an exactly axis-aligned / high-symmetry geometry the edge + frame's roll-gauge derivative is large; its cancellation in the analytic force is exact in fp32 + (fairchem) but leaks in bf16, giving a wrong force at the *exact* symmetric point (a ~0.03 A + rattle already removes it — the quaternion frame fixed most cases, this closes the residual for + molecules whose symmetry axis lands on a coordinate plane, e.g. planar molecules in the yz/xy + plane). The model is rotationally equivariant, so evaluating in a generic orientation and + rotating the force back is exact and deterministic, and moves every edge off the frame's special + directions. Applied to molecules only (periodic forces already match; rotating a cell is left + out to keep PBC/stress untouched).""" + ax = np.array([0.3, 0.5, 0.81]); ax = ax / np.linalg.norm(ax); th = 0.7 + K = np.array([[0.0, -ax[2], ax[1]], [ax[2], 0.0, -ax[0]], [-ax[1], ax[0], 0.0]]) + return np.eye(3) + np.sin(th) * K + (1.0 - np.cos(th)) * (K @ K) + + +def UMA(atoms, task=None, model="uma-s-1", charge=0, spin=1, refenv=None, checkpoint=None, + cache_dir=None, device=None, device_id=0, fast=False, trace=False, **kwargs): + """Zero-config entry point — the face of the library. + + from tt_atom import UMA + atoms.calc = UMA(atoms) # energy + forces on the card, nothing else to know + + Picks sensible defaults (``uma-s-1``; ``task`` inferred from periodicity — ``omat`` for a fully + periodic cell, else ``omol``; auto-build + composition-cache the bundle; auto-locate the + reference env; device 0) and returns a ready :class:`TTAtomCalculator`. Every default is a plain + keyword you can override (``task=``, ``charge=``/``spin=``, ``trace=``, ``fast=``, ``device_id=``, + ``refenv=``, ...). It is exactly :meth:`TTAtomCalculator.from_uma` with task inference on top — + reach for ``from_uma`` / ``TTAtomCalculator(bundle)`` directly only when you want to pin the + task or manage the bundle file yourself.""" + from . import bundle_cache as BC + + if task is None: + task = BC.infer_task(atoms) + return TTAtomCalculator.from_uma(model=model, task_name=task, atoms=atoms, charge=charge, + spin=spin, refenv=refenv, checkpoint=checkpoint, + cache_dir=cache_dir, device=device, device_id=device_id, + fast=fast, trace=trace, **kwargs) + + +class TTAtomCalculator(Calculator): + implemented_properties = ["energy", "energies", "free_energy", "forces", "stress"] + + def __init__(self, bundle, task_name=None, device=None, device_id=0, gamma=0.0, + fast=False, trace=False, trace_region_size=400_000_000, **kwargs): + """``bundle`` is a TT-Atom weight bundle (path or ``WeightBundle``) exported for a fixed + (composition, charge, spin, task): UMA's MoLE routing consumes the dataset token, so the + task is baked in at merge time and cannot be switched at runtime. ``task_name`` mirrors + ``fairchem.core.FAIRChemCalculator(task_name=...)``; when given it must match the bundle's + task (a mismatch raises, rather than silently using the wrong normalizer). + + ``trace=True`` captures the device forward+backward once and replays it each step for a + fixed topology (MD / relaxation): ~2x fewer host dispatches, bit-for-bit the same forces. + The neighbour list is rechecked every step and the trace is re-captured automatically if + an atom crosses the cutoff, so results are always correct. When passing your own + ``device`` with ``trace=True``, open it with a non-zero ``trace_region_size``.""" + super().__init__(**kwargs) + if isinstance(bundle, str): + bundle = WeightBundle.load(bundle) + self.bundle = bundle + self.cfg = bundle.config + self.C = self.cfg["sphere_channels"] + self.fast = fast + self.trace = trace + # evaluate molecules off the coordinate axes (bf16 exact-symmetry force fix; see _eval_rotation) + self._eval_rot = torch.tensor(_eval_rotation(), dtype=torch.float32) + if task_name is not None and task_name != bundle.task: + raise ValueError( + f"task_name={task_name!r} does not match this bundle's task {bundle.task!r}. " + f"UMA's MoLE routing bakes the task into the merged bundle; export a bundle for " + f"{task_name!r} (tools/export_weights.py --task {task_name}) to use that task.") + self._owns_device = device is None + self.device = device if device is not None else D.open_device( + device_id, trace_region_size=trace_region_size if trace else 0) + self._engine = None + self._engine_edges = None + self._engine_shift = None + self._batch_engine = None + self._batch_edges = None + w = bundle.weights + self.backbone = Backbone(w, self.device, self.cfg, bundle.to_grid_mat, + bundle.from_grid_mat, fast=fast) + self.geo = HostGeometry(w, self.cfg, bundle.to_m, bundle.gauss_offset, + bundle.gauss_coeff, gamma=gamma, + coefficient_index=bundle.coefficient_index) + self._w = w + # energy normalizer (real checkpoints: E = rmsd*E_raw + mean + sum_i refs[Z_i], + # F = rmsd*F_raw); identity for the random-weight bundles (rmsd=1, mean=0, refs=None) + self.scale_rmsd = bundle.scale_rmsd + self.scale_mean = bundle.scale_mean + self.elem_refs = bundle.elem_refs + self.task = self.task_name = bundle.task + + @classmethod + def from_uma(cls, model="uma-s-1", task_name="omol", atoms=None, charge=0, spin=1, + refenv=None, checkpoint=None, cache_dir=None, device=None, device_id=0, + fast=False, trace=False, **kwargs): + """fairchem-parallel entry point: return a ready calculator for ``atoms``, auto-building + and caching the composition-specific merged bundle on first use. + + Mirrors ``FAIRChemCalculator`` in spirit — you hand it a structure + task and get a + calculator back — but hides the two frictions inherent to running UMA on ttnn: + + * MoLE routing bakes one merged bundle per *(reduced composition, charge, spin, task)*, + so we hash those into a cache key under ``~/.cache/tt_atom/bundles`` (override with + ``$TT_ATOM_CACHE`` or ``cache_dir``). + * ttnn (numpy<2) and fairchem (numpy>=2) cannot share a process, so the *build* runs the + reference env as a subprocess. Resolution order: ``refenv`` arg > ``$TT_ATOM_REFENV`` > + ``~/.ttatom_run/refenv/bin/python``. A **cache hit needs no fairchem/refenv at all** — + it is a plain load, which is the common path. + + ``atoms`` is required (it determines the composition). When it carries ``info['charge']`` / + ``info['spin']`` those win over the args, so the bundle is merged with the exact charge/spin + the runtime will read back. First use per composition logs an honest one-time build notice. + """ + from . import bundle_cache as BC + + if atoms is None: + raise ValueError( + "from_uma needs `atoms` to determine the composition — MoLE bakes one bundle per " + "reduced composition/charge/spin/task, so there is no way to pick (or build) a " + "bundle without the structure. Pass the Atoms you want to run." + ) + # an explicit charge/spin on the atoms wins: it is what `calculate` reads back at runtime, + # so the merge must use the same value for a consistent result. + charge = atoms.info.get("charge", charge) + spin = atoms.info.get("spin", spin) + # ...and, symmetrically, stamp the resolved values back onto the atoms so `calculate` + # reads the *same* charge/spin the bundle was merged for. Without this, the flagship + # `UMA(atoms)` path (default charge=0, spin=1) merges a spin=1 bundle but `calculate` + # falls back to spin=0 — a silent mismatch between the baked MoLE routing and the runtime + # system embedding. `setdefault` respects an explicit value (which already won above). + atoms.info.setdefault("charge", charge) + atoms.info.setdefault("spin", spin) + path = BC.get_or_build(atoms, model=model, task=task_name, charge=charge, spin=spin, + refenv=refenv, checkpoint=checkpoint, cache_dir=cache_dir) + return cls(str(path), task_name=task_name, device=device, device_id=device_id, + fast=fast, trace=trace, **kwargs) + + def close(self): + if self._engine is not None: + self._engine.close() + self._engine = None + if self._batch_engine is not None: + self._batch_engine.close() + self._batch_engine = None + if self._owns_device and self.device is not None: + import ttnn + + ttnn.close_device(self.device) + self.device = None + + def calculate(self, atoms=None, properties=("energy", "forces"), system_changes=all_changes): + super().calculate(atoms, properties, system_changes) + pos = torch.tensor(np.asarray(atoms.get_positions()), dtype=torch.float32) + Z = torch.tensor(np.asarray(atoms.get_atomic_numbers()), dtype=torch.long) + chg = float(atoms.info.get("charge", 0.0)) + charge = torch.tensor([chg]) + spin = torch.tensor([float(atoms.info.get("spin", 0.0))]) + + pbc = np.asarray(atoms.get_pbc()) + cell = torch.tensor(np.asarray(atoms.get_cell()), dtype=torch.float32) if pbc.any() else None + # evaluate molecules in a generic orientation so no edge lands on the frame's bf16-sensitive + # axis (exact-symmetry force fix); equivariant, so the force is rotated back below. + rotate = cell is None + if rotate: + pos = pos @ self._eval_rot.T + edge_index, edge_cell_shift = radius_graph(pos, self.cfg["cutoff"], cell=cell, pbc=pbc) + if edge_index.shape[1] == 0: + raise ValueError("no edges within cutoff — system too sparse for this model") + sys_emb = csd_embedding(self._w, charge, spin, self.C, + dataset=self.task)[torch.zeros(Z.shape[0], dtype=torch.long)] + + # stress is autograd of energy wrt a symmetric strain — only meaningful for a periodic + # cell (variable-cell relaxation / NPT); request it when ASE asks or a cell is present. + want_stress = cell is not None and ("stress" in properties or pbc.all()) + virial = None + # The trace engine captures only the energy+force op stream — it has no stress readout — + # so when stress is *explicitly* requested (e.g. an ASE variable-cell filter) fall back to + # the eager stress path rather than silently dropping stress. A fully-periodic system that + # only wants energy/forces still enjoys the trace (stress is auto-computed but unrequested). + if self.trace and "stress" not in properties: + E, F = self._traced(pos, Z, edge_index, edge_cell_shift, sys_emb, charge=chg) + elif want_stress: + E, F, virial = Fmod.energy_and_forces(self.backbone, self.geo, pos, Z, edge_index, + sys_emb, edge_cell_shift=edge_cell_shift, + compute_stress=True, charge=chg) + else: + E, F = Fmod.energy_and_forces(self.backbone, self.geo, pos, Z, edge_index, sys_emb, + edge_cell_shift=edge_cell_shift, charge=chg) + # apply the per-task energy normalizer + element references (forces/virial scale by rmsd) + E = self.scale_rmsd * E + self.scale_mean + if self.elem_refs is not None: + E += float(self.elem_refs[Z].sum()) + F = self.scale_rmsd * F + if rotate: + F = F @ self._eval_rot # rotate the force back into the input frame + self.results["energy"] = E + self.results["free_energy"] = E + self.results["energies"] = np.full(len(atoms), E / len(atoms), dtype=np.float64) + self.results["forces"] = F.detach().numpy().astype(np.float64) + if virial is not None: + from ase.stress import full_3x3_to_voigt_6_stress + + # stress = (rmsd * dE_raw/dstrain) / V; fairchem's convention (uma/outputs.py) + stress = self.scale_rmsd * virial.detach().numpy().astype(np.float64) / atoms.get_volume() + self.results["stress"] = full_3x3_to_voigt_6_stress(stress) + + def evaluate_batch(self, systems, properties=("energy", "forces"), trace=False): + """Disjoint-union batched evaluation — K systems in ONE device forward (fairchem/PyG style). + + ``systems`` is a list of ASE ``Atoms`` (or ``(positions, atomic_numbers)`` / dicts). The + systems are concatenated into one block-diagonal graph, evaluated in a single device call, + and the per-system energies recovered by segment-sum; forces (when requested) come from the + one shared analytic backward (block-diagonal => each atom's own-system force). This is the + throughput path for the dispatch-bound regime of *many small systems*. + + Returns ``dict(energy=np.ndarray[K], forces=list[np.ndarray[N_k, 3]] | None)`` with the + per-system energy normalizer applied, mirroring the single-system ``calculate`` results. + + All systems must share this bundle's reduced composition: a merged uma-s-1 bundle bakes the + MoLE expert routing for one composition (fairchem's merged batched inference requires the + same), so the batch is e.g. conformers / an MD ensemble of one molecule. + + ``trace=True`` captures the batched device forward+backward once and replays it while the + batch topology (edge set) is unchanged — the throughput path for a *batched MD ensemble / + relaxation* of K fixed-composition replicas, where the sub-saturation batch forward is + host-dispatch-bound (measured ~2.5-3x for modest K). It re-captures whenever the neighbour + list changes, so results stay correct; leave it False for one-shot screening (a fresh batch + each call would re-capture every time, wasting the capture cost).""" + from . import disjoint + + # A merged uma-s-1 bundle bakes the MoLE expert routing for ONE (reduced composition, + # charge, spin), so a batch that mixes compositions/charge/spin is silently wrong. When the + # bundle carries a reference (its merge inputs), validate every system against it up front + # with a clear error rather than returning a plausible-but-wrong energy. (Reference-less + # random-weight bundles — the mechanism tests — skip this and stay composition-agnostic.) + ref = self.bundle.reference + if ref is not None: + from .bundle_cache import reduced_composition + + want_comp = reduced_composition(ref["atomic_numbers"]) + want_cs = (float(ref["charge"]), float(ref["spin"])) + for k, system in enumerate(systems): + _, Z_k, chg_k, spin_k, _, _ = disjoint._as_atoms_fields(system) + if reduced_composition(Z_k.tolist()) != want_comp: + raise ValueError( + f"batched system {k} has a different reduced composition than this bundle. " + "A merged uma-s-1 bundle bakes the MoLE routing for one reduced composition, " + "so every system in a batch must share it (e.g. conformers / an MD ensemble " + "of one molecule); evaluate other compositions with their own bundle." + ) + if (chg_k, spin_k) != want_cs: + raise ValueError( + f"batched system {k} has (charge, spin)=({chg_k}, {spin_k}) but this bundle " + f"was merged for {want_cs}. The MoLE routing bakes one charge/spin; every " + "system in a batch must share it." + ) + + bg = disjoint.assemble(systems, self.cfg["cutoff"], self._w, self.C, task=self.task) + # NB: unlike the single-system calculate(), the batched path does NOT apply the generic + # eval-rotation, so a batch member sitting at an EXACT high-symmetry geometry keeps the small + # residual bf16 symmetry-force error the quaternion frame doesn't fully cancel (throughput + # path — a 0.03 A rattle or the single-system path removes it). + want_forces = "forces" in properties + if trace and want_forces: + E_raw, F = self._traced_batch(bg) + else: + E_raw, F = Fmod.energy_and_forces_batch(self.backbone, self.geo, bg, + compute_forces=want_forces) + energies, forces_out, off = [], [], 0 + for k, n in enumerate(bg.natoms): + Ek = self.scale_rmsd * float(E_raw[k]) + self.scale_mean + if self.elem_refs is not None: + Ek += float(self.elem_refs[bg.Z[off:off + n]].sum()) + energies.append(Ek) + if want_forces: + Fk = (self.scale_rmsd * F[off:off + n]).detach().numpy().astype(np.float64) + forces_out.append(Fk) + off += n + return dict(energy=np.array(energies), forces=forces_out if want_forces else None) + + def _traced_batch(self, bg): + """Trace-replayed batched energy+forces; (re)captures on neighbour-list change. Returns + ``(E_raw: torch[K], F: torch[Ntot, 3])`` matching ``energy_and_forces_batch``.""" + from .trace import TracedEngine + + changed = (self._batch_engine is None + or self._batch_edges is None + or self._batch_edges.shape != bg.edge_index.shape + or not torch.equal(self._batch_edges, bg.edge_index)) + if changed: + if self._batch_engine is not None: + self._batch_engine.close() + self._batch_engine = TracedEngine( + self.backbone, self.geo, bg.Z, bg.edge_index, bg.sys_emb, + edge_cell_shift=bg.cell_shift, seg=bg.segment_matrix(), linear_scatter=True, + charge=bg.charge, system_natoms=bg.natoms) + self._batch_edges = bg.edge_index.clone() + return self._batch_engine(bg.pos) + + def _traced(self, pos, Z, edge_index, edge_cell_shift, sys_emb, charge=0.0): + """Trace-replayed energy+forces; (re)captures when the neighbour list changes. + + The captured trace bakes in ``edge_cell_shift`` (the per-edge periodic image offset), so a + changing cell at *fixed* topology — e.g. an NPT / variable-cell step that doesn't cross the + cutoff — must also trigger a re-capture, else the replay would silently use the stale + shift and return wrong forces. Hence the change test covers the cell shift, not just the + edge index.""" + from .trace import TracedEngine + + changed = (self._engine_edges is None or self._engine_edges.shape != edge_index.shape + or not torch.equal(self._engine_edges, edge_index) + or self._engine_shift is None or self._engine_shift.shape != edge_cell_shift.shape + or not torch.equal(self._engine_shift, edge_cell_shift)) + if changed: + if self._engine is not None: + self._engine.close() + self._engine = TracedEngine(self.backbone, self.geo, Z, edge_index, sys_emb, + edge_cell_shift=edge_cell_shift, charge=charge) + self._engine_edges = edge_index.clone() + self._engine_shift = edge_cell_shift.clone() + return self._engine(pos) diff --git a/tt_atom/cli.py b/tt_atom/cli.py new file mode 100644 index 0000000..a24b92b --- /dev/null +++ b/tt_atom/cli.py @@ -0,0 +1,318 @@ +"""``tt-atom`` console entry — the user-facing commands for the ttnn runtime environment. + + tt-atom run STRUCTURE --uma-s-1 [--task] [--charge --spin] [--relax|--md] [--trace] [--out] + tt-atom info BUNDLE # config / task / weight coverage + tt-atom verify BUNDLE # device parity vs the embedded fairchem reference + tt-atom relax BUNDLE [--input geom.xyz | --molecule NAME] [--trace] [--fmax --steps] + tt-atom md BUNDLE [--input geom.xyz | --molecule NAME] [--trace] [--steps --dt --temp] + tt-atom convert-checkpoint CKPT.pt --out BUNDLE.npz --molecule NAME [--task --charge --spin] + +``run`` is the fairchem-parallel one-shot: a structure file in, energy/relax/MD out, with the +composition-specific uma-s-1 bundle auto-built (first use per composition) and cached. All commands +run here (numpy<2 + ttnn); the one-time bundle build ``run`` triggers on a cache miss shells out to +the reference (fairchem, numpy>=2) environment. ``convert-checkpoint`` is the explicit/advanced +form of that build and detects a missing fairchem, printing the exact reference-env invocation. +""" +from __future__ import annotations + +import argparse +import pathlib +import sys + +import numpy as np + + +def _resolve_charge_spin(args, bundle): + """CLI charge/spin for a bundle-based command: an explicit --charge/--spin wins; otherwise + default to the (charge, spin) the bundle was merged for (its embedded reference), NOT a fixed + literal. A merged bundle bakes one charge/spin into the MoLE routing, so evaluating it with a + mismatched runtime value silently disagrees — an omol bundle is merged at spin=1, so the old + ``--spin`` default of 0 was wrong. Falls back to 0/0 for a reference-less bundle.""" + ref = bundle.reference + charge = args.charge + spin = args.spin + if charge is None: + charge = float(ref["charge"]) if ref is not None else 0.0 + if spin is None: + spin = float(ref["spin"]) if ref is not None else 0.0 + return charge, spin + + +def _atoms(args, bundle=None): + from ase.build import molecule + if args.input: + from ase.io import read + atoms = read(args.input) + else: + atoms = molecule(args.molecule) + if bundle is not None: + charge, spin = _resolve_charge_spin(args, bundle) + else: + charge = 0.0 if args.charge is None else args.charge + spin = 0.0 if args.spin is None else args.spin + atoms.info.setdefault("charge", charge) + atoms.info.setdefault("spin", spin) + return atoms + + +def _calc(args, bundle=None): + from .calculator import TTAtomCalculator + return TTAtomCalculator(bundle if bundle is not None else args.bundle, + device_id=args.device_id, fast=args.fast, + trace=getattr(args, "trace", False)) + + +def cmd_info(args): + from .weights import WeightBundle + b = WeightBundle.load(args.bundle) + ok, missing, present = b.verify_coverage() + print(f"bundle : {args.bundle}") + print(f"task : {b.task}") + print(f"config : {b.config}") + print(f"weights: {present} tensors, coverage {'OK' if ok else 'MISSING ' + str(missing[:5])}") + print(f"scale : rmsd={b.scale_rmsd} mean={b.scale_mean} elem_refs={'yes' if b.elem_refs is not None else 'no'}") + ref = b.reference + print(f"ref : {'embedded (E=%.5f eV, %d atoms)' % (ref['energy'], len(ref['atomic_numbers'])) if ref else 'none'}") + return 0 + + +def cmd_verify(args): + """Device parity vs the fairchem reference embedded in the bundle at convert time.""" + from ase import Atoms + from .weights import WeightBundle + from .calculator import TTAtomCalculator + + b = WeightBundle.load(args.bundle) + ref = b.reference + if ref is None: + print("bundle has no embedded reference (re-export with tools/export_weights.py --uma-s-1)") + return 2 + pbc = ref["pbc"] if ref["pbc"] is not None else False + atoms = Atoms(numbers=ref["atomic_numbers"], positions=ref["pos"], + cell=ref["cell"] if ref["cell"] is not None else None, pbc=pbc) + atoms.info.update(charge=ref["charge"], spin=ref["spin"]) + calc = TTAtomCalculator(b, device_id=args.device_id, fast=args.fast) + atoms.calc = calc + try: + E = atoms.get_potential_energy() + F = atoms.get_forces() + finally: + calc.close() + Eref, Fref = ref["energy"], ref["forces"] + rel = abs(E - Eref) / max(abs(Eref), 1e-9) + pcc = float(np.corrcoef(F.ravel(), Fref.ravel())[0, 1]) + mae = float(np.abs(F - Fref).mean()) + ok = rel < args.etol and pcc > args.fpcc + print(f"task={b.task} device E={E:.5f} ref E={Eref:.5f} rel={rel:.2e}") + print(f"force PCC={pcc:.5f} MAE={mae:.3e} eV/A |F|max dev/ref={np.abs(F).max():.4f}/{np.abs(Fref).max():.4f}") + print("PASS" if ok else f"FAIL (need rel<{args.etol}, PCC>{args.fpcc})") + return 0 if ok else 1 + + +def cmd_relax(args): + from ase.optimize import FIRE + + from .weights import WeightBundle + bundle = WeightBundle.load(args.bundle) if isinstance(args.bundle, str) else args.bundle + atoms = _atoms(args, bundle) + calc = _calc(args, bundle) + atoms.calc = calc + try: + e0 = atoms.get_potential_energy() + FIRE(atoms, logfile="-").run(fmax=args.fmax, steps=args.steps) + e1 = atoms.get_potential_energy() + fmax = float((atoms.get_forces() ** 2).sum(1).max() ** 0.5) + print(f"relax: E {e0:.6f} -> {e1:.6f} eV; fmax={fmax:.4f} (target {args.fmax}); " + f"converged={fmax <= args.fmax}") + if args.out: + from ase.io import write + write(args.out, atoms) + print(f"wrote {args.out}") + finally: + calc.close() + return 0 + + +def cmd_md(args): + from ase import units + from ase.md.langevin import Langevin + from ase.md.velocitydistribution import MaxwellBoltzmannDistribution + + from .weights import WeightBundle + bundle = WeightBundle.load(args.bundle) if isinstance(args.bundle, str) else args.bundle + atoms = _atoms(args, bundle) + calc = _calc(args, bundle) + atoms.calc = calc + try: + MaxwellBoltzmannDistribution(atoms, temperature_K=args.temp) + dyn = Langevin(atoms, timestep=args.dt * units.fs, temperature_K=args.temp, + friction=0.01 / units.fs) + e0 = atoms.get_potential_energy() + + def _log(): + ekin = atoms.get_kinetic_energy() + print(f" step {dyn.nsteps:4d} E={atoms.get_potential_energy():.5f} " + f"T={ekin / (1.5 * units.kB * len(atoms)):.1f} K") + + dyn.attach(_log, interval=max(1, args.steps // 10)) + dyn.run(args.steps) + print(f"md: {args.steps} steps ({args.dt} fs) at {args.temp} K; " + f"E {e0:.5f} -> {atoms.get_potential_energy():.5f} eV") + if args.out: + from ase.io import write + write(args.out, atoms) + print(f"wrote {args.out}") + finally: + calc.close() + return 0 + + +def cmd_run(args): + """One-shot: STRUCTURE -> auto-built/cached bundle -> single-point / relax / MD -> result. + + The fairchem-parallel entry point. Reads any ASE-readable structure, transparently builds and + caches the composition-specific uma-s-1 bundle on first use (via the reference env), then runs + on device. A cached composition needs no fairchem.""" + from ase.io import read + from . import bundle_cache as BC + from .calculator import TTAtomCalculator + + atoms = read(args.structure) + atoms.info.setdefault("charge", args.charge) + atoms.info.setdefault("spin", args.spin) + task = args.task or BC.infer_task(atoms) # zero-config: omat for a bulk cell, else omol + calc = TTAtomCalculator.from_uma(model="uma-s-1", task_name=task, atoms=atoms, + charge=args.charge, spin=args.spin, refenv=args.refenv, + device_id=args.device_id, fast=args.fast, trace=args.trace) + atoms.calc = calc + try: + e0 = atoms.get_potential_energy() + print(f"energy: {e0:.6f} eV ({len(atoms)} atoms, task={task}, " + f"charge={int(args.charge)}, spin={int(args.spin)})") + if args.relax: + from ase.optimize import FIRE + + FIRE(atoms, logfile="-").run(fmax=args.fmax, steps=args.steps) + e1 = atoms.get_potential_energy() + fmax = float((atoms.get_forces() ** 2).sum(1).max() ** 0.5) + print(f"relax: E {e0:.6f} -> {e1:.6f} eV; fmax={fmax:.4f} (target {args.fmax}); " + f"converged={fmax <= args.fmax}") + elif args.md: + from ase import units + from ase.md.langevin import Langevin + from ase.md.velocitydistribution import MaxwellBoltzmannDistribution + + MaxwellBoltzmannDistribution(atoms, temperature_K=args.temp) + dyn = Langevin(atoms, timestep=args.dt * units.fs, temperature_K=args.temp, + friction=0.01 / units.fs) + + def _log(): + ekin = atoms.get_kinetic_energy() + print(f" step {dyn.nsteps:4d} E={atoms.get_potential_energy():.5f} " + f"T={ekin / (1.5 * units.kB * len(atoms)):.1f} K") + + dyn.attach(_log, interval=max(1, args.steps // 10)) + dyn.run(args.steps) + print(f"md: {args.steps} steps ({args.dt} fs) at {args.temp} K; " + f"E {e0:.5f} -> {atoms.get_potential_energy():.5f} eV") + if args.out: + from ase.io import write + + write(args.out, atoms) + print(f"wrote {args.out}") + finally: + calc.close() + return 0 + + +def cmd_convert(args): + """Fairchem UMA checkpoint -> TT-Atom bundle. Needs the reference (fairchem) environment.""" + try: + import fairchem # noqa: F401 + except Exception: + tools = pathlib.Path(__file__).resolve().parent.parent / "tools" / "export_weights.py" + print("convert-checkpoint needs fairchem (numpy>=2), which cannot share this ttnn env.") + print("Run it in the reference environment, e.g.:\n") + print(f" HF_HUB_OFFLINE=1 /bin/python {tools} --uma-s-1 \\") + print(f" --checkpoint {args.checkpoint} --molecule {args.molecule} " + f"--task {args.task} --charge {args.charge} --spin {args.spin} --out {args.out}") + return 2 + import importlib.util + tools = pathlib.Path(__file__).resolve().parent.parent / "tools" / "export_weights.py" + spec = importlib.util.spec_from_file_location("_ttatom_export", tools) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + ns = argparse.Namespace(uma_s_1=True, checkpoint=args.checkpoint, molecule=args.molecule, + task=args.task, charge=args.charge, spin=args.spin, out=args.out) + mod.export_uma_s_1(ns) + if args.verify: + from .weights import WeightBundle + ok, missing, present = WeightBundle.load(args.out).verify_coverage() + print(f"roundtrip: reloaded bundle, coverage {'OK' if ok else 'MISSING ' + str(missing[:5])} " + f"({present} tensors). Run `tt-atom verify {args.out}` on device for numeric parity.") + return 0 + + +def main(argv=None): + ap = argparse.ArgumentParser(prog="tt-atom", description="TT-Atom: UMA MLIP inference on Tenstorrent") + ap.add_argument("--device-id", type=int, default=0) + ap.add_argument("--fast", action="store_true", help="bf8 weights (throughput; accuracy-safe)") + sub = ap.add_subparsers(dest="cmd", required=True) + + p = sub.add_parser("run", help="one-shot: structure -> auto-bundle -> single-point/relax/md") + p.add_argument("structure", help="ASE-readable structure (.xyz/.cif/.pdb/...)") + p.add_argument("--uma-s-1", action="store_true", help="use uma-s-1 (the default auto-build model)") + p.add_argument("--task", default=None, + help="dataset/task token (omol/omat/oc20/odac/omc); inferred from periodicity if unset") + p.add_argument("--charge", type=float, default=0.0) + p.add_argument("--spin", type=float, default=1.0) + p.add_argument("--refenv", default=None, help="fairchem python for the one-time bundle build") + p.add_argument("--trace", action="store_true", help="trace-captured device loop (~2x)") + g = p.add_mutually_exclusive_group() + g.add_argument("--relax", action="store_true", help="FIRE geometry relaxation") + g.add_argument("--md", action="store_true", help="Langevin molecular dynamics") + p.add_argument("--fmax", type=float, default=0.05) + p.add_argument("--steps", type=int, default=200) + p.add_argument("--dt", type=float, default=1.0) + p.add_argument("--temp", type=float, default=300.0) + p.add_argument("--out", help="write final geometry/trajectory here") + p.set_defaults(func=cmd_run) + + p = sub.add_parser("info", help="show bundle config/task/coverage"); p.add_argument("bundle") + p.set_defaults(func=cmd_info) + + p = sub.add_parser("verify", help="device parity vs the bundle's embedded reference") + p.add_argument("bundle"); p.add_argument("--etol", type=float, default=1e-3) + p.add_argument("--fpcc", type=float, default=0.99); p.set_defaults(func=cmd_verify) + + def _sys_args(p): + p.add_argument("bundle") + p.add_argument("--input", help="ASE-readable geometry file (.xyz/.cif/...)") + p.add_argument("--molecule", default="CH3CH2OH", help="ASE builtin molecule if no --input") + p.add_argument("--charge", type=float, default=None, + help="net charge (default: the value the bundle was merged for)") + p.add_argument("--spin", type=float, default=None, + help="spin multiplicity (default: the value the bundle was merged for)") + p.add_argument("--trace", action="store_true", help="trace-captured device loop (~2x)") + p.add_argument("--out", help="write final geometry here") + + p = sub.add_parser("relax", help="FIRE geometry relaxation"); _sys_args(p) + p.add_argument("--fmax", type=float, default=0.05); p.add_argument("--steps", type=int, default=200) + p.set_defaults(func=cmd_relax) + + p = sub.add_parser("md", help="Langevin molecular dynamics"); _sys_args(p) + p.add_argument("--steps", type=int, default=100); p.add_argument("--dt", type=float, default=1.0) + p.add_argument("--temp", type=float, default=300.0); p.set_defaults(func=cmd_md) + + p = sub.add_parser("convert-checkpoint", help="fairchem UMA .pt -> TT-Atom bundle") + p.add_argument("checkpoint"); p.add_argument("--out", required=True) + p.add_argument("--molecule", default="CH3CH2OH"); p.add_argument("--task", default="omol") + p.add_argument("--charge", type=int, default=0); p.add_argument("--spin", type=int, default=1) + p.add_argument("--verify", action="store_true"); p.set_defaults(func=cmd_convert) + + args = ap.parse_args(argv) + return args.func(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tt_atom/device.py b/tt_atom/device.py new file mode 100644 index 0000000..5f3d0dc --- /dev/null +++ b/tt_atom/device.py @@ -0,0 +1,154 @@ +"""Device + kernel-configuration helpers for TT-Atom. + +Everything that touches a Tenstorrent card goes through here so the numerics policy lives +in one place. The policy (validated on Blackhole p150): + + * matmul accumulation in fp32 (``fp32_dest_acc_en=True``) at ``HiFi4`` fidelity, with + ``packer_l1_acc=True`` -- this is what gives matmul PCC ~1.0 vs torch. + * weights default to ``bfloat16``; ``--fast`` mode may store weights as ``bfloat8_b`` but + keeps the fp32/HiFi4 accumulation above (only the operands get cheaper, not the math). + +``import ttnn`` is done lazily inside functions so that ``import tt_atom`` never opens or +probes a device. +""" +from __future__ import annotations + +import os +from contextlib import contextmanager + + +def device_ede() -> bool: + """Whether the edge-degree embedding (node init) is computed on device inside the trace. + + Default OFF (host torch, the original path). When ``TT_ATOM_DEVICE_EDE=1`` the radial MLP -> + rotate-back -> envelope -> scatter -> +l0 chain runs on device (see tt_atom/edge_degree.py), + removing the largest per-step host cost (the radial-MLP fwd+bw over E edges). Read at call + time so tests / benches can toggle it per run.""" + return os.environ.get("TT_ATOM_DEVICE_EDE") == "1" + + +def bf8_edge() -> bool: + """Whether the edgewise message dataflow runs in bfloat8_b (the E-sized [E,nsph*C] activations + that dominate the bandwidth-bound device replay). The device replay is DRAM-bandwidth bound on + these activations (bf8 halves the traffic -> ~2x on the fat matmuls; measured), NOT compute- + bound (HiFi4==LoFi) nor weight-bound (bf8 weights alone = 1.00x). The bf16<->bf8 boundary sits + at the SMALL N-sized node features (gather input / scatter output), so there is no per-edge + typecast overhead. Node residual stream + norms stay bf16. Requires the source-ttnn build whose + fused_rotate/gate/gc kernels accept bf8 I/O. Default OFF.""" + return os.environ.get("TT_ATOM_BF8_EDGE") == "1" + + +def edge_dtype(ttnn): + """bfloat8_b when bf8_edge() else bfloat16 — the working dtype of the edgewise E-sized flow.""" + return ttnn.bfloat8_b if bf8_edge() else ttnn.bfloat16 + + +# Budgets (bytes) for a single L1-resident intermediate. The L1-residency perf wins (grid, so2, +# norm chains) keep intermediates on-chip, but L1 is small (~1.4 MB/bank x 130) and op circular +# buffers grow with problem size, so at large N/E they consume nearly all L1 and an L1-resident +# activation OOMs (or fails trace capture with "writes not supported"). Residency is gated +# per-tensor: L1 only when the estimated tile-padded byte size fits the relevant budget, else +# DRAM-interleaved (the pre-optimization path -- correct at any size, just no L1 speedup). +# +# Two budgets because the crossovers differ (empirically verified on Blackhole p150): +# * SO2 edge tensors [E, nsph*Cin]: L1 fits at E<=2234 (N=128), fails at E=4834 (N=250). +# 12 MB -> edge cap ~2600 (see SO2Convolution.l1_max_edges). +# * GRID/NORM node tensors: pass the TRUE tile-padded feature width (the 3D [N,nsph,C]/[N,npts,C] +# tensors pad the coeff dim up to a tile: nsph 9->32, npts 42->64 -- a 1.5-3.5x blowup a naive +# nsph*C estimate misses). Verified real-byte crossovers: grid L1 fits at N=432 (~7MB), fails +# at N=686 (~11MB); norm fits at N=686 (~5.6MB), fails at N=1024 (~8.4MB). 8 MB covers both. +L1_RESIDENCY_BUDGET = 12_000_000 # SO(2) edge tensors [E, nsph*Cin] (flat, no coeff padding) +L1_NODE_BUDGET = 8_000_000 # grid / norm node tensors (pass tile-padded feature width) + + +def coeff_reshape(ttnn, t, shape): + """Reshape that collapses/expands the spherical-harmonic coefficient dim (nsph, e.g. 9), + which is NOT tile-aligned. A direct ``ttnn.reshape`` on a TILE tensor physically repacks the + tile padding (9 -> 32, a 3.5x data reorg -> ~18 ms at E~46k); routing through ROW_MAJOR + (contiguous, no coeff padding) is ~4x faster and bit-exact (a lossless layout round-trip, no + dtype change). For an already-ROW_MAJOR tensor a plain reshape is cheap, so pass through. + + Only use for reshapes that move the coefficient dim across the flat/3D boundary; a reshape + that only touches the batch (outer) dim never repacks and should stay a direct reshape.""" + try: + is_tile = t.layout == ttnn.TILE_LAYOUT + except Exception: + is_tile = True + if is_tile: + r = ttnn.to_layout(t, ttnn.ROW_MAJOR_LAYOUT) + r = ttnn.reshape(r, shape) + return ttnn.to_layout(r, ttnn.TILE_LAYOUT) + return ttnn.reshape(t, shape) + + +def l1_if_fits(ttnn, rows, width, *, dtype_bytes=2, budget=L1_RESIDENCY_BUDGET): + """Return an L1 memory config if a tile-padded ``[rows, width]`` tensor of ``dtype_bytes`` + fits the per-tensor L1 residency budget, else DRAM-interleaved. Guards the residency wins so + large systems degrade gracefully to DRAM instead of OOMing the trace.""" + rp = ((rows + 31) // 32) * 32 + wp = ((width + 31) // 32) * 32 + return ttnn.L1_MEMORY_CONFIG if rp * wp * dtype_bytes <= budget else ttnn.DRAM_MEMORY_CONFIG + + +def compute_kernel_config(fast: bool = False): + """The TT-Atom matmul/compute numerics policy. + + ``fast`` does not change the accumulation math here (operand dtype is chosen at weight + load time); HiFi4 + fp32 dest accumulation is kept in both modes for accuracy. + """ + import ttnn + + return ttnn.WormholeComputeKernelConfig( + math_fidelity=ttnn.MathFidelity.HiFi4, + math_approx_mode=False, + fp32_dest_acc_en=True, + packer_l1_acc=True, + ) + + +def open_device(device_id: int = 0, *, l1_small_size: int = 0, trace_region_size: int = 0): + """Open a single Tenstorrent device with the program cache enabled. + + The program cache is what makes warm calls cheap (kernels are compiled once); device + residency + trace capture build on top of it. + """ + import ttnn + + dev = ttnn.open_device( + device_id=device_id, + l1_small_size=l1_small_size, + trace_region_size=trace_region_size, + ) + dev.enable_program_cache() + return dev + + +def open_mesh(device_ids, *, l1_small_size: int = 0, trace_region_size: int = 0): + """Open a row mesh over ``device_ids`` for multi-card throughput. + + Requires ``TT_MESH_GRAPH_DESC_PATH`` to point at the matching fabric descriptor + (e.g. ``p150_x4_mesh_graph_descriptor.textproto`` for a 4-card QuietBox). + """ + import ttnn + + ids = list(device_ids) + mesh = ttnn.open_mesh_device( + ttnn.MeshShape(1, len(ids)), + l1_small_size=l1_small_size, + trace_region_size=trace_region_size, + device_ids=ids, + ) + mesh.enable_program_cache() + return mesh + + +@contextmanager +def device(device_id: int = 0, **kwargs): + """Context manager that opens a device and guarantees it is closed.""" + import ttnn + + dev = open_device(device_id, **kwargs) + try: + yield dev + finally: + ttnn.close_device(dev) diff --git a/tt_atom/disjoint.py b/tt_atom/disjoint.py new file mode 100644 index 0000000..26a8759 --- /dev/null +++ b/tt_atom/disjoint.py @@ -0,0 +1,140 @@ +"""Disjoint-union (block-diagonal) graph batching — the fairchem/PyG way. + +Concatenate K independent systems into one big block-diagonal graph so the device backbone +evaluates all K in a *single* forward. This is a throughput win in the dispatch-bound regime +(many small systems), where per-call host overhead — not device compute — dominates. + +Every eSCN-MD backbone op is per-node or per-edge: the SO(2) convs and rotations act edgewise, +the norms and grid/spectral feed-forwards act nodewise, and the scatter-add is the one-hot +matmul ``S[N, E] @ messages`` (see edgewise.py). Give each system's edges a per-system node +offset and that scatter matrix is automatically block-diagonal — a message on an edge lands only +on its own system's nodes. So the *entire* forward is batch-transparent as-is; the only change is +the energy readout, which becomes a segment-sum per system (``Backbone.energy_batch``). Forces +need no change either: the batch energy is the sum of per-system energies, and block-diagonality +makes ``dE_total/dx_n = dE_(system of n)/dx_n``, so the existing summed-energy backward yields +each atom's own-system force (the batched forces are just the concatenation). + +The single change that assembly must respect: build each system's neighbour list *separately* +(never on the concatenated positions), or a global radius graph would wire atoms across systems. + +Composition constraint: a merged uma-s-1 WeightBundle bakes the MoLE expert routing for one +reduced composition (fairchem's ``merge_MOLE_model`` asserts the same), so a batch that wants +*correct* energies shares that composition — e.g. conformers / an MD ensemble of one molecule. +The assembly itself is composition-agnostic (block-diagonal is block-diagonal); only the routing +baked into the weights is not. +""" +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch + +from .geometry import csd_embedding, radius_graph + + +def _as_atoms_fields(system): + """Accept an ASE ``Atoms`` or a ``(positions, atomic_numbers)`` / dict and return the fields + disjoint-union assembly needs: positions, atomic numbers, charge, spin, cell, pbc.""" + if hasattr(system, "get_positions"): # ASE Atoms + pos = torch.tensor(np.asarray(system.get_positions()), dtype=torch.float32) + Z = torch.tensor(np.asarray(system.get_atomic_numbers()), dtype=torch.long) + charge = float(system.info.get("charge", 0.0)) + spin = float(system.info.get("spin", 0.0)) + pbc = np.asarray(system.get_pbc()) + cell = (torch.tensor(np.asarray(system.get_cell()), dtype=torch.float32) + if pbc.any() else None) + return pos, Z, charge, spin, cell, pbc + if isinstance(system, dict): + pos = torch.as_tensor(system["pos"], dtype=torch.float32) + Z = torch.as_tensor(system["Z"], dtype=torch.long) + return pos, Z, float(system.get("charge", 0.0)), float(system.get("spin", 0.0)), \ + system.get("cell"), system.get("pbc") + pos, Z = system # (positions, atomic_numbers) + return (torch.as_tensor(pos, dtype=torch.float32), torch.as_tensor(Z, dtype=torch.long), + 0.0, 0.0, None, None) + + +@dataclass +class BatchedGraph: + """A disjoint-union of K systems ready for one device forward. + + ``edge_index`` carries per-system node offsets (block-diagonal scatter). ``batch`` maps each + atom to its system id; ``natoms`` are the per-system atom counts (to split concatenated + outputs back). ``sys_emb`` [Ntot, C] is the per-node system (charge/spin/dataset) embedding. + """ + pos: torch.Tensor # [Ntot, 3] + Z: torch.Tensor # [Ntot] + edge_index: torch.Tensor # [2, Etot], node offsets applied + cell_shift: torch.Tensor # [Etot, 3] + batch: torch.Tensor # [Ntot] atom -> system id + natoms: list # per-system atom counts (len K) + sys_emb: torch.Tensor # [Ntot, C] + charge: float = 0.0 # shared system charge (a batch is one composition => one charge) + + @property + def K(self): + return len(self.natoms) + + def segment_matrix(self): + """One-hot segment matrix ``seg[K, Ntot]`` (``seg[k, n] = 1`` iff atom n in system k) for + the segment-sum energy readout.""" + seg = torch.zeros(self.K, self.pos.shape[0]) + seg[self.batch, torch.arange(self.pos.shape[0])] = 1.0 + return seg + + +def assemble(systems, cutoff, weights, sphere_channels, task="omol"): + """Build the block-diagonal graph for ``systems`` (list of ASE ``Atoms`` / dicts / tuples). + + Each system's neighbour list is built independently (never on concatenated positions), then + edges are offset by the running node count so the union stays block-diagonal. Returns a + ``BatchedGraph``. Raises if any system has no edges within ``cutoff`` (too sparse for the model). + """ + if len(systems) == 0: + raise ValueError("empty batch") + pos_all, Z_all, ei_all, shift_all, batch_all, sys_all = [], [], [], [], [], [] + natoms = [] + charges = [] + node_off = 0 + for k, system in enumerate(systems): + pos, Z, charge, spin, cell, pbc = _as_atoms_fields(system) + charges.append(charge) + n = Z.shape[0] + ei, shift = radius_graph(pos, cutoff, cell=cell, pbc=pbc) + if ei.shape[1] == 0: + raise ValueError(f"system {k} has no edges within cutoff — too sparse for this model") + pos_all.append(pos) + Z_all.append(Z) + ei_all.append(ei + node_off) # per-system node offset + shift_all.append(shift) + batch_all.append(torch.full((n,), k, dtype=torch.long)) + se = csd_embedding(weights, torch.tensor([charge]), torch.tensor([spin]), + sphere_channels, dataset=task) # [1, C] + sys_all.append(se.expand(n, -1)) + natoms.append(n) + node_off += n + # charge_balanced_channels needs a per-system charge target; a merged bundle is one composition, + # so the batch shares one charge. + if len(set(charges)) != 1: + raise ValueError(f"evaluate_batch needs one shared charge; got {sorted(set(charges))}") + charge = charges[0] + # A charged batch's per-system target is charge/natoms — a per-system quantity the batched balance + # currently expresses as one scalar (charge/natoms[0]). That is only exact when all systems share + # an atom count, so require it (the reduced-composition guard permits e.g. CH-reducible systems of + # different sizes). Neutral batches (target 0) are unaffected. + if charge != 0.0 and len(set(natoms)) > 1: + raise ValueError( + "charged uma-s-1.2 batched forces need equal atom counts per system (the per-system " + f"charge/natoms target differs across sizes {sorted(set(natoms))}); batch equal-size " + "systems, or evaluate the charged systems one at a time.") + return BatchedGraph( + pos=torch.cat(pos_all, dim=0), + Z=torch.cat(Z_all, dim=0), + edge_index=torch.cat(ei_all, dim=1), + cell_shift=torch.cat(shift_all, dim=0), + batch=torch.cat(batch_all, dim=0), + natoms=natoms, + sys_emb=torch.cat(sys_all, dim=0), + charge=charge, + ) diff --git a/tt_atom/edge_degree.py b/tt_atom/edge_degree.py new file mode 100644 index 0000000..73e6163 --- /dev/null +++ b/tt_atom/edge_degree.py @@ -0,0 +1,93 @@ +"""On-device edge-degree embedding — the node initialisation, moved off host. + +The eSCN ``edge_degree_embedding`` builds the initial node features from the invariant edge +embedding: a per-edge radial MLP produces the m=0 block, it is rotated back into node spherical +harmonics with the (inverse) Wigner matrix, scaled by the radial envelope, scatter-added onto +target nodes, and the constant l=0 init (sphere embedding + system embedding) added on top. + +Historically this ran on host in ``geometry.HostGeometry`` (its autograd supplied ``dx_init/dpos`` +for the analytic force). At N=1000 the radial-MLP forward+backward over E~46k edges was the single +largest per-step host cost (~140 ms). Structurally it is identical to a message-passing layer, so +we run it on device inside the captured trace instead — a few small GEMMs + the fused rotation + +the linear scatter, all machinery the backbone already owns. The host then only differentiates the +cheap geometric terms (wigner_inv, x_edge, envelope), whose adjoints this module's backward +(:func:`edge_degree_bw`) accumulates into the same ``acc`` dict the block backward fills. + +Forward mirrors ``geometry.HostGeometry.__call__``'s node-init block; backward mirrors it exactly +in reverse (cf. ``forces.edgewise_bw``). Gated by ``device.device_ede()`` (``TT_ATOM_DEVICE_EDE=1``). +""" +from __future__ import annotations + +from .device import compute_kernel_config +from .so2 import RadialMLP + + +class EdgeDegreeEmbedding: + """radial MLP -> pad -> rotate-back -> envelope -> scatter/rescale -> + l0, on device.""" + + def __init__(self, weights, device, cfg, *, rescale): + import ttnn + + self.ttnn = ttnn + self.device = device + self.C = cfg["sphere_channels"] + self.lmax = cfg["lmax"] + self.m0 = cfg.get("mmax_m0_coeffs", self.lmax + 1) + self.rescale = float(rescale) + # fold the 1/rescale node-init scale into the radial MLP's final linear (fp32 -> bf16), so it + # lands in the matmul's fp32 accumulation instead of a lossy bf16 multiply by 0.2 downstream. + self.rad = RadialMLP(weights, "edge_degree_embedding.rad_func", device, ttnn.bfloat16, + out_scale=1.0 / self.rescale) + self.kcfg = compute_kernel_config() + + def __call__(self, graph, l0): + """``x_init`` [N, nsph, C] from ``graph.x_edge`` + rot_inv + envelope and the constant l0.""" + ttnn = self.ttnn + from . import rotation, scatter + + E, C = graph.E, self.C + nred, nsph, N = graph.nred, graph.nsph, graph.N + edm = self.rad(graph.x_edge) # [E, m0*C] + # place the m=0 block at the front of the reduced m-space (zeros in the tail coeffs): the + # flat layout is coeff-major, so padding columns m0*C -> nred*C is exactly F.pad(...,(0,0,0, + # nred-m0)) on [E, m0, C]. When nred==m0 (rare) this is a no-op. + if nred > self.m0: + edm = ttnn.pad(edm, [(0, 0), (0, (nred - self.m0) * C)], value=0.0) + self._cache_edm = edm # rotate input (for the VJP) + m_back = rotation.rotate(ttnn, edm, graph.rot_inv_ij, graph.rot_inv_coef, nred, C, + self.device, n_out=nsph) # [E, nsph*C] + m_env = ttnn.multiply(m_back, graph.edge_envelope_f) # [E, nsph*C] * [E,1] broadcast + self._cache_mback = m_back + if graph.linear_scatter: + node = scatter.segment_sum(ttnn, m_env, graph.tgt_gather, graph.Dmax_t, N, nsph * C) + else: + node = ttnn.matmul(graph.scatter, m_env, compute_kernel_config=self.kcfg) # [N, nsph*C] + node = ttnn.reshape(node, (N, nsph, C)) # 1/rescale already folded into the radial MLP + return ttnn.add(node, l0) + + +def edge_degree_bw(ede, graph, g_x_init, acc): + """VJP of :class:`EdgeDegreeEmbedding`. ``g_x_init`` [N,nsph,C] is the node adjoint after all + blocks; l0 is constant (pass-through). Accumulates the geometric adjoints (g rot_inv, + g_envelope) into ``acc`` and appends the radial adjoint so ``backbone_bw`` finishes g_x_edge.""" + ttnn = ede.ttnn + from . import rotation, scatter + + C = ede.C + N, nsph = g_x_init.shape[0], g_x_init.shape[1] + E, nred = graph.E, graph.nred + # x_init = node + l0 (1/rescale folded into the radial MLP); node = scatter(m_env) -> gather + # the node adjoint back to edges. + gnf = ttnn.to_layout(ttnn.reshape(g_x_init, (N, nsph * C)), ttnn.ROW_MAJOR_LAYOUT) + g_menv = ttnn.to_layout(ttnn.embedding(graph.tgt_idx, gnf), ttnn.TILE_LAYOUT) # [E, nsph*C] + # m_env = m_back * envelope + g_mback = ttnn.multiply(g_menv, graph.edge_envelope_f) + g_env = ttnn.sum(ttnn.multiply(g_menv, ede._cache_mback), dim=1, keepdim=True) # [E,1] + # inverse rotation backward (reduced m-space nred -> node SH nsph) + g_edm, g_rinv = rotation.rotate_bw(ttnn, ede._cache_edm, g_mback, graph.rot_inv_ij, + graph.rot_inv_coef, nred, C, ede.device, n_out=nsph) + if nred > ede.m0: + g_edm = ttnn.slice(g_edm, [0, 0], [E, ede.m0 * C]) # unpad -> the m=0 block + acc["rot_inv"] = g_rinv if acc["rot_inv"] is None else ttnn.add(acc["rot_inv"], g_rinv) + acc["envelope"] = g_env if acc["envelope"] is None else ttnn.add(acc["envelope"], g_env) + acc["g_rad"].append((ede, g_edm)) # backbone_bw does ede.rad.bw(g_edm) -> g_x_edge diff --git a/tt_atom/edgewise.py b/tt_atom/edgewise.py new file mode 100644 index 0000000..dde5070 --- /dev/null +++ b/tt_atom/edgewise.py @@ -0,0 +1,87 @@ +"""Edgewise message passing — the SO(2) message block of eSCN-MD. + +For every edge: gather source/target node features, rotate into the edge frame with the +host Wigner matrix, run the two SO(2) convolutions with a gate in between, apply the radial +envelope, rotate back, and scatter-add the messages onto target nodes. Gathers are row +selects (``ttnn.embedding``); the scatter-add is a fixed one-hot matmul ``S @ messages`` +(``S[n, e] = 1`` iff edge ``e`` targets node ``n``) — its transpose is exactly the gather, +which is what the analytic-force backward needs. + +Reference: ``fairchem ... escn_md_block.py:Edgewise.forward_chunk``. +""" +from __future__ import annotations + +from .device import compute_kernel_config +from .so2 import SO2Convolution +from .activation import GateActivation + + +class Edgewise: + def __init__(self, weights, prefix, device, *, sphere_channels, hidden_channels, + lmax, mmax, fast=False): + import ttnn + + self.ttnn = ttnn + self.device = device + self.C = sphere_channels + self.kcfg = compute_kernel_config() + extra = lmax * hidden_channels # gate scalar channels + self.so2_1 = SO2Convolution( + weights, f"{prefix}.so2_conv_1", device, + sphere_channels_in=2 * sphere_channels, m_output_channels=hidden_channels, + lmax=lmax, mmax=mmax, extra_m0_output_channels=extra, fast=fast) + self.gate = GateActivation(device, lmax=lmax, mmax=mmax, num_channels=hidden_channels) + self.so2_2 = SO2Convolution( + weights, f"{prefix}.so2_conv_2", device, + sphere_channels_in=hidden_channels, m_output_channels=sphere_channels, + lmax=lmax, mmax=mmax, extra_m0_output_channels=0, fast=fast) + + def __call__(self, x, graph): + """x: ttnn ``[N, nsph, C]``; ``graph`` carries the on-device geometric terms. + + Runs flat ``[E, nsph*C]`` end to end: gather (row select) -> per-coordinate concat of + source|target -> rotate to edge frame (sparse MAC) -> SO(2) conv x2 + gate -> radial + envelope -> rotate back -> scatter-add to targets (one-hot matmul).""" + ttnn = self.ttnn + from . import rotation + from .device import bf8_edge + N, nsph, C = x.shape[0], x.shape[1], self.C + E = graph.E + dev = self.device + _b8 = bf8_edge() + xf = ttnn.to_layout(ttnn.reshape(x, (N, nsph * C)), ttnn.ROW_MAJOR_LAYOUT) # gather operand + # Keep the src/tgt gathers ROW_MAJOR: the interleave (concat dim=2) + flatten to [E, nsph*2C] + # is done entirely in ROW_MAJOR (contiguous, no coeff-dim tile padding) with a single + # to_layout TILE at the end -- avoids the ~18 ms TILE 3D->2D repack of the 9-coeff dim. + xs = ttnn.reshape(ttnn.embedding(graph.src_idx, xf), (E, nsph, C)) # RM [E, nsph, C] + xt = ttnn.reshape(ttnn.embedding(graph.tgt_idx, xf), (E, nsph, C)) # RM + m_cat = ttnn.to_layout(ttnn.reshape(ttnn.concat([xs, xt], dim=2), (E, nsph * 2 * C)), + ttnn.TILE_LAYOUT) # flat [xs_i|xt_i] per coord + # bf8-edge boundary: from here the whole E-sized edge flow (rotate->so2->gate->so2-> + # rotate_back) runs bf8. bf8 can't be ROW_MAJOR so the cast lands after the RM gather. The + # device replay is DRAM-bandwidth bound on these [E,nsph*C] activations; bf8 halves traffic. + if _b8: + m_cat = ttnn.typecast(m_cat, ttnn.bfloat8_b) + + # rotate node SH (nsph) into the reduced m-space (nred); the SO(2) pipeline runs there + m_rot = rotation.rotate(ttnn, m_cat, graph.rot_fwd_ij, graph.rot_fwd_coef, nsph, 2 * C, + dev, n_out=graph.nred) + m, gating = self.so2_1(m_rot, graph.x_edge) # flat in/out throughout (reduced m-space) + m = self.gate(gating, m) + m_so2 = self.so2_2(m, graph.x_edge) # flat [E, nred*C] + m_env = ttnn.multiply(m_so2, graph.edge_envelope_f) # [E,1] broadcast + # rotate the reduced m-space message back to node SH (nsph) + m_back = rotation.rotate(ttnn, m_env, graph.rot_inv_ij, graph.rot_inv_coef, graph.nred, C, + dev, n_out=nsph) + self._cache_mcat, self._cache_mso2, self._cache_menv = m_cat, m_so2, m_env + + # scatter-add messages onto target nodes: dense one-hot matmul (small N) or linear O(E) + # gather+reduce (large N). See GraphContext.linear_scatter / tt_atom/scatter.py. + odt = ttnn.bfloat16 if bf8_edge() else None + if graph.linear_scatter: + from . import scatter + out = scatter.segment_sum(ttnn, m_back, graph.tgt_gather, graph.Dmax_t, N, nsph * C) + else: + # scatter (bf16 one-hot) @ m_back (bf8) -> bf16 node features (back on the bf16 stream) + out = ttnn.matmul(graph.scatter, m_back, dtype=odt, compute_kernel_config=self.kcfg) + return ttnn.reshape(out, (N, nsph, C)) diff --git a/tt_atom/forces.py b/tt_atom/forces.py new file mode 100644 index 0000000..23efb59 --- /dev/null +++ b/tt_atom/forces.py @@ -0,0 +1,653 @@ +"""Analytic forces — reverse-mode VJP through the device backbone (``F = -dE/dpos``). + +This is the production force path (not finite differences). The heavy ``dE/dfeature`` terms are +exactly the transposes of the forward GEMMs and run on device; the cheap geometric Jacobian +``d(geometric terms)/dpos`` is finished on host with torch autograd (it is <1% of the compute). + +The device VJP produces adjoints at the pos-dependent device inputs: + * ``g_x_init`` [N, nsph, C] + * ``g_wigner`` [E, nsph, nsph] (from the edge-frame rotation) + * ``g_wigner_inv`` [E, nsph, nsph] + * ``g_envelope`` [E, 1, 1] + * ``g_rad`` per radial conv (adjoint at the radial-MLP *output*) -> host autograd finishes + ``g_x_edge`` through the radial MLP, whose LayerNorms we deliberately keep off-device. + +Each VJP mirrors a forward module in ``tt_atom/`` and is unit-tested against ``tests/mirror.py`` +(a bit-exact torch transcription of the device forward) in ``tests/test_forces.py``. +""" +from __future__ import annotations + +import torch + + +def _mm(ttnn, g, W, kcfg, memory_config=None): + """grad wrt x of ``y = x @ W`` (W stored [in,out]): ``g @ W^T``. ``transpose_b`` folds the + transpose into the matmul (bit-identical), dropping an explicit transpose op per call — the + backward makes ~40 of these on constant weights, all in the captured trace. + + ``memory_config`` lets the BW-bound grid_bw chain keep its transpose-matmul outputs + L1-resident instead of round-tripping DRAM (same residency win as the forward grid module).""" + # bf8-edge: when the incoming adjoint is bf8, keep the transpose-matmul output bf8 so the + # backward edge flow stays bf8 (halves the [E,W] gradient traffic; the dominant bw cost). + edt = ttnn.bfloat8_b if g.dtype == ttnn.bfloat8_b else None + if memory_config is not None: + return ttnn.matmul(g, W, transpose_b=True, dtype=edt, compute_kernel_config=kcfg, + memory_config=memory_config) + return ttnn.matmul(g, W, transpose_b=True, dtype=edt, compute_kernel_config=kcfg) + + +# --------------------------------------------------------------------------- elementwise + + +def silu_bw(ttnn, g, x): + return ttnn.silu_bw(g, x)[0] + + +# --------------------------------------------------------------------------- RMS norm SH + + +def _rmsnorm_bw_flat(norm, g_out): + """Flat VJP of ``RMSNormSH`` ([N, nsph*C]) -- mirror of ``RMSNormSH._call_flat``. Uses the + cached flat centered input xc [N,nsph*C] and rsqrt scale inv [N,1].""" + ttnn = norm.ttnn + xc = norm._cache_xc # [N, nsph*C] + inv = norm._cache_inv # [N, 1] + N, C, nsph = g_out.shape[0], norm.C, norm.nsph + W = nsph * C + gf = ttnn.reshape(g_out, (N, W)) + # out = xc * (inv * awvec); bias additive (identity wrt input) + g_xc = ttnn.multiply(gf, ttnn.multiply(inv, norm.awvec)) # direct path + # g_inv = sum_j g_out_j * xc_j * awvec_j + g_inv = ttnn.sum(ttnn.multiply(ttnn.multiply(gf, xc), norm.awvec), dim=1, keepdim=True) # [N,1] + # inv=(ms+eps)^-1/2 -> g_ms = -0.5 inv^3 g_inv ; ms = sum_j wvec_j xc_j^2 -> g_xc += g_ms 2 wvec xc + g_ms = ttnn.multiply(g_inv, ttnn.multiply(ttnn.multiply(inv, inv), inv)) + g_ms = ttnn.multiply(g_ms, -0.5) + g_xc = ttnn.add(g_xc, ttnn.multiply(ttnn.multiply(ttnn.multiply(xc, norm.wvec), g_ms), 2.0)) + # centering backward on l0: g_x_l0 = g_xc_l0 - mean_C(g_xc_l0) + g_l0 = ttnn.slice(g_xc, [0, 0], [N, C]) + g_l0 = ttnn.subtract(g_l0, ttnn.mean(g_l0, dim=1, keepdim=True)) + out = ttnn.concat([g_l0, ttnn.slice(g_xc, [0, C], [N, W])], dim=1) + return ttnn.reshape(out, (N, nsph, C)) + + +def rmsnorm_bw(norm, g_out): + """VJP of ``RMSNormSH``. ``norm`` is the forward module (holds bdw, aw, ab, eps). + Reuses the centered input ``xc`` and rsqrt scale ``inv`` cached on the forward.""" + ttnn = norm.ttnn + if getattr(norm, "flat", False): + return _rmsnorm_bw_flat(norm, g_out) + xc = norm._cache_xc # centered input saved on forward + inv = norm._cache_inv # [N,1,1] rsqrt scale saved on forward + N, nsph, C = xc.shape + + # drop the affine bias on l0 (additive -> identity for grad) + s = ttnn.multiply(inv, norm.aw) # [N,nsph,C] via broadcast + g_xc = ttnn.multiply(g_out, s) # direct path + # g_inv = sum_{coeff,C}(g_out * xc * aw) + g_inv = ttnn.sum(ttnn.multiply(ttnn.multiply(g_out, xc), norm.aw), dim=1, keepdim=True) + g_inv = ttnn.sum(g_inv, dim=2, keepdim=True) # [N,1,1] + # inv = (ms+eps)^-1/2 -> d inv/d ms = -1/2 inv^3 + g_ms = ttnn.multiply(g_inv, ttnn.multiply(ttnn.multiply(inv, inv), inv)) + g_ms = ttnn.multiply(g_ms, -0.5) + # ms = mean_C(sum_coeff(xc^2 bdw)) -> d ms/d xc = (1/C) 2 xc bdw + g_xc = ttnn.add(g_xc, ttnn.multiply(ttnn.multiply(ttnn.multiply(xc, norm.bdw), g_ms), 2.0 / C)) + # centering backward on l0: g_x_l0 = g_xc_l0 - mean_C(g_xc_l0) + g_l0 = ttnn.slice(g_xc, [0, 0, 0], [N, 1, C]) + g_l0 = ttnn.subtract(g_l0, ttnn.mean(g_l0, dim=2, keepdim=True)) + g_rest = ttnn.slice(g_xc, [0, 1, 0], [N, nsph, C]) + return ttnn.concat([g_l0, g_rest], dim=1) + + +# --------------------------------------------------------------------------- gate + + +def gate_bw(gate, g_out): + """VJP of ``GateActivation`` (flat). Returns (g_gating[E,lmax*H], g_x flat[E,nsph*H]).""" + ttnn = gate.ttnn + gating = gate._cache_gating # pre-sigmoid [E, lmax*H] + x = gate._cache_x # pre-gate input flat [E, nsph*H] + gate_exp = gate._cache_gate # expanded sigmoid gate [E,(nsph-1)*H] (cached fwd) + E, H, lmax, ei = x.shape[0], gate.H, gate.lmax, gate.expand_index + + from .activation import _FUSED_GATE + if _FUSED_GATE and x.shape[1] % 32 == 0 and gate_exp.shape[1] % 32 == 0: + # one kernel: g_x = [g_out[:,:H]*silu'(x[:,:H]) | g_out[:,H:]*gate_exp] + op = ttnn._ttnn.operations.experimental.fused_gate + g_x = op(g_out, gate_exp, x, x.shape[1] // 32, gate_exp.shape[1] // 32, H // 32, 1) + g_vec = ttnn.slice(g_out, [0, H], [E, x.shape[1]]) + x_vec = ttnn.slice(x, [0, H], [E, x.shape[1]]) + else: + g_scalar = ttnn.slice(g_out, [0, 0], [E, H]) + g_vec = ttnn.slice(g_out, [0, H], [E, x.shape[1]]) + x_scalar = ttnn.slice(x, [0, 0], [E, H]) + x_vec = ttnn.slice(x, [0, H], [E, x.shape[1]]) + g_x = ttnn.concat([silu_bw(ttnn, g_scalar, x_scalar), ttnn.multiply(g_vec, gate_exp)], dim=1) + + g_gate_exp = ttnn.multiply(g_vec, x_vec) # [E, (nsph-1)*H] + # segment-sum the H-blocks back to [E, lmax*H]: transpose of the fwd expand matmul (one op) + g_sig = _mm(ttnn, g_gate_exp, gate.expand_w, gate.kcfg) # [E, lmax*H] + g_gating = ttnn.sigmoid_bw(g_sig, gating)[0] + return g_gating, g_x + + +# --------------------------------------------------------------------------- SO(2) conv + + +def so2_bw(conv, g_out, g_extra=None): + """VJP of ``SO2Convolution``. Returns (g_x flat ``[E, nsph*Cin]``, g_rad or None). + + ``g_rad`` is the adjoint at the radial-MLP *output* (the per-m multiplier), to be finished + on host. Matmul backward = transpose-matmul on device.""" + ttnn = conv.ttnn + kcfg = conv.kcfg + H = conv.H + lmax, mmax = conv.lmax, conv.mmax + nsph = (lmax + 1) ** 2 + E = g_out.shape[0] + + gf = g_out if len(g_out.shape) == 2 else ttnn.reshape(g_out, (E, nsph * H)) + + # per-m fused backward: m0 one transpose-matmul, each m>0 one transpose-matmul on the + # [[Wa,Wb],[-Wb,Wa]] block. Mirrors the collapsed forward (see so2.py _build_fused). + if getattr(conv, "fused_w", None) is not None: + ow = conv.fused_out_w + # split g_out into per-block adjoints [m0coef, m1(2Hh), m2(2Hh), ...] + seg, off = [], 0 + for wd in ow: + seg.append(ttnn.slice(gf, [0, off], [E, off + wd])); off += wd + g_full0 = ttnn.concat([g_extra, seg[0]], dim=1) if conv.extra else seg[0] + g_parts = [_mm(ttnn, g_full0, conv.fused_wm0, kcfg)] # [E, in0] + for m in range(1, mmax + 1): + g_parts.append(_mm(ttnn, seg[m], conv.fused_wm[m - 1], kcfg)) # [E, 2K] + g_xf = ttnn.concat(g_parts, dim=1) # [E, nsph*Cin] pre-radial-multiply + g_rad = None + if conv.has_radial: + xin = conv._cache_xin + mult = conv._cache_mult + # g_rad is the adjoint at the (duplicated) radial output; rad.bw's matmul with the + # duplicated net.6 weight sums the repeated real/imag rows -> the old collapse is implicit. + g_rad = ttnn.multiply(g_xf, xin) + g_xf = ttnn.multiply(g_xf, mult) + return g_xf, g_rad + + # split g into out-blocks: m0 (lmax+1 coeffs), then per m>0 (real Hh, imag Hh) + coeff_w = [] # column width per out-block (in H units) + coeff_w.append((lmax + 1) * H) + for m in range(1, mmax + 1): + Hh = conv.w_m[m - 1].shape[1] // 2 + coeff_w.append(Hh) # real + coeff_w.append(Hh) # imag + seg, off = [], 0 + for wd in coeff_w: + seg.append(ttnn.slice(gf, [0, off], [E, off + wd])); off += wd + + g_blocks = [] # adjoint per input m-block (flattened) + # m = 0 + g_lin = seg[0] + if conv.extra: + g_lin = ttnn.concat([g_extra, g_lin], dim=1) # [E, extra + H*(lmax+1)] + g_blocks.append(_mm(ttnn, g_lin, conv.w_m0, kcfg)) # [E, w0] + # m > 0 + si = 1 + for m in range(1, mmax + 1): + g_real = seg[si]; g_imag = seg[si + 1]; si += 2 # adjoints of out_real, out_imag + # fwd: out_real = r0 - i1, out_imag = i0 + r1, with [r0|r1]=real@W, [i0|i1]=imag@W. + # => g_fr = [g_real, g_imag], g_fi = [g_imag, -g_real]; g_{real,imag} = g_f @ W^T. + g_fr = ttnn.concat([g_real, g_imag], dim=1) # [E,2Hh] + g_fi = ttnn.concat([g_imag, ttnn.multiply(g_real, -1.0)], dim=1) + g_in_real = _mm(ttnn, g_fr, conv.w_m[m - 1], kcfg) # [E, nc*Cin] + g_in_imag = _mm(ttnn, g_fi, conv.w_m[m - 1], kcfg) + g_blocks.append(ttnn.concat([g_in_real, g_in_imag], dim=1)) # [E, 2*nc*Cin] + + g_xf = ttnn.concat(g_blocks, dim=1) # [E, nsph*Cin] + + g_rad = None + if conv.has_radial: + xin = conv._cache_xin # [E, nsph*Cin] pre-multiply + mult = conv._cache_mult # [E, nsph*Cin] + g_mult = ttnn.multiply(g_xf, xin) + g_xf = ttnn.multiply(g_xf, mult) + # collapse duplicated real/imag halves back to per-m radial channels + o = 0 + widths = [conv.rad_sizes[0]] + for m in range(1, mmax + 1): + widths += [conv.rad_sizes[m], conv.rad_sizes[m]] + segs = [] + for wd in widths: + segs.append(ttnn.slice(g_mult, [0, o], [E, o + wd])); o += wd + g_rad_parts = [segs[0]] + i = 1 + for m in range(1, mmax + 1): + g_rad_parts.append(ttnn.add(segs[i], segs[i + 1])); i += 2 + g_rad = ttnn.concat(g_rad_parts, dim=1) # [E, sum rad_sizes] + + return g_xf, g_rad # g_xf flat [E, nsph*Cin] + + +# --------------------------------------------------------------------------- grid atomwise + + +def grid_bw(grid, g_out): + ttnn = grid.ttnn + kcfg = grid.kcfg + N = g_out.shape[0] + from .device import l1_if_fits, L1_NODE_BUDGET # BW-bound [N,npts,C] chain -> L1 while it fits + _npts_pad = ((grid.npts + 31) // 32) * 32 # tile-padded point dim (3D tensor) + L1 = l1_if_fits(ttnn, N, _npts_pad * g_out.shape[2], budget=L1_NODE_BUDGET) + a1, a2 = grid._cache_a1, grid._cache_a2 # pre-silu activations [N,npts,H] + # from_grid backward: o = transpose(gt @ fg); gt = transpose(g_mlp_out) + # forward: gt=transpose(mlp,1,2); o=gt@fg; out=transpose(o,1,2) + g_o = ttnn.transpose(g_out, 1, 2, memory_config=L1) # [N,C,nsph] + g_gt = _mm(ttnn, g_o, grid.fg, kcfg, memory_config=L1) # [N,C,npts] + g_mlp = ttnn.transpose(g_gt, 1, 2, memory_config=L1) # [N,npts,C] + # mlp backward (no bias): a3 = s2@W4 ; s2=silu(a2); a2=s1@W2; s1=silu(a1); a1=g0@W0 + g_s2 = _mm(ttnn, g_mlp, grid.w4, kcfg, memory_config=L1) + g_a2 = silu_bw(ttnn, g_s2, a2) + g_s1 = _mm(ttnn, g_a2, grid.w2, kcfg, memory_config=L1) + g_a1 = silu_bw(ttnn, g_s1, a1) + g_g0 = _mm(ttnn, g_a1, grid.w0, kcfg, memory_config=L1) # [N,npts,C] + # to_grid backward: g0 = transpose(xt @ tg_T); xt=transpose(x,1,2) + g_g0t = ttnn.transpose(g_g0, 1, 2, memory_config=L1) # [N,C,npts] + g_xt = _mm(ttnn, g_g0t, grid.tg_T, kcfg) # [N,C,nsph] -> DRAM (feeds residual add) + return ttnn.transpose(g_xt, 1, 2) # [N,nsph,C] + + +# --------------------------------------------------------------------------- spectral atomwise + + +def _so3_linear_bw(sp, g_out, w_blocks, wf=None, cin=None): + """VJP of one ``SO3_Linear`` wrt its input. ``g_out`` [N,nsph,cout] -> g_x [N,nsph,cin]. + Per degree the GEMM is shared, so the input adjoint is ``g_out_block @ W_block^T`` (bias on + l=0 is additive -> identity wrt the input). When ``wf`` (the fused block-diagonal weight) is + given, this is one flat transpose-matmul (mirrors the fused forward).""" + ttnn = sp.ttnn + N = g_out.shape[0] + if wf is not None: + gf = ttnn.reshape(g_out, (N, sp.nsph * g_out.shape[2])) + gx = _mm(ttnn, gf, wf, sp.kcfg) # [N, nsph*cin] + return ttnn.reshape(gx, (N, sp.nsph, cin)) + outs, start = [], 0 + for l in range(sp.lmax + 1): + n = 2 * l + 1 + gb = ttnn.slice(g_out, [0, start, 0], [N, start + n, g_out.shape[2]]) + outs.append(_mm(ttnn, gb, w_blocks[l], sp.kcfg)) + start += n + return ttnn.concat(outs, dim=1) + + +def _spectral_bw_flat(sp, g_out): + """Fully-flat VJP of ``SpectralAtomwise`` -- mirror of ``SpectralAtomwise._call_flat``.""" + ttnn = sp.ttnn + N, H, C, nsph = g_out.shape[0], sp.H, sp.C, sp.nsph + xf = sp._cache_xf # [N, nsph*C] + a_scalar, gating, hf = sp._cache_a_scalar, sp._cache_gating, sp._cache_hf + gf = ttnn.reshape(g_out, (N, nsph * C)) + # so3_linear_2 backward (flat block-diagonal transpose-matmul) + g_g = _mm(ttnn, gf, sp.l2_wf, sp.kcfg) # [N, nsph*H] + # gate backward (flat): l0 SiLU; vector = h_vec * gate_exp + sg = ttnn.sigmoid(gating) + gate_exp = ttnn.matmul(sg, sp.gate_exp_w, compute_kernel_config=sp.kcfg) # [N,(nsph-1)*H] + g_scalar_h = silu_bw(ttnn, ttnn.slice(g_g, [0, 0], [N, H]), ttnn.slice(hf, [0, 0], [N, H])) + g_vec = ttnn.slice(g_g, [0, H], [N, nsph * H]) + g_h_vec = ttnn.multiply(g_vec, gate_exp) # g wrt h vector blocks + g_h = ttnn.concat([g_scalar_h, g_h_vec], dim=1) # [N, nsph*H] + # g wrt the (expanded) gate = g_vec * h_vec ; contract back to [N, lmax*H] via Ex^T + g_gate_exp = ttnn.multiply(g_vec, ttnn.slice(hf, [0, H], [N, nsph * H])) + g_sg = _mm(ttnn, g_gate_exp, sp.gate_exp_w, sp.kcfg) # [N, lmax*H] + g_gating = ttnn.sigmoid_bw(g_sg, gating)[0] + # so3_linear_1 backward -> g wrt x + g_x = _mm(ttnn, g_h, sp.l1_wf, sp.kcfg) # [N, nsph*C] + # scalar_mlp backward: add g_scalar onto x's l=0 block + g_a = silu_bw(ttnn, g_gating, a_scalar) + g_scalar = _mm(ttnn, g_a, sp.smlp_w, sp.kcfg) # [N, C] + g_x_l0 = ttnn.add(ttnn.slice(g_x, [0, 0], [N, C]), g_scalar) + out = ttnn.concat([g_x_l0, ttnn.slice(g_x, [0, C], [N, nsph * C])], dim=1) + return ttnn.reshape(out, (N, nsph, C)) + + +def spectral_bw(sp, g_out): + """VJP of ``SpectralAtomwise``. ``g_out`` [N,nsph,C] -> g wrt input x [N,nsph,C].""" + ttnn = sp.ttnn + if getattr(sp, "gate_exp_w", None) is not None: + return _spectral_bw_flat(sp, g_out) + N, H, C = g_out.shape[0], sp.H, sp.C + a_scalar = sp._cache_a_scalar + gating, h = sp._cache_gating, sp._cache_h + + # so3_linear_2 backward + g_g = _so3_linear_bw(sp, g_out, sp.l2_w, wf=sp.l2_wf, cin=H) # [N, nsph, H] + + # gate backward: l0 SiLU; l>=1 multiply by sigmoid(gating) per degree + sg = ttnn.sigmoid(gating) # [N, lmax*H] + g_h_parts = [silu_bw(ttnn, ttnn.slice(g_g, [0, 0, 0], [N, 1, H]), + ttnn.slice(h, [0, 0, 0], [N, 1, H]))] + g_sg_rows, start = [], 1 + for l in range(1, sp.lmax + 1): + n = 2 * l + 1 + g_gb = ttnn.slice(g_g, [0, start, 0], [N, start + n, H]) + h_b = ttnn.slice(h, [0, start, 0], [N, start + n, H]) + gl = ttnn.reshape(ttnn.slice(sg, [0, (l - 1) * H], [N, l * H]), (N, 1, H)) + g_h_parts.append(ttnn.multiply(g_gb, gl)) # g wrt h block + # g wrt the (broadcast) gate = sum over the n coeffs of g_gb * h_block + g_sg_rows.append(ttnn.sum(ttnn.multiply(g_gb, h_b), dim=1)) # [N, H] + start += n + g_h = ttnn.concat(g_h_parts, dim=1) # [N, nsph, H] + g_sg = ttnn.concat(g_sg_rows, dim=1) # [N, lmax*H] + g_gating = ttnn.sigmoid_bw(g_sg, gating)[0] # through sigmoid + + # so3_linear_1 backward -> g wrt x (l>=... all degrees) + g_x = _so3_linear_bw(sp, g_h, sp.l1_w, wf=sp.l1_wf, cin=sp.C) # [N, nsph, C] + + # scalar_mlp backward: gating = SiLU(scalar @ W + b); add g_scalar onto x's l=0 channel + g_a = silu_bw(ttnn, g_gating, a_scalar) # [N, lmax*H] + g_scalar = _mm(ttnn, g_a, sp.smlp_w, sp.kcfg) # [N, C] + g_scalar = ttnn.reshape(g_scalar, (N, 1, C)) + g_x_l0 = ttnn.add(ttnn.slice(g_x, [0, 0, 0], [N, 1, C]), g_scalar) + g_x_rest = ttnn.slice(g_x, [0, 1, 0], [N, sp.nsph, C]) + return ttnn.concat([g_x_l0, g_x_rest], dim=1) + + +# --------------------------------------------------------------------------- edgewise + + +def edgewise_bw(ew, graph, g_out, acc): + """VJP of ``Edgewise`` (flat MAC rotations). Returns g wrt node features [N,nsph,C]; + accumulates the geometric coefficient adjoints (g rot_fwd / rot_inv, g_envelope) and the + radial adjoint g_rad into ``acc``.""" + from . import rotation + ttnn = ew.ttnn + kcfg = ew.kcfg + C = ew.C + N, nsph = g_out.shape[0], g_out.shape[1] + E = graph.E + dev = ew.device + + from .device import bf8_edge + _b8 = bf8_edge() + # scatter backward: g_m_back[e] = g_out[tgt[e]] (gather by target), flat [E, nsph*C] + gof = ttnn.to_layout(ttnn.reshape(g_out, (N, nsph * C)), ttnn.ROW_MAJOR_LAYOUT) + g_mback = ttnn.to_layout(ttnn.embedding(graph.tgt_idx, gof), ttnn.TILE_LAYOUT) + if _b8: # bf8-edge boundary (bw): the reverse edge flow runs bf8 from here (see fwd in edgewise.py) + g_mback = ttnn.typecast(g_mback, ttnn.bfloat8_b) + # inverse rotation backward: forward mapped reduced m-space (nred) -> node SH (nsph) + g_menv, g_rinv = rotation.rotate_bw(ttnn, ew._cache_menv, g_mback, graph.rot_inv_ij, + graph.rot_inv_coef, graph.nred, C, dev, n_out=nsph) + # envelope: m_env = m_so2 * envelope (flat [E,9C] * [E,1]) + g_mso2 = ttnn.multiply(g_menv, graph.edge_envelope_f) + g_env = ttnn.sum(ttnn.multiply(g_menv, ew._cache_mso2), dim=1, keepdim=True) # [E,1] + # so2_2 -> gate -> so2_1 (all flat) + g_mgate, _ = so2_bw(ew.so2_2, g_mso2) + g_gating, g_mso1 = gate_bw(ew.gate, g_mgate) + g_mrot, g_rad = so2_bw(ew.so2_1, g_mso1, g_gating) # g_mrot flat [E, 9*2C] + # forward rotation backward: forward mapped node SH (nsph) -> reduced m-space (nred) + g_mcat, g_rfwd = rotation.rotate_bw(ttnn, ew._cache_mcat, g_mrot, graph.rot_fwd_ij, + graph.rot_fwd_coef, nsph, 2 * C, dev, n_out=graph.nred) + # gather backward: g_nodes = scatter_src(g_xs) + scatter_tgt(g_xt), where g_mcat per coord = + # [g_xs | g_xt] interleaved. For the matmul scatter path, matmul commutes with the channel + # slice, so scatter the FULL interleaved g_mcat then deinterleave the src/tgt channels on the + # N-sized [N,nsph*2C] result (~46x fewer rows than deinterleaving g_mcat at E) -- the E-sized + # RM<->TILE deinterleave was ~5 ms/block. Bit-identical (linearity). + if graph.linear_scatter: + # large-N linear path: deinterleave at E (RM round-trip, avoids the 9->32 tile-pad repack) + from . import scatter + W = nsph * C + g_mcat_rm = ttnn.to_layout(g_mcat, ttnn.ROW_MAJOR_LAYOUT) + g_mcat3 = ttnn.reshape(g_mcat_rm, (E, nsph, 2 * C)) + g_xs_f = ttnn.to_layout(ttnn.reshape(ttnn.slice(g_mcat3, [0, 0, 0], [E, nsph, C]), (E, W)), + ttnn.TILE_LAYOUT) + g_xt_f = ttnn.to_layout(ttnn.reshape(ttnn.slice(g_mcat3, [0, 0, C], [E, nsph, 2 * C]), (E, W)), + ttnn.TILE_LAYOUT) + g_nodes = ttnn.add(scatter.segment_sum(ttnn, g_xs_f, graph.src_gather, graph.Dmax_s, N, W), + scatter.segment_sum(ttnn, g_xt_f, graph.tgt_gather, graph.Dmax_t, N, W)) + g_nodes = ttnn.reshape(g_nodes, (N, nsph, C)) + else: + odt = ttnn.bfloat16 if _b8 else None # back to bf16 node grads + A = ttnn.matmul(graph.scatter_src, g_mcat, dtype=odt, compute_kernel_config=kcfg) # [N, nsph*2C] + B = ttnn.matmul(graph.scatter, g_mcat, dtype=odt, compute_kernel_config=kcfg) # [N, nsph*2C] + A_rm = ttnn.reshape(ttnn.to_layout(A, ttnn.ROW_MAJOR_LAYOUT), (N, nsph, 2 * C)) + B_rm = ttnn.reshape(ttnn.to_layout(B, ttnn.ROW_MAJOR_LAYOUT), (N, nsph, 2 * C)) + A_src = ttnn.to_layout(ttnn.reshape(ttnn.slice(A_rm, [0, 0, 0], [N, nsph, C]), (N, nsph * C)), + ttnn.TILE_LAYOUT) + B_tgt = ttnn.to_layout(ttnn.reshape(ttnn.slice(B_rm, [0, 0, C], [N, nsph, 2 * C]), (N, nsph * C)), + ttnn.TILE_LAYOUT) + g_nodes = ttnn.reshape(ttnn.add(A_src, B_tgt), (N, nsph, C)) + + acc["rot_fwd"] = g_rfwd if acc["rot_fwd"] is None else ttnn.add(acc["rot_fwd"], g_rfwd) + acc["rot_inv"] = g_rinv if acc["rot_inv"] is None else ttnn.add(acc["rot_inv"], g_rinv) + acc["envelope"] = g_env if acc["envelope"] is None else ttnn.add(acc["envelope"], g_env) + acc["g_rad"].append((ew.so2_1, g_rad)) + return g_nodes + + +# --------------------------------------------------------------------------- block / backbone + + +def block_bw(blk, graph, g_out, acc): + """VJP of ``_Block``. ``blk`` is the forward block module.""" + ttnn = blk.norm_1.ttnn + # x = atom_wise(n2) + x_res2 + g_n2 = (spectral_bw(blk.atom_wise, g_out) if getattr(blk, "ff_type", "grid") == "spectral" + else grid_bw(blk.atom_wise, g_out)) + g_after_edge = ttnn.add(rmsnorm_bw(blk.norm_2, g_n2), g_out) + # x = edge_wise(s) + x_res + g_s = edgewise_bw(blk.edge_wise, graph, g_after_edge, acc) + # s = n1 with l0 += sys_emb (identity wrt n1); add residual + g_x_in = ttnn.add(rmsnorm_bw(blk.norm_1, g_s), g_after_edge) + return g_x_in + + +def energy_bw(bb, node_emb): + """VJP of the energy head; returns g wrt node_emb [N,nsph,C] (only l=0 nonzero). + + The two constants (the ``dE/dE = 1`` seed and the l>=1 zero padding) are created once and + cached on ``bb`` so that no device buffer is allocated inside a captured trace region — the + ttnn trace machinery forbids allocations during capture (it hangs).""" + ttnn = bb.ttnn + kcfg = bb.kcfg + N, nsph, C = node_emb.shape + if getattr(bb, "_bw_seed", None) is None or tuple(bb._bw_seed.shape) != (N, 1): + bb._bw_seed = ttnn.ones((N, 1), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=bb.device) + bb._bw_zeros = ttnn.zeros((N, nsph - 1, C), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, + device=bb.device) + h = ttnn.reshape(ttnn.slice(node_emb, [0, 0, 0], [N, 1, C]), (N, C)) + a1 = ttnn.linear(h, bb.eh_w[0], bias=bb.eh_b[0], compute_kernel_config=kcfg) + s1 = ttnn.silu(a1) + a2 = ttnn.linear(s1, bb.eh_w[1], bias=bb.eh_b[1], compute_kernel_config=kcfg) + g_s2 = _mm(ttnn, bb._bw_seed, bb.eh_w[2], kcfg) + g_a2 = silu_bw(ttnn, g_s2, a2) + g_s1 = _mm(ttnn, g_a2, bb.eh_w[1], kcfg) + g_a1 = silu_bw(ttnn, g_s1, a1) + g_h = _mm(ttnn, g_a1, bb.eh_w[0], kcfg) # [N,C] + g_h = ttnn.reshape(g_h, (N, 1, C)) + return ttnn.concat([g_h, bb._bw_zeros], dim=1) + + +def backbone_bw(bb, graph, node_emb): + """Full reverse pass of the backbone+energy head. Returns a dict of device adjoints + (g_x_init, g_wigner, g_wigner_inv, g_envelope, g_x_edge). ``g_x_edge`` is finished on device + (radial-MLP backward), so the host only reads back adjoints and drives the geometric autograd.""" + from .model import balance_l0 + + ttnn = bb.ttnn + acc = {"rot_fwd": None, "rot_inv": None, "envelope": None, "g_rad": []} + g = energy_bw(bb, node_emb) + g = rmsnorm_bw(bb.final_norm, g) + do_bal = bb.ce > bb.cs + for blk in reversed(bb.blocks): + # balance_channels is self-adjoint (a projection I - mean); its VJP is the same op with + # zero target. It followed each block in the forward, so undo it before that block's VJP. + if do_bal: + g = balance_l0(bb.ttnn, g, graph.node_meanM, bb.cs, bb.ce, 0.0, bb.kcfg) + g = block_bw(blk, graph, g, acc) + acc["x_init"] = g + # device edge-degree embedding backward: the node adjoint ``g`` is backpropped through the + # on-device node init (scatter -> rotate-back -> envelope -> radial MLP), accumulating the + # geometric adjoints (rot_inv, envelope) and the radial adjoint into ``acc``. The host then + # never differentiates x_init -- only wigner_inv / x_edge / envelope (much cheaper). + if getattr(bb, "edge_degree", None) is not None: + from .edge_degree import edge_degree_bw + edge_degree_bw(bb.edge_degree, graph, g, acc) + # radial-MLP backward on device: each conv's radial adjoint (at the radial output) is finished + # to g wrt the shared invariant edge embedding x_edge on device (hand-written LN/SiLU VJP), + # summed across convs. This replaces the host torch.autograd radial finish (~100 ms at N=128) + # with a captured device pass -- only a single [E, x_edge] readback remains on host. + g_xe = None + for conv, g_rad in acc["g_rad"]: + # radial backward: rad.bw runs in fp32 by default (small hidden=128; a bf16 radial VJP + # mis-directs forces on OOD compressed heavy cells, so it upcasts internally). The opt-in + # fused_ln_bw kernel path stays bf16. Cast the bf8-edge g_rad back to bf16 for rad.bw's input. + if g_rad.dtype == ttnn.bfloat8_b: + g_rad = ttnn.typecast(g_rad, ttnn.bfloat16) + gc = conv.rad.bw(g_rad) + g_xe = gc if g_xe is None else ttnn.add(g_xe, gc) + acc["x_edge"] = g_xe + return acc + + +# --------------------------------------------------------------------------- full energy+force + + +def _forward(bb, geo, pos, atomic_numbers, edge_index, sys_node_embedding, edge_cell_shift, + requires_grad, compute_stress=False, force_linear_scatter=None, + charge=0.0, system_natoms=None): + """Shared forward: host geometry -> device-resident GraphContext + backbone node embedding. + Returns ``(node_emb, graph, t, pos_leaf, strain_leaf)``; ``pos_leaf`` tracks grad when + ``requires_grad``. ``strain_leaf`` is a zero symmetric 3x3 leaf (else None) that is applied to + the edge vectors as ``r' = r(I + sym(strain))`` — since ALL pos/cell dependence of the energy + flows through the edge vectors, ``dE/dstrain`` is exactly fairchem's symmetrized virial (the + combined position + cell contribution), so stress = dE/dstrain / volume.""" + import ttnn + + from .model import GraphContext + + device = bb.device + N, C = atomic_numbers.shape[0], geo.C + pos = pos.detach().clone().requires_grad_(requires_grad) + strain, edge_vec = None, None + if compute_stress: + src, tgt = edge_index[0], edge_index[1] + strain = torch.zeros(3, 3, dtype=pos.dtype, requires_grad=True) + ev = pos[src] - pos[tgt] + if edge_cell_shift is not None: + ev = ev + edge_cell_shift + sym = 0.5 * (strain + strain.transpose(0, 1)) + edge_vec = ev + ev @ sym # r' = r (I + sym(strain)) + t = geo(pos, atomic_numbers, edge_index, sys_node_embedding, edge_vec=edge_vec, + edge_cell_shift=edge_cell_shift) + + # the analytic-force backward keeps bf16 geometric operands (bf8 wigner would mix dtypes in + # the transpose-matmul adjoints); ``fast`` (bf8) is for the energy-throughput path. + graph = GraphContext(device, edge_index=edge_index, wigner=t["wigner"].detach(), + wigner_inv=t["wigner_inv"].detach(), x_edge=t["x_edge"].detach(), + edge_envelope=t["edge_envelope"].detach(), num_nodes=N, + linear_scatter=force_linear_scatter, + system_natoms=system_natoms, build_mean_op=(bb.ce > bb.cs)) + se3 = ttnn.from_torch(sys_node_embedding.reshape(N, 1, C), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + # with the device edge-degree embedding on, the backbone builds x_init on device and this + # operand is the constant l0 node init; otherwise it is the host-computed full x_init. + init = t["l0"] if bb.edge_degree is not None else t["x_init"] + x_init = ttnn.from_torch(init.detach(), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + # per-atom charge target = charge/natoms (uniform: one bundle is one composition/charge, and a + # same-composition batch shares natoms). 0 when neutral or when balancing is disabled. + per_sys_n = int(system_natoms[0]) if system_natoms else N + balance_add = (float(charge) / per_sys_n) if bb.ce > bb.cs else 0.0 + node_emb = bb.node_embedding(x_init, graph, se3, balance_add) + return node_emb, graph, t, pos, strain + + +def _forces(bb, geo, graph, node_emb, t, pos, strain=None): + """Reverse pass: device VJP ``dE/d{geometric inputs}`` finished by host autograd to ``-dE/dpos``. + + The energy seed is ``dE/dh = 1`` per node — the gradient of the *summed* energy. For a + disjoint-union batch that sum is ``sum_k E_k`` and block-diagonality makes each atom's + gradient ``-dE_(its system)/dx``, so the batched forces are the concatenation (no change). + + When ``strain`` is given (a zero symmetric 3x3 leaf that scaled the edge vectors in the + forward), the same host autograd also yields the virial ``dE/dstrain`` and this returns + ``(forces, virial)``; the caller divides the virial by the volume for the stress tensor.""" + import ttnn + + from . import rotation + + acc = backbone_bw(bb, graph, node_emb) + nsph, nred = graph.nsph, graph.nred + device_ede = bb.edge_degree is not None + # x_init adjoint is only differentiated on host when x_init is a host term; with the device + # edge-degree embedding it is consumed on device (its geometric adjoints land in acc below). + g_xi = None if device_ede else ttnn.to_torch(acc["x_init"]).float() + # scatter packed rotation-coefficient adjoints back to dense for the host dW/dpos autograd. + # wig_M is [E, nred, nsph] (fwd), wig_M_inv is [E, nsph, nred] (inv) — rectangular for uma-m. + g_wig = rotation.scatter_coef(ttnn.to_torch(acc["rot_fwd"]).float(), graph.rot_fwd_ij, nred, nsph) + g_winv = rotation.scatter_coef(ttnn.to_torch(acc["rot_inv"]).float(), graph.rot_inv_ij, nsph, nred) + g_env = ttnn.to_torch(acc["envelope"]).float().reshape(-1, 1, 1) # [E,1]->[E,1,1] + # radial finish is done on device (see backbone_bw); read back the single g_x_edge adjoint. + # x_edge = [gaussian(dist) | src_emb | tgt_emb]; only the gaussian block depends on pos, so the + # force VJP needs only its adjoint. Cast just that block bf16->f32 (the cast dominates readback); + # the pos-independent embedding columns contribute zero to dE/dpos. + ng = geo.offset.shape[0] + W = acc["x_edge"].shape[1] + gx = ttnn.to_torch(ttnn.slice(acc["x_edge"], [0, 0], [acc["x_edge"].shape[0], ng])) + g_xe = torch.zeros((gx.shape[0], W), dtype=torch.float32) + g_xe[:, :ng] = gx.float() + + outs = [t["wigner"], t["wigner_inv"], t["x_edge"], t["edge_envelope"]] + gouts = [g_wig, g_winv, g_xe, g_env] + if not device_ede: + outs = [t["x_init"]] + outs + gouts = [g_xi] + gouts + inputs = [pos] if strain is None else [pos, strain] + grads = torch.autograd.grad(outs, inputs, grad_outputs=gouts) + if strain is None: + return -grads[0] + return -grads[0], grads[1] # (forces, virial = dE/dstrain) + + +def energy_and_forces(bb, geo, pos, atomic_numbers, edge_index, sys_node_embedding, + edge_cell_shift=None, compute_stress=False, charge=0.0): + """Conservative energy + analytic forces ``F = -dE/dpos`` for one system. + + Device-resident forward + reverse VJP gives ``dE/d{geometric inputs}``; ``torch.autograd`` + through the host geometry supplies the cheap ``d(geometric)/dpos`` to finish the force. + ``edge_cell_shift`` [E, 3] carries the periodic image offsets (None for aperiodic systems). + Returns ``(energy: float, forces: torch.Tensor[N,3])``, or when ``compute_stress`` is set + ``(energy, forces, virial[3,3])`` where ``virial = dE/dstrain`` (the caller divides by the + cell volume for the stress tensor). + """ + import ttnn + + node_emb, graph, t, pos, strain = _forward( + bb, geo, pos, atomic_numbers, edge_index, sys_node_embedding, edge_cell_shift, + requires_grad=True, compute_stress=compute_stress, charge=charge) + E = float(ttnn.to_torch(bb.energy(node_emb)).reshape(-1)[0]) + if compute_stress: + F, virial = _forces(bb, geo, graph, node_emb, t, pos, strain=strain) + return E, F, virial + F = _forces(bb, geo, graph, node_emb, t, pos) + return E, F + + +def energy_and_forces_batch(bb, geo, bg, *, compute_forces=True): + """Disjoint-union batched energy (+ optional forces) for a ``disjoint.BatchedGraph`` ``bg``. + + One device forward over the concatenated block-diagonal graph; per-system energies come from + the segment-sum readout (``Backbone.energy_batch``) and forces — when requested — from the + single shared reverse pass (block-diagonal => per-system correct). Returns + ``(E_raw: torch[K], F: torch[Ntot, 3] or None)`` where ``E_raw`` is unnormalized (the caller + applies the per-system energy normalizer).""" + import ttnn + + # A disjoint-union batch is block-diagonal: the dense one-hot scatter S[Ntot,Etot] wastes + # O(K^2) on off-diagonal zeros, so force the linear O(Etot) gather+reduce for K>1 (measured + # +60% peak throughput at K=128 ethanol; see benchmarks/bench_batch.py). K=1 keeps the + # single-system node-count threshold (dense is ~5x faster at small N). + node_emb, graph, t, pos, _ = _forward(bb, geo, bg.pos, bg.Z, bg.edge_index, bg.sys_emb, + bg.cell_shift, requires_grad=compute_forces, + force_linear_scatter=(bg.K > 1 or None), + charge=bg.charge, system_natoms=bg.natoms) + seg = ttnn.from_torch(bg.segment_matrix(), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, + device=bb.device) + E = ttnn.to_torch(bb.energy_batch(node_emb, seg)).float().reshape(-1)[:bg.K] + F = _forces(bb, geo, graph, node_emb, t, pos) if compute_forces else None + return E, F diff --git a/tt_atom/geometry.py b/tt_atom/geometry.py new file mode 100644 index 0000000..c169992 --- /dev/null +++ b/tt_atom/geometry.py @@ -0,0 +1,290 @@ +"""Host geometry: the differentiable ``pos -> {geometric device inputs}`` map (torch, host). + +These are the per-edge geometric terms (Wigner rotation, radial edge embedding, envelope, the +edge-degree node init) that the device backbone consumes as fixed inputs. They are <1% of the +compute, so we keep them on host where ``torch.autograd`` supplies the cheap geometric Jacobian +``d(terms)/dpos`` for the analytic force. Nothing here imports fairchem (it must coexist with +ttnn / numpy<2); the pure-torch rotation helpers are vendored from fairchem (MIT) with the +e3nn 0.4.0 Wigner-D construction they themselves borrow. + +The roll angle ``gamma`` is a gauge the architecture is invariant to; we fix it (default 0) so +the geometry — and therefore the forward and the force — is deterministic. +""" +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +from . import quaternion + +EPS = 1e-7 + + +def csd_embedding(w, charge, spin, sphere_channels, dataset="omat"): + """System (charge/spin/dataset) embedding -> [nsys, C]. Mirrors fairchem ``csd_embedding``. + + Supports both charge/spin encodings: ``pos_emb`` (sin/cos of a random projection, the + self-chosen random-weight config) and ``rand_emb`` (a learned lookup table indexed by + charge+100 / spin, which uma-s-1 uses). The dataset token is the per-dataset embedding for + the active task (``omol``/``omat``/``oc20``/...). Dispatch is by which keys the bundle carries.""" + if "charge_embedding.rand_emb.weight" in w: # rand_emb (uma-s-1) + chg = F.embedding((charge.long() + 100), w["charge_embedding.rand_emb.weight"]) + sp = F.embedding(spin.long(), w["spin_embedding.rand_emb.weight"]) + else: # pos_emb (sin/cos) + def _cs(x, W, is_spin): + xp = x[:, None] * W[None, :] * 2 * math.pi + emb = torch.cat([torch.sin(xp), torch.cos(xp)], dim=-1) + if is_spin: + emb = emb.clone() + emb[torch.where(x == 0)[0]] = 0 + return emb + + chg = _cs(charge, w["charge_embedding.W"], False) + sp = _cs(spin, w["spin_embedding.W"], True) + ds = w[f"dataset_embedding.dataset_emb_dict.{dataset}.weight"][0].expand(charge.shape[0], -1) + return F.silu(F.linear(torch.cat([chg, sp, ds], dim=1), w["mix_csd.weight"], w["mix_csd.bias"])) + + +def radius_graph(pos, cutoff, cell=None, pbc=None): + """Brute-force O(N^2) neighbour list -> ``(edge_index[2, E], cell_shift[E, 3])``. + + ``cell_shift`` is the cartesian periodic image offset for each edge (zeros when aperiodic), + so the caller forms ``edge_vec = pos[src] - pos[tgt] + cell_shift``. Convention matches + fairchem ``radius_graph_pbc`` + ``get_pbc_distances`` exactly: an edge (src = imaged j, + tgt = i) has ``distance_vec = pos[j] - pos[i] + n·cell``. The graph is host-side and a + negligible fraction of the compute, so the O(N^2 · n_cells) brute force is fine for a cell. + + ``cell`` rows are the lattice vectors (ASE convention). ``pbc`` is a bool or length-3 mask; + an aperiodic graph results when ``cell``/``pbc`` are absent or ``pbc`` is all-False.""" + periodic = cell is not None and pbc is not None and bool(torch.as_tensor(pbc).any()) + if not periodic: + d = torch.linalg.norm(pos[:, None, :] - pos[None, :, :], dim=-1) + mask = (d < cutoff) & (d > 0) + src, tgt = torch.where(mask) + return torch.stack([src, tgt], dim=0), pos.new_zeros(int(mask.sum()), 3) + + cell = torch.as_tensor(cell, dtype=pos.dtype) # rows = lattice vectors + pbc = torch.as_tensor(pbc, dtype=torch.bool).reshape(-1).expand(3) + # perpendicular plane spacing along a_k is 1/||b_k|| (b = reciprocal rows); the number of + # image cells needed each way is ceil(cutoff · ||b_k||) — matches fairchem's rep_a{1,2,3}. + recip = torch.linalg.inv(cell).transpose(0, 1) # rows = reciprocal vectors + reps = [int(math.ceil(float(cutoff * torch.linalg.norm(recip[k])))) if bool(pbc[k]) else 0 + for k in range(3)] + ranges = [torch.arange(-r, r + 1, dtype=pos.dtype) for r in reps] + cells = torch.cartesian_prod(*ranges).reshape(-1, 3) # [n_cells, 3] integer offsets + shifts = cells @ cell # [n_cells, 3] cartesian + # disp[i, j, c] = pos[i] - (pos[j] + shift[c]); mask magnitude, keep (i, j, c) within cutoff + disp = pos[:, None, None, :] - (pos[None, :, None, :] + shifts[None, None, :, :]) + d2 = (disp ** 2).sum(-1) # [N, N, n_cells] + mask = (d2 <= cutoff * cutoff) & (d2 > 1e-8) + i, j, c = torch.where(mask) # i = receiver, j = imaged source + return torch.stack([j, i], dim=0), shifts[c] + + +# ----------------------------------------------------------- rotation (vendored, MIT/e3nn 0.4) + + +class _Safeacos(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x.clamp(-1 + EPS, 1 - EPS)) + return torch.acos(x) + + @staticmethod + def backward(ctx, g): + (xc,) = ctx.saved_tensors + return -g / torch.sqrt(1 - xc.pow(2)).clamp(min=EPS) + + +class _Safeatan2(torch.autograd.Function): + @staticmethod + def forward(ctx, y, x): + ctx.save_for_backward(y, x) + return torch.atan2(y, x) + + @staticmethod + def backward(ctx, g): + y, x = ctx.saved_tensors + denom = (x.pow(2) + y.pow(2)).clamp(min=EPS) + return g * x / denom, -g * y / denom + + +def _euler_angles(edge_vec, gamma_val): + xyz = F.normalize(edge_vec).clamp(-1.0, 1.0) + x, y, z = torch.split(xyz, 1, dim=1) + beta = _Safeacos.apply(y.squeeze(-1)) + alpha = _Safeatan2.apply(x.squeeze(-1), z.squeeze(-1)) + gamma = torch.full_like(alpha, gamma_val) + return -gamma, -beta, -alpha # intrinsic -> extrinsic + + +_ZROT_FREQS: dict = {} + + +def _z_rot_mat(angle, lv): + """Wigner z-rotation block: cos on the diagonal, sin on the anti-diagonal (frequency order + ``l .. -l``). Built functionally (``diag_embed`` + column-flip) instead of a Python loop of + per-element in-place writes — the matrix is identical (diagonal and anti-diagonal overlap only + at the centre, where ``cos(0) + sin(0) = 1``) but the autograd graph has far fewer nodes, so + the analytic-force VJP through the Wigner build is ~2x cheaper. Forward is bit-exact vs the + loop; the gradient differs only by float reduction order (~1e-6).""" + freqs = _ZROT_FREQS.get((lv, angle.dtype)) + if freqs is None: + freqs = torch.arange(lv, -lv - 1, -1, dtype=angle.dtype) + _ZROT_FREQS[(lv, angle.dtype)] = freqs + fa = freqs * angle[..., None] + return torch.diag_embed(torch.cos(fa)) + torch.diag_embed(torch.sin(fa)).flip(-1) + + +def _wigner_D(lv, alpha, beta, gamma, Jd): + alpha, beta, gamma = torch.broadcast_tensors(alpha, beta, gamma) + J = Jd[lv] + return _z_rot_mat(alpha, lv) @ J @ _z_rot_mat(beta, lv) @ J @ _z_rot_mat(gamma, lv) + + +def _eulers_to_wigner(eulers, lmax, Jd): + alpha, beta, gamma = eulers + size = (lmax + 1) ** 2 + wigner = torch.zeros(len(alpha), size, size, dtype=alpha.dtype) + start = 0 + for lv in range(lmax + 1): + blk = _wigner_D(lv, alpha, beta, gamma, Jd) + end = start + blk.shape[1] + wigner[:, start:end, start:end] = blk + start = end + return wigner + + +# --------------------------------------------------------------------------- radial MLP (host) + + +def radial_mlp(x, w, p): + x = F.linear(x, w[f"{p}.net.0.weight"], w[f"{p}.net.0.bias"]) + x = F.layer_norm(x, (x.shape[-1],), w[f"{p}.net.1.weight"], w[f"{p}.net.1.bias"], 1e-5) + x = F.silu(x) + x = F.linear(x, w[f"{p}.net.3.weight"], w[f"{p}.net.3.bias"]) + x = F.layer_norm(x, (x.shape[-1],), w[f"{p}.net.4.weight"], w[f"{p}.net.4.bias"], 1e-5) + x = F.silu(x) + return F.linear(x, w[f"{p}.net.6.weight"], w[f"{p}.net.6.bias"]) + + +# --------------------------------------------------------------------------- the geometry + + +class HostGeometry: + def __init__(self, weights, cfg, to_m, gauss_offset, gauss_coeff, *, gamma=0.0, + coefficient_index=None, use_quaternion=True): + self.w = weights + self.cfg = cfg + self.lmax = cfg["lmax"] + self.mmax = cfg.get("mmax", self.lmax) + self.C = cfg["sphere_channels"] + self.cutoff = cfg["cutoff"] + self.gamma = gamma + # edge->+Y frame: fairchem's default smooth two-chart QUATERNION (non-singular everywhere) + # vs the legacy ZYZ-Euler (singular on the +-Y axis -> wrong forces at exact symmetry). The + # Euler path (Jd below) is kept only so the A/B harness can still exercise the old behaviour. + self.use_quaternion = use_quaternion + self.qkernels = quaternion.WignerKernels(self.lmax) if use_quaternion else None + self.Jd = [weights[f"Jd_{l}"] for l in range(self.lmax + 1)] + self.to_m = to_m + # to_m is a permutation matrix (one 1 per row): the m-mapping einsums in _wigner are just a + # coefficient reorder. Precompute the permutation so we can index_select instead of a dense + # [E,nred,nsph] einsum -- bit-exact, and its autograd is a cheap index_add rather than a + # matmul backprop (a large chunk of the per-step host geometric-Jacobian cost at scale). + tm = torch.as_tensor(to_m) + self._is_perm = bool(((tm == 0) | (tm.abs() == 1)).all()) and bool((tm != 0).sum(1).max() == 1) + if self._is_perm: + self._to_m_perm = tm.abs().argmax(dim=1).long() # perm[m] = source coeff index + # coefficient subselection for mmaxnmj", self.to_m, wig) + wig_M_inv = torch.einsum("njk,mk->njm", wig_inv, self.to_m) + return wig_M, wig_M_inv + + def __call__(self, pos, atomic_numbers, edge_index, sys_node_embedding, edge_vec=None, + edge_cell_shift=None): + """Returns a dict of differentiable geometric device-inputs as functions of ``pos``. + + ``edge_cell_shift`` [E, 3] is the cartesian periodic image offset per edge (from + ``radius_graph``); it is constant wrt ``pos`` so the analytic force is unaffected.""" + w, C, lmax = self.w, self.C, self.lmax + src, tgt = edge_index[0], edge_index[1] + if edge_vec is None: + edge_vec = pos[src] - pos[tgt] # fairchem edge_distance_vec convention + if edge_cell_shift is not None: + edge_vec = edge_vec + edge_cell_shift + dist = torch.linalg.norm(edge_vec, dim=1) + + wig_M, wig_M_inv = self._wigner(edge_vec) + + # x_edge = [gaussian(dist) | source_emb[Z] | target_emb[Z]] + gauss = torch.exp(self.coeff * (dist.view(-1, 1) - self.offset.view(1, -1)) ** 2) + se = F.embedding(atomic_numbers[src], w["source_embedding.weight"]) + te = F.embedding(atomic_numbers[tgt], w["target_embedding.weight"]) + x_edge = torch.cat([gauss, se, te], dim=1) + + # envelope + ds = dist / self.cutoff + env_val = 1 + (ds ** self.env_p) * (self.env_a + ds * (self.env_b + self.env_c * ds)) + envelope = torch.where(ds < 1, env_val, torch.zeros_like(env_val)).reshape(-1, 1, 1) + + # edge-degree node init + N = atomic_numbers.shape[0] + nsph = (lmax + 1) ** 2 + # pos-independent l=0 init + l0 = F.embedding(atomic_numbers, w["sphere_embedding.weight"]) + sys_node_embedding + l0 = F.pad(l0.unsqueeze(1), (0, 0, 0, nsph - 1)) # [N,9,C] with only l0 set + + from .device import device_ede + if device_ede(): + # the full node init (radial MLP -> rotate-back -> envelope -> scatter) runs on device + # inside the trace; host only supplies the constant l0 and the geometric terms below. + return dict(x_init=None, l0=l0, wigner=wig_M, wigner_inv=wig_M_inv, + x_edge=x_edge, edge_envelope=envelope, edge_distance=dist, + edge_distance_vec=edge_vec) + + m0 = self.cfg["mmax_m0_coeffs"] if "mmax_m0_coeffs" in self.cfg else (lmax + 1) + edm = radial_mlp(x_edge, w, "edge_degree_embedding.rad_func").reshape(-1, m0, C) + # the radial output is the m=0 block; place it at the front of the reduced m-space + # (nred = nsph when mmax==lmax) and rotate back to node SH via wig_M_inv [E, nsph, nred] + edm = F.pad(edm, (0, 0, 0, self.nred - m0)) # [E, nred, C] + edm = torch.bmm(wig_M_inv, edm) * envelope # [E, nsph, C] + node = torch.zeros(N, nsph, C, dtype=pos.dtype) + node = node.index_add(0, tgt, edm / self.rescale) + x_init = node + l0 + + return dict(x_init=x_init, l0=l0, wigner=wig_M, wigner_inv=wig_M_inv, + x_edge=x_edge, edge_envelope=envelope, edge_distance=dist, + edge_distance_vec=edge_vec) diff --git a/tt_atom/grid.py b/tt_atom/grid.py new file mode 100644 index 0000000..b5f9cc7 --- /dev/null +++ b/tt_atom/grid.py @@ -0,0 +1,76 @@ +"""Grid feed-forward (``GridAtomwise``) — the per-node MLP in the eSCN-MD block. + +Spherical-harmonic coefficients are projected to a real-space S2 grid with a fixed linear +map, a pointwise 3-layer MLP runs over channels at every grid point, then the result is +projected back to coefficients. The two projections are constant buffers (topology-free), +so on device this is just two transpose-matmuls around a channelwise MLP. + +Reference: ``fairchem ... escn_md_block.py:GridAtomwise`` + ``common/so3.py:SO3_Grid``:: + + x_grid = einsum("bai, zic -> zbac", to_grid_mat, x) # to_grid + x_grid = grid_mlp(x_grid) # pointwise over channels + x = einsum("bai, zbac -> zic", from_grid_mat, x_grid) # from_grid +""" +from __future__ import annotations + +from .device import compute_kernel_config + + +def _to_dev(t, device, dtype): + import ttnn + + return ttnn.from_torch(t, dtype=dtype, layout=ttnn.TILE_LAYOUT, device=device) + + +class GridAtomwise: + def __init__(self, weights, prefix, device, to_grid_mat, from_grid_mat, *, fast=False): + import ttnn + + self.ttnn = ttnn + self.device = device + self.kcfg = compute_kernel_config() + wdtype = ttnn.bfloat8_b if fast else ttnn.bfloat16 + + # to_grid_mat / from_grid_mat: [b, a, nsph]; flatten the (b, a) grid -> [npts, nsph]. + b, a, nsph = to_grid_mat.shape + self.npts = b * a + self.nsph = nsph + tg = to_grid_mat.reshape(self.npts, nsph) # x_grid = tg @ x (over nsph) + fg = from_grid_mat.reshape(self.npts, nsph) # x = fg^T @ x_grid (over npts) + # device matmuls operate on [.., C, k] @ [k, n]; store the right-multiply operands. + self.tg_T = _to_dev(tg.T.contiguous(), device, wdtype) # [nsph, npts] + self.fg = _to_dev(fg.contiguous(), device, wdtype) # [npts, nsph] + + self.w0 = _to_dev(weights[f"{prefix}.grid_mlp.0.weight"].T.contiguous(), device, wdtype) + self.w2 = _to_dev(weights[f"{prefix}.grid_mlp.2.weight"].T.contiguous(), device, wdtype) + self.w4 = _to_dev(weights[f"{prefix}.grid_mlp.4.weight"].T.contiguous(), device, wdtype) + + def __call__(self, x): + """x: ttnn ``[N, nsph, C]`` -> ``[N, nsph, C]``.""" + ttnn = self.ttnn + from .device import l1_if_fits, L1_NODE_BUDGET + # gate residency on the grid tensor [N, npts, C] fitting L1 (else DRAM at large N). Use the + # tile-padded npts (42 -> next mult of 32) since the 3D grid tensor pads the point dim. + _npts_pad = ((self.npts + 31) // 32) * 32 + L1 = l1_if_fits(ttnn, x.shape[0], _npts_pad * x.shape[2], budget=L1_NODE_BUDGET) + # The whole grid transform (to_grid -> pointwise MLP -> from_grid) is a BW-bound chain of + # matmuls/transposes over the large [N, npts, C] real-space grid tensor. Keeping those + # intermediates L1-resident cuts the chain's DRAM traffic to just the input read + final + # write -- measured ~20% faster end-to-end at N>=128 (node PCC ~1.0 vs the DRAM path; the + # only change is the matmul reduction grid, bf16-equivalent). Final output spills to DRAM + # so the block's residual add stays on the small DRAM-interleaved node tensor. + # to_grid: x_grid[z, p, c] = sum_i tg[p, i] x[z, i, c] + xt = ttnn.transpose(x, 1, 2, memory_config=L1) # [N, C, nsph] + g = ttnn.matmul(xt, self.tg_T, compute_kernel_config=self.kcfg, memory_config=L1) # [N, C, npts] + g = ttnn.transpose(g, 1, 2, memory_config=L1) # [N, npts, C] + # pointwise MLP over channels (no bias, SiLU between) + a1 = ttnn.matmul(g, self.w0, compute_kernel_config=self.kcfg, memory_config=L1) + g = ttnn.silu(a1) + a2 = ttnn.matmul(g, self.w2, compute_kernel_config=self.kcfg, memory_config=L1) + g = ttnn.silu(a2) + self._cache_a1, self._cache_a2 = a1, a2 # for the analytic-force VJP + g = ttnn.matmul(g, self.w4, compute_kernel_config=self.kcfg, memory_config=L1) # [N, npts, C] + # from_grid: x[z, i, c] = sum_p fg[p, i] x_grid[z, p, c] + gt = ttnn.transpose(g, 1, 2, memory_config=L1) # [N, C, npts] + o = ttnn.matmul(gt, self.fg, compute_kernel_config=self.kcfg) # [N, C, nsph] + return ttnn.transpose(o, 1, 2) # [N, nsph, C] diff --git a/tt_atom/model.py b/tt_atom/model.py new file mode 100644 index 0000000..e923e9c --- /dev/null +++ b/tt_atom/model.py @@ -0,0 +1,305 @@ +"""eSEN / eSCN-MD backbone forward, device-resident on Tenstorrent. + +Assembles the ported modules (RMS-norm-SH, edgewise SO(2) message passing, grid feed-forward) +into the full backbone and an energy head. The geometric, per-edge terms (Wigner matrices, +radial edge embedding, envelope, the graph itself) are computed on host once — they are <1% of +the compute — and uploaded as the device-resident ``GraphContext``. Everything else (the dense +GEMM bulk) stays on device across the whole forward. + +Reference: ``fairchem ... escn_md.py:eSCNMDBackbone.forward`` + ``escn_md_block.py``. +""" +from __future__ import annotations + +import os + +import torch + +from .device import compute_kernel_config +from .norm import RMSNormSH +from .edgewise import Edgewise +from .grid import GridAtomwise +from .spectral import SpectralAtomwise + +# Above this node count the dense one-hot scatter matmul S[N,E]@m (O(N^2)) is replaced by the +# linear O(E) gather+reduce scatter (tt_atom/scatter.py). Small systems keep the dense path +# (one fat matmul, bit-identical to the golden mirror tests). Override with $TT_ATOM_SCATTER_THRESHOLD +# (set to 0 to force the linear path everywhere — used by the scatter parity test). +# The dense one-hot matmul scatter S[N,E]@m is the golden (bit-identical) path AND measured ~5x +# faster than the linear gather+reduce at N<=~1728 (segment_sum's RM-pad + gather dominate; the +# matmul is one fat bf16 op). The O(N*E) one-hot (92 MB @N=1000, 546 MB @N=1728) fits DRAM easily +# at MD sizes, so use it up to ~2048 nodes; only truly-large scale runs fall back to the linear +# O(E) path (scatter.py) to bound the O(N^2) memory. Was 384 (linear kicked in far too early). +SCATTER_LINEAR_THRESHOLD = int(os.environ.get("TT_ATOM_SCATTER_THRESHOLD", "2048")) + + +def _to_dev(t, device, dtype, layout=None): + import ttnn + + layout = layout or ttnn.TILE_LAYOUT + return ttnn.from_torch(t, dtype=dtype, layout=layout, device=device) + + +def balance_l0(ttnn, x, mean_op, cs, ce, add_scalar, kcfg): + """Charge/spin channel balancing (fairchem ``eSCNMDBackbone.balance_channels``): shift the + l=0 scalar part of channels ``[cs:ce]`` so each system's per-channel sum equals its target + (charge). ``mean_op`` is the [N,N] per-system mean operator (block-diagonal ``1/natoms``); + ``add_scalar = target/natoms`` (uniform per system — a bundle is one composition/charge, and + a same-composition batch shares natoms). The map is a projection ``I - mean`` on those + channels, hence self-adjoint: the identical call with ``add_scalar=0`` is its own VJP.""" + N, nsph, C = x.shape + l0 = ttnn.reshape(ttnn.slice(x, [0, 0, 0], [N, 1, C]), (N, C)) + ch = ttnn.slice(l0, [0, cs], [N, ce]) # [N, nch] + ch = ttnn.subtract(ch, ttnn.matmul(mean_op, ch, compute_kernel_config=kcfg)) + if add_scalar != 0.0: + ch = ttnn.add(ch, add_scalar) + parts = [] + if cs > 0: + parts.append(ttnn.slice(l0, [0, 0], [N, cs])) + parts.append(ch) + if ce < C: + parts.append(ttnn.slice(l0, [0, ce], [N, C])) + l0n = ttnn.reshape(parts[0] if len(parts) == 1 else ttnn.concat(parts, dim=1), (N, 1, C)) + rest = ttnn.slice(x, [0, 1, 0], [N, nsph, C]) + return ttnn.concat([l0n, rest], dim=1) + + +class GraphContext: + """Host-precomputed, device-resident geometric terms for one fixed topology.""" + + def __init__(self, device, *, edge_index, wigner, wigner_inv, x_edge, edge_envelope, + num_nodes, fast=False, linear_scatter=None, system_natoms=None, + build_mean_op=False): + import ttnn + + wdtype = ttnn.bfloat16 + E = edge_index.shape[1] + self.E = E + self.N = num_nodes + # per-system mean operator for charge/spin channel balancing: M[i,j] = 1/natoms(sys i) + # iff atoms i,j share a system, else 0. One system -> (1/N) ones[N,N]; a disjoint-union + # batch -> block-diagonal. Built on host per topology (so it is captured once in a trace), + # and only when a bundle actually balances channels (uma-s-1.2); None otherwise. + self.node_meanM = None + if build_mean_op: + if system_natoms is None: + M = torch.full((num_nodes, num_nodes), 1.0 / num_nodes) + else: + M = torch.zeros(num_nodes, num_nodes) + off = 0 + for n in system_natoms: + M[off:off + n, off:off + n] = 1.0 / int(n) + off += int(n) + self.node_meanM = _to_dev(M, device, wdtype) + src = edge_index[0].to(torch.int32) + tgt = edge_index[1].to(torch.int32) + self.src_idx = _to_dev(src, device, ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT) + self.tgt_idx = _to_dev(tgt, device, ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT) + # edge->node scatter-add (``out[n] = sum_{e:tgt[e]==n} m[e]``, and the src transpose used by + # the force VJP). Small systems: dense one-hot matmul S[N,E]@m — one fat op, bit-identical to + # the golden mirror tests. Large systems: linear O(E) gather+reduce (scatter.py) — the dense + # matmul is O(N^2) compute+memory (S alone is 92 MB at N=1000) and is why large-N scaling blew + # up; fairchem/PyG scatter_add is linear. See SCATTER_LINEAR_THRESHOLD. + # ``linear_scatter`` override: a disjoint-union BATCH is block-diagonal, so the dense + # one-hot S[Ntot,Etot] is mostly off-diagonal zeros — its cost is O(Ntot*Etot) ~ O(K^2) + # in the batch size while the linear gather+reduce stays O(Etot). So batches force the + # linear path (see energy_and_forces_batch); single systems keep the node-count threshold. + self.linear_scatter = (num_nodes > SCATTER_LINEAR_THRESHOLD if linear_scatter is None + else linear_scatter) + if self.linear_scatter: + from . import scatter as _sc + + tgt_g, self.Dmax_t = _sc.build_gather(tgt, num_nodes, E) + src_g, self.Dmax_s = _sc.build_gather(src, num_nodes, E) + self.tgt_gather = _to_dev(torch.from_numpy(tgt_g), device, ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT) + self.src_gather = _to_dev(torch.from_numpy(src_g), device, ttnn.uint32, ttnn.ROW_MAJOR_LAYOUT) + else: + # scatter one-hot stays bf16: bf8_b is block-float (shared per-tile exponent), so the + # 0/1 one-hot is NOT bit-exact in bf8 (measured Fpcc 0.98, no speed gain) — keep bf16. + S = torch.zeros(num_nodes, E) + S[tgt.long(), torch.arange(E)] = 1.0 + self.scatter = _to_dev(S, device, wdtype) + Ssrc = torch.zeros(num_nodes, E) + Ssrc[src.long(), torch.arange(E)] = 1.0 + self.scatter_src = _to_dev(Ssrc, device, wdtype) + # Wigner rotation as a flat sparse multiply-accumulate (see rotation.py): pack the dense + # per-edge matrices to their structural nonzeros. bf8 coefficients run faster and stay + # PCC-safe (the rotation is an orthogonal basis change) -> use in --fast. + from . import rotation + from .device import bf8_edge + # wigner (wig_M) is [E, nred, nsph], its inverse [E, nsph, nred]. nred is the reduced + # m-space (|m|<=mmax); nred == nsph when mmax==lmax (uma-s), nred < nsph for uma-m. + self.nred, self.nsph = wigner.shape[1], wigner.shape[2] + # bf8-edge: coef stays bf16 ROW_MAJOR here (bf8 can't be RM; RM is needed for the cheap + # per-step refresh). rotation.rotate casts the on-device TILE-expanded coef to bf8 to match + # its bf8 x input. bf8 rotation coef is parity-safe (orthogonal basis change, O(1) coefs). + _b8 = bf8_edge() + wig_dtype = ttnn.bfloat8_b if fast else wdtype + self.rot_fwd_ij, cf = rotation.pack(wigner) # node SH (nsph) -> reduced m-space (nred) + self.rot_inv_ij, ci = rotation.pack(wigner_inv) # reduced m-space (nred) -> node SH (nsph) + # coef stored ROW_MAJOR: the per-step refresh's from_torch of a [E, nnz] TILE tensor pays a + # tile-pad host tilize (~1.7-3.9 ms each; nnz pads to 32) vs ~0.04 ms RM. Consumers + # (rotation._coef_exp for the fused kernel) to_layout to TILE on device. Only affects the + # pos-dependent refresh cost -- topology buffers are unchanged. + self.rot_fwd_coef = _to_dev(cf, device, wig_dtype, ttnn.ROW_MAJOR_LAYOUT) + self.rot_inv_coef = _to_dev(ci, device, wig_dtype, ttnn.ROW_MAJOR_LAYOUT) + # x_edge is stored ROW_MAJOR: the per-step trace refresh's from_torch of a wide [E,320] TILE + # tensor does a slow host tilize (~24 ms vs ~1.4 ms RM); RadialMLP to_layouts it to TILE on + # device (~0.16 ms, inside the trace) instead. Only consumer is RadialMLP (so2 rad + edge_degree). + self.x_edge = _to_dev(x_edge, device, wdtype, ttnn.ROW_MAJOR_LAYOUT) + # only the flat [E,1] envelope is consumed on device (edgewise / edge-degree broadcast); the + # 3D [E,1,1] form tile-pads to [E,32,32] (a ~64 ms/step re-tilize on the trace refresh) and + # is read by nothing, so it is not materialised. + # Store ROW_MAJOR bf16 (like x_edge / rot coefs): a TILE (esp. bf8) host from_torch of the + # [E,1] envelope on the per-step refresh does a pathological host tilize + bf8 shared-exp + # pack (~7.8 ms/step for bf8, the single largest refresh cost). RM upload is ~0.1 ms; the + # forward tilizes (and, in bf8-edge mode, casts to bf8) ON DEVICE inside the trace via + # ``materialize_envelope`` — moving the whole cost to a tiny device op on the [E,1] tensor. + self._env_dtype = ttnn.bfloat8_b if _b8 else wdtype + self.edge_envelope_rm = _to_dev(edge_envelope.reshape(E, 1), device, wdtype, + ttnn.ROW_MAJOR_LAYOUT) + # materialize once here so the eager / per-module test path (which calls edge_wise without + # going through node_embedding) has a valid buffer; the traced forward re-materializes at + # its start so the tilize op reads the per-step-refreshed RM buffer (see node_embedding). + self.materialize_envelope() + + def materialize_envelope(self): + """Tilize (and, in bf8-edge mode, cast to bf8) the RM envelope on device. Called once at + the start of the backbone forward; the resulting device tensor is reused by every edgewise + block, the edge-degree init, and the backward. Captured in the trace so the per-step + refresh only writes the cheap RM buffer.""" + import ttnn + ev = ttnn.to_layout(self.edge_envelope_rm, ttnn.TILE_LAYOUT) + if self._env_dtype == ttnn.bfloat8_b: + ev = ttnn.typecast(ev, ttnn.bfloat8_b) + self.edge_envelope_f = ev + return ev + + +class _Block: + def __init__(self, weights, prefix, device, cfg, to_grid, from_grid, fast=False): + self.norm_1 = RMSNormSH(weights, f"{prefix}.norm_1", device, + lmax=cfg["lmax"], num_channels=cfg["sphere_channels"]) + self.edge_wise = Edgewise(weights, f"{prefix}.edge_wise", device, + sphere_channels=cfg["sphere_channels"], + hidden_channels=cfg["hidden_channels"], + lmax=cfg["lmax"], mmax=cfg["mmax"], fast=fast) + self.norm_2 = RMSNormSH(weights, f"{prefix}.norm_2", device, + lmax=cfg["lmax"], num_channels=cfg["sphere_channels"]) + self.ff_type = cfg.get("ff_type", "grid") + if self.ff_type == "spectral": + self.atom_wise = SpectralAtomwise(weights, f"{prefix}.atom_wise", device, + sphere_channels=cfg["sphere_channels"], + hidden_channels=cfg["hidden_channels"], + lmax=cfg["lmax"], mmax=cfg["mmax"], fast=fast) + else: + self.atom_wise = GridAtomwise(weights, f"{prefix}.atom_wise", device, + to_grid, from_grid, fast=fast) + + def __call__(self, x, graph, sys_node_embedding): + import ttnn + + C = sys_node_embedding.shape[-1] + N = x.shape[0] + x_res = x + x = self.norm_1(x) + # add system embedding at l=0 only + l0 = ttnn.add(ttnn.slice(x, [0, 0, 0], [N, 1, C]), sys_node_embedding) + x = ttnn.concat([l0, ttnn.slice(x, [0, 1, 0], [N, x.shape[1], C])], dim=1) + x = ttnn.add(self.edge_wise(x, graph), x_res) + x_res = x + x = self.norm_2(x) + x = ttnn.add(self.atom_wise(x), x_res) + return x + + +class Backbone: + """The eSCN-MD backbone forward + energy head, fully device-resident.""" + + def __init__(self, weights, device, cfg, to_grid_mat, from_grid_mat, *, fast=False): + import ttnn + + self.ttnn = ttnn + self.device = device + self.cfg = cfg + self.C = cfg["sphere_channels"] + # charge-balanced channels (fairchem charge_balanced_channels): l=0 scalar channels [cs:ce] + # are shifted after every block so their per-system sum equals the charge. cs==ce disables + # it (uma-s-1 / random-weight bundles); uma-s-1.2 uses [0:3]. + self.cs = int(cfg.get("charge_channel_start", 0)) + self.ce = int(cfg.get("charge_channel_end", 0)) + self.kcfg = compute_kernel_config() + wdtype = ttnn.bfloat16 + self.blocks = [ + _Block(weights, f"blocks.{i}", device, cfg, to_grid_mat, from_grid_mat, fast=fast) + for i in range(cfg["num_layers"]) + ] + self.final_norm = RMSNormSH(weights, "norm", device, + lmax=cfg["lmax"], num_channels=self.C) + # optional on-device edge-degree embedding (node init) — moves the largest per-step host + # cost (radial-MLP fwd+bw over E edges) onto the device inside the trace. When enabled the + # ``x_init`` operand passed to node_embedding is instead the CONSTANT l0 init and the full + # node init is computed on device from the graph's geometric terms. See tt_atom/edge_degree. + from .device import device_ede + if device_ede(): + from .edge_degree import EdgeDegreeEmbedding + self.edge_degree = EdgeDegreeEmbedding(weights, device, cfg, + rescale=cfg.get("edge_degree_rescale", 5.0)) + else: + self.edge_degree = None + # energy head: Linear-SiLU-Linear-SiLU-Linear on the l=0 channel + self.eh_w = [_to_dev(weights[f"energy_block.{i}.weight"].T.contiguous(), device, wdtype) + for i in (0, 2, 4)] + self.eh_b = [_to_dev(weights[f"energy_block.{i}.bias"], device, wdtype) + for i in (0, 2, 4)] + + def node_embedding(self, x_init, graph, sys_node_embedding, balance_add=0.0): + """Run the backbone; returns device node embedding ``[N, nsph, C]``. + + When the device edge-degree embedding is active, ``x_init`` is the constant l0 node init + and the full node init is built on device from the graph's geometric terms. + + ``balance_add`` is the per-atom charge target ``charge/natoms`` (0 for neutral or when + balancing is disabled); when ``cs self.cs + for blk in self.blocks: + x = blk(x, graph, sys_node_embedding) + if do_bal: + x = balance_l0(self.ttnn, x, graph.node_meanM, self.cs, self.ce, balance_add, self.kcfg) + return self.final_norm(x) + + def node_energy(self, node_emb): + """Per-node energy MLP (Linear-SiLU-Linear-SiLU-Linear) on the l=0 channel -> ``[N, 1]``.""" + ttnn = self.ttnn + N = node_emb.shape[0] + h = ttnn.slice(node_emb, [0, 0, 0], [N, 1, self.C]) + h = ttnn.reshape(h, (N, self.C)) + h = ttnn.silu(ttnn.linear(h, self.eh_w[0], bias=self.eh_b[0], compute_kernel_config=self.kcfg)) + h = ttnn.silu(ttnn.linear(h, self.eh_w[1], bias=self.eh_b[1], compute_kernel_config=self.kcfg)) + # fp32 output: the per-node energy (~1-2 eV once element references are subtracted) would + # otherwise be re-quantized to bf16 (~2^-8 rel), which is the dominant device energy error + # for large-|raw| systems (MgO, radicals). Does NOT affect forces — their VJP (energy_bw) + # seeds from the head weights, not this value. + return ttnn.linear(h, self.eh_w[2], bias=self.eh_b[2], compute_kernel_config=self.kcfg, + dtype=ttnn.float32) # [N,1] fp32 + + def energy(self, node_emb): + """Total energy of a single system: sum of the per-node energy (fp32).""" + return self.ttnn.sum(self.node_energy(node_emb), dim=0) + + def energy_batch(self, node_emb, seg): + """Per-system energies of a disjoint-union batch: segment-sum of the per-node energy by + the one-hot segment matrix ``seg`` [K, N] (``seg[k, n] = 1`` iff atom n is in system k), + expressed as the tile-friendly matmul ``seg @ node_energy`` -> ``[K, 1]``. Block-diagonal + batching leaves every backbone op within-system, so this reduction is the only change.""" + ttnn = self.ttnn + ne = self.node_energy(node_emb) # [N,1] fp32 + return ttnn.matmul(ttnn.typecast(seg, ttnn.float32), ne, compute_kernel_config=self.kcfg) + + def __call__(self, x_init, graph, sys_node_embedding, balance_add=0.0): + node_emb = self.node_embedding(x_init, graph, sys_node_embedding, balance_add) + return node_emb, self.energy(node_emb) diff --git a/tt_atom/norm.py b/tt_atom/norm.py new file mode 100644 index 0000000..4989f73 --- /dev/null +++ b/tt_atom/norm.py @@ -0,0 +1,111 @@ +"""Equivariant RMS layer norm over spherical-harmonic features (``rms_norm_sh``). + +Mirrors ``fairchem ... nn/layer_norm.py:EquivariantRMSNormArraySphericalHarmonicsV2``: +center the l=0 channel, compute a degree-balanced RMS over all coefficients, scale by a +per-degree affine weight, and add an l=0 bias. All reductions/elementwise -- no matmul. +""" +from __future__ import annotations + +import os + +import torch + +# RMSNormSH runs on 3D [N, nsph, C] whose tiny coefficient dim (nsph=9) tile-pads to 32 -- a ~3.5x +# blowup on every reduction/elementwise. The whole norm is a scalar-per-node RMS + per-(coeff,chan) +# affine, so it reformulates cleanly in flat [N, nsph*C]: the degree-balanced RMS folds into ONE +# weighted sum (wvec = bdw[coeff]/C) and the affine into flat multiplies -- no pad. Gated for A/B. +_NORM_FLAT = os.environ.get("TT_ATOM_NORM_FLAT", "1") == "1" + + +def _l_of_coeff(lmax): + return [l for l in range(lmax + 1) for _ in range(2 * l + 1)] + + +class RMSNormSH: + def __init__(self, weights, prefix, device, *, lmax, num_channels, eps=1e-5): + import ttnn + + self.ttnn = ttnn + self.device = device + self.lmax = lmax + self.C = num_channels + self.eps = eps + self.nsph = (lmax + 1) ** 2 + + lc = _l_of_coeff(lmax) + # degree-balance weight per coefficient: (1/(2l+1)) / (lmax+1) + bdw = torch.tensor([1.0 / (2 * l + 1) / (lmax + 1) for l in lc]).view(1, self.nsph, 1) + self.bdw = ttnn.from_torch(bdw, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + + aw = weights[f"{prefix}.affine_weight"] # [lmax+1, C] + aw_exp = aw[torch.tensor(lc)].view(1, self.nsph, self.C) # [1, nsph, C] + self.aw = ttnn.from_torch(aw_exp, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + ab = weights[f"{prefix}.affine_bias"].view(1, 1, self.C) + self.ab = ttnn.from_torch(ab, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + + # flat-layout constants (see module docstring): wvec[j] = bdw[coeff(j)]/C, awvec = aw flat, + # abf = affine bias as [1,C]. All [1, nsph*C] (or [1,C]) for broadcast against flat [N,nsph*C]. + self.flat = _NORM_FLAT + if self.flat: + wvec = torch.tensor([1.0 / (2 * l + 1) / (lmax + 1) for l in lc]).view(self.nsph, 1) + wvec = (wvec.expand(self.nsph, self.C).reshape(1, self.nsph * self.C) / self.C) + self.wvec = ttnn.from_torch(wvec.contiguous(), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + self.awvec = ttnn.from_torch(aw_exp.reshape(1, self.nsph * self.C).contiguous(), + dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + self.abf = ttnn.from_torch(weights[f"{prefix}.affine_bias"].view(1, self.C), + dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + + def __call__(self, x): + """x: ttnn ``[N, nsph, C]`` -> ``[N, nsph, C]``.""" + ttnn = self.ttnn + N = x.shape[0] + if self.flat: + return self._call_flat(x) + from .device import l1_if_fits, L1_NODE_BUDGET + # concat relocations only -> bit-identical (node PCC 1.0); keeps the norm's [N,nsph,C] + # working set on-chip while it fits L1 (falls back to DRAM at large N). Use the tile-padded + # width (nsph -> next mult of 32) since the 3D tensor pads the coeff dim. + L1 = l1_if_fits(ttnn, N, ((self.nsph + 31) // 32) * 32 * self.C, budget=L1_NODE_BUDGET) + # center l=0 across channels + l0 = ttnn.slice(x, [0, 0, 0], [N, 1, self.C]) + l0_mean = ttnn.mean(l0, dim=2, keepdim=True) # [N,1,1] + l0c = ttnn.subtract(l0, l0_mean) + rest = ttnn.slice(x, [0, 1, 0], [N, self.nsph, self.C]) + x = ttnn.concat([l0c, rest], dim=1, memory_config=L1) + + # degree-balanced component RMS + fn = ttnn.multiply(x, x) + fn = ttnn.multiply(fn, self.bdw) + fn = ttnn.sum(fn, dim=1, keepdim=True) # [N,1,C] + fn = ttnn.mean(fn, dim=2, keepdim=True) # [N,1,1] + fn = ttnn.rsqrt(ttnn.add(fn, self.eps)) + # cache the centered input + rsqrt scale for the analytic-force VJP (rmsnorm_bw then skips + # recomputing the centering + degree-balanced RMS -- fewer backward device ops, bit-exact) + self._cache_xc, self._cache_inv = x, fn + + out = ttnn.multiply(x, ttnn.multiply(fn, self.aw)) # broadcast [N,1,1]*[1,nsph,C] + # add bias to l=0 only + l0 = ttnn.add(ttnn.slice(out, [0, 0, 0], [N, 1, self.C]), self.ab) + rest = ttnn.slice(out, [0, 1, 0], [N, self.nsph, self.C]) + return ttnn.concat([l0, rest], dim=1, memory_config=L1) + + def _call_flat(self, x): + """Flat-layout RMSNormSH ([N, nsph*C]) -- no 3D coeff tile-pad. Bit-compatible with the 3D + path. Caches the centered flat input xc and rsqrt scale inv [N,1] for the flat backward.""" + ttnn = self.ttnn + N, C, W = x.shape[0], self.C, self.nsph * self.C + xf = ttnn.reshape(x, (N, W)) + # center l=0 (first C cols) across channels + l0 = ttnn.slice(xf, [0, 0], [N, C]) + l0c = ttnn.subtract(l0, ttnn.mean(l0, dim=1, keepdim=True)) # [N,C] + xc = ttnn.concat([l0c, ttnn.slice(xf, [0, C], [N, W])], dim=1) # [N, W] + # degree-balanced RMS as one weighted sum: ms = sum_j wvec_j * xc_j^2 + ms = ttnn.sum(ttnn.multiply(ttnn.multiply(xc, xc), self.wvec), dim=1, keepdim=True) # [N,1] + inv = ttnn.rsqrt(ttnn.add(ms, self.eps)) + self._cache_xc, self._cache_inv = xc, inv + out = ttnn.multiply(xc, ttnn.multiply(inv, self.awvec)) # [N,1]*[1,W] broadcast + # bias on l=0 only + l0o = ttnn.add(ttnn.slice(out, [0, 0], [N, C]), self.abf) + out = ttnn.concat([l0o, ttnn.slice(out, [0, C], [N, W])], dim=1) + return ttnn.reshape(out, (N, self.nsph, C)) diff --git a/tt_atom/quaternion.py b/tt_atom/quaternion.py new file mode 100644 index 0000000..f1d4cbd --- /dev/null +++ b/tt_atom/quaternion.py @@ -0,0 +1,223 @@ +"""Smooth two-chart quaternion edge-frame rotation (vendored from fairchem UMA, MIT). + +Replaces the singular ZYZ-Euler edge->+Y frame (``geometry._euler_angles``) with fairchem's DEFAULT +quaternion frame (``use_quaternion_wigner=True``), which is C-infinity across the whole sphere. The +Euler frame has a coordinate singularity on the +-Y axis: at axis-aligned/exactly-symmetric +geometries its ``d(wigner)/dpos`` is non-differentiable, so the analytic force comes out wrong +(the energy, being roll-gauge invariant, is fine). The quaternion frame removes that pole, so forces +are correct everywhere. It also removes the exact-zero Wigner entries that made the value-thresholded +``rotation.pack`` change width across a trajectory (the trace-buffer shape crash). + +Pure torch, no fairchem import (same stance as ``geometry.py``). Only the l<=4 kernels are vendored +(uma-s lmax=2, uma-m lmax=4); the l>=5 Ra/Rb machinery is never reached. Verbatim from +``fairchem/core/models/uma/common/quaternion/{quaternion_utils,wigner_d_custom_kernels, +wigner_d_hybrid}.py``; the ~29 KB coefficient table lives in ``assets/wigner_d_coefficients.pt``. +""" +from __future__ import annotations + +from pathlib import Path + +import torch + +_ASSET = Path(__file__).parent / "assets" / "wigner_d_coefficients.pt" + +# Blend region for the two-chart quaternion: ey in [-0.9, 0.9] +BLEND_START = -0.9 +BLEND_WIDTH = 1.8 + + +# ------------------------------------------------------------------ quaternion helpers + + +def _smooth_step_cinf(t: torch.Tensor) -> torch.Tensor: + """C-infinity smooth step (all derivatives 0 at t=0,1). step(t)=sigmoid((2t-1)/(t(1-t))).""" + t_clamped = t.clamp(0, 1) + eps = torch.finfo(t.dtype).eps + numerator = 2.0 * t_clamped - 1.0 + denom_safe = (t_clamped * (1.0 - t_clamped)).clamp(min=eps) + result = torch.sigmoid(numerator / denom_safe) + result = torch.where(t_clamped < eps, torch.zeros_like(result), result) + result = torch.where(t_clamped > 1 - eps, torch.ones_like(result), result) + return result + + +def quaternion_multiply(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: + """Hamilton product q1*q2, (w,x,y,z) convention.""" + w1, x1, y1, z1 = q1[..., 0], q1[..., 1], q1[..., 2], q1[..., 3] + w2, x2, y2, z2 = q2[..., 0], q2[..., 1], q2[..., 2], q2[..., 3] + w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 + x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 + y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2 + z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 + return torch.stack([w, x, y, z], dim=-1) + + +def quaternion_y_rotation(gamma: torch.Tensor) -> torch.Tensor: + """Quaternion for a rotation about +Y by angle gamma, shape (N,) -> (N,4).""" + half = gamma / 2 + return torch.stack([torch.cos(half), torch.zeros_like(gamma), + torch.sin(half), torch.zeros_like(gamma)], dim=-1) + + +def quaternion_nlerp(q1: torch.Tensor, q2: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Normalized linear interpolation normalize((1-t)q1 + t q2), sign-aligned.""" + dot = (q1 * q2).sum(dim=-1, keepdim=True) + q1_aligned = torch.where(dot < 0, -q1, q1) + t_exp = t.unsqueeze(-1) if t.dim() < q1.dim() else t + return torch.nn.functional.normalize((1.0 - t_exp) * q1_aligned + t_exp * q2, dim=-1) + + +def _quaternion_chart1_standard(ex, ey, ez): + """edge->+Y directly (half-vector); singular at edge=-Y (unused there, clamp detaches grad).""" + q = torch.stack([1.0 + ey, -ez, torch.zeros_like(ex), ex], dim=-1) + eps = torch.finfo(ex.dtype).eps + return q / torch.sqrt(torch.clamp(torch.sum(q ** 2, dim=-1, keepdim=True), min=eps)) + + +def _quaternion_chart2_via_minus_y(ex, ey, ez): + """edge->+Y via -Y; singular at edge=+Y (unused there, clamp detaches grad).""" + q = torch.stack([-ez, 1.0 - ey, ex, torch.zeros_like(ex)], dim=-1) + eps = torch.finfo(ex.dtype).eps + return q / torch.sqrt(torch.clamp(torch.sum(q ** 2, dim=-1, keepdim=True), min=eps)) + + +def quaternion_edge_to_y_stable(edge_vec: torch.Tensor) -> torch.Tensor: + """Two-chart edge->+Y quaternion with C-infinity NLERP blend (chart2 near -Y, chart1 near +Y). + ``edge_vec`` assumed normalized, shape (N,3) -> (N,4).""" + ex, ey, ez = edge_vec[..., 0], edge_vec[..., 1], edge_vec[..., 2] + q1 = _quaternion_chart1_standard(ex, ey, ez) + q2 = _quaternion_chart2_via_minus_y(ex, ey, ez) + t_smooth = _smooth_step_cinf((ey - BLEND_START) / BLEND_WIDTH) + return quaternion_nlerp(q2, q1, t_smooth) + + +# ------------------------------------------------------------------ quaternion -> Wigner-D (l<=4) + + +def _generate_monomials(n_vars: int, total_degree: int): + monomials = [] + + def gen(rv, rd, cur): + if rv == 1: + monomials.append(tuple(cur + [rd])); return + for i in range(rd + 1): + gen(rv - 1, rd - i, cur + [i]) + + gen(n_vars, total_degree, []) + return monomials + + +def _precompute_powers(w, x, y, z, max_power): + def pv(var): + p = {0: torch.ones_like(var), 1: var} + for i in range(2, max_power + 1): + p[i] = p[i // 2] * p[(i + 1) // 2] + return p + + return {0: pv(w), 1: pv(x), 2: pv(y), 3: pv(z)} + + +def quaternion_to_rotation_matrix(q: torch.Tensor) -> torch.Tensor: + """l=1 Wigner-D: quaternion (N,4) -> 3x3 rotation (N,3,3).""" + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + x2, y2, z2 = x * x, y * y, z * z + xy, xz, yz = x * y, x * z, y * z + wx, wy, wz = w * x, w * y, w * z + return torch.stack([ + torch.stack([1 - 2 * (y2 + z2), 2 * (xy - wz), 2 * (xz + wy)], dim=-1), + torch.stack([2 * (xy + wz), 1 - 2 * (x2 + z2), 2 * (yz - wx)], dim=-1), + torch.stack([2 * (xz - wy), 2 * (yz + wx), 1 - 2 * (x2 + y2)], dim=-1), + ], dim=-2) + + +def quaternion_to_wigner_d_l2_einsum(q: torch.Tensor, C_l2: torch.Tensor) -> torch.Tensor: + """l=2 Wigner-D via degree-4 polynomial einsum. C_l2: (5,5,4,4,4,4) -> (N,5,5).""" + C = C_l2.to(dtype=q.dtype, device=q.device) + q2 = q.unsqueeze(-1) * q.unsqueeze(-2) # (N,4,4) + q4 = q2.unsqueeze(-1).unsqueeze(-1) * q2.unsqueeze(-3).unsqueeze(-3) # (N,4,4,4,4) + return torch.einsum("nabcd,ijabcd->nij", q4, C) + + +def quaternion_to_wigner_d_matmul(q, ell, C, monomials): + """l=3 or l=4 standalone: D = M @ C^T. Returns (N,2ell+1,2ell+1).""" + C_cast = C.to(dtype=q.dtype, device=q.device) + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + powers = _precompute_powers(w, x, y, z, 2 * ell) + M = torch.stack([powers[0][a] * powers[1][b] * powers[2][c] * powers[3][d] + for a, b, c, d in monomials], dim=1) + size = 2 * ell + 1 + return (M @ C_cast.T).view(q.shape[0], size, size) + + +def quaternion_to_wigner_d_l3l4_batched(q, C_combined, monomials_l4): + """l=3 and l=4 in one degree-8 matmul. C_combined (130,165) -> (D_l3 (N,7,7), D_l4 (N,9,9)).""" + C_cast = C_combined.to(dtype=q.dtype, device=q.device) + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + powers = _precompute_powers(w, x, y, z, 8) + M = torch.stack([powers[0][a] * powers[1][b] * powers[2][c] * powers[3][d] + for a, b, c, d in monomials_l4], dim=1) + D_flat = M @ C_cast.T + N = q.shape[0] + return D_flat[:, :49].reshape(N, 7, 7), D_flat[:, 49:].reshape(N, 9, 9) + + +class WignerKernels: + """Loads the palette-compressed l=2,3,4 coefficient tables once (from the vendored asset). + Held on the ``HostGeometry`` and reused every step; device-independent (cast per call).""" + + def __init__(self, lmax: int, asset_path: Path = _ASSET): + raw = torch.load(asset_path, map_location="cpu", weights_only=True) + + def dec(ell): + k = f"C_l{ell}" + return raw[f"{k}_palette"][raw[f"{k}_indices"].long()].reshape(tuple(raw[f"{k}_shape"].tolist())) + + # the quaternion->Wigner kernels are provided up to l=4 (uma-s lmax=2, uma-m lmax=4); a + # higher-lmax checkpoint would be silently zero-filled for l>=5, so fail loudly instead. + if lmax > 4: + raise ValueError(f"quaternion edge frame supports lmax<=4, got lmax={lmax}; " + "use the Euler path (use_quaternion=False) for higher lmax") + self.lmax = lmax + self.C_l2 = dec(2) + if lmax >= 3: + self.C_l3 = dec(3) + self.monomials_l3 = _generate_monomials(4, 6) + if lmax >= 4: + self.C_l4 = dec(4) + self.monomials_l4 = _generate_monomials(4, 8) + self.C_combined_l3l4 = self._build_combined_l3l4() + + def _build_combined_l3l4(self): + """Lift l=3 (deg-6) to deg-8 by |q|^2=1 and stack with l=4 -> (130,165).""" + idx = {m: i for i, m in enumerate(self.monomials_l4)} + lifted = torch.zeros(self.C_l3.shape[0], len(self.monomials_l4), + dtype=self.C_l3.dtype, device=self.C_l3.device) + for j, (a, b, c, d) in enumerate(self.monomials_l3): + for m8 in [(a + 2, b, c, d), (a, b + 2, c, d), (a, b, c + 2, d), (a, b, c, d + 2)]: + lifted[:, idx[m8]] += self.C_l3[:, j] + return torch.cat([lifted, self.C_l4], dim=0) + + +def wigner_from_edge(edge_vec: torch.Tensor, lmax: int, kernels: WignerKernels, + gamma: float = 0.0) -> torch.Tensor: + """``edge_vec`` [E,3] -> block-diagonal Wigner-D [E,(lmax+1)^2,(lmax+1)^2] for the edge->+Y frame, + fully differentiable in ``edge_vec``. ``gamma`` is the deterministic roll (0 by default) that + keeps forces conservative (fairchem randomizes it for training augmentation; we fix it).""" + en = torch.nn.functional.normalize(edge_vec, dim=-1) + q = quaternion_edge_to_y_stable(en) + if gamma != 0.0: + q = quaternion_multiply(quaternion_y_rotation(en.new_full((q.shape[0],), gamma)), q) + size = (lmax + 1) ** 2 + D = q.new_zeros(q.shape[0], size, size) + D[:, 0, 0] = 1.0 + if lmax >= 1: + D[:, 1:4, 1:4] = quaternion_to_rotation_matrix(q) + if lmax >= 2: + D[:, 4:9, 4:9] = quaternion_to_wigner_d_l2_einsum(q, kernels.C_l2) + if lmax >= 4: + D3, D4 = quaternion_to_wigner_d_l3l4_batched(q, kernels.C_combined_l3l4, kernels.monomials_l4) + D[:, 9:16, 9:16] = D3 + D[:, 16:25, 16:25] = D4 + elif lmax >= 3: + D[:, 9:16, 9:16] = quaternion_to_wigner_d_matmul(q, 3, kernels.C_l3, kernels.monomials_l3) + return D diff --git a/tt_atom/rotation.py b/tt_atom/rotation.py new file mode 100644 index 0000000..ae224e0 --- /dev/null +++ b/tt_atom/rotation.py @@ -0,0 +1,270 @@ +"""Per-edge Wigner rotation as a flat sparse multiply-accumulate (the SO(2) frame change). + +The edgewise message block rotates node features into the edge frame with a per-edge Wigner +matrix ``W[e]`` (``m_out[e,i,:] = sum_j W[e,i,j] m_in[e,j,:]``). Done as a batched +``[E,9,9]x[E,9,C]`` matmul this is the single most expensive op: each edge is a tiny 9x9 (tile- +padded to 32x32) matmul, so it is launch/overhead bound (~2.9us/edge, flat in E, *not* flop +bound -- LoFi == HiFi4). + +But ``W`` has a FIXED sparsity pattern -- the same ``(i,j)`` nonzeros for every edge (block +diagonal in degree ``l``, permuted by the ``to_m`` reordering folded into ``W``). So the rotation +runs as the custom fused ``ttnn.experimental.fused_rotate`` kernel: one launch reads x once, keeps +all ``nnz`` multiply-accumulates in the dest registers (fp32 accumulate), writes out once -- +~4.3x faster than the ~35-dispatch addcmul MAC it replaces (7.01 -> 1.62 ms on the uma-s +rotate shape E=46016), at PCC 0.999995. tt-atom is the +custom-kernel-only uma-s build, so this kernel is the ALWAYS-ON path; a shape the kernel cannot +run (the rectangular reduced-m uma-m rotation) raises rather than silently falling back. + +Requires the source-ttnn build that carries the op (see ../custom_kernels/README.md and the +README Install section). +""" +from __future__ import annotations + +import torch + + +def _fused_op(ttnn): + return ttnn._ttnn.operations.experimental.fused_rotate + + +def _gc_op(ttnn): + return ttnn._ttnn.operations.experimental.fused_rotate_gc + + +# The 32 column-selector tiles ([32, 32*32]; tile c has column c all-ones). Pos/topology-independent +# constant -> built once per device. matmul(prod, sel[c]) = rowsum(prod) placed in output column c. +_SEL_CACHE: dict = {} + + +def _sel(ttnn, device): + key = id(device) + t = _SEL_CACHE.get(key) + if t is None: + s = torch.zeros(32, 32 * 32) + for c in range(32): + s[:, c * 32 + c] = 1.0 + t = ttnn.from_torch(s, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + _SEL_CACHE[key] = t + return t + + +# The rotate_bw coefficient adjoint (dE/dcoef) has two on-device forms. The fast one is the custom +# fused mul-reduce kernel (``fused_rotate_gc``): products L1-resident, one accumulating matmul does +# the W-reduction + column placement, reading gout+xin once. But its product CB (32*Wt = 512KB at +# W=256) leaves little L1 headroom, so it clashes with the trace's L1-resident backward tensors +# (grid/spectral) in the small/mid-N regime (measured clash at N=216 E~10k and N=512 E~23k; clash- +# free at large N where those tensors spill to DRAM). Since the gc kernel only wins where the device +# replay dominates (large graphs), gate it to large graphs and use the ones_bd segment-sum GEMM +# below otherwise -- the GEMM is correct at EVERY size and cheap in the small/mid-N dispatch-bound +# regime, so uma-s stays correct with no L1 clash at any size. (A batched-CB gc variant that fits +# small-N exists but its reload-add serial chain regressed large-N ~35%, so the fast full-tile +# kernel + this edge gate is the shippable choice.) Fixed constant, not an env toggle. +_GC_MIN_EDGES = 45000 + + +def _gc_kernel_ok(n_out, n_in, nnz, W, E) -> bool: + """L1 budget for the gc kernel CBs (see gc_program_factory): 2*(n_out+n_in)*Wt gout/xin + + 32 sel + 32*Wt prod + 2*ceil(nnz/32) out tiles. uma-s fits (~1.16MB); uma-m (W=256, 19x25) + overflows. Also require a large graph so the CBs co-reside with the trace's L1 tensors.""" + Wt = W // 32 + out_tiles = (nnz + 31) // 32 + tiles = 2 * (n_out + n_in) * Wt + 32 + 32 * Wt + 2 * out_tiles + return tiles * _TILE_BYTES <= _L1_CB_BUDGET and E >= _GC_MIN_EDGES + + +# Two hard limits gate the fused_rotate kernel to shapes it can run: +# (1) The compute kernel fans-in all `d` per-block products into dst[0..d-1] and sums them there, +# so the max fan-in degree must fit the fp32 DST register file (dst_full_sync_en -> 8 slots). +# (2) The program factory statically allocates double-buffered CBs for the x, coef and out tiles +# on each core; their total must fit L1 (1.5 MB). uma-m (W=256, n_in=25, n_out=19) blows past +# it (~2 MB) and TT_THROWs at program build; uma-s (W=128/256, 9x9) is ~0.4-0.7 MB and fits. +# uma-s always satisfies both (per-edge shape is size-independent); uma-m never does -> it raises. +_MAX_DST = 8 +_TILE_BYTES = 2048 # bf16 32x32 tile +_L1_CB_BUDGET = 1_400_000 # < 1.5 MB L1, leaving headroom for runtime/semaphores + + +def _cb_bytes(n_in, n_out, nnz, W) -> int: + Wt = W // 32 + return 2 * (n_in * Wt + n_out * Wt + nnz) * _TILE_BYTES + + +def _kernel_ok(deg, n_in, n_out, nnz, W) -> bool: + return (len(deg) > 0 and max(deg) <= _MAX_DST + and _cb_bytes(n_in, n_out, nnz, W) <= _L1_CB_BUDGET) + + +_UNSUPPORTED = ( + "tt-atom is the custom-kernel-only uma-s build: the fused rotation kernel cannot run this " + "shape (n_in={n_in}, n_out={n_out}, W={W}, nnz={nnz}). uma-s (square 9x9) is the validated, " + "supported target; a rectangular reduced-m checkpoint (uma-m: W=256, 19x25) overflows L1 and " + "is unsupported in this build." +) + + +_GROUP_CACHE: dict = {} + + +# id(coef) -> (coef_ref, coef_exp_dev). Holds the coef reference so its python id cannot be recycled +# onto a stale entry. Cleared past a step's worth; reset_expand_cache() drops it before a trace +# capture (see below). +_EXPAND_CACHE: dict = {} + + +def reset_expand_cache(): + """Drop the expanded-coef cache. The trace engine calls this right before capturing, so the + on-device coef expansion (repeat_interleave) is recorded INTO the trace and recomputed on every + replay from the in-place-refreshed compact coef -- instead of reusing a warmup-built tensor the + replay never refreshes (which would freeze the coefficients at the capture step).""" + _EXPAND_CACHE.clear() + + +def _coef_exp(ttnn, coef, dtype=None): + """Expand compact [E, nnz] device coefficients to [E, nnz*32] ON DEVICE (each nonzero's + coef broadcast across its 32-column tile). Cached by (id(coef), dtype); holds the coef + reference so its python id cannot be recycled onto a stale entry. ``dtype`` (bf8-edge) casts + the expanded coef to match the kernel's bf8 x input -- CACHED, so within a forward+backward + pass the expansion (and any typecast) runs once, not on every rotate call.""" + # coef is stored ROW_MAJOR (cheap refresh); tilize on device here for the kernel (see + # GraphContext). Cache keyed on the ORIGINAL (persistent) coef so repeat hits across calls. + orig = coef + key = (id(orig), dtype) + b = _EXPAND_CACHE.get(key) + if b is not None and b[0] is orig: + return b[1] + tiled = ttnn.to_layout(orig, ttnn.TILE_LAYOUT) if orig.layout != ttnn.TILE_LAYOUT else orig + ce_dev = ttnn.repeat_interleave(tiled, 32, dim=1) # [E, nnz*32] + if dtype is not None and ce_dev.dtype != dtype: + ce_dev = ttnn.typecast(ce_dev, dtype) + if len(_EXPAND_CACHE) > 64: + _EXPAND_CACHE.clear() + _EXPAND_CACHE[key] = (orig, ce_dev) + return ce_dev + + +def _group(ij, nblocks, by): + """Group the (i,j) nonzeros by output block for the fused kernel. ``by='i'`` (forward: + output block = i, fan-in over j) or ``by='j'`` (backward g_in: output block = j, fan-in + over i). Returns (deg[nblocks], ks[nnz] coef-tile index, js[nnz] the *other* index).""" + key = (id(ij), nblocks, by) + g = _GROUP_CACHE.get(key) + if g is not None and g[0] is ij: # identity guard: id(ij) can be recycled onto a stale + return g[1] # entry once the original list is gc'd (wrong nnz). + deg = [0] * nblocks + ks: list = [] + other: list = [] + for b in range(nblocks): + for k, (i, j) in enumerate(ij): + blk, oth = (i, j) if by == "i" else (j, i) + if blk == b: + ks.append(k); other.append(oth); deg[b] += 1 + res = (deg, ks, other) + if len(_GROUP_CACHE) > 64: + _GROUP_CACHE.clear() + _GROUP_CACHE[key] = (ij, res) # hold the ij reference so its id can't alias a stale entry + return res + + +# Cache of block-diagonal ones matrices [nnz*W, nnz] used to turn the rotate_bw coefficient-adjoint +# reductions (nnz per-nonzero dot products) into a single dense GEMM. Keyed by (nnz, W, device id). +_ONES_BD: dict = {} + + +def _ones_bd(ttnn, device, nnz, W): + """A [nnz*W, nnz] 0/1 block-diagonal selector: column k is 1 over rows [k*W:(k+1)*W]. Left- + multiplying a [E, nnz*W] tensor of per-nonzero products by it segment-sums each W-block, i.e. + computes the nnz row-wise dot products as ONE matmul (matrix engine, fp32 accum) instead of + nnz separate reductions. 1.0 is exact in bf16 so this is a plain fp32-accumulated sum + (bit-equivalent to ttnn.sum up to reduction order).""" + key = (nnz, W, id(device)) + t = _ONES_BD.get(key) + if t is None: + blk = torch.block_diag(*[torch.ones(W, 1) for _ in range(nnz)]) # [nnz*W, nnz] + t = ttnn.from_torch(blk, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + _ONES_BD[key] = t + return t + + +def gather_coef(wigner: torch.Tensor, ii: torch.Tensor, jj: torch.Tensor): + """Gather the packed coefficients ``[E, nnz]`` at a KNOWN (already-derived) sparsity pattern + ``(ii, jj)``. For a fixed topology the pattern never changes, so re-running :func:`pack`'s + ``amax`` reduction every step (trace refresh) is wasted work -- cache ``ii, jj`` once and call + this. Bit-identical to ``pack``'s coef output.""" + return wigner[:, ii, jj].contiguous() + + +def pack(wigner: torch.Tensor, tol: float = 1e-6): + """``[E, n_out, n_in]`` Wigner -> (``ij``: list of structural-nonzero (out,in) pairs, ``coef``: + ``[E, nnz]`` the per-edge values). The pattern is taken from ``amax`` over edges, so it + includes every entry nonzero for *any* edge in this topology (entries that are ~0 for all + edges contribute nothing, so dropping them is exact for this topology). + + The matrix is rectangular for mmax tol + idx = patt.nonzero(as_tuple=False) # [nnz,2], row-major (i outer, j inner) + coef = wigner[:, idx[:, 0], idx[:, 1]].contiguous() # [E, nnz] — vectorized gather (was a + ij = [(int(i), int(j)) for i, j in idx.tolist()] # per-nonzero torch.stack, ~2x slower) + return ij, coef + + +def rotate(ttnn, x_flat, ij, coef, n_in, W, device, n_out=None): + """``x_flat`` ``[E, n_in*W]`` (W channels per coordinate) -> rotated ``[E, n_out*W]`` via the + fused kernel. ``n_out`` defaults to ``n_in`` (square rotation, uma-s). Raises for a shape the + kernel cannot run (see :data:`_UNSUPPORTED`).""" + n_out = n_in if n_out is None else n_out + deg, ks, js = _group(ij, n_out, "i") + if not _kernel_ok(deg, n_in, n_out, len(ij), W): + raise RuntimeError(_UNSUPPORTED.format(n_in=n_in, n_out=n_out, W=W, nnz=len(ij))) + # bf8-edge: the kernel requires coef in the SAME dtype as its x input; the cast is done once + # inside _coef_exp and cached (not per rotate call). + ce_dev = _coef_exp(ttnn, coef, dtype=x_flat.dtype) + return _fused_op(ttnn)(x_flat, ce_dev, n_in, n_out, W, deg, ks, js) + + +def rotate_bw(ttnn, x_in_flat, g_out_flat, ij, coef, n_in, W, device, n_out=None): + """VJP of :func:`rotate`. Returns (g wrt ``x_in`` flat ``[E,n_in*W]``, g wrt the packed + coefficients ``[E, nnz]``). The coefficient adjoint is scattered back to a dense + ``[E, n_out, n_in]`` on host (:func:`scatter_coef`) to drive the geometric ``dW/dpos`` + autograd for the force.""" + n_out = n_in if n_out is None else n_out + E = x_in_flat.shape[0] + from .device import compute_kernel_config + + deg, ks, is_ = _group(ij, n_in, "j") + if not _kernel_ok(deg, n_in, n_out, len(ij), W): + raise RuntimeError(_UNSUPPORTED.format(n_in=n_in, n_out=n_out, W=W, nnz=len(ij))) + # g_in[j] = sum_{(i,j,k)} coef_k * gout_i is the SAME fused rotation with the pattern grouped + # by input block j (fan-in over the output rows i). + ce_dev = _coef_exp(ttnn, coef, dtype=g_out_flat.dtype) # bf8-edge: cached bf8 coef + g_in_flat = _fused_op(ttnn)(g_out_flat, ce_dev, n_out, n_in, W, deg, ks, is_) + + # coefficient adjoint gc[k] = sum_W(gout_i * in_j). Large graphs: the custom fused mul-reduce + # kernel (products L1-resident, one accumulating matmul reduces + places into gc[E,nnz]). + if _gc_kernel_ok(n_out, n_in, len(ij), W, E): + is_l = [i for i, j in ij]; js_l = [j for i, j in ij] + sel = _sel(ttnn, device) + if g_out_flat.dtype != sel.dtype: # bf8-edge: gc needs gout/xin/sel same dtype + sel = ttnn.typecast(sel, g_out_flat.dtype) + gc = _gc_op(ttnn)(g_out_flat, x_in_flat, sel, n_out, n_in, W, is_l, js_l) + return g_in_flat, gc + # small/mid-N: the gc kernel's product CB would clash with the trace's L1-resident tensors, so + # segment-sum the nnz per-nonzero dot products as ONE dense GEMM (block-diagonal ones matrix, + # matrix engine, fp32 accumulate). PCC ~1.0 vs the kernel; correct + cheap at these sizes. + in_cols = ttnn.split(x_in_flat, W, dim=1) + gout_cols = ttnn.split(g_out_flat, W, dim=1) + prods = [ttnn.multiply(gout_cols[i], in_cols[j]) for (i, j) in ij] + P = ttnn.concat(prods, dim=1) # [E, nnz*W] + gc = ttnn.matmul(P, _ones_bd(ttnn, device, len(ij), W), + compute_kernel_config=compute_kernel_config()) # [E, nnz] + return g_in_flat, gc + + +def scatter_coef(g_coef: torch.Tensor, ij, n_out: int, n_in: int = None) -> torch.Tensor: + """``[E, nnz]`` coefficient adjoints -> dense ``[E, n_out, n_in]`` (zeros off-pattern).""" + n_in = n_out if n_in is None else n_in + E = g_coef.shape[0] + g = torch.zeros(E, n_out, n_in, dtype=g_coef.dtype) + ii = torch.tensor([i for i, j in ij]); jj = torch.tensor([j for i, j in ij]) + g[:, ii, jj] = g_coef # vectorized scatter (was a py loop) + return g diff --git a/tt_atom/scatter.py b/tt_atom/scatter.py new file mode 100644 index 0000000..2c6e0d5 --- /dev/null +++ b/tt_atom/scatter.py @@ -0,0 +1,53 @@ +"""Linear O(E) edge->node scatter-add (replaces the dense one-hot matmul at scale). + +The edgewise message block aggregates per-edge messages onto their target node: +``out[n] = sum_{e : tgt[e]==n} m[e]``. The original device implementation is a dense one-hot +matmul ``S[N,E] @ m`` (and its transpose ``S_src`` in the force VJP). Since ``E ~= 46*N`` for a +6A-cutoff periodic graph, that matmul is O(N*E) = O(N^2) compute AND O(N^2) memory (the [N,E] +one-hot alone is 92 MB at N=1000) — the term that makes large-N scaling blow up, while +fairchem/PyG use a linear O(E) scatter_add. + +Here the scatter-add is done in O(E): group edges by node into a fixed-max-degree gather table +``gather[N, Dmax]`` (host, once per topology — sentinel ``E`` for the padding slots), gather the +messages into ``[N, Dmax, W]`` (a row-select via ``ttnn.embedding`` against the messages padded +with a single zero row), and reduce over the degree axis. Compute + memory are O(N*Dmax*W) = +O(E*W). Every op is a standard ttnn op that composes into a captured trace; the padding zero row +is produced on device (``multiply`` by 0.0) so no host constant write enters the trace. + +Not bit-identical to the matmul (the reduction sums in a different order), but the per-node sum +of ~46 O(1) terms matches to well within the force parity tolerance (PCC ~ 1.0). +""" +from __future__ import annotations + +import numpy as np +import torch + + +def build_gather(idx: torch.Tensor, num_nodes: int, E: int): + """``idx`` [E] (int, the src or tgt node of each edge) -> (``gather_flat`` [N*Dmax] int32 with + sentinel ``E`` in the padding slots, ``Dmax``). Row ``n`` of the [N, Dmax] table lists the edge + indices whose node is ``n``; the sentinel points at the zero pad row appended to the messages.""" + idx_np = idx.detach().cpu().numpy().astype(np.int64) + deg = np.bincount(idx_np, minlength=num_nodes) + Dmax = int(deg.max()) if E > 0 else 1 + gather = np.full((num_nodes, Dmax), E, dtype=np.int64) # sentinel -> zero pad row + order = np.argsort(idx_np, kind="stable") # edges grouped by node + node_of = idx_np[order] + starts = np.zeros(num_nodes, dtype=np.int64) + starts[1:] = np.cumsum(deg)[:-1] + slot = np.arange(E) - starts[node_of] # position within the node's group + gather[node_of, slot] = order # original edge index + return gather.reshape(-1).astype(np.int32), Dmax + + +def segment_sum(ttnn, msg, gather_dev, Dmax, N, W): + """``msg`` [E, W] (tile) -> ``out`` [N, W] with ``out[n] = sum over gathered edges of msg``. + + ``gather_dev`` is the [N*Dmax] uint32 table from :func:`build_gather` (sentinel ``E``); the + messages are padded with one on-device zero row so a sentinel gathers zero.""" + E = msg.shape[0] + zrow = ttnn.multiply(ttnn.slice(msg, [0, 0], [1, W]), 0.0) # [1,W] device zeros (trace-safe) + mpad = ttnn.to_layout(ttnn.concat([msg, zrow], dim=0), ttnn.ROW_MAJOR_LAYOUT) # [E+1, W] + g = ttnn.embedding(gather_dev, mpad) # [N*Dmax, W] row-select + g = ttnn.to_layout(ttnn.reshape(g, (N, Dmax, W)), ttnn.TILE_LAYOUT) + return ttnn.sum(g, dim=1) # [N, W] diff --git a/tt_atom/so2.py b/tt_atom/so2.py new file mode 100644 index 0000000..3d7b0e0 --- /dev/null +++ b/tt_atom/so2.py @@ -0,0 +1,354 @@ +"""SO(2) convolution — the compute heart of eSEN / eSCN-MD. + +The SO(2) trick turns the SO(3) tensor product into a set of per-order (per-m) dense GEMMs, +which is exactly why this architecture maps cleanly onto Tenstorrent. We keep everything in a +flattened 2D ``[E, (lmax+1)**2 * C]`` layout where each order ``m`` occupies a contiguous, +tile-aligned column block, so the whole module is column slices + matmuls + elementwise ops +(no awkward tile-dim-1 slicing). Validated to PCC ~1.0 against the fairchem reference. + +Reference: ``fairchem/core/models/uma/nn/so2_layers.py`` (``SO2_Convolution``). The l<->m +reordering (``to_m``) is folded into the host-side Wigner matrix, so on device we just split +features by m directly. +""" +from __future__ import annotations + +import os + +import torch + +from .device import compute_kernel_config + +# The whole SO(2) convolution (m=0 dense linear + every m>0 real/imag mixing) is ONE linear map +# from the post-radial input [E, nsph*Cin] to [extra | out]. The m>0 cross terms +# out_real = real@Wa - imag@Wb , out_imag = real@Wb + imag@Wa +# fold into a single constant block-structured weight, collapsing ~27 slice/matmul/combine ops +# (per conv, per pass) into one ttnn.linear. The device is op-count/DRAM-glue bound (matmul +# compute floor is ~5 ms/step), so the ~3x extra MACs from the block-diagonal zeros are cheap +# relative to the eliminated intermediate traffic. Bit-compatible column ordering; gated so the +# per-m path stays available for A/B. +_SO2_FUSED = os.environ.get("TT_ATOM_SO2_FUSED", "1") == "1" +# Route the radial-MLP LayerNorm backward (_ln_bw) through the custom fused reduction kernel +# (ttnn.experimental.fused_ln_bw): one kernel launch computes mean/rstd + dx with W L1-resident, +# vs ~15 ttnn ops. Biggest single fuseable glue (~14 ms/step, x~10 calls). Needs source-ttnn build. +_FUSED_LNBW = os.environ.get("TT_ATOM_FUSED_LNBW") == "1" +_RED_CACHE: dict = {} + + +def _red_tile(ttnn, device, W): + """[32,32] reduction selector: column 0 = 1/W (matmul rowsum-to-col0 = row mean). Cached.""" + key = (id(device), W) + t = _RED_CACHE.get(key) + if t is None: + r = torch.zeros(32, 32); r[:, 0] = 1.0 / W + t = ttnn.from_torch(r, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device) + _RED_CACHE[key] = t + return t + + +def _to_dev(t: torch.Tensor, device, dtype): + import ttnn + + return ttnn.from_torch(t, dtype=dtype, layout=ttnn.TILE_LAYOUT, device=device) + + +class RadialMLP: + """Linear -> (LayerNorm -> SiLU) x2 -> Linear. Produces per-m radial weights from the + invariant edge embedding. Mirrors ``fairchem ... nn/radial.py:RadialMLP``.""" + + def __init__(self, weights, prefix, device, wdtype, out_scale=1.0, dup_index=None, + out_dtype=None): + import ttnn + + self.ttnn = ttnn + self.device = device + self.eps = 1e-5 + # bf8-edge: emit the per-m multiplier directly in bf8 so the so2 xf*mult multiply needs no + # boundary typecast (the trunk stays bf16 for LN/SiLU precision; only net.6's output casts). + self.out_dtype = out_dtype + # ``dup_index`` (list of original output rows) duplicates/reorders net.6's output so the MLP + # emits the SO2 per-m multiplier ``mult`` [E, nsph*Cin] directly (real/imag blocks repeated), + # folding the so2 slice+concat mult-build into the constant weight. The backward is automatic: + # matmul with the duplicated weight sums the repeated rows' gradients (== the old collapse). + self._dup_index = dup_index + # net.0 Linear, net.1 LayerNorm, net.3 Linear, net.4 LayerNorm, net.6 Linear + self.w0 = _to_dev(weights[f"{prefix}.net.0.weight"].T.contiguous(), device, wdtype) + self.b0 = _to_dev(weights[f"{prefix}.net.0.bias"], device, wdtype) + self.ln1w = _to_dev(weights[f"{prefix}.net.1.weight"], device, ttnn.bfloat16) + self.ln1b = _to_dev(weights[f"{prefix}.net.1.bias"], device, ttnn.bfloat16) + self.w3 = _to_dev(weights[f"{prefix}.net.3.weight"].T.contiguous(), device, wdtype) + self.b3 = _to_dev(weights[f"{prefix}.net.3.bias"], device, wdtype) + self.ln4w = _to_dev(weights[f"{prefix}.net.4.weight"], device, ttnn.bfloat16) + self.ln4b = _to_dev(weights[f"{prefix}.net.4.bias"], device, ttnn.bfloat16) + # ``out_scale`` folds a downstream constant (e.g. the edge-degree 1/rescale) into the final + # linear in fp32 before the bf16 cast, so it lands inside the matmul's fp32 accumulation + # instead of a lossy bf16 elementwise multiply (0.2 is not representable in bf16). w6 is the + # scale factor applied to a linear's output, so scaling both weight and bias is exact. + self.out_scale = float(out_scale) + w6 = weights[f"{prefix}.net.6.weight"] * self.out_scale # [out, hidden] + b6 = weights[f"{prefix}.net.6.bias"] * self.out_scale # [out] + if dup_index is not None: + idx = torch.as_tensor(dup_index, dtype=torch.long) + w6 = w6[idx]; b6 = b6[idx] # duplicate/reorder rows + self.w6 = _to_dev(w6.T.contiguous(), device, wdtype) + self.b6 = _to_dev(b6, device, wdtype) + # broadcast copies of the LN scales for the hand-written backward ([1, n]) + n1 = weights[f"{prefix}.net.1.weight"].shape[0] + n4 = weights[f"{prefix}.net.4.weight"].shape[0] + self.ln1w_b = _to_dev(weights[f"{prefix}.net.1.weight"].reshape(1, n1), device, ttnn.bfloat16) + self.ln4w_b = _to_dev(weights[f"{prefix}.net.4.weight"].reshape(1, n4), device, ttnn.bfloat16) + # fp32 weight copies for the (default, non-fused) analytic-force backward. The radial MLP is + # tiny (edge-channel hidden), so its VJP runs in fp32 for a few % of one block's cost; a + # fully-bf16 backward mis-directs forces on out-of-distribution geometries (compressed heavy + # cells: el_Sn_cmp 230 meV/A, PCC 0.35 vs the fp64 oracle). Mirrors the fp32-accurate VJPs the + # SH-norm (rmsnorm_bw) and gate (gate_bw) backwards already use. w6_f mirrors the *transformed* + # w6 (out_scale + dup_index), so the backward matches the forward. + f32 = ttnn.float32 + self.w0_f = _to_dev(weights[f"{prefix}.net.0.weight"].T.contiguous(), device, f32) + self.w3_f = _to_dev(weights[f"{prefix}.net.3.weight"].T.contiguous(), device, f32) + self.w6_f = _to_dev(w6.T.contiguous(), device, f32) + self.ln1w_b_f = _to_dev(weights[f"{prefix}.net.1.weight"].reshape(1, n1), device, f32) + self.ln4w_b_f = _to_dev(weights[f"{prefix}.net.4.weight"].reshape(1, n4), device, f32) + self.kcfg = compute_kernel_config() + + def __call__(self, x_edge): + ttnn = self.ttnn + # x_edge arrives ROW_MAJOR (cheap trace refresh); tilize on device here (see GraphContext.x_edge) + if x_edge.layout != ttnn.TILE_LAYOUT: + x_edge = ttnn.to_layout(x_edge, ttnn.TILE_LAYOUT) + a0 = ttnn.linear(x_edge, self.w0, bias=self.b0, compute_kernel_config=self.kcfg) + n1 = ttnn.layer_norm(a0, weight=self.ln1w, bias=self.ln1b, epsilon=self.eps) + s1 = ttnn.silu(n1) + a3 = ttnn.linear(s1, self.w3, bias=self.b3, compute_kernel_config=self.kcfg) + n4 = ttnn.layer_norm(a3, weight=self.ln4w, bias=self.ln4b, epsilon=self.eps) + s2 = ttnn.silu(n4) + # cache pre-norm / pre-silu activations for the analytic-force VJP (device radial backward) + self._cache = (a0, n1, a3, n4) + return ttnn.linear(s2, self.w6, bias=self.b6, dtype=self.out_dtype, + compute_kernel_config=self.kcfg) + + def _silu_ln_bw(self, g, n, x, w_b): + """Fused SiLU-bw + LayerNorm-bw in ONE kernel launch: computes + ``dx = ln_bw( silu'(n) * g * gamma, x )`` where ``gy = g * silu'(n) * gamma`` is built + in-kernel (folds the external ``silu_bw`` op AND the affine-scale multiply). ``n`` is the + cached pre-silu activation (LN output); ``x`` the cached LN input; ``w_b`` the LN scale [1,W].""" + import struct + ttnn = self.ttnn + W = x.shape[-1] + red = _red_tile(ttnn, self.device, W) + eps_bits = struct.unpack(" full bandwidth win) and the edge + # activations/outputs run bf8; the edt output dtype keeps the flow bf8 for the next op. + wdtype = ttnn.bfloat8_b if (fast or self.bf8_edge) else ttnn.bfloat16 + self.edt = ttnn.bfloat8_b if self.bf8_edge else None + + self.num_coef = [lmax - m + 1 for m in range(mmax + 1)] # coeffs per order m + # flattened column offsets of each m-block in the [E, (lmax+1)^2 * Cin] input + self.in_offsets = [0] + w0 = (lmax + 1) * self.Cin + self.in_offsets.append(w0) + for m in range(1, mmax + 1): + self.in_offsets.append(self.in_offsets[-1] + 2 * self.num_coef[m] * self.Cin) + + self.has_radial = f"{prefix}.rad_func.net.0.weight" in weights + self.rad_prefix = f"{prefix}.rad_func" + # radial output is split per-m into widths num_coef[m]*Cin + self.rad_sizes = [self.num_coef[m] * self.Cin for m in range(mmax + 1)] + # dup_index makes the radial MLP emit the full mult [E, nsph*Cin] directly: m=0 block once, + # each m>0 block twice (real|imag). Maps output positions -> original radial-output rows. + # Only for the fused path (its backward relies on the dup weight summing repeated rows); the + # non-fused A/B path keeps the plain [E, sum rad_sizes] output + slice/concat mult build. + dup_index = None + self._rad_dup = _SO2_FUSED and self.has_radial + if self._rad_dup: + off, dup_index = 0, list(range(self.rad_sizes[0])) + off = self.rad_sizes[0] + for m in range(1, mmax + 1): + blk = list(range(off, off + self.rad_sizes[m])) + dup_index += blk + blk # real then imag + off += self.rad_sizes[m] + self.rad = (RadialMLP(weights, self.rad_prefix, device, wdtype, dup_index=dup_index, + out_dtype=self.edt) + if self.has_radial else None) + + # m=0 dense linear (has bias) + self.w_m0 = _to_dev(weights[f"{prefix}.fc_m0.weight"].T.contiguous(), device, wdtype) + self.b_m0 = _to_dev(weights[f"{prefix}.fc_m0.bias"], device, wdtype) + # m>0 linears (no bias) + self.w_m = [ + _to_dev(weights[f"{prefix}.so2_m_conv.{m-1}.fc.weight"].T.contiguous(), device, wdtype) + for m in range(1, mmax + 1) + ] + + # fused single-matmul weight (see module docstring). Built from the SAME weights, so the + # output column ordering is bit-identical to the per-m path: [extra | m0coeffs | m1r | m1i + # | ... ]. self.fused_extra_out is the extra-gating width sliced off the front. + self.fused_w = self.fused_b = None + self.fused_extra_out = extra_m0_output_channels + if _SO2_FUSED: + self._build_fused(weights, prefix, wdtype) + + def _build_fused(self, weights, prefix, wdtype): + """Per-m fused weights: m=0 one linear [in0->640]; each m>0 ONE dense matmul on the + contiguous [real|imag] input with the [[Wa,Wb],[-Wb,Wa]] block [2K->2Hh]. No block-diagonal + zeros (unlike a single whole-conv matmul) so no MAC blowup at large E, yet ~5 ops per conv + instead of ~15 (kills the real/imag slices + subtract/add combine).""" + lmax, mmax, Cin = self.lmax, self.mmax, self.Cin + w_m0 = weights[f"{prefix}.fc_m0.weight"].T.contiguous().float() # [in0, 640] + b_m0 = weights[f"{prefix}.fc_m0.bias"].float() + self.fused_wm0 = _to_dev(w_m0.contiguous(), self.device, wdtype) + self.fused_bm0 = _to_dev(b_m0.contiguous(), self.device, wdtype) + self.fused_m0_out = w_m0.shape[1] # 640 (extra + coeffs) + self.fused_wm = [] # [2K, 2Hh] per m>0 + self.fused_out_w = [self.fused_m0_out - self.extra] # coeff out width per block + for m in range(1, mmax + 1): + wm = weights[f"{prefix}.so2_m_conv.{m-1}.fc.weight"].T.contiguous().float() # [K, 2Hh] + K, twoHh = wm.shape[0], wm.shape[1] + Hh = twoHh // 2 + Wa = wm[:, :Hh]; Wb = wm[:, Hh:2 * Hh] + Wblk = torch.zeros(2 * K, 2 * Hh, dtype=torch.float32) + Wblk[0:K, 0:Hh] = Wa; Wblk[0:K, Hh:2 * Hh] = Wb # real -> [out_real|out_imag] + Wblk[K:2 * K, 0:Hh] = -Wb; Wblk[K:2 * K, Hh:2 * Hh] = Wa # imag -> [out_real|out_imag] + self.fused_wm.append(_to_dev(Wblk.contiguous(), self.device, wdtype)) + self.fused_out_w.append(2 * Hh) + self.fused_w = True # sentinel: fused path on + + def __call__(self, x, x_edge=None): + """x: ttnn ``[E, (lmax+1)**2, Cin]`` or flat ``[E, (lmax+1)**2 * Cin]``; returns flat + ``[E, (lmax+1)**2 * H]`` (+ extra_m0 gating features ``[E, extra]`` when configured).""" + ttnn = self.ttnn + E = x.shape[0] + nsph = (self.lmax + 1) ** 2 + xf = ttnn.reshape(x, (E, nsph * self.Cin)) if len(x.shape) == 3 else x + + if self.has_radial: + rad = self.rad(x_edge) + if self._rad_dup: + mult = rad # [E, nsph*Cin] directly (dup weight) + else: + off, rms = 0, [] + for m in range(self.mmax + 1): + rms.append(ttnn.slice(rad, [0, off], [E, off + self.rad_sizes[m]])) + off += self.rad_sizes[m] + mult = ttnn.concat([rms[0]] + sum(([rms[m], rms[m]] + for m in range(1, self.mmax + 1)), []), dim=1) + if self.edt is not None and mult.dtype != self.edt: + mult = ttnn.typecast(mult, self.edt) # match bf8 xf for the multiply + self._cache_xin, self._cache_mult = xf, mult # for the analytic-force VJP + xf = ttnn.multiply(xf, mult) # bf8*bf8 -> bf8 (keeps the edge flow bf8) + + # per-m fused path: m0 one linear + one dense matmul per m>0 (see _build_fused) + if self.fused_w is not None: + x0 = ttnn.slice(xf, [0, self.in_offsets[0]], [E, self.in_offsets[1]]) + full0 = ttnn.linear(x0, self.fused_wm0, bias=self.fused_bm0, dtype=self.edt, + compute_kernel_config=self.kcfg) # [E, 640] + extra = None + if self.extra: + extra = ttnn.slice(full0, [0, 0], [E, self.extra]) + m0 = ttnn.slice(full0, [0, self.extra], [E, self.fused_m0_out]) + else: + m0 = full0 + blocks = [m0] + for m in range(1, self.mmax + 1): + blk = ttnn.slice(xf, [0, self.in_offsets[m]], [E, self.in_offsets[m + 1]]) # [E,2K] + blocks.append(ttnn.matmul(blk, self.fused_wm[m - 1], dtype=self.edt, + compute_kernel_config=self.kcfg)) + out = ttnn.concat(blocks, dim=1) # [E, nsph*H] + return (out, extra) if self.extra else out + + out_blocks = [] + + # m = 0 + x0 = ttnn.slice(xf, [0, self.in_offsets[0]], [E, self.in_offsets[1]]) + x0 = ttnn.linear(x0, self.w_m0, bias=self.b_m0, compute_kernel_config=self.kcfg) + extra = None + if self.extra: + extra = ttnn.slice(x0, [0, 0], [E, self.extra]) + x0 = ttnn.slice(x0, [0, self.extra], [E, x0.shape[1]]) + out_blocks.append(x0) # [E, H*(lmax+1)] + + # m > 0 -- two flat 2D matmuls on the real/imag halves. The earlier [E,2,nc*Cin] + # reshape made the length-2 part-dim a tile dim (padded 2->32, a 16x data blowup and a + # per-edge batched matmul); slicing the halves and running two plain GEMMs is ~80x + # faster on device for the same math (validated) and keeps the radial layout intact. + for m in range(1, self.mmax + 1): + nc = self.num_coef[m] + K = nc * self.Cin # half width (real or imag) + Hh = self.w_m[m - 1].shape[1] // 2 # out_half = H*nc + blk = ttnn.slice(xf, [0, self.in_offsets[m]], [E, self.in_offsets[m + 1]]) # [E,2K] + real = ttnn.slice(blk, [0, 0], [E, K]) + imag = ttnn.slice(blk, [0, K], [E, 2 * K]) + fr = ttnn.matmul(real, self.w_m[m - 1], compute_kernel_config=self.kcfg) # [E,2Hh] + fi = ttnn.matmul(imag, self.w_m[m - 1], compute_kernel_config=self.kcfg) + r0 = ttnn.slice(fr, [0, 0], [E, Hh]) + r1 = ttnn.slice(fr, [0, Hh], [E, 2 * Hh]) + i0 = ttnn.slice(fi, [0, 0], [E, Hh]) + i1 = ttnn.slice(fi, [0, Hh], [E, 2 * Hh]) + out_blocks.append(ttnn.subtract(r0, i1)) # real coeffs + out_blocks.append(ttnn.add(i0, r1)) # imag coeffs + + out = ttnn.concat(out_blocks, dim=1) # flat [E, (lmax+1)^2 * H], m-primed + return (out, extra) if self.extra else out diff --git a/tt_atom/spectral.py b/tt_atom/spectral.py new file mode 100644 index 0000000..2f4e29e --- /dev/null +++ b/tt_atom/spectral.py @@ -0,0 +1,188 @@ +"""Spectral feed-forward (``SpectralAtomwise``) — the per-node FF used by uma-s-1. + +Where the grid FF (``grid.py``) projects to an S2 grid and runs a channelwise MLP, the spectral +FF stays in the spherical-harmonic basis: two ``SO3_Linear`` layers (a per-degree dense GEMM that +shares one weight matrix across the ``2l+1`` coefficients of each degree ``l``, with a bias on +``l=0`` only) with a gate nonlinearity between them. The gate is driven by scalar features +produced from the ``l=0`` channel by a small MLP. + +Reference: ``fairchem ... escn_md_block.py:SpectralAtomwise`` + ``nn/so3_layers.py:SO3_Linear``:: + + gating = SiLU(scalar_mlp(x[:, 0])) # [N, lmax*hidden] + x = so3_linear_1(x) # [N, nsph, hidden] + x = GateActivation(m_prime=False)(gating, x) # SiLU on l=0, sigmoid-gate on l>=1 + x = so3_linear_2(x) # [N, nsph, sphere_channels] + +Unlike the SO(2) path the coefficients here stay in natural (l, m) order (no m-primed Wigner +reorder), so the gate's per-degree expansion is the plain ``[0]*3 + [1]*5`` map for lmax=2. +""" +from __future__ import annotations + +import os + +import torch + +from .device import compute_kernel_config + +# SO3_Linear shares one [cin,cout] weight per degree l across that degree's 2l+1 coefficients. +# The per-degree path slices x into 3D [N, 2l+1, cin] blocks and runs a batched matmul -- but the +# tiny coefficient dim (1/3/5) tile-pads to 32, a ~6-32x row blowup that makes this tiny-N module +# cost ~100 ms/step. Folding it into ONE flat 2D matmul with a block-diagonal-by-coefficient +# weight [nsph*cin, nsph*cout] kills the padding entirely (bit-compatible ordering). Gated for A/B. +_SPECTRAL_FUSED = os.environ.get("TT_ATOM_SPECTRAL_FUSED", "1") == "1" + + +def _to_dev(t, device, dtype): + import ttnn + + return ttnn.from_torch(t, dtype=dtype, layout=ttnn.TILE_LAYOUT, device=device) + + +class SpectralAtomwise: + def __init__(self, weights, prefix, device, *, sphere_channels, hidden_channels, + lmax, mmax, fast=False): + import ttnn + + self.ttnn = ttnn + self.device = device + self.C = sphere_channels + self.H = hidden_channels + self.lmax = lmax + self.mmax = mmax + self.nsph = (lmax + 1) ** 2 + self.kcfg = compute_kernel_config() + wdtype = ttnn.bfloat8_b if fast else ttnn.bfloat16 + + # scalar_mlp: Linear(C -> lmax*H) + SiLU, on the l=0 channel only + self.smlp_w = _to_dev(weights[f"{prefix}.scalar_mlp.0.weight"].T.contiguous(), device, wdtype) + self.smlp_b = _to_dev(weights[f"{prefix}.scalar_mlp.0.bias"], device, wdtype) + + # SO3_Linear weight is [lmax+1, out, in]; store per-degree [in, out] for x @ W + def so3(name, cin, cout): + W = weights[f"{prefix}.{name}.weight"] # [lmax+1, out, in] + blocks = [_to_dev(W[l].T.contiguous(), device, wdtype) for l in range(lmax + 1)] + b = _to_dev(weights[f"{prefix}.{name}.bias"].view(1, 1, cout), device, wdtype) + return blocks, b + + self.l1_w, self.l1_b = so3("so3_linear_1", self.C, self.H) # C -> H + self.l2_w, self.l2_b = so3("so3_linear_2", self.H, self.C) # H -> C + + # fused block-diagonal-by-coefficient weights (see module docstring). One flat 2D matmul. + self.l1_wf = self.l1_bf = self.l2_wf = self.l2_bf = self.gate_exp_w = None + if _SPECTRAL_FUSED: + self.l1_wf, self.l1_bf = self._build_so3_fused("so3_linear_1", weights, prefix, + self.C, self.H, wdtype) + self.l2_wf, self.l2_bf = self._build_so3_fused("so3_linear_2", weights, prefix, + self.H, self.C, wdtype) + # gate expand (natural-order): sigmoid gate row (l-1) broadcasts over degree-l's 2l+1 + # coeffs. Ex [lmax*H, (nsph-1)*H] 0/1 selector -> one matmul (fwd) + transpose (bw), + # replacing the per-degree 3D [N,2l+1,H] slices (coeff tile-pad). + H = self.H + ex = torch.zeros(lmax * H, (self.nsph - 1) * H) + c = 0 + for l in range(1, lmax + 1): + for _ in range(2 * l + 1): + ex[(l - 1) * H:l * H, c * H:(c + 1) * H] = torch.eye(H) + c += 1 + self.gate_exp_w = ttnn.from_torch(ex.contiguous(), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=device) + + def _build_so3_fused(self, name, weights, prefix, cin, cout, wdtype): + """Block-diagonal weight [nsph*cin, nsph*cout]: coeff c (degree l(c)) uses W_l. Bias on + coeff 0 (l=0) only.""" + W = weights[f"{prefix}.{name}.weight"] # [lmax+1, out, in] + bias = weights[f"{prefix}.{name}.bias"].float() # [cout] + deg = [] # degree of each coefficient + for l in range(self.lmax + 1): + deg += [l] * (2 * l + 1) + Wbd = torch.zeros(self.nsph * cin, self.nsph * cout, dtype=torch.float32) + for c, l in enumerate(deg): + Wbd[c * cin:(c + 1) * cin, c * cout:(c + 1) * cout] = W[l].T.float() # [in,out] + bbd = torch.zeros(self.nsph * cout, dtype=torch.float32) + bbd[:cout] = bias # l=0 == coeff 0 + return _to_dev(Wbd.contiguous(), self.device, wdtype), _to_dev(bbd.contiguous(), self.device, wdtype) + + def _so3_linear_fused(self, x, wf, bf, cout): + """x [N, nsph, cin] -> [N, nsph, cout] via one flat 2D matmul on the block-diagonal weight.""" + ttnn = self.ttnn + N, cin = x.shape[0], x.shape[2] + xf = ttnn.reshape(x, (N, self.nsph * cin)) + out = ttnn.linear(xf, wf, bias=bf, compute_kernel_config=self.kcfg) + return ttnn.reshape(out, (N, self.nsph, cout)) + + def _so3_linear(self, x, w_blocks, bias): + """Per-degree SO3_Linear: x [N, nsph, cin] -> [N, nsph, cout]. One shared GEMM per l, + bias added on l=0 only.""" + ttnn = self.ttnn + N, cin = x.shape[0], x.shape[2] + outs, start = [], 0 + for l in range(self.lmax + 1): + n = 2 * l + 1 + xb = ttnn.slice(x, [0, start, 0], [N, start + n, cin]) # [N, n, cin] + ob = ttnn.matmul(xb, w_blocks[l], compute_kernel_config=self.kcfg) # [N, n, cout] + if l == 0: + ob = ttnn.add(ob, bias) # bias on l=0 only + outs.append(ob) + start += n + return ttnn.concat(outs, dim=1) + + def _gate(self, gating, x): + """GateActivation (m_prime=False): SiLU on l=0, sigmoid-gate per degree on l>=1. + gating: flat [N, lmax*H]; x: [N, nsph, H].""" + ttnn = self.ttnn + N, H = x.shape[0], self.H + sg = ttnn.sigmoid(gating) # [N, lmax*H] + scalar = ttnn.silu(ttnn.slice(x, [0, 0, 0], [N, 1, H])) # l=0 + parts, start = [scalar], 1 + for l in range(1, self.lmax + 1): + n = 2 * l + 1 + xb = ttnn.slice(x, [0, start, 0], [N, start + n, H]) # [N, n, H] + gl = ttnn.slice(sg, [0, (l - 1) * H], [N, l * H]) # [N, H] for degree l + gl = ttnn.reshape(gl, (N, 1, H)) + parts.append(ttnn.multiply(xb, gl)) # broadcast over the n coeffs + start += n + return ttnn.concat(parts, dim=1) + + def __call__(self, x): + """x: ttnn ``[N, nsph, C]`` -> ``[N, nsph, C]``.""" + ttnn = self.ttnn + N = x.shape[0] + if self.gate_exp_w is not None: + return self._call_flat(x) + scalar = ttnn.reshape(ttnn.slice(x, [0, 0, 0], [N, 1, self.C]), (N, self.C)) + a_scalar = ttnn.linear(scalar, self.smlp_w, bias=self.smlp_b, + compute_kernel_config=self.kcfg) # [N, lmax*H] pre-SiLU + gating = ttnn.silu(a_scalar) + if self.l1_wf is not None: + h = self._so3_linear_fused(x, self.l1_wf, self.l1_bf, self.H) + else: + h = self._so3_linear(x, self.l1_w, self.l1_b) # [N, nsph, H] + # cached for the analytic-force VJP (gating is post-SiLU = sigmoid input) + self._cache_x, self._cache_a_scalar = x, a_scalar + self._cache_gating, self._cache_h = gating, h + g = self._gate(gating, h) # [N, nsph, H] + if self.l2_wf is not None: + return self._so3_linear_fused(g, self.l2_wf, self.l2_bf, self.C) + return self._so3_linear(g, self.l2_w, self.l2_b) # [N, nsph, C] + + def _call_flat(self, x): + """Fully-flat SpectralAtomwise ([N, nsph*C]) -- so3_linears are block-diagonal flat matmuls + and the per-degree gate expands via one 0/1 matmul (no 3D coeff tile-pad). Bit-compatible. + Caches flat h and the gating for the flat backward (spectral_bw).""" + ttnn = self.ttnn + N, C, H, nsph = x.shape[0], self.C, self.H, self.nsph + xf = ttnn.reshape(x, (N, nsph * C)) + scalar = ttnn.slice(xf, [0, 0], [N, C]) # l=0 block + a_scalar = ttnn.linear(scalar, self.smlp_w, bias=self.smlp_b, compute_kernel_config=self.kcfg) + gating = ttnn.silu(a_scalar) # [N, lmax*H] + h = ttnn.linear(xf, self.l1_wf, bias=self.l1_bf, compute_kernel_config=self.kcfg) # [N, nsph*H] + self._cache_xf, self._cache_a_scalar = xf, a_scalar + self._cache_gating, self._cache_hf = gating, h + # gate: SiLU on l=0 block, sigmoid-gate (expanded per degree) on the vector blocks + sg = ttnn.sigmoid(gating) + scalar_h = ttnn.silu(ttnn.slice(h, [0, 0], [N, H])) + gate_exp = ttnn.matmul(sg, self.gate_exp_w, compute_kernel_config=self.kcfg) # [N,(nsph-1)*H] + vec = ttnn.multiply(ttnn.slice(h, [0, H], [N, nsph * H]), gate_exp) + g = ttnn.concat([scalar_h, vec], dim=1) # [N, nsph*H] + out = ttnn.linear(g, self.l2_wf, bias=self.l2_bf, compute_kernel_config=self.kcfg) + return ttnn.reshape(out, (N, nsph, C)) diff --git a/tt_atom/trace.py b/tt_atom/trace.py new file mode 100644 index 0000000..ec8cfd8 --- /dev/null +++ b/tt_atom/trace.py @@ -0,0 +1,201 @@ +"""Device-resident, trace-captured energy+forces for a FIXED topology (MD / relaxation). + +Profiling the eager path (real uma-s-1, ethanol, p150) shows the device forward+backward is +~96% of a ``calculate()`` call and is host-*dispatch*-bound for these small graphs (hundreds of +tiny ttnn ops), not compute-bound — the host geometry is only ~3 ms. So the e2e lever for an MD +or relaxation loop, where the topology (edge set) is fixed across steps, is a ttnn *trace*: the +device forward+backward instruction stream is captured once and replayed with zero per-op host +dispatch. Each step the cheap host geometry is recomputed from the new positions and its +pos-dependent device inputs (Wigner coefficients, radial edge embedding, envelope, node init) +are refreshed *in place* in the captured buffers; then the trace is replayed. + +Measured on p150 (see benchmarks/bench_trace.py): forward-only replay is ~2.6x the eager +forward; the full forward+backward trace gives the reported e2e MD speedup. Forces are bit-for- +bit the eager analytic forces (same op stream) — the trace only removes dispatch overhead. + +The engine assumes a fixed topology (fixed ``edge_index`` and edge count). The calculator falls +back to a re-capture whenever the neighbour list changes (an atom crosses the cutoff), so the +result is always correct; only steps that keep the topology enjoy the replay speedup. +""" +from __future__ import annotations + +import torch + +from . import rotation +from .model import GraphContext + + +def _host_like(ttnn, dev_tensor, torch_tensor): + """A HOST ttnn tensor matching ``dev_tensor``'s dtype/layout, holding ``torch_tensor``'s + data — the only operand ``copy_host_to_device_tensor`` accepts to overwrite a resident (and + trace-captured) buffer in place. + + Pre-convert a float32 source to bf16 in torch first: ``ttnn.from_torch(float32, dtype=bf16)`` + does a slow scalar host conversion (~3.8 ms for x_edge alone), whereas converting in torch + (SIMD, ~0.03 ms) then from_torch on the already-bf16 tensor is a plain memcpy — ~10x faster + per-step write. Both are round-to-nearest-even, so this is bit-identical.""" + if dev_tensor.dtype == ttnn.bfloat16 and torch_tensor.dtype == torch.float32: + torch_tensor = torch_tensor.to(torch.bfloat16) + return ttnn.from_torch(torch_tensor, dtype=dev_tensor.dtype, layout=dev_tensor.layout) + + +class TracedEngine: + """Captures the device forward+backward once for a fixed topology, then replays it. + + Construct with the same operands as ``forces.energy_and_forces`` (minus ``pos``); call with + successive positions. The first call captures the trace; later calls refresh + replay.""" + + def __init__(self, bb, geo, atomic_numbers, edge_index, sys_node_embedding, + edge_cell_shift=None, seg=None, linear_scatter=None, charge=0.0, + system_natoms=None): + import ttnn + + self.ttnn = ttnn + self.bb = bb + self.geo = geo + self.Z = atomic_numbers + self.edge_index = edge_index + self.shift = edge_cell_shift + self.se = sys_node_embedding + self.dev = bb.device + self.C = geo.C + self.N = atomic_numbers.shape[0] + self.tid = None + self._device_ede = bb.edge_degree is not None + # disjoint-union batch: ``seg`` [K, N] one-hot -> per-system energy readout inside the + # trace (block-diagonal, so forces are unchanged). ``linear_scatter`` forces the O(E) + # gather+reduce (the dense S[N,E] is O(K^2) for a block-diagonal batch). None => single. + self.seg = seg + self.K = seg.shape[0] if seg is not None else 1 + self._linear_scatter = linear_scatter + # charge_balanced_channels: the per-system mean operator uses the TRUE per-system atom counts + # (``system_natoms`` = bg.natoms for a batch; None => single system => (1/N) mean). A charged + # batch is required to have equal atom counts (disjoint.assemble guards this — the scalar + # target below can't express per-system counts), so charge/natoms[0] is exact there; a neutral + # batch has add=0. 0 when balancing is disabled. + self._system_natoms = list(system_natoms) if system_natoms else None + per_sys_n = int(system_natoms[0]) if system_natoms else self.N + self.balance_add = (float(charge) / per_sys_n) if bb.ce > bb.cs else 0.0 + + # ------------------------------------------------------------------ capture / refresh + + def _refresh(self, t): + """Overwrite the pos-dependent resident buffers in place (topology buffers untouched).""" + ttnn = self.ttnn + g = self.graph + # the sparsity pattern is fixed for the topology (cached on the graph at capture); only the + # coefficient values change per step, so gather them directly instead of re-running pack's + # amax reduction over [E,nred,nsph] twice per step. + if not hasattr(self, "_fwd_ii"): + fi = torch.tensor([i for i, j in g.rot_fwd_ij]); fj = torch.tensor([j for i, j in g.rot_fwd_ij]) + ii = torch.tensor([i for i, j in g.rot_inv_ij]); ij = torch.tensor([j for i, j in g.rot_inv_ij]) + self._fwd_ii, self._fwd_jj, self._inv_ii, self._inv_jj = fi, fj, ii, ij + cf = rotation.gather_coef(t["wigner"].detach(), self._fwd_ii, self._fwd_jj) + ci = rotation.gather_coef(t["wigner_inv"].detach(), self._inv_ii, self._inv_jj) + pairs = [ + (g.rot_fwd_coef, cf), + (g.rot_inv_coef, ci), + (g.x_edge, t["x_edge"].detach()), + # only the envelope [E,1] is consumed on device; the 3D edge_envelope [E,1,1] is dead + # (its tile pads to [E,32,32] -> ~64 ms/step to re-tilize), so it is not refreshed. + # Write the ROW_MAJOR bf16 buffer (~0.1 ms); the forward tilizes/casts on device (the + # bf8 TILE host from_torch here was ~7.8 ms/step -- the largest single refresh cost). + (g.edge_envelope_rm, t["edge_envelope"].detach().reshape(g.E, 1)), + ] + # x_init operand: the host full node init (pos-dependent, refresh) or, with the device + # edge-degree embedding, the CONSTANT l0 init (pos-independent -> uploaded once, never here). + if not self._device_ede: + pairs.append((self.x_init, t["x_init"].detach())) + for dev_t, src in pairs: + ttnn.copy_host_to_device_tensor(_host_like(ttnn, dev_t, src), dev_t) + + def _capture(self, t): + ttnn = self.ttnn + N, C = self.N, self.C + self.graph = GraphContext( + self.dev, edge_index=self.edge_index, wigner=t["wigner"].detach(), + wigner_inv=t["wigner_inv"].detach(), x_edge=t["x_edge"].detach(), + edge_envelope=t["edge_envelope"].detach(), num_nodes=N, + linear_scatter=self._linear_scatter, + system_natoms=self._system_natoms, build_mean_op=(self.bb.ce > self.bb.cs)) + self.se3 = ttnn.from_torch(self.se.reshape(N, 1, C), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=self.dev) + if self.seg is not None: + self.seg_dev = ttnn.from_torch(self.seg, dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=self.dev) + # x_init operand is the constant l0 node init (device edge-degree on) or the full host + # x_init (off). l0 is pos-independent so it is uploaded once here and never refreshed. + init = t["l0"] if self._device_ede else t["x_init"] + self.x_init = ttnn.from_torch(init.detach(), dtype=ttnn.bfloat16, + layout=ttnn.TILE_LAYOUT, device=self.dev) + from . import forces as Fmod + + def body(): + if self.seg is None: + node_emb, energy = self.bb(self.x_init, self.graph, self.se3, self.balance_add) + else: + node_emb = self.bb.node_embedding(self.x_init, self.graph, self.se3, self.balance_add) + energy = self.bb.energy_batch(node_emb, self.seg_dev) # [K, 1] + acc = Fmod.backbone_bw(self.bb, self.graph, node_emb) + return energy, acc + + body() # warmup: compile all kernels before capture + ttnn.synchronize_device(self.dev) + # Drop the warmup's expanded-coef cache so the capture below records the on-device coef + # expansion INTO the trace (refreshed from the compact coef on every replay) rather than + # freezing it at the capture step's values. See rotation.reset_expand_cache. + rotation.reset_expand_cache() + self.tid = ttnn.begin_trace_capture(self.dev, cq_id=0) + self.energy_t, self.acc = body() + ttnn.end_trace_capture(self.dev, self.tid, cq_id=0) + ttnn.synchronize_device(self.dev) + + # ------------------------------------------------------------------ evaluate + + def __call__(self, pos): + """Energy + analytic forces at ``pos`` (same topology as construction).""" + ttnn = self.ttnn + pos = pos.detach().clone().requires_grad_(True) + t = self.geo(pos, self.Z, self.edge_index, self.se, edge_cell_shift=self.shift) + + if self.tid is None: + self._capture(t) # records the op stream + leaves inputs set to this ``t`` + else: + self._refresh(t) + # trace capture records without executing, so a replay is needed to populate the outputs + # on the capture step too (the inputs already hold this ``t``'s data after _capture). + ttnn.execute_trace(self.dev, self.tid, cq_id=0, blocking=True) + + E_flat = ttnn.to_torch(self.energy_t).float().reshape(-1) + E = E_flat[:self.K].clone() if self.seg is not None else float(E_flat[0]) + acc = self.acc + nsph = self.graph.nsph + g_wig = rotation.scatter_coef(ttnn.to_torch(acc["rot_fwd"]).float(), + self.graph.rot_fwd_ij, nsph) + g_winv = rotation.scatter_coef(ttnn.to_torch(acc["rot_inv"]).float(), + self.graph.rot_inv_ij, nsph) + g_env = ttnn.to_torch(acc["envelope"]).float().reshape(-1, 1, 1) + # radial finish is done on device inside the captured trace (backbone_bw); read it back. + # Only the gaussian block of x_edge = [gaussian | src_emb | tgt_emb] depends on pos, so cast + # just that block bf16->f32 (the cast dominates readback); embedding cols add zero to dpos. + # only the gaussian block (first ng cols) of x_edge depends on pos; the src/tgt embedding + # columns are pos-independent (their adjoint is discarded). Slice on device -> read back + # only [E, ng] (5x less transfer+cast than the full [E, x_edge_width]). + ng = self.geo.offset.shape[0] + W = acc["x_edge"].shape[1] + gx = ttnn.to_torch(ttnn.slice(acc["x_edge"], [0, 0], [acc["x_edge"].shape[0], ng])) + g_xe = torch.zeros((gx.shape[0], W), dtype=torch.float32) + g_xe[:, :ng] = gx.float() + outs = [t["wigner"], t["wigner_inv"], t["x_edge"], t["edge_envelope"]] + gouts = [g_wig, g_winv, g_xe, g_env] + # host x_init adjoint only when x_init is a host term (device edge-degree consumes it on device) + if not self._device_ede: + outs = [t["x_init"]] + outs + gouts = [ttnn.to_torch(acc["x_init"]).float()] + gouts + g_pos = torch.autograd.grad(outs, pos, grad_outputs=gouts)[0] + return E, -g_pos + + def close(self): + if self.tid is not None: + self.ttnn.release_trace(self.dev, self.tid) + self.tid = None diff --git a/tt_atom/weights.py b/tt_atom/weights.py new file mode 100644 index 0000000..b36863b --- /dev/null +++ b/tt_atom/weights.py @@ -0,0 +1,155 @@ +"""Weight loading: a portable ``.npz`` bundle of the eSCN-MD parameters + fixed buffers. + +TT-Atom is the *implementation*; users bring their own fairchem checkpoint. Because ttnn +(numpy<2) and fairchem (numpy>=2) cannot share a process, real weights are *exported* once in a +fairchem environment (``tools/export_weights.py``) into a numpy bundle that this loader reads in +the ttnn environment. The bundle carries both the learned ``state_dict`` and the fixed geometric +buffers (Jd, to_m, SO3 grid matrices, gaussian basis) that are not all in a bare ``state_dict``. + +The bundle format is exactly the one the parity goldens already use, so tests and the calculator +share a single code path. ``WeightBundle.verify_coverage`` checks a real checkpoint is a drop-in +fit (every key the modules need is present with the right shape).""" +from __future__ import annotations + +import json +import pathlib + +import numpy as np +import torch + + +class WeightBundle: + def __init__(self, npz): + self._d = npz + self.config = json.loads(bytes(npz["config"]).decode()) + + @classmethod + def load(cls, path): + return cls(np.load(pathlib.Path(path))) + + def _t(self, key): + return torch.from_numpy(self._d[key].copy()) + + @property + def weights(self): + return {k[2:]: self._t(k).float() for k in self._d.files if k.startswith("w@")} + + def buffer(self, name): + return self._t(f"host@{name}").float() + + def has(self, key): + return key in self._d.files + + # --------------------------------------------------------------- energy normalizer / task + + @property + def task(self): + """Dataset token for the system embedding (omol/omat/oc20/...); omat for legacy bundles.""" + return self.config.get("task", "omat") + + @property + def scale_rmsd(self): + """Energy/force scale: real targets are ``rmsd * raw + mean`` (1.0 for legacy bundles).""" + return float(self._d["scale@rmsd"][0]) if self.has("scale@rmsd") else 1.0 + + @property + def scale_mean(self): + return float(self._d["scale@mean"][0]) if self.has("scale@mean") else 0.0 + + @property + def elem_refs(self): + """Per-element reference energies added back to the (denormed) energy, or None.""" + return self._t("scale@elem_refs").double() if self.has("scale@elem_refs") else None + + @property + def reference(self): + """Embedded fairchem reference (E/F + input) for this bundle's merge composition, so a + device-side roundtrip check needs no fairchem env. ``None`` if the bundle carries none.""" + if not self.has("ref@energy"): + return None + return dict( + energy=float(self._d["ref@energy"][0]), + forces=self._d["ref@forces"].copy(), + pos=self._d["ref@pos"].copy(), + atomic_numbers=self._d["ref@atomic_numbers"].copy(), + charge=float(self._d["ref@charge"][0]), + spin=float(self._d["ref@spin"][0]), + cell=self._d["ref@cell"].copy() if self.has("ref@cell") else None, + pbc=self._d["ref@pbc"].copy() if self.has("ref@pbc") else None, + ) + + # convenience accessors for the fixed geometry buffers + @property + def to_m(self): + return self.buffer("to_m") + + @property + def coefficient_index(self): + """Spherical-harmonic coefficient subselection for mmax