Release 0.1.4: GRACE-1L/2L with max_order 3, faster GPU kernels for GRACE-OFF, GRACE-2L RMSNorm epsilon

CHANGELOG.md

@@ -3,6 +3,72 @@
 Notable changes of grace-runtime, grace-runtime-cuda12 and grace-runtime-rocm6 (they share one version). Each
 release has a section "Changed defaults" when a default behaves differently.
 
+## [0.1.4](https://github.com/ICAMS/grace-runtime/releases/tag/0.1.4) (2026-11-05)
+
+GRACE-1L/2L models with `max_order` 3, faster GPU kernels for the GRACE-OFF models, and GRACE-2L results near the
+cutoff that match TensorFlow.
+
+### Added
+
+- **GRACE-1L and GRACE-2L models with `max_order` 3** (no AAAA / BBBB products; for example the small GRACE-OFF
+  models) run on the CPU and on NVIDIA and AMD GPUs, in Python and in LAMMPS, with the same results as TensorFlow.
+  Such models have no uncertainty head. GRACE-2L files of such models need `export_kokkos` of a tensorpotential newer
+  than 0.6.2.1; GRACE-1L files exported earlier load as they are. Older grace-runtime versions refuse such files
+  (`fc_AA2_w_left missing`).
+
+### Fixed
+
+- **GRACE-2L near the cutoff:** the RMSNorm epsilon is now 1e-10, the value of tensorpotential (was 1e-8). Energies and
+  forces change where atoms have few neighbors or bonds close to the cutoff (molecules, surfaces, dissociation), by up
+  to 2 eV for a dimer 0.05 Å inside the cutoff, and now agree with TensorFlow there. Bulk results are not visibly
+  affected.
+- CPU: after a call that stopped with the error for two atoms at the same position, the next call on the same
+  calculator returned wrong forces (the bond forces of the failed call were added; energies were right).
+
+### Changed
+
+- **Faster GPU kernels**, mainly for the GRACE-OFF models:
+  - GRACE-1L: GPU kernels of their own for FC widths 20 and 48 and for radial widths up to 32 and 48 (before:
+    zero-padded to 32 or 64).
+  - Models with lmax 3 get the wave-per-atom bond backward and, on AMD CDNA GPUs, the matrix-core basis kernels
+    (before: lmax 4 only).
+  - The specialised CG-product kernels are chosen product by product instead of for a whole model shape, so models
+    whose products match only partly (the GRACE-OFF models, GRACE-2L-OMAT, GRACE-2L-OAM) use them too.
+  - AMD GPUs: the GRACE-2L layer-2 backward of models with 48 layer-2 channels and lmax 4 fits more teams per compute
+    unit.
+
+Time per atom and step (µs; Python, fixed positions), about 5,000 atoms of water (GRACE-2L-OMAT: bcc W), before and
+after these kernels:
+
+| Model | RTX 4070 SUPER | MI250X (one GCD) |
+|---|---|---|
+| GRACE-OFF-1L-small | 0.85 → 0.54 (−36 %) | 1.35 → 0.70 (−48 %) |
+| GRACE-OFF-1L-large | 3.81 → 2.67 (−30 %) | 3.08 → 2.13 (−31 %) |
+| GRACE-OFF-2L-small | 2.04 → 1.96 (−4 %) | 2.84 → 2.28 (−20 %) |
+| GRACE-OFF-2L-medium | 3.56 → 3.45 (−3 %) | 4.50 → 3.74 (−17 %) |
+| GRACE-OFF-2L-large | 4.85 → 4.64 (−4 %) | 5.68 → 4.60 (−19 %) |
+| GRACE-2L-OMAT | 3.72 → 3.50 (−6 %) | 3.94 → 3.07 (−22 %) |
+
+Results change at the level of fp32 rounding only.
+
+### LAMMPS (`grace-runtime` branch of the LAMMPS fork)
+
+- Builds grace-runtime 0.1.4: the models, fixes and GPU kernels above.
+
+### Documentation
+
+- Models: fits with `max_order` 3; exporting a model published as `saved_model/` only; the GRACE-1L widths with GPU
+  kernels of their own.
+- Benchmarks: the time-per-atom chart names its benchmark system.
+
+### Changed defaults
+
+- NVIDIA GPUs: GRACE-1L models with lmax 3 use the wave-per-atom bond backward (`GRACE_FUSED_WAVE_BWD`, which still
+  switches it either way).
+- GPU padding of GRACE-1L fits: FC widths up to 20 pad to 20 (was 32), from 33 to 48 pad to 48 (was 64); radial
+  widths up to 32 run unpadded (was: padded to 32 unless they divide 256), from 33 to 48 pad to 48 (was 64). The
+  padding warning names the new widths.
+
 ## [0.1.3](https://github.com/ICAMS/grace-runtime/releases/tag/0.1.3) (2026-11-03)
 
 Fixes, and the LAMMPS fork on the current stable LAMMPS.

