perf: DPA4 .pt2 LAMMPS (C++ AOTI) inference 3-6x slower than Python (ASE) due to CPU-side edge construction and data marshalling

[SchrodingersCattt]

Summary

DPA4/SeZM .pt2 models exhibit a 3–6× throughput degradation when called from the C++ LAMMPS interface (DeepPotPTExpt::compute) compared to the Python interface (DeepEval._eval_model / ASE DP calculator), despite both paths executing the same AOTInductor-compiled graph from the same .pt2 archive.

At 65,536 atoms on an A800 GPU:

• Python (ASE): ~56 ms/step → 1.17M atom·step/s
• C++ (LAMMPS): ~507 ms/step → 130K atom·step/s (Pair component: ~350 ms)
• Gap: ~6.2× at large atom counts

For comparison, DPA1 (se_atten_v2, 0.52M params, .pth TorchScript) runs at ~128K atom·step/s in LAMMPS — meaning DPA4 with 115× fewer parameters (4,490 params) achieves essentially the same LAMMPS speed as DPA1, despite being 25× faster in Python/ASE.

Reproduction

Environment:

• deepmd-kit: 0.1.dev76+g9e08fcb0f (outisli branch)
• PyTorch: 2.7+ with CUDA 12.8
• LAMMPS: 22 Jul 2025 Update 2
• GPU: NVIDIA A800-SXM4-80GB

Model: DPA4 SeZM, n_blocks=1, channels=8, sel=181, lmax=2 (4,490 parameters total)

Benchmark protocol: Fe BCC, 20 warmup + 100 production NVE steps, atom counts from 250 to 101,306.

LAMMPS timing breakdown (from timer output, 100 production steps):

n_atoms	Loop (s)	Pair (s)	Neigh (s)	Pair−Neigh (s)	ASE (ms/step)	C++ model / Python model
250	0.27	0.19	0.08	0.11	3.5	0.3×
2,000	1.15	0.61	0.52	0.09	4.9	0.2×
9,826	5.75	3.42	2.27	1.15	11.1	1.0×
65,536	50.71	35.05	15.31	19.74	56.2	3.5×
101,306	77.88	54.73	22.66	32.07	88.4	3.6×

The Pair−Neigh column isolates the model inference + data marshalling time inside pair_deepmd. At small atom counts, C++ is actually faster than Python (JIT warmup amortized). At large atom counts, C++ is 3.5× slower — the gap grows with system size.

Key comparison with DPA1:

n_atoms=65,536	Params	LAMMPS Pair−Neigh (ms/step)
DPA1 .pth (TorchScript)	515,844	197
DPA4 .pt2 (AOTInductor)	4,490	197

The two models have identical C++ inference time despite 115× parameter difference. The DPA4 .pt2 AOTI graph is not delivering the expected speedup over TorchScript in the C++ call path.

Root cause analysis

We traced the full call chain in both paths:

Python path (DeepEval._eval_model, deep_eval.py):

numpy → torch.tensor(device=cuda)        # direct to GPU
     → VesinNeighborList.build()          # GPU O(N) cell list (vesin.torch)
     → edge_schema_from_extended()        # GPU tensor ops
     → self._pt2_runner(*model_inputs)    # AOTI graph (all inputs already on GPU)
     → .detach().cpu().numpy()            # single D2H at end

C++ path (DeepPotPTExpt::compute, DeepPotPTExpt.cc):

std::vector<double> coord → from_blob().clone().to(device)   # H2D #1
std::vector<int64_t> atype → from_blob().clone().to(device)  # H2D #2
LAMMPS InputNlist → nlist_data.copy_from_nlist()             # CPU copy
                  → shuffle_exclude_empty()                   # CPU remap
                  → mapping rebuild (O(nall) loop)            # CPU
                  → from_blob().clone().to(device)            # H2D #3
                  → createEdgeTensors()                       # CPU O(N×nnei) ← BOTTLENECK
                       → 4× from_blob().clone().to(device)   # H2D #4-7
loader->run(inputs)                                           # AOTI graph (same as Python)
output["energy"].to(kCPU)                                     # D2H #1 (implicit sync)
output["extended_force"].to(kCPU)                              # D2H #2
output["virial"].to(kCPU)                                      # D2H #3
select_map()                                                   # CPU remap

Three identified bottlenecks:

1. createEdgeTensors() runs entirely on CPU (~60-80ms at 65K atoms)

At 65,536 atoms × ~120 neighbors = ~7.8M edges, commonPT.h:183-330 iterates all neighbor pairs in C++ loops to compute edge indices and vectors, then copies to GPU with 4 separate .clone().to(device) calls.

