This issue comes from a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

The CUDA se_a tabulation gradient kernel folds sorted padding independently in each warp, so it can disagree with the CPU implementation when a sorted neighbor row has repeated padding entries.

The CPU gradient scans neighbors sequentially. Once it finds the sorted-padding sentinel, it folds the full tail into that first sentinel position:

• deepmd-kit/source/lib/src/tabulate.cc

  Lines 266 to 277 in 73de44b

  	FPTYPE ago = em_x[ii * nnei + nnei - 1];
  	bool unloop = false;
  	for (int jj = 0; jj < nnei; jj++) {
  	// construct the dy/dx
  	ll[0] = em[ii * nnei * 4 + jj * 4 + 0];
  	ll[1] = em[ii * nnei * 4 + jj * 4 + 1];
  	ll[2] = em[ii * nnei * 4 + jj * 4 + 2];
  	ll[3] = em[ii * nnei * 4 + jj * 4 + 3];
  	FPTYPE xx = em_x[ii * nnei + jj];
  	if (ago == xx && ll[1] == 0. && ll[2] == 0. && ll[3] == 0. && is_sorted) {
  	unloop = true;
  	}

• deepmd-kit/source/lib/src/tabulate.cc

  Lines 305 to 318 in 73de44b

  	FPTYPE dotllrr = dot(ll, rr);
  	if (unloop) {
  	grad += g * dotllrr * (nnei - jj);
  	dy_dem[ii * nnei * 4 + jj * 4 + 0] += res * rr[0] * (nnei - jj);
  	dy_dem[ii * nnei * 4 + jj * 4 + 1] += res * rr[1] * (nnei - jj);
  	dy_dem[ii * nnei * 4 + jj * 4 + 2] += res * rr[2] * (nnei - jj);
  	dy_dem[ii * nnei * 4 + jj * 4 + 3] += res * rr[3] * (nnei - jj);
  	if (enable_se_atten) {
  	// fill from jj to nnei
  	for (int jj2 = jj; jj2 < nnei; jj2++) {
  	dy_dtwo[ii * nnei * last_layer_size + jj2 * last_layer_size +
  	kk] += resold * dotllrr;
  	}
  	}

The CUDA gradient stripes neighbor indices by warp:

• deepmd-kit/source/lib/src/gpu/tabulate.cu

  Lines 397 to 413 in 73de44b

  	const int_64 block_idx = blockIdx.x; // nloc
  	const int thread_idx = threadIdx.x; // KTILE * WARP_SIZE, usually 128 here~
  	int warp_idx = GpuShuffleSync(0xffffffff, threadIdx.x / WARP_SIZE, 0);
  	int lane_idx = threadIdx.x % WARP_SIZE;
  	int breakpoint = nnei - 1;
  	bool unloop = false;
  	FPTYPE* iteratorA = (FPTYPE*)&_data[0]; // dy
  	for (int ii = 0; ii < MTILE; ii++) {
  	for (int jj = thread_idx; jj < last_layer_size; jj += blockDim.x) {
  	iteratorA[ii * last_layer_size + jj] =
  	dy[block_idx * MTILE * last_layer_size + ii * last_layer_size + jj];
  	}
  	}
  	__syncthreads();
  	FPTYPE ago = GpuShuffleSync(0xffffffff, em_x[block_idx * nnei + nnei - 1], 0);
  	for (int ii = warp_idx; ii < nnei; ii += KTILE) {
  	FPTYPE xx = em_x[block_idx * nnei + ii];

Each warp keeps its own breakpoint, and any warp that sees a padding sentinel applies its own tail multiplier:

• deepmd-kit/source/lib/src/gpu/tabulate.cu

  Lines 411 to 419 in 73de44b

  	FPTYPE ago = GpuShuffleSync(0xffffffff, em_x[block_idx * nnei + nnei - 1], 0);
  	for (int ii = warp_idx; ii < nnei; ii += KTILE) {
  	FPTYPE xx = em_x[block_idx * nnei + ii];
  	if (ago == xx && em[block_idx * nnei * 4 + ii * 4 + 1] == 0. &&
  	em[block_idx * nnei * 4 + ii * 4 + 2] == 0. &&
  	em[block_idx * nnei * 4 + ii * 4 + 3] == 0. && is_sorted) {
  	unloop = true;
  	breakpoint = ii;
  	}

• deepmd-kit/source/lib/src/gpu/tabulate.cu

  Lines 445 to 458 in 73de44b

  	for (int kk = 0; kk < MTILE; kk++) {
  	sum[kk] +=
  	(nnei - breakpoint) * iteratorA[kk * last_layer_size + jj] * res;
  	}
  	res = reg_em[0] * iteratorA[0 * last_layer_size + jj];
  	res += reg_em[1] * iteratorA[1 * last_layer_size + jj];
  	res += reg_em[2] * iteratorA[2 * last_layer_size + jj];
  	res += reg_em[3] * iteratorA[3 * last_layer_size + jj];
  	Csub += (nnei - breakpoint) * res_grad *
  	(enable_se_atten ? res * t + res : res);
  	if (enable_se_atten) {
  	// from ii to ii + (nnei - breakpoint)
  	for (int ii2 = ii; ii2 < ii + nnei - breakpoint; ii2++) {
  	dy_dtwo[block_idx * nnei * last_layer_size + ii2 * last_layer_size +

For example, if the first padding sentinel is at neighbor index p, CPU folds all entries [p, nnei) into p. CUDA can have a different warp reach a later sentinel first in its strided sequence and fold only [q, nnei) into q, leaving different dy_dem_x, dy_dem, and dy_dtwo values.

Impact

GPU training/inference gradients for compressed se_a models can silently diverge from CPU results when is_sorted=true and the neighbor list contains sorted padding. The same kernel is reached from both TensorFlow and PyTorch tabulation op wrappers.

Suggested fix

Compute the first padding breakpoint once per atom/block and share it across all warps before the gradient accumulation, or otherwise mirror the CPU sequential folding semantics. Add a CPU-vs-GPU regression for tabulate_fusion_se_a_grad with is_sorted=true, real neighbors followed by repeated padding sentinel values, and both two_embed == nullptr and non-null two_embed paths.