deepmodeling / deepmodeling/deepmd-kit

perf: NeighborGraph attention is 4-9x slower than dense and OOM-prone -- move to block-dense attention via the destination CSR

Open
#5,861 1 comment 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

## Summary

The NeighborGraph route is dense-competitive for attention-free descriptors, but **graph attention is 4-9x slower than the dense route and memory-bound to the point of OOM**, on both DPA-1 and DPA-2. The cause is the ragged pair-axis attention design; the proposed fix is block-dense attention through the existing destination CSR.

## Measurements (diamond at experimental density, a = 3.567 A, fp64)

Compiled `.pt2` AOTI inference, V100-32GB (6-layer DPA-2, g1=128, g2=32, 4 heads, repinit nsel 120 / rcut 6, repformer nsel 40 / rcut 4; ms/eval, 30 CUDA-synced evals):

| natoms | graph native | graph smooth | dense native | dense smooth | graph/dense |
|---:|---:|---:|---:|---:|---:|
| 64 | 54.4 | 68.9 | 13.7 | 14.5 | 4.7x |
| 216 | 201.4 | 259.7 | 29.4 | 31.4 | 8.3x |
| 512 | 725.2 | OOM (32 GB) | 57.3 | 62.0 | — |

Eager per-op attribution (T4, 64 atoms, `torch.no_grad`, CUDA-synced): DPA-2 `call_graph` = 172 ms vs dense `call` = 41 ms, of which LocalAtten 104 ms + Atten2Map 57 ms + Atten2MultiHeadApply 19 ms — every non-attention graph op (conv, grrg, symmetrization, pair enumeration) is <= 8 ms. With `smooth=False` the total only drops to 151 ms: **the cost is the graph attention structure, not the smooth slot-occupancy softmax** (which adds a secondary +14% eager / +20-29% compiled).

Cross-checks:

- DPA-1 `attn_layer=0`: graph 5.7 / 18.9 ms vs dense 4.7 / 12.5 ms (64 / 216 atoms) -> **1.2-1.5x, graph is fine without attention**.
- DPA-1 `attn_layer=2`: graph 176 ms vs dense 20.5 ms at 64 atoms (**8.6x**), and OOM at 216 atoms on a 15 GB T4 -- a single 5.14 GiB allocation (`P = E*nnei ~ 5.5M` pairs x `attn=128` channels, fp64) in `_graph_attention_one_layer`.
- The carry-all graph **builder is exonerated**: 4.1 / 4.4 / 8.1 / 17.9 ms at 64 / 216 / 512 / 1000 atoms -- near-linear and faster-scaling than the dense route's `extend_coord_with_ghosts` + `build_neighbor_list` (2.1 / 8.5 / 41 / 152 ms).

## Root cause

The ragged pair-axis attention design (NeighborGraph PR-D, "Approach A"): enumerate `P = sum(nnei^2)` center pairs, materialize `(P, channels)` intermediates, and run segment-softmax + `index_add` over them. This is gather/scatter- and allocation-bound (fp64 atomics included), while dense attention is a batched `(nloc, nnei, nnei)` cuBLAS GEMM, compute-bound. It explains the super-linear compiled scaling and both OOMs.

## Proposed fix: block-dense attention on the graph

Gather per-center compact `(nloc, nnei_max, c)` blocks through the already-existing destination CSR (`destination_order` / `destination_row_ptr`), run the DENSE attention math (batched GEMMs, per-row masks for ragged tails), and scatter back to the edge axis. Raggedness then lives only at the gather/scatter boundary; the `P` axis disappears entirely. Applies to:

- DPA-1 `se_atten` graph attention (`_graph_attention_one_layer`),
- DPA-2 `Atten2Map` / `Atten2MultiHeadApply` / `Atten2EquiVarApply` / `LocalAtten` graph twins.

Numerics note: the smooth-attention semantics (signed phantom count + slot-occupancy denominator) are layout-independent and carry over unchanged; per-op parity tests against the current ragged reference pin the transform.

Cheaper complements, independent of the above:

- fused attention kernels (the existing DPA-1 CUDA/Triton kernels deliberately cover only the attention-free factorizable path);
- per-forward caching of the slot-occupancy `theta` (geometry-only and layer-invariant: computed `nlayers x 2` times today, reclaiming most of the +20-29% smooth overhead);
- an fp32 attention path.

## Scope

Follow-up to #5779 (dpa2 graph-native full-stack) and #5715 (PR-D graph attention); performance-only, no semantics change. The dense route and attention-free graph configs are unaffected.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.