In contrast, the Python path uses vesin.torch which builds the neighbor/edge list directly on GPU with O(N) cell-list algorithm.

2. Multiple CPU↔GPU synchronization points (~30-50ms)

The C++ path has ~7 .clone().to(device) H2D transfers and ~3-5 .to(torch::kCPU) D2H transfers, each causing an implicit cudaStreamSynchronize. The Python path keeps all data GPU-resident with a single D2H at the end.

3. CPU-side data marshalling (~10-20ms)

select_real_atoms_coord, std::vector<double> type conversions, nlist_data.shuffle_exclude_empty(), mapping rebuild loop, and select_map for output remapping — all serial CPU work.

NOT contributing factors:

• CUDA Graphs: Disabled in both paths (triton.cudagraphs: False in compile options)
• Different compiled kernels: Same .pt2 archive, same AOTI compiled graph
• DP_TRITON_INFER: Baked at trace time, not a runtime difference
• Model architecture: Both paths call the same forward_common_lower_exportable

Proposed optimizations

1. GPU neighbor list / edge construction for C++: Use vesin C++ API with CUDA backend (or a custom CUDA kernel) to replace the CPU createEdgeTensors(). This would eliminate the dominant ~60-80ms bottleneck.

2. Pinned memory + async transfers: Use CUDA pinned memory for from_blob and cudaMemcpyAsync to overlap H2D transfers with GPU computation.

3. Batch output D2H: Instead of 3-5 separate .to(kCPU) calls (each triggering sync), concatenate outputs into a single contiguous GPU tensor and do one D2H transfer.

