Skip to content

perf(dpa4): opt so3grid (with pt_expt GridProduct wrapping fix) #5552

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
OutisLi merged 3 commits into from
Jun 18, 2026
Merged

perf(dpa4): opt so3grid (with pt_expt GridProduct wrapping fix) #5552

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4c6b538
perf(dpa4): opt so3grid
OutisLi Jun 11, 2026
b5471ee
fix
OutisLi Jun 14, 2026
f011367
fix(dpa4): make GridProduct pt_expt-wrappable
Jun 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions deepmd/dpmodel/descriptor/dpa4_nn/block.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,7 @@
not duplicated here): ``so2_attn_res``, ``so2_s2_activation``,
``node_wise_s2/so3``, ``message_node_s2/so3``, ``atten_f_mix``,
``atten_v_proj``, ``atten_o_proj`` (raised by ``SO2Convolution``) and
``ffn_so3_grid``, ``ffn_grid_mlp`` with the grid path active (raised by
``EquivariantFFN`` / ``S2GridNet``).
``ffn_so3_grid`` with the grid path active (raised by ``EquivariantFFN``).

The pt eval-time activation-checkpoint / nvtx instrumentation
(``DP_ACT_INFER``, ``DP_COMPILE_INFER``, ``nvtx_range``) is pt-runtime-only
Expand Down
11 changes: 3 additions & 8 deletions deepmd/dpmodel/descriptor/dpa4_nn/ffn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,6 @@

- ``ffn_so3_grid=True`` — the pt path instantiates ``SO3GridNet``
(pt ffn.py:209), which is not ported to dpmodel.
- ``grid_mlp=True`` together with the grid path active selects
``op_type='mlp'`` for ``S2GridNet`` (pt ffn.py:206); the delegate
``S2GridNet`` constructor raises for that op type (``GridMLP`` is not
ported), so no duplicate guard is added here.
"""

from __future__ import (
Expand Down Expand Up @@ -85,9 +81,9 @@ class EquivariantFFN(NativeOP):
kmax
Maximum Wigner-D frame order (|k|) used by the SO3 Wigner-D FFN grid.
grid_mlp
If True, select the polynomial grid MLP operation when the
block-internal FFN grid path is enabled. Not ported: the delegate
``S2GridNet`` raises ``NotImplementedError`` for ``op_type='mlp'``.
If True, select the polynomial grid MLP operation (``op_type='mlp'``)
when the block-internal FFN grid path is enabled. ``grid_branch`` takes
precedence when positive.
grid_branch
Number of scalar-routed polynomial product branches used when the
block-internal FFN grid path is enabled. ``0`` disables this branch
Expand Down Expand Up @@ -203,7 +199,6 @@ def __init__(
if self.use_grid_branch
else ("mlp" if self.use_grid_mlp else "glu")
)
# op_type='mlp' raises NotImplementedError inside S2GridNet
self.act: NativeOP = S2GridNet(
lmax=self.lmax,
channels=self.hidden_channels,
Expand Down
Loading
Loading