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
• require a sorted lower neighbor list while DPA1 geometric compression is active, matching the tabulation fold invariant
• add float32 and float64 regression coverage for unsorted forward, backward, and double-backward behavior, including a nonzero table derivative
• add compressed DPA1 forward_lower coverage for unsorted rcut+skin lists and an end-to-end plain se_atten case with exclude_types and set_davg_zero: true
• document why the non-attention se_a path intentionally keeps the sorted fold enabled

Fixes #5890

Historical origin

• The PyTorch custom-op forwarding bug was introduced by pt: add 4 tabulate_fusion op #3877, which added tabulate_fusion_se_atten while omitting the backend is_sorted argument in forward and first backward.
• The compressed DPA1 lower-nlist exposure became reachable in Add 4 pt descriptor compression #4227, which added PyTorch se_atten geometric compression while need_sorted_nlist_for_lower() still returned False. fix(pt): fix lammps nlist sort with large sel #3993 introduced that API/default earlier, but compression did not yet use the sorted fold.

Performance note

With non-empty exclude_types, the corrected is_sorted=False path traverses the full neighbor list instead of applying the sorted-padding fold. This may be measurably slower, but avoids silently folding interleaved excluded rows.

Testing

• CPU-only PyTorch deepmd_op_pt C++ target build
• python -m pytest source/tests/pt/test_tabulate_fusion_se_atten.py -v (12 passed)
• python -m pytest source/tests/pt/model/test_compressed_se_atten_forward_lower.py -v (1 passed)
• python -m pytest source/tests/pt/test_model_compression_se_atten.py::TestDeepPotATPBCExcludeTypes::test_1frame -v (1 passed)
• ruff format .
• ruff check on all modified Python files
• commit hooks, including Ruff, pylint, and clang-format
• dp --version
• dp --pt -h
• Python imports for deepmd and deepmd.pt

ruff check . was also run; it reports five pre-existing findings in deepmd/jax/jax_md/__init__.py and deepmd/tf/entrypoints/__init__.py, outside this PR.

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

[coderabbitai]

📝 Walkthrough

Walkthrough

PyTorch SE-Attention tabulation now propagates is_sorted through CPU/GPU forward and first-backward kernels. Compressed lower evaluation requests sorting, with new regression tests for unsorted neighbors and operator gradients.

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 pass is_sorted to CPU/GPU kernels, with autograd callers updated for standard and attention-specific paths.
Enable compressed lower sorting deepmd/pt/model/descriptor/se_atten.py, source/tests/pt/model/test_compressed_se_atten_forward_lower.py, source/tests/pt/test_model_compression_se_atten.py	Geometric compression requests sorted neighbor-list handling, and regression coverage exercises overcut unsorted lists and the updated compression configuration.
Validate unsorted operator behavior source/tests/pt/test_tabulate_fusion_se_atten.py	Parameterized tests compare sorted and unsorted forward results, first-order gradients, and second-order gradients.

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

Possibly related PRs

• deepmodeling/deepmd-kit#5537: Modifies the related PyTorch SE-Attention autograd path.
• deepmodeling/deepmd-kit#5844: Addresses sorted-padding behavior in SE-Attention gradients.
• deepmodeling/deepmd-kit#5904: Updates related sorted-padding gradient handling.

Suggested labels: bug, C++

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
Linked Issues check	✅ Passed	The PR fixes forward and first-backward is_sorted handling and adds regression coverage for unsorted compressed models.
Out of Scope Changes check	✅ Passed	The descriptor and model-compression changes support the unsorted-path fix and test coverage, with no unrelated scope apparent.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: propagating and honoring is_sorted in SE-Attention tabulation.

✨ 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.92%. Comparing base (e5fdff0) to head (8a3a023).
⚠️ Report is 15 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5905      +/-   ##
==========================================
+ Coverage   78.89%   78.92%   +0.02%     
==========================================
  Files        1054     1068      +14     
  Lines      121774   123852    +2078     
  Branches     4408     4522     +114     
==========================================
+ Hits        96076    97747    +1671     
- Misses      24121    24488     +367     
- Partials     1577     1617      +40     

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

[wanghan-iapcm]

