perf(pt): fused Triton zonal scatter for the GIE head (DP_GIE_SCATTER_INFER)

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

@@ -37,6 +37,7 @@
 )
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
+    triton_infer_level,
 )
 from deepmd.utils.version import (
     check_version_compatibility,
@@ -226,6 +227,41 @@ def __init__(
                 self.lmax, self.ebed_dim - 1, self.channels
             )
 
+        # === Fused Triton message-and-scatter operator ===
+        # Same boundary as the CUDA operator above, served by the Triton
+        # stack instead: the kernel specializes on the packed contiguous row
+        # layout, so anything else (or a non-fp32 compute) keeps the reference
+        # composition.
+        self._triton_scatter = None
+        # The fused Triton scatter rides the level-1 universal-kernel
+        # contract: no tunable configuration, strictly removes the
+        # (E, D-1, C) message round-trip, so no table selection is involved.
+        if triton_infer_level() >= 1 and self.dtype is torch.float32 and self.lmax == 3:
+            # The module owns its optional-triton import guard, so this
+            # import succeeds everywhere and the exported flag is the
+            # single source of truth for the capability check below.
+            from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import (
+                GIE_SCATTER_TRITON_AVAILABLE,
+                gie_zonal_scatter,
+            )
+
+            index_device = self.non_scalar_row_index.device
+            packed_rows = torch.arange(
+                1, self.ebed_dim, device=index_device, dtype=torch.long
+            )
+            counts = [2 * ell + 1 for ell in range(1, self.lmax + 1)]
+            packed_slots = torch.repeat_interleave(
+                torch.arange(self.lmax, device=index_device, dtype=torch.long),
+                torch.tensor(counts, device=index_device, dtype=torch.long),
+            )
+            if (
+                GIE_SCATTER_TRITON_AVAILABLE
+                and torch.cuda.is_available()
+                and torch.equal(self.non_scalar_row_index, packed_rows)
+                and torch.equal(self.radial_slot_index_for_row, packed_slots)
+            ):
+                self._triton_scatter = gie_zonal_scatter
+
     def forward(
         self,
         *,
@@ -281,16 +317,47 @@ def forward(
 
         # === Step 3. Broadcast radial features per row ===
         # Each non-scalar packed row reuses the radial feature of its degree l.
-        # The fused operator spans this broadcast and the scatter of Step 4, so
-        # it takes over whenever nothing else joins the message in between.
-        if (
-            self._can_fuse_scatter(zonal_coupling)
-            and spin_l1_message is None
-            and edge_cache.csr_cache is not None
-        ):
-            return self.forward_fused_scatter(
-                n_nodes, edge_cache, radial_feat, zonal_coupling
-            )
+        # One fused operator spans this broadcast and the scatter of Step 4,
+        # so it takes over whenever nothing else joins the message in
+        # between.  The eligibility conditions and the destination CSR are
+        # shared once here; only the backend selection differs (the CUDA
+        # operator folds the degree normalization into its kernel, the
+        # Triton one leaves the scale in the traced graph so its cotangent
+        # keeps flowing).  Training keeps the reference composition for the
+        # Triton backend: the fused backward registers no next-derivative
+        # formula, and force-loss training differentiates the backward
+        # itself.
+        if spin_l1_message is None and edge_cache.csr_cache is not None:
+            cuda_selected = self._can_fuse_scatter(zonal_coupling)
+            triton_selected = self._triton_scatter is not None and not self.training
+            if cuda_selected or triton_selected:
+                from .edge_cache import (
+                    cached_edge_csr,
+                )
+
+                # The destination view is only needed by a fused backend;
+                # the reference path below must not pay for its sort.
+                order, row_ptr = cached_edge_csr(edge_cache, "dst", n_nodes)
+            if cuda_selected:
+                return self.forward_fused_scatter(
+                    n_nodes, edge_cache, radial_feat, zonal_coupling, order, row_ptr
+                )
+            if triton_selected:
+                # The operator contract promises contiguous inputs: the
+                # radial slice and the gathered coupling are flattened in
+                # the traced graph before entering the op, like every
+                # other sezm_triton operator.
+                out = self._triton_scatter(
+                    zonal_coupling.contiguous(),
+                    radial_feat.contiguous(),
+                    edge_cache.dst,
+                    order,
+                    row_ptr,
+                    n_nodes,
+                    self.lmax,
+                )  # (N, D, C), l=0 zero, unnormalized
+                out.mul_(edge_cache.inv_sqrt_deg)
+                return out
 
         radial_value_for_row = radial_feat.index_select(
             1, self.radial_slot_index_for_row
@@ -333,6 +400,8 @@ def forward_fused_scatter(
         edge_cache: EdgeFeatureCache,
         radial_feat: torch.Tensor,
         zonal_coupling: torch.Tensor,
+        order: torch.Tensor,
+        row_ptr: torch.Tensor,
     ) -> torch.Tensor:
         """
         Build and reduce the geometric message with the fused CUDA operator.
@@ -359,13 +428,6 @@ def forward_fused_scatter(
             zonal_scatter,
         )
 
-        from .edge_cache import (
-            cached_edge_csr,
-        )
-
-        # === Step 1. Destination CSR, shared with every other edge consumer ===
-        order, row_ptr = cached_edge_csr(edge_cache, "dst", n_nodes)
-
         # === Step 2. Fused message build, reduction, padding and normalization ===
         # The operator emits the packed node layout already normalized, so the
         # scalar row and the degree scaling cost no extra pass. The scaling is

deepmd/pt_expt/kernels/triton/sezm/__init__.py

@@ -9,6 +9,9 @@
 from .force_assembly import (
     FORCE_ASSEMBLY_TRITON_AVAILABLE,
 )
+from .gie_scatter import (
+    GIE_SCATTER_TRITON_AVAILABLE,
+)
 from .radial_mix import (
     RADIAL_MIX_TRITON_AVAILABLE,
 )
@@ -39,8 +42,10 @@
     and STACK_FP16X3_TRITON_AVAILABLE
     and WIGNER_MONOMIALS_TRITON_AVAILABLE
     and FORCE_ASSEMBLY_TRITON_AVAILABLE
+    and GIE_SCATTER_TRITON_AVAILABLE
 )
 
 __all__ = [
+    "GIE_SCATTER_TRITON_AVAILABLE",
     "TRITON_AVAILABLE",
 ]