4. CUDA Graphs for inference: The .pt2 inference path has no autograd (the triton.cudagraphs: False rationale about stale autograd metadata doesn't apply). Wrapping loader->run() in a CUDA graph capture would eliminate per-kernel launch overhead.

5. Cache GPU tensors across steps: When ago > 0 (neighbor list unchanged), avoid re-creating coord/atype tensors from scratch — update them in-place on GPU.

Impact

This performance gap means DPA4's architectural efficiency advantage (compact SO(3)-equivariant representation) is completely invisible in LAMMPS production MD simulations. A DPA4 model with 4,490 params runs at the same speed as a DPA1 model with 515,844 params, making the model compression meaningless for the most common deployment scenario.

Related to #5574

[dosubot]

Great analysis. This issue is well-documented and aligns with findings from your related feature request in #5574 [1] — both point to the same fundamental problem: the C++ inference path performs CPU-side topology construction and multiple host↔device round-trips that negate the efficiency of the AOTI-compiled graph.

A few observations on the current state of the code:

Re: createEdgeTensors() bottleneck — This function was recently updated in PR #5580 (merged June 25) with the fold_to_local parameter for multi-rank edge indexing [2]. The function iterates all neighbor pairs in C++ to build edge_index, edge_vec, edge_scatter_index, and edge_mask, then copies them to GPU with separate .clone().to(device) calls — exactly the bottleneck you identified.

Re: vesin in C++ — The vesin library is currently only used in the Python path (deepmd/pt_expt/utils/vesin_neighbor_list.py). There's no vesin C++ integration, so the C++ path falls back to the serial CPU iteration you described.

Prior optimization work — PR #5407 (April 2026) achieved major .pt2 speedups for DPA1/DPA2 by replacing CPU-side buildTypeSortedNlist with on-device createNlistTensor() and caching mapping_tensor [3]. However, DPA4's edge-based path (createEdgeTensors) was introduced after that optimization and hasn't received the same GPU-side treatment.

Your proposed optimizations are well-aligned with the existing infrastructure:

1. GPU edge construction — The VesinNeighborList already works on GPU in Python; wrapping the vesin C++ API with CUDA backend (or porting edge_schema_from_neighbor_matrix to a CUDA kernel) would directly address the dominant bottleneck.
2. Pinned memory + async transfers — Straightforward given that the current code already uses from_blob().clone().to(device).
3. Batched output D2H — Currently energy, extended_force, and virial are each .to(kCPU) separately [4], so consolidating is feasible.
4. CUDA Graphs — Since AOTI graphs are static (no autograd), this is viable for the inference-only path.
5. Tensor caching — Edge tensors (edge_index_tensor, edge_index_ext_tensor) are already declared as cached member variables in DeepPotPTExpt [5], but the edge construction still runs on CPU every rebuild.

The compute_device API proposed in #5574 would also benefit from these same fixes — both issues share the need for a GPU-native edge construction path in C++.

To reply, just mention @dosu.

---

Share context across your team and agents. Try Dosu.

[SchrodingersCattt]

@dosubot Thanks for the thorough cross-referencing.

To clarify the priority: the dominant bottleneck is createEdgeTensors() running on CPU — introduced by #5562 (edge schema for DPA4). The nlist-path optimization in #5407 addressed DPA1/DPA2 but doesn't cover the edge-based path that DPA4/SeZM now uses.

We confirmed this with profiling: the pure AOTI loader->run() call takes only 2.9 ms at 4394 atoms, but LAMMPS Pair reports 12.3 ms — the remaining 9.4 ms is CPU edge construction + H2D/D2H marshalling, which is 3.2× the model inference itself.

A GPU-side edge construction (porting vesin or a custom CUDA kernel to the C++ path) would be the highest-impact fix.

[dosubot]

Thanks for the profiling confirmation — 3.2× overhead from CPU edge construction alone makes it clear where the priority should be.

A couple of relevant findings that may help with the fix:

1. Existing GPU nlist template: build_nlist_gpu() in source/lib/src/gpu/neighbor_list.cu already implements a three-phase GPU pipeline (distance computation → CUB parallel prefix scan for compaction → dense list fill) that could serve as a starting point for a GPU createEdgeTensors() [1]. The key operations to port — neighbor masking, index flattening, and edge-to-node scatter index construction — map well to this pattern.

2. Active work in PR feat(dpa4): multiple updates for DPA4/SeZM #5503: PR #5503 ("multiple updates for DPA4/SeZM") already includes GPU-accelerated O(N) neighbor list construction (nv_nlist backend) and edge cache utilities for the Python path [2]. This could potentially be extended or adapted for the C++ DeepPotPTExpt path as well.

The path forward seems straightforward: either expose the nv_nlist / vesin GPU construction from #5503 through a C++ interface that createEdgeTensors() can call, or implement a dedicated CUDA kernel following the CUB prefix-scan pattern from build_nlist_gpu() to directly produce edge_index, edge_vec, edge_scatter_index, and edge_mask on device — eliminating the CPU loop and the 4 separate .clone().to(device) transfers in one shot.

To reply, just mention @dosu.

---

Docs are dead. Just use Dosu.

[SchrodingersCattt]

@dosubot Good pointers on #5503 and build_nlist_gpu().

Looking at #5503 more closely: the nv_nlist GPU backend and edge_schema_from_extended() already produce edges on-device in Python — the exact data that the C++ path builds with CPU loops in createEdgeTensors(). The Python SeZMModel.build_neighbor_list() (sezm_model.py:2555) dispatches to _select_neighbor_builder() which picks vesin or nv_nlist, both GPU-native.

So the GPU edge construction code already exists in the Python path. The gap is specifically that DeepPotPTExpt::compute() in C++ doesn't use it — it receives a CPU nlist from LAMMPS and rebuilds edges with CPU C++ loops.

The cleanest fix might be: instead of porting vesin/nv_nlist to C++, pass the LAMMPS nlist to the GPU and run edge_schema_from_extended() (or its equivalent) as torch tensor ops on-device — the function in deepmd/pt_expt/utils/edge_schema.py is already pure tensor ops, no Python-specific logic.

[dosubot]

That approach makes a lot of sense — if edge_schema_from_extended() (or the equivalent edge construction logic in your branch) is purely torch tensor ops, it could be compiled into the .pt2 archive itself or called as a standalone traced function from C++ via LibTorch, avoiding the need to rewrite it as a CUDA kernel.

Two possible implementation routes:

1. Include edge construction in the exported graph: If edge_schema_from_extended() can be traced/exported as part of the .pt2 archive (i.e., bake it into the model's forward_lower so the AOTI graph accepts raw nlist + mapping and produces edges internally), then DeepPotPTExpt::compute() would just need to pass the nlist tensor to GPU and let the compiled graph handle the rest. This eliminates the C++ createEdgeTensors() entirely.

2. Export edge construction as a separate AOTI artifact: If including it in the main graph causes shape dynamism issues (edge count varies per step), export it as a standalone traced function that C++ calls before loader->run(). Still pure tensor ops, still on-device, no custom CUDA kernel needed.

Route 1 is cleaner if torch.export can handle the dynamic edge count (with dynamic shapes). Given that your .pt2 freeze already handles DPA4, this might be the most natural place to fold in the edge construction — making the C++ path match what Python already does implicitly through the model's own preprocessing.

The key win either way: moving the nlist → edge conversion from C++ CPU loops to compiled torch tensor ops on GPU, which is exactly the code that already works in the Python path.

To reply, just mention @dosu.

---

Share context across your team and agents. Try Dosu.