The plumbing is right. I walked all four autograd classes and every library entry point at head: is_sorted now reaches forward, first backward and grad-grad with the user's value in each; ctx->saved_data["is_sorted"] has exactly two stores and two matching reads, so no backward can read a key its forward never wrote; and since TORCH_LIBRARY_FRAGMENT registers a function pointer, the schema is inferred from the signature and cannot drift out of sync. The public signature is unchanged, so the TorchScript stub and the pt_expt fake kernel stay valid.

Worth stating explicitly what the pre-fix state was, because it is a bit worse than "a flag was ignored": the grad-grad helper already forwarded is_sorted correctly, while forward and first-backward did not. So the three autograd stages disagreed with each other -- the op was not consistently wrong, it was internally incoherent, which is why the second-order finite-difference check in the new test is such a sharp discriminator.

The hard-coded true on the non-attention path is correct rather than merely unchanged. is_sorted defaults to true in source/lib/include/tabulate.h, so behaviour is preserved exactly; and se_a applies exclusion per type-pair block with a separate op call per embedding index, so an excluded block is uniformly zero within one invocation and never interleaves. TF was already correct (it declares is_sorted as an op attr and reads it in every kernel), Paddle has no tabulate op, and pt_expt calls this same op so it is fixed transitively.

On why no existing test caught this. This is the part I found genuinely interesting, because there is already an end-to-end test that trains, freezes and compresses a se_atten model with exclude_types and compares against the uncompressed model at 1e-10 -- TestDeepPotATPBCExcludeTypes in source/tests/pt/test_model_compression_se_atten.py. It passes pre-fix, and not because of tolerance or because the op is bypassed. It passes because in that model the fold is never armed.

forward builds the env-mat from the unmasked nlist and applies the exclusion mask afterwards, so an excluded row is exactly all-zero including its em_x. Padding rows are not zero: prod_env_mat normalises without re-masking, so a padding row becomes (-avg_r/std_r, 0, 0, 0). The kernel's guard needs ago == xx, where ago is the last neighbour's scalar -- and se_atten_v2 defaults set_davg_zero=False, so that sentinel is non-zero while excluded rows are zero. The guard never matches an excluded row and fires only on the genuine trailing padding block, where the (nnei - jj) fold is exact. The nlist ordering is actually maximally adversarial there (the excluded O-H pairs are the nearest neighbours, sitting ahead of real ones), and it still cannot trigger.

So the untested cell is compressed se_atten, non-empty exclude_types, and an env-mat where the excluded row's em_x equals the padding sentinel -- which is what set_davg_zero=True gives you, and that is the default for plain se_atten rather than se_atten_v2. Nothing in the repo crosses it.

Your op-level test does construct exactly that configuration and is a real regression test: hand-derived closed-form expectations rather than the op compared against itself, all three autograd stages, both dtypes, and it runs in CI since the custom ops are built there. Pre-fix the forward returns 0.6*c where you assert 1.4*c, and dy_dem comes out [108, 0, 0] against your [36, 108, 27] -- not a tolerance question.

What would round it out is one end-to-end case crossing the cell above: a compressed se_atten (not v2) with exclude_types and set_davg_zero: True, asserted against the uncompressed model the way TestDeepPotATPBCExcludeTypes already does. That is the configuration a user would actually hit, and it is the one the current suite structurally cannot reach.

The one thing I would like your view on before approving. This PR makes the is_sorted=True fold genuinely live for se_atten for the first time, and the Python-side guarantee that keeps it safe was never added for that descriptor family:

deepmd-kit/deepmd/pt/model/descriptor/se_atten.py

Lines 813 to 817 in dba50c5

	def need_sorted_nlist_for_lower(self) -> bool:
	"""Returns whether the descriptor block needs sorted nlist when using `forward_lower`."""
	return False

Compare the se_e2_a version, where the method returns self.compress specifically so that the sorted-nlist precondition holds, with a comment explaining the coupling:

https://github.com/deepmodeling/deepmd-kit/blob/dba50c5a43a26d9697b4e186ba353fa106d6817c/deepmd/pt/model/descriptor/se_a.py#L863-L879