CMakeLists.txt

@@ -46,7 +46,7 @@ if(GRACE_WITH_KOKKOS)
   include(cmake/GraceHipAtomics.cmake)
   set(GRACE_KK_TESTS test_kk_1l_compile test_kk_2l_compile test_kk_3l_compile test_kk_csr test_kk_csr_3l
                      test_kk_2l_large test_kk_neighbors test_kk_coupling test_kk_large test_kk_gemm test_kk_device
-                     test_kk_launch_check test_kk_d3)
+                     test_kk_launch_check test_kk_d3 test_kk_order3 test_kk_cg_tables)
   foreach(t ${GRACE_KK_TESTS})
     add_executable(${t} tests/kk/${t}.cpp)
     target_link_libraries(${t} PRIVATE grace::libgrace Kokkos::kokkos)

docs/lammps.md

@@ -30,7 +30,7 @@ cd lammps
 ## 2. Build
 
 The pair styles are part of the `ML-PACE` package. When CMake configures the build, it downloads grace-runtime
-0.1.3 from GitHub, so the configure step needs network access. It also downloads the ML-PACE library. Pick one of the
+0.1.4 from GitHub, so the configure step needs network access. It also downloads the ML-PACE library. Pick one of the
 three recipes; run each from the `lammps` folder.
 
 ### CPU
