refactor(dpa4): triton rotation

Author: OutisLi

deepmd/pt/model/descriptor/sezm.py

@@ -33,7 +33,6 @@
 )
 
 import math
-import os
 from contextlib import (
     contextmanager,
 )
@@ -555,12 +554,6 @@ def __init__(
         self.layer_scale = bool(layer_scale)
         self.use_amp = bool(use_amp)  # and self.training
         self.trainable = bool(trainable)
-        self.use_triton = os.environ.get("DP_TRITON", "0").lower() in (
-            "1",
-            "true",
-            "yes",
-            "on",
-        )
         self.seed = seed
         self.random_gamma = bool(random_gamma)
         self.add_chg_spin_ebd = bool(add_chg_spin_ebd)
@@ -899,7 +892,6 @@ def __init__(
                     ffn_activation_function=self.ffn_activation_function,
                     ffn_glu_activation=self.ffn_glu_activation,
                     mlp_bias=self.mlp_bias,
-                    use_triton=self.use_triton,
                     eps=self.eps,
                     dtype=self.dtype,
                     seed=child_seed(seed_blocks, block_idx),
@@ -1128,7 +1120,6 @@ def forward(
                 # the model is roll-equivariant, so inference fixes gamma.
                 random_gamma=self.random_gamma and self.training,
                 wigner_calc=self.wigner_calc,
-                use_geometry_rbf_triton=(self.use_triton and not self.training),
             )
 
         ebed_dim_0 = self.node_ebed_dims[0]  # (node_lmax+1)^2

deepmd/pt/model/descriptor/sezm_nn/block.py

@@ -214,9 +214,6 @@ class SeZMInteractionBlock(nn.Module):
         - SO3Linear: l=0 bias
         - SO2Linear: l=0 bias
         - GatedActivation: gate linear bias
-    use_triton
-        If True, opt into fused Triton SO(2) rotation kernels inside
-        ``SO2Convolution`` when the runtime supports them.
     eps
         Small epsilon for numerical stability.
     dtype
@@ -275,7 +272,6 @@ def __init__(
         ffn_activation_function: str,
         ffn_glu_activation: bool = True,
         mlp_bias: bool = False,
-        use_triton: bool = False,
         eps: float = 1e-7,
         dtype: torch.dtype,
         seed: int | list[int] | None,
@@ -370,7 +366,6 @@ def __init__(
         self.ffn_activation_function = str(ffn_activation_function)
         self.ffn_glu_activation = bool(ffn_glu_activation)
         self.mlp_bias = bool(mlp_bias)
-        self.use_triton = bool(use_triton)
         self.eps = float(eps)
         self.dtype = dtype
         self.device = env.DEVICE
@@ -436,7 +431,6 @@ def __init__(
             lebedev_quadrature=self.so2_lebedev_quadrature,
             activation_function=self.so2_activation_function,
             mlp_bias=self.mlp_bias,
-            use_triton=self.use_triton,
             eps=self.eps,
             dtype=dtype,
             seed=seed_so2_conv,

deepmd/pt/model/descriptor/sezm_nn/cute/__init__.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""
+CuTe-DSL accelerated SO(2) rotation operators for SeZM / DPA4.
+
+This package provides a self-contained, ``torch.compile``-friendly implementation
+of the two fused gather + batched-GEMM operators used by the SeZM SO(2) edge
+convolution:
+
+* ``rotate_to_local``  : ``out[e] = wigner[e][coeff_index] @ x[src[e]]``
+* ``rotate_back``      : ``out[e] = wigner[e][:, coeff_index] @ x_local[e]``
+
+The kernels are written with the NVIDIA CuTe DSL (``cutlass.cute``) and fuse the
+Wigner-row/column gather and the source-node gather directly into the matmul, so
+the large ``D_to_m`` / ``x_src`` intermediates are never materialized. They are
+exposed through the modern ``torch.library.custom_op`` API (functional, with
+``register_fake`` + ``register_autograd``) so that they compose correctly with
+``torch.compile`` and autograd.
+
+The top-level entry points are re-exported here for convenience.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from .so2_rotation import (
+    SEZM_CUTE_AVAILABLE,
+    rotate_back_cute,
+    rotate_to_local_cute,
+)
+
+__all__ = [
+    "SEZM_CUTE_AVAILABLE",
+    "rotate_back_cute",
+    "rotate_to_local_cute",
+]