fix(dpa4): use triton_op to enable triton when freezing

[OutisLi]

use triton_op to enable triton when freezing and used in lammps

Summary by CodeRabbit

• Refactor
  • Updated GPU kernel compilation integration for improved compatibility with PyTorch's advanced compilation optimization framework.

[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: 47247fbc-f81a-4375-86eb-e3c41df4d786

📥 Commits

Reviewing files that changed from the base of the PR and between 0de53e9 and 1ae03fe.

📒 Files selected for processing (2)

• deepmd/pt/model/descriptor/sezm_nn/triton/radial_mix.py
• deepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.py

---

📝 Walkthrough

Walkthrough

This PR refactors operator registration in two Triton integration modules, switching from torch.library.custom_op to torch.library.triton_op to enable AOTInductor to bake Triton cubins into exported artifacts. Both modules add _has_no_edges() guards and wrap kernel launches with wrap_triton().

Changes

Triton Operator Registration and Kernel Launch Wrapping

Layer / File(s)	Summary
Radial Mix Triton Integration deepmd/pt/model/descriptor/sezm_nn/triton/radial_mix.py	Module docstring updated to reflect triton_op instances. wrap_triton imported. _has_no_edges() helper introduced; _launch_forward and _launch_backward use it for zero-edge fast paths and wrap kernel launches with wrap_triton(_radial_mix_fwd_kernel) and related. _radial_mix_op and _radial_mix_bwd_op switched from custom_op to triton_op registration.
SO2 Rotation Module Setup and Helper deepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.py	Module docstring updated to reflect functional triton_op instances for make_fx capture and AOTInductor baking. wrap_triton imported. _has_no_edges(n_edge: int) -> bool helper introduced to check zero-edge conditions without exporting-time symbolic tracing guards.
Dense SO2 Rotation Wrappers and Operators deepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.py	_launch_rotate_to_local_fwd, _launch_rotate_to_local_bwd, _launch_rotate_back_fwd, and _launch_rotate_back_bwd updated to use _has_no_edges() and wrap kernel launches with wrap_triton(). Dense rotation operator handles (_rotate_to_local_op, _rotate_to_local_bwd_op, _rotate_back_op) switched from custom_op to triton_op.
Block-Diagonal SO2 Rotation Wrappers and Operators deepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.py	_launch_bd_to_local_fwd, _launch_bd_to_local_bwd, _launch_bd_back_fwd, and _launch_bd_back_bwd updated to use _has_no_edges() and wrap kernel launches with wrap_triton(). Block-diagonal operator handles (_block_to_local_op, _block_to_local_bwd_op, _block_back_op, _block_back_bwd_op) switched from custom_op to triton_op.
Block-SO2 Rotation Wrappers and Operators deepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.py	_launch_bd_back_so2_fwd and _launch_bd_back_so2_bwd updated to use _has_no_edges() and wrap kernel launches with wrap_triton(). Block-SO2 operator handles (_block_back_so2_op, _block_back_so2_bwd_op) switched from custom_op to triton_op.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 27.78% 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 accurately describes the main change: switching from custom_op to triton_op registration to enable Triton kernel usage during model freezing.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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 and usage tips.}

[codecov]

Codecov Report

❌ Patch coverage is 27.58621% with 42 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.20%. Comparing base (d3834e2) to head (c833970).
⚠️ Report is 6 commits behind head on master.

Files with missing lines	Patch %	Lines
...pt/model/descriptor/sezm_nn/triton/so2_rotation.py	24.48%	37 Missing ⚠️
...d/pt/model/descriptor/sezm_nn/triton/radial_mix.py	44.44%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5521      +/-   ##
==========================================
- Coverage   82.21%   82.20%   -0.01%     
==========================================
  Files         892      892              
  Lines      101531   101573      +42     
  Branches     4240     4240              
==========================================
+ Hits        83475    83501      +26     
- Misses      16753    16769      +16     
  Partials     1303     1303              

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