CUDA Kernel Compatibility#

Overview#

DeePMD-kit's TensorFlow backend contains CUDA kernels in source/lib/src/gpu/ that implement GPU-accelerated neighbor-list construction and descriptor computation. A latent synchronization bug in the parallel_prefix_scan kernel causes GPU deadlocks on Blackwell-architecture GPUs (RTX 5090/5090D, SM 12.0). The same invalid CUDA pattern was present on older architectures (Ada/Ampere/Volta) but did not reliably manifest due to differences in code generation and barrier scheduling.

The Bug: Divergent Block-Wide Collective in parallel_prefix_scan#

File: source/lib/src/gpu/neighbor_list.cu

The parallel_prefix_scan kernel computes an exclusive prefix sum over the neighbor-list data for each atom block. It uses cub::BlockScan<int, THREADS_PER_BLOCK> (a CUB block-wide collective) together with __syncthreads() inside a per-thread loop:

for (int ii = threadIdx.x; ii < nall; ii += THREADS_PER_BLOCK) {
    ...
    BlockScan(temp_storage).ExclusiveSum(o_data, o_data, prefix_op);
    __syncthreads();
    ...
}

THREADS_PER_BLOCK is defined as TPB = 256 in source/lib/include/device.h. When nall is not a multiple of 256, the final loop iteration is entered only by threads where threadIdx.x < nall % 256. However, BlockScan and __syncthreads() are block-wide collectives — CUDA requires all non-exited threads in the block to participate. Calling them from only a subset of threads is undefined behavior and can cause the GPU to deadlock 1.

Why It Shows Up on Blackwell (SM 12.0) First#

The bug is architectural undefined behavior, not a Blackwell-specific regression. On pre-Blackwell GPUs (Volta/Ampere/Ada), the same invalid pattern appeared to work by accident due to differences in JIT code generation, CUB implementation internals, driver-level barrier scheduling, or warp-level timing. On RTX 5090 / SM 12.0, it reliably manifests as an infinite GPU spin at 100% utilization inside build_nlist_gpu 1.

Trigger Condition#

The deadlock is triggered when PBC (periodic boundary condition) coordinate expansion pushes nall past a 256-thread block boundary — i.e., when nall > 256 and nall % 256 != 0. The Hybrid Descriptor calls compute_input_stats() for each sub-descriptor with PBC-expanded coordinates, making it consistently hit this condition on non-trivial systems 2.

The Fix (PR #5575)#

PR #5575 restructures the loop so that every thread participates in every segment iteration, loading a sentinel value (-1) and skipping the output write for out-of-bounds indices:

for (int base = 0; base < nall; base += THREADS_PER_BLOCK) {
    int ii = base + threadIdx.x;
    int i_data = ii < nall ? temp_nlist[block_offset + ii] : -1;
    ...
    BlockScan(temp_storage).ExclusiveSum(o_data, o_data, prefix_op);
    __syncthreads();
    if (ii < nall && i_data != -1) {
        nei_order[block_offset + ii] = o_data;
    }
    ...
}

The fix was validated on RTX 5090 hardware with CUDA 12.9, TensorFlow 2.19.1, and CMAKE_CUDA_ARCHITECTURES=120. Training on the previously hanging input completed successfully 3. A regression test with nall = TPB + 68 was added to cover the tail-segment case 3.

Symptoms#

• Training hangs indefinitely at DEEPMD INFO data stating... (this step may take long time)
• GPU utilization stays at 100% with no forward progress
• py-spy / gdb shows the main thread blocked in Session.run() → ProdEnvMatAOp → build_nlist_gpu → cudaLaunchKernel 1
• Affects TensorFlow backend only; PyTorch backend is not affected
• Affects Hybrid Descriptor reliably; individual descriptors (se_e2_a, se_e2_r, se_e3) can work if their nall happens to be a multiple of 256 2

Workarounds (Pre-Fix)#

1. Use the PyTorch backend: Switch dp --tf train to dp --pt train. The PyTorch neighbor-list code path does not use this kernel 1.
2. Run on a non-Blackwell GPU: V100, A100, H100, RTX 4090 all avoid the deadlock in practice 2.
3. Build from source with the patch from PR #5575 applied, using CMAKE_CUDA_ARCHITECTURES=120.

General Rule: Block-Wide Collectives Require Full Participation#

This bug exemplifies a broader CUDA correctness rule: all non-exited threads in a block must call block-wide collectives (__syncthreads(), cub::BlockScan, cub::BlockReduce, etc.) on every iteration. Wrapping these calls inside a thread-indexed loop where only some threads reach the synchronization boundary is undefined behavior. The canonical fix is to iterate by segment base (not by threadIdx.x) and guard only the memory-access and output paths — exactly the pattern used in PR #5575.

Key References#