feat(lammps): support device-side model deviation with Kokkos

doc/model/dpa4c.md

@@ -276,6 +276,59 @@ Run under Kokkos with one GPU:
 lmp -k on g 1 -sf kk -in in.lammps
 ```
 
+### Model-deviation sampling with Kokkos
+
+`deepmd/kk` supports model-deviation sampling with one MPI rank. The first
+model supplies the energy, force, and virial that drive dynamics on every
+step. Each model is loaded once. Additional models are evaluated only at the
+existing `out_freq` interval, including step zero; `out_freq 0` disables
+sampling. The first model's already computed outputs are reused, so a
+two-model sample takes two model evaluations in total.
+
+```lammps
+atom_modify map yes
+pair_style deepmd/kk m0.pt2 m1.pt2 out_freq 1000 out_file model_devi.out
+pair_coeff * * H C N O P Cl
+```
+
+Run with `-k on g 1 -sf kk -pk kokkos neigh full newton off`. Multi-model
+sampling requires float32 compact canonical `.pt2` artifacts with identical
+cutoffs and complete type maps. Models must have zero frame and atomic
+parameter dimensions and no spin types. A compressed archive's fixed default
+charge/spin state is supported: its API may report a nonzero charge/spin
+width because an optional fold can reset that state, even though the compiled
+canonical forward has no runtime charge/spin input. Reference models must
+have matching charge/spin widths and identical frozen default states, including
+older compressed archives without a charge-state fold. Every model must export
+atomic virials, even when the `atomic` output option is disabled, because the
+device path uses them to compute the total virial. Missing or incompatible
+default states and missing atomic-virial outputs are rejected during model
+initialization.
+Explicit `charge_spin` settings and runtime parameters from `compute`, `fix`,
+or `ttm` are unsupported in this mode. Incompatible models are rejected;
+there is no fallback to host sampling. The single-rank periodic
+box restriction still applies: its thickness must exceed twice the cutoff in
+every periodic direction. Multi-model runs with more than one MPI rank are
+rejected; existing single-model edge/graph and multi-rank execution is
+unchanged.
+
+All models consume the same canonical graph on the GPU. Kokkos accumulates
+force and virial deviations on the GPU with population statistics; no
+coordinates, neighbor lists, raw forces, or raw virials are copied to the
+host for sampling. Without `atomic`, only six final statistics and a status
+flag are copied back, and the file contains seven columns including the
+timestep. `atomic` optionally adds the final per-atom force deviations.
+`relative`, `relative_v`, units, and restart output retain ordinary `deepmd`
+semantics. The mean force deviation is the arithmetic mean over atoms; the
+reported average virial deviation is the RMS over nine tensor components.
+
+Sampling runs after the first model's driving outputs have been consumed;
+reference outputs never replace its dynamics. The existing graph builder's
+node-map refresh and LAMMPS control/thermodynamic output remain unchanged.
+The device-resident sampling claim concerns the additional model evaluations
+and statistics, not every operation performed by LAMMPS. Include the extra
+reference inference and device statistics when comparing performance.
+
 ### CPU hosts
 
 DPA4C has a second set of hand-written operators for the CPU, so a compressed

source/api_c/include/c_api.h

@@ -13,7 +13,7 @@ extern "C" {
 /** C API version. Bumped whenever the API is changed.
  * @since API version 22
  */
-#define DP_C_API_VERSION 30
+#define DP_C_API_VERSION 31
 
 /**
  * @brief Neighbor list.
@@ -2614,6 +2614,34 @@ int DP_DeepPotGetDimAParam(DP_DeepPot* dp);
  */
 int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp);
 
+/**
+ * @brief Copy the current default charge/spin state, including frozen states.
+ * Supported by the PyTorch Exportable backend. Zero values identify an
+ * unconditioned model; an unknown state is an error.
+ * Clears any previous error before querying; DP_DeepPotCheckOK reports only
+ * the result of this query.
+ * @param[in] dp An initialized DP.
+ * @param[out] values Caller-owned buffer, or NULL with capacity zero to query
+ * the required length. No values are written on error.
+ * @param[in] capacity Buffer capacity in doubles; must fit the entire state.
+ * @return Required length, or -1 on error. Inspect DP_DeepPotCheckOK for
+ * errors. The length is independent of DP_DeepPotGetDimChgSpin.
+ * @since API version 31
+ */
+int DP_DeepPotGetDefaultChgSpin(DP_DeepPot* dp, double* values, int capacity);
+
+/**
+ * @brief Whether the model provides atomic virials.
+ * Supported by the PyTorch Exportable backend.
+ * Clears any previous error before querying; DP_DeepPotCheckOK reports only
+ * the result of this query, including successful queries returning false.
+ * @param[in] dp An initialized DP.
+ * @return True if available, false if unavailable or on error. Inspect
+ * DP_DeepPotCheckOK to distinguish an unsupported query from unavailability.
+ * @since API version 31
+ */
+bool DP_DeepPotHasAtomicVirial(DP_DeepPot* dp);
+
 /**
  * @brief Fix the charge/spin condition served for the rest of the run.
  * It becomes the condition of every later evaluation that is not given one

source/api_c/include/deepmd.hpp

@@ -1287,6 +1287,33 @@ class DeepPot : public DeepBaseModel {
     return dchgspin;
   }
 
+  /**
+   * @brief Get the current default charge/spin state, including frozen states.
+   * An empty vector identifies an unconditioned model. Throws for an unknown
+   * state or unsupported backend. The length need not equal dim_chg_spin().
+   */
+  std::vector<double> get_default_chg_spin() const {
+    assert(dp);
+    const int size = DP_DeepPotGetDefaultChgSpin(dp, nullptr, 0);
+    DP_CHECK_OK(DP_DeepPotCheckOK, dp);
+    std::vector<double> state(size);
+    if (size > 0) {
+      DP_DeepPotGetDefaultChgSpin(dp, state.data(), size);
+      DP_CHECK_OK(DP_DeepPotCheckOK, dp);
+    }
+    return state;
+  }
+
+  /**
+   * @brief Whether the model provides atomic virials; throws if unsupported.
+   */
+  bool has_atomic_virial() const {
+    assert(dp);
+    const bool result = DP_DeepPotHasAtomicVirial(dp);
+    DP_CHECK_OK(DP_DeepPotCheckOK, dp);
+    return result;
+  }
+
   /**
    * @brief Fix the charge/spin condition served for the rest of the run.
    * It becomes the condition of every later evaluation that is not given one

source/api_c/src/c_api.cc

@@ -2790,6 +2790,44 @@ int DP_DeepPotGetDimAParam(DP_DeepPot* dp) {
 
 int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp) { return dp->dp.dim_chg_spin(); }
 
+int DP_DeepPotGetDefaultChgSpin(DP_DeepPot* dp, double* values, int capacity) {
+  try {
+    dp->exception.clear();
+    if (capacity < 0 || (values == nullptr && capacity != 0)) {
+      throw deepmd::deepmd_exception(
+          "default charge/spin output requires a nonnegative capacity and "
+          "a buffer, or NULL with capacity zero for a length query");
+    }
+    const auto state = dp->dp.get_default_chg_spin();
+    if (state.size() >
+        static_cast<std::size_t>(std::numeric_limits<int>::max())) {
+      throw deepmd::deepmd_exception("default charge/spin state is too large");
+    }
+    const int size = static_cast<int>(state.size());
+    if (values != nullptr) {
+      if (capacity < size) {
+        throw deepmd::deepmd_exception(
+            "default charge/spin output buffer is too small");
+      }
+      std::copy(state.begin(), state.end(), values);
+    }
+    return size;
+  } catch (const deepmd::deepmd_exception& ex) {
+    dp->exception = ex.what();
+    return -1;
+  }
+}
+
+bool DP_DeepPotHasAtomicVirial(DP_DeepPot* dp) {
+  try {
+    dp->exception.clear();
+    return dp->dp.has_atomic_virial();
+  } catch (const deepmd::deepmd_exception& ex) {
+    dp->exception = ex.what();
+    return false;
+  }
+}
+
 void DP_DeepPotSetChargeSpin(DP_DeepPot* dp,
                              const double* charge_spin,
                              const int numb_chg_spin) {

source/api_cc/include/DeepPot.h

@@ -208,6 +208,19 @@ class DeepPotBackend : public DeepBaseModelBackend {
    **/
   virtual int dim_chg_spin() const { return 0; }
 
+  /**
+   * @brief Get the current default charge/spin state, including frozen states.
+   * An empty vector identifies an unconditioned model. Throws if the backend
+   * cannot establish the state. Its length need not equal dim_chg_spin().
+   */
+  virtual std::vector<double> get_default_chg_spin() const;
+
+  /**
+   * @brief Whether the loaded model provides atomic virials.
+   * Throws if the backend cannot query this capability.
+   */
+  virtual bool has_atomic_virial() const;
+
   /**
    * @brief Reject a charge/spin condition no embedding table can address.
    *
@@ -909,6 +922,12 @@ class DeepPot : public DeepBaseModel {
 
   int dim_chg_spin() const;
 
+  /** @copydoc DeepPotBackend::get_default_chg_spin */
+  std::vector<double> get_default_chg_spin() const;
+
+  /** @copydoc DeepPotBackend::has_atomic_virial */
+  bool has_atomic_virial() const;
+
   /**
    * @brief Fix the charge/spin condition served for the rest of the run.
    * @param[in] charge_spin The condition, of length ``dim_chg_spin()``.

source/api_cc/include/DeepPotPTExpt.h

@@ -413,6 +413,8 @@ class DeepPotPTExpt : public DeepPotBackend {
    */
   bool uses_fp32_edge_vectors() const override;
   bool uses_canonical_graph_inference() const override;
+  std::vector<double> get_default_chg_spin() const override;
+  bool has_atomic_virial() const override;
 
  private:
   template <typename EDGE_TYPE>
@@ -460,6 +462,8 @@ class DeepPotPTExpt : public DeepPotBackend {
   bool has_default_fparam_;
   std::vector<double> default_fparam_;
   std::vector<double> default_chg_spin_;
+  // Defer unavailable metadata to the query, preserving single-model loading.
+  std::string default_chg_spin_query_error_;
   /** Half-open row range of each charge-state value, from the archive. */
   std::vector<std::pair<double, double> > chg_spin_table_ranges_;
   double rcut;

source/api_cc/src/DeepPot.cc

@@ -692,6 +692,16 @@ bool DeepPotBackend::supports_device_edge_inference() const { return false; }
 
 bool DeepPotBackend::uses_canonical_graph_inference() const { return false; }
 
+std::vector<double> DeepPotBackend::get_default_chg_spin() const {
+  throw deepmd::deepmd_exception(
+      "this backend cannot query the model's default charge/spin state");
+}
+
+bool DeepPotBackend::has_atomic_virial() const {
+  throw deepmd::deepmd_exception(
+      "this backend cannot query atomic-virial availability");
+}
+
 void DeepPot::compute_edges_gpu(double* d_atom_energy,
                                 double* d_force,
                                 double* d_atom_virial,
@@ -778,6 +788,12 @@ bool DeepPot::uses_canonical_graph_inference() const {
 
 int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); }
 
+std::vector<double> DeepPot::get_default_chg_spin() const {
+  return dp->get_default_chg_spin();
+}
+
+bool DeepPot::has_atomic_virial() const { return dp->has_atomic_virial(); }
+
 void DeepPot::set_charge_spin(const std::vector<double>& charge_spin) {
   dp->set_charge_spin(charge_spin);
 }

source/api_cc/src/DeepPotPTExpt.cc

@@ -407,6 +407,15 @@ void DeepPotPTExpt::init(const std::string& model,
     check_charge_spin_domain(default_chg_spin_, chg_spin_table_ranges_);
   }
 
+  // A legacy compressed archive can carry a frozen default without shipping
+  // a fold. Retain that state without changing either conditioning width.
+  try {
+    default_chg_spin_ =
+        deepmd::ptexpt::read_default_chg_spin_for_query(metadata, dchgspin);
+  } catch (const deepmd::deepmd_exception& ex) {
+    default_chg_spin_query_error_ = ex.what();
+  }
+
   int num_intra_nthreads, num_inter_nthreads;
   get_env_nthreads(num_intra_nthreads, num_inter_nthreads);
   if (num_inter_nthreads) {
@@ -451,6 +460,7 @@ void DeepPotPTExpt::set_charge_spin(const std::vector<double>& charge_spin) {
                               *loader);
   }
   default_chg_spin_.swap(next_charge_spin);
+  default_chg_spin_query_error_.clear();
 }
 
 std::vector<torch::Tensor> DeepPotPTExpt::run_model(
@@ -2549,6 +2559,11 @@ void DeepPotPTExpt::compute_canonical_graph_gpu_impl(
     throw deepmd::deepmd_exception(
         "compute_canonical_graph_gpu requires a compact canonical artifact.");
   }
+  if (!has_atomic_virial()) {
+    throw deepmd::deepmd_exception(
+        "compute_canonical_graph_gpu requires atomic virials, but the model "
+        "was exported without them or does not declare atom_virial output.");
+  }
   if (!gpu_enabled) {
     throw deepmd::deepmd_exception(
         "compute_canonical_graph_gpu requires a CUDA device.");
@@ -2700,4 +2715,21 @@ bool DeepPotPTExpt::uses_canonical_graph_inference() const {
   return lower_input_is_canonical_;
 }
 
+std::vector<double> DeepPotPTExpt::get_default_chg_spin() const {
+  assert(inited);
+  if (!default_chg_spin_query_error_.empty()) {
+    throw deepmd::deepmd_exception(default_chg_spin_query_error_);
+  }
+  return default_chg_spin_;
+}
+
+bool DeepPotPTExpt::has_atomic_virial() const {
+  assert(inited);
+  // Canonical graphs expose public output names; other lower schemas use
+  // the internal energy derivative names consumed by compute().
+  const char* key = lower_input_is_canonical_ ? "atom_virial" : "energy_derv_c";
+  return do_atomic_virial && std::find(output_keys.begin(), output_keys.end(),
+                                       key) != output_keys.end();
+}
+
 #endif