perf(dpa1): optimize graph CUDA inference and deployment#5758
Open
OutisLi wants to merge 12 commits into
Open
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds high-performance inference paths for the attention-free DPA1 (se_atten, attn_layer == 0) and related graph-lower workflows, introducing Triton and hand-written CUDA/cuBLAS kernels and wiring them into both Python inference and exported .pt2 artifacts (including LAMMPS/Kokkos consumption).
Changes:
- Add Triton fused kernels for environment-matrix (
env_mat) and DPA1 fused environment convolution (se_conv/edge_conv) with parity/FX-tracing tests. - Add CUDA graph-lower fused operators for fitting (
graph_fitting), force/virial assembly (edge_force_virial), and an end-to-end DPA1 energy-force operator (dpa1_graph_energy_force), plus export/autotune plumbing. - Extend the C++ API + LAMMPS integration to support GPU-resident edge/graph inference with runtime params and multi-rank/message-passing routing.
Reviewed changes
Copilot reviewed 58 out of 58 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| source/tests/pt/model/test_env_mat_triton.py | New CUDA/Triton parity tests for fused dense env-mat (forward + force). |
| source/tests/pt_expt/utils/test_edge_env_mat_triton.py | New parity + FX-tracing tests for edge-stream env-mat Triton op. |
| source/tests/pt_expt/model/test_edge_energy_deriv.py | Adds tests for inference-only force_precision downcast behavior. |
| source/tests/pt_expt/model/test_dpa1_graph_lower.py | Adjusts force-derivative dtype assertion for graph vs dense precision behavior. |
| source/tests/pt_expt/descriptor/test_dpa1_triton.py | New tests for DPA1 Triton routing/parity and FX “opaque node” capture. |
| source/op/pt/graph_ops.h | New shared C++ declarations for fused graph-lower operator entrypoints. |
| source/op/pt/graph_fitting.cu | New fused cuBLAS-based graph-lower fitting op (+ backward). |
| source/op/pt/edge_force_virial.cu | New fused CUDA scatter/reduction op assembling force + virials from edge grads. |
| source/op/pt/dpa1_graph_energy_force.cu | New end-to-end fused DPA1 graph energy/force/virial CUDA op. |
| source/op/pt/CMakeLists.txt | Builds/links new CUDA operator sources and cuBLAS when CUDA toolkit enabled. |
| source/lmp/pair_deepspin.cpp | Emits model summary in derived ctor; improves comments. |
| source/lmp/pair_deepmd.h | Adds comm-only nlist builder helper for message-passing/device-resident paths. |
| source/lmp/pair_deepmd.cpp | Emits model summary in derived ctor; implements make_comm_nlist; comment cleanups. |
| source/lmp/pair_deepmd_kokkos.h | New Kokkos device pair style header for GPU-resident .pt2 inference integration. |
| source/lmp/pair_base.h | Switches base model members to references (avoid use-before-constructed copies). |
| source/lmp/pair_base.cpp | Moves summary emission out of base ctor; clarifies centroid-stress and restart behavior. |
| source/lmp/builtin.cmake | Links DeePMD C++ API and defines DP_USE_CXX_API when Kokkos package enabled. |
| source/api_cc/src/DeepPotPTExpt.cc | Clarifies multi-rank/message-passing dispatch and error messaging; comment updates. |
| source/api_cc/src/DeepPot.cc | Extends compute_edges_gpu API to accept runtime params + extended-node/comm info. |
| source/api_cc/include/DeepPotPTExpt.h | Updates compute_edges_gpu contract to support edge/graph inputs and comm metadata. |
| source/api_cc/include/DeepPot.h | Adds overload for runtime fparam/aparam and extended-node/comm support in edge GPU API. |
| deepmd/pt/model/descriptor/se_atten.py | Routes env-convolution tail through Triton fused ops for eligible inference configs. |
| deepmd/pt/model/descriptor/env_mat.py | Routes prod_env_mat through Triton fused env-mat in opt-in inference. |
| deepmd/pt_expt/utils/vesin_neighbor_list.py | Replaces small matmul with broadcast reduce to avoid inefficient GEMM. |
| deepmd/pt_expt/utils/serialization.py | Adds graph-lower auto selection, forces fused CUDA level for graph exports, adds autotune and better logging. |
| deepmd/pt_expt/utils/edge_schema.py | Same matmul→broadcast-reduce optimization for image offsets. |
| deepmd/pt_expt/model/make_model.py | Adds level-2 fused energy/force graph path and inference force/virial precision control. |
| deepmd/pt_expt/model/edge_transform_output.py | Adds force_precision inference cast + routes scatter through fused CUDA op when available. |
| deepmd/pt_expt/fitting/ener_fitting.py | Adds graph-native fitting routing through fused CUDA op when eligible. |
| deepmd/pt_expt/entrypoints/compress.py | Re-exports compressed graph-lower models directly to graph .pt2 when eligible. |
| deepmd/pt_expt/descriptor/se_atten_v2.py | Exposes fused eligibility and fused paths for v2 descriptor wrapper. |
| deepmd/kernels/utils.py | Documents/introduces DP_CUDA_INFER gating and expands Triton level documentation. |
| deepmd/kernels/triton/dpa1/tile_configs.py | New DPA1 Triton launch config tables (builtin + runtime registration). |
| deepmd/kernels/triton/dpa1/sweep_tile_configs.py | New sweep tool to generate per-device DPA1 Triton launch configs. |
| deepmd/kernels/triton/dpa1/gemm_fp16x3.py | New inference-only fp16x3 tensor-core GEMM for DPA1 embedding (level 3). |
| deepmd/kernels/triton/dpa1/activation.py | New shared activation codes + Triton activation helpers for DPA1 kernels. |
| deepmd/kernels/triton/dpa1/init.py | New package init exporting DPA1 Triton ops and availability. |
| deepmd/kernels/cuda/graph_fitting.py | New Python bindings + fake/CPU/autograd glue for deepmd::graph_fitting. |
| deepmd/kernels/cuda/edge_force_virial.py | New Python bindings + fake/CPU glue for deepmd::edge_force_virial. |
| deepmd/kernels/cuda/dpa1/graph_energy_force.py | New Python bindings + fake/CPU glue for deepmd::dpa1_graph_energy_force. |
| deepmd/kernels/cuda/dpa1/init.py | New CUDA DPA1 package init/docs. |
| deepmd/kernels/cuda/init.py | New CUDA kernels package init/docs. |
| deepmd/kernels/autotune.py | New freeze-time Triton autotune registry and driver. |
| deepmd/entrypoints/convert_backend.py | Passes lower_kind="auto" to backend hooks when supported; forwards atomic-virial flag when applicable. |
| deepmd/dpmodel/utils/neighbor_graph/derivatives.py | Optimizes per-frame virial reduction (reduce atom_virial by node frame). |
| deepmd/dpmodel/descriptor/dpa1.py | Expands graph-lower eligibility to geo-compressed strip DPA1 (attn_layer==0). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds fused CUDA and Triton inference paths for DPA1 descriptor, fitting, and force/virial computation, updates export/routing logic, and extends LAMMPS GPU edge inference to accept extended-node and communication inputs. ChangesFused CUDA/Triton inference kernels and Python routing
LAMMPS/api_cc GPU edge inference
Estimated code review effort: 5 (Critical) | ~150 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (3)
source/lmp/builtin.cmake (1)
80-84: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueGuard the include-dir property against NOTFOUND / generator expressions.
If
INTERFACE_INCLUDE_DIRECTORIESis unset,_deepmd_incbecomes..-NOTFOUNDand the loop still appends a bogus<...>-NOTFOUND/deepmdpath. If it contains generator expressions (e.g.$<BUILD_INTERFACE:...>), string-concatenating/deepmdproduces malformed paths. Consider skipping when unset and being aware of genexes.♻️ Suggested guard
get_target_property(_deepmd_inc DeePMD::deepmd_cc INTERFACE_INCLUDE_DIRECTORIES) if(_deepmd_inc) foreach(_inc IN LISTS _deepmd_inc) target_include_directories(lammps PUBLIC "${_inc}/deepmd") endforeach() endif()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/lmp/builtin.cmake` around lines 80 - 84, Guard the DeePMD include directory handling in builtin.cmake so `get_target_property(_deepmd_inc DeePMD::deepmd_cc INTERFACE_INCLUDE_DIRECTORIES)` does not feed a NOTFOUND value or generator expressions into `target_include_directories(lammps ...)`. Update the `_deepmd_inc`/`foreach(_inc IN LISTS _deepmd_inc)` block to skip when the property is unset and avoid blindly appending `/deepmd` to entries that may be generator expressions, using the existing `DeePMD::deepmd_cc` include-property logic as the place to fix.source/lmp/pair_deepmd_kokkos.h (1)
11-16: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
deepmd/kk/hostis registered but unconditionally errors atinit_style(in GPU Kokkos builds).In a Kokkos GPU build,
PairDeepMDKokkos<LMPHostType>::init_style()(pair_deepmd_kokkos.cpp lines 125-127) always callserror->all(...)since this style is GPU-only. Registering a pair style whose only behavior is to error may confuse users browsing available styles. Consider either gating thedeepmd/kk/hostPairStyleentry behind a build flag that only exists when a genuine host fallback is implemented, or adding a doc comment here noting it is a reserved/placeholder name.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/lmp/pair_deepmd_kokkos.h` around lines 11 - 16, The deepmd/kk/host PairStyle entry in pair_deepmd_kokkos.h is registered even though PairDeepMDKokkos<LMPHostType>::init_style() unconditionally errors in GPU Kokkos builds. Update the registration so deepmd/kk/host is only exposed when a real host implementation exists, or annotate the PairStyle block with a clear placeholder/reserved comment; use the PairStyle declarations and PairDeepMDKokkos<LMPHostType>::init_style() as the key symbols to locate the change.deepmd/kernels/utils.py (1)
75-86: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated parse/validate logic between
triton_infer_levelandcuda_infer_level.Both functions repeat the same "strip → int() → ValueError → membership-check → ValueError" pattern, differing only by env var name and levels tuple. Extracting a shared helper would avoid drift as new levels/knobs are added.
♻️ Proposed shared helper
+def _read_env_level(env_var: str, levels: tuple[int, ...]) -> int: + raw = os.environ.get(env_var, "0").strip() + try: + level = int(raw) + except ValueError: + raise ValueError( + f"{env_var} must be an integer in {levels}, got {raw!r}" + ) from None + if level not in levels: + raise ValueError(f"{env_var} must be one of {levels}, got {level}") + return level + + def triton_infer_level() -> int: ... - raw = os.environ.get("DP_TRITON_INFER", "0").strip() - try: - level = int(raw) - except ValueError: - raise ValueError( - f"DP_TRITON_INFER must be an integer in {TRITON_INFER_LEVELS}, got {raw!r}" - ) from None - if level not in TRITON_INFER_LEVELS: - raise ValueError( - f"DP_TRITON_INFER must be one of {TRITON_INFER_LEVELS}, got {level}" - ) - return level + return _read_env_level("DP_TRITON_INFER", TRITON_INFER_LEVELS) def cuda_infer_level() -> int: ... - raw = os.environ.get("DP_CUDA_INFER", "0").strip() - try: - level = int(raw) - except ValueError: - raise ValueError( - f"DP_CUDA_INFER must be an integer in {CUDA_INFER_LEVELS}, got {raw!r}" - ) from None - if level not in CUDA_INFER_LEVELS: - raise ValueError( - f"DP_CUDA_INFER must be one of {CUDA_INFER_LEVELS}, got {level}" - ) - return level + return _read_env_level("DP_CUDA_INFER", CUDA_INFER_LEVELS)Also applies to: 134-145
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/kernels/utils.py` around lines 75 - 86, The parsing and validation flow in triton_infer_level and cuda_infer_level is duplicated, so extract the shared “strip, convert to int, validate against allowed levels” logic into a common helper and have both functions call it with their env var name and levels tuple. Keep the existing ValueError messages/context aligned with the current behavior, and update both triton_infer_level and cuda_infer_level to use the shared helper so future knob changes stay consistent.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/kernels/cuda/dpa1/graph_energy_force.py`:
- Around line 248-257: The CPU fallback in graph_energy_force currently passes
g_e and edge_vec through to edge_force_virial without matching the expected
weight dtype, so this branch can return fp64 force/virial outputs instead of
w1.dtype like _fake, CUDA, and graph_compress._ef_cpu. Update the
graph_energy_force path to cast both g_e and edge_vec to w1.dtype before calling
torch.ops.deepmd.edge_force_virial, keeping the return contract consistent
across backends.
In `@deepmd/kernels/cuda/graph_fitting.py`:
- Around line 102-119: _trim the forward workspace saved by graph_fitting so
only the adot portion is preserved for backward. In _forward_fake, the returned
temporary buffer is currently sized for both h and adot, but
graph_fitting_backward only consumes the adot slice; update the forward
allocation and any corresponding saved state logic to keep h as local scratch
and return/save just the derivative chunks, using the existing _forward_fake and
graph_fitting_backward symbols to align the ABI._
In `@deepmd/pt_expt/descriptor/se_atten_v2.py`:
- Around line 60-79: DescrptSeAttenV2 is still using the inherited dpmodel
call_graph, so the accelerated dispatch path in DescrptDPA1.call_graph is never
reached. Add call_graph delegation in DescrptSeAttenV2, following the existing
pattern used by _fused_eligible, fused_energy_force_graph, and _call_triton, and
include the related _call_graph_* helper delegations that call_graph relies on.
In `@deepmd/pt_expt/model/edge_transform_output.py`:
- Around line 108-113: The fused scatter path in
fit_output_to_model_output_graph is currently selected based only on
cuda_infer_level() and fused_scatter_available(), which can bypass the
differentiable fallback during training. Update the guard around
fused_edge_force_virial to also check that create_graph is false (or otherwise
restrict it to inference-only use) so training with higher-order gradients keeps
the autodiff path. Use the fit_output_to_model_output_graph and
fused_edge_force_virial symbols to locate the branch and preserve the existing
fallback behavior when create_graph is enabled.
In `@source/api_cc/src/DeepPotPTExpt.cc`:
- Around line 1746-1754: The compute_edges_gpu() path needs the same fail-fast
guard as compute() when nall_nodes > nloc. In DeepPotPTExpt::compute_edges_gpu,
use the existing extended/use_with_comm checks to reject extended node sets
unless the model is graph-input or a message-passing model with the required
with-comm artifact and comm_nlist, so the code does not fall through into
run_model_graph() without ghost-feature exchange.
In `@source/op/pt/dpa1_graph_compress.cu`:
- Around line 786-791: The per-frame reduction in the energy aggregation block
is using a frame_id built from repeat_interleave(arange(nf), n_node), which only
matches real nodes and can misalign with the padded node axis used by
atom_energy. Update this logic in the Step 3 segment-sum path to use a
padding-safe frame mapping derived from the full node_capacity/flat node axis,
and zero out any padding contributions before calling index_add_ so the
reduction stays valid for padded frames.
In `@source/op/pt/graph_fitting.cu`:
- Around line 38-46: The singleton in cublas_handle() is shared across devices,
streams, and threads, but a cuBLAS handle is device-bound and can be retargeted
by concurrent users. Replace the single static handle with a per-device cache
(or equivalent thread-local/device-keyed storage) so each GPU gets its own
initialized handle, and keep the CUBLAS_PEDANTIC_MATH setup in the
initialization path for each cached handle.
---
Nitpick comments:
In `@deepmd/kernels/utils.py`:
- Around line 75-86: The parsing and validation flow in triton_infer_level and
cuda_infer_level is duplicated, so extract the shared “strip, convert to int,
validate against allowed levels” logic into a common helper and have both
functions call it with their env var name and levels tuple. Keep the existing
ValueError messages/context aligned with the current behavior, and update both
triton_infer_level and cuda_infer_level to use the shared helper so future knob
changes stay consistent.
In `@source/lmp/builtin.cmake`:
- Around line 80-84: Guard the DeePMD include directory handling in
builtin.cmake so `get_target_property(_deepmd_inc DeePMD::deepmd_cc
INTERFACE_INCLUDE_DIRECTORIES)` does not feed a NOTFOUND value or generator
expressions into `target_include_directories(lammps ...)`. Update the
`_deepmd_inc`/`foreach(_inc IN LISTS _deepmd_inc)` block to skip when the
property is unset and avoid blindly appending `/deepmd` to entries that may be
generator expressions, using the existing `DeePMD::deepmd_cc` include-property
logic as the place to fix.
In `@source/lmp/pair_deepmd_kokkos.h`:
- Around line 11-16: The deepmd/kk/host PairStyle entry in pair_deepmd_kokkos.h
is registered even though PairDeepMDKokkos<LMPHostType>::init_style()
unconditionally errors in GPU Kokkos builds. Update the registration so
deepmd/kk/host is only exposed when a real host implementation exists, or
annotate the PairStyle block with a clear placeholder/reserved comment; use the
PairStyle declarations and PairDeepMDKokkos<LMPHostType>::init_style() as the
key symbols to locate the change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 2eb6f406-977b-42c2-b12e-3e8d5ff303c1
📒 Files selected for processing (58)
deepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/utils/neighbor_graph/derivatives.pydeepmd/entrypoints/convert_backend.pydeepmd/kernels/autotune.pydeepmd/kernels/cuda/__init__.pydeepmd/kernels/cuda/dpa1/__init__.pydeepmd/kernels/cuda/dpa1/graph_compress.pydeepmd/kernels/cuda/dpa1/graph_descriptor.pydeepmd/kernels/cuda/dpa1/graph_energy_force.pydeepmd/kernels/cuda/edge_force_virial.pydeepmd/kernels/cuda/graph_fitting.pydeepmd/kernels/triton/dpa1/__init__.pydeepmd/kernels/triton/dpa1/activation.pydeepmd/kernels/triton/dpa1/edge_conv.pydeepmd/kernels/triton/dpa1/gemm_fp16x3.pydeepmd/kernels/triton/dpa1/se_conv.pydeepmd/kernels/triton/dpa1/sweep_tile_configs.pydeepmd/kernels/triton/dpa1/tile_configs.pydeepmd/kernels/triton/env_mat.pydeepmd/kernels/utils.pydeepmd/pt/model/descriptor/env_mat.pydeepmd/pt/model/descriptor/se_atten.pydeepmd/pt_expt/descriptor/dpa1.pydeepmd/pt_expt/descriptor/se_atten_v2.pydeepmd/pt_expt/entrypoints/compress.pydeepmd/pt_expt/fitting/ener_fitting.pydeepmd/pt_expt/model/edge_transform_output.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/utils/edge_schema.pydeepmd/pt_expt/utils/serialization.pydeepmd/pt_expt/utils/vesin_neighbor_list.pysource/api_cc/include/DeepPot.hsource/api_cc/include/DeepPotPTExpt.hsource/api_cc/src/DeepPot.ccsource/api_cc/src/DeepPotPTExpt.ccsource/lmp/builtin.cmakesource/lmp/pair_base.cppsource/lmp/pair_base.hsource/lmp/pair_deepmd.cppsource/lmp/pair_deepmd.hsource/lmp/pair_deepmd_kokkos.cppsource/lmp/pair_deepmd_kokkos.hsource/lmp/pair_deepspin.cppsource/op/pt/CMakeLists.txtsource/op/pt/dpa1_graph_common.cuhsource/op/pt/dpa1_graph_compress.cusource/op/pt/dpa1_graph_descriptor.cusource/op/pt/dpa1_graph_energy_force.cusource/op/pt/edge_force_virial.cusource/op/pt/graph_fitting.cusource/op/pt/graph_ops.hsource/tests/pt/model/test_descriptor_dpa1_triton.pysource/tests/pt/model/test_env_mat_triton.pysource/tests/pt_expt/descriptor/test_dpa1_cuda.pysource/tests/pt_expt/descriptor/test_dpa1_triton.pysource/tests/pt_expt/model/test_dpa1_graph_lower.pysource/tests/pt_expt/model/test_edge_energy_deriv.pysource/tests/pt_expt/utils/test_edge_env_mat_triton.py
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/pt/model/descriptor/se_atten.py (1)
647-654: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: extract the duplicated fused-path eligibility gate.
The same six-term condition is repeated verbatim here and at Lines 726-733. A single property (e.g.
self._use_se_conv(rr)) would keep the two branches in sync and make the TorchScript-first ordering intentional in one place.♻️ Sketch
def _se_conv_active(self, rr: torch.Tensor) -> bool: return ( not torch.jit.is_scripting() and self.use_triton_infer and self.se_conv_eligible and not self.training and not self.geo_compress and rr.is_cuda )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt/model/descriptor/se_atten.py` around lines 647 - 654, The fused-path eligibility condition is duplicated in the SE attention logic, so extract the repeated six-term gate into a single helper or property on the `SeAtten`/descriptor class (for example, a `_se_conv_active(rr)`-style method) and use it in both branches. Keep the TorchScript-first `torch.jit.is_scripting()` check and the existing flags (`use_triton_infer`, `se_conv_eligible`, `training`, `geo_compress`, and `rr.is_cuda`) inside that shared helper so both call sites stay in sync.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@deepmd/pt/model/descriptor/se_atten.py`:
- Around line 647-654: The fused-path eligibility condition is duplicated in the
SE attention logic, so extract the repeated six-term gate into a single helper
or property on the `SeAtten`/descriptor class (for example, a
`_se_conv_active(rr)`-style method) and use it in both branches. Keep the
TorchScript-first `torch.jit.is_scripting()` check and the existing flags
(`use_triton_infer`, `se_conv_eligible`, `training`, `geo_compress`, and
`rr.is_cuda`) inside that shared helper so both call sites stay in sync.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 2eb6f406-977b-42c2-b12e-3e8d5ff303c1
📒 Files selected for processing (6)
deepmd/kernels/cuda/dpa1/graph_compress.pydeepmd/kernels/triton/dpa1/sweep_tile_configs.pydeepmd/kernels/triton/env_mat.pydeepmd/pt/model/descriptor/se_atten.pysource/op/pt/CMakeLists.txtsource/op/pt/graph_fitting.cu
💤 Files with no reviewable changes (2)
- deepmd/kernels/cuda/dpa1/graph_compress.py
- deepmd/kernels/triton/env_mat.py
🚧 Files skipped from review as they are similar to previous changes (3)
- source/op/pt/CMakeLists.txt
- deepmd/kernels/triton/dpa1/sweep_tile_configs.py
- source/op/pt/graph_fitting.cu
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@source/install/docker_package_c.sh`:
- Around line 29-32: The CUDA repository setup in docker_package_c.sh still uses
an insecure http:// NVIDIA repo URL; update the repo URL used by the yum
config-manager call to https://, matching the secure pattern already used in the
CodeQL workflow. Also search for the same NVIDIA CUDA repository URL in
pyproject.toml and change that instance to HTTPS as well so both locations are
consistent.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: f7e18dfd-95d8-49f0-88ab-63d714101f62
📒 Files selected for processing (4)
.github/workflows/build_cc.yml.github/workflows/codeql.ymlpyproject.tomlsource/install/docker_package_c.sh
✅ Files skipped from review due to trivial changes (1)
- .github/workflows/build_cc.yml
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5758 +/- ##
==========================================
- Coverage 79.77% 78.38% -1.40%
==========================================
Files 1020 1040 +20
Lines 116872 119947 +3075
Branches 4308 4345 +37
==========================================
+ Hits 93239 94019 +780
- Misses 22095 24376 +2281
- Partials 1538 1552 +14 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz-bot
suggested changes
Jul 10, 2026
njzjz-bot
left a comment
njzjz-bot
left a comment
Contributor
There was a problem hiding this comment.
Requesting changes because the current fused/default routing can change model results and silently break existing export/build/API contracts. The inline comments focus on atomic-model finalization, convert-backend defaults, the public C++ virtual API, Kokkos build gating, and ownership of the graph-lower capability predicate.
OutisLi
force-pushed
the
pr/dpa1
branch
2 times, most recently
from
4e37fb8 to
c73ecd0
Compare
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.
OutisLi
changed the title
Introduce a source-only canonical graph ABI for eligible geometrically compressed DPA1 models, allowing LAMMPS Kokkos to pass cutoff-compacted dual-CSR topology directly into AOTInductor without rebuilding graph metadata each step. - replace the generic ten-tensor deployment graph with a dedicated eight-tensor ABI containing atom types, owned/total node counts, int64 source indices, FP32 edge vectors, destination/source row pointers, and source order - derive destinations implicitly from canonical CSR rows and remove deployment edge masks, destination indices, destination order, index casts, guard concatenation, and bridge-side CSR construction - add compact compressed-descriptor forward/backward and force/virial operators while sharing the existing CUDA kernels with the generic graph path - build destination and source CSR views directly in the Kokkos cutoff producer with capacity-backed persistent storage and checked size_t/int64 accounting - preserve owned/halo energy semantics, NULL-type compaction, empty-rank handling, centroid virial reverse communication, and standard C++ inference - select the compact ABI automatically when pt_expt compress exports an eligible DPA1 model, while retaining the full NeighborGraph ABI for attention, exclusions, masked topology, and unsupported configurations - reduce fitting workspaces by writing GEMMs directly into activation slots and converting backward derivatives in place - add dynamic export, E=0/1 guards, CUDA parity, C++ packing, DeepEval, single/multi-rank, and large-system regression coverage At 129,168 atoms the compact path reduces whole-step time by 16.8%, 14.8%, and 10.3% for S, M, and L. The M process GPU-memory peak drops by 1,034 MiB, and the S/M/L OOM ceilings increase to 9.52M, 7.53M, and 5.51M atoms.
Instantiate mask-free compressed descriptor forward and backward kernels for cutoff-compacted canonical graphs, while retaining authoritative mask checks for generic cached, excluded, and permutation graph inputs. Include graph validity mode in the adaptive tuning key so masked and compact workloads select launch configurations independently. This improves compact model-only inference by about 1.4% without changing outputs or the generic graph contract.
Preserve conversion and C++ API compatibility, gate the Kokkos pair style explicitly, and keep graph CSR data coherent with pair exclusions. Restore cross-platform C++ builds, centralize graph-lower capability checks, and remove redundant custom-op registration state.
Keep convert-backend command-free for graph-eligible artifacts and isolate CUDA stream headers from ROCm builds.
Use the lightweight c10 stream API so CUDA builds do not require cuSPARSE headers while preserving current-stream synchronization.
Avoid optional PyTorch CUDA headers in portable builds while preserving the device-output completion contract.
wanghan-iapcm
requested changes
Jul 12, 2026
5 tasks
This was referenced Jul 13, 2026
Use parallel order/row-pointer pairs consistently across Python, export, CUDA, and C++ interfaces, and make the standard force fallback branch explicit.
njzjz
reviewed
Jul 13, 2026
Provide pointer-based C entry points and header-only wrappers so built-in and distributed Kokkos integrations share the supported C API path.
Limit the CUDA 12.9 CCCL workaround to that release; CUDA 13 supports its current architecture set with -arch=all.
Preserve the C API device-graph additions alongside the latest API, loss, and neighbor-list fixes.
wanghan-iapcm
approved these changes
Jul 13, 2026
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
Collaborator
There was a problem hiding this comment.
All the review threads I raised are addressed and verified against the current HEAD:
neighbor_graph/graph.pyfield order: the full sweep to the parallel layoutdestination_order, destination_row_ptr, source_order, source_row_ptrlanded in 2540154 and is consistent across the dataclass, the dpmodel CSR builder, and the pt_expt torch twin.edge_transform_output.py: the standard array-API force/virial path is now the explicitelsebranch paired with the fused path, with a single shared return.DeepPotPTExpt.ccGPU-edges graph branch: the destination-major physical-edge contract is now documented at the call site, and the pair-exclusion case canonicalizes.deepmd/kernels/placement: deferred to tracking issue #5782 with the target (move underpt_expt/kernels/), the pt->pt_expt dependency direction, and thedpmodelmust-not-import invariant captured, as agreed.
Approving on that basis. To be clear about scope: this reflects that my review points are resolved — it is not a full independent audit of every new CUDA/Triton kernel and the Kokkos deployment path.
njzjz-bot
approved these changes
Jul 13, 2026
njzjz-bot
left a comment
njzjz-bot
left a comment
Contributor
There was a problem hiding this comment.
Reviewed current HEAD c2538f0. The blocking issues from my earlier review are addressed: model-level finalization is preserved in the fused path, the conversion behavior is intentional and covered, the public backend overload remains compatible, Kokkos now uses the C API in both built-in and plugin modes, and the graph capability predicate has been separated from training internals.
I also checked the updated dual-CSR ordering and validation, canonical graph packing, owned/halo masking, runtime fparam/aparam handling, C/C++ device API forwarding, and the Kokkos force/virial ownership and reverse-communication flow. I found no new blocking issue.
Validation performed:
- 42 neighbor-graph, builder, CSR, and canonical-graph tests passed, including 2 subtests.
- Randomized masked/canonical CSR invariants passed for varied node and edge counts.
- The full PR CI matrix is green, including Python, C++, CUDA, ROCm, wheels, CodeQL, and pre-commit.
- Heavier local pt_expt model tests were not runnable because the installed DeePMD operator was built against PyTorch 2.10.0 while the runtime is 2.12.1; the corresponding PR CI jobs passed.
Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh
njzjz
approved these changes
Jul 13, 2026
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 join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
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.