@@ -136,10 +136,10 @@ Other AMD GPUs ([GPU backends](backends.md#amd-rocm)):
 
 ### Without network access
 
-Download grace-runtime 0.1.3 on a machine with network access, copy it over and pass its path:
+Download grace-runtime 0.1.4 on a machine with network access, copy it over and pass its path:
 
 ```bash
-git clone --depth 1 -b 0.1.3 https://github.com/ICAMS/grace-runtime.git
+git clone --depth 1 -b 0.1.4 https://github.com/ICAMS/grace-runtime.git
 cmake -S cmake -B build-cpu ... -D LOCAL_GRACE_RUNTIME=$PWD/grace-runtime
 ```
 

docs/models.md

@@ -25,13 +25,30 @@ grace_utils -p model.yaml -c CHECKPOINT export_kokkos -o model.npz
 Fits of the early GRACE-1L/2L presets export the same way. Their files are format 2; an older grace-runtime refuses
 them and asks for an update.
 
+GRACE-1L and GRACE-2L fits with `max_order` 3 (no AAAA / BBBB products; for example the small GRACE-OFF models) run as
+well as those with `max_order` 4, with the same results; the exported file has no AAAA / BBBB arrays and grace-runtime
+adds an inactive stage when it loads it. Such models have no uncertainty head. Other values of `max_order` are refused
+by `export_kokkos`. GRACE-2L fits with `max_order` 3 need a tensorpotential newer than 0.6.2.1 to export.
+
+A model published as `model.yaml` and `saved_model/` without its training checkpoints needs a checkpoint first (with a
+tensorpotential newer than 0.6.2.1):
+
+```bash
+grace_utils -p model.yaml saved_model_to_checkpoint -s saved_model -o checkpoint
+grace_utils -p model.yaml -c checkpoint export_kokkos -o model.npz
+```
+
+(`export_kokkos` refuses a checkpoint that does not hold the model's weights, such as the files in
+`saved_model/variables/`.)
+
 On a GPU, GRACE-2L fits up to the widths of the registry models run on the fast kernels, narrower layers zero-padded
 with the same results and a warning (lmax at most 4; not models with an uncertainty head): a fit of the current presets
 with at most 42 layer-1 and 32 layer-2 radial channels, 64 layer-2 product functions and `n_mlp_dens` at most 16, and
 any fit with at most 32 layer-1 and 48 layer-2 radial channels, 48 layer-2 product functions and `n_mlp_dens` at most
-12. GRACE-1L fits with up to 64 radial channels and FC outputs run zero-padded the same way. Wider GRACE-2L fits stop
-with an error that names their widths, also with `device="auto"`; with `GRACE_GENERIC_KERNELS=1` set in the environment
-they run on generic GPU kernels, several times slower, with a warning. Please open an
+12. GRACE-1L fits with up to 64 radial channels and FC outputs run on the fast kernels too: FC widths 20, 32, 48 and 64
+and radial widths up to 32, 48 and 64 on kernels of their own, other widths zero-padded the same way. Wider GRACE-2L
+fits stop with an error that names their widths, also with `device="auto"`; with `GRACE_GENERIC_KERNELS=1` set in the
+environment they run on generic GPU kernels, several times slower, with a warning. Please open an
 [issue](https://github.com/ICAMS/grace-runtime/issues) with the widths the message names so that fast kernels can be
 added.
 

libgrace/include/grace/fused_cpu_types.h

@@ -148,11 +148,13 @@ struct RowLocks {
 };
 
 // per-thread force rows [nmax][3] (no atomics: neighbours of one thread's atoms are often owned by another).
-// Kept all-zero between computes: reduce() adds them to atom->f in contiguous chunks and clears them.
+// Kept all-zero between computes: reduce() adds them to atom->f in contiguous chunks and clears them. A compute that
+// throws between ensure() and reduce() (the zero-bond error) leaves them dirty; the next ensure() clears them.
 // Each thread first touches its own rows (SpreadBuf), which it then writes all compute long.
 struct ThreadForces {
   SpreadBuf<double> buf;
   int nth = 1, nmax = 0;
+  bool dirty = false;    // set by ensure(), cleared by reduce()
   void ensure(int nthreads, int nmax_)
   {
     if (nthreads != nth || nmax_ > nmax) {
@@ -161,7 +163,9 @@ struct ThreadForces {
       buf.assign((size_t) nthreads * nmax_ * 3, 0.0, nthreads);
       nth = nthreads;
       nmax = nmax_;
-    }
+    } else if (dirty)
+      buf.assign(buf.size(), 0.0, nthreads);
+    dirty = true;
   }
   double *of(int tid) { return buf.data() + (size_t) tid * nmax * 3; }
   void reduce(double *f, int nall)    // f: atom->f[0] (contiguous [nmax][3])
@@ -181,6 +185,7 @@ struct ThreadForces {
         std::memset(src + q0, 0, (q1 - q0) * sizeof(double));
       }
     }
+    dirty = false;
   }
 };
 

libgrace/include/grace/kk/fused_dense.h

@@ -18,6 +18,10 @@
 
 #include <algorithm>
 #include <array>
+#include <cmath>
+#include <initializer_list>
+#include <string>
+#include <utility>
 #include <type_traits>
 #include <vector>
 
@@ -436,6 +440,10 @@ void fc_fwd(const FCMetaT<DeviceType> &m, int cs, const V3T &in_l, const V3T &in
 template<class DeviceType, int NO, class V3T>
 void fc_fwd4(const FCMetaT<DeviceType> &m, int cs, const V3T &in_l, const V3T &in_r, const V3T &out)
 {
+  // halves of whole float4: NO a multiple of 8, else the two-atom kernel (this one is then not compiled)
+  if constexpr (NO % 8 != 0) {
+    fc_fwd<DeviceType, NO>(m, cs, in_l, in_r, out);
+  } else {
   using F = float;
   if (fc_fwd_on_mfma<DeviceType, NO>(m, cs, in_l, in_r, out)) return;
   constexpr int NH = NO / 2, NQ = NH / 4;
@@ -483,6 +491,7 @@ void fc_fwd4(const FCMetaT<DeviceType> &m, int cs, const V3T &in_l, const V3T &i
       }
     }
   });
+  }
 }
 
 // reverse FC, right part, from the transposed weights: thread = (two atoms, source, 1/H of the inputs);
@@ -561,7 +570,9 @@ void fc_bwd(const FCMetaT<DeviceType> &m, int cs, const V3T &dout, const V3T &di
   if (fc_bwd_on_mfma<DeviceType, NO, true>(m, cs, dout, din_r)) return;
   if (m.hr == 1 && m.nirp == 32) { fc_bwd_rt<DeviceType, NO, 32, 1>(m, cs, dout, din_r); return; }
   if (m.hr == 1 && m.nirp == 44) { fc_bwd_rt<DeviceType, NO, 44, 1>(m, cs, dout, din_r); return; }
+  if (m.hr == 1 && m.nirp == 48) { fc_bwd_rt<DeviceType, NO, 48, 1>(m, cs, dout, din_r); return; }
   if (m.hr == 2 && m.nirp == 64) { fc_bwd_rt<DeviceType, NO, 64, 2>(m, cs, dout, din_r); return; }
+  if (nir > NIRMAX) GRACE_KK_ERROR("reverse FC: {} right inputs, compiled for up to {}", nir, NIRMAX);
   Kokkos::parallel_for("gfused:RevFC_R", Kokkos::RangePolicy<DeviceType, hip_lb_t<DeviceType, 256, 2>>(0, cs * nsrc), KOKKOS_LAMBDA(const int q) {
     const int sidx = q / cs, ii = q - sidx * cs;
     if (soff(sidx) == soff(sidx + 1)) return;
@@ -643,6 +654,76 @@ void prod_rev(const ProdMetaT<DeviceType> &m, int cs, const V3T &L, const V3T &R
   }
 }
 
+/* ----------------------------------------------------------------------
+   Shape-specialised CG products (cg_tables.h): one struct per compiled CG table, a list of them per engine product.
+   Each product of a model is matched on its own: cgt_match gives the index of the compiled table equal to it (indices,
+   and coefficients up to the fp32 literals), -1 for none. cg_prod_fwd / cg_prod_rev run that table's kernel, or
+   prod_fwd / prod_rev for -1, with the same results: dL assigned (zero beyond the table) or, ACC, accumulated.
+------------------------------------------------------------------------- */
+
+template<template<class...> class List, class... T, class P>
+int cgt_match(List<T...>, const P &p)
+{
+  auto same = [&](const int nt, const int (*tab)[3], const double *cg) {
+    if ((int) p.cg_coeff.size() != nt) return false;
+    for (int t = 0; t < nt; t++)
+      if (p.left_ind[t] != tab[t][0] || p.right_ind[t] != tab[t][1] || p.m_sum_ind[t] != tab[t][2] ||
+          std::fabs(p.cg_coeff[t] - cg[t]) > 1e-6 * std::max(1.0, std::fabs(cg[t]))) return false;   // fp32 literals
+    return true;
+  };
+  int i = 0, id = -1;
+  auto one = [&](const int nt, const int (*tab)[3], const double *cg) { if (id < 0 && same(nt, tab, cg)) id = i; i++; };
+  (one(T::nt, T::tab, T::cg), ...);
+  return id;
+}
+
+// log line: "specialised: AA AAA, generic: AAAA" (names[k] with ids[k] >= 0, the others)
+inline std::string cgt_summary(std::initializer_list<std::pair<const char *, int>> prods)
+{
+  std::string sp, gen;
+  for (const auto &[name, id] : prods) (id >= 0 ? sp : gen) += std::string(" ") + name;
+  return "specialised:" + (sp.empty() ? std::string(" none") : sp) + ", generic:" +
+         (gen.empty() ? std::string(" none") : gen);
+}
+
+template<class DeviceType, class T, class V3T>
+void cgt_fwd_t(const char *label, int cs, const V3T &L, const V3T &R, const V3T &out)
+{
+  const int nc = (int) out.extent(1);
+  Kokkos::parallel_for(label, Kokkos::RangePolicy<DeviceType, hip_lb_t<DeviceType, 256, 2>>(0, cs * nc), KOKKOS_LAMBDA(const int q) {
+    const int k = q / cs, ii = q - k * cs;
+    T::fwd(L, R, out, ii, k);
+  });
+}
+
+template<class DeviceType, class T, bool ACC, class V3T>
+void cgt_rev_t(const char *label, int cs, const V3T &L, const V3T &R, const V3T &D, const V3T &dL, const V3T &dR)
+{
+  const int nc = (int) D.extent(1), nx = ACC ? 0 : (int) dL.extent(2);
+  Kokkos::parallel_for(label, Kokkos::RangePolicy<DeviceType, hip_lb_t<DeviceType, 256, 2>>(0, cs * nc), KOKKOS_LAMBDA(const int q) {
+    const int k = q / cs, ii = q - k * cs;
+    T::template rev<ACC>(L, R, D, dL, dR, ii, k);
+    for (int x = T::nl; x < nx; x++) dL(ii, k, x) = 0.0f;
+  });
+}
+
+template<template<class...> class List, class... T, class DeviceType, class V3T>
+void cg_prod_fwd(List<T...>, const int id, const ProdMetaT<DeviceType> &m, const char *label, int cs, const V3T &L,
+                 const V3T &R, const V3T &out)
+{
+  int i = 0;
+  if (!((id == i++ && (cgt_fwd_t<DeviceType, T>(label, cs, L, R, out), true)) || ...)) prod_fwd(m, cs, L, R, out);
+}
+
+template<bool ACC, template<class...> class List, class... T, class DeviceType, class V3T>
+void cg_prod_rev(List<T...>, const int id, const ProdMetaT<DeviceType> &m, const char *label, int cs, const V3T &L,
+                 const V3T &R, const V3T &D, const V3T &dL, const V3T &dR)
+{
+  int i = 0;
+  if (!((id == i++ && (cgt_rev_t<DeviceType, T, ACC>(label, cs, L, R, D, dL, dR), true)) || ...))
+    prod_rev(m, cs, L, R, D, dL, dR, ACC);
+}
+
 }    // namespace GraceFusedKK
 }    // namespace kk
 }    // namespace grace