Skip to content

fix(cuda): share se_a padding breakpoint #5844

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
njzjz-bot wants to merge 2 commits into
base: master
Choose a base branch
from
Open

fix(cuda): share se_a padding breakpoint #5844

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f647d0c
fix(cuda): share se_a padding breakpoint
Jul 16, 2026
21b73ab
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 76 additions & 55 deletions source/lib/src/gpu/tabulate.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -408,71 +407,93 @@ __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;
}
}
}
Expand Down
106 changes: 106 additions & 0 deletions source/lib/tests/test_tabulate_se_a.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<double> test_em_x = {0.04, 0.07, 0.11, 0.14};
test_em_x.resize(test_nnei, 0.19);
std::vector<double> 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<double> test_dy(4 * last_layer_size);
for (int ii = 0; ii < test_dy.size(); ++ii) {
test_dy[ii] = 0.01 * (ii + 1);
}
std::vector<double> 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<double>* two_embed_host) {
std::vector<double> expected_dy_dem_x(test_nnei);
std::vector<double> expected_dy_dem(test_nnei * 4);
std::vector<double> expected_dy_dtwo(test_nnei * last_layer_size);
deepmd::tabulate_fusion_se_a_grad_cpu<double>(
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<double> actual_dy_dem_x(test_nnei);
std::vector<double> actual_dy_dem(test_nnei * 4);
std::vector<double> 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<double>(
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
Loading