[Feature Request] Skip ghost nodes in the fused DPA4C operator on several ranks via a bucketed host-side local-node count #6091

Summary

On several MPI ranks, the fused compressed DPA4C operator still evaluates the descriptor readout, the fitting network and the node backward for the ghost nodes, whose energies the ownership mask then discards. Skipping them would shorten the multi-rank step noticeably. The obstacle is that the operator only knows the number of local nodes as a device tensor. A host-known, bucketed local-node count as part of the CUDA graph key would provide the shape that the skip needs.

Detailed Description

Current behavior. This concerns the compressed DPA4C CUDA path in LAMMPS (pair_style deepmd/kk, canonical graph, model replayed from a captured CUDA graph), which is in development. Its fused energy-force operator sizes the readout, fitting and node-backward kernels by the full node count, local plus ghost nodes. Ghost rows have no edges and their energies are masked out, so this work is wasted. The operator receives the local count only as a device tensor. Using it as a kernel extent would require a device-to-host read every step, and that read cannot happen inside a graph capture.

Estimated gain. This is an estimate from the time share of these kernels, not a measured result: on three ranks with about one million atoms, the step would be about 7% shorter for Nano and 14% to 16% shorter for Plus.

Suggested approach. This follows the batch-size buckets of the CUDA graph runner in DeepSeek's DeepInfer: a graph is captured per bucket, and the host pads the batch to the nearest bucket.

1. LAMMPS already knows nlocal on the host. The pair style passes it through the C API as a host integer.
2. The count is rounded up to a bucket, for example a multiple of 256 or 1,024 nodes, with the padded rows masked. The bucket joins the key of the graph replay.
3. The readout, fitting and node-backward kernels are launched over the bucketed local count instead of all nodes.

The graph then still replays across neighbor-list rebuilds as long as the bucket does not change.

Open points.

• Check whether the multi-rank path currently replays at all, since the model graph keeps its extents only without halo nodes. If it evaluates eagerly, bucketing would also bring replay to multiple ranks.
• The C API signature changes, so this needs a new entry point name and an API version bump.
• The bucket granularity trades padding waste against recapture frequency.