[pull] master from deepmodeling:master - #269
Merged
Merged
Conversation
DPA1 graph inference previously decomposed descriptor, fitting, and force/virial work into many generic tensor operations, retained large edge-scale autograd state, and lacked a shared sparse topology contract across Python export, the C++ runtime, and LAMMPS. Build an inference-oriented graph pipeline that keeps eligible attention-free DPA1 models and their neighbor data on the GPU while retaining explicit fallbacks for unsupported configurations. - Add inference-only CUDA operators for attention-free concat and strip DPA1. The uncompressed path fuses environment construction, the three-layer embedding network, type handling, moment reduction, Gram contraction, and analytical edge-gradient backward. The geometrically compressed strip path fuses quintic table interpolation, type-pair gating, moments, descriptors, and analytical gradients. It supports masked edges, canonical or permutation CSR, int32 or int64 addressing, width buckets from 8 through 256, non-bucket padding and slicing, and axis widths up to 16. - Add descriptor-agnostic fused energy fitting and force/virial operators. Run eligible fitting networks through pedantic FP32 cuBLAS GEMMs with TF32 disabled and fused bias, activation, timestep, residual, and derivative epilogues, while retaining atomic and per-frame energies in FP64. Reduce both edge incidences with one warp per node and no global floating-point atomics, then accumulate frame virials through FP64 partial reductions without materializing an (E, 9) outer-product tensor. - Define cumulative DP_CUDA_INFER levels. Level 1 uses separately registered descriptor, fitting, and CSR force operators with first-order autograd. Level 2 returns force and virial as values, using one opaque end-to-end operator for eligible uncompressed DPA1 and an explicit custom-operator chain for compressed DPA1. This removes the inference autograd tape, suppresses the unused rotation output, avoids retaining the descriptor only for shape metadata, and shortens saved-state lifetimes for memory reuse. - Add make_fx- and export-composable Triton kernels for the PyTorch se-family environment matrix and the dense and graph DPA1 environment convolutions. Their closed-form first backward preserves the force path while avoiding decomposed gather, switch, normalization, gating, and segment reduction chains. DP_TRITON_INFER=2 resolves per-GPU launch tables and freeze-time tuning fills uncovered model shapes; level 3 optionally enables the compensated FP16x3 final embedding GEMM. The default CUDA graph path remains FP32. - Select compressed-kernel resource policy at first uncaptured use from balanced or occupancy-oriented 128/256-thread launches, cache the result by device, direction, descriptor shape, topology, index width, and workload class, and retain architecture-specific fallbacks for small inputs, capture, or tuning failures. Build the native operators only with a CUDA-enabled PyTorch and include a lowest-supported portable PTX target alongside native code. - Extend NeighborGraph with keyword-only destination/source CSR orders, int64 row pointers, and an explicit destination_sorted property without changing its original positional constructor. Generic builders preserve payload order and omit CSR by default; consumers opt into CSR or stable destination-major canonicalization, which moves masked guards to the suffix and makes destination order the identity. Edge masks remain authoritative inside every row, and export validates permutations, row membership, bounds, and canonical identity before tracing. - Make graph-form AOTInductor deployment select the optimized path automatically for eligible .pt2 conversion and export compressed graph models directly. Keep frame, node, and edge axes dynamic; preserve custom operators through export; record the graph edge-vector dtype in metadata; and make per-atom virial part of the graph artifact contract required by the Kokkos consumer. Eligible compressed artifacts accept FP32 edge geometry directly, while generic and uncompressed graph artifacts retain the FP64 geometry ABI. Use int64 Inductor indexing and bounded Triton launch tiling for large dynamic graphs. - Extend DeepEval and the C++ API to consume the same canonical graph ABI. Add device-edge capability and dtype queries, FP32 and FP64 edge-vector overloads, runtime frame and atomic parameters, total-versus-owned node counts, and optional communication metadata. Host ingestion canonicalizes arbitrary payloads, while the device path constructs destination identity CSR and source order with histogram, prefix-sum, and counting scatter. - Add pair_style deepmd/kk for device-resident edge and graph .pt2 inference. Build compact model-cutoff edges from the Kokkos full neighbor list with count, scan, and fill passes; compact NULL-mapped atom types; emit FP32 or FP64 vectors according to artifact metadata; and scatter model-node outputs back on device. Single-rank execution folds periodic images onto local owners, whereas domain decomposition retains local-plus-halo nodes and explicitly reverse-communicates force and centroid per-atom virial through either device-aware or host-staged communication. Message-passing edge artifacts use their with-comm forward, including empty-rank phantom inputs; message-passing graph artifacts fail fast because that multi-rank ABI is not supported. - Preserve graph correctness at ownership and masking boundaries. Exclude halo fitting outputs from energy, zero halo atomic parameters, retain virtual-atom and pair-exclusion masks in fused level 2, handle zero-node graphs, and reduce frame virial from node virials instead of a contended edge scatter. Validate and broadcast multi-frame fparam/aparam layouts, fix wide atomic-parameter allocation when filtering virtual atoms, and avoid copying incompletely constructed LAMMPS model members. Reduce graph-builder overhead with bounded GPU neighbor-capacity estimates and direct length-three periodic-shift reductions. - Require regeneration of graph-form .pt2 artifacts because their positional ABI now includes n_local and both CSR views. Compressed graph export is limited to FP32, geometrically compressed, attention-free strip models without excluded type pairs, with output width at most 256 and axis_neuron <= 16. The uncompressed CUDA descriptor requires three FP32 embedding layers whose widths stay equal or double within its compiled bounds. Other configurations continue through Triton, reference graph, or nlist lowering. deepmd/kk additionally requires a GPU edge/graph artifact, one model, and a valid atom map. - Add regression coverage for CUDA and Triton forward/gradient parity, compressed and uncompressed end-to-end energy, force, global virial, atomic energy, and atom virial; smooth and non-smooth one- and two-sided type gates; residual and activation variants; non-power-of-two widths; int32/int64 addressing; masked cached edges; canonical and permuted CSR; ownership; zero-node and many-frame reductions; make_fx, torch.compile, and dynamic graph export. Extend DeepEval and C++ tests for dynamic edge counts, multi-rank graph ingestion, parameter validation and broadcasting, atomic outputs, CSR construction, and virtual-atom parameter preservation.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )