perf(pt/dpa4): triton fusion kernel

Author: OutisLi

deepmd/pt/entrypoints/freeze_pt2.py

@@ -857,7 +857,7 @@ def freeze_sezm_to_pt2(
         exported = move_to_device_pass(exported, target_device)
 
     out_path_str = str(out_path)
-    compile_options = build_inductor_compile_options()
+    compile_options = build_inductor_compile_options(inference=True)
     # Keep AOTInductor aligned with the eval compile path.  ``triton.max_tiles=1``
     # keeps data-dependent edge axes on Triton's x grid, whose bound is large
     # enough for production-scale neighbor lists.

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

@@ -1173,14 +1173,23 @@ def __init__(
         # ``DP_CUTE_INFER``. Each flag is read once at construction so it becomes a
         # compile-time constant in the traced (``make_fx``) graph, and each only
         # takes effect during inference. ``DP_TRITON_INFER`` replaces the dense
-        # ``bmm`` rotation with fused Triton kernels; ``DP_CUTE_INFER`` selects a
-        # fused CuTe value-path operator and is independent of the Triton flag. The
-        # CuTe value-path entry is bound at the end of construction (once every
-        # submodule exists) and stays ``None`` when the backend is unavailable or
-        # the block layout is unsupported.
+        # ``bmm`` rotation with fused Triton kernels; ``DP_CUTE_INFER`` selects
+        # the experimental CuTe value-path operator instead.  Both flags claim
+        # the same ``so2_message`` value path, so enabling them together has no
+        # coherent meaning and is rejected at construction.  The fused
+        # value-path entries are bound at the end of construction (once every
+        # submodule exists) and stay ``None`` when the backend is unavailable
+        # or the block layout is unsupported.
         self.use_triton_infer = use_triton_infer()
         self.use_cute_infer = use_cute_infer()
+        if self.use_triton_infer and self.use_cute_infer:
+            raise ValueError(
+                "DP_TRITON_INFER and DP_CUTE_INFER are mutually exclusive: "
+                "both select the fused SO(2) value-path backend. Enable "
+                "exactly one of them."
+            )
         self._cute_value_path = None
+        self._triton_value_path = None
 
         # === Step 1. Split deterministic seeds at the module top-level ===
         seed_so2_stack = child_seed(seed, 0)
@@ -1570,7 +1579,23 @@ def __init__(
             self._flash_atten_fn = flash_atten_aggregate
             self._build_row_ptr_fn = build_row_ptr
 
-        # === Step 13. Optional fused CuTe SO(2) value-path operator ===
+        # === Step 13. Optional fused Triton SO(2) value-path operators ===
+        # Fuses rotate-to-local, the radial degree mixing, the gated mixing
+        # stack, and the focus competition of ``so2_message`` into the
+        # ``sezm_triton::so2_rotate_mix`` / ``so2_mixing_stack`` operators.
+        # The factory validates the block layout (``mmax == 1``, gated stack
+        # with an identity final layer, supported focus widths) and returns
+        # ``None`` otherwise, leaving the reference path in charge.
+        if self.use_triton_infer:
+            from .triton.so2_value_path import (
+                make_triton_value_path,
+            )
+
+            self._triton_value_path = make_triton_value_path(self)
+
+        # === Step 14. Optional fused CuTe SO(2) value-path operator ===
+        # Experimental alternative backend; mutually exclusive with the Triton
+        # flag (enforced above).
         if self.use_cute_infer:
             from .cute import (
                 make_cute_value_path,
@@ -1922,7 +1947,15 @@ def so2_message(
         src, dst = edge_cache.src, edge_cache.dst
         n_edge = src.numel()
 
-        if self._cute_value_path is not None and not self.training:
+        if self._triton_value_path is not None and not self.training:
+            # === Steps 1-5 (fused Triton operators). ``so2_rotate_mix`` folds
+            # the rotation and the radial degree mixing into one edge-parallel
+            # kernel writing the focus-major layout; ``so2_mixing_stack`` runs
+            # the whole gated stack with the competition weight fused into its
+            # final store, keeping the inter-layer activations off the traced
+            # graph. ===
+            x_local, rad_feat = self._triton_value_path(x, edge_cache, radial_feat)
+        elif self._cute_value_path is not None and not self.training:
             # === Steps 1-5 (fused CuTe operator). The operator folds
             # rotate_to_local, radial degree mixing, the multi-layer gated SO(2)
             # stack, and the focus competition into the bucketed kernels; the

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

@@ -6,6 +6,9 @@
 descriptor; the package-level API only exposes availability.
 """
 
+from .force_assembly import (
+    FORCE_ASSEMBLY_TRITON_AVAILABLE,
+)
 from .radial_mix import (
     RADIAL_MIX_TRITON_AVAILABLE,
 )
@@ -15,6 +18,12 @@
 from .so2_rotation import (
     TRITON_ROTATION_AVAILABLE,
 )
+from .so2_value_path import (
+    SO2_VALUE_PATH_TRITON_AVAILABLE,
+)
+from .wigner_monomials import (
+    WIGNER_MONOMIALS_TRITON_AVAILABLE,
+)
 
 # Every kernel module guards its ``@triton.jit`` definitions behind a ``triton``
 # import, so the module-level checks are equivalent. Expose a single
@@ -23,6 +32,9 @@
     TRITON_ROTATION_AVAILABLE
     and RADIAL_MIX_TRITON_AVAILABLE
     and SO2_BLOCK_GEMM_TRITON_AVAILABLE
+    and SO2_VALUE_PATH_TRITON_AVAILABLE
+    and WIGNER_MONOMIALS_TRITON_AVAILABLE
+    and FORCE_ASSEMBLY_TRITON_AVAILABLE
 )
 
 __all__ = [