fix(cuda): share se_a padding breakpoint

[njzjz-bot]

Summary

• find the first sorted-padding sentinel once per atom block and share it across all warps
• process neighbors in uniform tile rounds so HIP's block-wide synchronization remains safe
• fold only the first sentinel by the full padding-tail length, matching the CPU implementation
• give the breakpoint warp sole ownership of the attention dy_dtwo tail
• add direct CPU/GPU parity coverage for both plain se_a and attention gradients

Why existing tests missed the bug

• The existing C++ fixture uses nnei = 4 without a padding sentinel, so each of the four warps handles one real neighbor and never chooses competing breakpoints.
• Existing CUDA attention tests do not copy dy_dtwo back to the host or assert it, leaving overlapping tail writes invisible.
• PyTorch fixtures also use four real neighbors without repeated sorted padding and do not assert the two-embedding gradient.
• The new fixture uses 12 neighbors with repeated padding beginning at index 4, forcing the old four warp-local scans to select different sentinels.

Validation

• cmake --build source/build --target deepmd_op_cuda runUnitTests_lib -j2
• source/build/lib/tests/runUnitTests_lib --gtest_filter='TestTabulateSeA.tabulate_fusion_se_a_grad_cpu' (passed)
• the new CPU/GPU parity regression compiles for both null and non-null two_embed paths
• local GPU execution was attempted but the CUDA 12.4 build lacks an sm_120 kernel image for the RTX 5090; CUDA CI will execute the regression
• ruff format .
• ruff check .

Fixes #5650

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

Summary by CodeRabbit

• Bug Fixes

  • Improved GPU gradient calculations for datasets containing sorted padding.
  • Prevented overlapping gradient writes and ensured correct handling of padding entries.
  • Maintained consistent behavior between CPU and GPU calculations.

• Tests

  • Added GPU coverage for sorted-padding gradient calculations, including descriptor and embedding gradients.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 93be4e85-5579-49ce-b0da-e5d37fb443c9

📥 Commits

Reviewing files that changed from the base of the PR and between 6c3b985 and 21b73ab.

📒 Files selected for processing (2)

• source/lib/src/gpu/tabulate.cu
• source/lib/tests/test_tabulate_se_a.cc

---

📝 Walkthrough

Walkthrough

The CUDA se_a gradient kernel now shares sorted-padding breakpoints across the block and uses uniform accumulation. A GPU regression test compares sorted-padding results with CPU gradients, including optional two_embed behavior and padding-tail expectations.

Changes

Sorted-padding gradient handling

Layer / File(s)	Summary
Shared breakpoint and uniform accumulation source/lib/src/gpu/tabulate.cu	The kernel computes one synchronized sorted-padding breakpoint per block, gates inactive iterations, applies repeat counts, and constrains dy_dtwo writes to avoid overlap.
CPU/GPU sorted-padding regression source/lib/tests/test_tabulate_se_a.cc	A GPU test compares descriptor and optional two-embedding gradients with CPU results and verifies padding-tail copying and zero descriptor gradients.

Estimated code review effort: 4 (Complex) | ~45 minutes

Suggested labels: CUDA

Suggested reviewers: njzjz

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: sharing the CUDA se_a padding breakpoint.
Linked Issues check	✅ Passed	The kernel now shares one first padding breakpoint per block and the new test covers CPU/GPU parity with repeated padding and two_embed paths.
Out of Scope Changes check	✅ Passed	The changes stay within the issue scope, affecting only the se_a CUDA kernel and its regression test.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.32%. Comparing base (6c3b985) to head (21b73ab).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5844      +/-   ##
==========================================
- Coverage   78.58%   78.32%   -0.27%     
==========================================
  Files        1050     1050              
  Lines      120637   120637              
  Branches     4356     4358       +2     
==========================================
- Hits        94801    94485     -316     
- Misses      24278    24585     +307     
- Partials     1558     1567       +9     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[njzjz]

Possible reviewers based on changed lines, exact file history, and exact-file review history:

• @wanghan-iapcm — 2 commits on changed files (source/lib/src/gpu/tabulate.cu, source/lib/tests/test_tabulate_se_a.cc).
• @iProzd — 1 commits on changed files (source/lib/tests/test_tabulate_se_a.cc).

No review request was made automatically.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh