perf(pt): replace SeZM CuTe inference kernels

Author: OutisLi

deepmd/dpmodel/descriptor/dpa4_nn/so2.py

@@ -1616,7 +1616,6 @@ def __init__(
         # value path: every hook stays ``None`` and ``so2_message`` takes the
         # dense branch. The ``pt_expt`` backend binds the selected implementation.
         self._triton_value_path = None
-        self._cute_value_path = None
         self._cutile_value_path = None
 
         # === Step 14. Optional fused training seams ===
@@ -2375,12 +2374,6 @@ def so2_message(
             if self._cached_edge_csr_fn is not None:
                 self._cached_edge_csr_fn(edge_cache, "src", x.shape[0])
             x_local, rad_feat = self._triton_value_path(x, edge_cache, radial_feat)
-        elif self._cute_value_path is not None and not 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
-            # per-edge focus-major intermediates stay resident on chip. ===
-            x_local, rad_feat = self._cute_value_path(x, edge_cache, radial_feat)
         else:
             # === Steps 1-3. Rotation, radial mixing and the focus-major cast ===
             x_local, rad_feat = self._rotate_mix(x, edge_cache, radial_feat)

deepmd/kernels/cute/neo/__init__.py

@@ -52,7 +52,6 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
-from deepmd.kernels.cute.neo import runtime_policy as cute_runtime_policy
 from deepmd.pt.utils import (
     env,
 )
@@ -66,6 +65,7 @@
 from deepmd.pt.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.pt_expt.kernels.cute.sezm import runtime_policy as cute_runtime_policy
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
     use_amp_infer,
@@ -1776,7 +1776,7 @@ def forward_with_edges(
     @torch.jit.unused
     def _run_output_readout(self, ffn_in: torch.Tensor) -> torch.Tensor:
         """Return the residual-inclusive scalar output with guarded CuTe routing."""
-        from deepmd.kernels.cute.neo.readout_l0 import (
+        from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import (
             run_neo_output_readout,
         )
 
@@ -1953,7 +1953,7 @@ def _packed_wigner_candidate(
         geometry_dtype: torch.dtype,
     ) -> bool:
         """Return whether all blocks satisfy packed K1's stable contract."""
-        from deepmd.kernels.cute.neo.k1 import (
+        from deepmd.pt_expt.kernels.cute.sezm.k1 import (
             is_packed_wigner_candidate,
         )
 

deepmd/pt/model/descriptor/sezm.py

@@ -26,16 +26,16 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
-from deepmd.kernels.cute.neo.runtime_policy import (
-    is_cute_infer_enabled,
-)
 from deepmd.pt.utils import (
     env,
 )
 from deepmd.pt.utils.env import (
     PRECISION_DICT,
     RESERVED_PRECISION_DICT,
 )
+from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import (
+    is_cute_infer_enabled,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -834,7 +834,7 @@ def _run_so2_unit_impl(
     ) -> torch.Tensor:
         """Run the SO(2) unit implementation."""
         if not self.training and is_cute_infer_enabled():
-            from deepmd.kernels.cute.neo.k1 import (
+            from deepmd.pt_expt.kernels.cute.sezm.k1 import (
                 maybe_run_cute_k1,
             )
 

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

@@ -25,7 +25,7 @@
     rearrange,
 )
 
-from deepmd.kernels.cute.neo.runtime_policy import (
+from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import (
     is_cute_infer_enabled,
 )
 
@@ -795,7 +795,7 @@ def _build_edge_wigner(
     if not build_full:
         return None, None, edge_quat
     if packed_wigner:
-        from deepmd.kernels.cute.neo.k4_wignerd import (
+        from deepmd.pt_expt.kernels.cute.sezm.k4_wignerd import (
             run_cute_wignerd,
         )
 
@@ -819,7 +819,7 @@ def _packed_wigner_edges_eligible(
     runtime_dtypes: tuple[torch.dtype, ...] = (),
 ) -> bool:
     """Finish packed eligibility from scalar shape and provenance metadata."""
-    from deepmd.kernels.cute.neo.k1 import (
+    from deepmd.pt_expt.kernels.cute.sezm.k1 import (
         packed_wigner_edges_eligible,
     )
 

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

@@ -22,10 +22,6 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
-from deepmd.kernels.cute.neo.k1_wigner_layout import (
-    PACKED_VALUE_COUNT,
-    ZONAL_PANEL_OFFSETS,
-)
 from deepmd.pt.model.network.mlp import (
     MLPLayer,
 )
@@ -39,6 +35,10 @@
 from deepmd.pt.utils.utils import (
     get_generator,
 )
+from deepmd.pt_expt.kernels.cute.sezm.k1_wigner_layout import (
+    PACKED_VALUE_COUNT,
+    ZONAL_PANEL_OFFSETS,
+)
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
 )
@@ -305,7 +305,7 @@ def forward(
 
         # === Step 3. Optional fused message construction and reduction ===
         if not self.training:
-            from deepmd.kernels.cute.neo.gie import (
+            from deepmd.pt_expt.kernels.cute.sezm.gie import (
                 is_cute_gie_enabled,
                 maybe_run_cute_gie,
             )

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

@@ -30,13 +30,13 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
-from deepmd.kernels.cute.neo import runtime_policy as cute_runtime_policy
 from deepmd.pt.utils import (
     env,
 )
 from deepmd.pt.utils.utils import (
     get_generator,
 )
+from deepmd.pt_expt.kernels.cute.sezm import runtime_policy as cute_runtime_policy
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
     triton_train_level,
@@ -1062,7 +1062,7 @@ def _pair_grid(
             and cute_runtime_policy.is_cute_infer_enabled()
             and _inference_mode_is_frozen(self)
         ):
-            from deepmd.kernels.cute.neo.output_grid_product import (
+            from deepmd.pt_expt.kernels.cute.sezm.output_grid_product import (
                 maybe_run_cute_output_grid_product,
             )
 

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

@@ -487,20 +487,20 @@ def _clear_sezm_compile_cache(model: nn.Module) -> None:
     crash or silently skip LoRA parameters.  Mirrors the pattern used in
     :meth:`SeZMModel.reset_head_for_mode`.
     """
-    from deepmd.kernels.cute.neo.k1 import (
-        invalidate_cute_k1_state,
-    )
     from deepmd.pt.model.model.sezm_model import (
         _clear_shared_sezm_compile_cache,
     )
+    from deepmd.pt_expt.kernels.cute.sezm.k1 import (
+        invalidate_cute_k1_state,
+    )
 
     readout_invalidator = None
 
     for m in model.modules():
         invalidate_cute_k1_state(m)
         if hasattr(m, "_neo_sm80_readout_input_fold_cache"):
             if readout_invalidator is None:
-                from deepmd.kernels.cute.neo.readout_l0 import (
+                from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import (
                     invalidate_neo_readout_input_fold,
                 )
 

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

@@ -37,7 +37,6 @@
     cuda_train_enabled,
     triton_infer_level,
     triton_train_level,
-    use_cute_infer,
     use_cutile_infer,
 )
 from deepmd.utils.version import (
@@ -1215,31 +1214,27 @@ def __init__(
         self.compute_dtype = get_promoted_dtype(self.dtype)
         # Opt-in fused fast paths, selected by ``DP_TRITON_INFER`` /
         # ``DP_TRITON_TRAIN`` (cumulative levels, see :func:`triton_infer_level`
-        # and :func:`triton_train_level`) and ``DP_CUTE_INFER``. Each is read
+        # and :func:`triton_train_level`). Each is read
         # once at construction so it becomes a compile-time constant in the
         # traced (``make_fx``) graph. Level 1 replaces the dense ``bmm``
         # rotation with universal Triton kernels; level 2 additionally binds the
         # table-configured fused value path; inference level 3 routes the mixing
         # stack through the fp16x3 tensor-core operator on swept shapes.
-        # ``DP_CUTE_INFER`` selects the experimental CuTe value-path operator
-        # instead; both gates claim the same ``so2_message`` value path, so
-        # enabling them together has no coherent meaning and is rejected at
-        # construction. The CuTe, cuTile and hand-written CUDA paths remain
-        # inference-only. 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.
+        # ``DP_CUTILE_INFER`` selects a complete alternative SO(2) path and is
+        # therefore mutually exclusive with Triton. CuTe dispatch is resolved at
+        # the enclosing SeZM block, where its exact-shape K1 kernel can coexist
+        # with these fallback paths. cuTile and the hand-written CUDA paths remain
+        # inference-only. Fused entries stay ``None`` when their backend is
+        # unavailable or the block layout is unsupported.
         self.triton_infer_level = triton_infer_level()
         self.triton_train_level = triton_train_level()
         self.use_triton_infer = self.triton_infer_level >= 1
-        self.use_cute_infer = use_cute_infer()
         self.use_cutile_infer = use_cutile_infer()
-        if sum((self.use_triton_infer, self.use_cute_infer, self.use_cutile_infer)) > 1:
+        if self.use_triton_infer and self.use_cutile_infer:
             raise ValueError(
-                "DP_TRITON_INFER, DP_CUTE_INFER and DP_CUTILE_INFER are mutually "
-                "exclusive: each selects a complete accelerated inference path. "
-                "Enable exactly one of them."
+                "DP_TRITON_INFER and DP_CUTILE_INFER are mutually exclusive: "
+                "each selects a complete accelerated SO(2) inference path."
             )
-        self._cute_value_path = None
         self._triton_value_path = None
         self._cutile_value_path = None
 
@@ -1753,19 +1748,9 @@ def __init__(
             if SEGMENT_SOFTMAX_TRITON_AVAILABLE:
                 self._segment_softmax_fn = segment_softmax
 
-        # === Step 17. 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 deepmd.pt_expt.kernels.cute.sezm import (
-                make_cute_value_path,
-            )
-
-            self._cute_value_path = make_cute_value_path(self)
-
-        # === Step 18. Optional fused cuTile SO(2) value-path operators ===
-        # Complete cuTile inference path, mutually exclusive with the two gates
-        # above. The factory validates the block layout and returns ``None``
+        # === Step 17. Optional fused cuTile SO(2) value-path operators ===
+        # Complete cuTile inference path, mutually exclusive with Triton. The
+        # factory validates the block layout and returns ``None``
         # otherwise, leaving the dense reference path in charge.
         if self.use_cutile_infer:
             from deepmd.pt_expt.kernels.cutile.sezm.so2_value_path import (
@@ -2460,12 +2445,6 @@ def so2_message(
             # inter-layer activations off the traced graph. ===
             cached_edge_csr(edge_cache, "src", x.shape[0])
             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
-            # per-edge focus-major intermediates stay resident on chip. ===
-            x_local, rad_feat = self._cute_value_path(x, edge_cache, radial_feat)
         elif self._triton_rotate_mix is not None and active_triton_level(self) >= 1:
             # === Steps 1-3 (fused rotate-mix operator). One edge-parallel
             # kernel gathers the source features, applies the block-diagonal

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

@@ -532,7 +532,7 @@
 
 def _neo_cute_infer_enabled() -> bool:
     """Return whether the model builder must destination-sort CuTe edges."""
-    from deepmd.kernels.cute.neo.runtime_policy import (
+    from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import (
         is_cute_infer_enabled,
     )
 
@@ -548,7 +548,7 @@ def _neo_cute_nlist_eager_island_enabled(device: torch.device) -> bool:
         compute_capability = tuple(torch.cuda.get_device_capability(device))
     except RuntimeError:
         return False
-    from deepmd.kernels.cute.neo.runtime_policy import (
+    from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import (
         is_k1_eager_island_enabled,
     )
 
@@ -929,7 +929,7 @@ def _invalidate_compiled_state_after_load(
             ):
                 continue
             if k1_invalidator is None:
-                from deepmd.kernels.cute.neo.k1 import (
+                from deepmd.pt_expt.kernels.cute.sezm.k1 import (
                     invalidate_cute_k1_state,
                 )
 
@@ -2156,15 +2156,17 @@ def trace_and_compile(
         # Register Python-owned K1 state before make_fx starts. The opt-in thin
         # path can then keep adjacent linears in this graph while the CuTe work
         # remains opaque behind its existing custom op.
-        from deepmd.kernels.cute.neo import runtime_policy as cute_runtime_policy
+        from deepmd.pt_expt.kernels.cute.sezm import (
+            runtime_policy as cute_runtime_policy,
+        )
 
         compute_capability = (
             tuple(torch.cuda.get_device_capability(coord.device))
             if coord.device.type == "cuda"
             else None
         )
         if not self.training and cute_runtime_policy.is_cute_infer_enabled():
-            from deepmd.kernels.cute.neo.readout_l0 import (
+            from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import (
                 maybe_prepare_sm80_readout_input_fold,
             )
 
@@ -2179,7 +2181,7 @@ def trace_and_compile(
             and cute_runtime_policy.is_cute_infer_enabled()
             and cute_runtime_policy.is_supported_k1_capability(compute_capability)
         ):
-            from deepmd.kernels.cute.neo.k1 import (
+            from deepmd.pt_expt.kernels.cute.sezm.k1 import (
                 prepare_cute_k1_blocks,
             )