torchref.base.targets.triton.nonbonded module

Triton forward + analytic backward for the non-bonded prolsq VDW target.

The forward fuses the gather + (optional) cartesian symmetry transform + prolsq shape energy + per-pair constant into one kernel. The backward chains through:

v(d) = max(0, d_vdw + buf − d) E(v) = c_rep · v^r_exp (when v > 0) dE/dv = c_rep · r_exp · v^(r_exp − 1) dv/dd = −1 (when v > 0, else 0) dd/dpos1 = −diff/d dd/dpos2 = +diff/d grad_pos2_source = R_cartᵀ · grad_pos2 (with symmetry)

Cartesian symmetry transforms (M·R·M⁻¹ and M·t) are precomputed once on the host, so the kernel only does a 3×3 matvec (forward) and 3×3 transposed matvec (backward) per pair.

torchref.base.targets.triton.nonbonded.nonbonded_heavy_math_triton(xyz, indices, min_distances, symop_indices, cell_offsets, symop_matrices, symop_translations, fractional_matrix, inv_fractional_matrix, c_rep, r_exp, buffer_, sigma_vdw)

Triton-backed heavy-heavy VDW prolsq NLL with analytic backward.