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 handling for inputs that include sorted padding sentinels.
  • Prevented overlapping gradient writes during the fusion backward pass and ensured correct update ranges.
  • Kept CPU and GPU results consistent for padding-related gradient behavior.

• Tests

  • Added GPU test coverage for sorted-padding gradient calculations, validating embedding/descriptor gradients and correct dy_dtwo propagation.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: c4713ac6-cabe-4a7f-8e04-416d9819968c

📥 Commits

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

📒 Files selected for processing (1)

• 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

Possibly related issues

• #5891: The changes modify sorted-padding two_embed/dy_dtwo handling and add related padding-gradient coverage in the same GPU backward kernel.

Suggested labels: CUDA

Suggested reviewers: wanghan-iapcm

🚥 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 is concise and accurately summarizes the main CUDA se_a padding breakpoint fix.
Linked Issues check	✅ Passed	The kernel change shares the first sorted-padding breakpoint and the new regression test covers the repeated-padding cases with and without two_embed.
Out of Scope Changes check	✅ Passed	The code changes stay focused on the reported se_a gradient bug and its parity test coverage.
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

[wanghan-iapcm]

Nice fix — the shared breakpoint restores the CPU tail-folding contract and, as a bonus, removes the divergent-break that could make warps hit the ROCm block-wide GpuSyncThreads() a different number of times. The new parity test is a genuine regression test for the striped-breakpoint case (pre-fix, warps 0-3 fold counts 8/7/6/5 and overwrite each other's dy_dtwo). One minor coverage note inline.