fix(pt): honor is_sorted in SE-Attention tabulation

[njzjz]

Summary

• forward is_sorted through the PyTorch SE-Attention forward helper to the CPU and GPU library entry points
• forward the saved flag through the first-backward helper so all autograd stages use the same sorting contract
• add float32 and float64 regression coverage for unsorted forward, backward, and double-backward behavior

Fixes #5890

Testing

• python -m pytest source/tests/pt/test_tabulate_fusion_se_atten.py -v (12 passed)
• ruff format .
• ruff check .
• pre-commit run clang-format --files source/op/pt/tabulate_multi_device.cc
• dp --version
• dp --pt -h
• Python imports for deepmd and deepmd.pt

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

Summary by CodeRabbit

• Bug Fixes

  • Corrected attention operations to respect whether input data is sorted or unsorted across forward and gradient calculations.
  • Ensured unsorted neighbor data contributes correctly without unintended folding behavior.

• Tests

  • Added coverage for forward, backward, and second-order gradient behavior with sorted and unsorted inputs.
  • Added validation across supported floating-point precisions.

[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: e5e7247a-a7a0-4f42-8c40-b36b3c60d194

📥 Commits

Reviewing files that changed from the base of the PR and between e5fdff0 and dba50c5.

📒 Files selected for processing (2)

• source/op/pt/tabulate_multi_device.cc
• source/tests/pt/test_tabulate_fusion_se_atten.py

---

📝 Walkthrough

Walkthrough

PyTorch SE-Attention tabulation now forwards is_sorted through CPU and GPU forward and first-backward kernels. New parameterized tests verify unsorted forward, gradient, and second-order behavior.

Changes

SE-Attention sorting propagation

Layer / File(s)	Summary
Propagate sorting flag through dispatch source/op/pt/tabulate_multi_device.cc	Forward and gradient helpers now accept and pass is_sorted to CPU/GPU kernels, with autograd callers updated for sorted and attention-specific paths.
Validate unsorted attention behavior source/tests/pt/test_tabulate_fusion_se_atten.py	Parameterized tests compare sorted and unsorted forward outputs, first-order gradients, and second-order gradients.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• deepmodeling/deepmd-kit#5537: Updates related PyTorch SE-Attention autograd sorting propagation.
• deepmodeling/deepmd-kit#5844: Fixes CUDA SE-Attention gradient behavior for sorted-padding cases.
• deepmodeling/deepmd-kit#5904: Adjusts sorted-padding gradient handling in SE-Attention kernels.

Suggested labels: bug, OP, Python

Suggested reviewers: copilot

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: honoring is_sorted in SE-Attention tabulation.
Linked Issues check	✅ Passed	The changes address the reported bug by threading is_sorted through forward and first-backward paths and adding unsorted regression tests.
Out of Scope Changes check	✅ Passed	The added test coverage and helper signature updates are directly related to the linked issue and show no unrelated changes.

✨ 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.}

[Copilot]

Pull request overview

This PR fixes a correctness bug in the PyTorch SE-Attention tabulation custom op where is_sorted was accepted/stored but not forwarded into the CPU/GPU backend entry points for forward and first-backward, causing unsorted neighbor rows to be treated as sorted. It ensures all autograd stages (forward, backward, double-backward) use the same sorting contract, matching expected backend behavior.

Changes:

• Thread is_sorted through the PyTorch SE-Attention forward helper into both CPU and GPU tabulate_fusion_se_a_* entry points.
• Thread is_sorted through the first-backward helper so backward uses the same sorting contract as forward (and aligns with grad-grad’s existing behavior).
• Add float32/float64 regression tests covering unsorted forward, backward, and second-order backward behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
source/op/pt/tabulate_multi_device.cc	Forwards is_sorted through SE-Attention forward and first-backward helper calls into CPU/GPU backend functions and preserves the flag across autograd stages.
source/tests/pt/test_tabulate_fusion_se_atten.py	Adds regression coverage ensuring unsorted neighbor rows affect forward/gradients correctly for both float32 and float64, including double-backward checks.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.64%. Comparing base (e5fdff0) to head (dba50c5).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5905      +/-   ##
==========================================
- Coverage   78.89%   78.64%   -0.26%     
==========================================
  Files        1054     1054              
  Lines      121774   121774              
  Branches     4408     4408              
==========================================
- Hits        96076    95771     -305     
- Misses      24121    24427     +306     
+ Partials     1577     1576       -1     

☔ 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.