Add UMA-S-1.2 support (charge-balanced channels) + edge-frame & energy-head fixes - #1

Open

janweinreich wants to merge 132 commits into

moritztng:masterfrom

janweinreich:uma-s-1.2-support

Commits on Jun 29, 2026

Scaffold: Apache-2.0 license, pyproject (tt_atom), NOTICE, package init, dirs

moritztng
and
claude
committed
042951e
Reference harness + golden fixtures (phase 2)

moritztng
and
claude
committed
e413ea0
Add device.py: TT device/mesh open + HiFi4/fp32-acc kernel config (phase 3 start)

moritztng
and
claude
committed
3fadb19
chore: gitignore ttnn-generated runtime artifacts (generated/)
moritztng
committed
50a2f4d
so2.py: device-resident SO(2) convolution + parity test (PCC 0.99999/1.0 on p150)

moritztng
and
claude
committed
6d14b07
norm.py: equivariant RMS-norm-SH + parity test (PCC 0.990 on p150)

moritztng
and
claude
committed
d40db25
golden: add SO3 grid matrices + sys_node_embedding (unblocks grid atomwise + assembly)

moritztng
and
claude
committed
4fead53
model: full eSCN-MD forward device-resident — end-to-end PCC 0.99998

moritztng
and
claude
committed
7014e93
forces: device reverse-mode VJP through backbone — adjoints PCC >= 0.9998

moritztng
and
claude
committed
8b36e82
geometry + forces: end-to-end analytic forces from positions — F PCC 0.99998

moritztng
and
claude
committed
cbc4537
weights + calculator + relaxation example (phases 5-6)

moritztng
and
claude
committed
df363a6
benchmarks: single-card throughput + CPU-vs-TT harness (phase 7 start)

moritztng
and
claude
committed
6e40801
perf: --fast bf8 path (weights + optional bf8 Wigner) — measured, accuracy-safe

moritztng
and
claude
committed
22e769e
perf: SO(2) per-m conv as flat 2D GEMMs — 12x on so2, 2.7x/block

moritztng
and
claude
committed
402cd08
perf: Wigner rotation as flat sparse multiply-accumulate (rotation.py)

moritztng
and
claude
committed
ca002d8
perf: fully-flat SO(2)/gate chain — drop all [E,9,C] tile-padded reshapes

moritztng
and
claude
committed
e12b4c4
feat: multi-card throughput fan-out (batch.py) — 3.95x on 4 cards

moritztng
and
claude
committed
c61dc4c
docs+bench: README, RELEASE, perf charts, accuracy column (phases 8-9)

moritztng
and
claude
committed
1b1a8fd

Commits on Jun 30, 2026

test: real uma-s-1 golden harness + MoLE-merge correctness anchor

moritztng
and
claude
committed
67bbc71
feat: spectral atomwise FF (ff_type=spectral) device-resident + analytic-force VJP

moritztng
and
claude
committed
2b72eb7
feat: rand_emb csd embedding + per-task energy normalizer + spectral coverage

moritztng
and
claude
committed
e189bce
test: tests/test_realweight.py — real uma-s-1 parity (auto-skips w/o bundle)

moritztng
and
claude
committed
8f0b9a2
docs+tool: real uma-s-1 accuracy numbers + --uma-s-1 export path

moritztng
and
claude
committed
de64f4b

Commits on Jul 1, 2026

feat: periodic (PBC) neighbor list — cell-aware min-image graph + parity

moritztng
and
claude
committed
66d2b41
feat: task generality — clean task_name API + oc20 slab parity (3 tasks)

moritztng
and
claude
committed
729c457
feat: honest mmax<lmax guard + uma-m-1p1 attempt documented

moritztng
and
claude
committed
493dc68
perf: trace-captured device-resident MD/relaxation path — 2.3x e2e

moritztng
and
claude
committed
d720388
feat: tt-atom CLI + one-command converter w/ device roundtrip verify + examples

moritztng
and
claude
committed
49f89b6
docs: README migration guide + honest task/periodic/trace tables

moritztng
and
claude
committed
6be3283
refactor: elegance pass — remove dead imports + unused locals

moritztng
and
claude
committed
c0ee074
docs: RELEASE.md — production numbers (periodic, 3 tasks, trace 2.33x, drop-in)

moritztng
and
claude
committed
2a93ba6

Commits on Jul 2, 2026

feat: disjoint-union (block-diagonal) graph batching — assembly + segment-sum energy + batched forces

moritztng
and
claude
committed
a692e5b
feat: evaluate_batch API + vs-fairchem batched parity (real uma-s-1)

moritztng
and
claude
committed
cd1e3be
bench: disjoint-union batching throughput (1 p150, energy-only)

moritztng
and
claude
committed
c1b7633
docs: README batched-inference section + evaluate_batch example + import cleanup

moritztng
and
claude
committed
8f83d79
Merge feat/batching: disjoint-union (block-diagonal) graph batching

moritztng
and
claude
committed
48de4b3
feat: stress/virial tensor — strain-on-edge_vec autograd + ASE Voigt + variable-cell relax

moritztng
and
claude
committed
4cd7dec
feat: odac + omc task parity — MgO framework + solid-CO2 goldens

moritztng
and
claude
committed
1028e90
feat: uma-m-1p1 support — rectangular reduced-m-space Wigner (lmax=4/mmax=2)

moritztng
and
claude
committed
0b22f6b
docs: README — stress + odac/omc + uma-m + fairchem-equivalence table & non-goals

moritztng
and
claude
committed
3dfd065
Merge feat/completeness: stress/virial + odac/omc tasks + uma-m-1p1

moritztng
and
claude
committed
6825e70

Commits on Jul 3, 2026

Commits on Jul 6, 2026

Commits on Jul 7, 2026

fix(rotate): gate fused kernel to L1-fitting shapes, MAC fallback for uma-m

moritztng
and
claude
committed
b718d10
probe(rotate): env-gated gc ablation + per-nonzero-sum A/B flags (perf profiling)

moritztng
and
claude
committed
4002201
cleanup(rotate): drop rejected gc experiments (bf8/3d/per-nonzero-sum)

moritztng
and
claude
committed
39d428d
perf(rotate): route rotate_bw coefficient adjoint through fused_rotate_gc kernel

moritztng
and
claude
committed
478c314
perf(rotate): gate gc kernel to large graphs (E>=45000), avoid small-N L1 clash

moritztng
and
claude
committed
d0b300b
perf(so2): fuse whole SO(2) conv into one block-structured matmul

moritztng
and
claude
committed
04db526
perf(spectral): fuse SO3_Linear into one block-diagonal matmul

moritztng
and
claude
committed
39ba041
perf(so2): per-m fused matmuls, drop block-diagonal MAC blowup

moritztng
and
claude
committed
4a7e9cb
perf(scatter): use one-hot matmul scatter up to N=2048 (was 384)

moritztng
and
claude
committed
100cda4
perf(norm): flat-layout RMSNormSH (kill 3D coeff tile-pad)

moritztng
and
claude
committed
88c56bd
perf(readback): slice x_edge adjoint to gaussian block before readback

moritztng
and
claude
committed
ed5da27
perf(refresh): store x_edge ROW_MAJOR, tilize on device in RadialMLP

moritztng
and
claude
committed
4c01653
perf(gate): expand gate rows via one matmul (0/1 selector) instead of slice+concat

moritztng
and
claude
committed
f4f6783
perf(spectral): fully-flat SpectralAtomwise (gate via matmul-expand)

moritztng
and
claude
committed
48c1dfa
perf(edgewise-bw): deinterleave src/tgt channels AFTER scatter (N-sized not E-sized)

moritztng
and
claude
committed
830e0df
perf(so2): bake real/imag duplication into radial net.6 weight (emit mult directly)

moritztng
and
claude
committed
21e3389
perf(refresh): store rotation coef ROW_MAJOR, tilize on device

moritztng
and
claude
committed
e5ecc0e
perf(radial-bw): route LN backward through custom fused_ln_bw kernel

moritztng
and
claude
committed
f517126
perf(radial-bw): fold SiLU-bw + LN affine scale into fused_ln_bw kernel call

moritztng
and
claude
committed
5deba37
perf(gate): route gate fwd/bw glue through custom fused_gate kernel

moritztng
and
claude
committed
a29c4cf
perf(edgewise): bf8_b edge dataflow (opt-in TT_ATOM_BF8_EDGE) — bandwidth-bound win

moritztng
and
claude
committed
b1293a7
perf(bf8-edge): emit so2 radial multiplier directly in bf8 (drop boundary typecast)

moritztng
and
claude
committed
322798c
perf(bf8-edge): cache the bf8 coef_exp typecast (was 16x/step -> 2x/step)

moritztng
and
claude
committed
309a219
perf(refresh): move envelope tilize+bf8-cast off host onto device trace

moritztng
and
claude
committed
26b7f0a

Commits on Jul 8, 2026

Commits on Jul 9, 2026

Commits on Jul 10, 2026

uma-s-1.2 support: charge_balanced_channels + fp32 energy head

janweinreich
and
claude
committed
509cc76
chore: explicitly track tt_atom/assets/*.pt (vendored Wigner-D table)

janweinreich
and
claude
committed
5f06a6f
fix: smooth quaternion edge frame + generic eval rotation

janweinreich
and
claude
committed
497141a
fix: fp32 radial-MLP backward (recovers forces on OOD heavy compressed cells)

janweinreich
and
claude
committed
d564ed3
test: gated uma-s-1.2 real-weight parity + local-checkpoint golden gen

janweinreich
and
claude
committed
b280a19
docs: uma-s-1.2 validation report + README usage note

janweinreich
and
claude
committed
deb0288