The exposure is a compressed DPA1 with empty exclude_types, which gives is_sorted = True, so the fold applies -- but need_sorted_nlist_for_lower() returning False means the forward_lower neighbour list from C++/LAMMPS, built at rcut+skin and not pre-sorted, can place an out-of-rcut neighbour with sw == 0 ahead of real ones. The guard then fires early and multiplies that row by (nnei - jj), silently producing a wrong descriptor. That is the same failure mode fixed for se_e2_a. se_t.py and se_t_tebd.py have the identical unconditional return False.

This hazard predates your PR -- before it the flag was ignored and the kernel behaved as sorted regardless, so the LAMMPS exposure was already there. I am not claiming you introduced it. But it is the other half of the same invariant, and this PR is what makes the flag meaningful, so I would rather we decide deliberately whether it belongs here or in an immediate follow-up than leave it implicit. return self.geo_compress would mirror the se_a fix.

Two smaller notes inline. Also worth a line in the description: with non-empty exclude_types the fold is now genuinely disabled, so compressed DPA1 pays a full nnei loop where it previously short-circuited. Correct, but a measurable cost that users may notice.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[njzjz]

Addressed in 8a3a023:

• DescrptBlockSeAtten.need_sorted_nlist_for_lower() now returns self.geo_compress, so only the geometric tabulation path requests the extra lower-nlist filtering/sort. Type-embedding-only compression remains unchanged.
• Added a functional compressed-DPA1 forward_lower regression with an unsorted rcut+skin list containing both in-cutoff and out-of-cutoff neighbors.
• Changed the existing exclude-types end-to-end case to plain se_atten with set_davg_zero: true and tebd_input_mode: strip, so it crosses the previously untested cell.
• Made the operator fixture's a1 coefficients nonzero, so the em_x gradient and the grad-grad dz_xx * var_grad term are exercised.
• Documented the two invariants supporting the non-attention se_a hard-coded true values.

I also audited se_t.py and se_t_tebd.py. Their compressed path uses tabulate_fusion_se_t, which does not have the is_sorted early-fold contract, so they do not need the analogous geo_compress override.

History-wise, there are two distinct origins:

• The bug fixed by [Code scan] PyTorch SE-Attention tabulation ignores is_sorted in forward and first backward #5890 was introduced by pt: add 4 tabulate_fusion op #3877, which added the PyTorch tabulation op but omitted is_sorted from the backend calls in forward and first backward.
• The lower-nlist exposure discussed here became reachable in Add 4 pt descriptor compression #4227, which added PyTorch se_atten geometric compression while need_sorted_nlist_for_lower() still returned False. fix(pt): fix lammps nlist sort with large sel #3993 introduced that method/default earlier, but there was no compressed se_atten fold at that point.

Validated with the rebuilt PyTorch C++ op, all 12 SE-Attention tabulation tests, the new forward_lower regression, and the targeted end-to-end compression comparison.

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

[wanghan-iapcm]

All five points from the previous round check out against 8a3a02381.

DescrptBlockSeAtten.need_sorted_nlist_for_lower() now returns self.geo_compress, which is the right predicate -- narrower than se_a's self.compress, since only the geometric path reaches the se_a tabulate op, and geo_compress is initialized in __init__ so uncompressed models are unaffected. The new forward_lower regression is a genuine one: it fails pre-fix on the need_sorted_nlist_for_lower() assertion alone, and lines 114-115 verify the constructed list really does contain both in-cutoff and out-of-cutoff real neighbors rather than assuming it.

On the operator fixture, I re-derived the expectations after the a1 change: every one of the nine previously-verified numbers scales by exactly 1.5, and the em_x assertion is now discriminating instead of vacuous. The reconfigured end-to-end case crosses the cell I described -- with set_davg_zero: true the padding rows normalize to all-zero and so match excluded rows exactly, which is what arms the guard; pre-fix that case folds an excluded row with real neighbors behind it and misses the 1e-10 comparison.

One correction on my side: my note about se_t.py and se_t_tebd.py was over-broad and your rebuttal is right. The ago == xx && ... && is_sorted fold and its break exist only in the three se_a kernels, on CPU (tabulate.cc 186/275/379) and GPU (tabulate.cu 329/414/518); se_t, se_t_tebd and se_r take no is_sorted parameter and have no fold, so an override there would have no effect. Nothing to change.

Thanks also for the performance note in the description.