deepmodeling / deepmodeling/deepmd-kit
perf: DPA4 .pt2 LAMMPS (C++ AOTI) inference 3-6x slower than Python (ASE) due to CPU-side edge construction and data marshalling
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## 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 coord → from_blob().clone().to(device) # H2D #1
std::vector 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` 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
Contributor guide
Assessment
This issue has not been assessed yet.