From f647d0cbc3be77cec002b6b1a0b8a36bd4260144 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 20:41:48 +0000 Subject: [PATCH 1/3] fix(cuda): share se_a padding breakpoint Compute the first sorted-padding sentinel once per atom and use uniform tile rounds so CUDA and ROCm gradients match the CPU tail-folding semantics. Add CPU/GPU parity coverage for plain and attention gradients, including the two-embedding tail. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/src/gpu/tabulate.cu | 133 +++++++++++++++---------- source/lib/tests/test_tabulate_se_a.cc | 106 ++++++++++++++++++++ 2 files changed, 184 insertions(+), 55 deletions(-) diff --git a/source/lib/src/gpu/tabulate.cu b/source/lib/src/gpu/tabulate.cu index a3d9cc0e3a..d0f0a23bd2 100644 --- a/source/lib/src/gpu/tabulate.cu +++ b/source/lib/src/gpu/tabulate.cu @@ -398,8 +398,7 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( 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; + __shared__ int breakpoint; 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) { @@ -408,71 +407,95 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( } } __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]; - 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; + + // Sorted padding must be folded at the first sentinel for the whole atom, + // exactly as in the sequential CPU implementation. A warp-local search can + // select several later sentinels because neighbor indices are striped over + // KTILE warps, producing duplicate tail contributions and overlapping + // dy_dtwo writes. + if (thread_idx == 0) { + breakpoint = nnei; // nnei means that no padding sentinel was found. + if (is_sorted) { + const FPTYPE ago = em_x[block_idx * nnei + nnei - 1]; + for (int ii = 0; ii < nnei; ++ii) { + if (ago == em_x[block_idx * nnei + ii] && + 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.) { + breakpoint = ii; + break; + } + } } + } + __syncthreads(); - int table_idx = 0; - FPTYPE reg_em[MTILE] = {em[block_idx * nnei * MTILE + ii * 4 + 0], - em[block_idx * nnei * MTILE + ii * 4 + 1], - em[block_idx * nnei * MTILE + ii * 4 + 2], - em[block_idx * nnei * MTILE + ii * 4 + 3]}; + // Keep the tile loop uniform across the block. GpuSyncThreads is a warp + // barrier on CUDA but a block barrier on ROCm, so warp-specific early exits + // would deadlock HIP when the shared breakpoint falls inside a tile. + for (int tile = 0; tile < nnei && tile <= breakpoint; tile += KTILE) { + const int ii = tile + warp_idx; + const bool active = ii < nnei && ii <= breakpoint; FPTYPE Csub = (FPTYPE)0.; FPTYPE sum[MTILE] = {(FPTYPE)0.}; - FPTYPE extrapolate_delta = (FPTYPE)0.; - locate_xx_se_a(xx, table_idx, lower, upper, max, stride0, stride1, - extrapolate_delta); + if (active) { + const int repeat_count = ii == breakpoint ? nnei - breakpoint : 1; + FPTYPE xx = em_x[block_idx * nnei + ii]; + int table_idx = 0; + FPTYPE reg_em[MTILE] = { + em[block_idx * nnei * MTILE + ii * MTILE + 0], + em[block_idx * nnei * MTILE + ii * MTILE + 1], + em[block_idx * nnei * MTILE + ii * MTILE + 2], + em[block_idx * nnei * MTILE + ii * MTILE + 3]}; + FPTYPE extrapolate_delta = (FPTYPE)0.; + locate_xx_se_a(xx, table_idx, lower, upper, max, stride0, stride1, + extrapolate_delta); - FPTYPE var[6]; - for (int jj = lane_idx; jj < last_layer_size; jj += WARP_SIZE) { - load_polynomial_params(var, table, table_idx, jj, last_layer_size); - FPTYPE res_grad = polynomial5_grad(var, xx); - FPTYPE res = polynomial5(var, xx) + res_grad * extrapolate_delta; - FPTYPE oldres = res; - FPTYPE t; - if (enable_se_atten) { - t = two_embed[block_idx * nnei * last_layer_size + - ii * last_layer_size + jj]; - res = res * t + res; - } + FPTYPE var[6]; + for (int jj = lane_idx; jj < last_layer_size; jj += WARP_SIZE) { + load_polynomial_params(var, table, table_idx, jj, last_layer_size); + FPTYPE res_grad = polynomial5_grad(var, xx); + FPTYPE res = polynomial5(var, xx) + res_grad * extrapolate_delta; + FPTYPE oldres = res; + FPTYPE t; + if (enable_se_atten) { + t = two_embed[block_idx * nnei * last_layer_size + + ii * last_layer_size + jj]; + res = res * t + res; + } - 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 + - jj] = oldres * res; + for (int kk = 0; kk < MTILE; kk++) { + sum[kk] += repeat_count * + 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 += repeat_count * res_grad * + (enable_se_atten ? res * t + res : res); + if (enable_se_atten) { + // A real neighbor owns one entry; the first sentinel owns the full + // padding tail. No other warp writes these dy_dtwo positions. + for (int ii2 = ii; ii2 < ii + repeat_count; ii2++) { + dy_dtwo[block_idx * nnei * last_layer_size + + ii2 * last_layer_size + jj] = oldres * res; + } } } } GpuSyncThreads(); - for (int kk = 0; kk < MTILE; kk++) { - warp_reduce(sum[kk]); - } - warp_reduce(Csub); - if (lane_idx == 0) { + if (active) { for (int kk = 0; kk < MTILE; kk++) { - dy_dem[block_idx * nnei * MTILE + ii * 4 + kk] = sum[kk]; + warp_reduce(sum[kk]); + } + warp_reduce(Csub); + if (lane_idx == 0) { + for (int kk = 0; kk < MTILE; kk++) { + dy_dem[block_idx * nnei * MTILE + ii * MTILE + kk] = sum[kk]; + } + dy_dem_x[block_idx * nnei + ii] = Csub; } - dy_dem_x[block_idx * nnei + ii] = Csub; - } - if (unloop) { - break; } } } diff --git a/source/lib/tests/test_tabulate_se_a.cc b/source/lib/tests/test_tabulate_se_a.cc index 66a77f41fd..1f5601f63c 100644 --- a/source/lib/tests/test_tabulate_se_a.cc +++ b/source/lib/tests/test_tabulate_se_a.cc @@ -856,4 +856,110 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu) { deepmd::delete_device_memory(dy_dev); deepmd::delete_device_memory(two_embed_dev); } + +TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_sorted_padding) { + constexpr int test_nloc = 1; + constexpr int test_nnei = 12; + constexpr int padding_begin = 4; + std::vector test_em_x = {0.04, 0.07, 0.11, 0.14}; + test_em_x.resize(test_nnei, 0.19); + std::vector test_em(test_nnei * 4, 0.0); + for (int ii = 0; ii < padding_begin; ++ii) { + test_em[ii * 4 + 0] = test_em_x[ii]; + test_em[ii * 4 + 1] = 0.1 * (ii + 1); + test_em[ii * 4 + 2] = -0.05 * (ii + 1); + test_em[ii * 4 + 3] = 0.025 * (ii + 1); + } + for (int ii = padding_begin; ii < test_nnei; ++ii) { + test_em[ii * 4] = test_em_x[ii]; + } + std::vector test_dy(4 * last_layer_size); + for (int ii = 0; ii < test_dy.size(); ++ii) { + test_dy[ii] = 0.01 * (ii + 1); + } + std::vector test_two_embed(test_nnei * last_layer_size); + for (int ii = 0; ii < test_two_embed.size(); ++ii) { + test_two_embed[ii] = 0.001 * (ii + 1); + } + + auto compare_cpu_gpu = [&](const std::vector* two_embed_host) { + std::vector expected_dy_dem_x(test_nnei); + std::vector expected_dy_dem(test_nnei * 4); + std::vector expected_dy_dtwo(test_nnei * last_layer_size); + deepmd::tabulate_fusion_se_a_grad_cpu( + expected_dy_dem_x.data(), expected_dy_dem.data(), + two_embed_host == nullptr ? nullptr : expected_dy_dtwo.data(), + table.data(), info.data(), test_em_x.data(), test_em.data(), + two_embed_host == nullptr ? nullptr : two_embed_host->data(), + test_dy.data(), test_nloc, test_nnei, last_layer_size, true); + + std::vector actual_dy_dem_x(test_nnei); + std::vector actual_dy_dem(test_nnei * 4); + std::vector actual_dy_dtwo(test_nnei * last_layer_size); + double *dy_dem_x_dev = nullptr, *dy_dem_dev = nullptr, + *dy_dtwo_dev = nullptr, *table_dev = nullptr, *em_x_dev = nullptr, + *em_dev = nullptr, *two_embed_dev = nullptr, *dy_dev = nullptr; + deepmd::malloc_device_memory_sync(dy_dem_x_dev, actual_dy_dem_x); + deepmd::malloc_device_memory_sync(dy_dem_dev, actual_dy_dem); + deepmd::malloc_device_memory_sync(table_dev, table); + deepmd::malloc_device_memory_sync(em_x_dev, test_em_x); + deepmd::malloc_device_memory_sync(em_dev, test_em); + deepmd::malloc_device_memory_sync(dy_dev, test_dy); + if (two_embed_host != nullptr) { + deepmd::malloc_device_memory_sync(dy_dtwo_dev, actual_dy_dtwo); + deepmd::malloc_device_memory_sync(two_embed_dev, *two_embed_host); + } + + deepmd::tabulate_fusion_se_a_grad_gpu( + dy_dem_x_dev, dy_dem_dev, dy_dtwo_dev, table_dev, info.data(), + em_x_dev, em_dev, two_embed_dev, dy_dev, test_nloc, test_nnei, + last_layer_size, true); + deepmd::memcpy_device_to_host(dy_dem_x_dev, actual_dy_dem_x); + deepmd::memcpy_device_to_host(dy_dem_dev, actual_dy_dem); + if (two_embed_host != nullptr) { + deepmd::memcpy_device_to_host(dy_dtwo_dev, actual_dy_dtwo); + } + + for (int ii = 0; ii < actual_dy_dem_x.size(); ++ii) { + EXPECT_NEAR(actual_dy_dem_x[ii], expected_dy_dem_x[ii], 1e-10); + } + for (int ii = 0; ii < actual_dy_dem.size(); ++ii) { + EXPECT_NEAR(actual_dy_dem[ii], expected_dy_dem[ii], 1e-10); + } + if (two_embed_host != nullptr) { + for (int ii = 0; ii < actual_dy_dtwo.size(); ++ii) { + EXPECT_NEAR(actual_dy_dtwo[ii], expected_dy_dtwo[ii], 1e-10); + } + // The CPU contract copies the first sentinel's two-embedding gradient + // across the complete sorted-padding tail. + for (int ii = padding_begin + 1; ii < test_nnei; ++ii) { + for (int jj = 0; jj < last_layer_size; ++jj) { + EXPECT_NEAR(actual_dy_dtwo[ii * last_layer_size + jj], + actual_dy_dtwo[padding_begin * last_layer_size + jj], + 1e-10); + } + } + } + // Later sentinels are padding, not independent neighbors owned by other + // warps, and therefore must retain zero descriptor gradients. + for (int ii = padding_begin + 1; ii < test_nnei; ++ii) { + EXPECT_DOUBLE_EQ(actual_dy_dem_x[ii], 0.0); + for (int jj = 0; jj < 4; ++jj) { + EXPECT_DOUBLE_EQ(actual_dy_dem[ii * 4 + jj], 0.0); + } + } + + deepmd::delete_device_memory(dy_dem_x_dev); + deepmd::delete_device_memory(dy_dem_dev); + deepmd::delete_device_memory(dy_dtwo_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(two_embed_dev); + deepmd::delete_device_memory(dy_dev); + }; + + compare_cpu_gpu(nullptr); + compare_cpu_gpu(&test_two_embed); +} #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM From 21b73ab303581f467bba0cae101a880c35cc2d82 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:48:26 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/lib/src/gpu/tabulate.cu | 20 +++++++++----------- source/lib/tests/test_tabulate_se_a.cc | 6 +++--- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/source/lib/src/gpu/tabulate.cu b/source/lib/src/gpu/tabulate.cu index d0f0a23bd2..d978b5643c 100644 --- a/source/lib/src/gpu/tabulate.cu +++ b/source/lib/src/gpu/tabulate.cu @@ -442,11 +442,10 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( const int repeat_count = ii == breakpoint ? nnei - breakpoint : 1; FPTYPE xx = em_x[block_idx * nnei + ii]; int table_idx = 0; - FPTYPE reg_em[MTILE] = { - em[block_idx * nnei * MTILE + ii * MTILE + 0], - em[block_idx * nnei * MTILE + ii * MTILE + 1], - em[block_idx * nnei * MTILE + ii * MTILE + 2], - em[block_idx * nnei * MTILE + ii * MTILE + 3]}; + FPTYPE reg_em[MTILE] = {em[block_idx * nnei * MTILE + ii * MTILE + 0], + em[block_idx * nnei * MTILE + ii * MTILE + 1], + em[block_idx * nnei * MTILE + ii * MTILE + 2], + em[block_idx * nnei * MTILE + ii * MTILE + 3]}; FPTYPE extrapolate_delta = (FPTYPE)0.; locate_xx_se_a(xx, table_idx, lower, upper, max, stride0, stride1, extrapolate_delta); @@ -465,21 +464,20 @@ __global__ void tabulate_fusion_se_a_grad_fifth_order_polynomial( } for (int kk = 0; kk < MTILE; kk++) { - sum[kk] += repeat_count * - iteratorA[kk * last_layer_size + jj] * res; + sum[kk] += repeat_count * 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 += repeat_count * res_grad * - (enable_se_atten ? res * t + res : res); + Csub += + repeat_count * res_grad * (enable_se_atten ? res * t + res : res); if (enable_se_atten) { // A real neighbor owns one entry; the first sentinel owns the full // padding tail. No other warp writes these dy_dtwo positions. for (int ii2 = ii; ii2 < ii + repeat_count; ii2++) { - dy_dtwo[block_idx * nnei * last_layer_size + - ii2 * last_layer_size + jj] = oldres * res; + dy_dtwo[block_idx * nnei * last_layer_size + ii2 * last_layer_size + + jj] = oldres * res; } } } diff --git a/source/lib/tests/test_tabulate_se_a.cc b/source/lib/tests/test_tabulate_se_a.cc index 1f5601f63c..2b5b8faa00 100644 --- a/source/lib/tests/test_tabulate_se_a.cc +++ b/source/lib/tests/test_tabulate_se_a.cc @@ -911,9 +911,9 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_sorted_padding) { } deepmd::tabulate_fusion_se_a_grad_gpu( - dy_dem_x_dev, dy_dem_dev, dy_dtwo_dev, table_dev, info.data(), - em_x_dev, em_dev, two_embed_dev, dy_dev, test_nloc, test_nnei, - last_layer_size, true); + dy_dem_x_dev, dy_dem_dev, dy_dtwo_dev, table_dev, info.data(), em_x_dev, + em_dev, two_embed_dev, dy_dev, test_nloc, test_nnei, last_layer_size, + true); deepmd::memcpy_device_to_host(dy_dem_x_dev, actual_dy_dem_x); deepmd::memcpy_device_to_host(dy_dem_dev, actual_dy_dem); if (two_embed_host != nullptr) { From f9c8a8af90e995a86e145f5ee597748ca5115bb1 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 23 Jul 2026 20:06:15 +0800 Subject: [PATCH 3/3] test(cuda): cover unsorted se_a tabulation Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/lib/tests/test_tabulate_se_a.cc | 37 +++++++++++++++----------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/source/lib/tests/test_tabulate_se_a.cc b/source/lib/tests/test_tabulate_se_a.cc index 2b5b8faa00..966909a60e 100644 --- a/source/lib/tests/test_tabulate_se_a.cc +++ b/source/lib/tests/test_tabulate_se_a.cc @@ -857,7 +857,7 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu) { deepmd::delete_device_memory(two_embed_dev); } -TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_sorted_padding) { +TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_padding_modes) { constexpr int test_nloc = 1; constexpr int test_nnei = 12; constexpr int padding_begin = 4; @@ -882,7 +882,8 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_sorted_padding) { test_two_embed[ii] = 0.001 * (ii + 1); } - auto compare_cpu_gpu = [&](const std::vector* two_embed_host) { + auto compare_cpu_gpu = [&](const std::vector* two_embed_host, + const bool is_sorted) { std::vector expected_dy_dem_x(test_nnei); std::vector expected_dy_dem(test_nnei * 4); std::vector expected_dy_dtwo(test_nnei * last_layer_size); @@ -891,7 +892,7 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_sorted_padding) { two_embed_host == nullptr ? nullptr : expected_dy_dtwo.data(), table.data(), info.data(), test_em_x.data(), test_em.data(), two_embed_host == nullptr ? nullptr : two_embed_host->data(), - test_dy.data(), test_nloc, test_nnei, last_layer_size, true); + test_dy.data(), test_nloc, test_nnei, last_layer_size, is_sorted); std::vector actual_dy_dem_x(test_nnei); std::vector actual_dy_dem(test_nnei * 4); @@ -913,7 +914,7 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_sorted_padding) { deepmd::tabulate_fusion_se_a_grad_gpu( dy_dem_x_dev, dy_dem_dev, dy_dtwo_dev, table_dev, info.data(), em_x_dev, em_dev, two_embed_dev, dy_dev, test_nloc, test_nnei, last_layer_size, - true); + is_sorted); deepmd::memcpy_device_to_host(dy_dem_x_dev, actual_dy_dem_x); deepmd::memcpy_device_to_host(dy_dem_dev, actual_dy_dem); if (two_embed_host != nullptr) { @@ -932,20 +933,24 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_sorted_padding) { } // The CPU contract copies the first sentinel's two-embedding gradient // across the complete sorted-padding tail. - for (int ii = padding_begin + 1; ii < test_nnei; ++ii) { - for (int jj = 0; jj < last_layer_size; ++jj) { - EXPECT_NEAR(actual_dy_dtwo[ii * last_layer_size + jj], - actual_dy_dtwo[padding_begin * last_layer_size + jj], - 1e-10); + if (is_sorted) { + for (int ii = padding_begin + 1; ii < test_nnei; ++ii) { + for (int jj = 0; jj < last_layer_size; ++jj) { + EXPECT_NEAR(actual_dy_dtwo[ii * last_layer_size + jj], + actual_dy_dtwo[padding_begin * last_layer_size + jj], + 1e-10); + } } } } // Later sentinels are padding, not independent neighbors owned by other // warps, and therefore must retain zero descriptor gradients. - for (int ii = padding_begin + 1; ii < test_nnei; ++ii) { - EXPECT_DOUBLE_EQ(actual_dy_dem_x[ii], 0.0); - for (int jj = 0; jj < 4; ++jj) { - EXPECT_DOUBLE_EQ(actual_dy_dem[ii * 4 + jj], 0.0); + if (is_sorted) { + for (int ii = padding_begin + 1; ii < test_nnei; ++ii) { + EXPECT_DOUBLE_EQ(actual_dy_dem_x[ii], 0.0); + for (int jj = 0; jj < 4; ++jj) { + EXPECT_DOUBLE_EQ(actual_dy_dem[ii * 4 + jj], 0.0); + } } } @@ -959,7 +964,9 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_sorted_padding) { deepmd::delete_device_memory(dy_dev); }; - compare_cpu_gpu(nullptr); - compare_cpu_gpu(&test_two_embed); + for (const bool is_sorted : {false, true}) { + compare_cpu_gpu(nullptr, is_sorted); + compare_cpu_gpu(&test_two_embed, is_sorted); + } } #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM