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

[wanghan-iapcm]

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.

[wanghan-iapcm]

Graph-native alternatives (preferred over the block-dense gather)

The block-dense gather above is one implementation; there are fully graph-native fixes that avoid padding entirely:

1. Online-softmax chunked attention in pure array API (all dtypes, export-safe): stream over neighbor slots in bounded tiles with a running max/denominator (the FlashAttention recurrence in xp ops). Never materializes the (P, channels) intermediates: memory drops from O(P·c) to O(E·c), eliminating both observed OOMs and the allocation storm; compute remains scalar-op based.

2. Fused varlen segment attention (the target): per-center attention over the CSR is variable-length flash attention — destination_row_ptr is cu_seqlens. FlashAttention-varlen / torch FlexAttention with a block mask computes softmax(QK^T)V per segment in tiles: O(E) memory, compute-bound, no padding, no P axis. LocalAtten maps directly (queries = centers, keys = edges); Atten2Map is per-center self-attention among the center's edges — also a standard varlen pattern.

The smooth slot-occupancy semantics compose cleanly with fused kernels: the occupancy-modified weights factor as w_j = softmax_j * (sum_j e_j / D_tilde), i.e. a per-segment scalar rescale of the vanilla softmax. The fused kernel computes the standard aggregation; D_tilde (phantom terms, theta, chi) is computed separately by the existing O(E) segment ops — which profiling shows are cheap (the native graph-attention cost is the pair machinery, not the softmax math). All smoothness/parity properties live in the rescale factor, untouched.

Constraints: flash-class kernels are fp32/bf16 only and not array-API. Both point at the dispatch architecture already established for the dpa1 factorizable path (_fused_eligible + DP_TRITON_INFER levels): the array-API online-softmax version becomes the reference implementation (fp64 / numpy / jax / export), with the varlen kernel as the opt-in CUDA fast path.

Suggested sequencing: (a) theta per-forward caching (trivial), (b) online-softmax reference rewrite of the attention ops (kills the OOMs and the P axis everywhere), (c) fused varlen kernel behind the existing fused-dispatch gates.