diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml
index 866eca155e..bc37eea241 100644
--- a/.github/workflows/test_cuda.yml
+++ b/.github/workflows/test_cuda.yml
@@ -28,7 +28,7 @@ jobs:
     # container:
     #   image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04
     #   options: --gpus all
-    if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
+    if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'Test CUDA') || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
     steps:
       - uses: actions/checkout@v7
       - uses: actions/setup-python@v7
@@ -58,7 +58,7 @@ jobs:
       - run: |
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
-          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
+          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,cute,torch,jax] mpi4py --reinstall-package deepmd-kit
           # See https://github.com/jax-ml/jax/issues/29042
           source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
         env:
@@ -66,6 +66,7 @@ jobs:
           DP_ENABLE_NATIVE_OPTIMIZATION: 1
           DP_ENABLE_PYTORCH: 1
       - run: dp --version
+      - run: python -c "import cutlass.cute"
       - run: python -m pytest source/tests --ignore=source/tests/pd
         env:
           NUM_WORKERS: 0
@@ -78,7 +79,7 @@ jobs:
   test_cc:
     name: Test C++ on CUDA
     runs-on: gpu
-    if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
+    if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'Test CUDA') || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
     steps:
       # Jobs run on separate runners, so the C++ job needs its own complete
       # CUDA toolchain and Python dependency installation.
diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py
index 59c027bbf9..2e793fe65e 100644
--- a/deepmd/dpmodel/descriptor/dpa4.py
+++ b/deepmd/dpmodel/descriptor/dpa4.py
@@ -1210,10 +1210,10 @@ def __init__(
         # Accelerated backends may replace the distance-to-radial chain and the
         # packed Wigner-D construction. The array-API reference leaves these
         # hooks unbound and always retains the dense Wigner matrices.
-        self._cuda_radial_fn = None
-        self._cuda_wigner_fn = None
-        self._wigner_free_conv = False
-        self._packed_wigner_train = False
+        self.cuda_infer_l_1_radial = None
+        self.cuda_infer_l_1_wigner = None
+        self.cuda_infer_l_2_covers_all_blocks = False
+        self.cuda_train_covers_all_blocks = False
 
         # === Optional descriptor-level attention residuals ===
         self.final_block_attn_res = None
@@ -1549,9 +1549,31 @@ def _run_graph(
             graph = apply_pair_exclusion(graph, atype_flat, self.emask)
         if n_out_nodes is None:
             n_out_nodes = atype_flat.shape[0]
+        packed_wigner_graph = self.prepare_packed_wigner_graph(
+            graph, atype_flat.shape[0]
+        )
+        packed_wigner = packed_wigner_graph is not None
+        if packed_wigner_graph is not None:
+            graph = packed_wigner_graph
         edge_index = graph.edge_index
         edge_vec = graph.edge_vec
         edge_mask = graph.edge_mask
+        # Graph-owned endpoint orderings remain aligned with the edge payload
+        # and are shared by every segmented consumer through the edge cache.
+        graph_csr_cache = None
+        if all(
+            value is not None
+            for value in (
+                graph.destination_order,
+                graph.destination_row_ptr,
+                graph.source_order,
+                graph.source_row_ptr,
+            )
+        ):
+            graph_csr_cache = {
+                "dst": (graph.destination_order, graph.destination_row_ptr),
+                "src": (graph.source_order, graph.source_row_ptr),
+            }
 
         xp = array_api_compat.array_namespace(edge_vec)
         device = array_api_compat.device(edge_vec)
@@ -1596,14 +1618,18 @@ def _run_graph(
             bridging_switch=self.bridging_switch,
             edge_envelope=self.edge_envelope,
             radial_basis=self.radial_basis,
-            fused_radial=None if training else self._cuda_radial_fn,
-            fused_wigner=None if training else self._cuda_wigner_fn,
+            fused_radial=None if training else self.cuda_infer_l_1_radial,
+            fused_wigner=None if training else self.cuda_infer_l_1_wigner,
             # Random local-Z roll is a training-only augmentation; the model
             # is roll-equivariant, so inference fixes gamma.
             random_gamma=self.random_gamma and training,
             wigner_calc=self.wigner_calc,
-            build_wigner=self._build_full_wigner(),
+            build_wigner=self._build_full_wigner() or packed_wigner,
             node_partial_exchange=node_partial_exchange,
+            packed_wigner=packed_wigner,
+            destinations_sorted=graph.destination_sorted,
+            packed_wigner_fn=self.build_packed_wigner,
+            csr_cache=graph_csr_cache,
         )
 
         ebed_dim_0 = self.node_init_dim  # (node_init_lmax+1)^2
@@ -1723,6 +1749,9 @@ def _run_graph(
             edge_cache = edge_cache_to_dtype(
                 edge_cache, get_xp_precision(xp, self.precision)
             )
+            edge_cache.cute_infer_so2_metadata = self.prepare_cute_infer_so2_metadata(
+                edge_cache, n_nodes
+            )
             x = self._forward_blocks(
                 x, edge_cache, rad_feat_per_block, comm_dict=comm_dict
             )
@@ -1901,10 +1930,30 @@ def _apply_readout(self, x: Array, n_rows: int) -> Array:
             )
         for layer in self.readout_pre_layers:
             x_ro = x_ro + layer(x_ro)
+        if not self.readout_pre_layers and self.so3_readout != "none":
+            accelerated = self.run_cute_infer_readout(x_ro)
+            if accelerated is not None:
+                return xp.reshape(accelerated, (n_rows, 1, 1, self.channels))
         if self.so3_readout == "none":
             return (x_ro + self.output_ffn(x_ro))[:, 0:1, :, :]
         return x_ro[:, 0:1, :, :] + self.output_ffn.call_scalar(x_ro)
 
+    def run_cute_infer_readout(self, ffn_in: Array) -> Array | None:
+        """Run the CuTe readout when its exact inference contract matches.
+
+        Parameters
+        ----------
+        ffn_in : Array
+            Equivariant readout input with shape ``(N, D, 1, C)``.
+
+        Returns
+        -------
+        Array or None
+            Residual-inclusive scalar output with shape ``(N, C)``, or ``None``
+            when the backend has no eligible implementation.
+        """
+        return None
+
     def _edge_quaternion(self, edge_cache: EdgeCache) -> Array:
         """
         Return the cached global->local edge quaternion, rebuilding if absent.
@@ -1935,8 +1984,73 @@ def _build_full_wigner(self) -> bool:
         if not self._need_full_wigner:
             return False
         if self._in_training_mode():
-            return not self._packed_wigner_train
-        return not self._wigner_free_conv
+            return not self.cuda_train_covers_all_blocks
+        return not self.cuda_infer_l_2_covers_all_blocks
+
+    def prepare_packed_wigner_graph(
+        self,
+        graph: NeighborGraph,
+        n_nodes: int,
+    ) -> NeighborGraph | None:
+        """Prepare an edge graph for backend packed-Wigner storage.
+
+        Parameters
+        ----------
+        graph : NeighborGraph
+            Edge graph supplied to the descriptor.
+        n_nodes : int
+            Number of nodes addressed by the graph.
+
+        Returns
+        -------
+        NeighborGraph or None
+            A graph satisfying the packed layout contract, or ``None`` when
+            the backend does not select that representation.
+        """
+        return None
+
+    def build_packed_wigner(
+        self,
+        edge_quat: Array,
+        wigner_calc: Any,
+    ) -> Array | None:
+        """Build backend-specific packed Wigner storage when available.
+
+        Parameters
+        ----------
+        edge_quat : Array
+            Global-to-local edge quaternions with shape ``(E, 4)``.
+        wigner_calc : Any
+            Wigner calculator carrying the degree and basis convention.
+
+        Returns
+        -------
+        Array or None
+            Packed per-edge Wigner storage, or ``None`` to retain dense storage.
+        """
+        return None
+
+    def prepare_cute_infer_so2_metadata(
+        self,
+        edge_cache: EdgeCache,
+        n_nodes: int,
+    ) -> tuple[Array, Array, Array] | None:
+        """Build the edge metadata consumed by the CuTe SO2 implementation.
+
+        Parameters
+        ----------
+        edge_cache : EdgeCache
+            Per-forward cache carrying the destination-major edge payload.
+        n_nodes : int
+            Number of nodes addressed by the edge payload.
+
+        Returns
+        -------
+        tuple[Array, Array, Array] or None
+            Destination row pointers, source order, and source row pointers, or
+            ``None`` when the backend does not select CuTe SO2.
+        """
+        return None
 
     def _shared_wigner_runs(
         self,
@@ -1985,6 +2099,8 @@ def _build_gie_zonal_coupling(
             the blocks are skipped (all-Cartesian model) the full coupling is
             reconstructed from the edge quaternion via the m=0-only path.
         """
+        if edge_cache.D_packed is not None:
+            return self.build_cute_infer_zonal_coupling(edge_cache)
         if edge_cache.Dt_full is None:
             calc = self.gie_zonal_wigner_calc or self.wigner_calc
             shared = self._shared_wigner_runs(edge_cache, calc.lmax)
@@ -2012,6 +2128,28 @@ def _build_gie_zonal_coupling(
         )
         return xp.concat([mp_coupling, extra_coupling], axis=1)
 
+    def build_cute_infer_zonal_coupling(self, edge_cache: EdgeCache) -> Array:
+        """Extract the GIE zonal coupling from packed Wigner storage.
+
+        Parameters
+        ----------
+        edge_cache : EdgeCache
+            Per-forward cache carrying backend packed Wigner storage.
+
+        Returns
+        -------
+        Array
+            Zonal coupling with shape ``(E, D_node - 1)``.
+
+        Raises
+        ------
+        NotImplementedError
+            If a backend supplies packed Wigner storage without this extractor.
+        """
+        raise NotImplementedError(
+            "packed Wigner storage requires a backend zonal-coupling implementation"
+        )
+
     def _apply_charge_spin_embedding(
         self,
         type_ebed: Array,
diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py
index ec51a3d23e..2c8a569099 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py
@@ -44,6 +44,7 @@
 )
 
 WignerCalculatorFn = Callable[[Any], "tuple[Any, Any]"]
+PackedWignerFn = Callable[[Any, WignerCalculatorFn], Any | None]
 # Distance and keep weight to the keep-weighted envelope and radial basis, the
 # fused replacement of applying the two modules separately.
 FusedRadialFn = Callable[[Any, Any], "tuple[Any, Any]"]
@@ -94,7 +95,11 @@ class EdgeCache:
     csr_cache
         Lazy cache for endpoint CSR views used by segmented accelerated
         operators, keyed by endpoint role (``"dst"`` or ``"src"``). Built once
-        per step and shared by every consumer.
+        per step and shared by every consumer. Permutations and row pointers
+        use int64 regardless of the input graph index dtype.
+    cute_infer_so2_metadata
+        Optional ``(destination_row_ptr, source_order, source_row_ptr)`` tuple
+        prepared once for a backend CuTe SO2 implementation.
     edge_src_gate
         Optional per-edge Source Freeze Propagation Gate (SFPG) weight with
         shape (E, 1). Equals ``eta[src]`` where
@@ -112,6 +117,12 @@ class EdgeCache:
         all slots are valid (e.g. the sparse
         :func:`_edge_cache_from_arrays` path, where masking is folded into
         the per-edge weights). This field has no pt counterpart.
+    destinations_sorted
+        Whether the edge payload is destination-major with masked entries
+        confined to the suffix.
+    D_packed
+        Optional backend-specific packed Wigner panel. Dense consumers observe
+        ``D_full`` and ``Dt_full`` as ``None`` when this storage is active.
     """
 
     src: Any
@@ -126,10 +137,13 @@ class EdgeCache:
     Dt_full: Any = None
     D_to_m_cache: dict[str, Any] = field(default_factory=dict)
     Dt_from_m_cache: dict[str, Any] = field(default_factory=dict)
-    csr_cache: dict[str, Any] | None = field(default_factory=dict)
+    csr_cache: dict[str, tuple[Any, Any]] | None = field(default_factory=dict)
     edge_src_gate: Any = None
     edge_quat: Any = None
     edge_mask: Any = None
+    cute_infer_so2_metadata: tuple[Any, Any, Any] | None = None
+    destinations_sorted: bool = False
+    D_packed: Any = None
 
 
 def compute_edge_src_gate(
@@ -279,6 +293,10 @@ def _edge_cache_from_arrays(
     node_partial_exchange: Callable[[Any], Any] | None = None,
     fused_radial: FusedRadialFn | None = None,
     fused_wigner: WignerCalculatorFn | None = None,
+    packed_wigner: bool = False,
+    destinations_sorted: bool = False,
+    packed_wigner_fn: PackedWignerFn | None = None,
+    csr_cache: dict[str, tuple[Any, Any]] | None = None,
 ) -> EdgeCache:
     """
     Build the global edge cache from a sparse edge list.
@@ -331,6 +349,17 @@ def _edge_cache_from_arrays(
     fused_wigner
         Optional fused replacement of ``wigner_calc`` that builds the packed
         pair in one kernel pass.
+    packed_wigner
+        Whether the backend-specific packed Wigner contract is satisfied.
+    destinations_sorted
+        Whether the edge payload is destination-major with masked entries in
+        a suffix.
+    packed_wigner_fn
+        Optional backend implementation that builds a packed Wigner panel.
+    csr_cache
+        Optional endpoint CSR views inherited from the input graph. The cache
+        owns the same edge axis as ``edge_index``. Permutations and row pointers
+        are normalized to int64 at this graph-to-cache boundary.
     gamma
         Optional per-edge roll angles with shape (E,), used only when
         ``random_gamma`` is True. When None, drawn with the backend's RNG
@@ -351,6 +380,18 @@ def _edge_cache_from_arrays(
     n_nodes = type_ebed.shape[0]
     src = xp.astype(edge_index[0, ...], xp.int64)
     dst = xp.astype(edge_index[1, ...], xp.int64)
+    if csr_cache is not None:
+        csr_cache = {
+            endpoint: (
+                order if order.dtype == xp.int64 else xp.astype(order, xp.int64),
+                (
+                    row_ptr
+                    if row_ptr.dtype == xp.int64
+                    else xp.astype(row_ptr, xp.int64)
+                ),
+            )
+            for endpoint, (order, row_ptr) in csr_cache.items()
+        }
 
     # === Step 1. Normalize mask ===
     edge_keep = xp.astype(edge_mask, xp.bool)
@@ -379,7 +420,7 @@ def _edge_cache_from_arrays(
         edge_rbf = radial_basis(edge_len) * edge_keep_f  # (E, n_radial)
 
     # === Step 4. Edge quaternion -> Wigner-D blocks ===
-    D_full, Dt_full, edge_quat = _build_edge_wigner(
+    D_full, Dt_full, D_packed, edge_quat = _build_edge_wigner(
         edge_vec=edge_vec,
         edge_len=edge_len,
         eps=eps,
@@ -387,7 +428,9 @@ def _edge_cache_from_arrays(
         wigner_calc=fused_wigner if fused_wigner is not None else wigner_calc,
         gamma=gamma,
         build_full=build_wigner,
-    )  # (E, D, D), (E, D, D), (E, 4)
+        packed_wigner=packed_wigner,
+        packed_wigner_fn=packed_wigner_fn,
+    )
 
     # === Step 5. Edge type features ===
     edge_type_feat = build_edge_type_feat(type_ebed, src, dst)
@@ -419,9 +462,12 @@ def _edge_cache_from_arrays(
         edge_env=edge_env,
         D_full=D_full,
         Dt_full=Dt_full,
+        D_packed=D_packed,
         edge_quat=edge_quat,
         deg_norm_floor=deg_norm_floor,
         edge_src_gate=edge_src_gate,
+        destinations_sorted=destinations_sorted,
+        csr_cache=csr_cache,
     )
 
 
@@ -434,7 +480,9 @@ def _build_edge_wigner(
     wigner_calc: WignerCalculatorFn,
     gamma: Any = None,
     build_full: bool = True,
-) -> tuple[Any, Any, Any]:
+    packed_wigner: bool = False,
+    packed_wigner_fn: PackedWignerFn | None = None,
+) -> tuple[Any, Any, Any, Any]:
     """
     Build packed Wigner-D blocks from edge vectors.
 
@@ -461,13 +509,16 @@ def _build_edge_wigner(
         False (all message-passing blocks take the Cartesian path), only the
         quaternion is returned and the blocks are ``None``; the geometric
         initial embedding reconstructs the zonal coupling from the quaternion.
+    packed_wigner
+        Whether the backend-specific packed Wigner contract is satisfied.
+    packed_wigner_fn
+        Optional backend implementation that builds a packed Wigner panel.
 
     Returns
     -------
-    tuple[Array, Array, Array]
-        Packed Wigner-D matrices ``(D_full, Dt_full)`` with shape ``(E, D, D)``
-        (or ``None`` when ``build_full`` is False) and the quaternion used to
-        build them with shape ``(E, 4)``.
+    tuple[Array, Array, Array, Array]
+        Dense Wigner-D matrices ``(D_full, Dt_full)``, an optional packed panel,
+        and the quaternion used to build them.
     """
     xp = array_api_compat.array_namespace(edge_vec)
     device = array_api_compat.device(edge_vec)
@@ -491,9 +542,13 @@ def _build_edge_wigner(
 
     # === Step 3. Convert quaternions to packed Wigner-D blocks ===
     if not build_full:
-        return None, None, edge_quat
+        return None, None, None, edge_quat
+    if packed_wigner and packed_wigner_fn is not None:
+        D_packed = packed_wigner_fn(edge_quat, wigner_calc)
+        if D_packed is not None:
+            return None, None, D_packed, edge_quat
     D_full, Dt_full = wigner_calc(edge_quat)
-    return D_full, Dt_full, edge_quat
+    return D_full, Dt_full, None, edge_quat
 
 
 def _finalize_edge_cache(
@@ -507,9 +562,12 @@ def _finalize_edge_cache(
     edge_env: Any,
     D_full: Any,
     Dt_full: Any,
+    D_packed: Any,
     edge_quat: Any,
     deg_norm_floor: float,
     edge_src_gate: Any = None,
+    destinations_sorted: bool = False,
+    csr_cache: dict[str, tuple[Any, Any]] | None = None,
 ) -> EdgeCache:
     """
     Assemble the shared `EdgeCache` layout.
@@ -536,6 +594,8 @@ def _finalize_edge_cache(
     Dt_full
         Transposed packed Wigner-D matrices with shape (E, D, D), or None
         when the full Wigner-D construction is skipped.
+    D_packed
+        Optional backend-specific packed Wigner panel.
     edge_quat
         Global-to-local quaternions used to build the Wigner-D matrices with
         shape (E, 4).
@@ -547,12 +607,18 @@ def _finalize_edge_cache(
     edge_src_gate
         Optional per-edge SFPG weight with shape (E, 1). ``None`` in
         non-bridging mode.
+    destinations_sorted
+        Whether ``dst`` is nondecreasing with masked entries in a suffix.
+    csr_cache
+        Optional endpoint CSR views aligned with the cache edge axis.
 
     Returns
     -------
     EdgeCache
         Finalized per-edge cache shared by eager and compile paths.
     """
+    if deg_norm_floor <= 0.0:
+        raise ValueError("deg_norm_floor must be positive")
     xp = array_api_compat.array_namespace(edge_vec, dst)
     device = array_api_compat.device(edge_vec)
     # === Step 1. Build smooth destination degrees ===
@@ -574,11 +640,13 @@ def _finalize_edge_cache(
         inv_sqrt_deg=inv_sqrt_deg,
         D_full=D_full,
         Dt_full=Dt_full,
+        D_packed=D_packed,
         D_to_m_cache={},
         Dt_from_m_cache={},
-        csr_cache={},
+        csr_cache={} if csr_cache is None else csr_cache,
         edge_src_gate=edge_src_gate,
         edge_quat=edge_quat,
+        destinations_sorted=destinations_sorted,
     )
 
 
@@ -639,16 +707,20 @@ def edge_cache_to_dtype(cache: EdgeCache, dtype: Any) -> EdgeCache:
     # Use local variables with explicit None check and assignment.
     _D_full = cache.D_full
     _Dt_full = cache.Dt_full
+    cached_D_packed = cache.D_packed
     _edge_src_gate = cache.edge_src_gate
     _edge_quat = cache.edge_quat
     D_full: Any = None
     Dt_full: Any = None
+    D_packed: Any = None
     edge_src_gate: Any = None
     edge_quat: Any = None
     if _D_full is not None:
         D_full = xp.astype(_D_full, dtype)
     if _Dt_full is not None:
         Dt_full = xp.astype(_Dt_full, dtype)
+    if cached_D_packed is not None:
+        D_packed = xp.astype(cached_D_packed, dtype)
     if _edge_src_gate is not None:
         edge_src_gate = xp.astype(_edge_src_gate, dtype)
     if _edge_quat is not None:
@@ -667,10 +739,13 @@ def edge_cache_to_dtype(cache: EdgeCache, dtype: Any) -> EdgeCache:
         inv_sqrt_deg=xp.astype(cache.inv_sqrt_deg, dtype),
         D_full=D_full,
         Dt_full=Dt_full,
+        D_packed=D_packed,
         D_to_m_cache=None if cache.D_to_m_cache is None else {},
         Dt_from_m_cache=None if cache.Dt_from_m_cache is None else {},
         csr_cache=None if cache.csr_cache is None else dict(cache.csr_cache),
         edge_src_gate=edge_src_gate,
         edge_quat=edge_quat,
         edge_mask=cache.edge_mask,
+        cute_infer_so2_metadata=cache.cute_infer_so2_metadata,
+        destinations_sorted=cache.destinations_sorted,
     )
diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py
index 9de135e641..ba138049e7 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py
@@ -315,17 +315,28 @@ def call(
                 axis=1,
             )  # (E, D-1)
 
-        # === Step 3. Broadcast radial features per row ===
+        # === Step 3. Optional backend message construction and reduction ===
+        accelerated = self.run_cute_infer_gie(
+            n_nodes=n_nodes,
+            edge_cache=edge_cache,
+            radial_feat=radial_feat,
+            zonal_coupling=zonal_coupling,
+            spin_l1_message=spin_l1_message,
+        )
+        if accelerated is not None:
+            return accelerated
+
+        # === Step 4. 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 5, so
+        # The fused operator spans this broadcast and the scatter of Step 6, so
         # it takes over whenever nothing else joins the message in between.
         if (
-            self._can_fuse_scatter(zonal_coupling)
+            self.can_run_cuda_infer_l_1_scatter(zonal_coupling)
             and spin_l1_message is None
             and edge_cache.edge_src_gate is None
             and edge_cache.csr_cache is not None
         ):
-            return self.forward_fused_scatter(
+            return self.run_cuda_infer_l_1_scatter(
                 n_nodes, edge_cache, radial_feat, zonal_coupling
             )
 
@@ -339,7 +350,7 @@ def call(
             zonal_coupling[:, :, None] * radial_value_for_row
         )  # (E, D-1, C)
 
-        # === Step 3b. Fold in the neighbor-spin l=1 message (native spin) ===
+        # === Step 4b. Fold in the neighbor-spin l=1 message (native spin) ===
         # The l=1 coefficients occupy the first three packed non-scalar rows, so
         # the neighbor-spin message joins the geometric message there and then
         # shares the source gate, scatter and degree normalization below.
@@ -352,7 +363,7 @@ def call(
                 non_scalar_message, 1, scatter_index, spin_l1_message
             )
 
-        # === Step 4. Source Freeze Propagation Gate (optional) ===
+        # === Step 5. Source Freeze Propagation Gate (optional) ===
         # Mute messages emitted by nodes whose local neighborhood enters
         # the frozen zone. ``edge_src_gate`` is ``None`` outside bridging
         # mode so this is a no-op in normal training.
@@ -362,7 +373,7 @@ def call(
                 xp.reshape(src_gate, (n_edge, 1, 1)), non_scalar_message.dtype
             )
 
-        # === Step 5. Scatter to nodes and normalize ===
+        # === Step 6. Scatter to nodes and normalize ===
         # Destination scatter-add over ``edge_cache.dst`` (pt ``index_add_``),
         # applied after the validity masking below. This reduction is
         # layout-agnostic: it is correct both for the padded ``call`` (row-major
@@ -399,18 +410,83 @@ def call(
         out = out * xp.astype(edge_cache.inv_sqrt_deg, out.dtype)
         return xp.astype(out, dtype)
 
-    def _can_fuse_scatter(self, zonal_coupling: Any) -> bool:
-        """Return whether a backend can run the fused scatter for this input."""
+    def run_cute_infer_gie(
+        self,
+        *,
+        n_nodes: int,
+        edge_cache: EdgeCache,
+        radial_feat: Any,
+        zonal_coupling: Any,
+        spin_l1_message: Any = None,
+    ) -> Any | None:
+        """Run a backend CuTe GIE implementation when its contract matches.
+
+        Parameters
+        ----------
+        n_nodes : int
+            Number of nodes addressed by the edge cache.
+        edge_cache : EdgeCache
+            Per-forward geometric edge cache.
+        radial_feat : Any
+            Radial features with shape ``(E, lmax, C)``.
+        zonal_coupling : Any
+            Zonal coupling with shape ``(E, D - 1)``.
+        spin_l1_message : Any, optional
+            Native-spin message with shape ``(E, 3, C)``.
+
+        Returns
+        -------
+        Any or None
+            Initial embedding with shape ``(N, D, C)``, or ``None`` when no
+            backend CuTe implementation serves the input.
+        """
+        return None
+
+    def can_run_cuda_infer_l_1_scatter(self, zonal_coupling: Any) -> bool:
+        """Return whether a backend can run the fused scatter for this input.
+
+        Parameters
+        ----------
+        zonal_coupling : Any
+            Zonal coupling with shape ``(E, D - 1)``.
+
+        Returns
+        -------
+        bool
+            Whether the CUDA inference level-one scatter is eligible.
+        """
         return False
 
-    def forward_fused_scatter(
+    def run_cuda_infer_l_1_scatter(
         self,
         n_nodes: int,
         edge_cache: EdgeCache,
         radial_feat: Any,
         zonal_coupling: Any,
     ) -> Any:
-        """Build and reduce the geometric message with a backend operator."""
+        """Build and reduce the geometric message with a backend operator.
+
+        Parameters
+        ----------
+        n_nodes : int
+            Number of nodes addressed by the edge cache.
+        edge_cache : EdgeCache
+            Per-forward geometric edge cache.
+        radial_feat : Any
+            Radial features with shape ``(E, lmax, C)``.
+        zonal_coupling : Any
+            Zonal coupling with shape ``(E, D - 1)``.
+
+        Returns
+        -------
+        Any
+            Initial embedding with shape ``(N, D, C)``.
+
+        Raises
+        ------
+        NotImplementedError
+            If the backend has no fused CUDA scatter implementation.
+        """
         raise NotImplementedError(
             "The fused GIE scatter is unavailable in the dpmodel reference"
         )
diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
index 22bf8569aa..c8f585c3d4 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
@@ -881,8 +881,8 @@ def __init__(
         # training hook differentiates the same expression inside the force
         # graph on the frame-packed operands, with analytic first and second
         # order.
-        self._grid_pair_fn = None
-        self._grid_pair_train_fn = None
+        self.cuda_infer_l_1_grid_pair = None
+        self.triton_train_l_1_grid_pair = None
         self._from_grid_t = np.ascontiguousarray(projector.from_grid_mat.T)
 
         self.scalar_act = SwiGLU()
@@ -954,7 +954,9 @@ def call_scalar(self, query: Any, context: Any = None) -> Any:
         because materializing a scalar-only fallback grid would be slower than
         that fused operator.
         """
-        if self._grid_pair_fn is not None and not getattr(self, "training", False):
+        if self.cuda_infer_l_1_grid_pair is not None and not getattr(
+            self, "training", False
+        ):
             return self._slice_scalar_layout(self.call(query, context))
         return self._forward(query, context, scalar_only=True)
 
@@ -1220,26 +1222,31 @@ def _pair_grid(self, left: Any, right: Any) -> Any | None:
         c_wide = left.shape[3] // self.n_frames
         if c_wide % 32 != 0 or left.shape != right.shape:
             return None
-        if getattr(self, "training", False):
+        training = getattr(self, "training", False)
+        if not training:
+            candidate = self.run_cute_infer_grid_pair(left, right)
+            if candidate is not None:
+                return candidate
+        if training:
             # Training form: frame-packed operands ride through unreshaped,
             # with analytic first and second order behind the call; under
             # autocast it runs the same reduced-precision regime as the dense
             # einsum composition it replaces.
-            if self._grid_pair_train_fn is None:
+            if self.triton_train_l_1_grid_pair is None:
                 return None
-            return self._grid_pair_train_fn(
+            return self.triton_train_l_1_grid_pair(
                 left,
                 right,
                 self.projector.to_grid_mat,
                 self._from_grid_t,
                 self.n_frames,
             )
-        if self._grid_pair_fn is None or left.shape[2] != 1:
+        if self.cuda_infer_l_1_grid_pair is None or left.shape[2] != 1:
             return None
         n_batch, coeff_dim = left.shape[0], left.shape[1]
         flat_p = coeff_dim * self.n_frames
         xp = array_api_compat.array_namespace(left, right)
-        out = self._grid_pair_fn(
+        out = self.cuda_infer_l_1_grid_pair(
             xp.reshape(left, (n_batch, flat_p, c_wide)),
             xp.reshape(right, (n_batch, flat_p, c_wide)),
             self.projector.to_grid_mat,
@@ -1247,6 +1254,24 @@ def _pair_grid(self, left: Any, right: Any) -> Any | None:
         )
         return xp.reshape(out, (n_batch, coeff_dim, 1, self.n_frames * c_wide))
 
+    def run_cute_infer_grid_pair(self, left: Any, right: Any) -> Any | None:
+        """Run a backend CuTe grid product when its contract matches.
+
+        Parameters
+        ----------
+        left : Any
+            Left coefficient operand with shape ``(N, D, F, n_frames * C)``.
+        right : Any
+            Right coefficient operand with the same shape as ``left``.
+
+        Returns
+        -------
+        Any or None
+            Coefficient result with the same shape as ``left``, or ``None``
+            when the backend has no eligible CuTe implementation.
+        """
+        return None
+
     def _project_pair_in_one_transform(
         self,
         left: Any,
diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py
index 323fe3bbfe..68907b49f9 100644
--- a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py
+++ b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py
@@ -1592,9 +1592,9 @@ def __init__(
         # attention layout without the optional focus-mix / value / output
         # projections (the deployed DPA4 configuration). Whichever of the
         # mutually exclusive inference gates is active then supplies the
-        # implementation, and ``self._flash_atten_fn`` being bound is what marks
+        # implementation, and ``self.flash_attention`` being bound is what marks
         # the fused path as live. The array-API reference leaves the hook unbound.
-        self._flash_atten_layout_ok = (
+        self.flash_attention_layout_supported = (
             self.n_atten_head > 0
             and self.mmax == 1
             and self.needs_local_frame
@@ -1604,9 +1604,10 @@ def __init__(
             and self.attn_o_proj is None
             and self.attn_focus_mix is None
         )
-        self._flash_atten_fn = None
-        self._cuda_conv_fn = None
-        self._cached_edge_csr_fn = None
+        self.flash_attention = None
+        self.flash_attention_supports_training = False
+        self.cuda_infer_l_2_conv = None
+        self.edge_csr = None
 
         # === Step 13. Optional fused SO(2) value-path seam ===
         # The fused value path folds the rotate-to-local projection, radial
@@ -1615,21 +1616,20 @@ def __init__(
         # array-API reference has no such kernel, so it never runs the fused
         # 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
+        self.triton_infer_l_2_value = None
+        self.cutile_infer_value = None
 
         # === Step 14. Optional fused training seams ===
         # Training differentiates the convolution twice under a force loss, so
         # its accelerated forms carry analytic backward and second-order
         # implementations of their own: one fused operator for the value stream
-        # up to the attention aggregation (``_cuda_value_train``), and the
+        # up to the attention aggregation (``cuda_train_value``), and the
         # segmented attention softmax / flash aggregation pair for the
-        # attention span (``_flash_atten_trains`` marks the bound aggregation
+        # attention span (``flash_attention_supports_training`` marks the bound
+        # aggregation
         # as training-capable). The array-API reference leaves every hook
         # unbound and trains through the dense expression.
-        self._cuda_value_train = None
-        self._flash_atten_trains = False
+        self.cuda_train_value = None
         self.trainable = bool(trainable)
 
     def call(
@@ -1773,23 +1773,25 @@ def forward_attention(
         # does not serve the bridging mode, whose source gate reshapes the
         # softmax normalization.
         training = getattr(self, "training", False)
-        run_cuda = (
-            self._cuda_conv_fn is not None
+        run_cuda_infer_l_2 = (
+            self.cuda_infer_l_2_conv is not None
             and not training
             and edge_cache.edge_src_gate is None
         )
-        run_flash = (
-            self._flash_atten_fn is not None
-            and (not training or self._flash_atten_trains)
-            and not run_cuda
+        run_flash_attention = (
+            self.flash_attention is not None
+            and (self.flash_attention_supports_training or not training)
+            and not run_cuda_infer_l_2
         )
-        if run_cuda:
-            return self.forward_attention_cuda(x, edge_cache, radial_feat, x_l0_node)
-        if run_flash:
+        if run_cuda_infer_l_2:
+            return self.forward_attention_cuda_infer_l_2(
+                x, edge_cache, radial_feat, x_l0_node
+            )
+        if run_flash_attention:
             return self.forward_attention_flash(x, edge_cache, radial_feat, x_l0_node)
         return self.forward_attention_dense(x, edge_cache, radial_feat, x_l0_node)
 
-    def forward_attention_cuda(
+    def forward_attention_cuda_infer_l_2(
         self,
         x: Array,
         edge_cache: EdgeCache,
@@ -1825,7 +1827,7 @@ def forward_attention_cuda(
         xp = array_api_compat.array_namespace(x)
 
         # === Step 1. Projected radial features ===
-        rad_feat = self._cuda_conv_fn.radial_features(
+        rad_feat = self.cuda_infer_l_2_conv.radial_features(
             radial_feat
         )  # (E, lmax+1, C_wide)
 
@@ -1836,7 +1838,9 @@ def forward_attention_cuda(
         head_gate = self.attention_head_gate(x_l0_node)  # (N, Fa, H)
 
         # === Step 4. Fused convolution ===
-        out = self._cuda_conv_fn(x, edge_cache, rad_feat, q_node, k_node, head_gate)
+        out = self.cuda_infer_l_2_conv(
+            x, edge_cache, rad_feat, q_node, k_node, head_gate
+        )
         return xp.astype(out, get_xp_precision(xp, self.precision))
 
     def forward_attention_flash(
@@ -1853,7 +1857,10 @@ def forward_attention_flash(
         kernel folds the block-diagonal rotate-back, the inverse-rotation
         rescale, the per-edge weighting and the destination reduction into a
         single atomic-free pass, so the transient rotate-back message and
-        weighted value tensors are never materialized.
+        weighted value tensors are never materialized. The value producer is
+        selected independently inside ``so2_message``: training uses
+        ``cuda_train_value`` when it is bound, while the flash operator remains
+        the aggregation consumer selected by the Triton training gate.
 
         Parameters
         ----------
@@ -1889,14 +1896,14 @@ def forward_attention_flash(
         # === Step 4. Fused rotate-back and weighted destination reduction ===
         # The destination CSR view is built once per step and shared by every
         # segment consumer of the graph.
-        if self._cached_edge_csr_fn is None:
+        if self.edge_csr is None:
             raise RuntimeError("The fused attention path requires a CSR builder")
         dst = edge_cache.dst
-        order, row_ptr = self._cached_edge_csr_fn(edge_cache, "dst", x.shape[0])
+        order, row_ptr = self.edge_csr(edge_cache, "dst", x.shape[0])
         rotation = edge_cache.Dt_full
         if rotation is None:
-            rotation = self._cuda_value_train.edge_runs(edge_cache)
-        pre_gate = self._flash_atten_fn(
+            rotation = self.cuda_train_value.edge_runs(edge_cache)
+        pre_gate = self.flash_attention(
             x_local,
             rotation,
             self.rotate_inv_rescale_full,
@@ -2349,22 +2356,22 @@ def so2_message(
         n_edge = src.shape[0]
 
         training = getattr(self, "training", False)
-        if self._cutile_value_path is not None and not training:
+        if self.cutile_infer_value is not None and not training:
             # === Steps 1-5 (fused cuTile operators). ``rotate_mix`` folds the
             # rotation and the radial degree mixing into one edge-parallel
             # kernel writing the focus-major layout; ``mixing_stack`` runs the
             # whole gated stack, keeping the inter-layer activations and the
             # gated-layer pre-activations off the traced graph entirely. ===
-            x_local, rad_feat = self._cutile_value_path(x, edge_cache, radial_feat)
-        elif self._cuda_value_train is not None and training:
+            x_local, rad_feat = self.cutile_infer_value(x, edge_cache, radial_feat)
+        elif self.cuda_train_value is not None and training:
             # === Steps 1-5 (one CUDA kernel, training). The whole value
             # stream up to the attention aggregation runs in a single launch
             # with analytic backward and second order; only the backward
             # anchors reach device memory. ===
-            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._cuda_value_train(x, edge_cache, radial_feat)
-        elif self._triton_value_path is not None and not training:
+            if self.edge_csr is not None:
+                self.edge_csr(edge_cache, "src", x.shape[0])
+            x_local, rad_feat = self.cuda_train_value(x, edge_cache, radial_feat)
+        elif self.triton_infer_l_2_value is not None and not 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
@@ -2372,15 +2379,9 @@ def so2_message(
             # final store, keeping the inter-layer activations off the traced
             # graph. The rotate-mix backward reduces through the source CSR
             # view, which is built once per step and kept on the edge cache. ===
-            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)
+            if self.edge_csr is not None:
+                self.edge_csr(edge_cache, "src", x.shape[0])
+            x_local, rad_feat = self.triton_infer_l_2_value(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)
diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
index 6895ce3e1d..24d9945f75 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
@@ -45,6 +45,7 @@
     NeighborGraph,
     append_isolated_frames,
     apply_pair_exclusion,
+    compact_edges,
     compact_nodes,
     expand_node_values,
     frame_id_from_n_node,
@@ -79,6 +80,7 @@
     "build_neighbor_graph_ase",
     "canonicalize_neighbor_graph",
     "center_edge_pairs",
+    "compact_edges",
     "compact_nodes",
     "edge_env_mat",
     "edge_force_virial",
diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py
index 7aca89f1b6..17509fc13e 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/graph.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py
@@ -315,6 +315,191 @@ def node_ownership_mask(n_node: Array, n_local: Array, n_total: int) -> Array:
     return index_in_frame < local_count
 
 
+def _compact_csr(
+    order: Array,
+    row_ptr: Array,
+    edge_mask: Array,
+    new_position: Array,
+) -> tuple[Array, Array]:
+    """Project an endpoint CSR view onto a compact edge axis.
+
+    Parameters
+    ----------
+    order
+        Edge permutation grouped by one endpoint with shape (E,).
+    row_ptr
+        CSR row pointers with shape (N + 1,).
+    edge_mask
+        Boolean survivor mask on the original edge axis with shape (E,).
+    new_position
+        Mapping from original edge indices to compact edge indices with shape
+        (E,). Removed edges map to -1.
+
+    Returns
+    -------
+    tuple[Array, Array]
+        The filtered edge permutation and rebuilt row pointers.
+    """
+    xp = array_api_compat.array_namespace(order, row_ptr, edge_mask, new_position)
+    ordered_keep = xp.take(edge_mask, order, axis=0)
+    (ordered_keep_position,) = xp.nonzero(ordered_keep)
+    kept_edge = xp.take(order, ordered_keep_position, axis=0)
+    compact_order = xp.take(new_position, kept_edge, axis=0)
+    compact_row_ptr = _compact_csr_row_ptr(ordered_keep, row_ptr)
+    return compact_order, compact_row_ptr
+
+
+def _compact_csr_row_ptr(
+    ordered_keep: Array,
+    row_ptr: Array,
+) -> Array:
+    """Rebuild CSR row pointers after filtering the ordered edge axis."""
+    xp = array_api_compat.array_namespace(ordered_keep, row_ptr)
+    prefix = xp.concat(
+        [
+            xp.zeros(
+                (1,),
+                dtype=row_ptr.dtype,
+                device=array_api_compat.device(row_ptr),
+            ),
+            xp.cumulative_sum(xp.astype(ordered_keep, row_ptr.dtype), axis=0),
+        ],
+        axis=0,
+    )
+    return xp.take(prefix, row_ptr, axis=0)
+
+
+def compact_edges(graph: NeighborGraph) -> NeighborGraph:
+    """Drop masked edges and remap derived indices onto the compact edge axis.
+
+    Complete destination and source CSR views are filtered in their own edge
+    order, remapped to the compact payload, and retained without sorting. This
+    preserves arbitrary mask placement inside CSR rows rather than assuming a
+    masked suffix.
+
+    Parameters
+    ----------
+    graph : NeighborGraph
+        Graph whose ``edge_mask`` selects the retained edge payload. Angle
+        fields must either both be present or both be absent.
+
+    Returns
+    -------
+    NeighborGraph
+        A graph containing only valid edges and angles. Complete CSR views are
+        retained; incomplete views are cleared.
+
+    Raises
+    ------
+    ValueError
+        If the angle fields are incomplete or have inconsistent shapes.
+    """
+    import dataclasses
+
+    xp = array_api_compat.array_namespace(graph.edge_mask)
+    has_angle_index = graph.angle_index is not None
+    has_angle_mask = graph.angle_mask is not None
+    if has_angle_index != has_angle_mask:
+        raise ValueError(
+            "compact_edges requires angle_index and angle_mask to both be set "
+            "or both be None"
+        )
+    if has_angle_index:
+        if graph.angle_index.ndim != 2 or graph.angle_index.shape[0] != 2:
+            raise ValueError(
+                "compact_edges requires angle_index with shape (2, A); got "
+                f"{tuple(graph.angle_index.shape)}"
+            )
+        if graph.angle_index.shape[1] != graph.angle_mask.shape[0]:
+            raise ValueError(
+                "compact_edges: angle_index and angle_mask disagree on A; got "
+                f"{graph.angle_index.shape[1]} and "
+                f"{graph.angle_mask.shape[0]}"
+            )
+
+    (keep_index,) = xp.nonzero(graph.edge_mask)
+    preserve_csr = all(
+        value is not None
+        for value in (
+            graph.destination_order,
+            graph.destination_row_ptr,
+            graph.source_order,
+            graph.source_row_ptr,
+        )
+    )
+    if preserve_csr or has_angle_index:
+        survivor = xp.astype(graph.edge_mask, graph.edge_index.dtype)
+        rank = xp.cumulative_sum(survivor, axis=0) - survivor
+        new_position = xp.where(
+            graph.edge_mask,
+            rank,
+            xp.full_like(rank, -1),
+        )
+    if preserve_csr:
+        if graph.destination_sorted:
+            destination_order = xp.arange(
+                keep_index.shape[0],
+                dtype=graph.destination_order.dtype,
+                device=array_api_compat.device(graph.destination_order),
+            )
+            destination_row_ptr = _compact_csr_row_ptr(
+                graph.edge_mask,
+                graph.destination_row_ptr,
+            )
+        else:
+            destination_order, destination_row_ptr = _compact_csr(
+                graph.destination_order,
+                graph.destination_row_ptr,
+                graph.edge_mask,
+                new_position,
+            )
+        source_order, source_row_ptr = _compact_csr(
+            graph.source_order,
+            graph.source_row_ptr,
+            graph.edge_mask,
+            new_position,
+        )
+    else:
+        destination_order = None
+        destination_row_ptr = None
+        source_order = None
+        source_row_ptr = None
+    fields = {
+        "edge_index": xp.take(graph.edge_index, keep_index, axis=1),
+        "edge_vec": xp.take(graph.edge_vec, keep_index, axis=0),
+        "edge_mask": xp.take(graph.edge_mask, keep_index, axis=0),
+        "destination_order": destination_order,
+        "destination_row_ptr": destination_row_ptr,
+        "source_order": source_order,
+        "source_row_ptr": source_row_ptr,
+        "destination_sorted": graph.destination_sorted and preserve_csr,
+    }
+    if has_angle_index:
+        # Angles address the original edge axis. An exclusive survivor count
+        # maps each retained edge to its compact position; angles touching a
+        # removed edge are discarded.
+        first = xp.take(new_position, graph.angle_index[0, :], axis=0)
+        second = xp.take(new_position, graph.angle_index[1, :], axis=0)
+        angle_keep = xp.logical_and(
+            xp.astype(graph.angle_mask, xp.bool),
+            xp.logical_and(first >= 0, second >= 0),
+        )
+        (angle_keep_index,) = xp.nonzero(angle_keep)
+        fields["angle_index"] = xp.stack(
+            [
+                xp.take(first, angle_keep_index, axis=0),
+                xp.take(second, angle_keep_index, axis=0),
+            ],
+            axis=0,
+        )
+        fields["angle_mask"] = xp.take(
+            graph.angle_mask,
+            angle_keep_index,
+            axis=0,
+        )
+    return dataclasses.replace(graph, **fields)
+
+
 def apply_pair_exclusion(
     graph: NeighborGraph,
     atype: Array,
@@ -387,61 +572,7 @@ def apply_pair_exclusion(
         destination_sorted=False,
     )
     if compact:
-        # Angle fields are a coupled pair (produced together by the angle
-        # builder): both present or both None. Fail fast on any inconsistent
-        # state — a partial or shape-mismatched pair is a caller bug that would
-        # otherwise remap silently wrong.
-        has_ai = out.angle_index is not None
-        has_am = out.angle_mask is not None
-        if has_ai != has_am:
-            raise ValueError(
-                "apply_pair_exclusion(compact=True): angle_index and angle_mask "
-                "must both be set or both be None; got "
-                f"angle_index={'set' if has_ai else 'None'}, "
-                f"angle_mask={'set' if has_am else 'None'}."
-            )
-        if has_ai:
-            if out.angle_index.ndim != 2 or out.angle_index.shape[0] != 2:
-                raise ValueError(
-                    "apply_pair_exclusion(compact=True): angle_index must have "
-                    f"shape (2, A); got {tuple(out.angle_index.shape)}."
-                )
-            if out.angle_index.shape[1] != out.angle_mask.shape[0]:
-                raise ValueError(
-                    "apply_pair_exclusion(compact=True): angle_index (2, A) and "
-                    f"angle_mask (A,) disagree on A: {out.angle_index.shape[1]} "
-                    f"vs {out.angle_mask.shape[0]}."
-                )
-        (keep_idx,) = xp.nonzero(out.edge_mask)
-        fields = {
-            "edge_index": out.edge_index[:, keep_idx],
-            "edge_vec": xp.take(out.edge_vec, keep_idx, axis=0),
-            "edge_mask": xp.take(out.edge_mask, keep_idx, axis=0),
-        }
-        if has_ai:
-            # Angles reference PRE-compaction edge positions; remap them to the
-            # compacted axis and drop any angle whose constituent edges were
-            # excluded. ``new_pos`` maps old edge position -> new position via an
-            # exclusive prefix sum over the survivors (-1 for dropped edges).
-            surv = xp.astype(out.edge_mask, out.edge_index.dtype)  # (E,) 0/1
-            rank = xp.cumulative_sum(surv, axis=0) - surv  # survivors before me
-            new_pos = xp.where(out.edge_mask, rank, xp.full_like(rank, -1))
-            a_new = xp.take(new_pos, out.angle_index[0, :], axis=0)
-            b_new = xp.take(new_pos, out.angle_index[1, :], axis=0)
-            both_survive = xp.logical_and(a_new >= 0, b_new >= 0)
-            angle_keep = xp.logical_and(
-                xp.astype(out.angle_mask, xp.bool), both_survive
-            )
-            (angle_keep_idx,) = xp.nonzero(angle_keep)
-            fields["angle_index"] = xp.stack(
-                [
-                    xp.take(a_new, angle_keep_idx, axis=0),
-                    xp.take(b_new, angle_keep_idx, axis=0),
-                ],
-                axis=0,
-            )
-            fields["angle_mask"] = xp.take(out.angle_mask, angle_keep_idx, axis=0)
-        out = dataclasses.replace(out, **fields)
+        out = compact_edges(out)
     return out
 
 
diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py
index b3cd33798a..26d30a5e75 100644
--- a/deepmd/pt/entrypoints/freeze_pt2.py
+++ b/deepmd/pt/entrypoints/freeze_pt2.py
@@ -379,7 +379,7 @@ def _tune_triton_configs(model: torch.nn.Module, target_device: torch.device) ->
     # tables, and must reflect the target device and any fresh registrations.
     for module in model.modules():
         if isinstance(module, SO2Convolution) and module.triton_infer_level >= 2:
-            module._triton_value_path = make_triton_value_path(module)
+            module.triton_infer_l_2_value = make_triton_value_path(module)
 
 
 # The trace-time sendlist for the with-comm artifact embeds the address of a
@@ -1005,16 +1005,16 @@ def _freeze_sezm_to_pt2(
     force_fused_scatter = target_device.type == "cuda"
     for module in model.modules():
         if isinstance(module, SO2Linear):
-            module._force_block_diag_matmul = force_block_diag
+            module.force_block_diag_matmul = force_block_diag
         if isinstance(module, GeometricInitialEmbedding):
-            module._force_fused_scatter = force_fused_scatter
+            module.force_cuda_infer_l_1_scatter = force_fused_scatter
 
     # Sweep any Triton launch-table keys this checkpoint needs that are not
     # covered for the local GPU, so the traced graph bakes tuned launches.
     _tune_triton_configs(model, target_device)
 
     has_triton_value_path = any(
-        getattr(module, "_triton_value_path", None) is not None
+        getattr(module, "triton_infer_l_2_value", None) is not None
         for module in model.modules()
     )
     if has_triton_value_path:
diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py
index f468379565..24457cd300 100644
--- a/deepmd/pt/infer/deep_eval.py
+++ b/deepmd/pt/infer/deep_eval.py
@@ -782,18 +782,19 @@ def _eval_lower_strategy(
                 list(inner.get_sel()),
                 return_mode="edges",
             )
-            predict = inner.forward_common_lower(
-                edge_schema.coord,
-                edge_schema.atype,
-                edge_schema.edge_index,
-                edge_schema.edge_vec,
-                edge_schema.edge_scatter_index,
-                edge_schema.edge_mask,
-                fparam=fparam,
-                aparam=aparam,
-                charge_spin=charge_spin,
-                input_prec=coord.dtype,
-            )
+            with self.dp._frozen_parameter_context():
+                predict = inner.forward_common_lower(
+                    edge_schema.coord,
+                    edge_schema.atype,
+                    edge_schema.edge_index,
+                    edge_schema.edge_vec,
+                    edge_schema.edge_scatter_index,
+                    edge_schema.edge_mask,
+                    fparam=fparam,
+                    aparam=aparam,
+                    charge_spin=charge_spin,
+                    input_prec=coord.dtype,
+                )
         else:
             ext_coord, ext_atype, nlist, mapping = self._nlist_builder.build(
                 coord,
@@ -802,16 +803,17 @@ def _eval_lower_strategy(
                 self.rcut,
                 list(inner.get_sel()),
             )
-            model_lower = inner.forward_common_lower(
-                ext_coord,
-                ext_atype,
-                nlist,
-                mapping,
-                fparam=fparam,
-                aparam=aparam,
-                do_atomic_virial=do_atomic_virial,
-                charge_spin=charge_spin,
-            )
+            with self.dp._frozen_parameter_context():
+                model_lower = inner.forward_common_lower(
+                    ext_coord,
+                    ext_atype,
+                    nlist,
+                    mapping,
+                    fparam=fparam,
+                    aparam=aparam,
+                    do_atomic_virial=do_atomic_virial,
+                    charge_spin=charge_spin,
+                )
             predict = communicate_extended_output(
                 model_lower,
                 self.output_def,
diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index 9e565bfb40..606956e284 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -65,6 +65,10 @@
 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.cute.sezm.so2.metadata import (
+    build_sorted_edge_index_metadata,
+)
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
     use_amp_infer,
@@ -1071,19 +1075,19 @@ def __init__(
         # The fused convolution paths consume only the three structural rows of
         # each Wigner degree block. Source-gated attention bypasses that fused
         # convolution, so its dense per-edge rotations remain available.
-        self._wigner_free_conv = (
+        self.cuda_infer_l_2_covers_all_blocks = (
             self.bridging_switch is None
             and bool(self.blocks)
             and all(
-                getattr(block.so2_conv, "_cuda_conv_fn", None) is not None
-                and not block.so2_conv._cuda_conv_fn._compete
+                getattr(block.so2_conv, "cuda_infer_l_2_conv", None) is not None
+                and not block.so2_conv.cuda_infer_l_2_conv.focus_compete
                 for block in self.blocks
             )
         )
-        self._packed_wigner_train = bool(self.blocks) and all(
-            getattr(block.so2_conv, "_cuda_value_train", None) is not None
-            and block.so2_conv._flash_atten_fn is not None
-            and block.so2_conv._flash_atten_trains
+        self.cuda_train_covers_all_blocks = bool(self.blocks) and all(
+            getattr(block.so2_conv, "cuda_train_value", None) is not None
+            and block.so2_conv.flash_attention is not None
+            and block.so2_conv.flash_attention_supports_training
             for block in self.blocks
         )
 
@@ -1091,8 +1095,8 @@ def __init__(
         # distance and are cheap enough that the compiler inlines them into
         # every consumer and re-evaluates them there. Behind an operator
         # boundary the chain runs once per step.
-        self._cuda_radial_fn = None
-        self._cuda_wigner_fn = None
+        self.cuda_infer_l_1_radial = None
+        self.cuda_infer_l_1_wigner = None
         if cuda_infer_level() >= 1:
             from deepmd.pt_expt.kernels.cuda.dpa4.edge_radial import (
                 make_cuda_edge_radial,
@@ -1101,13 +1105,13 @@ def __init__(
                 make_cuda_wigner_dense,
             )
 
-            self._cuda_radial_fn = make_cuda_edge_radial(
+            self.cuda_infer_l_1_radial = make_cuda_edge_radial(
                 self.edge_envelope, self.radial_basis
             )
             # The dense Wigner pair otherwise costs five full-size passes
             # over the (E, D, D) tensors; the fused build pays only the
             # output writes.
-            self._cuda_wigner_fn = make_cuda_wigner_dense(
+            self.cuda_infer_l_1_wigner = make_cuda_wigner_dense(
                 self.mp_init_lmax, self.compute_dtype
             )
 
@@ -1198,6 +1202,7 @@ def forward(
         force_embedding: torch.Tensor | None = None,
         charge_spin: torch.Tensor | None = None,
         spin: torch.Tensor | None = None,
+        edge_index_sorted_by_dst: bool = False,
     ) -> tuple[
         torch.Tensor,
         torch.Tensor,
@@ -1239,6 +1244,10 @@ def forward(
             initial SO(3) backbone state before the interaction blocks.
         charge_spin
             Frame-level charge and spin conditions with shape (nf, 2).
+        spin
+            Optional per-atom spin vectors.
+        edge_index_sorted_by_dst
+            Host-side provenance that ``edge_index[1]`` is nondecreasing.
 
         Returns
         -------
@@ -1274,6 +1283,7 @@ def forward(
                 edge_index=edge_index,
                 edge_vec=edge_vec,
                 edge_mask=edge_mask,
+                edge_index_sorted_by_dst=edge_index_sorted_by_dst,
                 force_embedding=force_embedding,
                 charge_spin=charge_spin,
                 spin=spin,
@@ -1334,6 +1344,11 @@ def forward(
         # sparse-edge path, so ``forward`` keeps the original
         # bridging-free aggregation semantics.
         with nvtx_range("build_edge_cache"):
+            packed_wigner_candidate = self.is_cute_infer_packed_wigner_candidate(
+                type_ebed.device,
+                type_ebed.dtype,
+                extended_coord.dtype,
+            )
             edge_cache = build_edge_cache(
                 type_ebed=type_ebed,
                 extended_coord=extended_coord,
@@ -1350,7 +1365,8 @@ def forward(
                 # the model is roll-equivariant, so inference fixes gamma.
                 random_gamma=self.random_gamma and self.training,
                 wigner_calc=self.wigner_calc,
-                build_wigner=self._build_full_wigner(),
+                packed_wigner_candidate=packed_wigner_candidate,
+                build_wigner=(self._build_full_wigner() or packed_wigner_candidate),
             )
 
         ebed_dim_0 = self.node_init_dim  # (node_init_lmax+1)^2
@@ -1428,28 +1444,41 @@ def forward(
 
         # === Step 10. Fuse edge type features into radial features (fp32+) ===
         with nvtx_range("radial_fuse"):
+            block_dtype = (
+                torch.float32 if edge_cache.D_packed is not None else self.dtype
+            )
             radial_feat = radial_feat + rearrange(
                 edge_cache.edge_type_feat, "E C -> E 1 C"
             )
-            radial_feat = radial_feat.to(dtype=self.dtype)
+            radial_feat = radial_feat.to(dtype=block_dtype)
             rad_feat_per_block = [
                 radial_feat[:, :rad_len, :] for rad_len in self.rad_sizes_per_block
             ]  # list of (E, lmax+1, C)
 
-        # === Step 11. Convert to self.dtype and run blocks ===
+        # === Step 11. Convert to the block runtime dtype and run blocks ===
         # The block stage is skipped entirely for the zero-block descriptor,
         # sparing the working edge-cache dtype cast that only the blocks consume.
         # A frame without valid edges takes the same path as any other, so an
         # isolated atom is one function of its features whether or not the
         # frame holds other edges.
         with nvtx_range("blocks"):
-            x = x.to(dtype=self.dtype)  # (N, D, 1, C)
+            x = x.to(dtype=block_dtype)  # (N, D, 1, C)
             if force_embedding is not None:
-                x = x + force_embedding.to(dtype=self.dtype)
+                x = x + force_embedding.to(dtype=block_dtype)
             if self.blocks:
-                edge_cache = edge_cache_to_dtype(edge_cache, self.dtype)
+                edge_cache = edge_cache_to_dtype(edge_cache, block_dtype)
+                edge_cache = edge_cache._replace(
+                    cute_infer_so2_metadata=self.prepare_cute_infer_so2_metadata(
+                        edge_cache,
+                        n_nodes,
+                    )
+                )
                 with self._compute_mode_ctx(extended_coord.device):
-                    x = self._forward_blocks(x, edge_cache, rad_feat_per_block)
+                    x = self._forward_blocks(
+                        x,
+                        edge_cache,
+                        rad_feat_per_block,
+                    )
 
         # === Step 12. Final l=0 output mixing ===
         with nvtx_range("output_ffn"):
@@ -1475,6 +1504,7 @@ def forward_with_edges(
         edge_index: torch.Tensor,
         edge_vec: torch.Tensor,
         edge_mask: torch.Tensor,
+        edge_index_sorted_by_dst: bool = False,
         force_embedding: torch.Tensor | None = None,
         charge_spin: torch.Tensor | None = None,
         spin: torch.Tensor | None = None,
@@ -1509,6 +1539,8 @@ def forward_with_edges(
             Edge vectors with shape (E, 3) in Å.
         edge_mask
             Edge mask with shape (E,).
+        edge_index_sorted_by_dst
+            Host-side provenance that ``edge_index[1]`` is nondecreasing.
         force_embedding
             Optional precomputed equivariant force embedding with shape
             ``(nf * nloc, D, 1, channels)``, where
@@ -1628,6 +1660,11 @@ def forward_with_edges(
             )
         # === Step 3. Build edge cache once (sparse edges) ===
         with nvtx_range("build_edge_cache"):
+            packed_wigner_candidate = self.is_cute_infer_packed_wigner_candidate(
+                type_ebed.device,
+                type_ebed.dtype,
+                extended_coord.dtype,
+            )
             edge_cache = build_edge_cache_from_edges(
                 type_ebed=type_ebed,
                 atype_flat=atype_flat,
@@ -1643,15 +1680,17 @@ def forward_with_edges(
                 bridging_switch=self.bridging_switch,
                 edge_envelope=self.edge_envelope,
                 radial_basis=self.radial_basis,
-                fused_radial=(None if self.training else self._cuda_radial_fn),
-                fused_wigner=(None if self.training else self._cuda_wigner_fn),
+                fused_radial=(None if self.training else self.cuda_infer_l_1_radial),
+                fused_wigner=(None if self.training else self.cuda_infer_l_1_wigner),
                 has_exclude_types=bool(self.exclude_types),
                 edge_type_keep_mask=self._edge_type_keep_mask,
                 # Random local-Z roll is a training-only augmentation;
                 # the model is roll-equivariant, so inference fixes gamma.
                 random_gamma=self.random_gamma and self.training,
                 wigner_calc=self.wigner_calc,
-                build_wigner=self._build_full_wigner(),
+                packed_wigner_candidate=packed_wigner_candidate,
+                destinations_sorted=edge_index_sorted_by_dst,
+                build_wigner=(self._build_full_wigner() or packed_wigner_candidate),
                 node_partial_exchange=node_partial_exchange,
             )
 
@@ -1725,26 +1764,38 @@ def forward_with_edges(
 
         # === Step 9. Fuse edge type features into radial features (fp32+) ===
         with nvtx_range("radial_fuse"):
-            radial_feat = radial_feat.to(dtype=self.dtype)
+            block_dtype = (
+                torch.float32 if edge_cache.D_packed is not None else self.dtype
+            )
+            radial_feat = radial_feat.to(dtype=block_dtype)
             radial_feat = radial_feat + rearrange(
-                edge_cache.edge_type_feat.to(dtype=self.dtype), "E C -> E 1 C"
+                edge_cache.edge_type_feat.to(dtype=block_dtype), "E C -> E 1 C"
             )
             rad_feat_per_block = [
                 radial_feat[:, :rad_len, :] for rad_len in self.rad_sizes_per_block
             ]
 
-        # === Step 10. Convert to self.dtype and run blocks ===
+        # === Step 10. Convert to the block runtime dtype and run blocks ===
         # The block stage is skipped entirely for the zero-block descriptor,
         # sparing the working edge-cache dtype cast that only the blocks consume.
         with nvtx_range("blocks"):
-            x = x.to(dtype=self.dtype)  # (N, D, 1, C)
+            x = x.to(dtype=block_dtype)  # (N, D, 1, C)
             if force_embedding is not None:
-                x = x + force_embedding.to(dtype=self.dtype)
+                x = x + force_embedding.to(dtype=block_dtype)
             if self.blocks:
-                edge_cache = edge_cache_to_dtype(edge_cache, self.dtype)
+                edge_cache = edge_cache_to_dtype(edge_cache, block_dtype)
+                edge_cache = edge_cache._replace(
+                    cute_infer_so2_metadata=self.prepare_cute_infer_so2_metadata(
+                        edge_cache,
+                        n_nodes,
+                    )
+                )
                 with self._compute_mode_ctx(extended_coord.device):
                     x = self._forward_blocks(
-                        x, edge_cache, rad_feat_per_block, comm_dict=comm_dict
+                        x,
+                        edge_cache,
+                        rad_feat_per_block,
+                        comm_dict=comm_dict,
                     )
 
         # === Step 11. Keep the owned-atom rows for the read-out ===
@@ -1781,6 +1832,73 @@ def forward_with_edges(
         )
         return descriptor, latent.contiguous(), vacuum
 
+    @torch.jit.unused
+    def run_cute_infer_readout(
+        self,
+        ffn_in: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Run the CuTe readout when its exact inference contract matches.
+
+        Parameters
+        ----------
+        ffn_in : torch.Tensor
+            Equivariant readout input with shape ``(N, D, 1, C)``.
+
+        Returns
+        -------
+        torch.Tensor or None
+            Residual-inclusive scalar output with shape ``(N, C)``, or ``None``
+            when the exact CuTe contract is not satisfied.
+        """
+        if self.training or not cute_runtime_policy.is_cute_infer_enabled():
+            return None
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import (
+            maybe_run_neo_readout_l0,
+        )
+
+        return maybe_run_neo_readout_l0(
+            self.output_ffn,
+            ffn_in,
+        )
+
+    def prepare_cute_infer_so2_metadata(
+        self,
+        edge_cache: EdgeFeatureCache,
+        n_nodes: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
+        """Build the edge metadata consumed by the CuTe SO2 implementation.
+
+        Parameters
+        ----------
+        edge_cache : EdgeFeatureCache
+            Per-forward cache carrying the destination-major edge payload.
+        n_nodes : int
+            Number of nodes addressed by the edge payload.
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor] or None
+            Destination row pointers, source order, and source row pointers, or
+            ``None`` when the exact CuTe contract is not satisfied.
+        """
+        if torch.jit.is_scripting():
+            return None
+        if (
+            self.training
+            or not edge_cache.destinations_sorted
+            or edge_cache.D_packed is None
+            or edge_cache.edge_src_gate is not None
+        ):
+            return None
+        if not cute_runtime_policy.is_cute_infer_enabled():
+            return None
+        return build_sorted_edge_index_metadata(
+            edge_cache.src,
+            edge_cache.dst,
+            n_nodes,
+            validate_sorted=cute_runtime_policy.is_cute_strict_enabled(),
+        )
+
     def _forward_blocks(
         self,
         x: torch.Tensor,
@@ -1897,6 +2015,49 @@ def node_l0_extractor(v: torch.Tensor) -> torch.Tensor:
         ).to(dtype=self.dtype)
         return x
 
+    def is_cute_infer_packed_wigner_candidate(
+        self,
+        device: torch.device,
+        dtype: torch.dtype,
+        geometry_dtype: torch.dtype,
+    ) -> bool:
+        """Return whether all blocks satisfy the packed CuTe SO2 contract.
+
+        Parameters
+        ----------
+        device : torch.device
+            Device that executes the descriptor blocks.
+        dtype : torch.dtype
+            Descriptor block dtype.
+        geometry_dtype : torch.dtype
+            Edge-geometry compute dtype.
+
+        Returns
+        -------
+        bool
+            Whether packed Wigner storage can replace dense storage for every
+            interaction block.
+        """
+        from deepmd.pt_expt.kernels.cute.sezm.so2.operation import (
+            is_packed_wigner_candidate,
+        )
+
+        return is_packed_wigner_candidate(
+            blocks=self.blocks,
+            training=self.training,
+            device=device,
+            dtype=dtype,
+            producer_modules=(
+                self.radial_basis,
+                self.radial_embedding,
+                self.edge_envelope,
+                *((self.inner_clamp,) if self.inner_clamp is not None else ()),
+                *((self.bridging_switch,) if self.bridging_switch is not None else ()),
+            ),
+            producer_dtypes=(dtype, geometry_dtype),
+            has_edge_src_gate=self.bridging_switch is not None,
+        )
+
     def _apply_readout(self, x: torch.Tensor, n_rows: int) -> torch.Tensor:
         """Fold the node tensor into the scalar (``l=0``) descriptor.
 
@@ -1935,6 +2096,14 @@ def _apply_readout(self, x: torch.Tensor, n_rows: int) -> torch.Tensor:
             x_ro = x[:, : self.node_readout_dim, :, :].to(dtype=self.compute_dtype)
         for layer in self.readout_pre_layers:
             x_ro = x_ro + layer(x_ro)
+        if (
+            not self.readout_pre_layers
+            and self.so3_readout != "none"
+            and not torch.jit.is_scripting()
+        ):
+            accelerated = self.run_cute_infer_readout(x_ro)
+            if accelerated is not None:
+                return accelerated.reshape(n_rows, 1, 1, self.channels)
         if self.so3_readout == "none":
             return (x_ro + self.output_ffn(x_ro))[:, 0:1, :, :]
         return x_ro[:, 0:1, :, :] + self.output_ffn.forward_scalar(x_ro)
@@ -1969,8 +2138,8 @@ def _build_full_wigner(self) -> bool:
         if not self._need_full_wigner:
             return False
         if self.training:
-            return not self._packed_wigner_train
-        return not self._wigner_free_conv
+            return not self.cuda_train_covers_all_blocks
+        return not self.cuda_infer_l_2_covers_all_blocks
 
     def _shared_wigner_runs(
         self,
@@ -2004,13 +2173,13 @@ def _shared_wigner_runs(
         if edge_cache.csr_cache is None:
             return None
         if self.training:
-            if not self._packed_wigner_train:
+            if not self.cuda_train_covers_all_blocks:
                 return None
-            fused = self.blocks[0].so2_conv._cuda_value_train
+            fused = self.blocks[0].so2_conv.cuda_train_value
         else:
-            if not self._wigner_free_conv:
+            if not self.cuda_infer_l_2_covers_all_blocks:
                 return None
-            fused = self.blocks[0].so2_conv._cuda_conv_fn
+            fused = self.blocks[0].so2_conv.cuda_infer_l_2_conv
         if fused is None or lmax > self.lmax:
             return None
         return fused.edge_runs(edge_cache)[:, 1 : (lmax + 1) ** 2]
@@ -2031,6 +2200,8 @@ def _build_gie_zonal_coupling(
             the blocks are skipped (all-Cartesian model) the full coupling is
             reconstructed from the edge quaternion via the m=0-only path.
         """
+        if edge_cache.D_packed is not None:
+            return self.build_cute_infer_zonal_coupling(edge_cache)
         if edge_cache.Dt_full is None:
             calc = self.gie_zonal_wigner_calc or self.wigner_calc
             shared = self._shared_wigner_runs(edge_cache, calc.lmax)
@@ -2053,6 +2224,39 @@ def _build_gie_zonal_coupling(
         )
         return torch.cat([mp_coupling, extra_coupling], dim=1)
 
+    def build_cute_infer_zonal_coupling(
+        self,
+        edge_cache: EdgeFeatureCache,
+    ) -> torch.Tensor:
+        """Extract the GIE zonal coupling from packed Wigner storage.
+
+        Parameters
+        ----------
+        edge_cache : EdgeFeatureCache
+            Per-forward cache carrying CuTe packed Wigner storage.
+
+        Returns
+        -------
+        torch.Tensor
+            Zonal coupling with shape ``(E, D_node - 1)``.
+
+        Raises
+        ------
+        ValueError
+            If the edge cache does not carry packed Wigner storage.
+        """
+        D_packed = edge_cache.D_packed
+        if D_packed is None:
+            raise ValueError("CuTe zonal coupling requires packed Wigner storage")
+        mp_coupling = D_packed.index_select(1, self.gie.packed_zonal_offsets)
+        if self.gie_zonal_wigner_calc is None:
+            return mp_coupling
+        extra_coupling = self.gie_zonal_wigner_calc.forward_zonal(
+            self._edge_quaternion(edge_cache),
+            lmin=self.lmax + 1,
+        )
+        return torch.cat([mp_coupling, extra_coupling], dim=1)
+
     def _apply_charge_spin_embedding(
         self,
         type_ebed: torch.Tensor,
diff --git a/deepmd/pt/model/descriptor/sezm_nn/block.py b/deepmd/pt/model/descriptor/sezm_nn/block.py
index d26f6797df..e09038a3b2 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/block.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/block.py
@@ -713,7 +713,13 @@ def forward(
             - full AttnRes path returns `(block_output, None, so2_unit_output, ffn_unit_outputs)`
             - block AttnRes path returns `(block_output, block_summary, None, None)`
         """
-        return self._forward_impl(x, edge_cache, radial_feat, unit_history, comm_dict)
+        return self._forward_impl(
+            x,
+            edge_cache,
+            radial_feat,
+            unit_history,
+            comm_dict,
+        )
 
     def _extract_l0_from_canonical(self, value: torch.Tensor) -> torch.Tensor:
         """
@@ -793,7 +799,11 @@ def _run_so2_unit(
                 use_reentrant=False,
                 preserve_rng_state=True,
             )
-        return self._run_so2_unit_impl(x, edge_cache, radial_feat)
+        return self._run_so2_unit_impl(
+            x,
+            edge_cache,
+            radial_feat,
+        )
 
     def _run_so2_unit_impl(
         self,
@@ -802,6 +812,45 @@ def _run_so2_unit_impl(
         radial_feat: torch.Tensor,
     ) -> torch.Tensor:
         """Run the SO(2) unit implementation."""
+        if edge_cache.D_packed is not None:
+            from deepmd.pt_expt.kernels.cute.sezm.so2.operation import (
+                maybe_run_cute_so2,
+            )
+
+            metadata = edge_cache.cute_infer_so2_metadata
+            destination_row_ptr, source_order, source_row_ptr = (
+                (None, None, None) if metadata is None else metadata
+            )
+            accelerated = maybe_run_cute_so2(
+                self,
+                x,
+                edge_cache,
+                radial_feat,
+                dst_ptr=destination_row_ptr,
+                source_order=source_order,
+                source_ptr=source_row_ptr,
+            )
+            if accelerated is not None:
+                return accelerated
+            from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import (
+                dense_wigner_for_fallback,
+            )
+
+            if edge_cache.edge_quat is None:
+                raise ValueError("packed Wigner fallback requires edge quaternions")
+            d_full, dt_full = dense_wigner_for_fallback(
+                edge_cache.edge_quat,
+                lmax=self.lmax,
+                eps=self.so2_conv.eps,
+            )
+            edge_cache = edge_cache._replace(
+                D_full=d_full,
+                Dt_full=dt_full,
+                D_packed=None,
+                D_to_m_cache=None,
+                Dt_from_m_cache=None,
+            )
+
         n_node = x.shape[0]
         channels = self.channels
         use_full_node = self.node_lmax == self.lmax
@@ -892,7 +941,12 @@ def _forward_with_residual_shortcuts(
             Tuple `(block_output, None, None, None)`.
         """
         with nvtx_range("so2_conv"):
-            so2_unit_output = self._run_so2_unit(x, edge_cache, radial_feat, comm_dict)
+            so2_unit_output = self._run_so2_unit(
+                x,
+                edge_cache,
+                radial_feat,
+                comm_dict,
+            )
             so2_state = x + so2_unit_output
 
         with nvtx_range("ffn"):
@@ -950,7 +1004,10 @@ def _forward_with_full_attn_res(
                     current_x=x,
                 )
             so2_unit_output = self._run_so2_unit(
-                so2_input, edge_cache, radial_feat, comm_dict
+                so2_input,
+                edge_cache,
+                radial_feat,
+                comm_dict,
             )
 
         with nvtx_range("ffn"):
@@ -1018,7 +1075,10 @@ def _forward_with_block_attn_res(
                     current_x=x,
                 )
             so2_unit_output = self._run_so2_unit(
-                so2_input, edge_cache, radial_feat, comm_dict
+                so2_input,
+                edge_cache,
+                radial_feat,
+                comm_dict,
             )
 
         with nvtx_range("ffn"):
diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
index 1953c2f620..9d3dabaf2f 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
@@ -25,6 +25,10 @@
     rearrange,
 )
 
+from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import (
+    is_cute_infer_enabled,
+)
+
 from .utils import (
     get_promoted_dtype,
     nvtx_range,
@@ -73,10 +77,11 @@ class EdgeFeatureCache(NamedTuple):
     inv_sqrt_deg
         Inverse square root smooth degree normalization with shape (N, 1, 1).
     D_full
-        Block-diagonal Wigner-D matrix with shape (E, D, D) where D=(lmax+1)^2.
-        Used for efficient batched rotation. None if not available.
+        Block-diagonal Wigner-D matrix with shape (E, D, D). None when dense
+        Wigner storage is skipped.
     Dt_full
-        Transpose of D_full with shape (E, D, D). None if not available.
+        Transpose of D_full with shape (E, D, D). None when dense Wigner
+        storage is skipped.
     edge_quat
         Per-edge global-to-local quaternion actually used to build ``D_full`` and
         ``Dt_full`` with shape (E, 4). Includes the optional random local-Z roll.
@@ -90,6 +95,9 @@ class EdgeFeatureCache(NamedTuple):
         Lazy cache for endpoint CSR views used by segmented accelerated
         operators, keyed by endpoint role (``"dst"`` or ``"src"``). Built once
         per step and shared by every consumer.
+    cute_infer_so2_metadata
+        Optional ``(destination_row_ptr, source_order, source_row_ptr)`` tuple
+        prepared once for the destination-sorted CuTe SO2 path.
     edge_src_gate
         Optional per-edge Source Freeze Propagation Gate (SFPG) weight with
         shape (E, 1). Equals ``eta[src]`` where
@@ -101,6 +109,12 @@ class EdgeFeatureCache(NamedTuple):
         by this gate to forbid any node whose local neighborhood enters
         the frozen zone from propagating information along its outgoing
         edges.
+    destinations_sorted
+        Host-side provenance indicating that ``dst`` is nondecreasing. CuTe SO2
+        may only consume caches carrying this guarantee.
+    D_packed
+        Opt-in Neo packed Wigner panel with shape (E, 46). This is kept
+        separate so the public dense fields have a stable rank.
     """
 
     src: torch.Tensor
@@ -118,6 +132,23 @@ class EdgeFeatureCache(NamedTuple):
     csr_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None
     edge_src_gate: torch.Tensor | None = None
     edge_quat: torch.Tensor | None = None
+    cute_infer_so2_metadata: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = (
+        None
+    )
+    destinations_sorted: bool = False
+    D_packed: torch.Tensor | None = None
+
+
+def _separate_packed_wigner(
+    D_full: torch.Tensor | None,
+    Dt_full: torch.Tensor | None,
+) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
+    """Move the opt-in 46-value panel out of the dense Wigner fields."""
+    if D_full is None or D_full.dim() != 2:
+        return D_full, Dt_full, None
+    if Dt_full is not D_full:
+        raise RuntimeError("packed Wigner forward and transpose must share storage")
+    return None, None, D_full
 
 
 def cached_edge_csr(
@@ -298,6 +329,7 @@ def build_edge_cache(
     radial_basis: Callable[[torch.Tensor], torch.Tensor],
     random_gamma: bool,
     wigner_calc: WignerCalculatorFn,
+    packed_wigner_candidate: bool = False,
     build_wigner: bool = True,
 ) -> EdgeFeatureCache:
     """
@@ -358,6 +390,11 @@ def build_edge_cache(
     wigner_calc
         Callable that converts edge-aligned quaternions into packed Wigner-D
         blocks.
+    packed_wigner_candidate
+        Whether descriptor-level strict-FP32 SO2 checks passed. Concrete edge
+        count and destination ordering are checked before panel generation.
+    build_wigner
+        Whether to materialize Wigner-D blocks for the SO(2) path.
 
     Returns
     -------
@@ -413,7 +450,18 @@ def build_edge_cache(
             random_gamma=random_gamma,
             wigner_calc=wigner_calc,
             build_full=build_wigner,
+            packed_wigner=(
+                build_wigner
+                and _packed_wigner_edges_eligible(
+                    packed_wigner_candidate,
+                    edge_count=dst.numel(),
+                    node_count=n_nodes,
+                    destinations_sorted=True,
+                    runtime_dtypes=(edge_vec.dtype, edge_env.dtype, edge_rbf.dtype),
+                )
+            ),
         )  # (E, D, D), (E, D, D), (E, 4)
+        D_full, Dt_full, D_packed = _separate_packed_wigner(D_full, Dt_full)
 
     edge_type_feat = build_edge_type_feat(type_ebed, src, dst)  # (E, C)
 
@@ -427,8 +475,10 @@ def build_edge_cache(
         edge_env=edge_env,
         D_full=D_full,
         Dt_full=Dt_full,
+        D_packed=D_packed,
         edge_quat=edge_quat,
         deg_norm_floor=deg_norm_floor,
+        destinations_sorted=True,
     )
 
 
@@ -451,6 +501,8 @@ def build_edge_cache_from_edges(
     edge_type_keep_mask: EdgeTypeKeepMaskFn,
     random_gamma: bool,
     wigner_calc: WignerCalculatorFn,
+    packed_wigner_candidate: bool = False,
+    destinations_sorted: bool = False,
     build_wigner: bool = True,
     node_partial_exchange: Callable[[torch.Tensor], torch.Tensor] | None = None,
     fused_radial: FusedRadialFn | None = None,
@@ -504,6 +556,13 @@ def build_edge_cache_from_edges(
     wigner_calc
         Callable that converts edge-aligned quaternions into packed Wigner-D
         blocks.
+    packed_wigner_candidate
+        Whether descriptor-level strict-FP32 SO2 checks passed. Concrete edge
+        count and destination ordering are checked before panel generation.
+    destinations_sorted
+        Host-side provenance that ``edge_index[1]`` is nondecreasing.
+    build_wigner
+        Whether to materialize Wigner-D blocks for the SO(2) path.
     fused_wigner
         Optional fused replacement of ``wigner_calc`` that builds the packed
         pair in one kernel pass.
@@ -516,6 +575,11 @@ def build_edge_cache_from_edges(
     n_nodes = type_ebed.shape[0]
     src = edge_index[0].to(dtype=torch.long)
     dst = edge_index[1].to(dtype=torch.long)
+    if is_cute_infer_enabled() and destinations_sorted and dst.numel() > 1:
+        torch._assert_async(
+            torch.all(dst[1:] >= dst[:-1]),
+            "destinations_sorted=True requires nondecreasing destination indices",
+        )
 
     # === Step 1. Normalize mask and apply type exclusions ===
     edge_keep = edge_mask.to(dtype=torch.bool)
@@ -553,8 +617,20 @@ def build_edge_cache_from_edges(
             eps=eps,
             random_gamma=random_gamma,
             wigner_calc=fused_wigner if fused_wigner is not None else wigner_calc,
+            packed_wigner=(
+                build_wigner
+                and bridging_switch is None
+                and _packed_wigner_edges_eligible(
+                    packed_wigner_candidate,
+                    edge_count=dst.numel(),
+                    node_count=n_nodes,
+                    destinations_sorted=destinations_sorted,
+                    runtime_dtypes=(edge_vec.dtype, edge_env.dtype, edge_rbf.dtype),
+                )
+            ),
             build_full=build_wigner,
         )  # (E, D, D), (E, D, D), (E, 4)
+        D_full, Dt_full, D_packed = _separate_packed_wigner(D_full, Dt_full)
 
     # === Step 5. Edge type features ===
     edge_type_feat = build_edge_type_feat(type_ebed, src, dst)
@@ -587,9 +663,11 @@ def build_edge_cache_from_edges(
         edge_env=edge_env,
         D_full=D_full,
         Dt_full=Dt_full,
+        D_packed=D_packed,
         edge_quat=edge_quat,
         deg_norm_floor=deg_norm_floor,
         edge_src_gate=edge_src_gate,
+        destinations_sorted=destinations_sorted,
     )
 
 
@@ -601,6 +679,7 @@ def _build_edge_wigner(
     random_gamma: bool,
     wigner_calc: WignerCalculatorFn,
     build_full: bool = True,
+    packed_wigner: bool = False,
 ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]:
     """
     Build packed Wigner-D blocks from edge vectors.
@@ -623,13 +702,14 @@ def _build_edge_wigner(
         False (all message-passing blocks take the Cartesian path), only the
         quaternion is returned and the blocks are ``None``; the geometric
         initial embedding reconstructs the zonal coupling from the quaternion.
+    packed_wigner
+        Whether the exact packed SO2 eligibility contract has passed.
 
     Returns
     -------
     tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]
-        Packed Wigner-D matrices ``(D_full, Dt_full)`` with shape ``(E, D, D)``
-        (or ``None`` when ``build_full`` is False) and the quaternion used to
-        build them with shape ``(E, 4)``.
+        Wigner data with dense shape ``(E, D, D)``, packed shape ``(E, 46)``,
+        or ``None`` when ``build_full`` is false, plus the edge quaternion.
     """
     # === Step 1. Build edge-aligned quaternions ===
     edge_quat = build_edge_quaternion(
@@ -650,10 +730,44 @@ def _build_edge_wigner(
     # === Step 3. Convert quaternions to packed Wigner-D blocks ===
     if not build_full:
         return None, None, edge_quat
+    if packed_wigner:
+        from deepmd.pt_expt.kernels.cute.sezm.wignerd import (
+            run_cute_wignerd,
+        )
+
+        cute_wigner = run_cute_wignerd(
+            edge_quat,
+            wigner_calc,
+            packed_wigner=packed_wigner,
+        )
+        if cute_wigner is not None:
+            return cute_wigner[0], cute_wigner[1], edge_quat
     D_full, Dt_full = wigner_calc(edge_quat)
     return D_full, Dt_full, edge_quat
 
 
+def _packed_wigner_edges_eligible(
+    candidate: bool,
+    *,
+    edge_count: int,
+    node_count: int,
+    destinations_sorted: bool,
+    runtime_dtypes: tuple[torch.dtype, ...] = (),
+) -> bool:
+    """Finish packed eligibility from scalar shape and provenance metadata."""
+    from deepmd.pt_expt.kernels.cute.sezm.so2.operation import (
+        packed_wigner_edges_eligible,
+    )
+
+    return packed_wigner_edges_eligible(
+        candidate=candidate,
+        edge_count=edge_count,
+        node_count=node_count,
+        destinations_sorted=destinations_sorted,
+        runtime_dtypes=runtime_dtypes,
+    )
+
+
 def _finalize_edge_cache(
     *,
     n_nodes: int,
@@ -665,9 +779,11 @@ def _finalize_edge_cache(
     edge_env: torch.Tensor,
     D_full: torch.Tensor | None,
     Dt_full: torch.Tensor | None,
+    D_packed: torch.Tensor | None,
     edge_quat: torch.Tensor,
     deg_norm_floor: float,
     edge_src_gate: torch.Tensor | None = None,
+    destinations_sorted: bool = False,
 ) -> EdgeFeatureCache:
     """
     Assemble the shared `EdgeFeatureCache` layout.
@@ -689,11 +805,14 @@ def _finalize_edge_cache(
     edge_env
         Smooth edge envelope weights with shape (E, 1).
     D_full
-        Packed Wigner-D matrices with shape (E, D, D), or None when the
+        Dense Wigner-D matrices with shape (E, D, D), or None when the
         full Wigner-D construction is skipped (all-Cartesian model).
     Dt_full
-        Transposed packed Wigner-D matrices with shape (E, D, D), or None
+        Transposed dense Wigner-D matrices with shape (E, D, D), or None
         when the full Wigner-D construction is skipped.
+    D_packed
+        Optional Neo packed Wigner panel with shape (E, 46). Dense consumers
+        continue to observe stable-rank ``D_full`` and ``Dt_full`` fields.
     edge_quat
         Global-to-local quaternions used to build the Wigner-D matrices with
         shape (E, 4).
@@ -705,12 +824,17 @@ def _finalize_edge_cache(
     edge_src_gate
         Optional per-edge SFPG weight with shape (E, 1). ``None`` in
         non-bridging mode.
+    destinations_sorted
+        Host-side provenance that ``dst`` is nondecreasing.
 
     Returns
     -------
     EdgeFeatureCache
         Finalized per-edge cache shared by eager and compile paths.
     """
+    if deg_norm_floor <= 0.0:
+        raise ValueError("deg_norm_floor must be positive")
+
     # === Step 1. Build smooth destination degrees ===
     with nvtx_range("degree"):
         deg = torch.zeros(n_nodes, dtype=edge_vec.dtype, device=edge_vec.device)  # (N,)
@@ -730,11 +854,13 @@ def _finalize_edge_cache(
         inv_sqrt_deg=inv_sqrt_deg,
         D_full=D_full,
         Dt_full=Dt_full,
+        D_packed=D_packed,
         D_to_m_cache={},
         Dt_from_m_cache={},
         csr_cache={},
         edge_src_gate=edge_src_gate,
         edge_quat=edge_quat,
+        destinations_sorted=destinations_sorted,
     )
 
 
@@ -890,16 +1016,20 @@ def edge_cache_to_dtype(
     # Use local variables with explicit None check and assignment.
     _D_full = cache.D_full
     _Dt_full = cache.Dt_full
+    cached_D_packed = cache.D_packed
     _edge_src_gate = cache.edge_src_gate
     _edge_quat = cache.edge_quat
     D_full: torch.Tensor | None = None
     Dt_full: torch.Tensor | None = None
+    D_packed: torch.Tensor | None = None
     edge_src_gate: torch.Tensor | None = None
     edge_quat: torch.Tensor | None = None
     if _D_full is not None:
         D_full = _D_full.to(dtype=dtype)
     if _Dt_full is not None:
         Dt_full = _Dt_full.to(dtype=dtype)
+    if cached_D_packed is not None:
+        D_packed = cached_D_packed.to(dtype=dtype)
     if _edge_src_gate is not None:
         edge_src_gate = _edge_src_gate.to(dtype=dtype)
     if _edge_quat is not None:
@@ -918,9 +1048,12 @@ def edge_cache_to_dtype(
         inv_sqrt_deg=cache.inv_sqrt_deg.to(dtype=dtype),
         D_full=D_full,
         Dt_full=Dt_full,
+        D_packed=D_packed,
         D_to_m_cache=None if cache.D_to_m_cache is None else {},
         Dt_from_m_cache=None if cache.Dt_from_m_cache is None else {},
         csr_cache=None if cache.csr_cache is None else dict(cache.csr_cache),
         edge_src_gate=edge_src_gate,
         edge_quat=edge_quat,
+        cute_infer_so2_metadata=cache.cute_infer_so2_metadata,
+        destinations_sorted=cache.destinations_sorted,
     )
diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
index f3215f3072..017f14a4da 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
@@ -35,6 +35,10 @@
 from deepmd.pt.utils.utils import (
     get_generator,
 )
+from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import (
+    PACKED_VALUE_COUNT,
+    ZONAL_PANEL_OFFSETS,
+)
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
 )
@@ -199,6 +203,16 @@ def __init__(
             node_radial_l_index,
             persistent=True,
         )
+        packed_zonal_offsets = torch.tensor(
+            ZONAL_PANEL_OFFSETS if self.lmax == 3 else (),
+            device=self.device,
+            dtype=torch.long,
+        )
+        self.register_buffer(
+            "packed_zonal_offsets",
+            packed_zonal_offsets,
+            persistent=False,
+        )
         # The l=1 coefficients (packed rows 1..3) are the first three entries of
         # the non-scalar sequence ``node_row_index = [1, 2, ..., D-1]``, so the
         # native neighbor-spin l=1 message folds in at these local positions.
@@ -212,17 +226,17 @@ def __init__(
         # The reference composition materializes the per-edge message, an
         # (E, D-1, C) tensor that dominates the cost of this module. The fused
         # operator keeps it in registers and reduces through the destination CSR.
-        self._cuda_scatter = False
+        self.cuda_infer_l_1_scatter = False
         # ``None`` keeps the runtime ``zonal_coupling.is_cuda`` dispatch; the
         # freeze pins it to the AOTI target because tracing always runs on CPU.
-        self._force_fused_scatter: bool | None = None
+        self.force_cuda_infer_l_1_scatter: bool | None = None
         if cuda_infer_level() >= 1 and self.dtype is torch.float32:
             from deepmd.pt_expt.kernels.cuda.dpa4.zonal_scatter import (
                 op_available,
                 supported,
             )
 
-            self._cuda_scatter = op_available() and supported(
+            self.cuda_infer_l_1_scatter = op_available() and supported(
                 self.lmax, self.ebed_dim - 1, self.channels
             )
 
@@ -259,37 +273,58 @@ def forward(
         torch.Tensor
             Initial features to add with shape (N, D, C). l=0 is guaranteed zero.
         """
-        # === Step 1. Initialize output ===
+        # === Step 1. Validate the non-scalar contract ===
         device = edge_cache.edge_vec.device
         dtype = edge_cache.edge_vec.dtype
-        out = torch.zeros(
-            n_nodes, self.ebed_dim, self.channels, device=device, dtype=dtype
-        )  # (N, D, C)
         if self.lmax == 0:
-            return out
+            return torch.zeros(
+                n_nodes, self.ebed_dim, self.channels, device=device, dtype=dtype
+            )
 
         # === Step 2. Gather all m=0 columns (l >= 1) in one shot ===
         # Advanced indexing pairs one packed non-scalar row with the zonal m=0 column
         # from the same degree block in Dt_full.
         if zonal_coupling is None:
-            Dt_full = edge_cache.Dt_full  # (E, D, D)
-            zonal_coupling = Dt_full[
-                :,
-                self.non_scalar_row_index,
-                self.zonal_m0_col_index_for_row,
-            ]  # (E, D-1)
-
-        # === Step 3. Broadcast radial features per row ===
+            D_packed = edge_cache.D_packed
+            if D_packed is not None:
+                if self.lmax != 3 or D_packed.shape[1] != PACKED_VALUE_COUNT:
+                    raise ValueError("packed Wigner zonal coupling requires Neo lmax=3")
+                zonal_coupling = D_packed.index_select(
+                    1,
+                    self.packed_zonal_offsets,
+                )
+            else:
+                Dt_full = edge_cache.Dt_full  # (E, D, D)
+                if Dt_full is None:
+                    raise RuntimeError("GIE requires dense or packed Wigner storage")
+                zonal_coupling = Dt_full[
+                    :,
+                    self.non_scalar_row_index,
+                    self.zonal_m0_col_index_for_row,
+                ]  # (E, D-1)
+
+        # === Step 3. Optional backend message construction and reduction ===
+        accelerated = self.run_cute_infer_gie(
+            n_nodes=n_nodes,
+            edge_cache=edge_cache,
+            radial_feat=radial_feat,
+            zonal_coupling=zonal_coupling,
+            spin_l1_message=spin_l1_message,
+        )
+        if accelerated is not None:
+            return accelerated
+
+        # === Step 4. Eager fallback: 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 5, so
         # it takes over whenever nothing else joins the message in between.
         if (
-            self._can_fuse_scatter(zonal_coupling)
+            self.can_run_cuda_infer_l_1_scatter(zonal_coupling)
             and spin_l1_message is None
             and edge_cache.edge_src_gate is None
             and edge_cache.csr_cache is not None
         ):
-            return self.forward_fused_scatter(
+            return self.run_cuda_infer_l_1_scatter(
                 n_nodes, edge_cache, radial_feat, zonal_coupling
             )
 
@@ -309,7 +344,7 @@ def forward(
                 1, self.l1_local_index, spin_l1_message
             )
 
-        # === Step 4. Source Freeze Propagation Gate (optional) ===
+        # === Step 5. Source Freeze Propagation Gate (optional) ===
         # Mute messages emitted by nodes whose local neighborhood enters
         # the frozen zone. ``edge_src_gate`` is ``None`` outside bridging
         # mode so this is a no-op in normal training.
@@ -319,7 +354,10 @@ def forward(
                 dtype=non_scalar_message.dtype
             ).unsqueeze(-1)
 
-        # === Step 5. Scatter to nodes and normalize ===
+        # === Step 6. Scatter to nodes and normalize ===
+        out = torch.zeros(
+            n_nodes, self.ebed_dim, self.channels, device=device, dtype=dtype
+        )  # (N, D, C)
         # Avoid advanced-index writeback (out[:, non_scalar_row_index, :]) which produces a copy.
         non_scalar_out = out.new_zeros(
             n_nodes, self.non_scalar_row_index.numel(), self.channels
@@ -329,16 +367,71 @@ def forward(
         out.mul_(edge_cache.inv_sqrt_deg)
         return out
 
-    def _can_fuse_scatter(self, zonal_coupling: torch.Tensor) -> bool:
-        """Return whether the fused scatter serves the runtime or trace target."""
+    def run_cute_infer_gie(
+        self,
+        *,
+        n_nodes: int | torch.SymInt,
+        edge_cache: EdgeFeatureCache,
+        radial_feat: torch.Tensor,
+        zonal_coupling: torch.Tensor,
+        spin_l1_message: torch.Tensor | None = None,
+    ) -> torch.Tensor | None:
+        """Run the CuTe GIE path when its exact inference contract matches.
+
+        Parameters
+        ----------
+        n_nodes : int or torch.SymInt
+            Number of nodes addressed by the edge cache.
+        edge_cache : EdgeFeatureCache
+            Per-forward geometric edge cache.
+        radial_feat : torch.Tensor
+            Radial features with shape ``(E, lmax, C)``.
+        zonal_coupling : torch.Tensor
+            Zonal coupling with shape ``(E, D - 1)``.
+        spin_l1_message : torch.Tensor, optional
+            Native-spin message with shape ``(E, 3, C)``.
+
+        Returns
+        -------
+        torch.Tensor or None
+            Initial embedding with shape ``(N, D, C)``, or ``None`` when the
+            exact CuTe contract is not satisfied.
+        """
+        if self.training or spin_l1_message is not None:
+            return None
+        from deepmd.pt_expt.kernels.cute.sezm.gie import (
+            maybe_run_cute_gie,
+        )
+
+        return maybe_run_cute_gie(
+            self,
+            n_nodes=n_nodes,
+            edge_cache=edge_cache,
+            radial_feat=radial_feat,
+            zonal_coupling=zonal_coupling,
+        )
+
+    def can_run_cuda_infer_l_1_scatter(self, zonal_coupling: torch.Tensor) -> bool:
+        """Return whether the fused scatter serves the runtime or trace target.
+
+        Parameters
+        ----------
+        zonal_coupling : torch.Tensor
+            Zonal coupling with shape ``(E, D - 1)``.
+
+        Returns
+        -------
+        bool
+            Whether the CUDA inference level-one scatter is eligible.
+        """
         target_is_cuda = (
             zonal_coupling.is_cuda
-            if self._force_fused_scatter is None
-            else self._force_fused_scatter
+            if self.force_cuda_infer_l_1_scatter is None
+            else self.force_cuda_infer_l_1_scatter
         )
-        return self._cuda_scatter and not self.training and target_is_cuda
+        return self.cuda_infer_l_1_scatter and not self.training and target_is_cuda
 
-    def forward_fused_scatter(
+    def run_cuda_infer_l_1_scatter(
         self,
         n_nodes: int | torch.SymInt,
         edge_cache: EdgeFeatureCache,
diff --git a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py
index 4a8b73f0a0..f65f80d273 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py
@@ -36,6 +36,7 @@
 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,
@@ -145,6 +146,13 @@ def _project_frames(
     )
 
 
+def _inference_mode_is_frozen(module: nn.Module) -> bool:
+    """Return whether first-order inference-only shortcuts are safe."""
+    return not module.training and not any(
+        parameter.requires_grad for parameter in module.parameters()
+    )
+
+
 def _project_pair_in_one_transform(
     left: torch.Tensor,
     right: torch.Tensor,
@@ -713,7 +721,7 @@ def __init__(
         # transposed so both matrices are read row-major by grid point.
         # The operator is instantiated per coefficient-slot count, which this
         # projector fixes, so the choice is made once here rather than per call.
-        self._grid_pair_fn = None
+        self.cuda_infer_l_1_grid_pair = None
         if (
             cuda_infer_level() >= 1
             and self.projector.to_grid_mat.dtype is torch.float32
@@ -726,7 +734,7 @@ def __init__(
 
             slots = int(self.projector.to_grid_mat.shape[1])
             if op_available() and slots in SUPPORTED_SLOTS:
-                self._grid_pair_fn = grid_pair
+                self.cuda_infer_l_1_grid_pair = grid_pair
         # The training form differentiates the same expression inside the
         # force graph: a Triton tensor-core sandwich (grid-axis blocks, one
         # resident output tile) with analytic first and second order, one
@@ -737,7 +745,7 @@ def __init__(
         # dense section is small and the operator's dispatch chain costs
         # more than its kernels save on the host-bound configurations, so
         # the narrow grids stay with the compiler.
-        self._grid_pair_train_fn = None
+        self.triton_train_l_1_grid_pair = None
         if (
             triton_train_level() >= 1
             and self.projector.to_grid_mat.dtype is torch.float32
@@ -749,7 +757,7 @@ def __init__(
             )
 
             if GRID_PAIR_TRITON_AVAILABLE:
-                self._grid_pair_train_fn = grid_pair_train
+                self.triton_train_l_1_grid_pair = grid_pair_train
         self.register_buffer(
             "_from_grid_t",
             self.projector.from_grid_mat.transpose(0, 1).contiguous(),
@@ -838,7 +846,7 @@ def forward_scalar(
         because materializing a scalar-only fallback grid would be slower than
         that fused operator.
         """
-        if self._grid_pair_fn is not None and not self.training:
+        if self.cuda_infer_l_1_grid_pair is not None and not self.training:
             return self._slice_scalar_layout(self.forward(query, context))
         return self._forward(query, context, scalar_only=True)
 
@@ -854,7 +862,6 @@ def _forward(
         input_dtype = query.dtype
         query_ndfc, shape_info = self._to_ndfc(query)
         left, right, scalar_pair = self._prepare_pair(query_ndfc, context)
-
         # === Step 2. Select the static projection plan and apply the grid op ===
         direct_scalar = scalar_only and self._scalar_product_weight is not None
         coeff_out = self.grid_op(
@@ -1057,26 +1064,30 @@ def _pair_grid(
         c_wide = left.shape[3] // self.n_frames
         if c_wide % 32 != 0 or left.shape != right.shape:
             return None
+        if not self.training:
+            candidate = self.run_cute_infer_grid_pair(left, right)
+            if candidate is not None:
+                return candidate
         if self.training:
             # Training form: frame-packed operands ride through unreshaped,
             # with analytic first and second order behind the call. The
             # operator carries an autocast rule, so under AMP it runs the
             # same bf16-with-fp32-accumulation regime as the dense einsum
             # composition it replaces.
-            if self._grid_pair_train_fn is None:
+            if self.triton_train_l_1_grid_pair is None:
                 return None
-            return self._grid_pair_train_fn(
+            return self.triton_train_l_1_grid_pair(
                 left,
                 right,
                 self.projector.to_grid_mat,
                 self._from_grid_t,
                 self.n_frames,
             )
-        if self._grid_pair_fn is None or left.shape[2] != 1:
+        if self.cuda_infer_l_1_grid_pair is None or left.shape[2] != 1:
             return None
         n_batch, coeff_dim = left.shape[0], left.shape[1]
         flat_p = coeff_dim * self.n_frames
-        out = self._grid_pair_fn(
+        out = self.cuda_infer_l_1_grid_pair(
             left.reshape(n_batch, flat_p, c_wide),
             right.reshape(n_batch, flat_p, c_wide),
             self.projector.to_grid_mat,
@@ -1084,6 +1095,43 @@ def _pair_grid(
         )
         return out.reshape(n_batch, coeff_dim, 1, self.n_frames * c_wide)
 
+    def run_cute_infer_grid_pair(
+        self,
+        left: torch.Tensor,
+        right: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Run the CuTe grid product when its exact inference contract matches.
+
+        Parameters
+        ----------
+        left : torch.Tensor
+            Left coefficient operand with shape ``(N, D, F, n_frames * C)``.
+        right : torch.Tensor
+            Right coefficient operand with the same shape as ``left``.
+
+        Returns
+        -------
+        torch.Tensor or None
+            Coefficient result with the same shape as ``left``, or ``None``
+            when the exact CuTe contract is not satisfied.
+        """
+        if (
+            not cute_runtime_policy.is_cute_infer_enabled()
+            or not _inference_mode_is_frozen(self)
+        ):
+            return None
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.product import (
+            maybe_run_cute_output_grid_product,
+        )
+
+        return maybe_run_cute_output_grid_product(
+            left,
+            right,
+            self.projector.to_grid_mat,
+            self.projector.from_grid_mat,
+            n_frames=self.n_frames,
+        )
+
     def _project_pair_in_one_transform(
         self,
         left: torch.Tensor,
diff --git a/deepmd/pt/model/descriptor/sezm_nn/lora.py b/deepmd/pt/model/descriptor/sezm_nn/lora.py
index 3db8befae2..9c41d0da7a 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/lora.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/lora.py
@@ -487,7 +487,28 @@ 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.pt.model.model.sezm_model import (
+        _clear_shared_sezm_compile_cache,
+    )
+    from deepmd.pt_expt.kernels.cute.sezm.so2.operation import (
+        invalidate_cute_so2_state,
+    )
+
+    readout_invalidator = None
+
     for m in model.modules():
+        invalidate_cute_so2_state(m)
+        if hasattr(m, "_neo_sm80_readout_input_fold_cache"):
+            if readout_invalidator is None:
+                from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import (
+                    invalidate_neo_readout_input_fold,
+                )
+
+                readout_invalidator = invalidate_neo_readout_input_fold
+            readout_invalidator(m)
+        for name in tuple(vars(m)):
+            if name.startswith("_deepmd_cute_"):
+                delattr(m, name)
         core_cache = getattr(m, "compiled_core_compute_cache", None)
         if isinstance(core_cache, dict):
             core_cache.clear()
@@ -501,6 +522,7 @@ def _clear_sezm_compile_cache(model: nn.Module) -> None:
                 m._dens_compiled = False
             if hasattr(m, "_dens_pending_compile_t0"):
                 m._dens_pending_compile_t0 = None
+    _clear_shared_sezm_compile_cache()
 
 
 def _swap_submodule(parent: nn.Module, attr: str, new_module: nn.Module) -> None:
diff --git a/deepmd/pt/model/descriptor/sezm_nn/so2.py b/deepmd/pt/model/descriptor/sezm_nn/so2.py
index 0a3418efcf..1048f7c758 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/so2.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/so2.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 (
@@ -371,7 +370,7 @@ def __init__(
         # Export override for the block-diagonal vs dense matmul branch below.
         # ``None`` keeps the runtime ``x_flat.is_cuda`` dispatch; the freeze sets
         # it so the AOTI graph follows the *target* device, not the CPU trace.
-        self._force_block_diag_matmul: bool | None = None
+        self.force_block_diag_matmul: bool | None = None
 
         # The assembled SO(2) weight is block-diagonal over |m| groups; the
         # forward contracts only the diagonal blocks (see _block_diagonal_matmul).
@@ -383,7 +382,7 @@ def __init__(
         # Triton BN=64 block-diagonal GEMM that consumes the strided operands
         # without a contiguity copy. Bound only when Triton is available and every
         # block width aligns to BN=64; otherwise the eager path is kept.
-        self._block_diag_gemm = None
+        self.triton_l_1_block_diag_gemm = None
         self.triton_infer_level = triton_infer_level()
         self.triton_train_level = triton_train_level()
         if max(self.triton_infer_level, self.triton_train_level) >= 1:
@@ -396,7 +395,7 @@ def __init__(
             if SO2_BLOCK_GEMM_TRITON_AVAILABLE and slices_supported(
                 self._block_diag_slices
             ):
-                self._block_diag_gemm = block_diag_gemm
+                self.triton_l_1_block_diag_gemm = block_diag_gemm
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
@@ -432,12 +431,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         # trips an Inductor AVX2 C++ codegen bug, so only CPU needs it. Every other
         # device uses the block-diagonal contraction, which skips the structural
         # off-|m| zeros. ``make_fx`` resolves this Python branch at trace time, so
-        # the freeze pins ``_force_block_diag_matmul`` to the AOTI target device
+        # the freeze pins ``force_block_diag_matmul`` to the AOTI target device
         # (tracing always runs on CPU regardless of where the artifact will run).
-        if self._force_block_diag_matmul is None:
+        if self.force_block_diag_matmul is None:
             use_block_diag = not x_flat.is_cpu
         else:
-            use_block_diag = self._force_block_diag_matmul
+            use_block_diag = self.force_block_diag_matmul
         if use_block_diag:
             out_flat = self._block_diagonal_matmul(x_flat, weight)
         else:
@@ -576,8 +575,13 @@ def _block_diagonal_matmul(
             Flattened output with shape ``(F, E, D_m*Cout)``.
         """
         weight = weight.permute(1, 0, 2)  # (F, D_m*Cin, D_m*Cout)
-        if self._block_diag_gemm is not None and active_triton_level(self) >= 1:
-            return self._block_diag_gemm(x_flat, weight, self._block_diag_slices)
+        if (
+            self.triton_l_1_block_diag_gemm is not None
+            and active_triton_level(self) >= 1
+        ):
+            return self.triton_l_1_block_diag_gemm(
+                x_flat, weight, self._block_diag_slices
+            )
         blocks = [
             torch.bmm(
                 x_flat[:, :, in0:in1],
@@ -727,7 +731,7 @@ def __init__(
         # of the ``degree_channel`` low-rank branch in the ``mmax == 1`` layout.
         self.triton_infer_level = triton_infer_level()
         self.triton_train_level = triton_train_level()
-        self._radial_mix_block = None
+        self.triton_l_1_radial_mix = None
         if (
             max(self.triton_infer_level, self.triton_train_level) >= 1
             and self.mode == "degree_channel"
@@ -738,7 +742,7 @@ def __init__(
                 radial_mix_block,
             )
 
-            self._radial_mix_block = radial_mix_block
+            self.triton_l_1_radial_mix = radial_mix_block
 
     def _build_dense_scatter_indices(self) -> tuple[torch.Tensor, torch.Tensor]:
         compact_indices: list[int] = []
@@ -832,8 +836,11 @@ def forward(self, x_local: torch.Tensor, radial_feat: torch.Tensor) -> torch.Ten
             compact = kernel_flat.view(
                 x_local.shape[0], self.degree_kernel_size, self.rank
             )
-            if self._radial_mix_block is not None and active_triton_level(self) >= 1:
-                return self._radial_mix_block(
+            if (
+                self.triton_l_1_radial_mix is not None
+                and active_triton_level(self) >= 1
+            ):
+                return self.triton_l_1_radial_mix(
                     compact, x_local, self.channel_basis, self.lmax
                 )
             kernel = self._scatter_rank_kernel(compact)
@@ -1213,35 +1220,27 @@ def __init__(
         self.device = env.DEVICE
         self.precision = RESERVED_PRECISION_DICT[dtype]
         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
-        # 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.
+        # Acceleration gates are read once at construction so they become
+        # compile-time constants in the traced (``make_fx``) graph. Triton level
+        # 1 replaces individual operators, level 2 binds the fused value path,
+        # and inference level 3 selects its fp16x3 implementation on swept
+        # shapes. ``DP_CUTILE_INFER`` selects a complete alternative SO(2) path
+        # and is mutually exclusive with Triton. ``DP_CUDA_INFER`` and
+        # ``DP_CUDA_TRAIN`` independently bind hand-written CUDA producers that
+        # take precedence over the corresponding narrower paths. CuTe dispatch
+        # is resolved at the enclosing block, where its exact-shape SO2 kernel
+        # can coexist with every fallback. Unsupported entries remain ``None``.
         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:
+        self.cuda_infer_level = cuda_infer_level()
+        self.cutile_infer_enabled = use_cutile_infer()
+        if self.triton_infer_level >= 1 and self.cutile_infer_enabled:
             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
+        self.triton_infer_l_2_value = None
+        self.cutile_infer_value = None
 
         # === Step 1. Split deterministic seeds at the module top-level ===
         seed_so2_stack = child_seed(seed, 0)
@@ -1620,9 +1619,9 @@ def __init__(
         # attention layout without the optional focus-mix / value / output
         # projections (the deployed DPA4 configuration). Whichever of the
         # mutually exclusive inference gates is active then supplies the
-        # implementation, and ``self._flash_atten_fn`` being bound is what marks
+        # implementation, and ``self.flash_attention`` being bound is what marks
         # the fused path as live.
-        self._flash_atten_layout_ok = (
+        self.flash_attention_layout_supported = (
             self.n_atten_head > 0
             and self.mmax == 1
             and self.needs_local_frame
@@ -1634,24 +1633,24 @@ def __init__(
         )
         # The cuTile aggregation is inference-only; the Triton one also serves
         # training, so it is bound whenever either gate asks for level 1.
-        self._flash_atten_fn = None
-        self._flash_atten_trains = False
-        if self._flash_atten_layout_ok and self.use_cutile_infer:
+        self.flash_attention = None
+        self.flash_attention_supports_training = False
+        if self.flash_attention_layout_supported and self.cutile_infer_enabled:
             from deepmd.pt_expt.kernels.cutile.sezm.flash_atten import (
                 flash_atten_aggregate,
             )
 
-            self._flash_atten_fn = flash_atten_aggregate
+            self.flash_attention = flash_atten_aggregate
         elif (
-            self._flash_atten_layout_ok
+            self.flash_attention_layout_supported
             and max(self.triton_infer_level, self.triton_train_level) >= 1
         ):
             from deepmd.pt_expt.kernels.triton.sezm.flash_atten import (
                 flash_atten_aggregate,
             )
 
-            self._flash_atten_fn = flash_atten_aggregate
-            self._flash_atten_trains = self.triton_train_level >= 1
+            self.flash_attention = flash_atten_aggregate
+            self.flash_attention_supports_training = self.triton_train_level >= 1
 
         # === Step 12. Optional fused Triton SO(2) value path (inference) ===
         # Fuses rotate-to-local, the radial degree mixing, the gated mixing
@@ -1672,7 +1671,7 @@ def __init__(
                 make_triton_value_path,
             )
 
-            self._triton_value_path = make_triton_value_path(self)
+            self.triton_infer_l_2_value = make_triton_value_path(self)
 
         # === Step 13. Optional fused rotate-to-local + radial degree mixing ===
         # Level-1 companion of the value path: only the rotation and the
@@ -1689,7 +1688,7 @@ def __init__(
         # hidden widths. Below the bound the separate rotation and radial-mix
         # kernels win end to end (a 64-wide stack loses ~10%, a 128-wide one
         # gains), so the binding follows the measured crossover.
-        self._triton_rotate_mix = None
+        self.triton_l_1_rotate_mix = None
         if (
             max(self.triton_infer_level, self.triton_train_level) >= 1
             and self.hidden_channels >= 128
@@ -1698,7 +1697,7 @@ def __init__(
                 make_triton_rotate_mix,
             )
 
-            self._triton_rotate_mix = make_triton_rotate_mix(self)
+            self.triton_l_1_rotate_mix = make_triton_rotate_mix(self)
 
         # === Step 14. Optional fused CUDA SO(2) convolution (inference) ===
         # One hand-written CUDA operator spans the complete per-edge path:
@@ -1708,13 +1707,13 @@ def __init__(
         # flash aggregation, and takes precedence over them when the block
         # matches its supported configuration. The factory returns ``None``
         # otherwise, leaving whichever narrower path is bound in charge.
-        self._cuda_conv_fn = None
-        if cuda_infer_level() >= 2 and self._flash_atten_layout_ok:
+        self.cuda_infer_l_2_conv = None
+        if self.cuda_infer_level >= 2 and self.flash_attention_layout_supported:
             from deepmd.pt_expt.kernels.cuda.dpa4 import (
                 make_cuda_so2_conv,
             )
 
-            self._cuda_conv_fn = make_cuda_so2_conv(self)
+            self.cuda_infer_l_2_conv = make_cuda_so2_conv(self)
 
         # === Step 15. Optional fused CUDA SO(2) value path (training) ===
         # One CUDA operator spans the training value stream up to the attention
@@ -1725,13 +1724,13 @@ def __init__(
         # behind the same differentiable boundary. The attention span stays on
         # the Triton operator composition inside the traced graph. Bound under
         # ``DP_CUDA_TRAIN=1``; ``so2_message`` dispatches to it in training mode.
-        self._cuda_value_train = None
+        self.cuda_train_value = None
         if cuda_train_enabled():
             from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import (
                 make_cuda_so2_value,
             )
 
-            self._cuda_value_train = make_cuda_so2_value(self)
+            self.cuda_train_value = make_cuda_so2_value(self)
 
         # === Step 16. Optional fused destination-segmented attention softmax ===
         # One CSR-segmented operator per direction replaces the
@@ -1740,7 +1739,7 @@ def __init__(
         # hand-derived second order keep the force-loss trace from expanding
         # the chain into materialized surfaces and serialized scatters. The
         # source-gated (SFPG) form and fp64 compute keep the reference path.
-        self._segment_softmax_fn = None
+        self.triton_l_1_segment_softmax = None
         if (
             max(self.triton_infer_level, self.triton_train_level) >= 1
             and self.compute_dtype is torch.float32
@@ -1752,28 +1751,18 @@ 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.triton_l_1_segment_softmax = segment_softmax
 
-            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:
+        if self.cutile_infer_enabled:
             from deepmd.pt_expt.kernels.cutile.sezm.so2_value_path import (
                 make_cutile_value_path,
             )
 
-            self._cutile_value_path = make_cutile_value_path(self)
+            self.cutile_infer_value = make_cutile_value_path(self)
 
     def forward(
         self,
@@ -1912,23 +1901,25 @@ def forward_attention(
         # The fused CUDA operator computes the attention weights itself, so it
         # does not serve the bridging mode, whose source gate reshapes the
         # softmax normalization.
-        run_cuda = (
-            self._cuda_conv_fn is not None
+        run_cuda_infer_l_2 = (
+            self.cuda_infer_l_2_conv is not None
             and not self.training
             and edge_cache.edge_src_gate is None
         )
-        run_flash = (
-            self._flash_atten_fn is not None
-            and (self._flash_atten_trains or not self.training)
-            and not run_cuda
+        run_flash_attention = (
+            self.flash_attention is not None
+            and (self.flash_attention_supports_training or not self.training)
+            and not run_cuda_infer_l_2
         )
-        if run_cuda:
-            return self.forward_attention_cuda(x, edge_cache, radial_feat, x_l0_node)
-        if run_flash:
+        if run_cuda_infer_l_2:
+            return self.forward_attention_cuda_infer_l_2(
+                x, edge_cache, radial_feat, x_l0_node
+            )
+        if run_flash_attention:
             return self.forward_attention_flash(x, edge_cache, radial_feat, x_l0_node)
         return self.forward_attention_dense(x, edge_cache, radial_feat, x_l0_node)
 
-    def forward_attention_cuda(
+    def forward_attention_cuda_infer_l_2(
         self,
         x: torch.Tensor,
         edge_cache: EdgeFeatureCache,
@@ -1962,7 +1953,7 @@ def forward_attention_cuda(
             Node update with shape (N, D, C_wide).
         """
         # === Step 1. Projected radial features ===
-        rad_feat = self._cuda_conv_fn.radial_features(
+        rad_feat = self.cuda_infer_l_2_conv.radial_features(
             radial_feat
         )  # (E, lmax+1, C_wide)
 
@@ -1973,7 +1964,9 @@ def forward_attention_cuda(
         head_gate = self.attention_head_gate(x_l0_node)  # (N, Fa, H)
 
         # === Step 4. Fused convolution ===
-        out = self._cuda_conv_fn(x, edge_cache, rad_feat, q_node, k_node, head_gate)
+        out = self.cuda_infer_l_2_conv(
+            x, edge_cache, rad_feat, q_node, k_node, head_gate
+        )
         return out.to(dtype=self.dtype)  # (N, D, C_wide)
 
     def forward_attention_flash(
@@ -1990,7 +1983,10 @@ def forward_attention_flash(
         kernel folds the block-diagonal rotate-back, the inverse-rotation
         rescale, the per-edge weighting and the destination reduction into a
         single atomic-free pass, so the transient rotate-back message and
-        weighted value tensors are never materialized.
+        weighted value tensors are never materialized. The value producer is
+        selected independently inside ``so2_message``: training uses
+        ``cuda_train_value`` when it is bound, while the flash operator remains
+        the aggregation consumer selected by the Triton training gate.
 
         Parameters
         ----------
@@ -2029,8 +2025,8 @@ def forward_attention_flash(
         order, row_ptr = cached_edge_csr(edge_cache, "dst", n_node)
         rotation = edge_cache.Dt_full
         if rotation is None:
-            rotation = self._cuda_value_train.edge_runs(edge_cache)
-        pre_gate = self._flash_atten_fn(
+            rotation = self.cuda_train_value.edge_runs(edge_cache)
+        pre_gate = self.flash_attention(
             x_local,
             rotation,
             self.rotate_inv_rescale_full,
@@ -2215,7 +2211,7 @@ def attention_weights(
         edge_src_gate = edge_cache.edge_src_gate
         n_nodes = x_l0_node.shape[0]
         if (
-            self._segment_softmax_fn is not None
+            self.triton_l_1_segment_softmax is not None
             and edge_src_gate is None
             and attn_logits.is_cuda
             and active_triton_level(self) >= 1
@@ -2233,7 +2229,7 @@ def attention_weights(
                 + float(self.eps)
             ).reshape(-1)  # (F * H,)
             n_channel = self.attn_n_focus * self.n_atten_head
-            alpha = self._segment_softmax_fn(
+            alpha = self.triton_l_1_segment_softmax(
                 attn_logits.reshape(n_edge, n_channel).to(dtype=torch.float32),
                 edge_cache.edge_env.reshape(n_edge).to(dtype=torch.float32),
                 null_logit,
@@ -2435,24 +2431,21 @@ def so2_message(
         src, dst = edge_cache.src, edge_cache.dst
         n_edge = src.numel()
 
-        if self._cutile_value_path is not None and not self.training:
+        if self.cutile_infer_value is not None and not self.training:
             # === Steps 1-5 (fused cuTile operators). ``rotate_mix`` folds the
             # rotation and the radial degree mixing into one edge-parallel
             # kernel writing the focus-major layout; ``mixing_stack`` runs the
             # whole gated stack, keeping the inter-layer activations and the
             # gated-layer pre-activations off the traced graph entirely. ===
-            x_local, rad_feat = self._cutile_value_path(x, edge_cache, radial_feat)
-        elif self._cuda_value_train is not None and self.training:
+            x_local, rad_feat = self.cutile_infer_value(x, edge_cache, radial_feat)
+        elif self.cuda_train_value is not None and self.training:
             # === Steps 1-5 (one CUDA kernel, training). The whole value
-            # stream up to the attention aggregation runs in a single launch;
-            # only the backward anchors reach device memory. ===
+            # stream up to the attention aggregation runs in a single launch
+            # with analytic backward and second order; only the backward
+            # anchors reach device memory. ===
             cached_edge_csr(edge_cache, "src", x.shape[0])
-            x_local, rad_feat = self._cuda_value_train(x, edge_cache, radial_feat)
-        elif (
-            self._triton_value_path is not None
-            and not self.training
-            and active_triton_level(self) >= 2
-        ):
+            x_local, rad_feat = self.cuda_train_value(x, edge_cache, radial_feat)
+        elif self.triton_infer_l_2_value is not None and not self.training:
             # === Steps 1-5 (fused Triton operators, inference).
             # ``so2_rotate_mix`` folds the rotation and the radial degree
             # mixing into one edge-parallel kernel writing the focus-major
@@ -2460,14 +2453,8 @@ def so2_message(
             # the competition weight fused into its final store, keeping the
             # 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:
+            x_local, rad_feat = self.triton_infer_l_2_value(x, edge_cache, radial_feat)
+        elif self.triton_l_1_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
             # Wigner rotation and the radial degree mixing, and writes the
@@ -2477,7 +2464,7 @@ def so2_message(
             # step and kept on the edge cache. ===
             with nvtx_range("SO2Conv/rotate_mix"):
                 cached_edge_csr(edge_cache, "src", x.shape[0])
-                u0, rad_feat = self._triton_rotate_mix(x, edge_cache, radial_feat)
+                u0, rad_feat = self.triton_l_1_rotate_mix(x, edge_cache, radial_feat)
                 x_local = u0.view(
                     self.n_focus, n_edge, self.reduced_dim, self.so2_focus_dim
                 )  # (F, E, D_m, Cf)
@@ -2496,14 +2483,14 @@ def so2_message(
                 D_full = edge_cache.D_full
                 x_dst_local: torch.Tensor | None = None
                 if active_triton_level(self) >= 1:
-                    # ``self._rotate_to_local_fn`` was bound in ``__init__`` (the
+                    # ``self.triton_l_1_rotate_to_local`` was bound in ``__init__`` (the
                     # block kernel for the m-major ``mmax == 1`` layout, dense
                     # otherwise).
-                    x_local = self._rotate_to_local_fn(
+                    x_local = self.triton_l_1_rotate_to_local(
                         x, src, D_full
                     )  # (E, D_m, C_wide)
                     if self.node_wise_grid_product is not None:
-                        x_dst_local = self._rotate_to_local_fn(
+                        x_dst_local = self.triton_l_1_rotate_to_local(
                             x, dst, D_full
                         )  # (E, D_m, C_wide)
                 else:
@@ -2734,7 +2721,9 @@ def _so2_rotate_back(
             if active_triton_level(self) >= 1 and self.mmax == 1:
                 # The block kernel consumes the (E, F, D_m, Cf) focus layout in
                 # place, folding the inverse transpose into its channel addressing.
-                x_message = self._rotate_back_fn(x_local, Dt_full)  # (E, D, C_wide)
+                x_message = self.triton_l_1_rotate_back(
+                    x_local, Dt_full
+                )  # (E, D, C_wide)
             else:
                 # Restore reduced global layout (E, D_m, C_wide) for inverse rotation.
                 x_local = (
@@ -2743,7 +2732,9 @@ def _so2_rotate_back(
                     .reshape(n_edge, self.reduced_dim, self.hidden_channels)
                 )
                 if active_triton_level(self) >= 1:
-                    x_message = self._rotate_back_fn(x_local, Dt_full)  # (E, D, C_wide)
+                    x_message = self.triton_l_1_rotate_back(
+                        x_local, Dt_full
+                    )  # (E, D, C_wide)
                 else:
                     Dt_from_m = project_Dt_from_m(
                         Dt_full=Dt_full,
@@ -2894,8 +2885,8 @@ def _build_so2_mixing(
         self.reduced_dim = int(coeff_index_m.numel())
 
         # === Step 2. Triton rotation kernels: block for mmax == 1, dense otherwise ===
-        self._rotate_to_local_fn = None
-        self._rotate_back_fn = None
+        self.triton_l_1_rotate_to_local = None
+        self.triton_l_1_rotate_back = None
         if max(self.triton_infer_level, self.triton_train_level) >= 1:
             from deepmd.pt_expt.kernels.triton.sezm.so2_rotation import (
                 rotate_back_block_so2,
@@ -2905,20 +2896,22 @@ def _build_so2_mixing(
             )
 
             if self.mmax == 1:
-                self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_block(
-                    x, src, wigner, self.lmax
+                self.triton_l_1_rotate_to_local = lambda x, src, wigner: (
+                    rotate_to_local_block(x, src, wigner, self.lmax)
                 )
                 # The block kernel reads the (E, F, D_m, Cf) focus layout directly,
                 # so the rotate-back path passes ``x_local`` before the global
                 # reshape and the transpose-back copy is skipped.
-                self._rotate_back_fn = lambda x_local, wigner: rotate_back_block_so2(
-                    x_local, wigner, self.lmax
+                self.triton_l_1_rotate_back = lambda x_local, wigner: (
+                    rotate_back_block_so2(x_local, wigner, self.lmax)
                 )
             else:
-                self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_dense(
-                    x, src, wigner, self.coeff_index_m, self.ebed_dim_full
+                self.triton_l_1_rotate_to_local = lambda x, src, wigner: (
+                    rotate_to_local_dense(
+                        x, src, wigner, self.coeff_index_m, self.ebed_dim_full
+                    )
                 )
-                self._rotate_back_fn = lambda x_local, wigner: rotate_back_dense(
+                self.triton_l_1_rotate_back = lambda x_local, wigner: rotate_back_dense(
                     x_local, wigner, self.coeff_index_m, self.ebed_dim_full
                 )
 
diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index b283e3a7e5..f75efe572c 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -459,9 +459,9 @@ def __init__(
                     ]
             # The monomial basis routes through whichever accelerated backend
             # is selected; the two gates are mutually exclusive.
-            self._use_cutile_monomials = use_cutile_infer()
-            self._use_triton_monomials = triton_infer_level() >= 1
-            self._use_triton_train_monomials = triton_train_level() >= 1
+            self.cutile_infer_monomials = use_cutile_infer()
+            self.triton_infer_l_1_monomials = triton_infer_level() >= 1
+            self.triton_train_l_1_monomials = triton_train_level() >= 1
             # The l = 2 contraction tensor collapsed onto the 35 unique
             # degree-4 monomials: column m of the coefficient matrix sums
             # C_l2[:, :, p] over the 4^4 index tuples p whose component
@@ -1150,12 +1150,12 @@ def _monomial_matrix(
             and (
                 (
                     not self.training
-                    and (self._use_triton_monomials or self._use_cutile_monomials)
+                    and (self.triton_infer_l_1_monomials or self.cutile_infer_monomials)
                 )
-                or (self.training and self._use_triton_train_monomials)
+                or (self.training and self.triton_train_l_1_monomials)
             )
         ):
-            if not self.training and self._use_cutile_monomials:
+            if not self.training and self.cutile_infer_monomials:
                 from deepmd.pt_expt.kernels.cutile.sezm.wigner_monomials import (
                     wigner_monomials as monomial_basis,
                 )
@@ -1191,12 +1191,12 @@ def _compute_l2_block(self, edge_quaternion: torch.Tensor) -> torch.Tensor:
             and (
                 (
                     not self.training
-                    and (self._use_triton_monomials or self._use_cutile_monomials)
+                    and (self.triton_infer_l_1_monomials or self.cutile_infer_monomials)
                 )
-                or (self.training and self._use_triton_train_monomials)
+                or (self.training and self.triton_train_l_1_monomials)
             )
         ):
-            if not self.training and self._use_cutile_monomials:
+            if not self.training and self.cutile_infer_monomials:
                 from deepmd.pt_expt.kernels.cutile.sezm.wigner_monomials import (
                     wigner_monomials as monomial_basis,
                 )
diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py
index 0711e77b1d..c2fe9eb000 100644
--- a/deepmd/pt/model/model/sezm_model.py
+++ b/deepmd/pt/model/model/sezm_model.py
@@ -258,7 +258,8 @@
 * ``triton.cudagraphs=False``
       cudagraphs capture autograd metadata only once.  Higher-order
       gradients need fresh metadata per call, so cudagraphs would feed
-      stale autograd state into the second backward.
+      stale autograd state into the second backward. Packed SO2 also retains
+      Python-owned forward state, so direct graph capture remains disabled.
 * ``max_fusion_size=DP_FUSION_SIZE`` (default 8)
       Caps kernel fusion complexity so Inductor's scheduler does not
       time out on the large edge-level reductions inside the
@@ -353,7 +354,8 @@
 NOTE 10 -- Tail dummy edges
 ---------------------------
 
-The edge-schema builders append two masked edges at the end of every batch.
+The edge-schema builders append two masked edges to every batch. They remain
+trailing unless a later destination sort permutes them with the real edges.
 Real edge compaction happens via
 ``torch.nonzero(valid_mask)``, whose output length is data-dependent
 and can be zero in sparse or single-atom systems (e.g. isolated-atom
@@ -363,9 +365,9 @@
 ``dynamic=True``.  A pair of dummy slots also gives Inductor's batched
 matmul lowering a static ``E >= 2`` edge-axis bound, avoiding
 data-dependent layout guards on ``E == 1`` that would otherwise cause
-an extra recompile when the first batch contains no real edges.  Each
-dummy's ``edge_mask`` is ``False`` so it contributes exactly zero to
-every downstream sum or gather.
+an extra recompile when the first batch contains no real edges. Each dummy's
+``edge_mask`` is ``False`` wherever sorting places it, so it contributes
+exactly zero to every downstream sum or gather.
 
 NOTE 11 -- Edge-vector leaf (gather outside the AD region)
 ----------------------------------------------------------
@@ -528,6 +530,131 @@
 SeZMModel_ = make_model(SeZMAtomicModel)
 
 
+def _neo_cute_infer_enabled() -> bool:
+    """Return whether the model builder must destination-sort CuTe edges."""
+    from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import (
+        is_cute_infer_enabled,
+    )
+
+    return is_cute_infer_enabled()
+
+
+@torch.compiler.assume_constant_result
+def _neo_cute_nlist_eager_island_enabled(device: torch.device) -> bool:
+    """Return whether Neo requests an eager Toolkit-Ops neighbor-list call."""
+    if not _neo_cute_infer_enabled() or device.type != "cuda":
+        return False
+    try:
+        compute_capability = tuple(torch.cuda.get_device_capability(device))
+    except RuntimeError:
+        return False
+    from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import (
+        is_so2_eager_island_enabled,
+    )
+
+    return is_so2_eager_island_enabled(compute_capability)
+
+
+@torch.compiler.disable
+def _build_neo_neighbor_list_eager_island(
+    builder: NeighborList,
+    coord: torch.Tensor,
+    atype: torch.Tensor,
+    box: torch.Tensor | None,
+    rcut: float,
+    sel: list[int],
+    *,
+    return_mode: str,
+) -> Any:
+    """Build Neo neighbors outside Dynamo when the runtime policy requests it."""
+    return builder.build(
+        coord,
+        atype,
+        box,
+        rcut,
+        sel,
+        return_mode=return_mode,
+    )
+
+
+def _build_neo_neighbor_list(
+    builder: NeighborList,
+    coord: torch.Tensor,
+    atype: torch.Tensor,
+    box: torch.Tensor | None,
+    rcut: float,
+    sel: list[int],
+    *,
+    return_mode: str,
+) -> Any:
+    """Apply Neo runtime policy around a general neighbor-list strategy."""
+    if isinstance(builder, NvNeighborList) and _neo_cute_nlist_eager_island_enabled(
+        coord.device
+    ):
+        return _build_neo_neighbor_list_eager_island(
+            builder,
+            coord,
+            atype,
+            box,
+            rcut,
+            sel,
+            return_mode=return_mode,
+        )
+    return builder.build(
+        coord,
+        atype,
+        box,
+        rcut,
+        sel,
+        return_mode=return_mode,
+    )
+
+
+def _neo_cute_so2_requires_sorted_edges(
+    descriptor: Any,
+    *,
+    training: bool,
+    device: torch.device,
+) -> bool:
+    """Return whether this descriptor can use packed Wigner-D on sorted edges.
+
+    ``forward_with_edges`` promotes coordinates and edge vectors to the
+    descriptor compute dtype before constructing Wigner data.  Mirror that
+    post-cast contract here; model inputs may still be FP64 at this boundary.
+    """
+    compute_dtype = descriptor.compute_dtype
+    return bool(
+        not training
+        and _neo_cute_infer_enabled()
+        and descriptor.is_cute_infer_packed_wigner_candidate(
+            device,
+            compute_dtype,
+            compute_dtype,
+        )
+    )
+
+
+def _sort_edge_tensors_by_destination(
+    edge_index: torch.Tensor,
+    edge_vec: torch.Tensor,
+    edge_mask: torch.Tensor,
+    edge_scatter_index: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Stably sort aligned edge tensors by destination, then source."""
+    if edge_index.shape[1] == 0:
+        return edge_index, edge_vec, edge_mask, edge_scatter_index
+    src = edge_index[0].to(dtype=torch.long)
+    dst = edge_index[1].to(dtype=torch.long)
+    source_stride = src.max().clamp_min(0) + 1
+    permutation = torch.argsort(dst * source_stride + src, stable=True)
+    return (
+        edge_index.index_select(1, permutation).contiguous(),
+        edge_vec.index_select(0, permutation).contiguous(),
+        edge_mask.index_select(0, permutation).contiguous(),
+        edge_scatter_index.index_select(1, permutation).contiguous(),
+    )
+
+
 def _select_neighbor_builder(nf: int, device: torch.device) -> NeighborList:
     """Select the O(N) neighbor builder for the given batch shape and device.
 
@@ -586,6 +713,13 @@ def _select_neighbor_builder(nf: int, device: torch.device) -> NeighborList:
 # knows which buffers were promoted and in what order.
 _SEZM_TASK_BUF_ORDER: dict[tuple[Any, ...], tuple[str, ...]] = {}
 
+
+def _clear_shared_sezm_compile_cache() -> None:
+    """Drop shared graphs that may contain frozen parameter constants."""
+    _SEZM_COMPILE_CACHE.clear()
+    _SEZM_TASK_BUF_ORDER.clear()
+
+
 _ENV_BOOL_CHOICES = {
     "1": True,
     "true": True,
@@ -736,6 +870,9 @@ def __init__(
         # Maps cache_key -> task_buf_order for this instance so forward()
         # knows which buffers to pass and in what order.
         object.__setattr__(self, "_task_buf_order_cache", {})
+        self.register_load_state_dict_post_hook(
+            self._invalidate_compiled_state_after_load
+        )
 
         # Training follows `use_compile`. Evaluation/inference samples env
         # policy at init time so path and precision stay fixed per model.
@@ -771,6 +908,41 @@ def __init__(
             else None
         )
 
+    @staticmethod
+    def _invalidate_compiled_state_after_load(
+        module: SeZMModel,
+        incompatible_keys: Any,
+    ) -> None:
+        """Discard graphs that may have captured state from before a load."""
+        del incompatible_keys
+        module.compiled_core_compute_cache.clear()
+        module._task_buf_order_cache.clear()
+        object.__setattr__(module, "compiled_embedding", None)
+        object.__setattr__(module, "_embedding_task_buf_order", None)
+        object.__setattr__(module, "compiled_dens_compute", None)
+        module._dens_compiled = False
+        module._core_compute_pending_compile_t0 = None
+        module._core_compute_pending_compile_key = None
+        module._dens_pending_compile_t0 = None
+        so2_invalidator = None
+        for child in module.modules():
+            if not (
+                hasattr(child, "_deepmd_cute_so2_state")
+                or hasattr(child, "_deepmd_cute_gate_expand_contract")
+            ):
+                continue
+            if so2_invalidator is None:
+                from deepmd.pt_expt.kernels.cute.sezm.so2.operation import (
+                    invalidate_cute_so2_state,
+                )
+
+                so2_invalidator = invalidate_cute_so2_state
+            so2_invalidator(child)
+        # Shared callables can contain make_fx get_attr constants, including
+        # prepared CuTe readout folds. A checkpoint load therefore invalidates
+        # the process-level cache as well as this instance's local slots.
+        _clear_shared_sezm_compile_cache()
+
     # =========================================================================
     # Forward Methods
     # =========================================================================
@@ -1492,6 +1664,21 @@ def core_compute(
         descriptor_model = self.atomic_model.descriptor
 
         # === Step 1. Establish the force-autograd endpoint ===
+        sort_edges_by_dst = _neo_cute_so2_requires_sorted_edges(
+            descriptor_model,
+            training=self.training,
+            device=edge_vec.device,
+        )
+        if sort_edges_by_dst:
+            edge_index, edge_vec, edge_mask, edge_scatter_index = (
+                _sort_edge_tensors_by_destination(
+                    edge_index,
+                    edge_vec,
+                    edge_mask,
+                    edge_scatter_index,
+                )
+            )
+
         # Neighbor-list construction and periodic-image resolution are explicit
         # caller responsibilities.  Once the edge displacements are supplied,
         # SeZM differentiates only the pure map ``(edge_vec, theta) -> E``.
@@ -1563,6 +1750,7 @@ def core_compute(
                 edge_index=edge_index,
                 edge_vec=edge_vec,
                 edge_mask=edge_mask,
+                edge_index_sorted_by_dst=sort_edges_by_dst,
                 charge_spin=charge_spin,
                 spin=spin,
                 comm_dict=comm_dict,
@@ -1746,10 +1934,16 @@ def core_compute_dens(
         descriptor_model = self.atomic_model.descriptor
 
         # === Step 1. Build compact sparse edges ===
+        sort_edges_by_dst = _neo_cute_so2_requires_sorted_edges(
+            descriptor_model,
+            training=self.training,
+            device=extended_coord.device,
+        )
         edge_index, edge_vec, edge_mask, _ = self.build_edge_list_from_nlist(
             extended_coord=extended_coord,
             nlist=nlist,
             mapping=mapping,
+            sort_by_destination=sort_edges_by_dst,
         )
 
         # === Step 2. Force embedding ===
@@ -1767,6 +1961,7 @@ def core_compute_dens(
                 edge_index=edge_index,
                 edge_vec=edge_vec,
                 edge_mask=edge_mask,
+                edge_index_sorted_by_dst=sort_edges_by_dst,
                 force_embedding=force_embedding,
                 charge_spin=charge_spin,
                 vacuum_conditions=self.atomic_model.vacuum_conditions()
@@ -1963,6 +2158,53 @@ def trace_and_compile(
             )
             return
 
+        # Register Python-owned SO2 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.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.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import (
+                maybe_prepare_sm80_readout_input_fold,
+            )
+
+            maybe_prepare_sm80_readout_input_fold(
+                self.atomic_model.descriptor.output_ffn,
+                compute_capability,
+            )
+        prepared_cute_so2 = False
+        if (
+            not self.training
+            and compute_capability is not None
+            and cute_runtime_policy.is_cute_infer_enabled()
+            and cute_runtime_policy.is_supported_so2_capability(compute_capability)
+        ):
+            from deepmd.pt_expt.kernels.cute.sezm.so2.operation import (
+                prepare_cute_so2_blocks,
+            )
+
+            prepared_cute_so2 = prepare_cute_so2_blocks(
+                self.atomic_model.descriptor.blocks,
+                training=self.training,
+                device=coord.device,
+                dtype=self.atomic_model.descriptor.dtype,
+            )
+        if prepared_cute_so2 and cute_runtime_policy.is_so2_thin_wrapper_enabled(
+            compute_capability
+        ):
+            from ..network.mlp import (
+                enable_neo_cute_compile_visible_linears,
+            )
+
+            enable_neo_cute_compile_visible_linears(self)
+
         log.info(
             "SeZM: start tracing and compiling (mode=%s, coord_corr=%s)",
             mode,
@@ -2770,8 +3012,10 @@ def build_neighbor_list(
     ) -> EdgeNeighborList:
         """Build the unified edge-vector schema for the ``forward`` entry."""
         nf, nloc = atype.shape[:2]
-        return _select_neighbor_builder(nf, coord.device).build(
-            coord.view(nf, nloc, 3),
+        coord = coord.view(nf, nloc, 3)
+        return _build_neo_neighbor_list(
+            _select_neighbor_builder(nf, coord.device),
+            coord,
             atype,
             box,
             self.get_rcut(),
@@ -2797,8 +3041,10 @@ def build_extended_neighbor_list(
         contract order ``(extended_coord, extended_atype, nlist, mapping)``.
         """
         nf, nloc = atype.shape[:2]
-        return _select_neighbor_builder(nf, coord.device).build(
-            coord.view(nf, nloc, 3),
+        coord = coord.view(nf, nloc, 3)
+        return _build_neo_neighbor_list(
+            _select_neighbor_builder(nf, coord.device),
+            coord,
             atype,
             box,
             self.get_rcut(),
@@ -2812,6 +3058,7 @@ def build_edge_list_from_nlist(
         extended_coord: torch.Tensor,
         nlist: torch.Tensor,
         mapping: torch.Tensor | None,
+        sort_by_destination: bool | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         """
         Build a compact edge list from a DeePMD padded neighbor list.
@@ -2836,6 +3083,9 @@ def build_edge_list_from_nlist(
             DeePMD padded neighbor list with shape (nf, nloc, nsel).
         mapping
             Extended-to-local mapping with shape (nf, nall), or ``None``.
+        sort_by_destination
+            Whether to sort edges by destination and source. ``None`` preserves
+            the environment-controlled CuTe inference behavior.
 
         Returns
         -------
@@ -2845,7 +3095,10 @@ def build_edge_list_from_nlist(
         edge_vec
             Edge vectors with shape (E+2, 3).
         edge_mask
-            Boolean mask with shape (E+2,).  The two trailing elements are ``False``.
+            Boolean mask with shape (E+2,). The two padded elements are
+            ``False``. They are trailing only when sorting is disabled;
+            destination sorting may move them into the interior, so consumers
+            must select valid edges through this mask rather than by position.
         edge_scatter_index
             Scatter-domain (src, dst) indices with shape (2, E+2), aligned
             1:1 with ``edge_index`` and ``edge_vec``.
@@ -2862,12 +3115,24 @@ def build_edge_list_from_nlist(
             nlist,
             mapping,
         )
-        return (
+        edge_tensors = (
             edge_schema.edge_index,
             edge_schema.edge_vec,
             edge_schema.edge_mask,
             edge_schema.edge_scatter_index,
         )
+        should_sort = (
+            _neo_cute_so2_requires_sorted_edges(
+                self.atomic_model.descriptor,
+                training=self.training,
+                device=extended_coord.device,
+            )
+            if sort_by_destination is None
+            else bool(sort_by_destination)
+        )
+        if should_sort:
+            return _sort_edge_tensors_by_destination(*edge_tensors)
+        return edge_tensors
 
     # =========================================================================
     # Input Canonicalization
diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py
index 13ea438f4f..705d05c362 100644
--- a/deepmd/pt/model/network/mlp.py
+++ b/deepmd/pt/model/network/mlp.py
@@ -46,6 +46,37 @@ def empty_t(shape: tuple[int, ...], precision: torch.dtype) -> torch.Tensor:
     return torch.empty(shape, dtype=precision, device=device)
 
 
+@torch.compiler.assume_constant_result
+def _use_so2_compile_visible_linear(
+    input_device: torch.device | None = None,
+) -> bool:
+    """Use the same device-aware thin-wrapper policy as SO2 dispatch."""
+    from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import (
+        is_so2_thin_wrapper_enabled,
+    )
+
+    capability = None
+    if (
+        input_device is None or input_device.type == "cuda"
+    ) and torch.cuda.is_available():
+        try:
+            capability = tuple(torch.cuda.get_device_capability(input_device))
+        except RuntimeError:
+            pass
+    # A sentinel avoids querying another device when the input is on the CPU.
+    return is_so2_thin_wrapper_enabled(capability or (-1, -1))
+
+
+def _matmul_bias(
+    value: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+) -> torch.Tensor:
+    """Avoid eager addmm's expanded-bias copy and expose the add to Inductor."""
+    output = torch.matmul(value, weight)
+    return output if bias is None else output + bias
+
+
 class Identity(nn.Module):
     def __init__(self) -> None:
         super().__init__()
@@ -86,6 +117,7 @@ def __init__(
     ) -> None:
         super().__init__()
         self.trainable = trainable
+        self._deepmd_cute_compile_visible_linear = False
         # only use_timestep when skip connection is established.
         self.use_timestep = use_timestep and (
             num_out == num_in or num_out == num_in * 2
@@ -204,7 +236,20 @@ def forward(
         ori_prec = xx.dtype
         if not env.DP_DTYPE_PROMOTION_STRICT:
             xx = xx.to(self.prec)
-        yy = F.linear(xx, self.matrix.t(), self.bias)
+        if torch.jit.is_scripting():
+            yy = F.linear(xx, self.matrix.t(), self.bias)
+        elif (
+            not self.training
+            and xx.dtype == torch.float32
+            and self.matrix.dtype == torch.float32
+            and (self.bias is None or self.bias.dtype == torch.float32)
+            and not torch.is_autocast_enabled(xx.device.type)
+            and self._deepmd_cute_compile_visible_linear
+            and _use_so2_compile_visible_linear(xx.device)
+        ):
+            yy = _matmul_bias(xx, self.matrix, self.bias)
+        else:
+            yy = F.linear(xx, self.matrix.t(), self.bias)
         yy = self.activate(yy)
         yy = yy * self.idt if self.idt is not None else yy
         if self.resnet:
@@ -278,6 +323,13 @@ def check_load_param(ss: str) -> nn.Parameter | None:
         return obj
 
 
+def enable_neo_cute_compile_visible_linears(module: nn.Module) -> None:
+    """Select the alternate eval linear topology only inside one Neo model."""
+    for child in module.modules():
+        if isinstance(child, MLPLayer):
+            child._deepmd_cute_compile_visible_linear = True
+
+
 MLP_ = make_multilayer_network(MLPLayer, nn.Module)
 
 
diff --git a/deepmd/pt_expt/descriptor/dpa4.py b/deepmd/pt_expt/descriptor/dpa4.py
index db843ffb01..34e893a5bf 100644
--- a/deepmd/pt_expt/descriptor/dpa4.py
+++ b/deepmd/pt_expt/descriptor/dpa4.py
@@ -16,6 +16,11 @@
     C3CutoffEnvelope as C3CutoffEnvelopeDP,
 )
 from deepmd.dpmodel.descriptor.dpa4_nn.radial import InnerClamp as InnerClampDP
+from deepmd.dpmodel.utils.neighbor_graph import (
+    NeighborGraph,
+    attach_edge_csr,
+    compact_edges,
+)
 from deepmd.pt_expt.common import (
     register_dpmodel_mapping,
     torch_module,
@@ -216,19 +221,19 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # The fused convolution paths consume only the three structural rows of
         # each Wigner degree block. Source-gated attention bypasses that fused
         # convolution, so its dense per-edge rotations remain available.
-        self._wigner_free_conv = (
+        self.cuda_infer_l_2_covers_all_blocks = (
             self.bridging_switch is None
             and bool(self.blocks)
             and all(
-                getattr(block.so2_conv, "_cuda_conv_fn", None) is not None
-                and not block.so2_conv._cuda_conv_fn._compete
+                getattr(block.so2_conv, "cuda_infer_l_2_conv", None) is not None
+                and not block.so2_conv.cuda_infer_l_2_conv.focus_compete
                 for block in self.blocks
             )
         )
-        self._packed_wigner_train = bool(self.blocks) and all(
-            getattr(block.so2_conv, "_cuda_value_train", None) is not None
-            and block.so2_conv._flash_atten_fn is not None
-            and block.so2_conv._flash_atten_trains
+        self.cuda_train_covers_all_blocks = bool(self.blocks) and all(
+            getattr(block.so2_conv, "cuda_train_value", None) is not None
+            and block.so2_conv.flash_attention is not None
+            and block.so2_conv.flash_attention_supports_training
             for block in self.blocks
         )
 
@@ -236,8 +241,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # distance and are cheap enough that the compiler inlines them into
         # every consumer and re-evaluates them there. Behind an operator
         # boundary the chain runs once per step.
-        self._cuda_radial_fn = None
-        self._cuda_wigner_fn = None
+        self.cuda_infer_l_1_radial = None
+        self.cuda_infer_l_1_wigner = None
         if cuda_infer_level() >= 1:
             from deepmd.pt_expt.kernels.cuda.dpa4.edge_radial import (
                 make_cuda_edge_radial,
@@ -246,13 +251,13 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                 make_cuda_wigner_dense,
             )
 
-            self._cuda_radial_fn = make_cuda_edge_radial(
+            self.cuda_infer_l_1_radial = make_cuda_edge_radial(
                 self.edge_envelope, self.radial_basis
             )
             # The dense Wigner pair otherwise costs five full-size passes
             # over the (E, D, D) tensors; the fused build pays only the
             # output writes.
-            self._cuda_wigner_fn = make_cuda_wigner_dense(
+            self.cuda_infer_l_1_wigner = make_cuda_wigner_dense(
                 self.mp_init_lmax,
                 get_xp_precision(torch, self.compute_precision),
             )
@@ -310,17 +315,258 @@ def _shared_wigner_runs(self, edge_cache: Any, lmax: int) -> torch.Tensor | None
         if edge_cache.csr_cache is None:
             return None
         if self.training:
-            if not self._packed_wigner_train:
+            if not self.cuda_train_covers_all_blocks:
                 return None
-            fused = self.blocks[0].so2_conv._cuda_value_train
+            fused = self.blocks[0].so2_conv.cuda_train_value
         else:
-            if not self._wigner_free_conv:
+            if not self.cuda_infer_l_2_covers_all_blocks:
                 return None
-            fused = self.blocks[0].so2_conv._cuda_conv_fn
+            fused = self.blocks[0].so2_conv.cuda_infer_l_2_conv
         if fused is None or lmax > self.lmax:
             return None
         return fused.edge_runs(edge_cache)[:, 1 : (lmax + 1) ** 2]
 
+    def is_cute_infer_packed_wigner_candidate(
+        self,
+        device: torch.device,
+        dtype: torch.dtype,
+        geometry_dtype: torch.dtype,
+    ) -> bool:
+        """Return whether all blocks satisfy the packed CuTe SO2 contract.
+
+        Parameters
+        ----------
+        device : torch.device
+            Device that executes the descriptor blocks.
+        dtype : torch.dtype
+            Descriptor block dtype.
+        geometry_dtype : torch.dtype
+            Edge-geometry compute dtype.
+
+        Returns
+        -------
+        bool
+            Whether packed Wigner storage can replace dense storage for every
+            interaction block.
+        """
+        from deepmd.pt_expt.kernels.cute.sezm.so2.operation import (
+            is_packed_wigner_candidate,
+        )
+
+        return is_packed_wigner_candidate(
+            blocks=self.blocks,
+            training=self.training,
+            device=device,
+            dtype=dtype,
+            producer_modules=(
+                self.radial_basis,
+                self.radial_embedding,
+                self.edge_envelope,
+                *((self.inner_clamp,) if self.inner_clamp is not None else ()),
+                *((self.bridging_switch,) if self.bridging_switch is not None else ()),
+            ),
+            producer_dtypes=(dtype, geometry_dtype),
+            has_edge_src_gate=self.bridging_switch is not None,
+        )
+
+    def prepare_packed_wigner_graph(
+        self,
+        graph: NeighborGraph,
+        n_nodes: int,
+    ) -> NeighborGraph | None:
+        """Prepare a compact destination-major graph for CuTe packed SO2.
+
+        Parameters
+        ----------
+        graph : NeighborGraph
+            Edge graph supplied to the descriptor.
+        n_nodes : int
+            Number of nodes addressed by the graph.
+
+        Returns
+        -------
+        NeighborGraph or None
+            Canonical graph without masked edges, or ``None`` when the exact
+            CuTe contract is not satisfied.
+        """
+        dtype = get_xp_precision(torch, self.precision)
+        if not self.is_cute_infer_packed_wigner_candidate(
+            graph.edge_vec.device,
+            dtype,
+            graph.edge_vec.dtype,
+        ):
+            return None
+        # Graph-native inputs retain canonical CSR through compaction. Dense
+        # adapters and exclusion transforms construct it once at this boundary.
+        graph = compact_edges(graph)
+        if not graph.destination_sorted:
+            graph = attach_edge_csr(graph, n_nodes, canonicalize=True)
+        from deepmd.pt_expt.kernels.cute.sezm.so2.operation import (
+            packed_wigner_edges_eligible,
+        )
+
+        if not packed_wigner_edges_eligible(
+            candidate=True,
+            edge_count=graph.edge_index.shape[1],
+            node_count=n_nodes,
+            destinations_sorted=graph.destination_sorted,
+            runtime_dtypes=(
+                dtype,
+                get_xp_precision(torch, self.compute_precision),
+            ),
+        ):
+            return None
+        return graph
+
+    def build_packed_wigner(
+        self,
+        edge_quat: torch.Tensor,
+        wigner_calc: Any,
+    ) -> torch.Tensor | None:
+        """Build the CuTe packed Wigner panel for an eligible graph.
+
+        Parameters
+        ----------
+        edge_quat : torch.Tensor
+            Global-to-local edge quaternions with shape ``(E, 4)``.
+        wigner_calc : Any
+            Wigner calculator carrying the degree and basis convention.
+
+        Returns
+        -------
+        torch.Tensor or None
+            CuTe packed Wigner storage with shape ``(E, 46)``, or ``None`` when
+            the kernel declines the runtime input.
+        """
+        from deepmd.pt_expt.kernels.cute.sezm.wignerd import (
+            run_cute_wignerd,
+        )
+
+        result = run_cute_wignerd(
+            edge_quat,
+            wigner_calc,
+            packed_wigner=True,
+        )
+        return None if result is None else result[0]
+
+    def prepare_cute_infer_so2_metadata(
+        self,
+        edge_cache: Any,
+        n_nodes: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
+        """Adapt graph-owned CSR metadata for the CuTe SO2 implementation.
+
+        Parameters
+        ----------
+        edge_cache : EdgeCache
+            Per-forward cache carrying the destination-major edge payload.
+        n_nodes : int
+            Number of nodes addressed by the edge payload.
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor] or None
+            Destination row pointers, source order, and source row pointers, or
+            ``None`` when the exact CuTe contract is not satisfied.
+        """
+        if (
+            self.training
+            or not edge_cache.destinations_sorted
+            or edge_cache.D_packed is None
+            or edge_cache.edge_src_gate is not None
+        ):
+            return None
+        from deepmd.pt_expt.kernels.cute.sezm import (
+            runtime_policy,
+        )
+
+        if not runtime_policy.is_cute_infer_enabled():
+            return None
+        csr_cache = edge_cache.csr_cache
+        if csr_cache is None or "dst" not in csr_cache or "src" not in csr_cache:
+            raise RuntimeError("packed CuTe SO2 requires destination/source graph CSR")
+        destination_row_ptr = csr_cache["dst"][1]
+        source_order, source_row_ptr = csr_cache["src"]
+        if (
+            destination_row_ptr.shape != (n_nodes + 1,)
+            or source_order.shape != edge_cache.src.shape
+            or source_row_ptr.shape != (n_nodes + 1,)
+        ):
+            raise RuntimeError("packed CuTe SO2 received inconsistent graph CSR shapes")
+        if runtime_policy.is_cute_strict_enabled() and edge_cache.dst.numel() > 1:
+            torch._assert_async(
+                torch.all(edge_cache.dst[1:] >= edge_cache.dst[:-1]),
+                "Neo SO2 destinations_sorted=True requires monotonically "
+                "nondecreasing destination indices",
+            )
+        return (
+            destination_row_ptr.to(dtype=torch.int32).contiguous(),
+            source_order.to(dtype=torch.int32).contiguous(),
+            source_row_ptr.to(dtype=torch.int32).contiguous(),
+        )
+
+    def build_cute_infer_zonal_coupling(self, edge_cache: Any) -> torch.Tensor:
+        """Extract the GIE zonal coupling from packed Wigner storage.
+
+        Parameters
+        ----------
+        edge_cache : EdgeCache
+            Per-forward cache carrying CuTe packed Wigner storage.
+
+        Returns
+        -------
+        torch.Tensor
+            Zonal coupling with shape ``(E, D_node - 1)``.
+
+        Raises
+        ------
+        ValueError
+            If the edge cache does not carry packed Wigner storage.
+        """
+        D_packed = edge_cache.D_packed
+        if D_packed is None:
+            raise ValueError("CuTe zonal coupling requires packed Wigner storage")
+        mp_coupling = D_packed.index_select(1, self.gie.packed_zonal_offsets)
+        if self.gie_zonal_wigner_calc is None:
+            return mp_coupling
+        extra_coupling = self.gie_zonal_wigner_calc.forward_zonal(
+            self._edge_quaternion(edge_cache),
+            lmin=self.lmax + 1,
+        )
+        return torch.cat([mp_coupling, extra_coupling], dim=1)
+
+    def run_cute_infer_readout(
+        self,
+        ffn_in: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Run the CuTe readout when its exact inference contract matches.
+
+        Parameters
+        ----------
+        ffn_in : torch.Tensor
+            Equivariant readout input with shape ``(N, D, 1, C)``.
+
+        Returns
+        -------
+        torch.Tensor or None
+            Residual-inclusive scalar output with shape ``(N, C)``, or ``None``
+            when the exact CuTe contract is not satisfied.
+        """
+        from deepmd.pt_expt.kernels.cute.sezm import (
+            runtime_policy,
+        )
+
+        if self.training or not runtime_policy.is_cute_infer_enabled():
+            return None
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import (
+            maybe_run_neo_readout_l0,
+        )
+
+        return maybe_run_neo_readout_l0(
+            self.output_ffn,
+            ffn_in,
+        )
+
     @classmethod
     def deserialize(cls, data: dict) -> "DescrptDPA4":
         # deserialize assigns numpy arrays after __init__, which demotes
diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/block.py b/deepmd/pt_expt/descriptor/dpa4_nn/block.py
index af14de1b49..732c4154f1 100644
--- a/deepmd/pt_expt/descriptor/dpa4_nn/block.py
+++ b/deepmd/pt_expt/descriptor/dpa4_nn/block.py
@@ -150,6 +150,54 @@ def _run_so2_unit(
             )
         return self._run_so2_unit_impl(x, edge_cache, radial_feat)
 
+    def _run_so2_unit_impl(
+        self,
+        x: torch.Tensor,
+        edge_cache: EdgeCache,
+        radial_feat: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run the SO(2) unit implementation."""
+        if edge_cache.D_packed is not None:
+            from deepmd.pt_expt.kernels.cute.sezm.so2.operation import (
+                maybe_run_cute_so2,
+            )
+
+            metadata = edge_cache.cute_infer_so2_metadata
+            destination_row_ptr, source_order, source_row_ptr = (
+                (None, None, None) if metadata is None else metadata
+            )
+            accelerated = maybe_run_cute_so2(
+                self,
+                x,
+                edge_cache,
+                radial_feat,
+                dst_ptr=destination_row_ptr,
+                source_order=source_order,
+                source_ptr=source_row_ptr,
+            )
+            if accelerated is not None:
+                return accelerated
+            from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import (
+                dense_wigner_for_fallback,
+            )
+
+            if edge_cache.edge_quat is None:
+                raise ValueError("packed Wigner fallback requires edge quaternions")
+            d_full, dt_full = dense_wigner_for_fallback(
+                edge_cache.edge_quat,
+                lmax=self.lmax,
+                eps=self.so2_conv.eps,
+            )
+            edge_cache = dataclasses.replace(
+                edge_cache,
+                D_full=d_full,
+                Dt_full=dt_full,
+                D_packed=None,
+                D_to_m_cache=None,
+                Dt_from_m_cache=None,
+            )
+        return super()._run_so2_unit_impl(x, edge_cache, radial_feat)
+
     def _run_ffn_unit(self, x: torch.Tensor, unit_idx: int) -> torch.Tensor:
         if self._use_infer_activation_checkpoint(x):
             return checkpoint(
diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/embedding.py b/deepmd/pt_expt/descriptor/dpa4_nn/embedding.py
index 3dd60102d2..4fc9d156cf 100644
--- a/deepmd/pt_expt/descriptor/dpa4_nn/embedding.py
+++ b/deepmd/pt_expt/descriptor/dpa4_nn/embedding.py
@@ -20,6 +20,9 @@
 from deepmd.pt_expt.common import (
     torch_module,
 )
+from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import (
+    ZONAL_PANEL_OFFSETS,
+)
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
 )
@@ -35,14 +38,23 @@ class GeometricInitialEmbedding(GeometricInitialEmbeddingDP):
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
+        self.register_buffer(
+            "packed_zonal_offsets",
+            torch.tensor(
+                ZONAL_PANEL_OFFSETS if self.lmax == 3 else (),
+                device=self.non_scalar_row_index.device,
+                dtype=torch.long,
+            ),
+            persistent=False,
+        )
         # === Fused message-and-scatter operator ===
         # The reference composition materializes the per-edge message, an
         # (E, D-1, C) tensor that dominates the cost of this module. The fused
         # operator keeps it in registers and reduces through the destination CSR.
-        self._cuda_scatter = False
+        self.cuda_infer_l_1_scatter = False
         # ``None`` keeps the runtime ``zonal_coupling.is_cuda`` dispatch; the
         # freeze pins it to the AOTI target because tracing always runs on CPU.
-        self._force_fused_scatter: bool | None = None
+        self.force_cuda_infer_l_1_scatter: bool | None = None
         if (
             cuda_infer_level() >= 1
             and get_xp_precision(torch, self.precision) is torch.float32
@@ -52,20 +64,75 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                 supported,
             )
 
-            self._cuda_scatter = op_available() and supported(
+            self.cuda_infer_l_1_scatter = op_available() and supported(
                 self.lmax, self.ebed_dim - 1, self.channels
             )
 
-    def _can_fuse_scatter(self, zonal_coupling: torch.Tensor) -> bool:
-        """Return whether the fused scatter serves the runtime or trace target."""
+    def run_cute_infer_gie(
+        self,
+        *,
+        n_nodes: int | torch.SymInt,
+        edge_cache: Any,
+        radial_feat: torch.Tensor,
+        zonal_coupling: torch.Tensor,
+        spin_l1_message: torch.Tensor | None = None,
+    ) -> torch.Tensor | None:
+        """Run the CuTe GIE path when its exact inference contract matches.
+
+        Parameters
+        ----------
+        n_nodes : int or torch.SymInt
+            Number of nodes addressed by the edge cache.
+        edge_cache : EdgeCache
+            Per-forward geometric edge cache.
+        radial_feat : torch.Tensor
+            Radial features with shape ``(E, lmax, C)``.
+        zonal_coupling : torch.Tensor
+            Zonal coupling with shape ``(E, D - 1)``.
+        spin_l1_message : torch.Tensor, optional
+            Native-spin message with shape ``(E, 3, C)``.
+
+        Returns
+        -------
+        torch.Tensor or None
+            Initial embedding with shape ``(N, D, C)``, or ``None`` when the
+            exact CuTe contract is not satisfied.
+        """
+        if self.training or spin_l1_message is not None:
+            return None
+        from deepmd.pt_expt.kernels.cute.sezm.gie import (
+            maybe_run_cute_gie,
+        )
+
+        return maybe_run_cute_gie(
+            self,
+            n_nodes=n_nodes,
+            edge_cache=edge_cache,
+            radial_feat=radial_feat,
+            zonal_coupling=zonal_coupling,
+        )
+
+    def can_run_cuda_infer_l_1_scatter(self, zonal_coupling: torch.Tensor) -> bool:
+        """Return whether the fused scatter serves the runtime or trace target.
+
+        Parameters
+        ----------
+        zonal_coupling : torch.Tensor
+            Zonal coupling with shape ``(E, D - 1)``.
+
+        Returns
+        -------
+        bool
+            Whether the CUDA inference level-one scatter is eligible.
+        """
         target_is_cuda = (
             zonal_coupling.is_cuda
-            if self._force_fused_scatter is None
-            else self._force_fused_scatter
+            if self.force_cuda_infer_l_1_scatter is None
+            else self.force_cuda_infer_l_1_scatter
         )
-        return self._cuda_scatter and not self.training and target_is_cuda
+        return self.cuda_infer_l_1_scatter and not self.training and target_is_cuda
 
-    def forward_fused_scatter(
+    def run_cuda_infer_l_1_scatter(
         self,
         n_nodes: int | torch.SymInt,
         edge_cache: Any,
diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/grid_net.py b/deepmd/pt_expt/descriptor/dpa4_nn/grid_net.py
index 909ea7ca82..b88da2b9bc 100644
--- a/deepmd/pt_expt/descriptor/dpa4_nn/grid_net.py
+++ b/deepmd/pt_expt/descriptor/dpa4_nn/grid_net.py
@@ -7,6 +7,9 @@
 
 import torch
 
+from deepmd.dpmodel.common import (
+    get_xp_precision,
+)
 from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import S2GridNet as S2GridNetDP
 from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import SO3GridNet as SO3GridNetDP
 from deepmd.pt_expt.common import (
@@ -18,7 +21,7 @@
 )
 
 
-def _bind_grid_pair(module: Any) -> None:
+def bind_grid_pair_operators(module: Any) -> None:
     """Bind the fused coefficient-grid pair operators that serve the layout.
 
     The inference operator (CUDA, register-resident walk) and the training
@@ -28,7 +31,13 @@ def _bind_grid_pair(module: Any) -> None:
     below 75 slots the dense section is small and the operator's dispatch
     chain costs more than its kernels save on the host-bound
     configurations, so the narrow grids stay with the compiler.
+
+    Parameters
+    ----------
+    module : Any
+        PT-expt grid module receiving the eligible operator bindings.
     """
+    module.dtype = get_xp_precision(torch, module.precision)
     if module.projector.to_grid_mat.dtype is not torch.float32:
         return
     slots = int(module.projector.to_grid_mat.shape[1])
@@ -40,7 +49,7 @@ def _bind_grid_pair(module: Any) -> None:
         )
 
         if op_available() and slots in SUPPORTED_SLOTS:
-            module._grid_pair_fn = grid_pair
+            module.cuda_infer_l_1_grid_pair = grid_pair
     if triton_train_level() >= 1 and slots >= 75:
         from deepmd.pt_expt.kernels.triton.sezm.grid_pair import (
             GRID_PAIR_TRITON_AVAILABLE,
@@ -48,7 +57,52 @@ def _bind_grid_pair(module: Any) -> None:
         )
 
         if GRID_PAIR_TRITON_AVAILABLE:
-            module._grid_pair_train_fn = grid_pair_train
+            module.triton_train_l_1_grid_pair = grid_pair_train
+
+
+def run_cute_infer_grid_pair(
+    module: Any,
+    left: torch.Tensor,
+    right: torch.Tensor,
+) -> torch.Tensor | None:
+    """Run the CuTe grid product when its exact inference contract matches.
+
+    Parameters
+    ----------
+    module : Any
+        Grid module carrying the projector and frozen inference state.
+    left : torch.Tensor
+        Left coefficient operand with shape ``(N, D, F, n_frames * C)``.
+    right : torch.Tensor
+        Right coefficient operand with the same shape as ``left``.
+
+    Returns
+    -------
+    torch.Tensor or None
+        Coefficient result with the same shape as ``left``, or ``None`` when
+        the exact CuTe contract is not satisfied.
+    """
+    from deepmd.pt_expt.kernels.cute.sezm import (
+        runtime_policy,
+    )
+
+    if (
+        not runtime_policy.is_cute_infer_enabled()
+        or module.training
+        or any(parameter.requires_grad for parameter in module.parameters())
+    ):
+        return None
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid.product import (
+        maybe_run_cute_output_grid_product,
+    )
+
+    return maybe_run_cute_output_grid_product(
+        left,
+        right,
+        module.projector.to_grid_mat,
+        module.projector.from_grid_mat,
+        n_frames=module.n_frames,
+    )
 
 
 @torch_module
@@ -57,7 +111,29 @@ class S2GridNet(S2GridNetDP):
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
-        _bind_grid_pair(self)
+        bind_grid_pair_operators(self)
+
+    def run_cute_infer_grid_pair(
+        self,
+        left: torch.Tensor,
+        right: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Run the CuTe grid product when its exact inference contract matches.
+
+        Parameters
+        ----------
+        left : torch.Tensor
+            Left coefficient operand with shape ``(N, D, F, n_frames * C)``.
+        right : torch.Tensor
+            Right coefficient operand with the same shape as ``left``.
+
+        Returns
+        -------
+        torch.Tensor or None
+            Coefficient result with the same shape as ``left``, or ``None``
+            when the exact CuTe contract is not satisfied.
+        """
+        return run_cute_infer_grid_pair(self, left, right)
 
 
 @torch_module
@@ -66,4 +142,26 @@ class SO3GridNet(SO3GridNetDP):
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
-        _bind_grid_pair(self)
+        bind_grid_pair_operators(self)
+
+    def run_cute_infer_grid_pair(
+        self,
+        left: torch.Tensor,
+        right: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """Run the CuTe grid product when its exact inference contract matches.
+
+        Parameters
+        ----------
+        left : torch.Tensor
+            Left coefficient operand with shape ``(N, D, F, n_frames * C)``.
+        right : torch.Tensor
+            Right coefficient operand with the same shape as ``left``.
+
+        Returns
+        -------
+        torch.Tensor or None
+            Coefficient result with the same shape as ``left``, or ``None``
+            when the exact CuTe contract is not satisfied.
+        """
+        return run_cute_infer_grid_pair(self, left, right)
diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py
index 87e3344e7c..f200becec1 100644
--- a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py
+++ b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py
@@ -1,19 +1,19 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""pt_expt SO(2) runtime bindings for accelerated inference kernels.
+"""PT-expt SO(2) runtime bindings for accelerated kernels.
 
-The dpmodel SO(2) modules are array-API only. These wrappers inject the
-reference PT inference paths around three hot paths, mirroring
+The array-API dpmodel modules define the shared execution seams. These wrappers
+bind the same accelerated implementations and dispatch names used by
 ``deepmd.pt.model.descriptor.sezm_nn.so2``:
 
 - the block-diagonal GEMM of :class:`SO2Linear`,
-- the two rotation hot paths of :class:`SO2Convolution`, and
+- the rotation, value, and attention paths of :class:`SO2Convolution`, and
 - the low-rank branch of :class:`DynamicRadialDegreeMixer`.
 
-Triton, CuTe, and cuTile are mutually exclusive complete SO(2) paths. The
-hand-written CUDA operators form an independent cumulative layer and take
-precedence where their factories bind. Every gate is resolved at construction
-so export records a static dispatch choice; training and unsupported layouts
-retain the dpmodel reference path.
+Triton and cuTile are mutually exclusive complete SO(2) paths. Hand-written
+CUDA operators form an independent cumulative layer, while CuTe replaces exact
+Neo spans at the enclosing block. Every gate is resolved at construction so
+tracing records a static dispatch choice; unsupported layouts retain the
+dpmodel reference path.
 """
 
 from __future__ import (
@@ -27,6 +27,9 @@
 
 import torch
 
+from deepmd.dpmodel.common import (
+    get_xp_precision,
+)
 from deepmd.dpmodel.descriptor.dpa4_nn.so2 import (
     DynamicRadialDegreeMixer as DynamicRadialDegreeMixerDP,
 )
@@ -40,7 +43,6 @@
     cuda_train_enabled,
     triton_infer_level,
     triton_train_level,
-    use_cute_infer,
     use_cutile_infer,
 )
 
@@ -49,13 +51,16 @@
 )
 
 
-def _active_triton_level(module: Any) -> int:
-    """Return the Triton dispatch level governing the module's current mode.
+def active_triton_level(module: Any) -> int:
+    """
+    Return the Triton acceleration level that applies to a module's mode.
 
-    The levels are read at construction to decide which kernels to bind, but
-    consulted here at call time: inference and training are separate gates,
-    and a module built with both bound must follow whichever one matches the
-    mode it is being called in.
+    Inference and training are gated independently: an operator qualifies for
+    inference as soon as it reproduces the forward and the coordinate gradient
+    with the parameters held fixed, whereas training additionally requires the
+    gradient of every parameter it consumes and a second derivative of its own
+    backward, which the force loss traverses. Both levels are captured at
+    construction, so the branch this drives resolves at trace time.
 
     Parameters
     ----------
@@ -65,7 +70,7 @@ def _active_triton_level(module: Any) -> int:
     Returns
     -------
     int
-        The training level in training mode, the inference level otherwise.
+        The level in effect for the module's current mode.
     """
     return module.triton_train_level if module.training else module.triton_infer_level
 
@@ -85,7 +90,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # Export override for the block-diagonal vs dense matmul branch below.
         # ``None`` keeps the runtime ``x_flat.is_cuda`` dispatch; the freeze sets
         # it so the AOTI graph follows the *target* device, not the CPU trace.
-        self._force_block_diag_matmul: bool | None = None
+        self.force_block_diag_matmul: bool | None = None
 
         # Fast path (``DP_TRITON_INFER >= 1`` or ``DP_TRITON_TRAIN >= 1``):
         # the per-|m|-block batched bmm + cat of ``_block_diagonal_matmul`` is
@@ -98,7 +103,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # constants in the traced (``make_fx``) graph.
         self.triton_infer_level = triton_infer_level()
         self.triton_train_level = triton_train_level()
-        self._block_diag_gemm = None
+        self.triton_l_1_block_diag_gemm = None
         if max(self.triton_infer_level, self.triton_train_level) >= 1:
             from deepmd.pt_expt.kernels.triton.sezm.so2_block_gemm import (
                 SO2_BLOCK_GEMM_TRITON_AVAILABLE,
@@ -109,7 +114,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             if SO2_BLOCK_GEMM_TRITON_AVAILABLE and slices_supported(
                 self._block_diag_slices
             ):
-                self._block_diag_gemm = block_diag_gemm
+                self.triton_l_1_block_diag_gemm = block_diag_gemm
 
     def _block_diagonal_matmul(
         self, x_flat: torch.Tensor, weight: torch.Tensor
@@ -118,22 +123,27 @@ def _block_diagonal_matmul(
         # trips an Inductor AVX2 C++ codegen bug, so only CPU needs it. Every other
         # device uses the block-diagonal contraction, which skips the structural
         # off-|m| zeros. ``make_fx`` resolves this Python branch at trace time, so
-        # the freeze pins ``_force_block_diag_matmul`` to the AOTI target device
+        # the freeze pins ``force_block_diag_matmul`` to the AOTI target device
         # (tracing always runs on CPU regardless of where the artifact will run).
-        if self._force_block_diag_matmul is None:
+        if self.force_block_diag_matmul is None:
             use_block_diag = not x_flat.is_cpu
         else:
-            use_block_diag = self._force_block_diag_matmul
+            use_block_diag = self.force_block_diag_matmul
         if not use_block_diag:
             return torch.einsum("fei,ifo->feo", x_flat, weight)
-        if self._block_diag_gemm is not None and _active_triton_level(self) >= 1:
+        if (
+            self.triton_l_1_block_diag_gemm is not None
+            and active_triton_level(self) >= 1
+        ):
             # The fused GEMM consumes the ``(F, D_m*Cin, D_m*Cout)`` presentation
             # directly from the strided weight, so the permute is applied here and
             # the contiguity copy the dpmodel ``bmm`` cat path would need is
             # skipped. The eager fallback permutes ``weight`` internally, so it is
             # passed the stored ``(D_m*Cin, F, D_m*Cout)`` layout untouched.
             weight = weight.permute(1, 0, 2)  # (F, D_m*Cin, D_m*Cout)
-            return self._block_diag_gemm(x_flat, weight, self._block_diag_slices)
+            return self.triton_l_1_block_diag_gemm(
+                x_flat, weight, self._block_diag_slices
+            )
         return super()._block_diagonal_matmul(x_flat, weight)
 
 
@@ -152,7 +162,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # constants in the traced (``make_fx``) graph.
         self.triton_infer_level = triton_infer_level()
         self.triton_train_level = triton_train_level()
-        self._radial_mix_block = None
+        self.triton_l_1_radial_mix = None
         if (
             max(self.triton_infer_level, self.triton_train_level) >= 1
             and self.mode == "degree_channel"
@@ -163,13 +173,13 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                 radial_mix_block,
             )
 
-            self._radial_mix_block = radial_mix_block
+            self.triton_l_1_radial_mix = radial_mix_block
 
     def _mix_rank_compact(
         self, compact: torch.Tensor, x_local: torch.Tensor
     ) -> torch.Tensor:
-        if self._radial_mix_block is not None and _active_triton_level(self) >= 1:
-            return self._radial_mix_block(
+        if self.triton_l_1_radial_mix is not None and active_triton_level(self) >= 1:
+            return self.triton_l_1_radial_mix(
                 compact, x_local, self.channel_basis, self.lmax
             )
         return super()._mix_rank_compact(compact, x_local)
@@ -177,37 +187,41 @@ def _mix_rank_compact(
 
 @torch_module
 class SO2Convolution(SO2ConvolutionDP):
-    """SO(2) convolution with opt-in accelerated inference kernels."""
+    """SO(2) convolution with opt-in inference and training kernels."""
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
-        # The inference gates are read once at construction so they become
-        # compile-time constants in the traced (``make_fx``) graph. Triton,
-        # CuTe and cuTile claim the same SO(2) value path and are mutually
-        # exclusive; the hand-written CUDA operators form an independent,
-        # cumulative layer and take precedence where their factories bind.
+        self.dtype = get_xp_precision(torch, self.precision)
+        self.compute_dtype = get_xp_precision(torch, self.compute_precision)
+        # Acceleration gates are read once at construction so they become
+        # compile-time constants in the traced (``make_fx``) graph. Triton level
+        # 1 replaces individual operators, level 2 binds the fused value path,
+        # and inference level 3 selects its fp16x3 implementation on swept
+        # shapes. ``DP_CUTILE_INFER`` selects a complete alternative SO(2) path
+        # and is mutually exclusive with Triton. ``DP_CUDA_INFER`` and
+        # ``DP_CUDA_TRAIN`` independently bind hand-written CUDA producers that
+        # take precedence over the corresponding narrower paths. CuTe dispatch
+        # is resolved at the enclosing block, where its exact-shape SO2 kernel
+        # can coexist with every fallback. Unsupported entries remain ``None``.
         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:
+        self.cuda_infer_level = cuda_infer_level()
+        self.cutile_infer_enabled = use_cutile_infer()
+        if self.triton_infer_level >= 1 and self.cutile_infer_enabled:
             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
-        self._cached_edge_csr_fn = cached_edge_csr
+        self.triton_infer_l_2_value = None
+        self.cutile_infer_value = None
+        self.edge_csr = cached_edge_csr
 
         # === Triton rotation kernels: block for mmax == 1, dense otherwise ===
         # The rotation operators carry differentiable backwards (the force
         # loss traverses them twice), so the training gate binds them as well;
-        # ``_active_triton_level`` then selects the path per mode.
-        self._rotate_to_local_fn = None
-        self._rotate_back_fn = None
+        # ``active_triton_level`` then selects the path per mode.
+        self.triton_l_1_rotate_to_local = None
+        self.triton_l_1_rotate_back = None
         if max(self.triton_infer_level, self.triton_train_level) >= 1:
             from deepmd.pt_expt.kernels.triton.sezm.so2_rotation import (
                 rotate_back_block_so2,
@@ -217,24 +231,26 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             )
 
             if self.mmax == 1:
-                self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_block(
-                    x, src, wigner, self.lmax
+                self.triton_l_1_rotate_to_local = lambda x, src, wigner: (
+                    rotate_to_local_block(x, src, wigner, self.lmax)
                 )
                 # The block kernel reads the (E, F, D_m, Cf) focus layout directly,
                 # so the rotate-back path passes ``x_local`` before the global
                 # reshape and the transpose-back copy is skipped.
-                self._rotate_back_fn = lambda x_local, wigner: rotate_back_block_so2(
-                    x_local, wigner, self.lmax
+                self.triton_l_1_rotate_back = lambda x_local, wigner: (
+                    rotate_back_block_so2(x_local, wigner, self.lmax)
                 )
             else:
-                self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_dense(
-                    x, src, wigner, self.coeff_index_m, self.ebed_dim_full
+                self.triton_l_1_rotate_to_local = lambda x, src, wigner: (
+                    rotate_to_local_dense(
+                        x, src, wigner, self.coeff_index_m, self.ebed_dim_full
+                    )
                 )
-                self._rotate_back_fn = lambda x_local, wigner: rotate_back_dense(
+                self.triton_l_1_rotate_back = lambda x_local, wigner: rotate_back_dense(
                     x_local, wigner, self.coeff_index_m, self.ebed_dim_full
                 )
 
-        # === Step 12. Optional fused flash-attention aggregation kernel ===
+        # === Step 11. Optional fused flash-attention aggregation kernel ===
         # Folds the entire ``n_atten_head > 0`` value aggregation -- block-diagonal
         # rotate-back, inverse-rotation rescale, envelope-gated softmax weighting,
         # and the destination scatter -- into a single destination-segmented
@@ -248,29 +264,30 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # attention layout without the optional focus-mix / value / output
         # projections (the deployed DPA4 configuration). Whichever of the
         # mutually exclusive inference gates is active then supplies the
-        # implementation, and ``self._flash_atten_fn`` being bound is what marks
+        # implementation, and ``self.flash_attention`` being bound is what marks
         # the fused path as live.
         # The cuTile aggregation is inference-only; the Triton one also serves
-        # training (analytic backward and second order), so it is bound
-        # whenever either gate asks for level 1 and ``_flash_atten_trains``
-        # marks it as training-capable for the dpmodel dispatch.
-        if self._flash_atten_layout_ok and self.use_cutile_infer:
+        # training, so it is bound whenever either gate asks for level 1.
+        self.flash_attention = None
+        self.flash_attention_supports_training = False
+        if self.flash_attention_layout_supported and self.cutile_infer_enabled:
             from deepmd.pt_expt.kernels.cutile.sezm.flash_atten import (
                 flash_atten_aggregate,
             )
 
-            self._flash_atten_fn = flash_atten_aggregate
-        elif self._flash_atten_layout_ok and (
-            self.use_triton_infer or self.triton_train_level >= 1
+            self.flash_attention = flash_atten_aggregate
+        elif (
+            self.flash_attention_layout_supported
+            and max(self.triton_infer_level, self.triton_train_level) >= 1
         ):
             from deepmd.pt_expt.kernels.triton.sezm.flash_atten import (
                 flash_atten_aggregate,
             )
 
-            self._flash_atten_fn = flash_atten_aggregate
-            self._flash_atten_trains = self.triton_train_level >= 1
+            self.flash_attention = flash_atten_aggregate
+            self.flash_attention_supports_training = self.triton_train_level >= 1
 
-        # === Step 13. Optional fused Triton SO(2) value-path operators ===
+        # === Step 12. Optional fused Triton SO(2) value path (inference) ===
         # 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.
@@ -281,66 +298,32 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # it engages at ``DP_TRITON_INFER >= 2``; at level 3 the factory
         # additionally routes the mixing stack through the fp16x3 tensor-core
         # operator on shapes whose configuration passed the fp64 validation
-        # sweep.
+        # sweep. Training never composes this path: the training value stream
+        # is either the level-1 operator composition (step 13 and the dense
+        # mixing stack) or the fused CUDA value path (step 15).
         if self.triton_infer_level >= 2:
             from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import (
                 make_triton_value_path,
             )
 
-            self._triton_value_path = make_triton_value_path(self)
+            self.triton_infer_l_2_value = make_triton_value_path(self)
 
-        # === Step 13b. Optional fused CUDA SO(2) convolution ===
-        # One hand-written CUDA operator spans the complete per-edge path:
-        # rotate-to-local, the radial degree mixing, the gated mixing stack, the
-        # inverse rotation, the attention weighting and the destination
-        # reduction. It therefore supersedes both the fused value path and the
-        # flash aggregation, and takes precedence over them when the block
-        # matches its supported configuration. The factory returns ``None``
-        # otherwise, leaving whichever narrower path is bound in charge.
-        self._cuda_conv_fn = None
-        if cuda_infer_level() >= 2 and self._flash_atten_layout_ok:
-            from deepmd.pt_expt.kernels.cuda.dpa4 import (
-                make_cuda_so2_conv,
-            )
-
-            self._cuda_conv_fn = make_cuda_so2_conv(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 deepmd.pt_expt.kernels.cute.sezm import (
-                make_cute_value_path,
-            )
-
-            self._cute_value_path = make_cute_value_path(self)
-
-        # === Step 15. 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``
-        # otherwise, leaving the dense reference path in charge.
-        if self.use_cutile_infer:
-            from deepmd.pt_expt.kernels.cutile.sezm.so2_value_path import (
-                make_cutile_value_path,
-            )
-
-            self._cutile_value_path = make_cutile_value_path(self)
-
-        # === Step 16. Optional fused rotate-mix operator ===
-        # One edge-parallel kernel gathers the source features, applies the
-        # block-diagonal Wigner rotation and the radial degree mixing, and
-        # writes the focus-major mixing input directly; the degree-expanded
-        # local intermediate and its relayout never reach the traced graph.
-        # The operator carries a differentiable backward and a hand-derived
-        # second order, so it serves force-loss training.
+        # === Step 13. Optional fused rotate-to-local + radial degree mixing ===
+        # Level-1 companion of the value path: only the rotation and the
+        # degree mixing fuse into one edge-parallel operator writing the
+        # focus-major mixing input directly, while the mixing stack itself
+        # stays with the compiler. This removes the degree-expanded local
+        # intermediate and its relayout from the traced graph; the operator's
+        # backward reduces through the source CSR view and carries a
+        # hand-derived second order, so it serves force-loss training.
         #
         # The operator is quadrilinear, so a force loss re-enters its forward
         # and backward several times for the second order. That fixed cost is
         # repaid only where the materialization it removes is large: the wide
         # hidden widths. Below the bound the separate rotation and radial-mix
-        # kernels win end to end, so the binding follows the measured
-        # crossover (see :func:`_rotate_mix_supported`).
-        self._triton_rotate_mix = None
+        # kernels win end to end (a 64-wide stack loses ~10%, a 128-wide one
+        # gains), so the binding follows the measured crossover.
+        self.triton_l_1_rotate_mix = None
         if (
             max(self.triton_infer_level, self.triton_train_level) >= 1
             and self.hidden_channels >= 128
@@ -349,9 +332,25 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                 make_triton_rotate_mix,
             )
 
-            self._triton_rotate_mix = make_triton_rotate_mix(self)
+            self.triton_l_1_rotate_mix = make_triton_rotate_mix(self)
+
+        # === Step 14. Optional fused CUDA SO(2) convolution (inference) ===
+        # One hand-written CUDA operator spans the complete per-edge path:
+        # rotate-to-local, the radial degree mixing, the gated mixing stack, the
+        # inverse rotation, the attention weighting and the destination
+        # reduction. It therefore supersedes both the fused value path and the
+        # flash aggregation, and takes precedence over them when the block
+        # matches its supported configuration. The factory returns ``None``
+        # otherwise, leaving whichever narrower path is bound in charge.
+        self.cuda_infer_l_2_conv = None
+        if self.cuda_infer_level >= 2 and self.flash_attention_layout_supported:
+            from deepmd.pt_expt.kernels.cuda.dpa4 import (
+                make_cuda_so2_conv,
+            )
+
+            self.cuda_infer_l_2_conv = make_cuda_so2_conv(self)
 
-        # === Step 17. Optional fused CUDA SO(2) value path (training) ===
+        # === Step 15. Optional fused CUDA SO(2) value path (training) ===
         # One CUDA operator spans the training value stream up to the attention
         # aggregation: rotate-to-local, radial degree mixing, the cross-focus
         # competition weight, the whole gated mixing stack and the final
@@ -360,21 +359,22 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         # behind the same differentiable boundary. The attention span stays on
         # the Triton operator composition inside the traced graph. Bound under
         # ``DP_CUDA_TRAIN=1``; ``so2_message`` dispatches to it in training mode.
+        self.cuda_train_value = None
         if cuda_train_enabled():
             from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import (
                 make_cuda_so2_value,
             )
 
-            self._cuda_value_train = make_cuda_so2_value(self)
+            self.cuda_train_value = make_cuda_so2_value(self)
 
-        # === Step 18. Optional fused destination-segmented attention softmax ===
+        # === Step 16. Optional fused destination-segmented attention softmax ===
         # One CSR-segmented operator per direction replaces the
         # scatter/gather softmax chain of the attention weights, sharing the
         # destination-sorted view with the flash aggregation; its backward and
         # hand-derived second order keep the force-loss trace from expanding
         # the chain into materialized surfaces and serialized scatters. The
         # source-gated (SFPG) form and fp64 compute keep the reference path.
-        self._segment_softmax_fn = None
+        self.triton_l_1_segment_softmax = None
         if (
             max(self.triton_infer_level, self.triton_train_level) >= 1
             and self.compute_precision == "float32"
@@ -386,7 +386,18 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             )
 
             if SEGMENT_SOFTMAX_TRITON_AVAILABLE:
-                self._segment_softmax_fn = segment_softmax
+                self.triton_l_1_segment_softmax = segment_softmax
+
+        # === 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.cutile_infer_enabled:
+            from deepmd.pt_expt.kernels.cutile.sezm.so2_value_path import (
+                make_cutile_value_path,
+            )
+
+            self.cutile_infer_value = make_cutile_value_path(self)
 
     def _rotate_mix(
         self,
@@ -394,11 +405,11 @@ def _rotate_mix(
         edge_cache: EdgeCache,
         radial_feat: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        if self._triton_rotate_mix is not None and _active_triton_level(self) >= 1:
+        if self.triton_l_1_rotate_mix is not None and active_triton_level(self) >= 1:
             # The operator's backward reduces through the source CSR view,
             # built once per step and kept on the edge cache.
             cached_edge_csr(edge_cache, "src", x.shape[0])
-            u0, rad_feat = self._triton_rotate_mix(x, edge_cache, radial_feat)
+            u0, rad_feat = self.triton_l_1_rotate_mix(x, edge_cache, radial_feat)
             x_local = u0.view(
                 self.n_focus,
                 edge_cache.src.shape[0],
@@ -414,9 +425,9 @@ def _attention_softmax(
         edge_cache: EdgeCache,
         n_nodes: int,
     ) -> torch.Tensor:
-        active_level = _active_triton_level(self)
+        active_level = active_triton_level(self)
         if (
-            self._segment_softmax_fn is not None
+            self.triton_l_1_segment_softmax is not None
             and edge_cache.edge_src_gate is None
             and attn_logits.is_cuda
             and active_level >= 1
@@ -435,7 +446,7 @@ def _attention_softmax(
                 + float(self.eps)
             ).reshape(-1)  # (F * H,)
             n_channel = self.attn_n_focus * self.n_atten_head
-            alpha = self._segment_softmax_fn(
+            alpha = self.triton_l_1_segment_softmax(
                 attn_logits.reshape(n_edge, n_channel).to(dtype=torch.float32),
                 edge_cache.edge_env.reshape(n_edge).to(dtype=torch.float32),
                 null_logit,
@@ -448,38 +459,47 @@ def _attention_softmax(
             )
         return super()._attention_softmax(attn_logits, edge_cache, n_nodes)
 
-    def _rotation_active(self) -> bool:
-        """Whether the bound rotation kernels serve the current mode."""
-        return self._rotate_to_local_fn is not None and _active_triton_level(self) >= 1
+    def rotation_kernel_active(self) -> bool:
+        """Return whether the bound rotation kernels serve the current mode.
+
+        Returns
+        -------
+        bool
+            Whether level-one Triton rotation is active.
+        """
+        return (
+            self.triton_l_1_rotate_to_local is not None
+            and active_triton_level(self) >= 1
+        )
 
     def _rotate_to_local(
         self, x: torch.Tensor, edge_cache: EdgeCache
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        if self._rotation_active():
-            # ``self._rotate_to_local_fn`` was bound in ``__init__`` (the block
-            # kernel for the m-major ``mmax == 1`` layout, dense otherwise).
+        if self.rotation_kernel_active():
+            # The block kernel serves the m-major ``mmax == 1`` layout; the
+            # dense kernel serves every other supported layout.
             D_full = edge_cache.D_full
-            x_local = self._rotate_to_local_fn(x, edge_cache.src, D_full)
+            x_local = self.triton_l_1_rotate_to_local(x, edge_cache.src, D_full)
             x_dst_local: torch.Tensor | None = None
             if self.node_wise_grid_product is not None:
-                x_dst_local = self._rotate_to_local_fn(x, edge_cache.dst, D_full)
+                x_dst_local = self.triton_l_1_rotate_to_local(x, edge_cache.dst, D_full)
             return x_local, x_dst_local
         return super()._rotate_to_local(x, edge_cache)
 
     def _rotate_back(
         self, x_local: torch.Tensor, edge_cache: EdgeCache, n_edge: int
     ) -> torch.Tensor:
-        if self._rotation_active():
+        if self.rotation_kernel_active():
             Dt_full = edge_cache.Dt_full
             if self.mmax == 1:
                 # The block kernel consumes the (E, F, D_m, Cf) focus layout in
                 # place, folding the inverse transpose into its channel addressing.
-                return self._rotate_back_fn(x_local, Dt_full)
+                return self.triton_l_1_rotate_back(x_local, Dt_full)
             # Restore reduced global layout (E, D_m, C_wide) for the dense kernel.
             x_std = (
                 x_local.transpose(1, 2)
                 .contiguous()
                 .reshape(n_edge, self.reduced_dim, self.hidden_channels)
             )
-            return self._rotate_back_fn(x_std, Dt_full)
+            return self.triton_l_1_rotate_back(x_std, Dt_full)
         return super()._rotate_back(x_local, edge_cache, n_edge)
diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/wignerd.py b/deepmd/pt_expt/descriptor/dpa4_nn/wignerd.py
index aa6be46d21..64ea10c738 100644
--- a/deepmd/pt_expt/descriptor/dpa4_nn/wignerd.py
+++ b/deepmd/pt_expt/descriptor/dpa4_nn/wignerd.py
@@ -98,8 +98,8 @@ def __init__(
                     ]
             # The monomial basis routes through whichever accelerated backend
             # is selected; the two gates are mutually exclusive.
-            self._use_cutile_monomials = use_cutile_infer()
-            self._use_triton_monomials = triton_infer_level() >= 1
+            self.cutile_infer_monomials = use_cutile_infer()
+            self.triton_infer_l_1_monomials = triton_infer_level() >= 1
             # The l = 2 contraction tensor collapsed onto the 35 unique
             # degree-4 monomials: column m of the coefficient matrix sums
             # C_l2[:, :, p] over the 4^4 index tuples p whose component
@@ -185,9 +185,9 @@ def _monomial_matrix(
             exponents is not None
             and edge_quaternion.is_cuda
             and not self.training
-            and (self._use_triton_monomials or self._use_cutile_monomials)
+            and (self.triton_infer_l_1_monomials or self.cutile_infer_monomials)
         ):
-            if self._use_cutile_monomials:
+            if self.cutile_infer_monomials:
                 from deepmd.pt_expt.kernels.cutile.sezm.wigner_monomials import (
                     wigner_monomials as monomial_basis,
                 )
@@ -212,9 +212,9 @@ def _compute_l2_block(self, edge_quaternion: torch.Tensor) -> torch.Tensor:
             exponents is not None
             and edge_quaternion.is_cuda
             and not self.training
-            and (self._use_triton_monomials or self._use_cutile_monomials)
+            and (self.triton_infer_l_1_monomials or self.cutile_infer_monomials)
         ):
-            if self._use_cutile_monomials:
+            if self.cutile_infer_monomials:
                 from deepmd.pt_expt.kernels.cutile.sezm.wigner_monomials import (
                     wigner_monomials as monomial_basis,
                 )
diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py b/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py
index de407189ed..0b8b69c174 100644
--- a/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py
+++ b/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py
@@ -569,7 +569,7 @@ def __init__(self, conv: Any) -> None:
         self._conv = conv
         mixer = conv.radial_degree_mixer
         self._rank = 0 if mixer is None else int(mixer.rank)
-        self._compete = bool(conv.focus_compete and conv.n_focus > 1)
+        self.focus_compete = bool(conv.focus_compete and conv.n_focus > 1)
         # Packed-run polynomial tables, fitted once per degree and materialized
         # on the compute device at the first call.
         self._tables_cpu = wigner_run_tables(conv.lmax)
@@ -749,7 +749,7 @@ def __call__(
         conv = self._conv
         n_node = x.shape[0]
         kc, cb = self._degree_kernel(rad_feat)
-        if self._compete:
+        if self.focus_compete:
             fscale = self._focus_scale(x, edge_cache, rad_feat, kc, cb)  # (E, F)
         else:
             fscale = x.new_empty(0)
diff --git a/deepmd/pt_expt/kernels/cute/sezm/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/__init__.py
index 70f5cdca4b..56baf26081 100644
--- a/deepmd/pt_expt/kernels/cute/sezm/__init__.py
+++ b/deepmd/pt_expt/kernels/cute/sezm/__init__.py
@@ -1,54 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""
-CuTe-DSL fused SO(2) value-path operator for SeZM / DPA4.
-
-This package hosts a single bucketed CuTe operator that folds the entire per-edge
-value path of :class:`~deepmd.pt_expt.descriptor.dpa4_nn.so2.SO2Convolution`
-(``rotate_to_local`` -> radial degree mix -> the three-layer gated SO(2) mixing
-stack -> focus competition) into a fused forward kernel and a matching
-recompute backward kernel, keeping the per-edge intermediates on chip. It is an
-opt-in inference path enabled by ``DP_CUTE_INFER``; the final local features are
-handed to the committed flash-attention aggregation for rotate-back and scatter.
-Kernel entry points are internal implementation details of the SeZM descriptor;
-the package-level API only exposes availability and the value-path factory.
-
-Current limitations
--------------------
-Performance
-    On H20 / fp32 the operator is about 2.8x slower than the compiled Triton +
-    flash-attention path (roughly 489 / 724 ms versus 174 / 262 ms per force step
-    at 2000 / 4000 atoms). Peak memory is at parity with, or marginally below,
-    the compiled path (about 0.5 / 0.8 GB lower) and roughly 1.68x below the
-    eager path. The bottleneck is the recompute backward, which dominates the
-    kernel time: its occupancy is capped by the block-diagonal weight held
-    resident in shared memory, and both the forward and backward GEMMs run at the
-    hand-written plateau of about 21% of fp32 peak (versus about 52% for cuBLAS).
-
-Deployment
-    This is a Python-inference-only path. The ``cutlass.cute`` kernels are
-    nvcc / NVRTC JIT-compiled at runtime and do not bake into the AOTInductor
-    ``.pt2`` artifact, so the operator is unavailable to the LAMMPS / GPUMD C++
-    inference path. ``DP_CUTE_INFER`` and ``DP_TRITON_INFER`` both claim the
-    fused SO(2) value path and are mutually exclusive; enabling both is
-    rejected at construction.
-
-Correctness
-    The force is bit-exact against the eager reference (energy relative error
-    about 1e-9, force relative error about 5e-7 in fp32).
-"""
-
-from __future__ import (
-    annotations,
-)
-
-from .forward import (
-    SEZM_CUTE_AVAILABLE,
-)
-from .operator import (
-    make_cute_value_path,
-)
-
-__all__ = [
-    "SEZM_CUTE_AVAILABLE",
-    "make_cute_value_path",
-]
+"""CuTe kernels and PyTorch integration for supported SeZM layouts."""
diff --git a/deepmd/pt_expt/kernels/cute/sezm/backward.py b/deepmd/pt_expt/kernels/cute/sezm/backward.py
deleted file mode 100644
index cf1227e683..0000000000
--- a/deepmd/pt_expt/kernels/cute/sezm/backward.py
+++ /dev/null
@@ -1,692 +0,0 @@
-# SPDX-License-Identifier: LGPL-3.0-or-later
-# pyright: reportMissingImports=false
-# ruff: noqa: ANN001, ANN201, ANN202, ANN204, ANN205
-"""
-CuTe-DSL fused recompute-backward kernel for the SeZM SO(2) value path.
-
-Given the upstream gradients ``g_out`` (w.r.t. the pre-focus-compete local
-features ``x_local``) and ``g_fgate`` (w.r.t. the pre-mixing ``l = 0`` scalar),
-one bucketed kernel recomputes the forward value path from the small saved
-inputs (``x``, ``D_to_m``, ``Kc``) entirely on chip and backpropagates it,
-emitting the position-path gradients that carry the ``edge_vec`` -> force
-dependence::
-
-    grad_x       node-feature gradient, scattered to source nodes
-    grad_D_to_m  Wigner-rotation gradient, atomically summed over focus
-    grad_Kc      radial degree-kernel gradient, atomically summed over focus
-
-The weights are frozen on the inference force path. No ``E x D_m x C``
-intermediate reaches DRAM: the kernel recomputes the forward storing only the two
-gated-layer pre-activations in shared memory, then backpropagates the residual
-stack (gated-activation backward in registers/smem), the radial degree mix, and
-the rotation. This kernel is specialized to the three-layer ``[gated, gated,
-identity]`` mixing stack of the deployed configuration.
-
-Buffers per CTA (one focus of a bucket of ``B`` edges): four ``B x (D_m*Cf)``
-scratch tensors plus one ``max_block^2`` weight/gate scratch, inside the sm_90
-shared-memory limit at ``B = 16``. All accumulation is fp32 IEEE.
-"""
-
-from __future__ import (
-    annotations,
-)
-
-import torch
-
-from .forward import (
-    SEZM_CUTE_AVAILABLE,
-)
-
-if SEZM_CUTE_AVAILABLE:
-    import cuda.bindings.driver as cuda
-    import cutlass
-    import cutlass.cute as cute
-    import cutlass.cute.math as cmath
-    from cutlass.cute.runtime import (
-        from_dlpack,
-    )
-
-    from .forward import (
-        ForwardRunner,
-    )
-
-    class BackwardProgram:
-        """Bucketed fused recompute-backward program for the SO(2) value path.
-
-        Parameters
-        ----------
-        lmax, mmax, cf, n_focus, n_layers, bucket, threads, rb, rn
-            Kernel configuration (see :class:`.forward.ForwardProgram`).
-        """
-
-        def __init__(
-            self,
-            *,
-            lmax: int,
-            mmax: int,
-            cf: int,
-            n_focus: int,
-            n_layers: int,
-            bucket: int,
-            threads: int,
-            rb: int,
-            rn: int,
-        ) -> None:
-            self.lmax, self.mmax, self.cf, self.nf, self.nl = (
-                lmax,
-                mmax,
-                cf,
-                n_focus,
-                n_layers,
-            )
-            self._B, self._T, self._RB, self._RN = bucket, threads, rb, rn
-            self.D = (lmax + 1) ** 2
-            self.Dm = (lmax + 1) + sum(2 * (lmax - m + 1) for m in range(1, mmax + 1))
-            self.Cw = n_focus * cf
-            self.gate_out = lmax * cf
-            self.ngroup = 1 + 2 * mmax
-            self.FLAT = self.Dm * cf
-            groups = [lmax + 1] + [2 * (lmax - m + 1) for m in range(1, mmax + 1)]
-            self._blocks: list[tuple[int, int]] = []
-            off = 0
-            for g in groups:
-                self._blocks.append((off, g * cf))
-                off += g * cf
-            self._max_sb = max(sb for _, sb in self._blocks)
-            self._scr = max(
-                self._max_sb * self._max_sb,
-                cf * self.gate_out + 2 * bucket * self.gate_out,
-            )
-            assert self._B % rb == 0
-            for _, sb in self._blocks:
-                assert sb % rn == 0
-
-        @cute.jit
-        def __call__(
-            self,
-            mGout,
-            mGfg,
-            mX,
-            mSrc,
-            mDtoM,
-            mKc,
-            mCB,
-            mW,
-            mGW,
-            mExpand,
-            mGx,
-            mGD,
-            mGKc,
-            n_edge: cutlass.Int32,
-            n_bucket: cutlass.Int32,
-            stream: cuda.CUstream,
-        ):
-            self.kernel(
-                mGout,
-                mGfg,
-                mX,
-                mSrc,
-                mDtoM,
-                mKc,
-                mCB,
-                mW,
-                mGW,
-                mExpand,
-                mGx,
-                mGD,
-                mGKc,
-                n_edge,
-            ).launch(grid=[n_bucket, self.nf, 1], block=[self._T, 1, 1], stream=stream)
-
-        @cute.kernel
-        def kernel(
-            self,
-            mGout,
-            mGfg,
-            mX,
-            mSrc,
-            mDtoM,
-            mKc,
-            mCB,
-            mW,
-            mGW,
-            mExpand,
-            mGx,
-            mGD,
-            mGKc,
-            n_edge: cutlass.Int32,
-        ):
-            D = cutlass.const_expr(self.D)
-            Dm = cutlass.const_expr(self.Dm)
-            CF = cutlass.const_expr(self.cf)
-            FLAT = cutlass.const_expr(self.FLAT)
-            B = cutlass.const_expr(self._B)
-            T = cutlass.const_expr(self._T)
-
-            tidx, _, _ = cute.arch.thread_idx()
-            bucket, focus, _ = cute.arch.block_idx()
-            e0 = bucket * B
-
-            smem = cutlass.utils.SmemAllocator()
-            b0 = smem.allocate_tensor(
-                cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16
-            )
-            b1 = smem.allocate_tensor(
-                cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16
-            )
-            b2 = smem.allocate_tensor(
-                cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16
-            )
-            b3 = smem.allocate_tensor(
-                cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16
-            )
-            s_w = smem.allocate_tensor(
-                cutlass.Float32,
-                cute.make_layout((cutlass.const_expr(self._scr),), stride=(1,)),
-                16,
-            )
-
-            # === Step 1. Forward recompute: store z_0 -> b1, z_1 -> b2 ===
-            # b0 carries the running layer input h_l; b3 is the activation temp.
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                dm = rem // CF
-                c = rem % CF
-                e = e0 + b
-                eq = e if e < n_edge else n_edge - 1
-                src = mSrc[eq]
-                acc = cutlass.Float32(0.0)
-                for k in cutlass.range_constexpr(D):
-                    acc += mDtoM[eq, dm, k] * mX[src, k, focus * CF + c]
-                b1[b, rem] = acc
-                i += T
-            cute.arch.sync_threads()
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                o = rem // CF
-                c = rem % CF
-                e = e0 + b
-                eq = e if e < n_edge else n_edge - 1
-                acc = cutlass.Float32(0.0)
-                for ii in cutlass.range_constexpr(Dm):
-                    acc += mKc[eq, o, ii] * b1[b, ii * CF + c]
-                b0[b, rem] = acc * mCB[focus * CF + c]
-                i += T
-            cute.arch.sync_threads()
-            for lyr in cutlass.range_constexpr(self.nl - 1):
-                zbuf = b1 if lyr == 0 else b2
-                self._gemm_fwd(b0, zbuf, s_w, mW, lyr, focus, tidx)
-                self._gated_fwd(zbuf, b3, s_w, mGW, mExpand, lyr, focus, tidx)
-                i = tidx
-                while i < B * FLAT:
-                    b = i // FLAT
-                    rem = i % FLAT
-                    b0[b, rem] = b0[b, rem] + b3[b, rem]
-                    i += T
-                cute.arch.sync_threads()
-
-            # === Step 2. Reverse the residual stack; b0 accumulates grad_h ===
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                o = rem // CF
-                c = rem % CF
-                b0[b, rem] = mGout[e0 + b, focus, o, c]
-                i += T
-            cute.arch.sync_threads()
-            # layer 2 (identity): grad_z = grad_out; grad_h += W_2^T @ grad_z
-            self._gemm_bwd(b0, b3, s_w, mW, self.nl - 1, focus, tidx)
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                b0[b, rem] = b0[b, rem] + b3[b, rem]
-                i += T
-            cute.arch.sync_threads()
-            # layer 1 (gated): grad_z_1 -> b2 in place; grad_h += W_1^T @ grad_z_1
-            self._gated_bwd(b2, b0, s_w, mGW, mExpand, 1, focus, tidx)
-            self._gemm_bwd(b2, b3, s_w, mW, 1, focus, tidx)
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                b0[b, rem] = b0[b, rem] + b3[b, rem]
-                i += T
-            cute.arch.sync_threads()
-            # layer 0 (gated): grad_z_0 -> b1 in place; grad_h += W_0^T @ grad_z_0
-            self._gated_bwd(b1, b0, s_w, mGW, mExpand, 0, focus, tidx)
-            self._gemm_bwd(b1, b3, s_w, mW, 0, focus, tidx)
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                b0[b, rem] = b0[b, rem] + b3[b, rem]
-                i += T
-            cute.arch.sync_threads()
-            # add the focus-competition gradient into the l=0 row
-            i = tidx
-            while i < B * CF:
-                b = i // CF
-                c = i % CF
-                b0[b, c] = b0[b, c] + mGfg[e0 + b, focus, c]
-                i += T
-            cute.arch.sync_threads()
-
-            # === Step 3. Radial + rotation backward; b0 holds grad_h0 ===
-            # recompute x_rot -> b1
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                dm = rem // CF
-                c = rem % CF
-                e = e0 + b
-                eq = e if e < n_edge else n_edge - 1
-                src = mSrc[eq]
-                acc = cutlass.Float32(0.0)
-                for k in cutlass.range_constexpr(D):
-                    acc += mDtoM[eq, dm, k] * mX[src, k, focus * CF + c]
-                b1[b, rem] = acc
-                i += T
-            cute.arch.sync_threads()
-            # grad_Kc[o, ii] = sum_c channel_basis[c] * grad_h0[o, c] * x_rot[ii, c]
-            i = tidx
-            while i < B * Dm * Dm:
-                b = i // (Dm * Dm)
-                rem = i % (Dm * Dm)
-                o = rem // Dm
-                ii = rem % Dm
-                e = e0 + b
-                acc = cutlass.Float32(0.0)
-                for c in cutlass.range_constexpr(CF):
-                    acc += mCB[focus * CF + c] * b0[b, o * CF + c] * b1[b, ii * CF + c]
-                cute.arch.atomic_add(mGKc.iterator + mGKc.layout((e, o, ii)), acc)
-                i += T
-            # grad_x_rot[ii, c] = channel_basis[c] * sum_o Kc[o, ii] * grad_h0[o, c] -> b2
-            i = tidx
-            while i < B * Dm * CF:
-                b = i // (Dm * CF)
-                rem = i % (Dm * CF)
-                ii = rem // CF
-                c = rem % CF
-                e = e0 + b
-                eq = e if e < n_edge else n_edge - 1
-                acc = cutlass.Float32(0.0)
-                for o in cutlass.range_constexpr(Dm):
-                    acc += mKc[eq, o, ii] * b0[b, o * CF + c]
-                b2[b, ii * CF + c] = acc * mCB[focus * CF + c]
-                i += T
-            cute.arch.sync_threads()
-            # grad_x[src, k, c] += sum_ii D_to_m[ii, k] * grad_x_rot[ii, c]
-            i = tidx
-            while i < B * D * CF:
-                b = i // (D * CF)
-                rem = i % (D * CF)
-                k = rem // CF
-                c = rem % CF
-                e = e0 + b
-                eq = e if e < n_edge else n_edge - 1
-                src = mSrc[eq]
-                acc = cutlass.Float32(0.0)
-                for ii in cutlass.range_constexpr(Dm):
-                    acc += mDtoM[eq, ii, k] * b2[b, ii * CF + c]
-                cute.arch.atomic_add(
-                    mGx.iterator + mGx.layout((src, k, focus * CF + c)), acc
-                )
-                i += T
-            # grad_D_to_m[ii, k] = sum_c grad_x_rot[ii, c] * x_src[k, c]
-            i = tidx
-            while i < B * Dm * D:
-                b = i // (Dm * D)
-                rem = i % (Dm * D)
-                ii = rem // D
-                k = rem % D
-                e = e0 + b
-                eq = e if e < n_edge else n_edge - 1
-                src = mSrc[eq]
-                acc = cutlass.Float32(0.0)
-                for c in cutlass.range_constexpr(CF):
-                    acc += b2[b, ii * CF + c] * mX[src, k, focus * CF + c]
-                cute.arch.atomic_add(mGD.iterator + mGD.layout((e, ii, k)), acc)
-                i += T
-
-        @cute.jit
-        def _gemm_fwd(self, hbuf, zbuf, s_w, mW, lyr, focus, tidx):
-            """Recompute ``zbuf = hbuf @ W[lyr, focus]`` (block-diagonal)."""
-            T = cutlass.const_expr(self._T)
-            RB = cutlass.const_expr(self._RB)
-            RN = cutlass.const_expr(self._RN)
-            B = cutlass.const_expr(self._B)
-            for ob, sb in cutlass.const_expr(self._blocks):
-                j = tidx
-                while j < sb * sb:
-                    s_w[(j // sb) * sb + (j % sb)] = mW[
-                        lyr, focus, ob + (j // sb), ob + (j % sb)
-                    ]
-                    j += T
-                cute.arch.sync_threads()
-                n_mt = cutlass.const_expr((B // RB) * (sb // RN))
-                racc = cute.make_rmem_tensor(
-                    cute.make_layout((RB, RN)), cutlass.Float32
-                )
-                mt = tidx
-                while mt < n_mt:
-                    bi = (mt // (sb // RN)) * RB
-                    nj = (mt % (sb // RN)) * RN
-                    for r in cutlass.range_constexpr(RB):
-                        for s in cutlass.range_constexpr(RN):
-                            racc[r, s] = cutlass.Float32(0.0)
-                    for k in cutlass.range(sb):
-                        for r in cutlass.range_constexpr(RB):
-                            a_rk = hbuf[bi + r, ob + k]
-                            for s in cutlass.range_constexpr(RN):
-                                racc[r, s] += a_rk * s_w[k * sb + nj + s]
-                    for r in cutlass.range_constexpr(RB):
-                        for s in cutlass.range_constexpr(RN):
-                            zbuf[bi + r, ob + nj + s] = racc[r, s]
-                    mt += T
-                cute.arch.sync_threads()
-
-        @cute.jit
-        def _gemm_bwd(self, gzbuf, ghbuf, s_w, mW, lyr, focus, tidx):
-            """Compute ``ghbuf = W[lyr, focus]^T @ gzbuf`` (block-diagonal)."""
-            T = cutlass.const_expr(self._T)
-            RB = cutlass.const_expr(self._RB)
-            RN = cutlass.const_expr(self._RN)
-            B = cutlass.const_expr(self._B)
-            for ob, sb in cutlass.const_expr(self._blocks):
-                j = tidx
-                while j < sb * sb:
-                    s_w[(j // sb) * sb + (j % sb)] = mW[
-                        lyr, focus, ob + (j // sb), ob + (j % sb)
-                    ]
-                    j += T
-                cute.arch.sync_threads()
-                n_mt = cutlass.const_expr((B // RB) * (sb // RN))
-                racc = cute.make_rmem_tensor(
-                    cute.make_layout((RB, RN)), cutlass.Float32
-                )
-                mt = tidx
-                while mt < n_mt:
-                    bi = (mt // (sb // RN)) * RB
-                    kj = (mt % (sb // RN)) * RN  # W in-index (grad_h output column)
-                    for r in cutlass.range_constexpr(RB):
-                        for s in cutlass.range_constexpr(RN):
-                            racc[r, s] = cutlass.Float32(0.0)
-                    for n in cutlass.range(sb):  # sum over the W out-index
-                        for r in cutlass.range_constexpr(RB):
-                            gz_rn = gzbuf[bi + r, ob + n]
-                            for s in cutlass.range_constexpr(RN):
-                                racc[r, s] += gz_rn * s_w[(kj + s) * sb + n]
-                    for r in cutlass.range_constexpr(RB):
-                        for s in cutlass.range_constexpr(RN):
-                            ghbuf[bi + r, ob + kj + s] = racc[r, s]
-                    mt += T
-                cute.arch.sync_threads()
-
-        @cute.jit
-        def _gated_fwd(self, zbuf, abuf, s_w, mGW, mExpand, lyr, focus, tidx):
-            """Recompute ``abuf = GatedActivation(zbuf)`` (silu l=0, gate l>0)."""
-            CF = cutlass.const_expr(self.cf)
-            GO = cutlass.const_expr(self.gate_out)
-            Dm = cutlass.const_expr(self.Dm)
-            B = cutlass.const_expr(self._B)
-            T = cutlass.const_expr(self._T)
-            SIG_OFF = cutlass.const_expr(self.cf * self.gate_out)
-            j = tidx
-            while j < CF * GO:
-                s_w[(j // GO) * GO + (j % GO)] = mGW[lyr, focus, j // GO, j % GO]
-                j += T
-            cute.arch.sync_threads()
-            j = tidx
-            while j < B * GO:
-                b = j // GO
-                o = j % GO
-                acc = cutlass.Float32(0.0)
-                for ii in cutlass.range_constexpr(CF):
-                    acc += zbuf[b, ii] * s_w[ii * GO + o]
-                s_w[SIG_OFF + b * GO + o] = cutlass.Float32(1.0) / (
-                    cutlass.Float32(1.0) + cmath.exp(-acc)
-                )
-                j += T
-            cute.arch.sync_threads()
-            j = tidx
-            while j < B * CF:
-                b = j // CF
-                c = j % CF
-                z = zbuf[b, c]
-                abuf[b, c] = z / (cutlass.Float32(1.0) + cmath.exp(-z))
-                j += T
-            j = tidx
-            while j < B * (Dm - 1) * CF:
-                b = j // ((Dm - 1) * CF)
-                rem = j % ((Dm - 1) * CF)
-                d1 = rem // CF
-                c = rem % CF
-                lidx = mExpand[d1]
-                abuf[b, (d1 + 1) * CF + c] = (
-                    zbuf[b, (d1 + 1) * CF + c] * s_w[SIG_OFF + b * GO + lidx * CF + c]
-                )
-                j += T
-            cute.arch.sync_threads()
-
-        @cute.jit
-        def _gated_bwd(self, zbuf, gabuf, s_w, mGW, mExpand, lyr, focus, tidx):
-            """Backprop the gated activation in place: ``zbuf`` (z_l) -> grad_z_l.
-
-            With ``g_a = gabuf`` the incoming gradient and the recomputed gate
-            sigmoids ``sig``::
-
-                grad_z[dm, c] = g_a[dm, c] * sig[expand[dm-1], c]              (dm > 0)
-                g_sig[L, c]   = sum_{dm: expand[dm-1]=L} g_a[dm, c] * z[dm, c]
-                grad_z[0, i]  = g_a[0, i] * silu'(z[0, i])
-                                + sum_{o'} Wg[i, o'] * g_sig * sig*(1-sig)
-            """
-            CF = cutlass.const_expr(self.cf)
-            GO = cutlass.const_expr(self.gate_out)
-            Dm = cutlass.const_expr(self.Dm)
-            LMAX = cutlass.const_expr(self.lmax)
-            NG = cutlass.const_expr(self.ngroup)
-            B = cutlass.const_expr(self._B)
-            T = cutlass.const_expr(self._T)
-            SIG_OFF = cutlass.const_expr(self.cf * self.gate_out)
-            GGL_OFF = cutlass.const_expr(self.cf * self.gate_out + B * self.gate_out)
-            j = tidx
-            while j < CF * GO:
-                s_w[(j // GO) * GO + (j % GO)] = mGW[lyr, focus, j // GO, j % GO]
-                j += T
-            cute.arch.sync_threads()
-            j = tidx
-            while j < B * GO:
-                b = j // GO
-                o = j % GO
-                acc = cutlass.Float32(0.0)
-                for ii in cutlass.range_constexpr(CF):
-                    acc += zbuf[b, ii] * s_w[ii * GO + o]
-                s_w[SIG_OFF + b * GO + o] = cutlass.Float32(1.0) / (
-                    cutlass.Float32(1.0) + cmath.exp(-acc)
-                )
-                j += T
-            cute.arch.sync_threads()
-            # g_gl[L, c] = (sum over the (1 + 2*mmax) |m| groups) * sigmoid'(gate).
-            # For degree l = L + 1 the contributing coefficients are dm = 1 + L + k*lmax.
-            j = tidx
-            while j < B * GO:
-                b = j // GO
-                o = j % GO
-                gate_l = o // CF
-                c = o % CF
-                gsig = cutlass.Float32(0.0)
-                for kk in cutlass.range_constexpr(NG):
-                    dm = 1 + gate_l + kk * LMAX
-                    gsig += gabuf[b, dm * CF + c] * zbuf[b, dm * CF + c]
-                s = s_w[SIG_OFF + b * GO + o]
-                s_w[GGL_OFF + b * GO + o] = gsig * s * (cutlass.Float32(1.0) - s)
-                j += T
-            cute.arch.sync_threads()
-            # grad_z[dm>0]: reads sig, writes disjoint from the l=0 slice.
-            j = tidx
-            while j < B * (Dm - 1) * CF:
-                b = j // ((Dm - 1) * CF)
-                rem = j % ((Dm - 1) * CF)
-                d1 = rem // CF
-                c = rem % CF
-                lidx = mExpand[d1]
-                zbuf[b, (d1 + 1) * CF + c] = (
-                    gabuf[b, (d1 + 1) * CF + c] * s_w[SIG_OFF + b * GO + lidx * CF + c]
-                )
-                j += T
-            # grad_z[0]: reads z[0, i] before overwriting it.
-            j = tidx
-            while j < B * CF:
-                b = j // CF
-                ii = j % CF
-                z0 = zbuf[b, ii]
-                sg = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cmath.exp(-z0))
-                silup = sg + z0 * sg * (cutlass.Float32(1.0) - sg)
-                acc = gabuf[b, ii] * silup
-                for o in cutlass.range_constexpr(GO):
-                    acc += s_w[ii * GO + o] * s_w[GGL_OFF + b * GO + o]
-                zbuf[b, ii] = acc
-                j += T
-            cute.arch.sync_threads()
-
-    class BackwardRunner:
-        """Compile-once driver for :class:`BackwardProgram` (force-path gradients).
-
-        Parameters
-        ----------
-        weights
-            Packed weights (see :class:`.forward.ForwardRunner`).
-        lmax, mmax, cf, n_focus, n_layers, bucket, threads, rb, rn
-            Kernel configuration.
-        """
-
-        def __init__(
-            self,
-            weights,
-            *,
-            lmax: int,
-            mmax: int,
-            cf: int,
-            n_focus: int,
-            n_layers: int,
-            bucket: int,
-            threads: int,
-            rb: int,
-            rn: int,
-        ) -> None:
-            self.op = BackwardProgram(
-                lmax=lmax,
-                mmax=mmax,
-                cf=cf,
-                n_focus=n_focus,
-                n_layers=n_layers,
-                bucket=bucket,
-                threads=threads,
-                rb=rb,
-                rn=rn,
-            )
-            self._B = bucket
-            self.nf, self.cf, self.Dm, self.D = n_focus, cf, self.op.Dm, self.op.D
-            self._compiled = None
-            fr = ForwardRunner(
-                weights,
-                lmax=lmax,
-                mmax=mmax,
-                cf=cf,
-                n_focus=n_focus,
-                n_layers=n_layers,
-                bucket=bucket,
-                threads=threads,
-                rb=rb,
-                rn=rn,
-            )
-            self.m_w, self.m_gw, self.m_cb, self.m_expand = (
-                fr.m_w,
-                fr.m_gw,
-                fr.m_cb,
-                fr.m_expand,
-            )
-
-        @staticmethod
-        def _dyn(t: torch.Tensor, leading: int):
-            return from_dlpack(t, assumed_align=16).mark_layout_dynamic(
-                leading_dim=leading
-            )
-
-        def __call__(
-            self,
-            x: torch.Tensor,
-            src: torch.Tensor,
-            d_to_m: torch.Tensor,
-            kc: torch.Tensor,
-            g_out: torch.Tensor,
-            g_fgate: torch.Tensor,
-        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            """Emit the value-path gradients.
-
-            Parameters
-            ----------
-            x, src, d_to_m, kc
-                Forward inputs (see :meth:`.forward.ForwardRunner.__call__`).
-            g_out : torch.Tensor
-                Gradient w.r.t. ``x_local`` with shape (E, F, D_m, Cf).
-            g_fgate : torch.Tensor
-                Gradient w.r.t. ``focus_gate`` with shape (E, F, Cf).
-
-            Returns
-            -------
-            grad_x : torch.Tensor
-                Node-feature gradient with shape (N, D, C_wide).
-            grad_d_to_m : torch.Tensor
-                Wigner-rotation gradient with shape (E, D_m, D).
-            grad_kc : torch.Tensor
-                Radial degree-kernel gradient with shape (E, D_m, D_m).
-            """
-            n_edge = src.shape[0]
-            n_bucket = (n_edge + self._B - 1) // self._B
-            n_pad = n_bucket * self._B
-            s32 = src.to(torch.int32)
-            g_out_p, g_fg_p = g_out, g_fgate
-            if n_pad > n_edge:
-                s32 = torch.cat([s32, s32.new_zeros(n_pad - n_edge)])
-                g_out_p = torch.cat(
-                    [g_out, g_out.new_zeros(n_pad - n_edge, *g_out.shape[1:])]
-                )
-                g_fg_p = torch.cat(
-                    [g_fgate, g_fgate.new_zeros(n_pad - n_edge, *g_fgate.shape[1:])]
-                )
-            grad_x = torch.zeros_like(x)
-            grad_d = torch.zeros(
-                n_pad, self.Dm, self.D, device=x.device, dtype=torch.float32
-            )
-            grad_kc = torch.zeros(
-                n_pad, self.Dm, self.Dm, device=x.device, dtype=torch.float32
-            )
-            views = (
-                self._dyn(g_out_p, 3),
-                self._dyn(g_fg_p, 2),
-                self._dyn(x, 2),
-                self._dyn(s32, 0),
-                self._dyn(d_to_m, 2),
-                self._dyn(kc, 2),
-                from_dlpack(self.m_cb, assumed_align=16),
-                from_dlpack(self.m_w, assumed_align=16),
-                from_dlpack(self.m_gw, assumed_align=16),
-                from_dlpack(self.m_expand, assumed_align=16),
-                self._dyn(grad_x, 2),
-                self._dyn(grad_d, 2),
-                self._dyn(grad_kc, 2),
-            )
-            args = (*views, cutlass.Int32(n_edge), cutlass.Int32(n_bucket))
-            stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream)
-            if self._compiled is None:
-                self._compiled = cute.compile(self.op, *args, stream=stream)
-            self._compiled(*args, stream=stream)
-            return grad_x, grad_d[:n_edge], grad_kc[:n_edge]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/compile_cache.py b/deepmd/pt_expt/kernels/cute/sezm/compile_cache.py
new file mode 100644
index 0000000000..0a9b9b89b5
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/compile_cache.py
@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Device-aware caching for architecture-specific CuTe compilation."""
+
+from __future__ import (
+    annotations,
+)
+
+from collections.abc import (
+    Callable,
+)
+from contextlib import (
+    nullcontext,
+)
+from functools import (
+    lru_cache,
+    wraps,
+)
+from typing import (
+    Any,
+    TypeVar,
+    cast,
+)
+
+_T = TypeVar("_T", bound=Callable[..., Any])
+
+
+def current_cuda_compile_identity() -> tuple[int, int, int]:
+    """Return the current CUDA device and its compute capability."""
+    import torch
+
+    device_index = torch.cuda.current_device()
+    major, minor = torch.cuda.get_device_capability(device_index)
+    return device_index, major, minor
+
+
+def device_aware_lru_cache(
+    *,
+    maxsize: int,
+    identity_getter: Callable[[], tuple[int, int, int]] = current_cuda_compile_identity,
+) -> Callable[[_T], _T]:
+    """Cache a compile factory separately for each CUDA device architecture."""
+
+    def decorate(function: _T) -> _T:
+        @lru_cache(maxsize=maxsize)
+        def cached(
+            identity: tuple[int, int, int],
+            args: tuple[Any, ...],
+            kwargs: tuple[tuple[str, Any], ...],
+        ) -> Any:
+            import torch
+
+            device_index = identity[0]
+            device_count = getattr(torch.cuda, "device_count", None)
+            device_is_visible = device_count is None or device_index < device_count()
+            compile_device = (
+                torch.cuda.device(device_index)
+                if torch.cuda.is_available() and device_is_visible
+                else nullcontext()
+            )
+            with compile_device:
+                return function(*args, **dict(kwargs))
+
+        @wraps(function)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return cached(
+                identity_getter(),
+                args,
+                tuple(sorted(kwargs.items())),
+            )
+
+        wrapper.cache_clear = cached.cache_clear
+        wrapper.cache_info = cached.cache_info
+        wrapper.cache_parameters = cached.cache_parameters
+        wrapper._deepmd_cute_cached = True
+        return cast("_T", wrapper)
+
+    return decorate
diff --git a/deepmd/pt_expt/kernels/cute/sezm/forward.py b/deepmd/pt_expt/kernels/cute/sezm/forward.py
deleted file mode 100644
index 0d6ca5ec8a..0000000000
--- a/deepmd/pt_expt/kernels/cute/sezm/forward.py
+++ /dev/null
@@ -1,444 +0,0 @@
-# SPDX-License-Identifier: LGPL-3.0-or-later
-# pyright: reportMissingImports=false
-# ruff: noqa: ANN001, ANN201, ANN204, ANN205
-"""
-CuTe-DSL fused forward kernel for the SeZM SO(2) value path.
-
-One bucketed kernel folds the entire per-edge value path of ``SO2Convolution``
-into a single launch::
-
-    gather x[src] -> rotate_to_local (D_to_m)                 [prologue, in smem]
-      -> radial degree mix (Kc, channel_basis)                [prologue, in smem]
-      -> 3x (block-diagonal SO2Linear + GatedActivation + residual)   [in smem]
-      -> x_local (E, F, D_m, Cf)                              [+ pre-mixing l=0 scalar]
-
-The block-diagonal ``SO2Linear`` weight is staged in shared memory once per
-bucket and reused across the bucket's ``B`` edges (register-blocked ``RB x RN``
-FMA micro-tile), so the ``E x D_m x C`` intermediates of all three mixing layers
-stay resident on chip and never reach DRAM. The focus competition (a per-edge
-softmax of the pre-mixing ``l = 0`` feature) is applied outside the kernel from
-the returned ``focus_gate`` scalar.
-
-Grid ``(n_bucket, n_focus)``; one CTA owns ``B`` edges of one focus stream.
-All accumulation is fp32 IEEE (no TF32) to keep the potential-energy surface
-smooth.
-"""
-
-from __future__ import (
-    annotations,
-)
-
-import torch
-
-try:
-    import cuda.bindings.driver as cuda
-    import cutlass
-    import cutlass.cute as cute
-    import cutlass.cute.math as cmath
-    from cutlass.cute.runtime import (
-        from_dlpack,
-    )
-
-    SEZM_CUTE_AVAILABLE = True
-except Exception:  # pragma: no cover - import guard for non-CuTe environments
-    SEZM_CUTE_AVAILABLE = False
-
-
-if SEZM_CUTE_AVAILABLE:
-
-    class ForwardProgram:
-        """Bucketed fused forward program for the SO(2) value path.
-
-        Parameters
-        ----------
-        lmax : int
-            Maximum spherical harmonic degree.
-        mmax : int
-            Maximum SO(2) order retained in the reduced layout.
-        cf : int
-            Per-focus channel width ``Cf``.
-        n_focus : int
-            Number of focus streams ``F``.
-        n_layers : int
-            Number of SO(2) mixing layers.
-        bucket : int
-            Edges processed per CTA ``B``.
-        threads : int
-            Threads per CTA.
-        rb, rn : int
-            Register micro-tile dimensions (``RB`` bucket rows, ``RN`` output
-            columns per thread) of the block-diagonal GEMM.
-        """
-
-        def __init__(
-            self,
-            *,
-            lmax: int,
-            mmax: int,
-            cf: int,
-            n_focus: int,
-            n_layers: int,
-            bucket: int,
-            threads: int,
-            rb: int,
-            rn: int,
-        ) -> None:
-            self.lmax, self.mmax, self.cf, self.nf, self.nl = (
-                lmax,
-                mmax,
-                cf,
-                n_focus,
-                n_layers,
-            )
-            self._B, self._T, self._RB, self._RN = bucket, threads, rb, rn
-            self.D = (lmax + 1) ** 2
-            self.Dm = (lmax + 1) + sum(2 * (lmax - m + 1) for m in range(1, mmax + 1))
-            self.Cw = n_focus * cf
-            self.gate_out = lmax * cf
-            self.FLAT = self.Dm * cf
-            # Block-diagonal |m| block widths in the flattened coeff*channel axis:
-            # m = 0 spans (lmax + 1) coefficients, each |m| > 0 spans 2*(lmax-m+1).
-            groups = [lmax + 1] + [2 * (lmax - m + 1) for m in range(1, mmax + 1)]
-            self._blocks: list[tuple[int, int]] = []
-            off = 0
-            for g in groups:
-                self._blocks.append((off, g * cf))
-                off += g * cf
-            self._max_sb = max(sb for _, sb in self._blocks)
-            # Shared scratch reused for the resident weight block and, during the
-            # gated activation, the gate weight plus per-edge sigmoid buffer.
-            self._scr = max(
-                self._max_sb * self._max_sb, cf * self.gate_out + bucket * self.gate_out
-            )
-            assert self._B % rb == 0
-            for _, sb in self._blocks:
-                assert sb % rn == 0
-
-        @cute.jit
-        def __call__(
-            self,
-            mX,
-            mSrc,
-            mDtoM,
-            mKc,
-            mCB,
-            mW,
-            mGW,
-            mExpand,
-            mOut,
-            mFocusGate,
-            n_edge: cutlass.Int32,
-            n_bucket: cutlass.Int32,
-            stream: cuda.CUstream,
-        ):
-            self.kernel(
-                mX, mSrc, mDtoM, mKc, mCB, mW, mGW, mExpand, mOut, mFocusGate, n_edge
-            ).launch(grid=[n_bucket, self.nf, 1], block=[self._T, 1, 1], stream=stream)
-
-        @cute.kernel
-        def kernel(
-            self,
-            mX,
-            mSrc,
-            mDtoM,
-            mKc,
-            mCB,
-            mW,
-            mGW,
-            mExpand,
-            mOut,
-            mFocusGate,
-            n_edge: cutlass.Int32,
-        ):
-            D = cutlass.const_expr(self.D)
-            Dm = cutlass.const_expr(self.Dm)
-            CF = cutlass.const_expr(self.cf)
-            FLAT = cutlass.const_expr(self.FLAT)
-            GO = cutlass.const_expr(self.gate_out)
-            B = cutlass.const_expr(self._B)
-            T = cutlass.const_expr(self._T)
-            RB = cutlass.const_expr(self._RB)
-            RN = cutlass.const_expr(self._RN)
-            SCR = cutlass.const_expr(self._scr)
-            GATE_OFF = cutlass.const_expr(self.cf * self.gate_out)
-            NGATE = cutlass.const_expr(self.nl - 1)
-
-            tidx, _, _ = cute.arch.thread_idx()
-            bucket, focus, _ = cute.arch.block_idx()
-            e0 = bucket * B
-
-            smem = cutlass.utils.SmemAllocator()
-            s_cur = smem.allocate_tensor(
-                cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16
-            )
-            s_tmp = smem.allocate_tensor(
-                cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16
-            )
-            s_scr = smem.allocate_tensor(
-                cutlass.Float32, cute.make_layout((SCR,), stride=(1,)), 16
-            )
-
-            # === Step 1. rotate_to_local: s_tmp[b, dm*Cf+c] = sum_k D_to_m x[src] ===
-            # Padding edges (e >= n_edge) clamp their read index; their output rows
-            # are sliced off by the caller.
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                dm = rem // CF
-                c = rem % CF
-                e = e0 + b
-                eq = e if e < n_edge else n_edge - 1
-                src = mSrc[eq]
-                acc = cutlass.Float32(0.0)
-                for k in cutlass.range_constexpr(D):
-                    acc += mDtoM[eq, dm, k] * mX[src, k, focus * CF + c]
-                s_tmp[b, rem] = acc
-                i += T
-            cute.arch.sync_threads()
-
-            # === Step 2. radial degree mix: s_cur = channel_basis * (Kc @ x_rot) ===
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                o = rem // CF
-                c = rem % CF
-                e = e0 + b
-                eq = e if e < n_edge else n_edge - 1
-                acc = cutlass.Float32(0.0)
-                for ii in cutlass.range_constexpr(Dm):
-                    acc += mKc[eq, o, ii] * s_tmp[b, ii * CF + c]
-                s_cur[b, rem] = acc * mCB[focus * CF + c]
-                i += T
-            cute.arch.sync_threads()
-
-            # === Step 3. Emit the pre-mixing l=0 scalar for the focus competition ===
-            i = tidx
-            while i < B * CF:
-                b = i // CF
-                c = i % CF
-                mFocusGate[e0 + b, focus, c] = s_cur[b, c]
-                i += T
-            cute.arch.sync_threads()
-
-            # === Step 4. Multi-layer gated SO(2) mixing (block-diagonal, residual) ===
-            for lyr in cutlass.range_constexpr(self.nl):
-                # --- SO2Linear: s_tmp = s_cur @ W[lyr, focus] (per |m| block) ---
-                for ob, sb in cutlass.const_expr(self._blocks):
-                    j = tidx
-                    while j < sb * sb:
-                        k = j // sb
-                        n = j % sb
-                        s_scr[k * sb + n] = mW[lyr, focus, ob + k, ob + n]
-                        j += T
-                    cute.arch.sync_threads()
-                    n_mt = cutlass.const_expr((B // RB) * (sb // RN))
-                    racc = cute.make_rmem_tensor(
-                        cute.make_layout((RB, RN)), cutlass.Float32
-                    )
-                    mt = tidx
-                    while mt < n_mt:
-                        bi = (mt // (sb // RN)) * RB
-                        nj = (mt % (sb // RN)) * RN
-                        for r in cutlass.range_constexpr(RB):
-                            for s in cutlass.range_constexpr(RN):
-                                racc[r, s] = cutlass.Float32(0.0)
-                        for k in cutlass.range(sb):
-                            for r in cutlass.range_constexpr(RB):
-                                a_rk = s_cur[bi + r, ob + k]
-                                for s in cutlass.range_constexpr(RN):
-                                    racc[r, s] += a_rk * s_scr[k * sb + nj + s]
-                        for r in cutlass.range_constexpr(RB):
-                            for s in cutlass.range_constexpr(RN):
-                                s_tmp[bi + r, ob + nj + s] = racc[r, s]
-                        mt += T
-                    cute.arch.sync_threads()
-
-                # --- GatedActivation (gated layers) or identity (last layer) ---
-                if lyr < NGATE:
-                    # gate FocusLinear weight -> s_scr[0 : Cf*GO]
-                    j = tidx
-                    while j < CF * GO:
-                        ii = j // GO
-                        o = j % GO
-                        s_scr[ii * GO + o] = mGW[lyr, focus, ii, o]
-                        j += T
-                    cute.arch.sync_threads()
-                    # gate sigmoids from the l=0 scalar -> s_scr[GATE_OFF + b*GO + o]
-                    j = tidx
-                    while j < B * GO:
-                        b = j // GO
-                        o = j % GO
-                        acc = cutlass.Float32(0.0)
-                        for ii in cutlass.range_constexpr(CF):
-                            acc += s_tmp[b, ii] * s_scr[ii * GO + o]
-                        s_scr[GATE_OFF + b * GO + o] = cutlass.Float32(1.0) / (
-                            cutlass.Float32(1.0) + cmath.exp(-acc)
-                        )
-                        j += T
-                    cute.arch.sync_threads()
-                    # l=0: silu(z) = z / (1 + exp(-z))
-                    j = tidx
-                    while j < B * CF:
-                        b = j // CF
-                        c = j % CF
-                        z = s_tmp[b, c]
-                        s_tmp[b, c] = z / (cutlass.Float32(1.0) + cmath.exp(-z))
-                        j += T
-                    # l>0: z * sigmoid(gate[expand[dm-1]])
-                    j = tidx
-                    while j < B * (Dm - 1) * CF:
-                        b = j // ((Dm - 1) * CF)
-                        rem = j % ((Dm - 1) * CF)
-                        d1 = rem // CF
-                        c = rem % CF
-                        lidx = mExpand[d1]
-                        z = s_tmp[b, (d1 + 1) * CF + c]
-                        s_tmp[b, (d1 + 1) * CF + c] = (
-                            z * s_scr[GATE_OFF + b * GO + lidx * CF + c]
-                        )
-                        j += T
-                    cute.arch.sync_threads()
-
-                # --- residual add: s_cur += activation(s_tmp) ---
-                i = tidx
-                while i < B * FLAT:
-                    b = i // FLAT
-                    rem = i % FLAT
-                    s_cur[b, rem] = s_cur[b, rem] + s_tmp[b, rem]
-                    i += T
-                cute.arch.sync_threads()
-
-            # === Step 5. Write the pre-focus-compete local features (E, F, D_m, Cf) ===
-            i = tidx
-            while i < B * FLAT:
-                b = i // FLAT
-                rem = i % FLAT
-                o = rem // CF
-                c = rem % CF
-                mOut[e0 + b, focus, o, c] = s_cur[b, rem]
-                i += T
-
-    class ForwardRunner:
-        """Compile-once driver for :class:`ForwardProgram`.
-
-        Prepares the static packed weights on construction and dispatches the
-        compiled kernel over dynamic edge counts.
-
-        Parameters
-        ----------
-        weights
-            Packed weights exposing ``so2_w`` (L, F, D_m*Cf, D_m*Cf),
-            ``gate_w`` (L, Cf, F, lmax*Cf), ``has_gate`` (L,), and
-            ``channel_basis`` (C_wide,).
-        lmax, mmax, cf, n_focus, n_layers, bucket, threads, rb, rn
-            Kernel configuration (see :class:`ForwardProgram`).
-        """
-
-        def __init__(
-            self,
-            weights,
-            *,
-            lmax: int,
-            mmax: int,
-            cf: int,
-            n_focus: int,
-            n_layers: int,
-            bucket: int,
-            threads: int,
-            rb: int,
-            rn: int,
-        ) -> None:
-            self.op = ForwardProgram(
-                lmax=lmax,
-                mmax=mmax,
-                cf=cf,
-                n_focus=n_focus,
-                n_layers=n_layers,
-                bucket=bucket,
-                threads=threads,
-                rb=rb,
-                rn=rn,
-            )
-            self._B = bucket
-            self.nf, self.cf, self.Dm = n_focus, cf, self.op.Dm
-            self._compiled = None
-            self._pack(weights, lmax, cf, n_focus, n_layers)
-
-        def _pack(self, w, lmax: int, cf: int, nf: int, nl: int) -> None:
-            dev = w.so2_w.device
-            self.m_w = w.so2_w.detach().contiguous()
-            gate = torch.zeros(nl, nf, cf, lmax * cf, device=dev, dtype=torch.float32)
-            for layer in range(nl):
-                if bool(w.has_gate[layer]):
-                    gate[layer] = w.gate_w[layer].detach().permute(1, 0, 2).contiguous()
-            self.m_gw = gate.contiguous()
-            self.m_cb = w.channel_basis.detach().contiguous()
-            # m-major degree index l(dm); the gate expand maps dm>0 -> (l-1).
-            l_index = list(range(lmax + 1))
-            for m in range(1, self.op.mmax + 1):
-                l_index += list(range(m, lmax + 1)) * 2
-            self.m_expand = torch.tensor(
-                [li - 1 for li in l_index[1:]], device=dev, dtype=torch.int32
-            ).contiguous()
-
-        @staticmethod
-        def _dyn(t: torch.Tensor, leading: int):
-            return from_dlpack(t, assumed_align=16).mark_layout_dynamic(
-                leading_dim=leading
-            )
-
-        def __call__(
-            self,
-            x: torch.Tensor,
-            src: torch.Tensor,
-            d_to_m: torch.Tensor,
-            kc: torch.Tensor,
-        ) -> tuple[torch.Tensor, torch.Tensor]:
-            """Run the fused forward.
-
-            Parameters
-            ----------
-            x : torch.Tensor
-                Node features with shape (N, D, C_wide).
-            src : torch.Tensor
-                Per-edge source-node indices with shape (E,).
-            d_to_m : torch.Tensor
-                Row-projected Wigner-D with shape (E, D_m, D).
-            kc : torch.Tensor
-                Radial degree kernel with shape (E, D_m, D_m).
-
-            Returns
-            -------
-            x_local : torch.Tensor
-                Pre-focus-compete local features with shape (E, F, D_m, Cf).
-            focus_gate : torch.Tensor
-                Pre-mixing l=0 scalar with shape (E, F, Cf).
-            """
-            n_edge = src.shape[0]
-            n_bucket = (n_edge + self._B - 1) // self._B
-            n_pad = n_bucket * self._B
-            s32 = src.to(torch.int32)
-            if n_pad > n_edge:
-                s32 = torch.cat([s32, s32.new_zeros(n_pad - n_edge)])
-            out = x.new_empty(n_pad, self.nf, self.Dm, self.cf)
-            fgate = x.new_empty(n_pad, self.nf, self.cf)
-            views = (
-                self._dyn(x, 2),
-                self._dyn(s32, 0),
-                self._dyn(d_to_m, 2),
-                self._dyn(kc, 2),
-                from_dlpack(self.m_cb, assumed_align=16),
-                from_dlpack(self.m_w, assumed_align=16),
-                from_dlpack(self.m_gw, assumed_align=16),
-                from_dlpack(self.m_expand, assumed_align=16),
-                self._dyn(out, 3),
-                self._dyn(fgate, 2),
-            )
-            args = (*views, cutlass.Int32(n_edge), cutlass.Int32(n_bucket))
-            stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream)
-            if self._compiled is None:
-                self._compiled = cute.compile(self.op, *args, stream=stream)
-            self._compiled(*args, stream=stream)
-            return out[:n_edge], fgate[:n_edge]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/gie.py b/deepmd/pt_expt/kernels/cute/sezm/gie.py
new file mode 100644
index 0000000000..972ece4f88
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/gie.py
@@ -0,0 +1,739 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+# pyright: reportMissingImports=false
+# ruff: noqa: ANN001
+"""Opt-in CuTe fusion for the DPA4 geometric initial embedding.
+
+The eager implementation materializes both ``radial_value_for_row`` and
+``non_scalar_message`` with shape ``(E, D - 1, C)`` before reducing by
+destination. This module computes the same strict-FP32 expression directly
+into ``(N, D, C)`` using the destination-sorted edge list. Its first backward
+produces radial, zonal/Wigner, source-gate, and degree-normalization gradients
+without an edge-by-row-by-channel temporary.
+
+``DP_CUTE_INFER`` is the master opt-in. The SM80/SM86 path enables this
+fusion by default; ``DP_CUTE_GIE=0`` disables it explicitly.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+import threading
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+import torch
+from torch import (
+    Tensor,
+)
+
+from .so2.metadata import (
+    build_destination_row_ptr,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+try:
+    import cutlass
+    import cutlass.cute as cute
+    import cutlass.torch as cutlass_torch
+    from cuda.bindings.driver import CUstream  # noqa: TC002
+    from cutlass.cute.runtime import (
+        from_dlpack,
+    )
+
+    SEZM_CUTE_GIE_AVAILABLE = True
+except Exception:  # pragma: no cover - import guard for non-CuTe environments
+    SEZM_CUTE_GIE_AVAILABLE = False
+
+
+def is_cute_gie_enabled(device: torch.device | None = None) -> bool:
+    """Return whether the architecture-selected GIE path is enabled."""
+    from .runtime_policy import (
+        is_gie_enabled,
+    )
+
+    if device is None:
+        if not torch.cuda.is_available():
+            return False
+        device_index = torch.cuda.current_device()
+    else:
+        if device.type != "cuda":
+            return False
+        device_index = device.index
+        if device_index is None:
+            device_index = torch.cuda.current_device()
+    return is_gie_enabled(tuple(torch.cuda.get_device_capability(device_index)))
+
+
+def _backward_compile_key(
+    device_identity: tuple[int, int, int],
+    lmax: int,
+    channels: int,
+    has_gate: bool,
+    radial_stride: tuple[int, ...],
+    dst_dtype: torch.dtype,
+) -> tuple[Any, ...]:
+    """Build an ABI-complete GIE backward compilation key."""
+    return (
+        "gie_bwd",
+        *device_identity,
+        lmax,
+        channels,
+        has_gate,
+        tuple(radial_stride),
+        dst_dtype,
+    )
+
+
+def _degree_slots(lmax: int, *, device: torch.device) -> Tensor:
+    degrees = torch.arange(1, lmax + 1, device=device, dtype=torch.long)
+    return torch.repeat_interleave(degrees - 1, 2 * degrees + 1)
+
+
+def _standard_index_contract(module: Any, lmax: int, row_count: int) -> bool:
+    rows = getattr(module, "non_scalar_row_index", None)
+    slots = getattr(module, "radial_slot_index_for_row", None)
+    if not isinstance(rows, Tensor) or not isinstance(slots, Tensor):
+        return False
+    if rows.numel() != row_count or slots.numel() != row_count:
+        return False
+    if rows.dtype != torch.long or slots.dtype != torch.long:
+        return False
+
+    # These buffers are constructor-owned and immutable. Validate their values
+    # when host-resident without introducing a CUDA synchronization.
+    if rows.device.type == "cpu" and slots.device.type == "cpu":
+        expected_rows = torch.arange(
+            1, row_count + 1, dtype=torch.long, device=torch.device("cpu")
+        )
+        expected_slots = _degree_slots(lmax, device=torch.device("cpu"))
+        return bool(
+            torch.equal(rows, expected_rows) and torch.equal(slots, expected_slots)
+        )
+    return True
+
+
+def validate_gie_contract(
+    module: Any,
+    n_nodes: int,
+    edge_cache: Any,
+    radial: Tensor,
+    zonal: Tensor,
+) -> bool:
+    """Validate the shape/layout contract without inspecting dynamic edge values."""
+    lmax = int(getattr(module, "lmax", -1))
+    channels = int(getattr(module, "channels", -1))
+    if lmax <= 0 or channels <= 0 or n_nodes <= 0:
+        return False
+    row_count = (lmax + 1) ** 2 - 1
+    dst = getattr(edge_cache, "dst", None)
+    inv_sqrt_deg = getattr(edge_cache, "inv_sqrt_deg", None)
+    gate = getattr(edge_cache, "edge_src_gate", None)
+    if not bool(getattr(edge_cache, "destinations_sorted", False)):
+        return False
+    if not isinstance(dst, Tensor) or not isinstance(inv_sqrt_deg, Tensor):
+        return False
+    if radial.dim() != 3 or zonal.dim() != 2 or dst.dim() != 1:
+        return False
+    edge_count = radial.shape[0]
+    if edge_count == 0:
+        return False
+    if (
+        zonal.shape != (edge_count, row_count)
+        or radial.shape[1:] != (lmax, channels)
+        or dst.shape[0] != edge_count
+        or inv_sqrt_deg.shape != (n_nodes, 1, 1)
+    ):
+        return False
+    if radial.dtype != torch.float32 or zonal.dtype != torch.float32:
+        return False
+    if inv_sqrt_deg.dtype != torch.float32:
+        return False
+    if dst.dtype not in (torch.int32, torch.int64):
+        return False
+    if not (radial.device == zonal.device == dst.device == inv_sqrt_deg.device):
+        return False
+    if radial.stride(-1) != 1 or radial.stride(-2) != channels:
+        return False
+    if (
+        not zonal.is_contiguous()
+        or not dst.is_contiguous()
+        or not inv_sqrt_deg.is_contiguous()
+    ):
+        return False
+    if gate is not None:
+        if not isinstance(gate, Tensor):
+            return False
+        if gate.shape not in ((edge_count,), (edge_count, 1)):
+            return False
+        if gate.dtype != torch.float32 or gate.device != radial.device:
+            return False
+        if not gate.is_contiguous():
+            return False
+    return _standard_index_contract(module, lmax, row_count)
+
+
+if SEZM_CUTE_GIE_AVAILABLE:
+    _F32 = cutlass.Float32
+    _I32 = cutlass.Int32
+    _WARPS_PER_BLOCK = 4
+    _LANES = 32
+
+    def _build_forward(lmax: int, channels: int, has_gate: bool) -> Callable:
+        row_count = (lmax + 1) ** 2 - 1
+
+        @cute.kernel
+        def kernel(m_radial, m_zonal, m_inv, m_dst_ptr, m_gate, m_out) -> None:
+            node, _, _ = cute.arch.block_idx()
+            lane, warp, _ = cute.arch.thread_idx()
+
+            if warp == 0:
+                for channel in cutlass.range(lane, channels, _LANES, unroll=1):
+                    m_out[node, channel] = _F32(0.0)
+
+            lo = m_dst_ptr[node].to(_I32)
+            hi = m_dst_ptr[node + 1].to(_I32)
+            for row in cutlass.range(warp, row_count, _WARPS_PER_BLOCK, unroll=1):
+                radial_slot = _I32(0)
+                for degree in cutlass.range_constexpr(lmax):
+                    start = (degree + 1) * (degree + 1) - 1
+                    stop = (degree + 2) * (degree + 2) - 1
+                    if row >= start and row < stop:
+                        radial_slot = degree
+                for channel in cutlass.range(lane, channels, _LANES, unroll=1):
+                    acc = _F32(0.0)
+                    for edge in cutlass.range(lo, hi, 1, unroll=1):
+                        scale = _F32(1.0)
+                        if has_gate:
+                            scale = m_gate[edge].to(_F32)
+                        acc += (
+                            m_zonal[edge, row].to(_F32)
+                            * m_radial[edge, radial_slot * channels + channel].to(_F32)
+                            * scale
+                        )
+                    m_out[node, (row + 1) * channels + channel] = acc * m_inv[
+                        node, 0
+                    ].to(_F32)
+
+        @cute.jit
+        def host(
+            m_radial,
+            m_zonal,
+            m_inv,
+            m_dst_ptr,
+            m_gate,
+            m_out,
+            stream: CUstream,
+        ) -> None:
+            nodes, _ = m_out.shape
+            kernel(m_radial, m_zonal, m_inv, m_dst_ptr, m_gate, m_out).launch(
+                grid=[nodes, 1, 1],
+                block=[_LANES, _WARPS_PER_BLOCK, 1],
+                stream=stream,
+            )
+
+        return host
+
+    def _build_backward(lmax: int, channels: int, has_gate: bool) -> Callable:
+        row_count = (lmax + 1) ** 2 - 1
+        full_width = (row_count + 1) * channels
+
+        @cute.kernel
+        def edge_kernel(
+            m_grad_out,
+            m_radial,
+            m_zonal,
+            m_inv,
+            m_dst,
+            m_gate,
+            m_grad_radial,
+            m_grad_zonal,
+            m_grad_gate,
+        ) -> None:
+            block_edge, _, _ = cute.arch.block_idx()
+            lane, warp, _ = cute.arch.thread_idx()
+            edge = block_edge * _WARPS_PER_BLOCK + warp
+            edge_count, _ = m_radial.shape
+            load_edge = edge
+            if edge >= edge_count:
+                load_edge = 0
+            node = m_dst[load_edge].to(_I32)
+            norm = m_inv[node, 0].to(_F32)
+            gate_value = _F32(1.0)
+            if has_gate:
+                gate_value = m_gate[load_edge].to(_F32)
+            grad_gate_acc = _F32(0.0)
+
+            if channels <= _LANES:
+                active_channel = lane < channels
+                for degree in cutlass.range_constexpr(lmax):
+                    start = (degree + 1) * (degree + 1) - 1
+                    stop = (degree + 2) * (degree + 2) - 1
+                    grad_radial_acc = _F32(0.0)
+                    for row in cutlass.range_constexpr(start, stop, 1):
+                        grad_value = _F32(0.0)
+                        radial_value = _F32(0.0)
+                        zonal_value = _F32(0.0)
+                        if active_channel:
+                            grad_value = (
+                                m_grad_out[node, (row + 1) * channels + lane].to(_F32)
+                                * norm
+                            )
+                            radial_value = m_radial[
+                                load_edge, degree * channels + lane
+                            ].to(_F32)
+                            zonal_value = m_zonal[load_edge, row].to(_F32)
+                            grad_radial_acc += grad_value * zonal_value * gate_value
+                            grad_gate_acc += grad_value * zonal_value * radial_value
+                        grad_zonal_value = cute.arch.warp_reduction_sum(
+                            grad_value * radial_value * gate_value
+                        )
+                        if lane == 0 and edge < edge_count:
+                            m_grad_zonal[edge, row] = grad_zonal_value
+                    if active_channel and edge < edge_count:
+                        m_grad_radial[edge, degree * channels + lane] = grad_radial_acc
+            else:
+                for degree in cutlass.range_constexpr(lmax):
+                    start = (degree + 1) * (degree + 1) - 1
+                    stop = (degree + 2) * (degree + 2) - 1
+                    for channel in cutlass.range(lane, channels, _LANES, unroll=1):
+                        grad_radial_acc = _F32(0.0)
+                        for row in cutlass.range_constexpr(start, stop, 1):
+                            grad_value = (
+                                m_grad_out[node, (row + 1) * channels + channel].to(
+                                    _F32
+                                )
+                                * norm
+                            )
+                            radial_value = m_radial[
+                                load_edge, degree * channels + channel
+                            ].to(_F32)
+                            zonal_value = m_zonal[load_edge, row].to(_F32)
+                            grad_radial_acc += grad_value * zonal_value * gate_value
+                            grad_gate_acc += grad_value * zonal_value * radial_value
+                        if edge < edge_count:
+                            m_grad_radial[edge, degree * channels + channel] = (
+                                grad_radial_acc
+                            )
+                    for row in cutlass.range_constexpr(start, stop, 1):
+                        grad_zonal_acc = _F32(0.0)
+                        for channel in cutlass.range(lane, channels, _LANES, unroll=1):
+                            grad_value = (
+                                m_grad_out[node, (row + 1) * channels + channel].to(
+                                    _F32
+                                )
+                                * norm
+                            )
+                            radial_value = m_radial[
+                                load_edge, degree * channels + channel
+                            ].to(_F32)
+                            grad_zonal_acc += grad_value * radial_value * gate_value
+                        grad_zonal_value = cute.arch.warp_reduction_sum(grad_zonal_acc)
+                        if lane == 0 and edge < edge_count:
+                            m_grad_zonal[edge, row] = grad_zonal_value
+
+            if has_gate:
+                grad_gate_value = cute.arch.warp_reduction_sum(grad_gate_acc)
+                if lane == 0 and edge < edge_count:
+                    m_grad_gate[edge] = grad_gate_value
+
+        @cute.kernel
+        def inv_kernel(m_grad_out, m_out, m_inv, m_grad_inv) -> None:
+            block_node, _, _ = cute.arch.block_idx()
+            lane, warp, _ = cute.arch.thread_idx()
+            node = block_node * _WARPS_PER_BLOCK + warp
+            node_count, _ = m_out.shape
+            load_node = node
+            if node >= node_count:
+                load_node = 0
+            acc = _F32(0.0)
+            for idx in cutlass.range(channels + lane, full_width, _LANES, unroll=1):
+                acc += m_grad_out[load_node, idx].to(_F32) * m_out[load_node, idx].to(
+                    _F32
+                )
+            acc = cute.arch.warp_reduction_sum(acc)
+            if lane == 0 and node < node_count:
+                m_grad_inv[node, 0] = acc / m_inv[node, 0].to(_F32)
+
+        @cute.jit
+        def host(
+            m_grad_out,
+            m_radial,
+            m_zonal,
+            m_inv,
+            m_dst,
+            m_gate,
+            m_out,
+            m_grad_radial,
+            m_grad_zonal,
+            m_grad_inv,
+            m_grad_gate,
+            stream: CUstream,
+        ) -> None:
+            edge_count, _ = m_radial.shape
+            node_count, _ = m_out.shape
+            edge_kernel(
+                m_grad_out,
+                m_radial,
+                m_zonal,
+                m_inv,
+                m_dst,
+                m_gate,
+                m_grad_radial,
+                m_grad_zonal,
+                m_grad_gate,
+            ).launch(
+                grid=[cute.ceil_div(edge_count, _WARPS_PER_BLOCK), 1, 1],
+                block=[_LANES, _WARPS_PER_BLOCK, 1],
+                stream=stream,
+            )
+            inv_kernel(m_grad_out, m_out, m_inv, m_grad_inv).launch(
+                grid=[cute.ceil_div(node_count, _WARPS_PER_BLOCK), 1, 1],
+                block=[_LANES, _WARPS_PER_BLOCK, 1],
+                stream=stream,
+            )
+
+        return host
+
+    _compile_lock = threading.Lock()
+    _compiled: dict[tuple[Any, ...], Any] = {}
+
+    def _as_cute(tensor: Tensor) -> Any:
+        value = from_dlpack(tensor)
+        if tensor.dim() <= 1:
+            return value.mark_layout_dynamic()
+        return value.mark_layout_dynamic(leading_dim=tensor.dim() - 1)
+
+    def _device_key(tensor: Tensor) -> tuple[int, int, int]:
+        index = tensor.device.index
+        if index is None:
+            index = torch.cuda.current_device()
+        major, minor = torch.cuda.get_device_capability(index)
+        return index, major, minor
+
+    def _get_compiled(
+        key: tuple[Any, ...],
+        builder: Callable[[], Callable],
+        example_args: tuple[Any, ...],
+    ) -> Any:
+        compiled = _compiled.get(key)
+        if compiled is not None:
+            return compiled
+        with _compile_lock:
+            compiled = _compiled.get(key)
+            if compiled is None:
+                compiled = cute.compile(builder(), *example_args)
+                _compiled[key] = compiled
+        return compiled
+
+    def _flat_radial(radial: Tensor) -> Tensor:
+        return radial.view(radial.shape[0], radial.shape[1] * radial.shape[2])
+
+    def _flat_node(tensor: Tensor) -> Tensor:
+        return tensor.view(tensor.shape[0], -1)
+
+    def _launch_forward(
+        radial: Tensor,
+        zonal: Tensor,
+        inv_sqrt_deg: Tensor,
+        dst_ptr: Tensor,
+        gate: Tensor,
+        lmax: int,
+        has_gate: bool,
+    ) -> Tensor:
+        channels = radial.shape[2]
+        out = torch.empty(
+            inv_sqrt_deg.shape[0],
+            (lmax + 1) ** 2,
+            channels,
+            device=radial.device,
+            dtype=radial.dtype,
+        )
+        radial_flat = _flat_radial(radial.detach())
+        inv_flat = _flat_node(inv_sqrt_deg.detach())
+        gate_flat = gate.detach().view(-1)
+        out_flat = _flat_node(out)
+        args = tuple(
+            _as_cute(value)
+            for value in (
+                radial_flat,
+                zonal.detach(),
+                inv_flat,
+                dst_ptr.detach(),
+                gate_flat,
+                out_flat,
+            )
+        )
+        device_identity = _device_key(radial)
+        with torch.cuda.device(device_identity[0]):
+            stream = cutlass_torch.current_stream()
+            key = (
+                "gie_fwd",
+                *device_identity,
+                lmax,
+                channels,
+                has_gate,
+                tuple(radial_flat.stride()),
+                dst_ptr.dtype,
+            )
+            compiled = _get_compiled(
+                key,
+                lambda: _build_forward(lmax, channels, has_gate),
+                (*args, stream),
+            )
+            compiled(*args, stream)
+        return out
+
+    def _launch_backward(
+        grad_out: Tensor,
+        radial: Tensor,
+        zonal: Tensor,
+        inv_sqrt_deg: Tensor,
+        dst: Tensor,
+        gate: Tensor,
+        out: Tensor,
+        lmax: int,
+        has_gate: bool,
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        channels = radial.shape[2]
+        grad_out_flat = _flat_node(grad_out.detach().contiguous())
+        radial_flat = _flat_radial(radial.detach())
+        inv_flat = _flat_node(inv_sqrt_deg.detach())
+        gate_flat = gate.detach().view(-1)
+        out_flat = _flat_node(out.detach())
+        grad_radial = torch.empty(
+            radial.shape,
+            device=radial.device,
+            dtype=radial.dtype,
+            memory_format=torch.contiguous_format,
+        )
+        grad_zonal = torch.empty_like(zonal, memory_format=torch.contiguous_format)
+        grad_inv = torch.empty_like(inv_sqrt_deg, memory_format=torch.contiguous_format)
+        grad_gate = torch.empty_like(gate, memory_format=torch.contiguous_format)
+        grad_radial_flat = _flat_radial(grad_radial)
+        grad_inv_flat = _flat_node(grad_inv)
+        grad_gate_flat = grad_gate.view(-1)
+        args = tuple(
+            _as_cute(value)
+            for value in (
+                grad_out_flat,
+                radial_flat,
+                zonal.detach(),
+                inv_flat,
+                dst.detach(),
+                gate_flat,
+                out_flat,
+                grad_radial_flat,
+                grad_zonal,
+                grad_inv_flat,
+                grad_gate_flat,
+            )
+        )
+        device_identity = _device_key(radial)
+        with torch.cuda.device(device_identity[0]):
+            stream = cutlass_torch.current_stream()
+            key = _backward_compile_key(
+                device_identity,
+                lmax,
+                channels,
+                has_gate,
+                tuple(radial_flat.stride()),
+                dst.dtype,
+            )
+            compiled = _get_compiled(
+                key,
+                lambda: _build_backward(lmax, channels, has_gate),
+                (*args, stream),
+            )
+            compiled(*args, stream)
+        return grad_radial, grad_zonal, grad_inv, grad_gate
+
+    @torch.library.custom_op(
+        "sezm_cute::gie_fused", mutates_args=(), device_types="cuda"
+    )
+    def _gie_op(
+        radial: Tensor,
+        zonal: Tensor,
+        inv_sqrt_deg: Tensor,
+        dst: Tensor,
+        dst_ptr: Tensor,
+        gate: Tensor,
+        lmax: int,
+        has_gate: bool,
+    ) -> Tensor:
+        del dst
+        return _launch_forward(
+            radial,
+            zonal,
+            inv_sqrt_deg,
+            dst_ptr,
+            gate,
+            int(lmax),
+            bool(has_gate),
+        )
+
+    @_gie_op.register_fake
+    def _(
+        radial: Tensor,
+        zonal: Tensor,
+        inv_sqrt_deg: Tensor,
+        dst: Tensor,
+        dst_ptr: Tensor,
+        gate: Tensor,
+        lmax: int,
+        has_gate: bool,
+    ) -> Tensor:
+        del zonal, dst, dst_ptr, gate, has_gate
+        return radial.new_empty(
+            (inv_sqrt_deg.shape[0], (int(lmax) + 1) ** 2, radial.shape[2])
+        )
+
+    @torch.library.custom_op(
+        "sezm_cute::gie_fused_bwd", mutates_args=(), device_types="cuda"
+    )
+    def _gie_bwd_op(
+        grad_out: Tensor,
+        radial: Tensor,
+        zonal: Tensor,
+        inv_sqrt_deg: Tensor,
+        dst: Tensor,
+        gate: Tensor,
+        out: Tensor,
+        lmax: int,
+        has_gate: bool,
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        return _launch_backward(
+            grad_out,
+            radial,
+            zonal,
+            inv_sqrt_deg,
+            dst,
+            gate,
+            out,
+            int(lmax),
+            bool(has_gate),
+        )
+
+    @_gie_bwd_op.register_fake
+    def _(
+        grad_out: Tensor,
+        radial: Tensor,
+        zonal: Tensor,
+        inv_sqrt_deg: Tensor,
+        dst: Tensor,
+        gate: Tensor,
+        out: Tensor,
+        lmax: int,
+        has_gate: bool,
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        del grad_out, dst, out, lmax, has_gate
+        return (
+            torch.empty_like(radial, memory_format=torch.contiguous_format),
+            torch.empty_like(zonal, memory_format=torch.contiguous_format),
+            torch.empty_like(inv_sqrt_deg, memory_format=torch.contiguous_format),
+            torch.empty_like(gate, memory_format=torch.contiguous_format),
+        )
+
+    def _gie_setup_context(
+        ctx: Any,
+        inputs: tuple[Any, ...],
+        output: Tensor,
+    ) -> None:
+        radial, zonal, inv_sqrt_deg, dst, _dst_ptr, gate, lmax, has_gate = inputs
+        ctx.save_for_backward(radial, zonal, inv_sqrt_deg, dst, gate, output)
+        ctx.lmax = int(lmax)
+        ctx.has_gate = bool(has_gate)
+
+    def _gie_backward(ctx: Any, grad_out: Tensor) -> tuple[Any, ...]:
+        radial, zonal, inv_sqrt_deg, dst, gate, out = ctx.saved_tensors
+        grad_radial, grad_zonal, grad_inv, grad_gate = _gie_bwd_op(
+            grad_out,
+            radial,
+            zonal,
+            inv_sqrt_deg,
+            dst,
+            gate,
+            out,
+            ctx.lmax,
+            ctx.has_gate,
+        )
+        return grad_radial, grad_zonal, grad_inv, None, None, grad_gate, None, None
+
+    _gie_op.register_autograd(_gie_backward, setup_context=_gie_setup_context)
+
+
+def gie_fused_cuda(
+    radial: Tensor,
+    zonal: Tensor,
+    inv_sqrt_deg: Tensor,
+    dst: Tensor,
+    gate: Tensor,
+    *,
+    n_nodes: int,
+    lmax: int,
+) -> Tensor:
+    """Run the fused CUDA path after the caller has validated its contract."""
+    if not SEZM_CUTE_GIE_AVAILABLE:
+        raise RuntimeError("CuTe DSL is unavailable")
+    dst_ptr = build_destination_row_ptr(dst, n_nodes)
+    has_gate = gate.numel() != 0
+    kernel_gate = gate if has_gate else radial.new_ones((1,))
+    return _gie_op(
+        radial,
+        zonal,
+        inv_sqrt_deg,
+        dst,
+        dst_ptr,
+        kernel_gate,
+        int(lmax),
+        has_gate,
+    )
+
+
+def maybe_run_cute_gie(
+    module: Any,
+    *,
+    n_nodes: int,
+    edge_cache: Any,
+    radial_feat: Tensor,
+    zonal_coupling: Tensor,
+) -> Tensor | None:
+    """Run the opt-in path or return ``None`` for the eager fallback."""
+    if (
+        not is_cute_gie_enabled(radial_feat.device)
+        or not SEZM_CUTE_GIE_AVAILABLE
+        or bool(getattr(module, "training", True))
+        or not radial_feat.is_cuda
+        or not validate_gie_contract(
+            module, n_nodes, edge_cache, radial_feat, zonal_coupling
+        )
+    ):
+        return None
+    gate = getattr(edge_cache, "edge_src_gate", None)
+    if gate is None:
+        gate = radial_feat.new_empty((0,))
+    return gie_fused_cuda(
+        radial_feat,
+        zonal_coupling,
+        edge_cache.inv_sqrt_deg,
+        edge_cache.dst,
+        gate,
+        n_nodes=n_nodes,
+        lmax=int(module.lmax),
+    )
+
+
+__all__ = [
+    "SEZM_CUTE_GIE_AVAILABLE",
+    "gie_fused_cuda",
+    "is_cute_gie_enabled",
+    "maybe_run_cute_gie",
+    "validate_gie_contract",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/operator.py b/deepmd/pt_expt/kernels/cute/sezm/operator.py
deleted file mode 100644
index a5a116f301..0000000000
--- a/deepmd/pt_expt/kernels/cute/sezm/operator.py
+++ /dev/null
@@ -1,324 +0,0 @@
-# SPDX-License-Identifier: LGPL-3.0-or-later
-"""
-Autograd operator and entry point wiring the fused CuTe kernels into the SeZM
-SO(2) convolution value path.
-
-:class:`_SO2ValuePathFunction` runs the fused forward kernel (saving only the
-small inputs ``x``, ``D_to_m``, ``Kc``) and, on the force path, recomputes the
-value path in the fused backward kernel -- so the per-edge ``E x D_m x C``
-intermediates stay off DRAM across the whole autograd graph. :func:`make_cute_value_path`
-builds the per-convolution entry :class:`_CuteSO2ValuePath`, which computes the
-radial-/scalar-only tensors (``D_to_m``, ``Kc``, ``rad_feat``) in ordinary
-autograd, invokes the operator to produce the pre-focus-compete local features
-and the pre-mixing ``l = 0`` scalar, and applies the focus competition -- exactly
-as the reference path does. The packed weights are extracted lazily on the first
-call so they reflect the loaded checkpoint.
-"""
-
-from __future__ import (
-    annotations,
-)
-
-from dataclasses import (
-    dataclass,
-)
-from typing import (
-    TYPE_CHECKING,
-)
-
-import torch
-
-from .forward import (
-    SEZM_CUTE_AVAILABLE,
-)
-
-if TYPE_CHECKING:
-    from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import (
-        EdgeCache,
-    )
-    from deepmd.dpmodel.descriptor.dpa4_nn.so2 import (
-        SO2Convolution,
-    )
-
-if SEZM_CUTE_AVAILABLE:
-    from .backward import (
-        BackwardRunner,
-    )
-    from .forward import (
-        ForwardRunner,
-    )
-
-# Validated configuration of the fused operator: the deployed three-layer
-# ``[gated, gated, identity]`` mixing stack in the ``lmax = 3, mmax = 1`` layout.
-_SUPPORTED_LMAX = 3
-_SUPPORTED_MMAX = 1
-_SUPPORTED_LAYERS = 3
-
-
-@dataclass
-class _PackedWeights:
-    """Static weights of the SO(2) value path, packed for the fused kernels.
-
-    Attributes
-    ----------
-    so2_w : torch.Tensor
-        Assembled block-diagonal SO2Linear weight per layer with shape
-        (L, F, D_m*Cf, D_m*Cf), in ``(in, out)`` convention.
-    gate_w : torch.Tensor
-        GatedActivation ``FocusLinear`` weight per layer with shape
-        (L, Cf, F, lmax*Cf); zero for non-gated layers.
-    has_gate : torch.Tensor
-        Boolean per-layer flag with shape (L,).
-    channel_basis : torch.Tensor
-        Radial degree-mixer channel basis with shape (C_wide,).
-    """
-
-    so2_w: torch.Tensor
-    gate_w: torch.Tensor
-    has_gate: torch.Tensor
-    channel_basis: torch.Tensor
-
-
-def _pack_weights(conv: SO2Convolution) -> _PackedWeights:
-    """Extract and pack the SO(2) value-path weights from a convolution block."""
-    n_layers = conv.mixing_layers
-    so2_w = torch.stack(
-        [
-            conv.so2_linears[layer]._build_so2_weight().permute(1, 0, 2).contiguous()
-            for layer in range(n_layers)
-        ]
-    )
-    gate_w, has_gate = [], []
-    for layer in range(n_layers):
-        non_linear = conv.non_linearities[layer]
-        if type(non_linear).__name__ == "GatedActivation" and non_linear.lmax > 0:
-            gate_w.append(
-                non_linear.gate_linear.weight.view(
-                    conv.so2_focus_dim, conv.n_focus, conv.lmax * conv.so2_focus_dim
-                ).contiguous()
-            )
-            has_gate.append(True)
-        else:
-            gate_w.append(
-                torch.zeros_like(gate_w[0])
-                if gate_w
-                else torch.zeros(
-                    conv.so2_focus_dim,
-                    conv.n_focus,
-                    conv.lmax * conv.so2_focus_dim,
-                    device=so2_w.device,
-                    dtype=so2_w.dtype,
-                )
-            )
-            has_gate.append(False)
-    return _PackedWeights(
-        so2_w=so2_w,
-        gate_w=torch.stack(gate_w),
-        has_gate=torch.tensor(has_gate, device=so2_w.device, dtype=torch.bool),
-        channel_basis=conv.radial_degree_mixer.channel_basis.reshape(-1).contiguous(),
-    )
-
-
-class _SO2ValuePathFunction(torch.autograd.Function):
-    """Fused CuTe forward with a recompute backward for the force path."""
-
-    @staticmethod
-    def forward(ctx, x, d_to_m, kc, src, fwd_runner, bwd_runner):  # noqa: ANN001, ANN205
-        with torch.no_grad():
-            x_local, focus_gate = fwd_runner(
-                x.detach(), src, d_to_m.detach(), kc.detach()
-            )
-        ctx.save_for_backward(x, d_to_m, kc)
-        ctx.src = src
-        ctx.bwd_runner = bwd_runner
-        return x_local, focus_gate
-
-    @staticmethod
-    def backward(ctx, grad_local, grad_focus_gate):  # noqa: ANN001, ANN205
-        x, d_to_m, kc = ctx.saved_tensors
-        need = ctx.needs_input_grad
-        grad_x, grad_d_to_m, grad_kc = ctx.bwd_runner(
-            x.detach(),
-            ctx.src,
-            d_to_m.detach(),
-            kc.detach(),
-            grad_local.detach().contiguous(),
-            grad_focus_gate.detach().contiguous(),
-        )
-        return (
-            grad_x if need[0] else None,
-            grad_d_to_m if need[1] else None,
-            grad_kc if need[2] else None,
-            None,
-            None,
-            None,
-        )
-
-
-class _CuteSO2ValuePath:
-    """Per-convolution entry that runs the value path through the fused kernels.
-
-    The convolution is held by reference so the packed weights are extracted
-    lazily on the first call (after the checkpoint is loaded) and the kernels are
-    compiled on first use.
-
-    Parameters
-    ----------
-    conv : SO2Convolution
-        The owning convolution block.
-    bucket_fwd, threads_fwd, bucket_bwd, threads_bwd, rb, rn : int
-        Fused-kernel launch configuration.
-    """
-
-    def __init__(
-        self,
-        conv: SO2Convolution,
-        *,
-        bucket_fwd: int = 32,
-        threads_fwd: int = 1024,
-        bucket_bwd: int = 16,
-        threads_bwd: int = 512,
-        rb: int = 4,
-        rn: int = 4,
-    ) -> None:
-        self._conv = conv
-        self._cfg = {
-            "lmax": conv.lmax,
-            "mmax": conv.mmax,
-            "cf": conv.so2_focus_dim,
-            "n_focus": conv.n_focus,
-            "n_layers": conv.mixing_layers,
-            "rb": rb,
-            "rn": rn,
-        }
-        self._launch = {
-            "bucket_fwd": bucket_fwd,
-            "threads_fwd": threads_fwd,
-            "bucket_bwd": bucket_bwd,
-            "threads_bwd": threads_bwd,
-        }
-        self._fwd_runner = None
-        self._bwd_runner = None
-
-    def _build(self) -> None:
-        weights = _pack_weights(self._conv)
-        self._fwd_runner = ForwardRunner(
-            weights,
-            bucket=self._launch["bucket_fwd"],
-            threads=self._launch["threads_fwd"],
-            **self._cfg,
-        )
-        self._bwd_runner = BackwardRunner(
-            weights,
-            bucket=self._launch["bucket_bwd"],
-            threads=self._launch["threads_bwd"],
-            **self._cfg,
-        )
-
-    def __call__(
-        self,
-        x: torch.Tensor,
-        edge_cache: EdgeCache,
-        radial_feat: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Compute the SO(2) local features and radial features via the fused op.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Node features with shape (N, D, C_wide).
-        edge_cache : EdgeCache
-            Precomputed edge cache (provides ``src`` and the Wigner ``D_full``).
-        radial_feat : torch.Tensor
-            Per-edge radial features with shape (E, lmax+1, C).
-
-        Returns
-        -------
-        x_local : torch.Tensor
-            Post-focus-compete local features with shape (E, F, D_m, Cf).
-        rad_feat : torch.Tensor
-            Projected radial features with shape (E, D_m, C_wide); its ``l = 0``
-            slice is consumed by the attention aggregation.
-        """
-        if self._fwd_runner is None:
-            self._build()
-        conv = self._conv
-        src = edge_cache.src
-
-        # === Step 1. Radial-/scalar-only tensors (kept in ordinary autograd) ===
-        d_to_m = edge_cache.D_full[
-            :, : conv.ebed_dim_full, : conv.ebed_dim_full
-        ].index_select(
-            1,
-            conv.coeff_index_m,
-        )
-        rad_feat = radial_feat[:, conv.degree_index_m, :]
-        rad_feat = conv.radial_hidden_proj(rad_feat)
-        mixer = conv.radial_degree_mixer
-        kernel_flat = mixer._project_radial(rad_feat)
-        compact = kernel_flat.view(src.shape[0], mixer.degree_kernel_size, mixer.rank)
-        kc = mixer._scatter_rank_kernel(compact).squeeze(-1)
-
-        # === Step 2. Fused value path -> pre-focus-compete local + l=0 scalar ===
-        x_local, focus_gate = _SO2ValuePathFunction.apply(
-            x, d_to_m, kc, src, self._fwd_runner, self._bwd_runner
-        )
-
-        # === Step 3. Cross-focus softmax competition (rotation-free scalars) ===
-        if conv.focus_compete and conv.n_focus > 1:
-            alpha = conv._focus_alpha(focus_gate)
-            x_local = x_local * alpha.to(dtype=x_local.dtype).unsqueeze(-1).unsqueeze(
-                -1
-            )
-        return x_local, rad_feat
-
-
-def _is_supported(conv: SO2Convolution) -> bool:
-    """Return whether ``conv`` matches the validated fused-operator configuration."""
-    if (
-        conv.lmax != _SUPPORTED_LMAX
-        or conv.mmax != _SUPPORTED_MMAX
-        or conv.mixing_layers != _SUPPORTED_LAYERS
-        or conv.node_wise_grid_product is not None
-        or conv.use_so2_attn_res
-        or conv.layer_scale
-        or conv.radial_degree_mixer is None
-        or conv.radial_hidden_proj is None
-    ):
-        return False
-    if any(type(norm).__name__ != "Identity" for norm in conv.so2_inter_norms):
-        return False
-    if any(linear.bias0 is not None for linear in conv.so2_linears):
-        return False
-    non_linears = conv.non_linearities
-    if any(
-        type(non_linears[layer]).__name__ != "GatedActivation"
-        or (
-            getattr(non_linears[layer].scalar_act, "activation", None)
-            or getattr(non_linears[layer], "activation_function", None)
-        )
-        != "silu"
-        for layer in range(_SUPPORTED_LAYERS - 1)
-    ):
-        return False
-    return type(non_linears[-1]).__name__ == "Identity"
-
-
-def make_cute_value_path(conv: SO2Convolution) -> _CuteSO2ValuePath | None:
-    """Build the fused CuTe value-path entry for a convolution block.
-
-    Parameters
-    ----------
-    conv : SO2Convolution
-        The convolution block to accelerate.
-
-    Returns
-    -------
-    _CuteSO2ValuePath or None
-        The entry callable when the CuTe backend is available and ``conv`` matches
-        the validated configuration; otherwise ``None`` (the caller falls back to
-        the reference path).
-    """
-    if not SEZM_CUTE_AVAILABLE or not _is_supported(conv):
-        return None
-    return _CuteSO2ValuePath(conv)
diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/__init__.py
new file mode 100644
index 0000000000..bae1a1c36e
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Output-grid orchestration and kernels for the CuTe SeZM path."""
diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/__init__.py
new file mode 100644
index 0000000000..4cde7864c0
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""CuTe DSL kernels for SeZM output-grid contractions."""
diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/readout_l0.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/readout_l0.py
new file mode 100644
index 0000000000..e6b9048455
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/readout_l0.py
@@ -0,0 +1,839 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Tiled strict-FP32 Neo output readout for degree zero only."""
+
+from __future__ import (
+    annotations,
+)
+
+from collections.abc import (
+    Callable,
+)
+from functools import (
+    lru_cache,
+)
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from .tiled_product import (
+    FAKE_TENSOR_KW,
+    GRID_SIZE,
+    PACKED_COEFF_DIM,
+    STAGES,
+    THREADS,
+    TILE_K,
+    TILE_M,
+    TILE_N,
+)
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN201, ANN202, ANN204, TC002, TC003
+
+
+HIDDEN_CHANNELS = 192
+GRAM_FORWARD_THREADS = HIDDEN_CHANNELS
+GRAM_ELEMENTS = PACKED_COEFF_DIM * PACKED_COEFF_DIM
+
+
+class TiledReadoutL0GramForward:
+    """Evaluate the channelwise Gram bilinear with one CTA per node."""
+
+    @cute.jit
+    def __call__(
+        self,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        gram: cute.Tensor,
+        out: cute.Tensor,
+        stream: CUstream,
+    ):
+        gram_layout = cute.make_layout(
+            (PACKED_COEFF_DIM, PACKED_COEFF_DIM),
+            stride=(PACKED_COEFF_DIM, 1),
+        )
+        right_layout = cute.make_layout(
+            (PACKED_COEFF_DIM, HIDDEN_CHANNELS),
+            stride=(HIDDEN_CHANNELS, 1),
+        )
+        self.kernel(
+            left,
+            right,
+            gram,
+            out,
+            gram_layout,
+            right_layout,
+        ).launch(
+            grid=(left.shape[0], 1, 1),
+            block=[GRAM_FORWARD_THREADS, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        gram: cute.Tensor,
+        out: cute.Tensor,
+        gram_layout: cute.Layout,
+        right_layout: cute.Layout,
+    ):
+        channel, _, _ = cute.arch.thread_idx()
+        node, _, _ = cute.arch.block_idx()
+
+        smem = cutlass.utils.SmemAllocator()
+        s_gram = smem.allocate_tensor(cutlass.Float32, gram_layout, 16)
+        s_right = smem.allocate_tensor(cutlass.Float32, right_layout, 16)
+
+        for linear in cutlass.range(
+            channel,
+            GRAM_ELEMENTS,
+            GRAM_FORWARD_THREADS,
+            unroll=1,
+        ):
+            row = linear // PACKED_COEFF_DIM
+            col = linear - row * PACKED_COEFF_DIM
+            s_gram[row, col] = gram[row, col].to(cutlass.Float32)
+        for coeff in cutlass.range(0, PACKED_COEFF_DIM, 1, unroll=1):
+            s_right[coeff, channel] = right[node, coeff, channel].to(cutlass.Float32)
+        cute.arch.sync_threads()
+
+        value = cutlass.Float32(0.0)
+        for row in cutlass.range(0, PACKED_COEFF_DIM, 1, unroll=1):
+            transformed_right = cutlass.Float32(0.0)
+            for col in cutlass.range(0, PACKED_COEFF_DIM, 1, unroll=1):
+                transformed_right += s_gram[row, col] * s_right[col, channel]
+            value += left[node, row, channel].to(cutlass.Float32) * transformed_right
+        out[node, channel] = value.to(out.element_type)
+
+
+class TiledReadoutL0GramBackward:
+    """Apply the frozen 48x48 Gram matrix to one 64-channel tile."""
+
+    def __init__(self) -> None:
+        self.cta_tiler = (TILE_M, TILE_N, TILE_K)
+        self.channel_tiles = HIDDEN_CHANNELS // TILE_N
+        self.cta_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=THREADS,
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        dq0: cute.Tensor,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        gram: cute.Tensor,
+        grad_left: cute.Tensor,
+        grad_right: cute.Tensor,
+        stream: CUstream,
+    ):
+        sA_layout = cute.make_layout(
+            (TILE_M, TILE_K, STAGES),
+            stride=(1, TILE_M + 4, TILE_K * (TILE_M + 4)),
+        )
+        sB_layout = cute.make_layout(
+            (TILE_N, TILE_K, STAGES),
+            stride=(1, TILE_N, TILE_K * TILE_N),
+        )
+        dq0_layout = cute.make_layout((TILE_N,), stride=(1,))
+
+        copy_a_atom = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            left.element_type,
+            num_bits_per_copy=left.element_type.width,
+        )
+        copy_a_layout = cute.make_layout(
+            (THREADS // TILE_K, TILE_K),
+            stride=(TILE_K, 1),
+        )
+        tiled_copy_A = cute.make_tiled_copy_tv(
+            copy_a_atom,
+            copy_a_layout,
+            cute.make_layout((1, 1)),
+        )
+
+        vector = 4
+        copy_b_atom = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            left.element_type,
+            num_bits_per_copy=left.element_type.width * vector,
+        )
+        copy_b_major = TILE_N // vector
+        copy_b_layout = cute.make_layout(
+            (copy_b_major, THREADS // copy_b_major),
+            stride=(1, copy_b_major),
+        )
+        tiled_copy_B = cute.make_tiled_copy_tv(
+            copy_b_atom,
+            copy_b_layout,
+            cute.make_layout((vector, 1)),
+        )
+
+        atoms_layout = cute.make_layout(
+            (THREADS // 16, 16, 1),
+            stride=(16, 1, 0),
+        )
+        permutation_m = cute.make_layout(
+            (atoms_layout.shape[0], 4),
+            stride=(4, 1),
+        )
+        permutation_n = cute.make_layout(
+            (atoms_layout.shape[1], 4),
+            stride=(4, 1),
+        )
+        tiled_mma = cute.make_tiled_mma(
+            cute.nvgpu.MmaUniversalOp(cutlass.Float32),
+            atoms_layout,
+            permutation_mnk=(permutation_m, permutation_n, None),
+        )
+
+        self.kernel(
+            dq0,
+            left,
+            right,
+            gram,
+            grad_left,
+            grad_right,
+            sA_layout,
+            sB_layout,
+            dq0_layout,
+            tiled_copy_A,
+            tiled_copy_B,
+            tiled_mma,
+        ).launch(
+            grid=(left.shape[0], self.channel_tiles, 1),
+            block=[THREADS, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        dq0: cute.Tensor,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        gram: cute.Tensor,
+        grad_left: cute.Tensor,
+        grad_right: cute.Tensor,
+        sA_layout: cute.Layout,
+        sB_layout: cute.Layout,
+        dq0_layout: cute.Layout,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        node, channel_tile, _ = cute.arch.block_idx()
+
+        matrix_b_layout = cute.make_layout(
+            (HIDDEN_CHANNELS, PACKED_COEFF_DIM),
+            stride=(1, HIDDEN_CHANNELS),
+        )
+        left_b = cute.make_tensor(
+            left[node, None, None].iterator,
+            matrix_b_layout,
+        )
+        right_b = cute.make_tensor(
+            right[node, None, None].iterator,
+            matrix_b_layout,
+        )
+        gram_t = cute.make_tensor(
+            gram.iterator,
+            cute.make_layout(
+                (PACKED_COEFF_DIM, PACKED_COEFF_DIM),
+                stride=(1, PACKED_COEFF_DIM),
+            ),
+        )
+
+        smem = cutlass.utils.SmemAllocator()
+        sA_left = smem.allocate_tensor(cutlass.Float32, sA_layout, 16)
+        sA_right = smem.allocate_tensor(cutlass.Float32, sA_layout, 16)
+        sB_left = smem.allocate_tensor(cutlass.Float32, sB_layout, 16)
+        sB_right = smem.allocate_tensor(cutlass.Float32, sB_layout, 16)
+        sDq0 = smem.allocate_tensor(cutlass.Float32, dq0_layout, 16)
+
+        for local_channel in cutlass.range(
+            tidx,
+            TILE_N,
+            THREADS,
+            unroll=1,
+        ):
+            channel = channel_tile * TILE_N + local_channel
+            sDq0[local_channel] = dq0[node, channel].to(cutlass.Float32)
+        cute.arch.sync_threads()
+
+        self._dual_gram_adjoint(
+            gram,
+            gram_t,
+            right_b,
+            left_b,
+            grad_left[node, None, None],
+            grad_right[node, None, None],
+            sA_left,
+            sA_right,
+            sB_left,
+            sB_right,
+            sDq0,
+            tiled_copy_A,
+            tiled_copy_B,
+            tiled_mma,
+            tidx,
+            channel_tile,
+        )
+
+    @cute.jit
+    def _dual_gram_adjoint(
+        self,
+        mA_left: cute.Tensor,
+        mA_right: cute.Tensor,
+        mB_left: cute.Tensor,
+        mB_right: cute.Tensor,
+        mOut_left: cute.Tensor,
+        mOut_right: cute.Tensor,
+        sA_left: cute.Tensor,
+        sA_right: cute.Tensor,
+        sB_left: cute.Tensor,
+        sB_right: cute.Tensor,
+        sDq0: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        channel_tile: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        gA_left = cute.local_tile(
+            mA_left,
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(1, None, 1),
+        )
+        gA_right = cute.local_tile(
+            mA_right,
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(1, None, 1),
+        )
+        gB_left = cute.local_tile(
+            mB_left,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(None, 1, 1),
+        )
+        gB_right = cute.local_tile(
+            mB_right,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(None, 1, 1),
+        )
+        gOut_left = cute.local_tile(
+            mOut_left,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(1, 1, None),
+        )
+        gOut_right = cute.local_tile(
+            mOut_right,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(1, 1, None),
+        )
+
+        thr_copy_A = tiled_copy_A.get_slice(tidx)
+        thr_copy_B = tiled_copy_B.get_slice(tidx)
+        tAgA_left = thr_copy_A.partition_S(gA_left)
+        tAgA_right = thr_copy_A.partition_S(gA_right)
+        tAsA_left = thr_copy_A.partition_D(sA_left)
+        tAsA_right = thr_copy_A.partition_D(sA_right)
+        tBgB_left = thr_copy_B.partition_S(gB_left)
+        tBgB_right = thr_copy_B.partition_S(gB_right)
+        tBsB_left = thr_copy_B.partition_D(sB_left)
+        tBsB_right = thr_copy_B.partition_D(sB_right)
+
+        cA = cute.local_tile(
+            cute.make_identity_tensor(mA_left.shape),
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(1, None, 1),
+        )
+        tAcA = thr_copy_A.partition_S(cA)
+        tApA = cute.make_rmem_tensor(
+            cute.make_layout(
+                (
+                    tAsA_left.shape[0][1],
+                    cute.size(tAsA_left, mode=[1]),
+                    cute.size(tAsA_left, mode=[2]),
+                ),
+                stride=(cute.size(tAsA_left, mode=[1]), 1, 0),
+            ),
+            cutlass.Boolean,
+        )
+        for rest_v in range(tApA.shape[0]):
+            for row in range(tApA.shape[1]):
+                tApA[rest_v, row, 0] = cute.elem_less(
+                    tAcA[(0, rest_v), row, 0, 0][0],
+                    PACKED_COEFF_DIM,
+                )
+
+        k_tile_count = cute.size(tAgA_left, mode=[3])
+        gmem_pipe_read = cutlass.Int32(0)
+        cute.copy(
+            tiled_copy_A,
+            tAgA_left[None, None, None, gmem_pipe_read],
+            tAsA_left[None, None, None, 0],
+            pred=tApA,
+        )
+        cute.copy(
+            tiled_copy_A,
+            tAgA_right[None, None, None, gmem_pipe_read],
+            tAsA_right[None, None, None, 0],
+            pred=tApA,
+        )
+        cute.copy(
+            tiled_copy_B,
+            tBgB_left[None, None, None, gmem_pipe_read],
+            tBsB_left[None, None, None, 0],
+        )
+        cute.copy(
+            tiled_copy_B,
+            tBgB_right[None, None, None, gmem_pipe_read],
+            tBsB_right[None, None, None, 0],
+        )
+        cute.arch.cp_async_commit_group()
+        gmem_pipe_read = gmem_pipe_read + 1
+        for stage in range(1, STAGES - 1):
+            cute.copy(
+                tiled_copy_A,
+                tAgA_left[None, None, None, gmem_pipe_read],
+                tAsA_left[None, None, None, stage],
+                pred=tApA,
+            )
+            cute.copy(
+                tiled_copy_A,
+                tAgA_right[None, None, None, gmem_pipe_read],
+                tAsA_right[None, None, None, stage],
+                pred=tApA,
+            )
+            cute.copy(
+                tiled_copy_B,
+                tBgB_left[None, None, None, gmem_pipe_read],
+                tBsB_left[None, None, None, stage],
+            )
+            cute.copy(
+                tiled_copy_B,
+                tBgB_right[None, None, None, gmem_pipe_read],
+                tBsB_right[None, None, None, stage],
+            )
+            cute.arch.cp_async_commit_group()
+            gmem_pipe_read = gmem_pipe_read + 1
+
+        tCsA_left = thr_mma.partition_A(sA_left)
+        tCsA_right = thr_mma.partition_A(sA_right)
+        tCsB_left = thr_mma.partition_B(sB_left)
+        tCsB_right = thr_mma.partition_B(sB_right)
+        tCgOut_left = thr_mma.partition_C(gOut_left)
+        tCgOut_right = thr_mma.partition_C(gOut_right)
+        tCrA_left = tiled_mma.make_fragment_A(tCsA_left[None, None, None, 0])
+        tCrA_right = tiled_mma.make_fragment_A(tCsA_right[None, None, None, 0])
+        tCrB_left = tiled_mma.make_fragment_B(tCsB_left[None, None, None, 0])
+        tCrB_right = tiled_mma.make_fragment_B(tCsB_right[None, None, None, 0])
+        tCrOut_left = tiled_mma.make_fragment_C(tCgOut_left)
+        tCrOut_right = tiled_mma.make_fragment_C(tCgOut_right)
+        tCrOut_left.fill(0.0)
+        tCrOut_right.fill(0.0)
+
+        smem_pipe_read = cutlass.Int32(0)
+        smem_pipe_write = cutlass.Int32(STAGES - 1)
+        tiles_issued = cutlass.Int32(STAGES - 1)
+        tCsA_left_p = tCsA_left[None, None, None, smem_pipe_read]
+        tCsA_right_p = tCsA_right[None, None, None, smem_pipe_read]
+        tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read]
+        tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read]
+        k_block_max = cute.size(tCrA_left, mode=[2])
+        if k_block_max > 1:
+            cute.arch.cp_async_wait_group(STAGES - 2)
+            self.cta_sync_barrier.arrive_and_wait()
+            cute.autovec_copy(
+                tCsA_left_p[None, None, 0],
+                tCrA_left[None, None, 0],
+            )
+            cute.autovec_copy(
+                tCsA_right_p[None, None, 0],
+                tCrA_right[None, None, 0],
+            )
+            cute.autovec_copy(
+                tCsB_left_p[None, None, 0],
+                tCrB_left[None, None, 0],
+            )
+            cute.autovec_copy(
+                tCsB_right_p[None, None, 0],
+                tCrB_right[None, None, 0],
+            )
+
+        for _ in range(k_tile_count):
+            for k_block in range(k_block_max, unroll_full=True):
+                if k_block == k_block_max - 1:
+                    tCsA_left_p = tCsA_left[None, None, None, smem_pipe_read]
+                    tCsA_right_p = tCsA_right[None, None, None, smem_pipe_read]
+                    tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read]
+                    tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read]
+                    cute.arch.cp_async_wait_group(STAGES - 2)
+                    self.cta_sync_barrier.arrive_and_wait()
+                k_block_next = (k_block + 1) % k_block_max
+                cute.autovec_copy(
+                    tCsA_left_p[None, None, k_block_next],
+                    tCrA_left[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsA_right_p[None, None, k_block_next],
+                    tCrA_right[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsB_left_p[None, None, k_block_next],
+                    tCrB_left[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsB_right_p[None, None, k_block_next],
+                    tCrB_right[None, None, k_block_next],
+                )
+                if k_block == 0 and tiles_issued < k_tile_count:
+                    cute.copy(
+                        tiled_copy_A,
+                        tAgA_left[None, None, None, gmem_pipe_read],
+                        tAsA_left[None, None, None, smem_pipe_write],
+                        pred=tApA,
+                    )
+                    cute.copy(
+                        tiled_copy_A,
+                        tAgA_right[None, None, None, gmem_pipe_read],
+                        tAsA_right[None, None, None, smem_pipe_write],
+                        pred=tApA,
+                    )
+                    cute.copy(
+                        tiled_copy_B,
+                        tBgB_left[None, None, None, gmem_pipe_read],
+                        tBsB_left[None, None, None, smem_pipe_write],
+                    )
+                    cute.copy(
+                        tiled_copy_B,
+                        tBgB_right[None, None, None, gmem_pipe_read],
+                        tBsB_right[None, None, None, smem_pipe_write],
+                    )
+                cute.gemm(
+                    tiled_mma,
+                    tCrOut_left,
+                    tCrA_left[None, None, k_block],
+                    tCrB_left[None, None, k_block],
+                    tCrOut_left,
+                )
+                cute.gemm(
+                    tiled_mma,
+                    tCrOut_right,
+                    tCrA_right[None, None, k_block],
+                    tCrB_right[None, None, k_block],
+                    tCrOut_right,
+                )
+                if k_block == 0:
+                    cute.arch.cp_async_commit_group()
+                    tiles_issued = tiles_issued + 1
+                    smem_pipe_write = smem_pipe_read
+                    smem_pipe_read = smem_pipe_read + 1
+                    if smem_pipe_read == STAGES:
+                        smem_pipe_read = cutlass.Int32(0)
+                    gmem_pipe_read = (
+                        gmem_pipe_read + 1
+                        if gmem_pipe_read + 1 < k_tile_count
+                        else cutlass.Int32(0)
+                    )
+
+        cute.arch.cp_async_wait_group(0)
+        self.cta_sync_barrier.arrive_and_wait()
+        cOut = cute.make_identity_tensor(gOut_left.shape)
+        tCpOut = thr_mma.partition_C(cOut)
+        pred = cute.make_rmem_tensor(tCrOut_left.layout, cutlass.Boolean)
+        for idx in range(cute.size(tCrOut_left.shape)):
+            pred[idx] = cute.elem_less(
+                tCpOut[idx],
+                (PACKED_COEFF_DIM, TILE_N),
+            )
+            if pred[idx]:
+                local_channel = tCpOut[idx][1]
+                scale = sDq0[local_channel].to(cutlass.Float32)
+                tCrOut_left[idx] = tCrOut_left[idx].to(cutlass.Float32) * scale
+                tCrOut_right[idx] = tCrOut_right[idx].to(cutlass.Float32) * scale
+        atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mOut_left.element_type,
+        )
+        cute.copy(atom, tCrOut_left, tCgOut_left, pred=pred)
+        cute.copy(atom, tCrOut_right, tCgOut_right, pred=pred)
+
+
+def compile_readout_l0_gram_forward(
+    device_index: int | None = None,
+    compute_capability: tuple[int, int] | None = None,
+) -> Callable:
+    """Compile the dense-Gram forward with symbolic runtime node count."""
+    import torch
+
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    actual_capability = tuple(torch.cuda.get_device_capability(device_index))
+    if (
+        compute_capability is not None
+        and tuple(compute_capability) != actual_capability
+    ):
+        raise ValueError("compile target does not match the selected CUDA device")
+    with torch.cuda.device(device_index):
+        nodes = cute.sym_int64()
+        fake_coeff = make_fake_compact_tensor(
+            cutlass.Float32,
+            (nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS),
+            stride_order=(2, 1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_gram = make_fake_compact_tensor(
+            cutlass.Float32,
+            (PACKED_COEFF_DIM, PACKED_COEFF_DIM),
+            stride_order=(1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_out = make_fake_compact_tensor(
+            cutlass.Float32,
+            (nodes, HIDDEN_CHANNELS),
+            stride_order=(1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+        return cute.compile(
+            TiledReadoutL0GramForward(),
+            fake_coeff,
+            fake_coeff,
+            fake_gram,
+            fake_out,
+            fake_stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+def compile_readout_l0_gram_backward(
+    device_index: int | None = None,
+    compute_capability: tuple[int, int] | None = None,
+) -> Callable:
+    """Compile the dense-Gram first-backward artifact."""
+    import torch
+
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    actual_capability = tuple(torch.cuda.get_device_capability(device_index))
+    if (
+        compute_capability is not None
+        and tuple(compute_capability) != actual_capability
+    ):
+        raise ValueError("compile target does not match the selected CUDA device")
+    with torch.cuda.device(device_index):
+        nodes = cute.sym_int64()
+        fake_coeff = make_fake_compact_tensor(
+            cutlass.Float32,
+            (nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS),
+            stride_order=(2, 1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_q0 = make_fake_compact_tensor(
+            cutlass.Float32,
+            (nodes, HIDDEN_CHANNELS),
+            stride_order=(1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_gram = make_fake_compact_tensor(
+            cutlass.Float32,
+            (PACKED_COEFF_DIM, PACKED_COEFF_DIM),
+            stride_order=(1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+        return cute.compile(
+            TiledReadoutL0GramBackward(),
+            fake_q0,
+            fake_coeff,
+            fake_coeff,
+            fake_gram,
+            fake_coeff,
+            fake_coeff,
+            fake_stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+@lru_cache(maxsize=16)
+def _compiled_readout_l0_gram_forward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> Callable:
+    return compile_readout_l0_gram_forward(device_index, compute_capability)
+
+
+@lru_cache(maxsize=16)
+def _compiled_readout_l0_gram_backward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> Callable:
+    return compile_readout_l0_gram_backward(device_index, compute_capability)
+
+
+def _compile_key(tensor) -> tuple[int, tuple[int, int]]:
+    import torch
+
+    device_index = tensor.device.index
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    compute_capability = tuple(torch.cuda.get_device_capability(device_index))
+    return int(device_index), compute_capability
+
+
+def _validate_tensors(left, right, to_grid, from_grid, out=None) -> None:
+    import torch
+
+    tensors = (left, right, to_grid, from_grid)
+    if (
+        any(not tensor.is_cuda for tensor in tensors)
+        or any(tensor.dtype != torch.float32 for tensor in tensors)
+        or any(tensor.device != left.device for tensor in tensors)
+        or left.ndim != 3
+        or left.shape[0] <= 0
+        or tuple(left.shape[1:]) != (PACKED_COEFF_DIM, HIDDEN_CHANNELS)
+        or right.shape != left.shape
+        or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM)
+        or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE)
+        or any(not tensor.is_contiguous() for tensor in tensors)
+        or any(tensor.data_ptr() % 16 != 0 for tensor in tensors)
+        or torch.cuda.get_device_capability(left.device)[0] < 8
+    ):
+        raise ValueError(
+            "readout l=0 requires contiguous CUDA FP32 left/right=(N,48,192), "
+            "to_grid=(152,48), and from_grid=(48,152) tensors on compute "
+            "capability 8.0+"
+        )
+    if out is not None and (
+        tuple(out.shape) != (left.shape[0], HIDDEN_CHANNELS)
+        or out.device != left.device
+        or out.dtype != left.dtype
+        or not out.is_contiguous()
+        or out.data_ptr() % 16 != 0
+    ):
+        raise ValueError("readout l=0 output must be contiguous with shape (N,192)")
+
+
+def _validate_gram_tensors(left, right, gram, out=None) -> None:
+    import torch
+
+    tensors = (left, right, gram)
+    if (
+        any(not tensor.is_cuda for tensor in tensors)
+        or any(tensor.dtype != torch.float32 for tensor in tensors)
+        or any(tensor.device != left.device for tensor in tensors)
+        or left.ndim != 3
+        or left.shape[0] <= 0
+        or tuple(left.shape[1:]) != (PACKED_COEFF_DIM, HIDDEN_CHANNELS)
+        or right.shape != left.shape
+        or tuple(gram.shape) != (PACKED_COEFF_DIM, PACKED_COEFF_DIM)
+        or any(not tensor.is_contiguous() for tensor in tensors)
+        or any(tensor.data_ptr() % 16 != 0 for tensor in tensors)
+        or torch.cuda.get_device_capability(left.device)[0] < 8
+    ):
+        raise ValueError(
+            "Gram readout l=0 requires contiguous CUDA FP32 left/right="
+            "(N,48,192) and gram=(48,48) tensors on compute capability 8.0+"
+        )
+    if out is not None and (
+        tuple(out.shape) != (left.shape[0], HIDDEN_CHANNELS)
+        or out.device != left.device
+        or out.dtype != left.dtype
+        or not out.is_contiguous()
+        or out.data_ptr() % 16 != 0
+    ):
+        raise ValueError("readout l=0 output must be contiguous with shape (N,192)")
+
+
+def run_readout_l0_gram(left, right, gram):
+    """Run ``left[:, :, h]^T G right[:, :, h]`` in strict FP32."""
+    import torch
+
+    _validate_gram_tensors(left, right, gram)
+    out = torch.empty(
+        (left.shape[0], HIDDEN_CHANNELS),
+        dtype=left.dtype,
+        device=left.device,
+    )
+    _compiled_readout_l0_gram_forward(*_compile_key(left))(
+        left,
+        right,
+        gram,
+        out,
+    )
+    return out
+
+
+def run_readout_l0(left, right, to_grid, from_grid):
+    """Run the strict-FP32 row-zero readout forward."""
+    _validate_tensors(left, right, to_grid, from_grid)
+    from ..readout_l0 import (
+        build_readout_l0_gram,
+    )
+
+    return run_readout_l0_gram(
+        left,
+        right,
+        build_readout_l0_gram(to_grid, from_grid),
+    )
+
+
+def run_readout_l0_backward(dq0, left, right, to_grid, from_grid):
+    """Run first backward and return full `(N,48,192)` input adjoints."""
+    import torch
+
+    _validate_tensors(left, right, to_grid, from_grid)
+    if (
+        tuple(dq0.shape) != (left.shape[0], HIDDEN_CHANNELS)
+        or dq0.device != left.device
+        or dq0.dtype != torch.float32
+        or not dq0.is_contiguous()
+    ):
+        raise ValueError("dq0 must be contiguous CUDA FP32 with shape (N,192)")
+    grad_left = torch.empty_like(left)
+    grad_right = torch.empty_like(right)
+    compile_key = _compile_key(left)
+    from ..readout_l0 import (
+        build_readout_l0_gram,
+    )
+
+    gram = build_readout_l0_gram(to_grid, from_grid)
+    _compiled_readout_l0_gram_backward(*compile_key)(
+        dq0,
+        left,
+        right,
+        gram,
+        grad_left,
+        grad_right,
+    )
+    return grad_left, grad_right
+
+
+__all__ = [
+    "TiledReadoutL0GramBackward",
+    "TiledReadoutL0GramForward",
+    "run_readout_l0",
+    "run_readout_l0_backward",
+    "run_readout_l0_gram",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/tiled_product.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/tiled_product.py
new file mode 100644
index 0000000000..ccdaf239f7
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/tiled_product.py
@@ -0,0 +1,2915 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Tiled strict-FP32 Neo output-grid product forward and first backward."""
+
+from __future__ import (
+    annotations,
+)
+
+from collections.abc import (
+    Callable,
+)
+from functools import (
+    lru_cache,
+)
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ... import (
+    runtime_policy,
+)
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN201, ANN202, ANN204, TC002, TC003
+
+
+PACKED_COEFF_DIM = 48
+GRID_SIZE = 152
+SUPPORTED_HIDDEN_CHANNELS = (96, 192)
+TILE_M = 64
+TILE_N = 64
+TILE_K = 8
+C96_TAIL_TILE_N = 32
+C96_TAIL_CHANNEL_TILE = 2
+SM80_C96_TILE_N = 48
+SM80_C96_THREADS = 128
+MMA_ATOMS_N = 16
+THREADS = 128
+STAGES = 3
+GRID_TILES = 3
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+
+def _validate_hidden_channels(hidden_channels: int) -> int:
+    hidden_channels = int(hidden_channels)
+    if hidden_channels not in SUPPORTED_HIDDEN_CHANNELS:
+        raise ValueError(
+            "tiled output-grid channel width must be one of "
+            f"{SUPPORTED_HIDDEN_CHANNELS}, got {hidden_channels}"
+        )
+    return hidden_channels
+
+
+class TiledOutputGridProductForward:
+    """Dual tiled projections, shared product, and tiled backprojection."""
+
+    def __init__(
+        self,
+        hidden_channels: int = 192,
+        *,
+        tile_n: int = TILE_N,
+        channel_tile_start: int = 0,
+        channel_tile_count: int | None = None,
+    ) -> None:
+        self.hidden_channels = _validate_hidden_channels(hidden_channels)
+        tile_n = int(tile_n)
+        if tile_n not in (C96_TAIL_TILE_N, SM80_C96_TILE_N, TILE_N):
+            raise ValueError("output-grid forward tile_n must be 32, 48, or 64")
+        if tile_n == SM80_C96_TILE_N and self.hidden_channels != 96:
+            raise ValueError("output-grid forward N=48 specializes C=96")
+        channel_tile_start = int(channel_tile_start)
+        total_channel_tiles = (self.hidden_channels + tile_n - 1) // tile_n
+        if channel_tile_count is None:
+            channel_tile_count = total_channel_tiles - channel_tile_start
+        channel_tile_count = int(channel_tile_count)
+        if (
+            channel_tile_start < 0
+            or channel_tile_count <= 0
+            or channel_tile_start + channel_tile_count > total_channel_tiles
+        ):
+            raise ValueError("invalid output-grid forward channel-tile range")
+        if tile_n == C96_TAIL_TILE_N and (
+            self.hidden_channels != 96
+            or channel_tile_start != C96_TAIL_CHANNEL_TILE
+            or channel_tile_count != 1
+        ):
+            raise ValueError("output-grid forward N=32 specializes the C96 tail panel")
+        self.cta_tiler = (TILE_M, tile_n, TILE_K)
+        self.channel_tile_start = channel_tile_start
+        self.channel_tiles = channel_tile_count
+        self.has_channel_residue = self.hidden_channels % tile_n != 0
+        self.cta_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=THREADS,
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        to_grid: cute.Tensor,
+        from_grid: cute.Tensor,
+        out: cute.Tensor,
+        stream: CUstream,
+    ):
+        tile_n = self.cta_tiler[1]
+        sA_layout = cute.make_layout(
+            (TILE_M, TILE_K, STAGES),
+            stride=(1, TILE_M + 4, TILE_K * (TILE_M + 4)),
+        )
+        sB_layout = cute.make_layout(
+            (tile_n, TILE_K, STAGES),
+            stride=(1, tile_n, TILE_K * tile_n),
+        )
+        product_layout = cute.make_layout(
+            (GRID_SIZE, tile_n),
+            stride=(tile_n, 1),
+        )
+        copy_a_atom = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            left.element_type,
+            num_bits_per_copy=left.element_type.width,
+        )
+        copy_a_layout = cute.make_layout(
+            (THREADS // TILE_K, TILE_K),
+            stride=(TILE_K, 1),
+        )
+        tiled_copy_A = cute.make_tiled_copy_tv(
+            copy_a_atom,
+            copy_a_layout,
+            cute.make_layout((1, 1)),
+        )
+        vector = 2 if cutlass.const_expr(tile_n == C96_TAIL_TILE_N) else 4
+        if cutlass.const_expr(tile_n == SM80_C96_TILE_N):
+            copy_b_atom = cute.make_copy_atom(
+                cute.nvgpu.cpasync.CopyG2SOp(),
+                left.element_type,
+                num_bits_per_copy=left.element_type.width,
+            )
+            copy_b_major = THREADS // TILE_K
+            copy_b_layout = cute.make_layout(
+                (copy_b_major, TILE_K),
+                stride=(1, copy_b_major),
+            )
+            copy_b_value_layout = cute.make_layout(
+                (tile_n // copy_b_major, 1),
+            )
+        else:
+            copy_b_atom = cute.make_copy_atom(
+                cute.nvgpu.cpasync.CopyG2SOp(),
+                left.element_type,
+                num_bits_per_copy=left.element_type.width * vector,
+            )
+            copy_b_major = tile_n // vector
+            copy_b_layout = cute.make_layout(
+                (copy_b_major, THREADS // copy_b_major),
+                stride=(1, copy_b_major),
+            )
+            copy_b_value_layout = cute.make_layout((vector, 1))
+        tiled_copy_B = cute.make_tiled_copy_tv(
+            copy_b_atom,
+            copy_b_layout,
+            copy_b_value_layout,
+        )
+        atoms_layout = cute.make_layout(
+            (THREADS // 16, 16, 1),
+            stride=(16, 1, 0),
+        )
+        permutation_m = cute.make_layout(
+            (atoms_layout.shape[0], 4),
+            stride=(4, 1),
+        )
+        values_n = tile_n // MMA_ATOMS_N
+        permutation_n = cute.make_layout(
+            (atoms_layout.shape[1], values_n),
+            stride=(values_n, 1),
+        )
+        tiled_mma = cute.make_tiled_mma(
+            cute.nvgpu.MmaUniversalOp(cutlass.Float32),
+            atoms_layout,
+            permutation_mnk=(permutation_m, permutation_n, None),
+        )
+        self.kernel(
+            left,
+            right,
+            to_grid,
+            from_grid,
+            out,
+            sA_layout,
+            sB_layout,
+            product_layout,
+            tiled_copy_A,
+            tiled_copy_B,
+            tiled_mma,
+        ).launch(
+            grid=(left.shape[0], self.channel_tiles, 1),
+            block=[THREADS, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        to_grid: cute.Tensor,
+        from_grid: cute.Tensor,
+        out: cute.Tensor,
+        sA_layout: cute.Layout,
+        sB_layout: cute.Layout,
+        product_layout: cute.Layout,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        node, channel_tile, _ = cute.arch.block_idx()
+        channel_tile = channel_tile + self.channel_tile_start
+
+        left_node = left[node, None, None]
+        right_node = right[node, None, None]
+        out_node = out[node, None, None]
+        matrix_b_layout = cute.make_layout(
+            (self.hidden_channels, PACKED_COEFF_DIM),
+            stride=(1, self.hidden_channels),
+        )
+        left_b = cute.make_tensor(left_node.iterator, matrix_b_layout)
+        right_b = cute.make_tensor(right_node.iterator, matrix_b_layout)
+
+        smem = cutlass.utils.SmemAllocator()
+        sA = smem.allocate_tensor(cutlass.Float32, sA_layout, 16)
+        sB_left = smem.allocate_tensor(cutlass.Float32, sB_layout, 16)
+        sB_right = smem.allocate_tensor(cutlass.Float32, sB_layout, 16)
+        product = smem.allocate_tensor(cutlass.Float32, product_layout, 16)
+
+        for grid_tile in cutlass.range_constexpr(GRID_TILES):
+            self._dual_projection_product(
+                to_grid,
+                left_b,
+                right_b,
+                product,
+                sA,
+                sB_left,
+                sB_right,
+                tiled_copy_A,
+                tiled_copy_B,
+                tiled_mma,
+                tidx,
+                grid_tile,
+                channel_tile,
+            )
+
+        product_b = cute.make_tensor(
+            product.iterator,
+            cute.make_layout(
+                (self.cta_tiler[1], GRID_SIZE),
+                stride=(1, self.cta_tiler[1]),
+            ),
+        )
+        self._backproject(
+            from_grid,
+            product_b,
+            out_node,
+            sA,
+            tiled_copy_A,
+            tiled_mma,
+            tidx,
+            channel_tile,
+        )
+
+    @cute.jit
+    def _dual_projection_product(
+        self,
+        mA: cute.Tensor,
+        mB_left: cute.Tensor,
+        mB_right: cute.Tensor,
+        mProduct: cute.Tensor,
+        sA: cute.Tensor,
+        sB_left: cute.Tensor,
+        sB_right: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        grid_tile: cutlass.Constexpr,
+        channel_tile: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        gA = cute.local_tile(
+            mA,
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, None, 1),
+        )
+        gB_left = cute.local_tile(
+            mB_left,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(None, 1, 1),
+        )
+        gB_right = cute.local_tile(
+            mB_right,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(None, 1, 1),
+        )
+        gProduct = cute.local_tile(
+            mProduct,
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, 1, None),
+        )
+
+        thr_copy_A = tiled_copy_A.get_slice(tidx)
+        thr_copy_B = tiled_copy_B.get_slice(tidx)
+        tAgA = thr_copy_A.partition_S(gA)
+        tAsA = thr_copy_A.partition_D(sA)
+        tBgB_left = thr_copy_B.partition_S(gB_left)
+        tBgB_right = thr_copy_B.partition_S(gB_right)
+        tBsB_left = thr_copy_B.partition_D(sB_left)
+        tBsB_right = thr_copy_B.partition_D(sB_right)
+
+        if cutlass.const_expr(self.has_channel_residue):
+            cB = cute.local_tile(
+                cute.make_identity_tensor(mB_left.shape),
+                tiler=self.cta_tiler,
+                coord=(0, channel_tile, None),
+                proj=(None, 1, 1),
+            )
+            tBcB = thr_copy_B.partition_S(cB)
+            tBpB = cute.make_rmem_tensor(
+                cute.make_layout(
+                    (
+                        tBsB_left.shape[0][1],
+                        cute.size(tBsB_left, mode=[1]),
+                        cute.size(tBsB_left, mode=[2]),
+                    ),
+                    stride=(cute.size(tBsB_left, mode=[1]), 1, 0),
+                ),
+                cutlass.Boolean,
+            )
+            for rest_v in range(tBpB.shape[0]):
+                for channel in range(tBpB.shape[1]):
+                    tBpB[rest_v, channel, 0] = cute.elem_less(
+                        tBcB[(0, rest_v), channel, 0, 0][0],
+                        mB_left.shape[0],
+                    )
+
+        cA = cute.local_tile(
+            cute.make_identity_tensor(mA.shape),
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, None, 1),
+        )
+        tAcA = thr_copy_A.partition_S(cA)
+        tApA = cute.make_rmem_tensor(
+            cute.make_layout(
+                (
+                    tAsA.shape[0][1],
+                    cute.size(tAsA, mode=[1]),
+                    cute.size(tAsA, mode=[2]),
+                ),
+                stride=(cute.size(tAsA, mode=[1]), 1, 0),
+            ),
+            cutlass.Boolean,
+        )
+        for rest_v in range(tApA.shape[0]):
+            for row in range(tApA.shape[1]):
+                tApA[rest_v, row, 0] = cute.elem_less(
+                    tAcA[(0, rest_v), row, 0, 0][0],
+                    mA.shape[0],
+                )
+
+        k_pipe_max = cute.size(tAsA, mode=[3])
+        k_tile_count = cute.size(tAgA, mode=[3])
+        gmem_pipe_read = cutlass.Int32(0)
+        cute.copy(
+            tiled_copy_A,
+            tAgA[None, None, None, gmem_pipe_read],
+            tAsA[None, None, None, 0],
+            pred=tApA,
+        )
+        if cutlass.const_expr(self.has_channel_residue):
+            cute.copy(
+                tiled_copy_B,
+                tBgB_left[None, None, None, gmem_pipe_read],
+                tBsB_left[None, None, None, 0],
+                pred=tBpB,
+            )
+        else:
+            cute.copy(
+                tiled_copy_B,
+                tBgB_left[None, None, None, gmem_pipe_read],
+                tBsB_left[None, None, None, 0],
+            )
+        if cutlass.const_expr(self.has_channel_residue):
+            cute.copy(
+                tiled_copy_B,
+                tBgB_right[None, None, None, gmem_pipe_read],
+                tBsB_right[None, None, None, 0],
+                pred=tBpB,
+            )
+        else:
+            cute.copy(
+                tiled_copy_B,
+                tBgB_right[None, None, None, gmem_pipe_read],
+                tBsB_right[None, None, None, 0],
+            )
+        cute.arch.cp_async_commit_group()
+        gmem_pipe_read = gmem_pipe_read + 1
+        for stage in range(1, STAGES - 1):
+            cute.copy(
+                tiled_copy_A,
+                tAgA[None, None, None, gmem_pipe_read],
+                tAsA[None, None, None, stage],
+                pred=tApA,
+            )
+            if cutlass.const_expr(self.has_channel_residue):
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB_left[None, None, None, gmem_pipe_read],
+                    tBsB_left[None, None, None, stage],
+                    pred=tBpB,
+                )
+            else:
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB_left[None, None, None, gmem_pipe_read],
+                    tBsB_left[None, None, None, stage],
+                )
+            if cutlass.const_expr(self.has_channel_residue):
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB_right[None, None, None, gmem_pipe_read],
+                    tBsB_right[None, None, None, stage],
+                    pred=tBpB,
+                )
+            else:
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB_right[None, None, None, gmem_pipe_read],
+                    tBsB_right[None, None, None, stage],
+                )
+            cute.arch.cp_async_commit_group()
+            gmem_pipe_read = gmem_pipe_read + 1
+
+        tCsA = thr_mma.partition_A(sA)
+        tCsB_left = thr_mma.partition_B(sB_left)
+        tCsB_right = thr_mma.partition_B(sB_right)
+        tCgProduct = thr_mma.partition_C(gProduct)
+        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
+        tCrB_left = tiled_mma.make_fragment_B(tCsB_left[None, None, None, 0])
+        tCrB_right = tiled_mma.make_fragment_B(tCsB_right[None, None, None, 0])
+        tCrLeft = tiled_mma.make_fragment_C(tCgProduct)
+        tCrRight = tiled_mma.make_fragment_C(tCgProduct)
+        tCrLeft.fill(0.0)
+        tCrRight.fill(0.0)
+
+        smem_pipe_read = cutlass.Int32(0)
+        smem_pipe_write = cutlass.Int32(STAGES - 1)
+        tiles_issued = cutlass.Int32(STAGES - 1)
+        tCsA_p = tCsA[None, None, None, smem_pipe_read]
+        tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read]
+        tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read]
+        k_block_max = cute.size(tCrA, mode=[2])
+        if k_block_max > 1:
+            cute.arch.cp_async_wait_group(STAGES - 2)
+            self.cta_sync_barrier.arrive_and_wait()
+            cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
+            cute.autovec_copy(
+                tCsB_left_p[None, None, 0],
+                tCrB_left[None, None, 0],
+            )
+            cute.autovec_copy(
+                tCsB_right_p[None, None, 0],
+                tCrB_right[None, None, 0],
+            )
+
+        for _ in range(k_tile_count):
+            for k_block in range(k_block_max, unroll_full=True):
+                if k_block == k_block_max - 1:
+                    tCsA_p = tCsA[None, None, None, smem_pipe_read]
+                    tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read]
+                    tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read]
+                    cute.arch.cp_async_wait_group(STAGES - 2)
+                    self.cta_sync_barrier.arrive_and_wait()
+
+                k_block_next = (k_block + 1) % k_block_max
+                cute.autovec_copy(
+                    tCsA_p[None, None, k_block_next],
+                    tCrA[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsB_left_p[None, None, k_block_next],
+                    tCrB_left[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsB_right_p[None, None, k_block_next],
+                    tCrB_right[None, None, k_block_next],
+                )
+                if k_block == 0 and tiles_issued < k_tile_count:
+                    cute.copy(
+                        tiled_copy_A,
+                        tAgA[None, None, None, gmem_pipe_read],
+                        tAsA[None, None, None, smem_pipe_write],
+                        pred=tApA,
+                    )
+
+                cute.gemm(
+                    tiled_mma,
+                    tCrLeft,
+                    tCrA[None, None, k_block],
+                    tCrB_left[None, None, k_block],
+                    tCrLeft,
+                )
+                cute.gemm(
+                    tiled_mma,
+                    tCrRight,
+                    tCrA[None, None, k_block],
+                    tCrB_right[None, None, k_block],
+                    tCrRight,
+                )
+
+                if k_block == 0:
+                    if tiles_issued < k_tile_count:
+                        if cutlass.const_expr(self.has_channel_residue):
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB_left[None, None, None, gmem_pipe_read],
+                                tBsB_left[None, None, None, smem_pipe_write],
+                                pred=tBpB,
+                            )
+                        else:
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB_left[None, None, None, gmem_pipe_read],
+                                tBsB_left[None, None, None, smem_pipe_write],
+                            )
+                        if cutlass.const_expr(self.has_channel_residue):
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB_right[None, None, None, gmem_pipe_read],
+                                tBsB_right[None, None, None, smem_pipe_write],
+                                pred=tBpB,
+                            )
+                        else:
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB_right[None, None, None, gmem_pipe_read],
+                                tBsB_right[None, None, None, smem_pipe_write],
+                            )
+                    cute.arch.cp_async_commit_group()
+                    tiles_issued = tiles_issued + 1
+                    smem_pipe_write = smem_pipe_read
+                    smem_pipe_read = smem_pipe_read + 1
+                    if smem_pipe_read == STAGES:
+                        smem_pipe_read = cutlass.Int32(0)
+                    gmem_pipe_read = (
+                        gmem_pipe_read + 1
+                        if gmem_pipe_read + 1 < k_tile_count
+                        else cutlass.Int32(0)
+                    )
+
+        cute.arch.cp_async_wait_group(0)
+        self.cta_sync_barrier.arrive_and_wait()
+        cProduct = cute.make_identity_tensor(gProduct.shape)
+        tCpProduct = thr_mma.partition_C(cProduct)
+        pred = cute.make_rmem_tensor(tCrLeft.layout, cutlass.Boolean)
+        residue_m = GRID_SIZE - TILE_M * grid_tile
+        for idx in range(cute.size(tCrLeft.shape)):
+            pred[idx] = cute.elem_less(
+                tCpProduct[idx],
+                (residue_m, self.cta_tiler[1]),
+            )
+            if pred[idx]:
+                tCrLeft[idx] = tCrLeft[idx].to(cutlass.Float32) * tCrRight[idx].to(
+                    cutlass.Float32
+                )
+        atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mProduct.element_type,
+        )
+        cute.copy(atom, tCrLeft, tCgProduct, pred=pred)
+        cute.arch.sync_threads()
+
+    @cute.jit
+    def _backproject(
+        self,
+        mA: cute.Tensor,
+        mB_shared: cute.Tensor,
+        mOut: cute.Tensor,
+        sA: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        channel_tile: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        gA = cute.local_tile(
+            mA,
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(1, None, 1),
+        )
+        sB = cute.local_tile(
+            mB_shared,
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(None, 1, 1),
+        )
+        gOut = cute.local_tile(
+            mOut,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(1, 1, None),
+        )
+        thr_copy_A = tiled_copy_A.get_slice(tidx)
+        tAgA = thr_copy_A.partition_S(gA)
+        tAsA = thr_copy_A.partition_D(sA)
+
+        cA = cute.local_tile(
+            cute.make_identity_tensor(mA.shape),
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(1, None, 1),
+        )
+        tAcA = thr_copy_A.partition_S(cA)
+        tApA = cute.make_rmem_tensor(
+            cute.make_layout(
+                (
+                    tAsA.shape[0][1],
+                    cute.size(tAsA, mode=[1]),
+                    cute.size(tAsA, mode=[2]),
+                ),
+                stride=(cute.size(tAsA, mode=[1]), 1, 0),
+            ),
+            cutlass.Boolean,
+        )
+        for rest_v in range(tApA.shape[0]):
+            for row in range(tApA.shape[1]):
+                tApA[rest_v, row, 0] = cute.elem_less(
+                    tAcA[(0, rest_v), row, 0, 0][0],
+                    PACKED_COEFF_DIM,
+                )
+
+        k_pipe_max = cute.size(tAsA, mode=[3])
+        k_tile_count = cute.size(tAgA, mode=[3])
+        gmem_pipe_read = cutlass.Int32(0)
+        cute.copy(
+            tiled_copy_A,
+            tAgA[None, None, None, gmem_pipe_read],
+            tAsA[None, None, None, 0],
+            pred=tApA,
+        )
+        cute.arch.cp_async_commit_group()
+        gmem_pipe_read = gmem_pipe_read + 1
+        for stage in range(1, STAGES - 1):
+            cute.copy(
+                tiled_copy_A,
+                tAgA[None, None, None, gmem_pipe_read],
+                tAsA[None, None, None, stage],
+                pred=tApA,
+            )
+            cute.arch.cp_async_commit_group()
+            gmem_pipe_read = gmem_pipe_read + 1
+
+        tCsA = thr_mma.partition_A(sA)
+        tSsB = thr_mma.partition_B(sB)
+        tCgOut = thr_mma.partition_C(gOut)
+        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
+        tCrB = tiled_mma.make_fragment_B(tSsB[None, None, None, 0])
+        tCrOut = tiled_mma.make_fragment_C(tCgOut)
+        tCrOut.fill(0.0)
+
+        smem_pipe_read = cutlass.Int32(0)
+        smem_pipe_write = cutlass.Int32(STAGES - 1)
+        tiles_issued = cutlass.Int32(STAGES - 1)
+        logical_k_tile = cutlass.Int32(0)
+        tCsA_p = tCsA[None, None, None, smem_pipe_read]
+        k_block_max = cute.size(tCrA, mode=[2])
+        if k_block_max > 1:
+            cute.arch.cp_async_wait_group(STAGES - 2)
+            self.cta_sync_barrier.arrive_and_wait()
+            cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
+            cute.autovec_copy(
+                tSsB[None, None, 0, logical_k_tile],
+                tCrB[None, None, 0],
+            )
+
+        for _ in range(k_tile_count):
+            for k_block in range(k_block_max, unroll_full=True):
+                if k_block == k_block_max - 1:
+                    tCsA_p = tCsA[None, None, None, smem_pipe_read]
+                    cute.arch.cp_async_wait_group(STAGES - 2)
+                    self.cta_sync_barrier.arrive_and_wait()
+                k_block_next = (k_block + 1) % k_block_max
+                fragment_k_tile = logical_k_tile
+                if k_block_max > 1:
+                    if k_block == k_block_max - 1:
+                        fragment_k_tile = (
+                            logical_k_tile + 1
+                            if logical_k_tile + 1 < k_tile_count
+                            else logical_k_tile
+                        )
+                cute.autovec_copy(
+                    tCsA_p[None, None, k_block_next],
+                    tCrA[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tSsB[None, None, k_block_next, fragment_k_tile],
+                    tCrB[None, None, k_block_next],
+                )
+                if k_block == 0 and tiles_issued < k_tile_count:
+                    cute.copy(
+                        tiled_copy_A,
+                        tAgA[None, None, None, gmem_pipe_read],
+                        tAsA[None, None, None, smem_pipe_write],
+                        pred=tApA,
+                    )
+                cute.gemm(
+                    tiled_mma,
+                    tCrOut,
+                    tCrA[None, None, k_block],
+                    tCrB[None, None, k_block],
+                    tCrOut,
+                )
+                if k_block == 0:
+                    cute.arch.cp_async_commit_group()
+                    tiles_issued = tiles_issued + 1
+                    smem_pipe_write = smem_pipe_read
+                    smem_pipe_read = smem_pipe_read + 1
+                    if smem_pipe_read == STAGES:
+                        smem_pipe_read = cutlass.Int32(0)
+                    gmem_pipe_read = (
+                        gmem_pipe_read + 1
+                        if gmem_pipe_read + 1 < k_tile_count
+                        else cutlass.Int32(0)
+                    )
+            logical_k_tile = logical_k_tile + 1
+
+        cute.arch.cp_async_wait_group(0)
+        self.cta_sync_barrier.arrive_and_wait()
+        pred = cute.make_rmem_tensor(tCrOut.layout, cutlass.Boolean)
+        if cutlass.const_expr(self.has_channel_residue):
+            cOut = cute.local_tile(
+                cute.make_identity_tensor(mOut.shape),
+                tiler=self.cta_tiler,
+                coord=(0, channel_tile, None),
+                proj=(1, 1, None),
+            )
+            tCpOut = thr_mma.partition_C(cOut)
+            for idx in range(cute.size(tCrOut.shape)):
+                pred[idx] = cute.elem_less(tCpOut[idx], mOut.shape)
+        else:
+            cOut = cute.make_identity_tensor(gOut.shape)
+            tCpOut = thr_mma.partition_C(cOut)
+            for idx in range(cute.size(tCrOut.shape)):
+                pred[idx] = cute.elem_less(
+                    tCpOut[idx],
+                    (PACKED_COEFF_DIM, self.cta_tiler[1]),
+                )
+        atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mOut.element_type,
+        )
+        cute.copy(atom, tCrOut, tCgOut, pred=pred)
+
+
+class TiledOutputGridProductBackward:
+    """Tiled dP, projection recomputation, and dual coefficient adjoints."""
+
+    def __init__(
+        self,
+        hidden_channels: int = 192,
+        *,
+        tile_n: int = TILE_N,
+        channel_tile_start: int = 0,
+        channel_tile_count: int | None = None,
+    ) -> None:
+        self.hidden_channels = _validate_hidden_channels(hidden_channels)
+        self.tile_k = TILE_K
+        tile_n = int(tile_n)
+        if tile_n not in (C96_TAIL_TILE_N, SM80_C96_TILE_N, TILE_N):
+            raise ValueError("output-grid backward tile_n must be 32, 48, or 64")
+        if tile_n == SM80_C96_TILE_N and self.hidden_channels != 96:
+            raise ValueError("output-grid N=48 specializes C=96")
+        channel_tile_start = int(channel_tile_start)
+        total_channel_tiles = (self.hidden_channels + tile_n - 1) // tile_n
+        if channel_tile_count is None:
+            channel_tile_count = total_channel_tiles - channel_tile_start
+        channel_tile_count = int(channel_tile_count)
+        if (
+            channel_tile_start < 0
+            or channel_tile_count <= 0
+            or channel_tile_start + channel_tile_count > total_channel_tiles
+        ):
+            raise ValueError("invalid output-grid backward channel-tile range")
+        if tile_n == C96_TAIL_TILE_N and (
+            self.hidden_channels != 96
+            or channel_tile_start != C96_TAIL_CHANNEL_TILE
+            or channel_tile_count != 1
+        ):
+            raise ValueError(
+                "output-grid backward N=32 specializes the C96 K=8 tail panel"
+            )
+        threads = SM80_C96_THREADS if tile_n == SM80_C96_TILE_N else THREADS
+        self.sm80_c96_n48_panel = tile_n == SM80_C96_TILE_N
+        self.cta_tiler = (TILE_M, tile_n, self.tile_k)
+        self.channel_tile_start = channel_tile_start
+        self.channel_tiles = channel_tile_count
+        self.has_channel_residue = self.hidden_channels % tile_n != 0
+        self.cta_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=threads,
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        grad_out: cute.Tensor,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        to_grid: cute.Tensor,
+        from_grid: cute.Tensor,
+        grad_left: cute.Tensor,
+        grad_right: cute.Tensor,
+        stream: CUstream,
+    ):
+        tile_n = self.cta_tiler[1]
+        threads = SM80_C96_THREADS if tile_n == SM80_C96_TILE_N else THREADS
+        sA_row_layout = cute.make_layout(
+            (TILE_M, self.tile_k, STAGES),
+            stride=(1, TILE_M + 4, self.tile_k * (TILE_M + 4)),
+        )
+        sA_col_layout = cute.make_layout(
+            (TILE_M, self.tile_k, STAGES),
+            stride=(1, TILE_M, self.tile_k * TILE_M),
+        )
+        sB_layout = cute.make_layout(
+            (tile_n, self.tile_k, STAGES),
+            stride=(1, tile_n, self.tile_k * tile_n),
+        )
+        grid_layout = cute.make_layout(
+            (GRID_SIZE, tile_n),
+            stride=(tile_n, 1),
+        )
+        panel_layout = cute.make_layout(
+            (TILE_M, tile_n),
+            stride=(tile_n, 1),
+        )
+
+        copy_a_row_atom = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            left.element_type,
+            num_bits_per_copy=left.element_type.width,
+        )
+        copy_a_row_layout = cute.make_layout(
+            (threads // self.tile_k, self.tile_k),
+            stride=(self.tile_k, 1),
+        )
+        tiled_copy_A_row = cute.make_tiled_copy_tv(
+            copy_a_row_atom,
+            copy_a_row_layout,
+            cute.make_layout((1, 1)),
+        )
+
+        vector = 4
+        copy_a_col_atom = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            left.element_type,
+            num_bits_per_copy=left.element_type.width * vector,
+        )
+        copy_a_col_major = TILE_M // vector
+        copy_a_col_layout = cute.make_layout(
+            (copy_a_col_major, threads // copy_a_col_major),
+            stride=(1, copy_a_col_major),
+        )
+        tiled_copy_A_col = cute.make_tiled_copy_tv(
+            copy_a_col_atom,
+            copy_a_col_layout,
+            cute.make_layout((vector, 1)),
+        )
+
+        if cutlass.const_expr(tile_n == SM80_C96_TILE_N):
+            # Keep all 128 threads in the copy/MMA contract. Each thread
+            # issues three scalar N copies, exactly covering 48x8 without a
+            # partial thread layout or an out-of-bounds 64-column staging tile.
+            copy_b_atom = cute.make_copy_atom(
+                cute.nvgpu.cpasync.CopyG2SOp(),
+                left.element_type,
+                num_bits_per_copy=left.element_type.width,
+            )
+            copy_b_major = threads // self.tile_k
+            copy_b_layout = cute.make_layout(
+                (copy_b_major, self.tile_k),
+                stride=(1, copy_b_major),
+            )
+            copy_b_value_layout = cute.make_layout(
+                (tile_n // copy_b_major, 1),
+            )
+        else:
+            vector_b = 2 if cutlass.const_expr(tile_n == C96_TAIL_TILE_N) else vector
+            copy_b_atom = cute.make_copy_atom(
+                cute.nvgpu.cpasync.CopyG2SOp(),
+                left.element_type,
+                num_bits_per_copy=left.element_type.width * vector_b,
+            )
+            copy_b_major = tile_n // vector_b
+            copy_b_layout = cute.make_layout(
+                (copy_b_major, threads // copy_b_major),
+                stride=(1, copy_b_major),
+            )
+            copy_b_value_layout = cute.make_layout((vector_b, 1))
+        tiled_copy_B = cute.make_tiled_copy_tv(
+            copy_b_atom,
+            copy_b_layout,
+            copy_b_value_layout,
+        )
+
+        # Follow the Ampere SGEMM thread topology: the universal-FMA atom is
+        # always tiled over 16 threads in N. N=48 assigns three consecutive N
+        # values to each thread. With 128 threads this produces a 32x48 MMA
+        # tile, which divides the 64x48 CTA exactly; a 96-thread 24x48 tile
+        # would create an unpredicated shared-memory fragment for rows 64..71.
+        atoms_n = MMA_ATOMS_N
+        atoms_m = threads // atoms_n
+        values_n = tile_n // atoms_n
+        atoms_layout = cute.make_layout(
+            (atoms_m, atoms_n, 1),
+            stride=(atoms_n, 1, 0),
+        )
+        permutation_m = cute.make_layout(
+            (atoms_layout.shape[0], 4),
+            stride=(4, 1),
+        )
+        permutation_n = cute.make_layout(
+            (atoms_layout.shape[1], values_n),
+            stride=(values_n, 1),
+        )
+        tiled_mma = cute.make_tiled_mma(
+            cute.nvgpu.MmaUniversalOp(cutlass.Float32),
+            atoms_layout,
+            permutation_mnk=(permutation_m, permutation_n, None),
+        )
+
+        self.kernel(
+            grad_out,
+            left,
+            right,
+            to_grid,
+            from_grid,
+            grad_left,
+            grad_right,
+            sA_row_layout,
+            sA_col_layout,
+            sB_layout,
+            grid_layout,
+            panel_layout,
+            tiled_copy_A_row,
+            tiled_copy_A_col,
+            tiled_copy_B,
+            tiled_mma,
+        ).launch(
+            grid=(left.shape[0], self.channel_tiles, 1),
+            block=[threads, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        grad_out: cute.Tensor,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        to_grid: cute.Tensor,
+        from_grid: cute.Tensor,
+        grad_left: cute.Tensor,
+        grad_right: cute.Tensor,
+        sA_row_layout: cute.Layout,
+        sA_col_layout: cute.Layout,
+        sB_layout: cute.Layout,
+        grid_layout: cute.Layout,
+        panel_layout: cute.Layout,
+        tiled_copy_A_row: cute.TiledCopy,
+        tiled_copy_A_col: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        node, channel_tile, _ = cute.arch.block_idx()
+        channel_tile = channel_tile + self.channel_tile_start
+
+        matrix_b_layout = cute.make_layout(
+            (self.hidden_channels, PACKED_COEFF_DIM),
+            stride=(1, self.hidden_channels),
+        )
+        grad_out_b = cute.make_tensor(
+            grad_out[node, None, None].iterator,
+            matrix_b_layout,
+        )
+        left_b = cute.make_tensor(
+            left[node, None, None].iterator,
+            matrix_b_layout,
+        )
+        right_b = cute.make_tensor(
+            right[node, None, None].iterator,
+            matrix_b_layout,
+        )
+        grad_left_node = grad_left[node, None, None]
+        grad_right_node = grad_right[node, None, None]
+        from_grid_t = cute.make_tensor(
+            from_grid.iterator,
+            cute.make_layout(
+                (GRID_SIZE, PACKED_COEFF_DIM),
+                stride=(1, GRID_SIZE),
+            ),
+        )
+        to_grid_t = cute.make_tensor(
+            to_grid.iterator,
+            cute.make_layout(
+                (PACKED_COEFF_DIM, GRID_SIZE),
+                stride=(1, PACKED_COEFF_DIM),
+            ),
+        )
+
+        smem = cutlass.utils.SmemAllocator()
+        sA_storage = smem.allocate_tensor(
+            cutlass.Float32,
+            sA_row_layout,
+            16,
+        )
+        sA_row = sA_storage
+        sA_col = cute.make_tensor(sA_storage.iterator, sA_col_layout)
+        if cutlass.const_expr(self.sm80_c96_n48_panel):
+            sB = smem.allocate_tensor(cutlass.Float32, sB_layout, 16)
+            adjoint_panel = smem.allocate_tensor(
+                cutlass.Float32,
+                panel_layout,
+                16,
+            )
+            self._panel_adjoint_backward(
+                grad_out_b,
+                left_b,
+                right_b,
+                to_grid,
+                from_grid_t,
+                to_grid_t,
+                grad_left_node,
+                grad_right_node,
+                adjoint_panel,
+                sA_row,
+                sA_col,
+                sB,
+                tiled_copy_A_row,
+                tiled_copy_A_col,
+                tiled_copy_B,
+                tiled_mma,
+                tidx,
+                channel_tile,
+            )
+        else:
+            sB_left = smem.allocate_tensor(cutlass.Float32, sB_layout, 16)
+            sB_right = smem.allocate_tensor(cutlass.Float32, sB_layout, 16)
+            grad_left_grid = smem.allocate_tensor(cutlass.Float32, grid_layout, 16)
+            grad_right_grid = smem.allocate_tensor(cutlass.Float32, grid_layout, 16)
+
+            for grid_tile in cutlass.range_constexpr(GRID_TILES):
+                self._single_projection_to_shared(
+                    from_grid_t,
+                    grad_out_b,
+                    grad_left_grid,
+                    sA_col,
+                    sB_left,
+                    tiled_copy_A_col,
+                    tiled_copy_B,
+                    tiled_mma,
+                    tidx,
+                    grid_tile,
+                    channel_tile,
+                )
+
+            for grid_tile in cutlass.range_constexpr(GRID_TILES):
+                self._dual_projection_adjoint(
+                    to_grid,
+                    left_b,
+                    right_b,
+                    grad_left_grid,
+                    grad_right_grid,
+                    sA_row,
+                    sB_left,
+                    sB_right,
+                    tiled_copy_A_row,
+                    tiled_copy_B,
+                    tiled_mma,
+                    tidx,
+                    grid_tile,
+                    channel_tile,
+                )
+
+            grad_left_grid_b = cute.make_tensor(
+                grad_left_grid.iterator,
+                cute.make_layout(
+                    (self.cta_tiler[1], GRID_SIZE),
+                    stride=(1, self.cta_tiler[1]),
+                ),
+            )
+            grad_right_grid_b = cute.make_tensor(
+                grad_right_grid.iterator,
+                cute.make_layout(
+                    (self.cta_tiler[1], GRID_SIZE),
+                    stride=(1, self.cta_tiler[1]),
+                ),
+            )
+            self._dual_backproject(
+                to_grid_t,
+                grad_left_grid_b,
+                grad_right_grid_b,
+                grad_left_node,
+                grad_right_node,
+                sA_col,
+                tiled_copy_A_col,
+                tiled_mma,
+                tidx,
+                channel_tile,
+            )
+
+    @cute.jit
+    def _panel_adjoint_backward(
+        self,
+        grad_out_b: cute.Tensor,
+        left_b: cute.Tensor,
+        right_b: cute.Tensor,
+        to_grid: cute.Tensor,
+        from_grid_t: cute.Tensor,
+        to_grid_t: cute.Tensor,
+        grad_left_node: cute.Tensor,
+        grad_right_node: cute.Tensor,
+        adjoint_panel: cute.Tensor,
+        sA_row: cute.Tensor,
+        sA_col: cute.Tensor,
+        sB: cute.Tensor,
+        tiled_copy_A_row: cute.TiledCopy,
+        tiled_copy_A_col: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        channel_tile: cutlass.Int32,
+    ):
+        """Keep dP in registers and reuse one shared adjoint panel."""
+        thr_mma = tiled_mma.get_slice(tidx)
+        gOut_left = cute.local_tile(
+            grad_left_node,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(1, 1, None),
+        )
+        gOut_right = cute.local_tile(
+            grad_right_node,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(1, 1, None),
+        )
+        tCgOut_left = thr_mma.partition_C(gOut_left)
+        tCgOut_right = thr_mma.partition_C(gOut_right)
+        tCrOut_left = tiled_mma.make_fragment_C(tCgOut_left)
+        tCrOut_right = tiled_mma.make_fragment_C(tCgOut_right)
+        tCrOut_left.fill(0.0)
+        tCrOut_right.fill(0.0)
+
+        panel_b = cute.make_tensor(
+            adjoint_panel.iterator,
+            cute.make_layout(
+                (self.cta_tiler[1], TILE_M),
+                stride=(1, self.cta_tiler[1]),
+            ),
+        )
+        for grid_tile in cutlass.range_constexpr(GRID_TILES):
+            tCrDP = self._projection_fragment(
+                from_grid_t,
+                grad_out_b,
+                adjoint_panel,
+                sA_col,
+                sB,
+                tiled_copy_A_col,
+                tiled_copy_B,
+                tiled_mma,
+                tidx,
+                grid_tile,
+                channel_tile,
+            )
+
+            tCrRight_grid = self._projection_fragment(
+                to_grid,
+                right_b,
+                adjoint_panel,
+                sA_row,
+                sB,
+                tiled_copy_A_row,
+                tiled_copy_B,
+                tiled_mma,
+                tidx,
+                grid_tile,
+                channel_tile,
+            )
+            self._store_adjoint_panel(
+                tCrDP,
+                tCrRight_grid,
+                adjoint_panel,
+                tiled_mma,
+                tidx,
+                grid_tile,
+                channel_tile,
+            )
+            self._backproject_panel_accumulate(
+                to_grid_t,
+                panel_b,
+                sA_col,
+                tiled_copy_A_col,
+                tiled_mma,
+                tidx,
+                grid_tile,
+                tCrOut_left,
+            )
+
+            tCrLeft_grid = self._projection_fragment(
+                to_grid,
+                left_b,
+                adjoint_panel,
+                sA_row,
+                sB,
+                tiled_copy_A_row,
+                tiled_copy_B,
+                tiled_mma,
+                tidx,
+                grid_tile,
+                channel_tile,
+            )
+            self._store_adjoint_panel(
+                tCrDP,
+                tCrLeft_grid,
+                adjoint_panel,
+                tiled_mma,
+                tidx,
+                grid_tile,
+                channel_tile,
+            )
+            self._backproject_panel_accumulate(
+                to_grid_t,
+                panel_b,
+                sA_col,
+                tiled_copy_A_col,
+                tiled_mma,
+                tidx,
+                grid_tile,
+                tCrOut_right,
+            )
+
+        pred = cute.make_rmem_tensor(tCrOut_left.layout, cutlass.Boolean)
+        if cutlass.const_expr(self.has_channel_residue):
+            cOut = cute.local_tile(
+                cute.make_identity_tensor(grad_left_node.shape),
+                tiler=self.cta_tiler,
+                coord=(0, channel_tile, None),
+                proj=(1, 1, None),
+            )
+            tCpOut = thr_mma.partition_C(cOut)
+            for idx in range(cute.size(tCrOut_left.shape)):
+                pred[idx] = cute.elem_less(tCpOut[idx], grad_left_node.shape)
+        else:
+            cOut = cute.make_identity_tensor(gOut_left.shape)
+            tCpOut = thr_mma.partition_C(cOut)
+            for idx in range(cute.size(tCrOut_left.shape)):
+                pred[idx] = cute.elem_less(
+                    tCpOut[idx],
+                    (PACKED_COEFF_DIM, self.cta_tiler[1]),
+                )
+        atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            grad_left_node.element_type,
+        )
+        cute.copy(atom, tCrOut_left, tCgOut_left, pred=pred)
+        cute.copy(atom, tCrOut_right, tCgOut_right, pred=pred)
+
+    @cute.jit
+    def _projection_fragment(
+        self,
+        mA: cute.Tensor,
+        mB: cute.Tensor,
+        mC_layout: cute.Tensor,
+        sA: cute.Tensor,
+        sB: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        grid_tile: cutlass.Constexpr,
+        channel_tile: cutlass.Int32,
+    ):
+        """Project one 64-row grid panel and return its register fragment."""
+        thr_mma = tiled_mma.get_slice(tidx)
+        gA = cute.local_tile(
+            mA,
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, None, 1),
+        )
+        gB = cute.local_tile(
+            mB,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(None, 1, 1),
+        )
+        thr_copy_A = tiled_copy_A.get_slice(tidx)
+        thr_copy_B = tiled_copy_B.get_slice(tidx)
+        tAgA = thr_copy_A.partition_S(gA)
+        tAsA = thr_copy_A.partition_D(sA)
+        tBgB = thr_copy_B.partition_S(gB)
+        tBsB = thr_copy_B.partition_D(sB)
+
+        if cutlass.const_expr(self.has_channel_residue):
+            cB = cute.local_tile(
+                cute.make_identity_tensor(mB.shape),
+                tiler=self.cta_tiler,
+                coord=(0, channel_tile, None),
+                proj=(None, 1, 1),
+            )
+            tBcB = thr_copy_B.partition_S(cB)
+            tBpB = cute.make_rmem_tensor(
+                cute.make_layout(
+                    (
+                        tBsB.shape[0][1],
+                        cute.size(tBsB, mode=[1]),
+                        cute.size(tBsB, mode=[2]),
+                    ),
+                    stride=(cute.size(tBsB, mode=[1]), 1, 0),
+                ),
+                cutlass.Boolean,
+            )
+            for rest_v in range(tBpB.shape[0]):
+                for channel in range(tBpB.shape[1]):
+                    tBpB[rest_v, channel, 0] = cute.elem_less(
+                        tBcB[(0, rest_v), channel, 0, 0][0],
+                        mB.shape[0],
+                    )
+
+        cA = cute.local_tile(
+            cute.make_identity_tensor(mA.shape),
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, None, 1),
+        )
+        tAcA = thr_copy_A.partition_S(cA)
+        tApA = cute.make_rmem_tensor(
+            cute.make_layout(
+                (
+                    tAsA.shape[0][1],
+                    cute.size(tAsA, mode=[1]),
+                    cute.size(tAsA, mode=[2]),
+                ),
+                stride=(cute.size(tAsA, mode=[1]), 1, 0),
+            ),
+            cutlass.Boolean,
+        )
+        for rest_v in range(tApA.shape[0]):
+            for row in range(tApA.shape[1]):
+                tApA[rest_v, row, 0] = cute.elem_less(
+                    tAcA[(0, rest_v), row, 0, 0][0],
+                    mA.shape[0],
+                )
+
+        k_tile_count = cute.size(tAgA, mode=[3])
+        gmem_pipe_read = cutlass.Int32(0)
+        cute.copy(
+            tiled_copy_A,
+            tAgA[None, None, None, gmem_pipe_read],
+            tAsA[None, None, None, 0],
+            pred=tApA,
+        )
+        if cutlass.const_expr(self.has_channel_residue):
+            cute.copy(
+                tiled_copy_B,
+                tBgB[None, None, None, gmem_pipe_read],
+                tBsB[None, None, None, 0],
+                pred=tBpB,
+            )
+        else:
+            cute.copy(
+                tiled_copy_B,
+                tBgB[None, None, None, gmem_pipe_read],
+                tBsB[None, None, None, 0],
+            )
+        cute.arch.cp_async_commit_group()
+        gmem_pipe_read = gmem_pipe_read + 1
+        for stage in range(1, STAGES - 1):
+            cute.copy(
+                tiled_copy_A,
+                tAgA[None, None, None, gmem_pipe_read],
+                tAsA[None, None, None, stage],
+                pred=tApA,
+            )
+            if cutlass.const_expr(self.has_channel_residue):
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB[None, None, None, gmem_pipe_read],
+                    tBsB[None, None, None, stage],
+                    pred=tBpB,
+                )
+            else:
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB[None, None, None, gmem_pipe_read],
+                    tBsB[None, None, None, stage],
+                )
+            cute.arch.cp_async_commit_group()
+            gmem_pipe_read = gmem_pipe_read + 1
+
+        tCsA = thr_mma.partition_A(sA)
+        tCsB = thr_mma.partition_B(sB)
+        tCgC = thr_mma.partition_C(mC_layout)
+        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
+        tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])
+        tCrC = tiled_mma.make_fragment_C(tCgC)
+        tCrC.fill(0.0)
+
+        smem_pipe_read = cutlass.Int32(0)
+        smem_pipe_write = cutlass.Int32(STAGES - 1)
+        tiles_issued = cutlass.Int32(STAGES - 1)
+        tCsA_p = tCsA[None, None, None, smem_pipe_read]
+        tCsB_p = tCsB[None, None, None, smem_pipe_read]
+        k_block_max = cute.size(tCrA, mode=[2])
+        if k_block_max > 1:
+            cute.arch.cp_async_wait_group(STAGES - 2)
+            self.cta_sync_barrier.arrive_and_wait()
+            cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
+            cute.autovec_copy(tCsB_p[None, None, 0], tCrB[None, None, 0])
+
+        for _ in range(k_tile_count):
+            for k_block in range(k_block_max, unroll_full=True):
+                if k_block == k_block_max - 1:
+                    tCsA_p = tCsA[None, None, None, smem_pipe_read]
+                    tCsB_p = tCsB[None, None, None, smem_pipe_read]
+                    cute.arch.cp_async_wait_group(STAGES - 2)
+                    self.cta_sync_barrier.arrive_and_wait()
+                k_block_next = (k_block + 1) % k_block_max
+                cute.autovec_copy(
+                    tCsA_p[None, None, k_block_next],
+                    tCrA[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsB_p[None, None, k_block_next],
+                    tCrB[None, None, k_block_next],
+                )
+                if k_block == 0 and tiles_issued < k_tile_count:
+                    cute.copy(
+                        tiled_copy_A,
+                        tAgA[None, None, None, gmem_pipe_read],
+                        tAsA[None, None, None, smem_pipe_write],
+                        pred=tApA,
+                    )
+                cute.gemm(
+                    tiled_mma,
+                    tCrC,
+                    tCrA[None, None, k_block],
+                    tCrB[None, None, k_block],
+                    tCrC,
+                )
+                if k_block == 0:
+                    if tiles_issued < k_tile_count:
+                        if cutlass.const_expr(self.has_channel_residue):
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB[None, None, None, gmem_pipe_read],
+                                tBsB[None, None, None, smem_pipe_write],
+                                pred=tBpB,
+                            )
+                        else:
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB[None, None, None, gmem_pipe_read],
+                                tBsB[None, None, None, smem_pipe_write],
+                            )
+                    cute.arch.cp_async_commit_group()
+                    tiles_issued = tiles_issued + 1
+                    smem_pipe_write = smem_pipe_read
+                    smem_pipe_read = smem_pipe_read + 1
+                    if smem_pipe_read == STAGES:
+                        smem_pipe_read = cutlass.Int32(0)
+                    gmem_pipe_read = (
+                        gmem_pipe_read + 1
+                        if gmem_pipe_read + 1 < k_tile_count
+                        else cutlass.Int32(0)
+                    )
+
+        cute.arch.cp_async_wait_group(0)
+        self.cta_sync_barrier.arrive_and_wait()
+        return tCrC
+
+    @cute.jit
+    def _store_adjoint_panel(
+        self,
+        tCrDP: cute.Tensor,
+        tCrBranch: cute.Tensor,
+        panel: cute.Tensor,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        grid_tile: cutlass.Constexpr,
+        channel_tile: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        tCgPanel = thr_mma.partition_C(panel)
+        cPanel = cute.make_identity_tensor(panel.shape)
+        tCpPanel = thr_mma.partition_C(cPanel)
+        pred = cute.make_rmem_tensor(tCrBranch.layout, cutlass.Boolean)
+        residue_m = GRID_SIZE - TILE_M * grid_tile
+        residue_n = self.hidden_channels - self.cta_tiler[1] * channel_tile
+        for idx in range(cute.size(tCrBranch.shape)):
+            pred[idx] = cute.elem_less(
+                tCpPanel[idx],
+                (residue_m, residue_n),
+            )
+            if pred[idx]:
+                tCrBranch[idx] = tCrDP[idx].to(cutlass.Float32) * tCrBranch[idx].to(
+                    cutlass.Float32
+                )
+        atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            panel.element_type,
+        )
+        cute.copy(atom, tCrBranch, tCgPanel, pred=pred)
+        cute.arch.sync_threads()
+
+    @cute.jit
+    def _backproject_panel_accumulate(
+        self,
+        mA: cute.Tensor,
+        panel_b: cute.Tensor,
+        sA: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        grid_tile: cutlass.Constexpr,
+        tCrOut: cute.Tensor,
+    ):
+        """Accumulate one 64-row adjoint panel into coefficient registers."""
+        thr_mma = tiled_mma.get_slice(tidx)
+        gA = cute.local_tile(
+            mA,
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(1, None, 1),
+        )
+        tAgA = tiled_copy_A.get_slice(tidx).partition_S(gA)
+        tAsA = tiled_copy_A.get_slice(tidx).partition_D(sA)
+        # Expose the K-tile mode required by the MMA B-fragment contract.
+        sB = cute.local_tile(
+            panel_b,
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(None, 1, 1),
+        )
+        tSsB = thr_mma.partition_B(sB)
+
+        cA = cute.local_tile(
+            cute.make_identity_tensor(mA.shape),
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(1, None, 1),
+        )
+        tAcA = tiled_copy_A.get_slice(tidx).partition_S(cA)
+        tApA = cute.make_rmem_tensor(
+            cute.make_layout(
+                (
+                    tAsA.shape[0][1],
+                    cute.size(tAsA, mode=[1]),
+                    cute.size(tAsA, mode=[2]),
+                ),
+                stride=(cute.size(tAsA, mode=[1]), 1, 0),
+            ),
+            cutlass.Boolean,
+        )
+        for rest_v in range(tApA.shape[0]):
+            for row in range(tApA.shape[1]):
+                tApA[rest_v, row, 0] = cute.elem_less(
+                    tAcA[(0, rest_v), row, 0, 0][0],
+                    PACKED_COEFF_DIM,
+                )
+
+        panel_k_tiles = TILE_M // self.tile_k
+        if cutlass.const_expr(grid_tile == GRID_TILES - 1):
+            panel_k_tiles = (GRID_SIZE - TILE_M * grid_tile) // self.tile_k
+        panel_k_start = grid_tile * (TILE_M // self.tile_k)
+        gmem_pipe_read = cutlass.Int32(panel_k_start)
+        cute.copy(
+            tiled_copy_A,
+            tAgA[None, None, None, gmem_pipe_read],
+            tAsA[None, None, None, 0],
+            pred=tApA,
+        )
+        cute.arch.cp_async_commit_group()
+        gmem_pipe_read = gmem_pipe_read + 1
+        for stage in range(1, STAGES - 1):
+            cute.copy(
+                tiled_copy_A,
+                tAgA[None, None, None, gmem_pipe_read],
+                tAsA[None, None, None, stage],
+                pred=tApA,
+            )
+            cute.arch.cp_async_commit_group()
+            gmem_pipe_read = gmem_pipe_read + 1
+
+        tCsA = thr_mma.partition_A(sA)
+        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
+        tCrB = tiled_mma.make_fragment_B(tSsB[None, None, None, 0])
+        smem_pipe_read = cutlass.Int32(0)
+        smem_pipe_write = cutlass.Int32(STAGES - 1)
+        tiles_issued = cutlass.Int32(STAGES - 1)
+        logical_k_tile = cutlass.Int32(0)
+        tCsA_p = tCsA[None, None, None, smem_pipe_read]
+        k_block_max = cute.size(tCrA, mode=[2])
+        if k_block_max > 1:
+            cute.arch.cp_async_wait_group(STAGES - 2)
+            self.cta_sync_barrier.arrive_and_wait()
+            cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
+            cute.autovec_copy(
+                tSsB[None, None, 0, logical_k_tile],
+                tCrB[None, None, 0],
+            )
+
+        for _ in range(panel_k_tiles):
+            for k_block in range(k_block_max, unroll_full=True):
+                if k_block == k_block_max - 1:
+                    tCsA_p = tCsA[None, None, None, smem_pipe_read]
+                    cute.arch.cp_async_wait_group(STAGES - 2)
+                    self.cta_sync_barrier.arrive_and_wait()
+                k_block_next = (k_block + 1) % k_block_max
+                fragment_k_tile = logical_k_tile
+                if k_block_max > 1 and k_block == k_block_max - 1:
+                    fragment_k_tile = (
+                        logical_k_tile + 1
+                        if logical_k_tile + 1 < panel_k_tiles
+                        else logical_k_tile
+                    )
+                cute.autovec_copy(
+                    tCsA_p[None, None, k_block_next],
+                    tCrA[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tSsB[
+                        None,
+                        None,
+                        k_block_next,
+                        fragment_k_tile,
+                    ],
+                    tCrB[None, None, k_block_next],
+                )
+                if k_block == 0 and tiles_issued < panel_k_tiles:
+                    cute.copy(
+                        tiled_copy_A,
+                        tAgA[None, None, None, gmem_pipe_read],
+                        tAsA[None, None, None, smem_pipe_write],
+                        pred=tApA,
+                    )
+                cute.gemm(
+                    tiled_mma,
+                    tCrOut,
+                    tCrA[None, None, k_block],
+                    tCrB[None, None, k_block],
+                    tCrOut,
+                )
+                if k_block == 0:
+                    cute.arch.cp_async_commit_group()
+                    tiles_issued = tiles_issued + 1
+                    smem_pipe_write = smem_pipe_read
+                    smem_pipe_read = smem_pipe_read + 1
+                    if smem_pipe_read == STAGES:
+                        smem_pipe_read = cutlass.Int32(0)
+                    gmem_pipe_read = (
+                        gmem_pipe_read + 1
+                        if gmem_pipe_read + 1 < panel_k_start + panel_k_tiles
+                        else cutlass.Int32(panel_k_start)
+                    )
+            logical_k_tile = logical_k_tile + 1
+
+        cute.arch.cp_async_wait_group(0)
+        self.cta_sync_barrier.arrive_and_wait()
+
+    @cute.jit
+    def _single_projection_to_shared(
+        self,
+        mA: cute.Tensor,
+        mB: cute.Tensor,
+        mGrid: cute.Tensor,
+        sA: cute.Tensor,
+        sB: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        grid_tile: cutlass.Constexpr,
+        channel_tile: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        gA = cute.local_tile(
+            mA,
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, None, 1),
+        )
+        gB = cute.local_tile(
+            mB,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(None, 1, 1),
+        )
+        gGrid = cute.local_tile(
+            mGrid,
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, 1, None),
+        )
+        thr_copy_A = tiled_copy_A.get_slice(tidx)
+        thr_copy_B = tiled_copy_B.get_slice(tidx)
+        tAgA = thr_copy_A.partition_S(gA)
+        tAsA = thr_copy_A.partition_D(sA)
+        tBgB = thr_copy_B.partition_S(gB)
+        tBsB = thr_copy_B.partition_D(sB)
+
+        if cutlass.const_expr(self.has_channel_residue):
+            cB = cute.local_tile(
+                cute.make_identity_tensor(mB.shape),
+                tiler=self.cta_tiler,
+                coord=(0, channel_tile, None),
+                proj=(None, 1, 1),
+            )
+            tBcB = thr_copy_B.partition_S(cB)
+            tBpB = cute.make_rmem_tensor(
+                cute.make_layout(
+                    (
+                        tBsB.shape[0][1],
+                        cute.size(tBsB, mode=[1]),
+                        cute.size(tBsB, mode=[2]),
+                    ),
+                    stride=(cute.size(tBsB, mode=[1]), 1, 0),
+                ),
+                cutlass.Boolean,
+            )
+            for rest_v in range(tBpB.shape[0]):
+                for channel in range(tBpB.shape[1]):
+                    tBpB[rest_v, channel, 0] = cute.elem_less(
+                        tBcB[(0, rest_v), channel, 0, 0][0],
+                        mB.shape[0],
+                    )
+
+        cA = cute.local_tile(
+            cute.make_identity_tensor(mA.shape),
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, None, 1),
+        )
+        tAcA = thr_copy_A.partition_S(cA)
+        tApA = cute.make_rmem_tensor(
+            cute.make_layout(
+                (
+                    tAsA.shape[0][1],
+                    cute.size(tAsA, mode=[1]),
+                    cute.size(tAsA, mode=[2]),
+                ),
+                stride=(cute.size(tAsA, mode=[1]), 1, 0),
+            ),
+            cutlass.Boolean,
+        )
+        for rest_v in range(tApA.shape[0]):
+            for row in range(tApA.shape[1]):
+                tApA[rest_v, row, 0] = cute.elem_less(
+                    tAcA[(0, rest_v), row, 0, 0][0],
+                    mA.shape[0],
+                )
+
+        k_tile_count = cute.size(tAgA, mode=[3])
+        gmem_pipe_read = cutlass.Int32(0)
+        cute.copy(
+            tiled_copy_A,
+            tAgA[None, None, None, gmem_pipe_read],
+            tAsA[None, None, None, 0],
+            pred=tApA,
+        )
+        if cutlass.const_expr(self.has_channel_residue):
+            cute.copy(
+                tiled_copy_B,
+                tBgB[None, None, None, gmem_pipe_read],
+                tBsB[None, None, None, 0],
+                pred=tBpB,
+            )
+        else:
+            cute.copy(
+                tiled_copy_B,
+                tBgB[None, None, None, gmem_pipe_read],
+                tBsB[None, None, None, 0],
+            )
+        cute.arch.cp_async_commit_group()
+        gmem_pipe_read = gmem_pipe_read + 1
+        for stage in range(1, STAGES - 1):
+            cute.copy(
+                tiled_copy_A,
+                tAgA[None, None, None, gmem_pipe_read],
+                tAsA[None, None, None, stage],
+                pred=tApA,
+            )
+            if cutlass.const_expr(self.has_channel_residue):
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB[None, None, None, gmem_pipe_read],
+                    tBsB[None, None, None, stage],
+                    pred=tBpB,
+                )
+            else:
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB[None, None, None, gmem_pipe_read],
+                    tBsB[None, None, None, stage],
+                )
+            cute.arch.cp_async_commit_group()
+            gmem_pipe_read = gmem_pipe_read + 1
+
+        tCsA = thr_mma.partition_A(sA)
+        tCsB = thr_mma.partition_B(sB)
+        tCgGrid = thr_mma.partition_C(gGrid)
+        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
+        tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])
+        tCrGrid = tiled_mma.make_fragment_C(tCgGrid)
+        tCrGrid.fill(0.0)
+
+        smem_pipe_read = cutlass.Int32(0)
+        smem_pipe_write = cutlass.Int32(STAGES - 1)
+        tiles_issued = cutlass.Int32(STAGES - 1)
+        tCsA_p = tCsA[None, None, None, smem_pipe_read]
+        tCsB_p = tCsB[None, None, None, smem_pipe_read]
+        k_block_max = cute.size(tCrA, mode=[2])
+        if k_block_max > 1:
+            cute.arch.cp_async_wait_group(STAGES - 2)
+            self.cta_sync_barrier.arrive_and_wait()
+            cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
+            cute.autovec_copy(tCsB_p[None, None, 0], tCrB[None, None, 0])
+
+        for _ in range(k_tile_count):
+            for k_block in range(k_block_max, unroll_full=True):
+                if k_block == k_block_max - 1:
+                    tCsA_p = tCsA[None, None, None, smem_pipe_read]
+                    tCsB_p = tCsB[None, None, None, smem_pipe_read]
+                    cute.arch.cp_async_wait_group(STAGES - 2)
+                    self.cta_sync_barrier.arrive_and_wait()
+                k_block_next = (k_block + 1) % k_block_max
+                cute.autovec_copy(
+                    tCsA_p[None, None, k_block_next],
+                    tCrA[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsB_p[None, None, k_block_next],
+                    tCrB[None, None, k_block_next],
+                )
+                if k_block == 0 and tiles_issued < k_tile_count:
+                    cute.copy(
+                        tiled_copy_A,
+                        tAgA[None, None, None, gmem_pipe_read],
+                        tAsA[None, None, None, smem_pipe_write],
+                        pred=tApA,
+                    )
+                cute.gemm(
+                    tiled_mma,
+                    tCrGrid,
+                    tCrA[None, None, k_block],
+                    tCrB[None, None, k_block],
+                    tCrGrid,
+                )
+                if k_block == 0:
+                    if tiles_issued < k_tile_count:
+                        if cutlass.const_expr(self.has_channel_residue):
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB[None, None, None, gmem_pipe_read],
+                                tBsB[None, None, None, smem_pipe_write],
+                                pred=tBpB,
+                            )
+                        else:
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB[None, None, None, gmem_pipe_read],
+                                tBsB[None, None, None, smem_pipe_write],
+                            )
+                    cute.arch.cp_async_commit_group()
+                    tiles_issued = tiles_issued + 1
+                    smem_pipe_write = smem_pipe_read
+                    smem_pipe_read = smem_pipe_read + 1
+                    if smem_pipe_read == STAGES:
+                        smem_pipe_read = cutlass.Int32(0)
+                    gmem_pipe_read = (
+                        gmem_pipe_read + 1
+                        if gmem_pipe_read + 1 < k_tile_count
+                        else cutlass.Int32(0)
+                    )
+
+        cute.arch.cp_async_wait_group(0)
+        self.cta_sync_barrier.arrive_and_wait()
+        cGrid = cute.make_identity_tensor(gGrid.shape)
+        tCpGrid = thr_mma.partition_C(cGrid)
+        pred = cute.make_rmem_tensor(tCrGrid.layout, cutlass.Boolean)
+        residue_m = GRID_SIZE - TILE_M * grid_tile
+        for idx in range(cute.size(tCrGrid.shape)):
+            pred[idx] = cute.elem_less(
+                tCpGrid[idx],
+                (residue_m, self.cta_tiler[1]),
+            )
+        atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mGrid.element_type,
+        )
+        cute.copy(atom, tCrGrid, tCgGrid, pred=pred)
+        cute.arch.sync_threads()
+
+    @cute.jit
+    def _dual_projection_adjoint(
+        self,
+        mA: cute.Tensor,
+        mB_left: cute.Tensor,
+        mB_right: cute.Tensor,
+        mGrad_left_grid: cute.Tensor,
+        mGrad_right_grid: cute.Tensor,
+        sA: cute.Tensor,
+        sB_left: cute.Tensor,
+        sB_right: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        grid_tile: cutlass.Constexpr,
+        channel_tile: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        gA = cute.local_tile(
+            mA,
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, None, 1),
+        )
+        gB_left = cute.local_tile(
+            mB_left,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(None, 1, 1),
+        )
+        gB_right = cute.local_tile(
+            mB_right,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(None, 1, 1),
+        )
+        gGrad_left = cute.local_tile(
+            mGrad_left_grid,
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, 1, None),
+        )
+        gGrad_right = cute.local_tile(
+            mGrad_right_grid,
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, 1, None),
+        )
+
+        thr_copy_A = tiled_copy_A.get_slice(tidx)
+        thr_copy_B = tiled_copy_B.get_slice(tidx)
+        tAgA = thr_copy_A.partition_S(gA)
+        tAsA = thr_copy_A.partition_D(sA)
+        tBgB_left = thr_copy_B.partition_S(gB_left)
+        tBgB_right = thr_copy_B.partition_S(gB_right)
+        tBsB_left = thr_copy_B.partition_D(sB_left)
+        tBsB_right = thr_copy_B.partition_D(sB_right)
+
+        if cutlass.const_expr(self.has_channel_residue):
+            cB = cute.local_tile(
+                cute.make_identity_tensor(mB_left.shape),
+                tiler=self.cta_tiler,
+                coord=(0, channel_tile, None),
+                proj=(None, 1, 1),
+            )
+            tBcB = thr_copy_B.partition_S(cB)
+            tBpB = cute.make_rmem_tensor(
+                cute.make_layout(
+                    (
+                        tBsB_left.shape[0][1],
+                        cute.size(tBsB_left, mode=[1]),
+                        cute.size(tBsB_left, mode=[2]),
+                    ),
+                    stride=(cute.size(tBsB_left, mode=[1]), 1, 0),
+                ),
+                cutlass.Boolean,
+            )
+            for rest_v in range(tBpB.shape[0]):
+                for channel in range(tBpB.shape[1]):
+                    tBpB[rest_v, channel, 0] = cute.elem_less(
+                        tBcB[(0, rest_v), channel, 0, 0][0],
+                        mB_left.shape[0],
+                    )
+
+        cA = cute.local_tile(
+            cute.make_identity_tensor(mA.shape),
+            tiler=self.cta_tiler,
+            coord=(grid_tile, 0, None),
+            proj=(1, None, 1),
+        )
+        tAcA = thr_copy_A.partition_S(cA)
+        tApA = cute.make_rmem_tensor(
+            cute.make_layout(
+                (
+                    tAsA.shape[0][1],
+                    cute.size(tAsA, mode=[1]),
+                    cute.size(tAsA, mode=[2]),
+                ),
+                stride=(cute.size(tAsA, mode=[1]), 1, 0),
+            ),
+            cutlass.Boolean,
+        )
+        for rest_v in range(tApA.shape[0]):
+            for row in range(tApA.shape[1]):
+                tApA[rest_v, row, 0] = cute.elem_less(
+                    tAcA[(0, rest_v), row, 0, 0][0],
+                    mA.shape[0],
+                )
+
+        k_tile_count = cute.size(tAgA, mode=[3])
+        gmem_pipe_read = cutlass.Int32(0)
+        cute.copy(
+            tiled_copy_A,
+            tAgA[None, None, None, gmem_pipe_read],
+            tAsA[None, None, None, 0],
+            pred=tApA,
+        )
+        if cutlass.const_expr(self.has_channel_residue):
+            cute.copy(
+                tiled_copy_B,
+                tBgB_left[None, None, None, gmem_pipe_read],
+                tBsB_left[None, None, None, 0],
+                pred=tBpB,
+            )
+        else:
+            cute.copy(
+                tiled_copy_B,
+                tBgB_left[None, None, None, gmem_pipe_read],
+                tBsB_left[None, None, None, 0],
+            )
+        if cutlass.const_expr(self.has_channel_residue):
+            cute.copy(
+                tiled_copy_B,
+                tBgB_right[None, None, None, gmem_pipe_read],
+                tBsB_right[None, None, None, 0],
+                pred=tBpB,
+            )
+        else:
+            cute.copy(
+                tiled_copy_B,
+                tBgB_right[None, None, None, gmem_pipe_read],
+                tBsB_right[None, None, None, 0],
+            )
+        cute.arch.cp_async_commit_group()
+        gmem_pipe_read = gmem_pipe_read + 1
+        for stage in range(1, STAGES - 1):
+            cute.copy(
+                tiled_copy_A,
+                tAgA[None, None, None, gmem_pipe_read],
+                tAsA[None, None, None, stage],
+                pred=tApA,
+            )
+            if cutlass.const_expr(self.has_channel_residue):
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB_left[None, None, None, gmem_pipe_read],
+                    tBsB_left[None, None, None, stage],
+                    pred=tBpB,
+                )
+            else:
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB_left[None, None, None, gmem_pipe_read],
+                    tBsB_left[None, None, None, stage],
+                )
+            if cutlass.const_expr(self.has_channel_residue):
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB_right[None, None, None, gmem_pipe_read],
+                    tBsB_right[None, None, None, stage],
+                    pred=tBpB,
+                )
+            else:
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB_right[None, None, None, gmem_pipe_read],
+                    tBsB_right[None, None, None, stage],
+                )
+            cute.arch.cp_async_commit_group()
+            gmem_pipe_read = gmem_pipe_read + 1
+
+        tCsA = thr_mma.partition_A(sA)
+        tCsB_left = thr_mma.partition_B(sB_left)
+        tCsB_right = thr_mma.partition_B(sB_right)
+        tCgGrad_left = thr_mma.partition_C(gGrad_left)
+        tCgGrad_right = thr_mma.partition_C(gGrad_right)
+        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
+        tCrB_left = tiled_mma.make_fragment_B(tCsB_left[None, None, None, 0])
+        tCrB_right = tiled_mma.make_fragment_B(tCsB_right[None, None, None, 0])
+        tCrLeft = tiled_mma.make_fragment_C(tCgGrad_left)
+        tCrRight = tiled_mma.make_fragment_C(tCgGrad_right)
+        tCrLeft.fill(0.0)
+        tCrRight.fill(0.0)
+
+        smem_pipe_read = cutlass.Int32(0)
+        smem_pipe_write = cutlass.Int32(STAGES - 1)
+        tiles_issued = cutlass.Int32(STAGES - 1)
+        tCsA_p = tCsA[None, None, None, smem_pipe_read]
+        tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read]
+        tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read]
+        k_block_max = cute.size(tCrA, mode=[2])
+        if k_block_max > 1:
+            cute.arch.cp_async_wait_group(STAGES - 2)
+            self.cta_sync_barrier.arrive_and_wait()
+            cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
+            cute.autovec_copy(
+                tCsB_left_p[None, None, 0],
+                tCrB_left[None, None, 0],
+            )
+            cute.autovec_copy(
+                tCsB_right_p[None, None, 0],
+                tCrB_right[None, None, 0],
+            )
+
+        for _ in range(k_tile_count):
+            for k_block in range(k_block_max, unroll_full=True):
+                if k_block == k_block_max - 1:
+                    tCsA_p = tCsA[None, None, None, smem_pipe_read]
+                    tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read]
+                    tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read]
+                    cute.arch.cp_async_wait_group(STAGES - 2)
+                    self.cta_sync_barrier.arrive_and_wait()
+                k_block_next = (k_block + 1) % k_block_max
+                cute.autovec_copy(
+                    tCsA_p[None, None, k_block_next],
+                    tCrA[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsB_left_p[None, None, k_block_next],
+                    tCrB_left[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsB_right_p[None, None, k_block_next],
+                    tCrB_right[None, None, k_block_next],
+                )
+                if k_block == 0 and tiles_issued < k_tile_count:
+                    cute.copy(
+                        tiled_copy_A,
+                        tAgA[None, None, None, gmem_pipe_read],
+                        tAsA[None, None, None, smem_pipe_write],
+                        pred=tApA,
+                    )
+                cute.gemm(
+                    tiled_mma,
+                    tCrLeft,
+                    tCrA[None, None, k_block],
+                    tCrB_left[None, None, k_block],
+                    tCrLeft,
+                )
+                cute.gemm(
+                    tiled_mma,
+                    tCrRight,
+                    tCrA[None, None, k_block],
+                    tCrB_right[None, None, k_block],
+                    tCrRight,
+                )
+                if k_block == 0:
+                    if tiles_issued < k_tile_count:
+                        if cutlass.const_expr(self.has_channel_residue):
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB_left[None, None, None, gmem_pipe_read],
+                                tBsB_left[None, None, None, smem_pipe_write],
+                                pred=tBpB,
+                            )
+                        else:
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB_left[None, None, None, gmem_pipe_read],
+                                tBsB_left[None, None, None, smem_pipe_write],
+                            )
+                        if cutlass.const_expr(self.has_channel_residue):
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB_right[None, None, None, gmem_pipe_read],
+                                tBsB_right[None, None, None, smem_pipe_write],
+                                pred=tBpB,
+                            )
+                        else:
+                            cute.copy(
+                                tiled_copy_B,
+                                tBgB_right[None, None, None, gmem_pipe_read],
+                                tBsB_right[None, None, None, smem_pipe_write],
+                            )
+                    cute.arch.cp_async_commit_group()
+                    tiles_issued = tiles_issued + 1
+                    smem_pipe_write = smem_pipe_read
+                    smem_pipe_read = smem_pipe_read + 1
+                    if smem_pipe_read == STAGES:
+                        smem_pipe_read = cutlass.Int32(0)
+                    gmem_pipe_read = (
+                        gmem_pipe_read + 1
+                        if gmem_pipe_read + 1 < k_tile_count
+                        else cutlass.Int32(0)
+                    )
+
+        cute.arch.cp_async_wait_group(0)
+        self.cta_sync_barrier.arrive_and_wait()
+        cGrid = cute.make_identity_tensor(gGrad_left.shape)
+        tCpGrid = thr_mma.partition_C(cGrid)
+        pred = cute.make_rmem_tensor(tCrLeft.layout, cutlass.Boolean)
+        tCrDP = tiled_mma.make_fragment_C(tCgGrad_left)
+        tCrDP.fill(0.0)
+        residue_m = GRID_SIZE - TILE_M * grid_tile
+        for idx in range(cute.size(tCrLeft.shape)):
+            pred[idx] = cute.elem_less(
+                tCpGrid[idx],
+                (residue_m, self.cta_tiler[1]),
+            )
+        atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mGrad_left_grid.element_type,
+        )
+        cute.copy(atom, tCgGrad_left, tCrDP, pred=pred)
+        for idx in range(cute.size(tCrLeft.shape)):
+            if pred[idx]:
+                left_value = tCrLeft[idx].to(cutlass.Float32)
+                right_value = tCrRight[idx].to(cutlass.Float32)
+                grad_product = tCrDP[idx].to(cutlass.Float32)
+                tCrLeft[idx] = grad_product * right_value
+                tCrRight[idx] = grad_product * left_value
+        cute.copy(atom, tCrLeft, tCgGrad_left, pred=pred)
+        cute.copy(atom, tCrRight, tCgGrad_right, pred=pred)
+        cute.arch.sync_threads()
+
+    @cute.jit
+    def _dual_backproject(
+        self,
+        mA: cute.Tensor,
+        mB_left_shared: cute.Tensor,
+        mB_right_shared: cute.Tensor,
+        mOut_left: cute.Tensor,
+        mOut_right: cute.Tensor,
+        sA: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        channel_tile: cutlass.Int32,
+    ):
+        thr_mma = tiled_mma.get_slice(tidx)
+        gA = cute.local_tile(
+            mA,
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(1, None, 1),
+        )
+        sB_left = cute.local_tile(
+            mB_left_shared,
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(None, 1, 1),
+        )
+        sB_right = cute.local_tile(
+            mB_right_shared,
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(None, 1, 1),
+        )
+        gOut_left = cute.local_tile(
+            mOut_left,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(1, 1, None),
+        )
+        gOut_right = cute.local_tile(
+            mOut_right,
+            tiler=self.cta_tiler,
+            coord=(0, channel_tile, None),
+            proj=(1, 1, None),
+        )
+        thr_copy_A = tiled_copy_A.get_slice(tidx)
+        tAgA = thr_copy_A.partition_S(gA)
+        tAsA = thr_copy_A.partition_D(sA)
+
+        cA = cute.local_tile(
+            cute.make_identity_tensor(mA.shape),
+            tiler=self.cta_tiler,
+            coord=(0, 0, None),
+            proj=(1, None, 1),
+        )
+        tAcA = thr_copy_A.partition_S(cA)
+        tApA = cute.make_rmem_tensor(
+            cute.make_layout(
+                (
+                    tAsA.shape[0][1],
+                    cute.size(tAsA, mode=[1]),
+                    cute.size(tAsA, mode=[2]),
+                ),
+                stride=(cute.size(tAsA, mode=[1]), 1, 0),
+            ),
+            cutlass.Boolean,
+        )
+        for rest_v in range(tApA.shape[0]):
+            for row in range(tApA.shape[1]):
+                tApA[rest_v, row, 0] = cute.elem_less(
+                    tAcA[(0, rest_v), row, 0, 0][0],
+                    PACKED_COEFF_DIM,
+                )
+
+        k_tile_count = cute.size(tAgA, mode=[3])
+        gmem_pipe_read = cutlass.Int32(0)
+        cute.copy(
+            tiled_copy_A,
+            tAgA[None, None, None, gmem_pipe_read],
+            tAsA[None, None, None, 0],
+            pred=tApA,
+        )
+        cute.arch.cp_async_commit_group()
+        gmem_pipe_read = gmem_pipe_read + 1
+        for stage in range(1, STAGES - 1):
+            cute.copy(
+                tiled_copy_A,
+                tAgA[None, None, None, gmem_pipe_read],
+                tAsA[None, None, None, stage],
+                pred=tApA,
+            )
+            cute.arch.cp_async_commit_group()
+            gmem_pipe_read = gmem_pipe_read + 1
+
+        tCsA = thr_mma.partition_A(sA)
+        tSsB_left = thr_mma.partition_B(sB_left)
+        tSsB_right = thr_mma.partition_B(sB_right)
+        tCgOut_left = thr_mma.partition_C(gOut_left)
+        tCgOut_right = thr_mma.partition_C(gOut_right)
+        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
+        tCrB_left = tiled_mma.make_fragment_B(tSsB_left[None, None, None, 0])
+        tCrB_right = tiled_mma.make_fragment_B(tSsB_right[None, None, None, 0])
+        tCrOut_left = tiled_mma.make_fragment_C(tCgOut_left)
+        tCrOut_right = tiled_mma.make_fragment_C(tCgOut_right)
+        tCrOut_left.fill(0.0)
+        tCrOut_right.fill(0.0)
+
+        smem_pipe_read = cutlass.Int32(0)
+        smem_pipe_write = cutlass.Int32(STAGES - 1)
+        tiles_issued = cutlass.Int32(STAGES - 1)
+        logical_k_tile = cutlass.Int32(0)
+        tCsA_p = tCsA[None, None, None, smem_pipe_read]
+        k_block_max = cute.size(tCrA, mode=[2])
+        if k_block_max > 1:
+            cute.arch.cp_async_wait_group(STAGES - 2)
+            self.cta_sync_barrier.arrive_and_wait()
+            cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
+            cute.autovec_copy(
+                tSsB_left[None, None, 0, logical_k_tile],
+                tCrB_left[None, None, 0],
+            )
+            cute.autovec_copy(
+                tSsB_right[None, None, 0, logical_k_tile],
+                tCrB_right[None, None, 0],
+            )
+
+        for _ in range(k_tile_count):
+            for k_block in range(k_block_max, unroll_full=True):
+                if k_block == k_block_max - 1:
+                    tCsA_p = tCsA[None, None, None, smem_pipe_read]
+                    cute.arch.cp_async_wait_group(STAGES - 2)
+                    self.cta_sync_barrier.arrive_and_wait()
+                k_block_next = (k_block + 1) % k_block_max
+                fragment_k_tile = logical_k_tile
+                if k_block_max > 1:
+                    if k_block == k_block_max - 1:
+                        fragment_k_tile = (
+                            logical_k_tile + 1
+                            if logical_k_tile + 1 < k_tile_count
+                            else logical_k_tile
+                        )
+                cute.autovec_copy(
+                    tCsA_p[None, None, k_block_next],
+                    tCrA[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tSsB_left[
+                        None,
+                        None,
+                        k_block_next,
+                        fragment_k_tile,
+                    ],
+                    tCrB_left[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tSsB_right[
+                        None,
+                        None,
+                        k_block_next,
+                        fragment_k_tile,
+                    ],
+                    tCrB_right[None, None, k_block_next],
+                )
+                if k_block == 0 and tiles_issued < k_tile_count:
+                    cute.copy(
+                        tiled_copy_A,
+                        tAgA[None, None, None, gmem_pipe_read],
+                        tAsA[None, None, None, smem_pipe_write],
+                        pred=tApA,
+                    )
+                cute.gemm(
+                    tiled_mma,
+                    tCrOut_left,
+                    tCrA[None, None, k_block],
+                    tCrB_left[None, None, k_block],
+                    tCrOut_left,
+                )
+                cute.gemm(
+                    tiled_mma,
+                    tCrOut_right,
+                    tCrA[None, None, k_block],
+                    tCrB_right[None, None, k_block],
+                    tCrOut_right,
+                )
+                if k_block == 0:
+                    cute.arch.cp_async_commit_group()
+                    tiles_issued = tiles_issued + 1
+                    smem_pipe_write = smem_pipe_read
+                    smem_pipe_read = smem_pipe_read + 1
+                    if smem_pipe_read == STAGES:
+                        smem_pipe_read = cutlass.Int32(0)
+                    gmem_pipe_read = (
+                        gmem_pipe_read + 1
+                        if gmem_pipe_read + 1 < k_tile_count
+                        else cutlass.Int32(0)
+                    )
+            logical_k_tile = logical_k_tile + 1
+
+        cute.arch.cp_async_wait_group(0)
+        self.cta_sync_barrier.arrive_and_wait()
+        pred = cute.make_rmem_tensor(tCrOut_left.layout, cutlass.Boolean)
+        if cutlass.const_expr(self.has_channel_residue):
+            cOut = cute.local_tile(
+                cute.make_identity_tensor(mOut_left.shape),
+                tiler=self.cta_tiler,
+                coord=(0, channel_tile, None),
+                proj=(1, 1, None),
+            )
+            tCpOut = thr_mma.partition_C(cOut)
+            for idx in range(cute.size(tCrOut_left.shape)):
+                pred[idx] = cute.elem_less(tCpOut[idx], mOut_left.shape)
+        else:
+            cOut = cute.make_identity_tensor(gOut_left.shape)
+            tCpOut = thr_mma.partition_C(cOut)
+            for idx in range(cute.size(tCrOut_left.shape)):
+                pred[idx] = cute.elem_less(
+                    tCpOut[idx],
+                    (PACKED_COEFF_DIM, self.cta_tiler[1]),
+                )
+        atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mOut_left.element_type,
+        )
+        cute.copy(atom, tCrOut_left, tCgOut_left, pred=pred)
+        cute.copy(atom, tCrOut_right, tCgOut_right, pred=pred)
+
+
+def compile_tiled_output_grid_product(
+    device_index: int | None = None,
+    compute_capability: tuple[int, int] | None = None,
+    hidden_channels: int = 192,
+    tile_n: int = TILE_N,
+    channel_tile_start: int = 0,
+    channel_tile_count: int | None = None,
+) -> Callable:
+    """Compile one forward artifact with symbolic runtime node count."""
+    import torch
+
+    hidden_channels = _validate_hidden_channels(hidden_channels)
+    tile_n = int(tile_n)
+    if tile_n not in (C96_TAIL_TILE_N, SM80_C96_TILE_N, TILE_N):
+        raise ValueError("output-grid forward tile_n must be 32, 48, or 64")
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    actual_capability = tuple(torch.cuda.get_device_capability(device_index))
+    if (
+        compute_capability is not None
+        and tuple(compute_capability) != actual_capability
+    ):
+        raise ValueError("compile target does not match the selected CUDA device")
+    if tile_n == SM80_C96_TILE_N and (
+        actual_capability not in runtime_policy.SM80_PROFILE_CAPABILITIES
+        or hidden_channels != 96
+    ):
+        raise ValueError("output-grid forward N=48 requires SM80-family and C=96")
+    if tile_n == C96_TAIL_TILE_N and (
+        actual_capability != (9, 0)
+        or hidden_channels != 96
+        or int(channel_tile_start) != C96_TAIL_CHANNEL_TILE
+        or int(channel_tile_count or 0) != 1
+    ):
+        raise ValueError("output-grid forward N=32 tail requires sm90 and C=96")
+    if (
+        tile_n == TILE_N
+        and channel_tile_count is not None
+        and (
+            actual_capability != (9, 0)
+            or hidden_channels != 96
+            or int(channel_tile_start) != 0
+            or int(channel_tile_count) != 1
+        )
+    ):
+        raise ValueError(
+            "partial output-grid forward N=64 launch requires sm90 C=96 base panel"
+        )
+    with torch.cuda.device(device_index):
+        nodes = cute.sym_int64()
+        fake_coeff = make_fake_compact_tensor(
+            cutlass.Float32,
+            (nodes, PACKED_COEFF_DIM, hidden_channels),
+            stride_order=(2, 1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_to_grid = make_fake_compact_tensor(
+            cutlass.Float32,
+            (GRID_SIZE, PACKED_COEFF_DIM),
+            stride_order=(1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_from_grid = make_fake_compact_tensor(
+            cutlass.Float32,
+            (PACKED_COEFF_DIM, GRID_SIZE),
+            stride_order=(1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+        return cute.compile(
+            TiledOutputGridProductForward(
+                hidden_channels,
+                tile_n=tile_n,
+                channel_tile_start=channel_tile_start,
+                channel_tile_count=channel_tile_count,
+            ),
+            fake_coeff,
+            fake_coeff,
+            fake_to_grid,
+            fake_from_grid,
+            fake_coeff,
+            fake_stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+def compile_tiled_output_grid_product_backward(
+    device_index: int | None = None,
+    compute_capability: tuple[int, int] | None = None,
+    hidden_channels: int = 192,
+    tile_n: int = TILE_N,
+    channel_tile_start: int = 0,
+    channel_tile_count: int | None = None,
+) -> Callable:
+    """Compile one first-backward artifact with symbolic node count."""
+    import torch
+
+    hidden_channels = _validate_hidden_channels(hidden_channels)
+    tile_n = int(tile_n)
+    if tile_n not in (C96_TAIL_TILE_N, SM80_C96_TILE_N, TILE_N):
+        raise ValueError("output-grid backward tile_n must be 32, 48, or 64")
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    actual_capability = tuple(torch.cuda.get_device_capability(device_index))
+    if (
+        compute_capability is not None
+        and tuple(compute_capability) != actual_capability
+    ):
+        raise ValueError("compile target does not match the selected CUDA device")
+    if tile_n == SM80_C96_TILE_N and (
+        actual_capability not in runtime_policy.SM80_PROFILE_CAPABILITIES
+        or hidden_channels != 96
+    ):
+        raise ValueError(
+            "output-grid N=48 panel adjoint requires SM80-family, C=96, and K=8"
+        )
+    if tile_n == C96_TAIL_TILE_N and (
+        actual_capability != (9, 0)
+        or hidden_channels != 96
+        or int(channel_tile_start) != C96_TAIL_CHANNEL_TILE
+        or int(channel_tile_count or 0) != 1
+    ):
+        raise ValueError("output-grid backward N=32 tail requires sm90, C=96, and K=8")
+    if (
+        tile_n == TILE_N
+        and channel_tile_count is not None
+        and (
+            actual_capability != (9, 0)
+            or hidden_channels != 96
+            or int(channel_tile_start) != 0
+            or int(channel_tile_count) != 1
+        )
+    ):
+        raise ValueError(
+            "partial output-grid backward N=64 launch requires sm90 C=96 K=8 base panel"
+        )
+    with torch.cuda.device(device_index):
+        nodes = cute.sym_int64()
+        fake_coeff = make_fake_compact_tensor(
+            cutlass.Float32,
+            (nodes, PACKED_COEFF_DIM, hidden_channels),
+            stride_order=(2, 1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_to_grid = make_fake_compact_tensor(
+            cutlass.Float32,
+            (GRID_SIZE, PACKED_COEFF_DIM),
+            stride_order=(1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_from_grid = make_fake_compact_tensor(
+            cutlass.Float32,
+            (PACKED_COEFF_DIM, GRID_SIZE),
+            stride_order=(1, 0),
+            **FAKE_TENSOR_KW,
+        )
+        fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+        return cute.compile(
+            TiledOutputGridProductBackward(
+                hidden_channels,
+                tile_n=tile_n,
+                channel_tile_start=channel_tile_start,
+                channel_tile_count=channel_tile_count,
+            ),
+            fake_coeff,
+            fake_coeff,
+            fake_coeff,
+            fake_to_grid,
+            fake_from_grid,
+            fake_coeff,
+            fake_coeff,
+            fake_stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+@lru_cache(maxsize=16)
+def _compiled_tiled_forward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+    hidden_channels: int,
+    tile_n: int,
+    channel_tile_start: int,
+    channel_tile_count: int | None,
+) -> Callable:
+    return compile_tiled_output_grid_product(
+        device_index,
+        compute_capability,
+        hidden_channels,
+        tile_n,
+        channel_tile_start,
+        channel_tile_count,
+    )
+
+
+@lru_cache(maxsize=16)
+def _compiled_tiled_backward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+    hidden_channels: int,
+    tile_n: int,
+    channel_tile_start: int,
+    channel_tile_count: int | None,
+) -> Callable:
+    return compile_tiled_output_grid_product_backward(
+        device_index,
+        compute_capability,
+        hidden_channels,
+        tile_n,
+        channel_tile_start,
+        channel_tile_count,
+    )
+
+
+def _compile_key(tensor) -> tuple[int, tuple[int, int], int]:
+    import torch
+
+    device_index = tensor.device.index
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    compute_capability = tuple(torch.cuda.get_device_capability(device_index))
+    return int(device_index), compute_capability, int(tensor.shape[2])
+
+
+def _validate_tensors(left, right, to_grid, from_grid, out=None) -> None:
+    import torch
+
+    tensors = (left, right, to_grid, from_grid)
+    if (
+        any(not tensor.is_cuda for tensor in tensors)
+        or any(tensor.dtype != torch.float32 for tensor in tensors)
+        or any(tensor.device != left.device for tensor in tensors)
+        or left.ndim != 3
+        or left.shape[0] <= 0
+        or int(left.shape[1]) != PACKED_COEFF_DIM
+        or int(left.shape[2]) not in SUPPORTED_HIDDEN_CHANNELS
+        or right.shape != left.shape
+        or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM)
+        or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE)
+        or any(not tensor.is_contiguous() for tensor in tensors)
+        or torch.cuda.get_device_capability(left.device)[0] < 8
+    ):
+        raise ValueError(
+            "tiled output-grid product requires contiguous CUDA FP32 "
+            "left/right=(N,48,C) with C in {96,192}, to_grid=(152,48), and "
+            "from_grid=(48,152) tensors on compute capability 8.0+"
+        )
+    if out is not None and (
+        out.shape != left.shape
+        or out.device != left.device
+        or out.dtype != left.dtype
+        or not out.is_contiguous()
+    ):
+        raise ValueError("tiled output-grid output must match left")
+
+
+def run_tiled_output_grid_product(
+    left,
+    right,
+    to_grid,
+    from_grid,
+    *,
+    use_sm80_c96_n48: bool = False,
+    use_sm90_c96_asymmetric_panels: bool = False,
+):
+    """Run the fused strict-FP32 tiled output-grid forward."""
+    import torch
+
+    _validate_tensors(left, right, to_grid, from_grid)
+    out = torch.empty_like(left)
+    return run_tiled_output_grid_product_out(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        out,
+        use_sm80_c96_n48=use_sm80_c96_n48,
+        use_sm90_c96_asymmetric_panels=use_sm90_c96_asymmetric_panels,
+    )
+
+
+def run_tiled_output_grid_product_out(
+    left,
+    right,
+    to_grid,
+    from_grid,
+    out,
+    *,
+    use_sm80_c96_n48: bool = False,
+    use_sm90_c96_asymmetric_panels: bool = False,
+):
+    """Run the fused forward into a caller-provided output tensor."""
+    _validate_tensors(left, right, to_grid, from_grid, out)
+    compile_key = _compile_key(left)
+    if use_sm80_c96_n48 and use_sm90_c96_asymmetric_panels:
+        raise ValueError("output-grid panel specializations are mutually exclusive")
+    if use_sm80_c96_n48 and (
+        compile_key[1] not in runtime_policy.SM80_PROFILE_CAPABILITIES
+        or compile_key[2] != 96
+    ):
+        raise ValueError("output-grid forward N=48 requires SM80-family and C=96")
+    if use_sm90_c96_asymmetric_panels and (
+        compile_key[1] != (9, 0) or compile_key[2] != 96
+    ):
+        raise ValueError("output-grid asymmetric forward requires sm90 and C=96")
+    if use_sm90_c96_asymmetric_panels:
+        _compiled_tiled_forward(
+            *compile_key,
+            TILE_N,
+            0,
+            1,
+        )(
+            left,
+            right,
+            to_grid,
+            from_grid,
+            out,
+        )
+        _compiled_tiled_forward(
+            *compile_key,
+            C96_TAIL_TILE_N,
+            C96_TAIL_CHANNEL_TILE,
+            1,
+        )(
+            left,
+            right,
+            to_grid,
+            from_grid,
+            out,
+        )
+        return out
+    tile_n = SM80_C96_TILE_N if use_sm80_c96_n48 else TILE_N
+    _compiled_tiled_forward(*compile_key, tile_n, 0, None)(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        out,
+    )
+    return out
+
+
+def run_tiled_output_grid_product_backward(
+    grad_out,
+    left,
+    right,
+    to_grid,
+    from_grid,
+    *,
+    use_sm80_c96_n48_panel: bool = False,
+    use_sm90_c96_asymmetric_panels: bool = False,
+):
+    """Run the complete fused first backward for left and right inputs."""
+    import torch
+
+    _validate_tensors(left, right, to_grid, from_grid)
+    if (
+        grad_out.shape != left.shape
+        or grad_out.device != left.device
+        or grad_out.dtype != left.dtype
+        or not grad_out.is_contiguous()
+    ):
+        raise ValueError("grad_out must be contiguous and match left")
+    compile_key = _compile_key(left)
+    if use_sm80_c96_n48_panel and use_sm90_c96_asymmetric_panels:
+        raise ValueError("output-grid panel specializations are mutually exclusive")
+    if use_sm80_c96_n48_panel and (
+        compile_key[1] not in runtime_policy.SM80_PROFILE_CAPABILITIES
+        or compile_key[2] != 96
+    ):
+        raise ValueError(
+            "output-grid N=48 panel adjoint requires SM80-family, C=96, and K=8"
+        )
+    if use_sm90_c96_asymmetric_panels and (
+        compile_key[1] != (9, 0) or compile_key[2] != 96
+    ):
+        raise ValueError("output-grid asymmetric backward requires sm90, C=96, and K=8")
+    grad_left = torch.empty_like(left)
+    grad_right = torch.empty_like(right)
+    if use_sm90_c96_asymmetric_panels:
+        _compiled_tiled_backward(
+            *compile_key,
+            TILE_N,
+            0,
+            1,
+        )(
+            grad_out,
+            left,
+            right,
+            to_grid,
+            from_grid,
+            grad_left,
+            grad_right,
+        )
+        _compiled_tiled_backward(
+            *compile_key,
+            C96_TAIL_TILE_N,
+            C96_TAIL_CHANNEL_TILE,
+            1,
+        )(
+            grad_out,
+            left,
+            right,
+            to_grid,
+            from_grid,
+            grad_left,
+            grad_right,
+        )
+        return grad_left, grad_right
+    tile_n = SM80_C96_TILE_N if use_sm80_c96_n48_panel else TILE_N
+    _compiled_tiled_backward(
+        *compile_key,
+        tile_n,
+        0,
+        None,
+    )(
+        grad_out,
+        left,
+        right,
+        to_grid,
+        from_grid,
+        grad_left,
+        grad_right,
+    )
+    return grad_left, grad_right
+
+
+__all__ = [
+    "TiledOutputGridProductBackward",
+    "TiledOutputGridProductForward",
+    "run_tiled_output_grid_product",
+    "run_tiled_output_grid_product_backward",
+    "run_tiled_output_grid_product_out",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid/product.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/product.py
new file mode 100644
index 0000000000..de4e343bf6
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/product.py
@@ -0,0 +1,329 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Strict-FP32 CuTe middle contractions for supported Neo grid MLPs."""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    Any,
+)
+
+import torch
+
+from .. import (
+    runtime_policy,
+)
+from ..runtime_policy import (
+    PORTABLE_TILED_BACKEND,
+    SUPPORTED_HIDDEN_CHANNELS,
+    select_output_grid_backend,
+)
+
+COEFF_DIM = 16
+N_FRAMES = 3
+PACKED_COEFF_DIM = COEFF_DIM * N_FRAMES
+GRID_SIZE = 152
+
+
+def _exact_hidden_channels(
+    left: torch.Tensor,
+    n_frames: int,
+) -> int | None:
+    if left.ndim != 4 or left.shape[0] <= 0 or int(n_frames) != N_FRAMES:
+        return None
+    for hidden_channels in SUPPORTED_HIDDEN_CHANNELS:
+        if tuple(left.shape[1:]) == (
+            COEFF_DIM,
+            1,
+            N_FRAMES * hidden_channels,
+        ):
+            return hidden_channels
+    return None
+
+
+def _has_exact_contract(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+    n_frames: int,
+) -> bool:
+    hidden_channels = _exact_hidden_channels(left, n_frames)
+    if (
+        hidden_channels is None
+        or not left.is_cuda
+        or left.dtype != torch.float32
+        or right.device != left.device
+        or right.dtype != left.dtype
+        or to_grid.device != left.device
+        or from_grid.device != left.device
+        or to_grid.dtype != left.dtype
+        or from_grid.dtype != left.dtype
+        or right.shape != left.shape
+        or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM)
+        or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE)
+        or not left.is_contiguous()
+        or not right.is_contiguous()
+        or not to_grid.is_contiguous()
+        or not from_grid.is_contiguous()
+        or to_grid.requires_grad
+        or from_grid.requires_grad
+        or torch.is_autocast_enabled("cuda")
+        or not runtime_policy.uses_strict_fp32_matmul()
+    ):
+        return False
+    max_intermediate_values = left.shape[0] * GRID_SIZE * hidden_channels
+    if max_intermediate_values > runtime_policy.INT32_MAX:
+        return False
+    compute_capability = tuple(torch.cuda.get_device_capability(left.device))
+    return (
+        select_output_grid_backend(compute_capability, hidden_channels)
+        == PORTABLE_TILED_BACKEND
+    )
+
+
+def _validate_exact_contract(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+    n_frames: int,
+) -> int:
+    hidden_channels = _exact_hidden_channels(left, n_frames)
+    if hidden_channels is None or not _has_exact_contract(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        n_frames,
+    ):
+        raise ValueError(
+            "the fused output GridMLP kernel requires contiguous CUDA FP32 "
+            "tensors with Neo's (D=16, F=1, frames=3, G=152, "
+            "C in {96, 192}) contract"
+        )
+    return hidden_channels
+
+
+def _output_grid_product_impl(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+    n_frames: int,
+) -> torch.Tensor:
+    hidden_channels = _validate_exact_contract(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        n_frames,
+    )
+    from .kernels.tiled_product import (
+        run_tiled_output_grid_product,
+    )
+
+    nodes = left.shape[0]
+    left_flat = left.detach().view(nodes, PACKED_COEFF_DIM, hidden_channels)
+    right_flat = right.detach().view(nodes, PACKED_COEFF_DIM, hidden_channels)
+    compute_capability = tuple(torch.cuda.get_device_capability(left.device))
+    use_sm80_c96_n48 = (
+        hidden_channels == 96
+        and compute_capability in runtime_policy.SM80_PROFILE_CAPABILITIES
+        and runtime_policy.is_output_grid_fwd_sm80_c96_n48_enabled(compute_capability)
+    )
+    use_sm90_c96_asymmetric_panels = (
+        hidden_channels == 96
+        and compute_capability == (9, 0)
+        and runtime_policy.is_output_grid_sm90_c96_asymmetric_panels_enabled(
+            compute_capability
+        )
+    )
+    out = run_tiled_output_grid_product(
+        left_flat,
+        right_flat,
+        to_grid.detach(),
+        from_grid.detach(),
+        use_sm80_c96_n48=use_sm80_c96_n48,
+        use_sm90_c96_asymmetric_panels=use_sm90_c96_asymmetric_panels,
+    )
+    return out.view_as(left)
+
+
+def _output_grid_product_bwd_impl(
+    grad_out: torch.Tensor,
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+    n_frames: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    hidden_channels = _validate_exact_contract(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        n_frames,
+    )
+    if (
+        grad_out.shape != left.shape
+        or grad_out.dtype != left.dtype
+        or grad_out.device != left.device
+    ):
+        raise ValueError("grad_out must match the fused output GridMLP output")
+    from .kernels.tiled_product import (
+        run_tiled_output_grid_product_backward,
+    )
+
+    nodes = left.shape[0]
+    compute_capability = tuple(torch.cuda.get_device_capability(left.device))
+    use_sm80_c96_n48_panel = (
+        hidden_channels == 96
+        and compute_capability in runtime_policy.SM80_PROFILE_CAPABILITIES
+        and runtime_policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(
+            compute_capability
+        )
+    )
+    use_sm90_c96_asymmetric_panels = (
+        hidden_channels == 96
+        and compute_capability == (9, 0)
+        and runtime_policy.is_output_grid_sm90_c96_asymmetric_panels_enabled(
+            compute_capability
+        )
+    )
+    grad_left, grad_right = run_tiled_output_grid_product_backward(
+        grad_out.detach()
+        .contiguous()
+        .view(
+            nodes,
+            PACKED_COEFF_DIM,
+            hidden_channels,
+        ),
+        left.detach().view(nodes, PACKED_COEFF_DIM, hidden_channels),
+        right.detach().view(nodes, PACKED_COEFF_DIM, hidden_channels),
+        to_grid.detach(),
+        from_grid.detach(),
+        use_sm80_c96_n48_panel=use_sm80_c96_n48_panel,
+        use_sm90_c96_asymmetric_panels=use_sm90_c96_asymmetric_panels,
+    )
+    return grad_left.view_as(left), grad_right.view_as(right)
+
+
+_output_grid_product_op = torch.library.custom_op(
+    "sezm_cute::output_grid_product",
+    mutates_args=(),
+)(_output_grid_product_impl)
+_output_grid_product_bwd_op = torch.library.custom_op(
+    "sezm_cute::output_grid_product_bwd",
+    mutates_args=(),
+)(_output_grid_product_bwd_impl)
+
+
+@_output_grid_product_op.register_fake
+def _output_grid_product_fake(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+    n_frames: int,
+) -> torch.Tensor:
+    del right, to_grid, from_grid, n_frames
+    return torch.empty(left.shape, dtype=left.dtype, device=left.device)
+
+
+@_output_grid_product_bwd_op.register_fake
+def _output_grid_product_bwd_fake(
+    grad_out: torch.Tensor,
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+    n_frames: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    del grad_out, to_grid, from_grid, n_frames
+    return (
+        torch.empty(left.shape, dtype=left.dtype, device=left.device),
+        torch.empty(right.shape, dtype=right.dtype, device=right.device),
+    )
+
+
+def _setup_context(
+    ctx: Any,
+    inputs: tuple,
+    output: torch.Tensor,
+) -> None:
+    del output
+    left, right, to_grid, from_grid, n_frames = inputs
+    ctx.save_for_backward(left, right, to_grid, from_grid)
+    ctx.n_frames = int(n_frames)
+
+
+def _backward(ctx: Any, grad_out: torch.Tensor) -> tuple:
+    left, right, to_grid, from_grid = ctx.saved_tensors
+    grad_left, grad_right = _output_grid_product_bwd_op(
+        grad_out,
+        left,
+        right,
+        to_grid,
+        from_grid,
+        ctx.n_frames,
+    )
+    return grad_left, grad_right, None, None, None
+
+
+_output_grid_product_op.register_autograd(
+    _backward,
+    setup_context=_setup_context,
+)
+
+
+def output_grid_product_cute(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+    *,
+    n_frames: int,
+) -> torch.Tensor:
+    """Run one supported exact-shape fused grid contraction."""
+    return _output_grid_product_op(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        int(n_frames),
+    )
+
+
+def maybe_run_cute_output_grid_product(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+    *,
+    n_frames: int,
+) -> torch.Tensor | None:
+    """Return ``None`` unless the master gate and exact Neo contract match."""
+    if not runtime_policy.is_cute_infer_enabled() or not _has_exact_contract(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        n_frames,
+    ):
+        return None
+    return output_grid_product_cute(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        n_frames=n_frames,
+    )
+
+
+__all__ = [
+    "maybe_run_cute_output_grid_product",
+    "output_grid_product_cute",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid/readout_l0.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/readout_l0.py
new file mode 100644
index 0000000000..0a0af62277
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/readout_l0.py
@@ -0,0 +1,803 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Strict-FP32 degree-zero readout for the exact Neo output GridMLP."""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+import torch
+
+from .. import (
+    runtime_policy,
+)
+from ..runtime_policy import (
+    PORTABLE_TILED_BACKEND,
+    select_output_grid_backend,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+COEFF_DIM = 16
+N_FRAMES = 3
+PACKED_COEFF_DIM = 48
+GRID_SIZE = 152
+HIDDEN_CHANNELS = 192
+PACKED_WIDTH = N_FRAMES * HIDDEN_CHANNELS
+_READOUT_INPUT_FOLD_CACHE = "_neo_sm80_readout_input_fold_cache"
+_READOUT_INPUT_FOLD_HOOK = "_neo_sm80_readout_input_fold_hook"
+_READOUT_INPUT_FOLD_LEFT = "_neo_sm80_readout_input_fold_left"
+_READOUT_INPUT_FOLD_RIGHT = "_neo_sm80_readout_input_fold_right"
+_READOUT_INPUT_FOLD_SCALAR = "_neo_sm80_readout_input_fold_scalar"
+_READOUT_INPUT_FOLD_BUFFERS = (
+    _READOUT_INPUT_FOLD_LEFT,
+    _READOUT_INPUT_FOLD_RIGHT,
+    _READOUT_INPUT_FOLD_SCALAR,
+)
+_READOUT_INPUT_FOLD_PREPARE_ERROR = (
+    "the SM80 readout input fold is stale or missing; call "
+    "prepare_sm80_readout_input_fold(output_ffn) after loading, replacing, "
+    "mutating, or moving model state and before torch.compile"
+)
+
+
+def _has_exact_product_shape(left: torch.Tensor) -> bool:
+    return (
+        left.ndim == 4
+        and left.shape[0] > 0
+        and tuple(left.shape[1:]) == (COEFF_DIM, 1, PACKED_WIDTH)
+    )
+
+
+def _has_exact_product_contract(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+) -> bool:
+    tensors = (left, right, to_grid, from_grid)
+    if (
+        not _has_exact_product_shape(left)
+        or right.shape != left.shape
+        or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM)
+        or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE)
+        or any(not tensor.is_cuda for tensor in tensors)
+        or any(tensor.dtype != torch.float32 for tensor in tensors)
+        or any(tensor.device != left.device for tensor in tensors)
+        or any(not tensor.is_contiguous() for tensor in tensors)
+        or to_grid.requires_grad
+        or from_grid.requires_grad
+    ):
+        return False
+    compute_capability = tuple(torch.cuda.get_device_capability(left.device))
+    return (
+        select_output_grid_backend(compute_capability, HIDDEN_CHANNELS)
+        == PORTABLE_TILED_BACKEND
+    )
+
+
+def _validate_product_contract(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+) -> None:
+    if not _has_exact_product_contract(left, right, to_grid, from_grid):
+        raise ValueError(
+            "the readout l=0 kernel requires contiguous CUDA FP32 tensors "
+            "with Neo's left/right=(N,16,1,576), to_grid=(152,48), and "
+            "from_grid=(48,152) contract"
+        )
+
+
+def build_readout_l0_gram(
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+) -> torch.Tensor:
+    """Collapse the frozen row-zero projector into a dense FP32 Gram matrix."""
+    if (
+        tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM)
+        or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE)
+        or to_grid.dtype != torch.float32
+        or from_grid.dtype != torch.float32
+        or to_grid.device != from_grid.device
+        or not to_grid.is_contiguous()
+        or not from_grid.is_contiguous()
+        or to_grid.requires_grad
+        or from_grid.requires_grad
+    ):
+        raise ValueError(
+            "readout l=0 Gram construction requires frozen contiguous FP32 "
+            "to_grid=(152,48) and from_grid=(48,152) tensors"
+        )
+    with torch.no_grad():
+        return torch.matmul(
+            to_grid.T,
+            from_grid[0, :, None] * to_grid,
+        ).contiguous()
+
+
+def has_neo_readout_contract(output_ffn: Any) -> bool:
+    """Return whether an output FFN matches the fused Neo readout contract.
+
+    Parameters
+    ----------
+    output_ffn : Any
+        Candidate equivariant output FFN.
+
+    Returns
+    -------
+    bool
+        Whether the FFN and its grid projection expose the required structure.
+    """
+    if type(output_ffn).__name__ != "EquivariantFFN":
+        return False
+    grid_net = output_ffn.act
+    if (
+        type(grid_net).__name__ != "SO3GridNet"
+        or type(grid_net.grid_op).__name__ != "GridMLP"
+    ):
+        return False
+    grid_op = grid_net.grid_op
+    projector = grid_net.projector
+    return (
+        output_ffn.lmax == 3
+        and output_ffn.channels == 32
+        and output_ffn.hidden_channels == 96
+        and output_ffn.kmax == 1
+        and output_ffn.grid_n_frames == N_FRAMES
+        and output_ffn.use_grid_net
+        and output_ffn.use_grid_mlp
+        and not output_ffn.use_grid_branch
+        and output_ffn.ffn_so3_grid
+        and not output_ffn.s2_activation
+        and not output_ffn.mlp_bias
+        and grid_net.lmax == 3
+        and grid_net.channels == 96
+        and grid_net.n_focus == 1
+        and grid_net.n_frames == N_FRAMES
+        and grid_net.mode == "self"
+        and grid_net.op_type == "mlp"
+        and grid_net.layout == "ndfc"
+        and grid_net.frame_zero_index == 0
+        and grid_net.frames == [0, -1, 1]
+        and grid_net.frame_expand is None
+        and grid_net.frame_contract is None
+        and grid_net.residual_scale is None
+        and grid_op.mode == "self"
+        and grid_op.channels == 96
+        and grid_op.hidden_channels == HIDDEN_CHANNELS
+        and grid_op.n_frames == N_FRAMES
+        and tuple(output_ffn.so3_linear_1.weight.shape) == (4, 32, PACKED_WIDTH)
+        and tuple(grid_op.left_proj.weight.shape) == (HIDDEN_CHANNELS, HIDDEN_CHANNELS)
+        and tuple(grid_op.right_proj.weight.shape) == (HIDDEN_CHANNELS, HIDDEN_CHANNELS)
+        and tuple(grid_op.out_proj.weight.shape) == (HIDDEN_CHANNELS, 96)
+        and tuple(grid_net.scalar_gate.weight.shape) == (HIDDEN_CHANNELS, 96)
+        and tuple(output_ffn.so3_linear_2.weight.shape) == (4, 288, 32)
+        and output_ffn.so3_linear_1.bias is None
+        and grid_op.left_proj.bias is None
+        and grid_op.right_proj.bias is None
+        and grid_op.out_proj.bias is None
+        and grid_net.scalar_gate.bias is None
+        and output_ffn.so3_linear_2.bias is None
+        and tuple(projector.to_grid_mat.shape) == (GRID_SIZE, PACKED_COEFF_DIM)
+        and tuple(projector.from_grid_mat.shape) == (PACKED_COEFF_DIM, GRID_SIZE)
+    )
+
+
+def _state_uses_strict_fp32(output_ffn: Any, device: torch.device) -> bool:
+    for tensor in (*output_ffn.parameters(), *output_ffn.buffers()):
+        if tensor.is_floating_point() and (
+            tensor.dtype != torch.float32
+            or tensor.device != device
+            or not tensor.is_contiguous()
+        ):
+            return False
+    return True
+
+
+def _inference_mode_is_frozen(output_ffn: Any) -> bool:
+    return not output_ffn.training and not any(
+        parameter.requires_grad for parameter in output_ffn.parameters()
+    )
+
+
+@torch.compiler.assume_constant_result
+def _uses_strict_fp32_matmul() -> bool:
+    """Preserve the tested private helper while sharing the runtime policy."""
+    return runtime_policy.uses_strict_fp32_matmul()
+
+
+def _has_exact_neo_readout_contract(
+    output_ffn: Any,
+    ffn_in: torch.Tensor,
+) -> bool:
+    if (
+        not has_neo_readout_contract(output_ffn)
+        or not _inference_mode_is_frozen(output_ffn)
+        or ffn_in.ndim != 4
+        or ffn_in.shape[0] <= 0
+        or tuple(ffn_in.shape[1:]) != (COEFF_DIM, 1, 32)
+        or not ffn_in.is_cuda
+        or ffn_in.dtype != torch.float32
+        or not ffn_in.is_contiguous()
+        or not _state_uses_strict_fp32(output_ffn, ffn_in.device)
+        or torch.is_autocast_enabled("cuda")
+        or not _uses_strict_fp32_matmul()
+    ):
+        return False
+    compute_capability = tuple(torch.cuda.get_device_capability(ffn_in.device))
+    return (
+        select_output_grid_backend(compute_capability, HIDDEN_CHANNELS)
+        == PORTABLE_TILED_BACKEND
+    )
+
+
+def _can_use_sm80_readout_input_fold(
+    output_ffn: Any,
+    ffn_in: torch.Tensor,
+) -> bool:
+    """Fail closed unless the exact frozen strict-FP32 SM80 path is active."""
+    if not _has_exact_neo_readout_contract(output_ffn, ffn_in):
+        return False
+    compute_capability = tuple(torch.cuda.get_device_capability(ffn_in.device))
+    return runtime_policy.is_readout_input_fold_enabled(compute_capability)
+
+
+def _readout_input_fold_sources(output_ffn: Any) -> tuple[torch.Tensor, ...]:
+    grid_net = output_ffn.act
+    grid_op = grid_net.grid_op
+    return (
+        output_ffn.so3_linear_1.weight,
+        grid_op.left_proj.weight,
+        grid_op.right_proj.weight,
+        grid_net.scalar_gate.weight,
+    )
+
+
+def _project_frames(
+    coeff: torch.Tensor,
+    projection: Any,
+    n_frames: int,
+) -> torch.Tensor:
+    """Apply one channel projection independently to every frame."""
+    n_batch, coefficient_dim, n_focus, _ = coeff.shape
+    projected = projection(
+        coeff.reshape(n_batch, coefficient_dim, n_focus, n_frames, -1)
+    )
+    return projected.reshape(n_batch, coefficient_dim, n_focus, -1)
+
+
+def _readout_input_fold_cache_key(
+    sources: tuple[torch.Tensor, ...],
+) -> tuple[tuple[Any, ...], ...]:
+    return tuple(
+        (
+            tensor.data_ptr(),
+            tensor._version,
+            tensor.dtype,
+            tensor.device,
+            tuple(tensor.shape),
+            tuple(tensor.stride()),
+            tensor.storage_offset(),
+        )
+        for tensor in sources
+    )
+
+
+def _readout_input_fold_cache_matches(
+    output_ffn: Any,
+    sources: tuple[torch.Tensor, ...],
+    cache_key: tuple[tuple[Any, ...], ...],
+) -> bool:
+    cache = getattr(output_ffn, _READOUT_INPUT_FOLD_CACHE, None)
+    return (
+        cache is not None
+        and len(cache) == 2
+        and len(cache[0]) == len(sources)
+        and all(
+            cached is current for cached, current in zip(cache[0], sources, strict=True)
+        )
+        and cache[1] == cache_key
+        and all(
+            isinstance(getattr(output_ffn, name, None), torch.Tensor)
+            for name in _READOUT_INPUT_FOLD_BUFFERS
+        )
+    )
+
+
+def _invalidate_sm80_readout_input_fold(output_ffn: Any) -> None:
+    setattr(output_ffn, _READOUT_INPUT_FOLD_CACHE, None)
+    for name in _READOUT_INPUT_FOLD_BUFFERS:
+        if name in output_ffn._buffers:
+            setattr(output_ffn, name, None)
+
+
+def invalidate_neo_readout_input_fold(output_ffn: Any) -> None:
+    """Invalidate frozen readout weights after parameter topology changes."""
+    _invalidate_sm80_readout_input_fold(output_ffn)
+
+
+def _invalidate_sm80_readout_input_fold_after_load(
+    output_ffn: Any,
+    incompatible_keys: Any,
+) -> None:
+    del incompatible_keys
+    _invalidate_sm80_readout_input_fold(output_ffn)
+
+
+def _ensure_sm80_readout_input_fold_load_hook(output_ffn: Any) -> None:
+    if getattr(output_ffn, _READOUT_INPUT_FOLD_HOOK, False):
+        return
+    output_ffn.register_load_state_dict_post_hook(
+        _invalidate_sm80_readout_input_fold_after_load
+    )
+    setattr(output_ffn, _READOUT_INPUT_FOLD_HOOK, True)
+
+
+def _set_nonpersistent_buffer(
+    module: Any,
+    name: str,
+    tensor: torch.Tensor,
+) -> None:
+    if name in module._buffers:
+        setattr(module, name, tensor)
+    else:
+        module.register_buffer(name, tensor, persistent=False)
+
+
+def _synchronize_sm80_readout_input_fold_build(
+    folded_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+) -> None:
+    """Make a newly built CUDA cache safe for every later consumer stream."""
+    device = folded_weights[0].device
+    if device.type != "cuda":
+        return
+    if torch.cuda.is_current_stream_capturing():
+        raise RuntimeError(
+            "prepare the SM80 readout input fold before CUDA graph capture"
+        )
+    ready = torch.cuda.Event()
+    ready.record(torch.cuda.current_stream(device))
+    ready.synchronize()
+
+
+def _build_sm80_readout_input_fold(
+    output_ffn: Any,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Compose the frozen maps without crossing either readout nonlinearity."""
+    input_weight, left_weight, right_weight, scalar_gate_weight = (
+        _readout_input_fold_sources(output_ffn)
+    )
+    grid_net = output_ffn.act
+    with torch.no_grad():
+        split_weight = input_weight.reshape(
+            output_ffn.lmax + 1,
+            output_ffn.channels,
+            2,
+            N_FRAMES,
+            96,
+        )
+        left_input = split_weight[:, :, 0]
+        right_input = split_weight[:, :, 1]
+        per_frame_input = torch.cat((left_input, right_input), dim=-1)
+        left_fold = torch.matmul(per_frame_input, left_weight)
+        right_fold = torch.matmul(per_frame_input, right_weight)
+        projected_left_weight = left_fold.reshape(
+            output_ffn.lmax + 1,
+            output_ffn.channels,
+            -1,
+        ).contiguous()
+        projected_right_weight = right_fold.reshape(
+            output_ffn.lmax + 1,
+            output_ffn.channels,
+            -1,
+        ).contiguous()
+
+        frame_zero = grid_net.frame_zero_index
+        scalar_pair_weight = torch.cat(
+            (
+                left_input[0, :, frame_zero],
+                right_input[0, :, frame_zero],
+            ),
+            dim=-1,
+        )
+        scalar_gate_fold = torch.matmul(scalar_pair_weight, scalar_gate_weight)
+        scalar_aux_weight = torch.cat(
+            (scalar_pair_weight, scalar_gate_fold),
+            dim=-1,
+        ).contiguous()
+    return (
+        projected_left_weight.detach().clone(),
+        projected_right_weight.detach().clone(),
+        scalar_aux_weight.detach().clone(),
+    )
+
+
+@torch.compiler.disable
+def prepare_sm80_readout_input_fold(
+    output_ffn: Any,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Prepare immutable folded weights before compiling the frozen readout."""
+    if torch.compiler.is_compiling():
+        raise RuntimeError(_READOUT_INPUT_FOLD_PREPARE_ERROR)
+    if not has_neo_readout_contract(output_ffn) or not _inference_mode_is_frozen(
+        output_ffn
+    ):
+        raise ValueError("readout input folding requires the exact frozen Neo FFN")
+
+    sources = _readout_input_fold_sources(output_ffn)
+    device = sources[0].device
+    if any(
+        tensor.dtype != torch.float32
+        or tensor.device != device
+        or not tensor.is_contiguous()
+        for tensor in sources
+    ):
+        raise ValueError(
+            "readout input folding requires contiguous FP32 source weights "
+            "on one device"
+        )
+    cache_key = _readout_input_fold_cache_key(sources)
+    if _readout_input_fold_cache_matches(output_ffn, sources, cache_key):
+        return (
+            getattr(output_ffn, _READOUT_INPUT_FOLD_LEFT),
+            getattr(output_ffn, _READOUT_INPUT_FOLD_RIGHT),
+            getattr(output_ffn, _READOUT_INPUT_FOLD_SCALAR),
+        )
+
+    _invalidate_sm80_readout_input_fold(output_ffn)
+    folded_weights = _build_sm80_readout_input_fold(output_ffn)
+    _synchronize_sm80_readout_input_fold_build(folded_weights)
+    _ensure_sm80_readout_input_fold_load_hook(output_ffn)
+    for name, tensor in zip(
+        _READOUT_INPUT_FOLD_BUFFERS,
+        folded_weights,
+        strict=True,
+    ):
+        _set_nonpersistent_buffer(output_ffn, name, tensor)
+    setattr(
+        output_ffn,
+        _READOUT_INPUT_FOLD_CACHE,
+        (sources, cache_key),
+    )
+    return folded_weights
+
+
+@torch.compiler.disable
+def maybe_prepare_sm80_readout_input_fold(
+    output_ffn: Any,
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Prepare only when a supported frozen Neo readout contract matches."""
+    if not has_neo_readout_contract(output_ffn) or not _inference_mode_is_frozen(
+        output_ffn
+    ):
+        return False
+    sources = _readout_input_fold_sources(output_ffn)
+    device = sources[0].device
+    if device.type != "cuda":
+        return False
+    if compute_capability is None:
+        compute_capability = tuple(torch.cuda.get_device_capability(device))
+    if not runtime_policy.is_readout_input_fold_enabled(compute_capability):
+        return False
+    if any(
+        tensor.dtype != torch.float32
+        or tensor.device != device
+        or not tensor.is_contiguous()
+        for tensor in sources
+    ):
+        return False
+    prepare_sm80_readout_input_fold(output_ffn)
+    return True
+
+
+def _get_prepared_sm80_readout_input_fold(
+    output_ffn: Any,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    cache = getattr(output_ffn, _READOUT_INPUT_FOLD_CACHE, None)
+    folded_weights = (
+        getattr(output_ffn, _READOUT_INPUT_FOLD_LEFT, None),
+        getattr(output_ffn, _READOUT_INPUT_FOLD_RIGHT, None),
+        getattr(output_ffn, _READOUT_INPUT_FOLD_SCALAR, None),
+    )
+    if (
+        cache is None
+        or len(cache) != 2
+        or len(cache[0]) != 4
+        or any(not isinstance(weight, torch.Tensor) for weight in folded_weights)
+    ):
+        raise RuntimeError(_READOUT_INPUT_FOLD_PREPARE_ERROR)
+
+    sources = _readout_input_fold_sources(output_ffn)
+    cached_sources, cache_key = cache
+    if (
+        any(
+            cached is not current
+            for cached, current in zip(cached_sources, sources, strict=True)
+        )
+        or any(
+            source.dtype != cached[2]
+            or source.device != cached[3]
+            or tuple(source.shape) != cached[4]
+            for source, cached in zip(sources, cache_key, strict=True)
+        )
+        or any(
+            weight.dtype != torch.float32
+            or weight.device != sources[0].device
+            or not weight.is_contiguous()
+            for weight in folded_weights
+        )
+    ):
+        raise RuntimeError(_READOUT_INPUT_FOLD_PREPARE_ERROR)
+
+    # Do not emit source ``_version`` counters into the graph. AOTAutograd
+    # represents them as unbacked symbolic integers and cannot lower the
+    # resulting assertion. Eager preparation validates versions before trace;
+    # SeZM's load-state hook invalidates both local and shared compiled graphs.
+    return folded_weights
+
+
+def _get_sm80_readout_input_fold(
+    output_ffn: Any,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Return prepared degree weights for operands and scalar auxiliaries."""
+    if torch.compiler.is_compiling():
+        return _get_prepared_sm80_readout_input_fold(output_ffn)
+    return prepare_sm80_readout_input_fold(output_ffn)
+
+
+def _maybe_prepare_sm80_readout_input_fold(
+    output_ffn: Any,
+    ffn_in: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None:
+    """Project directly from the 32-channel input to both product operands."""
+    if not _can_use_sm80_readout_input_fold(output_ffn, ffn_in):
+        return None
+
+    left_weight, right_weight, scalar_weight = _get_sm80_readout_input_fold(output_ffn)
+    expanded_left_weight = left_weight.index_select(
+        0,
+        output_ffn.so3_linear_1.expand_index,
+    )
+    expanded_right_weight = right_weight.index_select(
+        0,
+        output_ffn.so3_linear_1.expand_index,
+    )
+    left = torch.einsum("ndfi,dio->ndfo", ffn_in, expanded_left_weight).contiguous()
+    right = torch.einsum("ndfi,dio->ndfo", ffn_in, expanded_right_weight).contiguous()
+    scalar_aux = torch.einsum("nfi,io->nfo", ffn_in[:, 0], scalar_weight)
+    scalar_pair, scalar_gate_logits = torch.split(scalar_aux, (192, 96), dim=-1)
+    return left, right, scalar_pair, scalar_gate_logits
+
+
+def _run_neo_readout_l0(
+    output_ffn: Any,
+    ffn_in: torch.Tensor,
+    grid_product: Callable[
+        [torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+        torch.Tensor,
+    ],
+) -> torch.Tensor:
+    """Complete the exact Neo readout around one row-zero grid product."""
+    if not has_neo_readout_contract(output_ffn):
+        raise ValueError("readout l=0 completion requires the exact Neo output FFN")
+    grid_net = output_ffn.act
+    grid_op = grid_net.grid_op
+
+    prepared = _maybe_prepare_sm80_readout_input_fold(output_ffn, ffn_in)
+    if prepared is None:
+        projected = output_ffn.so3_linear_1(ffn_in)
+        left, right, scalar_pair = grid_net._prepare_self_pair(projected)
+        shape = (*left.shape[:-1], N_FRAMES, -1)
+        fused = torch.cat(
+            [left.reshape(shape), right.reshape(shape)],
+            dim=-1,
+        ).reshape(*left.shape[:-1], -1)
+        left = _project_frames(fused, grid_op.left_proj, N_FRAMES)
+        right = _project_frames(fused, grid_op.right_proj, N_FRAMES)
+        scalar_gate_logits = None
+    else:
+        left, right, scalar_pair, scalar_gate_logits = prepared
+
+    q0 = grid_product(
+        left,
+        right,
+        grid_net.projector.to_grid_mat,
+        grid_net.projector.from_grid_mat,
+    )
+    if tuple(q0.shape) != (ffn_in.shape[0], HIDDEN_CHANNELS):
+        raise ValueError("readout l=0 grid product must return shape (N,192)")
+
+    q0 = torch.matmul(q0, grid_op.out_proj.weight)
+    scalar_out = grid_net.scalar_act(scalar_pair)[:, 0, :]
+    if scalar_gate_logits is None:
+        scalar_gate_logits = grid_net.scalar_gate(scalar_pair)
+    scalar_gate = torch.sigmoid(scalar_gate_logits)[:, 0, :]
+    scalar_coeff = q0 * scalar_gate + scalar_out
+    output_weight = output_ffn.so3_linear_2.weight[0, :96, :]
+    return ffn_in[:, 0, 0, :] + torch.matmul(scalar_coeff, output_weight)
+
+
+def maybe_run_neo_readout_l0(
+    output_ffn: Any,
+    ffn_in: torch.Tensor,
+) -> torch.Tensor | None:
+    """Return the optimized final `[N,32]` readout or ``None`` for fallback."""
+    if not runtime_policy.is_cute_infer_enabled():
+        return None
+    if not _inference_mode_is_frozen(output_ffn):
+        return None
+    if not _has_exact_neo_readout_contract(output_ffn, ffn_in):
+        return None
+    return _run_neo_readout_l0(output_ffn, ffn_in, readout_l0_product_cute)
+
+
+def run_neo_output_readout(
+    output_ffn: Any,
+    ffn_in: torch.Tensor,
+    *,
+    parameters_frozen: bool = True,
+) -> torch.Tensor:
+    """Return the residual-inclusive `[N,32]` output with generic fallback."""
+    if parameters_frozen:
+        candidate = maybe_run_neo_readout_l0(output_ffn, ffn_in)
+        if candidate is not None:
+            return candidate
+    forward_scalar = getattr(output_ffn, "forward_scalar", None)
+    if forward_scalar is None:
+        forward_scalar = output_ffn.call_scalar
+    return (ffn_in[:, 0:1, :, :] + forward_scalar(ffn_in)).reshape(
+        ffn_in.shape[0], output_ffn.channels
+    )
+
+
+def _readout_l0_impl(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+) -> torch.Tensor:
+    _validate_product_contract(left, right, to_grid, from_grid)
+    from .kernels.readout_l0 import (
+        run_readout_l0,
+    )
+
+    nodes = left.shape[0]
+    q0 = run_readout_l0(
+        left.detach().view(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS),
+        right.detach().view(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS),
+        to_grid.detach(),
+        from_grid.detach(),
+    )
+    return q0
+
+
+def _readout_l0_bwd_impl(
+    dq0: torch.Tensor,
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    _validate_product_contract(left, right, to_grid, from_grid)
+    if (
+        tuple(dq0.shape) != (left.shape[0], HIDDEN_CHANNELS)
+        or dq0.dtype != left.dtype
+        or dq0.device != left.device
+        or not dq0.is_contiguous()
+    ):
+        raise ValueError("dq0 must be contiguous and have shape (N,192)")
+    from .kernels.readout_l0 import (
+        run_readout_l0_backward,
+    )
+
+    nodes = left.shape[0]
+    grad_left, grad_right = run_readout_l0_backward(
+        dq0.detach(),
+        left.detach().view(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS),
+        right.detach().view(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS),
+        to_grid.detach(),
+        from_grid.detach(),
+    )
+    return grad_left.view_as(left), grad_right.view_as(right)
+
+
+_readout_l0_op = torch.library.custom_op(
+    "sezm_cute::readout_l0",
+    mutates_args=(),
+)(_readout_l0_impl)
+_readout_l0_bwd_op = torch.library.custom_op(
+    "sezm_cute::readout_l0_bwd",
+    mutates_args=(),
+)(_readout_l0_bwd_impl)
+
+
+@_readout_l0_op.register_fake
+def _readout_l0_fake(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+) -> torch.Tensor:
+    del right, to_grid, from_grid
+    return torch.empty(
+        (left.shape[0], HIDDEN_CHANNELS),
+        dtype=left.dtype,
+        device=left.device,
+    )
+
+
+@_readout_l0_bwd_op.register_fake
+def _readout_l0_bwd_fake(
+    dq0: torch.Tensor,
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    del dq0, to_grid, from_grid
+    return (
+        torch.empty(left.shape, dtype=left.dtype, device=left.device),
+        torch.empty(right.shape, dtype=right.dtype, device=right.device),
+    )
+
+
+def _setup_context(
+    ctx: Any,
+    inputs: tuple,
+    output: torch.Tensor,
+) -> None:
+    del output
+    left, right, to_grid, from_grid = inputs
+    ctx.save_for_backward(left, right, to_grid, from_grid)
+
+
+def _backward(ctx: Any, dq0: torch.Tensor) -> tuple:
+    left, right, to_grid, from_grid = ctx.saved_tensors
+    grad_left, grad_right = _readout_l0_bwd_op(
+        dq0.contiguous(),
+        left,
+        right,
+        to_grid,
+        from_grid,
+    )
+    return grad_left, grad_right, None, None
+
+
+_readout_l0_op.register_autograd(
+    _backward,
+    setup_context=_setup_context,
+)
+
+
+def readout_l0_product_cute(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+) -> torch.Tensor:
+    """Run the exact-shape C=192 degree-zero grid contraction."""
+    return _readout_l0_op(left, right, to_grid, from_grid)
+
+
+__all__ = [
+    "build_readout_l0_gram",
+    "has_neo_readout_contract",
+    "invalidate_neo_readout_input_fold",
+    "maybe_prepare_sm80_readout_input_fold",
+    "maybe_run_neo_readout_l0",
+    "prepare_sm80_readout_input_fold",
+    "readout_l0_product_cute",
+    "run_neo_output_readout",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py b/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py
new file mode 100644
index 0000000000..320381d149
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py
@@ -0,0 +1,291 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Shared runtime policy for the opt-in Neo CuTe inference path."""
+
+from __future__ import (
+    annotations,
+)
+
+import os
+
+import torch
+
+from deepmd.pt_expt.kernels.utils import (
+    use_cute_infer,
+)
+
+_TRUE_VALUES = frozenset({"1", "true", "yes", "on"})
+_FALSE_VALUES = frozenset({"0", "false", "no", "off"})
+SM80_PROFILE_CAPABILITIES = frozenset({(8, 0), (8, 6)})
+SM90_CAPABILITY = (9, 0)
+FUSED_SO2_GATE_CAPABILITIES = frozenset({(8, 9), (12, 0)})
+OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS_ENV = (
+    "DP_CUTE_OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS"
+)
+READOUT_INPUT_FOLD_SM90_ENV = "DP_CUTE_READOUT_INPUT_FOLD_SM90"
+SUPPORTED_SO2_CAPABILITIES = SM80_PROFILE_CAPABILITIES | frozenset(
+    {(8, 9), (9, 0), (10, 0), (12, 0)}
+)
+_GIE_DEFAULT_CAPABILITIES = SM80_PROFILE_CAPABILITIES
+INT32_MAX = (1 << 31) - 1
+SO2_VALUES_PER_EDGE = 10 * 64
+SO2_VALUES_PER_NODE = 16 * 64
+PORTABLE_TILED_BACKEND = "portable_tiled"
+PYTORCH_BACKEND = "pytorch"
+SUPPORTED_HIDDEN_CHANNELS = (96, 192)
+_OUTPUT_GRID_ARCH_BACKENDS = {
+    "sm80": {
+        96: PORTABLE_TILED_BACKEND,
+        192: PORTABLE_TILED_BACKEND,
+    },
+    "sm90": {
+        96: PORTABLE_TILED_BACKEND,
+        192: PORTABLE_TILED_BACKEND,
+    },
+}
+
+
+def _env_override(name: str) -> bool | None:
+    value = os.environ.get(name)
+    if value is None or not value.strip():
+        return None
+    normalized = value.strip().lower()
+    if normalized in _TRUE_VALUES:
+        return True
+    if normalized in _FALSE_VALUES:
+        return False
+    raise ValueError(
+        f"{name} must be a boolean value (0/1, false/true, no/yes, off/on); "
+        f"got {value!r}"
+    )
+
+
+def is_cute_infer_enabled() -> bool:
+    """Return whether ``DP_CUTE_INFER`` enables the SeZM CuTe path."""
+    return use_cute_infer()
+
+
+def _current_compute_capability() -> tuple[int, int] | None:
+    if not torch.cuda.is_available():
+        return None
+    try:
+        return tuple(torch.cuda.get_device_capability())
+    except RuntimeError:
+        return None
+
+
+def output_grid_arch_key(compute_capability: tuple[int, int]) -> str:
+    """Return the architecture-family key used for output-grid dispatch."""
+    if tuple(compute_capability) in SM80_PROFILE_CAPABILITIES:
+        return "sm80"
+    major, minor = compute_capability
+    return f"sm{int(major)}{int(minor)}"
+
+
+def select_output_grid_backend(
+    compute_capability: tuple[int, int],
+    hidden_channels: int,
+) -> str:
+    """Select a width-specific CuTe or PyTorch output-grid backend."""
+    hidden_channels = int(hidden_channels)
+    if hidden_channels not in SUPPORTED_HIDDEN_CHANNELS:
+        return PYTORCH_BACKEND
+    architecture_backends = _OUTPUT_GRID_ARCH_BACKENDS.get(
+        output_grid_arch_key(compute_capability)
+    )
+    if architecture_backends is None:
+        return PYTORCH_BACKEND
+    return architecture_backends.get(hidden_channels, PYTORCH_BACKEND)
+
+
+def is_sm80_profile_enabled(
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Return whether the shared SM80/SM86 profile is selected."""
+    if compute_capability is None:
+        compute_capability = _current_compute_capability()
+    return (
+        is_cute_infer_enabled()
+        and compute_capability is not None
+        and tuple(compute_capability) in SM80_PROFILE_CAPABILITIES
+    )
+
+
+def _sm80_profile_feature(
+    name: str,
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Apply an SM80-family default with an explicit disable override."""
+    if not is_sm80_profile_enabled(compute_capability):
+        return False
+    return _env_override(name) is not False
+
+
+def _profile_or_explicit_feature(
+    name: str,
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Default on for the SM80 profile; otherwise require explicit opt-in."""
+    if is_sm80_profile_enabled(compute_capability):
+        return _env_override(name) is not False
+    return is_cute_infer_enabled() and _env_override(name) is True
+
+
+@torch.compiler.assume_constant_result
+def is_so2_thin_wrapper_enabled(
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Select compile-visible SO2 dispatch for the SM80 profile."""
+    return _profile_or_explicit_feature(
+        "DP_CUTE_SO2_THIN_WRAPPER",
+        compute_capability,
+    )
+
+
+@torch.compiler.assume_constant_result
+def is_cute_strict_enabled() -> bool:
+    """Return whether expensive CuTe contract assertions are requested."""
+    return _env_override("DP_CUTE_STRICT") is True
+
+
+def is_output_grid_bwd_sm80_c96_n48_panel_enabled(
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Select the C=96, N=48 SM80 panel adjoint."""
+    return _sm80_profile_feature(
+        "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL",
+        compute_capability,
+    )
+
+
+def is_output_grid_fwd_sm80_c96_n48_enabled(
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Select the C=96, N=48 SM80 forward."""
+    return _sm80_profile_feature(
+        "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48",
+        compute_capability,
+    )
+
+
+@torch.compiler.assume_constant_result
+def is_output_grid_sm90_c96_asymmetric_panels_enabled(
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Select the C96 N64+N32 panels only on exact SM90."""
+    if compute_capability is None:
+        compute_capability = _current_compute_capability()
+    return _master_gated_feature(
+        OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS_ENV,
+        default=compute_capability is not None
+        and tuple(compute_capability) == SM90_CAPABILITY,
+    )
+
+
+@torch.compiler.assume_constant_result
+def is_readout_input_fold_sm80_enabled(
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Select the frozen C=192 Neo readout fold on SM80."""
+    return _sm80_profile_feature(
+        "DP_CUTE_READOUT_INPUT_FOLD_SM80",
+        compute_capability,
+    )
+
+
+@torch.compiler.assume_constant_result
+def is_readout_input_fold_sm90_enabled(
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Select the frozen C=192 Neo readout fold only on exact SM90."""
+    if compute_capability is None:
+        compute_capability = _current_compute_capability()
+    return _master_gated_feature(
+        READOUT_INPUT_FOLD_SM90_ENV,
+        default=compute_capability is not None
+        and tuple(compute_capability) == SM90_CAPABILITY,
+    )
+
+
+def is_readout_input_fold_enabled(
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Select the architecture-specific frozen readout fold."""
+    return is_readout_input_fold_sm80_enabled(
+        compute_capability
+    ) or is_readout_input_fold_sm90_enabled(compute_capability)
+
+
+def is_supported_so2_capability(compute_capability: tuple[int, int]) -> bool:
+    """Return whether SO2 supports this compute capability."""
+    return tuple(compute_capability) in SUPPORTED_SO2_CAPABILITIES
+
+
+def so2_int32_indexing_is_safe(
+    *,
+    edge_count: int,
+    node_count: int,
+) -> bool:
+    """Check every flattened SO2 offset represented with signed Int32."""
+    if edge_count < 0 or node_count < 0:
+        return False
+    return (
+        edge_count <= INT32_MAX // SO2_VALUES_PER_EDGE
+        and node_count <= INT32_MAX // SO2_VALUES_PER_NODE
+    )
+
+
+@torch.compiler.assume_constant_result
+def uses_strict_fp32_matmul() -> bool:
+    """Read CUDA matmul precision outside Dynamo and fail closed on TF32."""
+    matmul = torch.backends.cuda.matmul
+    try:
+        precision = matmul.fp32_precision
+    except AttributeError:
+        precision = None
+    except RuntimeError:
+        return False
+    if precision is not None and precision != "none":
+        return precision == "ieee"
+    try:
+        return not matmul.allow_tf32
+    except RuntimeError:
+        return False
+
+
+def _master_gated_feature(name: str, *, default: bool) -> bool:
+    if not is_cute_infer_enabled() or not default:
+        return False
+    override = _env_override(name)
+    return override is not False
+
+
+def is_gie_enabled(compute_capability: tuple[int, int]) -> bool:
+    """Select the optimized geometric-initial-embedding path."""
+    return _master_gated_feature(
+        "DP_CUTE_GIE",
+        default=compute_capability in _GIE_DEFAULT_CAPABILITIES,
+    )
+
+
+def is_packed_wigner_enabled(compute_capability: tuple[int, int]) -> bool:
+    """Select packed Wigner storage required by the optimized SO2 profiles."""
+    return _master_gated_feature(
+        "DP_CUTE_SO2_PACKED_WIGNER",
+        default=compute_capability in SUPPORTED_SO2_CAPABILITIES,
+    )
+
+
+@torch.compiler.assume_constant_result
+def is_so2_eager_island_enabled(
+    compute_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Select the SM80 neighbor-list eager island for the Neo SO2 path."""
+    if not is_cute_infer_enabled():
+        return False
+    override = _env_override("DP_CUTE_SO2_EAGER_ISLANDS")
+    if override is not None:
+        return override
+    if compute_capability is None:
+        compute_capability = _current_compute_capability()
+    return compute_capability is not None and tuple(compute_capability) == (8, 0)
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/so2/__init__.py
new file mode 100644
index 0000000000..896b4aa2d9
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""SO(2) orchestration and kernels for the CuTe SeZM path."""
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/__init__.py
new file mode 100644
index 0000000000..910740da8d
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Portable CuTe DSL kernels for the Neo SO(2) path."""
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/combined_gate_forward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/combined_gate_forward.py
new file mode 100644
index 0000000000..c1654c898c
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/combined_gate_forward.py
@@ -0,0 +1,780 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""One-launch strict-FP32 Neo SO2 gate/residual forward."""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as cute_utils
+import torch
+from cutlass.cute.runtime import (
+    make_fake_stream,
+    make_fake_tensor,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+from ...runtime_policy import (
+    FUSED_SO2_GATE_CAPABILITIES,
+)
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN201, ANN202, ANN204
+
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+TILE_M = 64
+TILE_K = 16
+THREADS = 256
+STAGES = 3
+FOCUS_COUNT = 2
+M0_WIDTH = 4 * 32
+PAIR_WIDTH = 6 * 32
+FULL_WIDTH = M0_WIDTH + PAIR_WIDTH
+DEFAULT_STREAM = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT)
+
+
+def _supports_combined_forward(compute_capability: tuple[int, int]) -> bool:
+    return compute_capability in FUSED_SO2_GATE_CAPABILITIES
+
+
+def _require_16_byte_alignment(tensors: tuple[torch.Tensor, ...]) -> None:
+    if any(tensor.data_ptr() % 16 for tensor in tensors):
+        raise ValueError("combined Neo SO2 gate tensors must be 16-byte aligned")
+
+
+@cute.jit
+def _sigmoid(value):
+    return cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-value))
+
+
+def _fake_focus_tensor(width: int):
+    return make_fake_tensor(
+        cutlass.Float32,
+        (cute.sym_int32(), FOCUS_COUNT, width),
+        (FOCUS_COUNT * width, width, 1),
+        assumed_align=16,
+    )
+
+
+def _fake_focus_weight(width: int):
+    return make_fake_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, width, width),
+        (width * width, width, 1),
+        assumed_align=16,
+    )
+
+
+def prepare_neo_so2_gate_combined_weights(
+    w0: torch.Tensor,
+    wp: torch.Tensor,
+    gate_weight: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Pack immutable weights into contiguous focus-major tensors."""
+    return (
+        w0.transpose(1, 2).contiguous(),
+        wp.transpose(1, 2).contiguous(),
+        gate_weight.permute(1, 0, 2).contiguous(),
+    )
+
+
+class CuteNeoSO2GateCombined:
+    """M64/SO26/T256/S3 SIMT SGEMMs with one CTA-resident gate epilogue."""
+
+    def __init__(self) -> None:
+        self.cta_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=THREADS,
+        )
+
+    @cute.jit
+    def __call__(
+        self,
+        mX: cute.Tensor,
+        mResidual: cute.Tensor,
+        mW0: cute.Tensor,
+        mWP: cute.Tensor,
+        mGate: cute.Tensor,
+        mY: cute.Tensor,
+        mOut: cute.Tensor,
+        stream: cuda.CUstream = DEFAULT_STREAM,
+    ):
+        sA_layout = cute.make_layout(
+            (TILE_M, TILE_K, STAGES),
+            stride=(1, TILE_M + 4, TILE_K * (TILE_M + 4)),
+        )
+        sB_pair_layout = cute.make_layout(
+            (PAIR_WIDTH, TILE_K, STAGES),
+            stride=(1, PAIR_WIDTH + 4, TILE_K * (PAIR_WIDTH + 4)),
+        )
+        sB_m0_layout = cute.make_layout(
+            (M0_WIDTH, TILE_K, STAGES),
+            stride=(1, PAIR_WIDTH + 4, TILE_K * (PAIR_WIDTH + 4)),
+        )
+        sY0_layout = cute.make_layout(
+            (TILE_M, 32),
+            stride=(32, 1),
+        )
+        sGate_layout = cute.make_layout(
+            (TILE_M, 3 * 32),
+            stride=(3 * 32, 1),
+        )
+
+        copy_layout = cute.make_layout(
+            (THREADS // TILE_K, TILE_K),
+            stride=(TILE_K, 1),
+        )
+        copy_value = cute.make_layout((1, 1))
+        copy_a = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            mX.element_type,
+            num_bits_per_copy=mX.element_type.width,
+        )
+        copy_b = cute.make_copy_atom(
+            cute.nvgpu.cpasync.CopyG2SOp(),
+            mW0.element_type,
+            num_bits_per_copy=mW0.element_type.width,
+        )
+        tiled_copy_A = cute.make_tiled_copy_tv(copy_a, copy_layout, copy_value)
+        tiled_copy_B = cute.make_tiled_copy_tv(copy_b, copy_layout, copy_value)
+
+        atoms_layout = cute.make_layout(
+            (THREADS // 16, 16, 1),
+            stride=(16, 1, 0),
+        )
+        permutation_m = cute.make_layout(
+            (atoms_layout.shape[0], 4),
+            stride=(4, 1),
+        )
+        permutation_n = cute.make_layout(
+            (atoms_layout.shape[1], 4),
+            stride=(4, 1),
+        )
+        m0_op = cute.nvgpu.MmaUniversalOp(cutlass.Float32)
+        pair_op = cute.nvgpu.MmaUniversalOp(cutlass.Float32)
+        tiled_mma_m0 = cute.make_tiled_mma(
+            m0_op,
+            atoms_layout,
+            permutation_mnk=(permutation_m, permutation_n, None),
+        )
+        tiled_mma_pair = cute.make_tiled_mma(
+            pair_op,
+            atoms_layout,
+            permutation_mnk=(permutation_m, permutation_n, None),
+        )
+
+        self.kernel(
+            mX,
+            mResidual,
+            mW0,
+            mWP,
+            mGate,
+            mY,
+            mOut,
+            sA_layout,
+            sB_m0_layout,
+            sB_pair_layout,
+            sY0_layout,
+            sGate_layout,
+            tiled_copy_A,
+            tiled_copy_B,
+            tiled_mma_m0,
+            tiled_mma_pair,
+        ).launch(
+            grid=(cute.ceil_div(mY.shape[0], TILE_M), FOCUS_COUNT, 1),
+            block=[THREADS, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mX: cute.Tensor,
+        mResidual: cute.Tensor,
+        mW0: cute.Tensor,
+        mWP: cute.Tensor,
+        mGate: cute.Tensor,
+        mY: cute.Tensor,
+        mOut: cute.Tensor,
+        sA_layout: cute.Layout,
+        sB_m0_layout: cute.Layout,
+        sB_pair_layout: cute.Layout,
+        sY0_layout: cute.Layout,
+        sGate_layout: cute.Layout,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma_m0: cute.TiledMma,
+        tiled_mma_pair: cute.TiledMma,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        edge_tile, focus, _ = cute.arch.block_idx()
+
+        x_focus = mX[None, focus, None]
+        residual_focus = mResidual[None, focus, None]
+        y_focus = mY[None, focus, None]
+        out_focus = mOut[None, focus, None]
+        matrix_layout_m0 = cute.make_layout(
+            (mY.shape[0], M0_WIDTH),
+            stride=(FOCUS_COUNT * FULL_WIDTH, 1),
+        )
+        matrix_layout_pair = cute.make_layout(
+            (mY.shape[0], PAIR_WIDTH),
+            stride=(FOCUS_COUNT * FULL_WIDTH, 1),
+        )
+        mA0 = cute.make_tensor(x_focus.iterator, matrix_layout_m0)
+        mAPair = cute.make_tensor(
+            x_focus.iterator + M0_WIDTH,
+            matrix_layout_pair,
+        )
+        mR0 = cute.make_tensor(residual_focus.iterator, matrix_layout_m0)
+        mRPair = cute.make_tensor(
+            residual_focus.iterator + M0_WIDTH,
+            matrix_layout_pair,
+        )
+        mY0 = cute.make_tensor(y_focus.iterator, matrix_layout_m0)
+        mYPair = cute.make_tensor(
+            y_focus.iterator + M0_WIDTH,
+            matrix_layout_pair,
+        )
+        mOut0 = cute.make_tensor(out_focus.iterator, matrix_layout_m0)
+        mOutPair = cute.make_tensor(
+            out_focus.iterator + M0_WIDTH,
+            matrix_layout_pair,
+        )
+        w0_focus = mW0[focus, None, None]
+        wp_focus = mWP[focus, None, None]
+        gate_focus = mGate[focus, None, None]
+
+        smem = cute_utils.SmemAllocator()
+        sA = smem.allocate_tensor(cutlass.Float32, sA_layout, 16)
+        sB = smem.allocate_tensor(cutlass.Float32, sB_pair_layout, 16)
+        sY0 = smem.allocate_tensor(cutlass.Float32, sY0_layout, 16)
+        sGate = smem.allocate_tensor(cutlass.Float32, sGate_layout, 16)
+
+        self._run_m0(
+            mA0,
+            w0_focus,
+            mR0,
+            gate_focus,
+            mY0,
+            mOut0,
+            sA,
+            sB,
+            sY0,
+            sGate,
+            sB_m0_layout,
+            tiled_copy_A,
+            tiled_copy_B,
+            tiled_mma_m0,
+            tidx,
+            edge_tile,
+        )
+        cute.arch.sync_threads()
+        self._run_pair(
+            mAPair,
+            wp_focus,
+            mRPair,
+            mYPair,
+            mOutPair,
+            sA,
+            sB,
+            sGate,
+            tiled_copy_A,
+            tiled_copy_B,
+            tiled_mma_pair,
+            tidx,
+            edge_tile,
+        )
+
+    @cute.jit
+    def _run_m0(
+        self,
+        mA: cute.Tensor,
+        mB: cute.Tensor,
+        mR: cute.Tensor,
+        mGate: cute.Tensor,
+        mY: cute.Tensor,
+        mOut: cute.Tensor,
+        sA: cute.Tensor,
+        sB: cute.Tensor,
+        sY0: cute.Tensor,
+        sGate: cute.Tensor,
+        sB_m0_layout: cute.Layout,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        edge_tile: cutlass.Int32,
+    ):
+        sB_m0 = cute.make_tensor(sB.iterator, sB_m0_layout)
+        self._run_gemm(
+            mA,
+            mB,
+            mR,
+            mGate,
+            mY,
+            mOut,
+            sA,
+            sB_m0,
+            sY0,
+            sGate,
+            tiled_copy_A,
+            tiled_copy_B,
+            tiled_mma,
+            tidx,
+            edge_tile,
+            m0_block=True,
+        )
+
+    @cute.jit
+    def _run_pair(
+        self,
+        mA: cute.Tensor,
+        mB: cute.Tensor,
+        mR: cute.Tensor,
+        mY: cute.Tensor,
+        mOut: cute.Tensor,
+        sA: cute.Tensor,
+        sB: cute.Tensor,
+        sGate: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        edge_tile: cutlass.Int32,
+    ):
+        self._run_gemm(
+            mA,
+            mB,
+            mR,
+            mB,
+            mY,
+            mOut,
+            sA,
+            sB,
+            sA,
+            sGate,
+            tiled_copy_A,
+            tiled_copy_B,
+            tiled_mma,
+            tidx,
+            edge_tile,
+            m0_block=False,
+        )
+
+    @cute.jit
+    def _run_gemm(
+        self,
+        mA: cute.Tensor,
+        mB: cute.Tensor,
+        mR: cute.Tensor,
+        mGate: cute.Tensor,
+        mY: cute.Tensor,
+        mOut: cute.Tensor,
+        sA: cute.Tensor,
+        sB: cute.Tensor,
+        sY0: cute.Tensor,
+        sGate: cute.Tensor,
+        tiled_copy_A: cute.TiledCopy,
+        tiled_copy_B: cute.TiledCopy,
+        tiled_mma: cute.TiledMma,
+        tidx: cutlass.Int32,
+        edge_tile: cutlass.Int32,
+        m0_block: cutlass.Constexpr[bool],
+    ):
+        width = M0_WIDTH if cutlass.const_expr(m0_block) else PAIR_WIDTH
+        cta_tiler = (TILE_M, width, TILE_K)
+        tiler_coord = (edge_tile, 0, None)
+        thr_mma = tiled_mma.get_slice(tidx)
+
+        gA = cute.local_tile(
+            mA,
+            tiler=cta_tiler,
+            coord=tiler_coord,
+            proj=(1, None, 1),
+        )
+        gB = cute.local_tile(
+            mB,
+            tiler=cta_tiler,
+            coord=tiler_coord,
+            proj=(None, 1, 1),
+        )
+        gR = cute.local_tile(
+            mR,
+            tiler=cta_tiler,
+            coord=tiler_coord,
+            proj=(1, 1, None),
+        )
+        gY = cute.local_tile(
+            mY,
+            tiler=cta_tiler,
+            coord=tiler_coord,
+            proj=(1, 1, None),
+        )
+        gOut = cute.local_tile(
+            mOut,
+            tiler=cta_tiler,
+            coord=tiler_coord,
+            proj=(1, 1, None),
+        )
+
+        thr_copy_A = tiled_copy_A.get_slice(tidx)
+        thr_copy_B = tiled_copy_B.get_slice(tidx)
+        tAgA = thr_copy_A.partition_S(gA)
+        tAsA = thr_copy_A.partition_D(sA)
+        tBgB = thr_copy_B.partition_S(gB)
+        tBsB = thr_copy_B.partition_D(sB)
+
+        mcA = cute.make_identity_tensor(mA.shape)
+        cA = cute.local_tile(
+            mcA,
+            tiler=cta_tiler,
+            coord=tiler_coord,
+            proj=(1, None, 1),
+        )
+        tAcA = thr_copy_A.partition_S(cA)
+        tApA = cute.make_rmem_tensor(
+            cute.make_layout(
+                (
+                    tAsA.shape[0][1],
+                    cute.size(tAsA, mode=[1]),
+                    cute.size(tAsA, mode=[2]),
+                ),
+                stride=(cute.size(tAsA, mode=[1]), 1, 0),
+            ),
+            cutlass.Boolean,
+        )
+        for rest_v in range(tApA.shape[0]):
+            for row in range(tApA.shape[1]):
+                tApA[rest_v, row, 0] = cute.elem_less(
+                    tAcA[(0, rest_v), row, 0, 0][0],
+                    mA.shape[0],
+                )
+
+        k_pipe_max = cute.size(tAsA, mode=[3])
+        k_tile_count = cute.size(tAgA, mode=[3])
+        gmem_pipe_read = cutlass.Int32(0)
+        cute.copy(
+            tiled_copy_A,
+            tAgA[None, None, None, gmem_pipe_read],
+            tAsA[None, None, None, 0],
+            pred=tApA,
+        )
+        cute.copy(
+            tiled_copy_B,
+            tBgB[None, None, None, gmem_pipe_read],
+            tBsB[None, None, None, 0],
+        )
+        cute.arch.cp_async_commit_group()
+        gmem_pipe_read = (
+            gmem_pipe_read + 1
+            if gmem_pipe_read + 1 < k_tile_count
+            else cutlass.Int32(0)
+        )
+        for k_tile in range(1, k_pipe_max - 1):
+            if k_tile < k_tile_count:
+                cute.copy(
+                    tiled_copy_A,
+                    tAgA[None, None, None, gmem_pipe_read],
+                    tAsA[None, None, None, k_tile],
+                    pred=tApA,
+                )
+                cute.copy(
+                    tiled_copy_B,
+                    tBgB[None, None, None, gmem_pipe_read],
+                    tBsB[None, None, None, k_tile],
+                )
+            gmem_pipe_read = (
+                gmem_pipe_read + 1
+                if gmem_pipe_read + 1 < k_tile_count
+                else cutlass.Int32(0)
+            )
+            cute.arch.cp_async_commit_group()
+
+        tCsA = thr_mma.partition_A(sA)
+        tCsB = thr_mma.partition_B(sB)
+        tCgR = thr_mma.partition_C(gR)
+        tCgY = thr_mma.partition_C(gY)
+        tCgOut = thr_mma.partition_C(gOut)
+        tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0])
+        tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0])
+        tCrC = tiled_mma.make_fragment_C(tCgOut)
+        tCrC.fill(0.0)
+
+        smem_pipe_read = cutlass.Int32(0)
+        smem_pipe_write = cutlass.Int32(k_pipe_max - 1)
+        tiles_issued = cutlass.Int32(k_pipe_max - 1)
+        tCsA_p = tCsA[None, None, None, smem_pipe_read]
+        tCsB_p = tCsB[None, None, None, smem_pipe_read]
+        k_block_max = cute.size(tCrA, mode=[2])
+
+        if k_block_max > 1:
+            cute.arch.cp_async_wait_group(k_pipe_max - 2)
+            self.cta_sync_barrier.arrive_and_wait()
+            cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
+            cute.autovec_copy(tCsB_p[None, None, 0], tCrB[None, None, 0])
+
+        for _ in range(k_tile_count):
+            for k_block in range(k_block_max, unroll_full=True):
+                if k_block == k_block_max - 1:
+                    tCsA_p = tCsA[None, None, None, smem_pipe_read]
+                    tCsB_p = tCsB[None, None, None, smem_pipe_read]
+                    cute.arch.cp_async_wait_group(k_pipe_max - 2)
+                    self.cta_sync_barrier.arrive_and_wait()
+
+                k_block_next = (k_block + 1) % k_block_max
+                cute.autovec_copy(
+                    tCsA_p[None, None, k_block_next],
+                    tCrA[None, None, k_block_next],
+                )
+                cute.autovec_copy(
+                    tCsB_p[None, None, k_block_next],
+                    tCrB[None, None, k_block_next],
+                )
+                if k_block == 0:
+                    if tiles_issued < k_tile_count:
+                        cute.copy(
+                            tiled_copy_A,
+                            tAgA[None, None, None, gmem_pipe_read],
+                            tAsA[None, None, None, smem_pipe_write],
+                            pred=tApA,
+                        )
+                cute.gemm(
+                    tiled_mma,
+                    tCrC,
+                    tCrA[None, None, k_block],
+                    tCrB[None, None, k_block],
+                    tCrC,
+                )
+                if k_block == 0:
+                    if tiles_issued < k_tile_count:
+                        cute.copy(
+                            tiled_copy_B,
+                            tBgB[None, None, None, gmem_pipe_read],
+                            tBsB[None, None, None, smem_pipe_write],
+                        )
+                    cute.arch.cp_async_commit_group()
+                    tiles_issued = tiles_issued + 1
+                    smem_pipe_write = smem_pipe_read
+                    smem_pipe_read = smem_pipe_read + 1
+                    if smem_pipe_read == k_pipe_max:
+                        smem_pipe_read = cutlass.Int32(0)
+                    gmem_pipe_read = (
+                        gmem_pipe_read + 1
+                        if gmem_pipe_read + 1 < k_tile_count
+                        else cutlass.Int32(1)
+                    )
+
+        cute.arch.cp_async_wait_group(0)
+        self.cta_sync_barrier.arrive_and_wait()
+        tCrC.store(tCrC.load())
+
+        cC = cute.make_identity_tensor(gOut.shape)
+        tCpC = thr_mma.partition_C(cC)
+        predC = cute.make_rmem_tensor(tCrC.layout, cutlass.Boolean)
+        residue_m = mOut.shape[0] - cutlass.Int32(TILE_M) * edge_tile
+        for idx in range(cute.size(tCrC.shape)):
+            predC[idx] = cute.elem_less(tCpC[idx], (residue_m, width))
+
+        atom = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mOut.element_type,
+        )
+        cute.copy(atom, tCrC, tCgY, pred=predC)
+        tCrR = tiled_mma.make_fragment_C(tCgR)
+        tCrR.fill(0.0)
+        cute.copy(atom, tCgR, tCrR, pred=predC)
+
+        if cutlass.const_expr(m0_block):
+            for idx in range(cute.size(tCrC.shape)):
+                if predC[idx]:
+                    local_row = tCpC[idx][0]
+                    local_col = tCpC[idx][1]
+                    if local_col < 32:
+                        sY0[local_row, local_col] = tCrC[idx].to(cutlass.Float32)
+            cute.arch.sync_threads()
+
+            gate_slots = (TILE_M * 3 * 32 + THREADS - 1) // THREADS
+            for slot in cutlass.range_constexpr(gate_slots):
+                linear_idx = tidx + slot * THREADS
+                if linear_idx < TILE_M * 3 * 32:
+                    local_row = linear_idx // (3 * 32)
+                    gate_idx = linear_idx - local_row * (3 * 32)
+                    global_row = edge_tile * TILE_M + local_row
+                    if global_row < mOut.shape[0]:
+                        gate_logit = cutlass.Float32(0.0)
+                        for k in cutlass.range_constexpr(32):
+                            gate_logit += sY0[local_row, k] * mGate[k, gate_idx]
+                        sGate[local_row, gate_idx] = _sigmoid(gate_logit)
+            cute.arch.sync_threads()
+
+        for idx in range(cute.size(tCrC.shape)):
+            if predC[idx]:
+                local_row = tCpC[idx][0]
+                local_col = tCpC[idx][1]
+                value = tCrC[idx].to(cutlass.Float32)
+                if cutlass.const_expr(m0_block):
+                    if local_col < 32:
+                        value = value * _sigmoid(value)
+                    else:
+                        value = value * sGate[local_row, local_col - 32]
+                else:
+                    gate_idx = local_col
+                    if gate_idx >= 3 * 32:
+                        gate_idx = gate_idx - 3 * 32
+                    value = value * sGate[local_row, gate_idx]
+                tCrC[idx] = value + tCrR[idx].to(cutlass.Float32)
+
+        cute.copy(atom, tCrC, tCgOut, pred=predC)
+
+
+@device_aware_lru_cache(maxsize=8)
+def _compile_combined_forward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> Callable:
+    if not _supports_combined_forward(compute_capability):
+        raise RuntimeError("combined forward requires a supported compute capability")
+    with torch.cuda.device(device_index):
+        fake_x = _fake_focus_tensor(FULL_WIDTH)
+        fake_residual = _fake_focus_tensor(FULL_WIDTH)
+        fake_w0 = _fake_focus_weight(M0_WIDTH)
+        fake_wp = _fake_focus_weight(PAIR_WIDTH)
+        fake_gate = make_fake_tensor(
+            cutlass.Float32,
+            (FOCUS_COUNT, 32, 3 * 32),
+            (32 * 3 * 32, 3 * 32, 1),
+            assumed_align=16,
+        )
+        fake_y = _fake_focus_tensor(FULL_WIDTH)
+        fake_out = _fake_focus_tensor(FULL_WIDTH)
+        operation = CuteNeoSO2GateCombined()
+        fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False)
+        return cute.compile(
+            operation,
+            fake_x,
+            fake_residual,
+            fake_w0,
+            fake_wp,
+            fake_gate,
+            fake_y,
+            fake_out,
+            stream=fake_stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+class CuteNeoSO2GateCombinedFwdRunner:
+    """Prebuilt one-launch forward that writes full y and no global aux."""
+
+    def __init__(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor,
+        y: torch.Tensor,
+        out: torch.Tensor,
+        *,
+        packed_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+        packed_weights_ready: tuple[torch.cuda.Event, int],
+    ) -> None:
+        expected = (x.shape[0], FOCUS_COUNT, 10, 32)
+        if tuple(x.shape) != expected:
+            raise ValueError(f"x must have shape {expected}, got {tuple(x.shape)}")
+        if residual.shape != x.shape or y.shape != x.shape or out.shape != x.shape:
+            raise ValueError("residual, y, and out must match x")
+        if x.shape[0] <= 0:
+            raise ValueError("combined Neo SO2 gate forward requires E > 0")
+        tensors = (x, residual, y, out)
+        if any(
+            tensor.dtype != torch.float32 or not tensor.is_cuda for tensor in tensors
+        ):
+            raise TypeError(
+                "combined Neo SO2 gate forward requires CUDA float32 tensors"
+            )
+        if any(tensor.device != x.device for tensor in tensors):
+            raise ValueError(
+                "all combined Neo SO2 gate forward tensors must share x.device"
+            )
+        if any(not tensor.is_contiguous() for tensor in tensors):
+            raise ValueError(
+                "combined Neo SO2 gate forward requires canonical contiguous tensors"
+            )
+        _require_16_byte_alignment(tensors)
+        device_index = x.device.index
+        if device_index is None:
+            raise RuntimeError("combined Neo SO2 gate forward requires a CUDA index")
+        compute_capability = tuple(torch.cuda.get_device_capability(device_index))
+        if not _supports_combined_forward(compute_capability):
+            raise RuntimeError(
+                "combined forward requires a supported compute capability"
+            )
+
+        packed_w0, packed_wp, packed_gate = packed_weights
+        packed = (packed_w0, packed_wp, packed_gate)
+        if any(
+            tensor.dtype != torch.float32 or not tensor.is_cuda for tensor in packed
+        ):
+            raise TypeError(
+                "packed combined forward weights must be CUDA float32 tensors"
+            )
+        if any(tensor.device != x.device for tensor in packed):
+            raise ValueError("all packed combined forward weights must share x.device")
+        if any(not tensor.is_contiguous() for tensor in packed):
+            raise ValueError("packed combined forward weights must be contiguous")
+        _require_16_byte_alignment(packed)
+        if tuple(packed_w0.shape) != (FOCUS_COUNT, M0_WIDTH, M0_WIDTH):
+            raise ValueError("packed w0 must have shape (2,128,128)")
+        if tuple(packed_wp.shape) != (FOCUS_COUNT, PAIR_WIDTH, PAIR_WIDTH):
+            raise ValueError("packed wp must have shape (2,192,192)")
+        if tuple(packed_gate.shape) != (FOCUS_COUNT, 32, 3 * 32):
+            raise ValueError("packed gate weight must have shape (2,32,96)")
+
+        with torch.cuda.device(x.device):
+            self._compiled = _compile_combined_forward(
+                device_index,
+                compute_capability,
+            )
+        self._args = (
+            x.reshape(x.shape[0], FOCUS_COUNT, FULL_WIDTH),
+            residual.reshape(x.shape[0], FOCUS_COUNT, FULL_WIDTH),
+            packed_w0,
+            packed_wp,
+            packed_gate,
+            y.reshape(x.shape[0], FOCUS_COUNT, FULL_WIDTH),
+            out.reshape(x.shape[0], FOCUS_COUNT, FULL_WIDTH),
+        )
+        self._device = x.device
+        self._packed_weights_ready = packed_weights_ready
+        self._packed_weights_waited_streams: set[int] = set()
+        self.y = y
+        self.out = out
+
+    def __call__(self) -> torch.Tensor:
+        with torch.cuda.device(self._device):
+            torch_stream = torch.cuda.current_stream(self._device)
+            ready_event, producer_stream = self._packed_weights_ready
+            if (
+                torch_stream.cuda_stream != producer_stream
+                and torch_stream.cuda_stream not in self._packed_weights_waited_streams
+            ):
+                torch_stream.wait_event(ready_event)
+                self._packed_weights_waited_streams.add(torch_stream.cuda_stream)
+            stream = cuda.CUstream(torch_stream.cuda_stream)
+            self._compiled(*self._args, stream=stream)
+        return self.out
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/envelope_softmax.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/envelope_softmax.py
new file mode 100644
index 0000000000..d709e3eff2
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/envelope_softmax.py
@@ -0,0 +1,238 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""CuTe DSL kernel for envelope-gated segmented softmax.
+
+The input contract is intentionally strict: destination edges must already be
+sorted and represented by CSR row pointers before calling the CuTe kernel.
+"""
+
+# ruff: noqa: ANN001, ANN201, ANN204, TC002
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+@dataclass(frozen=True)
+class EnvelopeSoftmaxFwdParams:
+    threads: int
+    logits: cute.Tensor
+    edge_gate: cute.Tensor
+    dst_ptr: cute.Tensor
+    z_bias_raw: cute.Tensor
+    out: cute.Tensor
+    group_max: cute.Tensor
+    denom: cute.Tensor
+
+
+class EnvelopeSoftmaxForward:
+    def __init__(self, threads: int):
+        if threads % 32 != 0:
+            raise ValueError("threads must be a multiple of 32")
+        self.threads = threads
+        self.warps = threads // 32
+        self.dtype = cutlass.Float32
+
+    @cute.jit
+    def warp_sum(self, value):
+        return cute.arch.warp_reduction_sum(value)
+
+    @cute.jit
+    def warp_max(self, value):
+        return cute.arch.warp_reduction_max(value)
+
+    @cute.jit
+    def cta_sum(self, value, scratch, tidx):
+        lane = tidx % 32
+        warp = tidx // 32
+        value = self.warp_sum(value)
+        if lane == 0:
+            scratch[warp] = value
+        cute.arch.sync_threads()
+
+        total = self.dtype(0.0)
+        if tidx < self.warps:
+            total = scratch[tidx]
+        total = self.warp_sum(total)
+        if tidx == 0:
+            scratch[0] = total
+        cute.arch.sync_threads()
+        return scratch[0]
+
+    @cute.jit
+    def cta_max(self, value, scratch, tidx):
+        lane = tidx % 32
+        warp = tidx // 32
+        value = self.warp_max(value)
+        if lane == 0:
+            scratch[warp] = value
+        cute.arch.sync_threads()
+
+        neg_large = self.dtype(-3.4028234663852886e38)
+        total = neg_large
+        if tidx < self.warps:
+            total = scratch[tidx]
+        total = self.warp_max(total)
+        if tidx == 0:
+            scratch[0] = total
+        cute.arch.sync_threads()
+        return scratch[0]
+
+    @cute.jit
+    def softplus(self, value):
+        zero = self.dtype(0.0)
+        positive = cute.arch.fmax(value, zero)
+        magnitude = cute.arch.fmax(value, -value)
+        return positive + cute.log(self.dtype(1.0) + cute.exp(-magnitude))
+
+    @cute.kernel
+    def kernel(self, params: EnvelopeSoftmaxFwdParams, eps: cutlass.Constexpr[float]):
+        tidx, _, _ = cute.arch.thread_idx()
+        node, group, _ = cute.arch.block_idx()
+
+        smem = cutlass.utils.SmemAllocator()
+        scratch = smem.allocate_tensor(self.dtype, self.warps)
+
+        lo = params.dst_ptr[node]
+        hi = params.dst_ptr[node + 1]
+        null_mass = self.softplus(params.z_bias_raw[group].to(self.dtype)) + self.dtype(
+            eps
+        )
+        local_max = cute.log(null_mass)
+        for edge in cutlass.range(lo + tidx, hi, self.threads, unroll=1):
+            gate = params.edge_gate[edge].to(self.dtype)
+            if gate < self.dtype(0.0):
+                gate = self.dtype(0.0)
+            if gate > self.dtype(0.0):
+                value = params.logits[edge, group].to(self.dtype) + self.dtype(
+                    2.0
+                ) * cute.log(gate)
+                if value > local_max:
+                    local_max = value
+
+        group_max = self.cta_max(local_max, scratch, tidx)
+        # Every warp must consume scratch[0] before cta_sum reuses it.
+        cute.arch.sync_threads()
+
+        local_sum = self.dtype(0.0)
+        for edge in cutlass.range(lo + tidx, hi, self.threads, unroll=1):
+            gate = params.edge_gate[edge].to(self.dtype)
+            if gate < self.dtype(0.0):
+                gate = self.dtype(0.0)
+            if gate > self.dtype(0.0):
+                effective_logit = params.logits[edge, group].to(
+                    self.dtype
+                ) + self.dtype(2.0) * cute.log(gate)
+                local_sum += cute.exp(effective_logit - group_max)
+
+        denom_sum = self.cta_sum(local_sum, scratch, tidx)
+        denom = denom_sum + null_mass * cute.exp(-group_max)
+
+        if tidx == 0:
+            params.group_max[node, group] = group_max.to(params.group_max.element_type)
+            params.denom[node, group] = denom.to(params.denom.element_type)
+        cute.arch.sync_threads()
+
+        for edge in cutlass.range(lo + tidx, hi, self.threads, unroll=1):
+            gate = params.edge_gate[edge].to(self.dtype)
+            if gate < self.dtype(0.0):
+                gate = self.dtype(0.0)
+            alpha = self.dtype(0.0)
+            if gate > self.dtype(0.0):
+                effective_logit = params.logits[edge, group].to(
+                    self.dtype
+                ) + self.dtype(2.0) * cute.log(gate)
+                num = cute.exp(effective_logit - group_max)
+                alpha = num / denom
+            params.out[edge, group] = alpha.to(params.out.element_type)
+
+
+@cute.jit
+def envelope_softmax_forward_jit(
+    logits: cute.Tensor,
+    edge_gate: cute.Tensor,
+    dst_ptr: cute.Tensor,
+    z_bias_raw: cute.Tensor,
+    out: cute.Tensor,
+    group_max: cute.Tensor,
+    denom: cute.Tensor,
+    threads: cutlass.Constexpr[int],
+    eps: cutlass.Constexpr[float],
+    stream: CUstream,
+):
+    params = EnvelopeSoftmaxFwdParams(
+        threads=threads,
+        logits=logits,
+        edge_gate=edge_gate,
+        dst_ptr=dst_ptr,
+        z_bias_raw=z_bias_raw,
+        out=out,
+        group_max=group_max,
+        denom=denom,
+    )
+    n_nodes, groups = denom.shape
+    EnvelopeSoftmaxForward(threads).kernel(params, eps).launch(
+        grid=[n_nodes, groups, 1],
+        block=[threads, 1, 1],
+        stream=stream,
+    )
+
+
+@device_aware_lru_cache(maxsize=16)
+def compile_envelope_softmax_forward(threads: int, eps: float = 1.0e-7) -> Callable:
+    e = cute.sym_int64()
+    n = cute.sym_int64()
+    g = cute.sym_int64()
+    fake_logits = make_fake_compact_tensor(cutlass.Float32, (e, g), stride_order=(1, 0))
+    fake_gate = make_fake_compact_tensor(cutlass.Float32, (e,), stride_order=(0,))
+    fake_dst_ptr = make_fake_compact_tensor(
+        cutlass.Int32, (cute.sym_int64(),), stride_order=(0,)
+    )
+    fake_z = make_fake_compact_tensor(cutlass.Float32, (g,), stride_order=(0,))
+    fake_out = make_fake_compact_tensor(cutlass.Float32, (e, g), stride_order=(1, 0))
+    fake_group_max = make_fake_compact_tensor(
+        cutlass.Float32, (n, g), stride_order=(1, 0)
+    )
+    fake_denom = make_fake_compact_tensor(cutlass.Float32, (n, g), stride_order=(1, 0))
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        envelope_softmax_forward_jit,
+        fake_logits,
+        fake_gate,
+        fake_dst_ptr,
+        fake_z,
+        fake_out,
+        fake_group_max,
+        fake_denom,
+        threads,
+        eps,
+        fake_stream,
+        options="--enable-tvm-ffi",
+    )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py
new file mode 100644
index 0000000000..061e6c2179
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py
@@ -0,0 +1,275 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035
+"""CuTe forward for Neo attention-prelude source features.
+
+The forward kernel fuses two independent PyTorch producer chains in one launch:
+
+    optional focus RMSNorm -> two-focus logits -> softmax -> label smoothing
+    scalar Q/K RMSNorm -> Q projection + K projection
+
+The kernel is specialized to the Neo SO2 shape `(E, F=2, C=32)`.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from functools import (
+    lru_cache,
+)
+from typing import (
+    Callable,
+)
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+
+@cute.jit
+def _warp_sum(value):
+    return cute.arch.warp_reduction_sum(value)
+
+
+@cute.jit
+def neo_attention_prelude_forward_jit(
+    focus: cute.Tensor,
+    x_l0: cute.Tensor,
+    focus_weight: cute.Tensor,
+    focus_scale: cute.Tensor,
+    q_weight: cute.Tensor,
+    k_weight: cute.Tensor,
+    qk_scale: cute.Tensor,
+    focus_alpha: cute.Tensor,
+    q_node: cute.Tensor,
+    k_node: cute.Tensor,
+    stream: CUstream,
+    focus_eps: cutlass.Float32,
+    qk_eps: cutlass.Float32,
+    tau: cutlass.Float32,
+    label_smoothing: cutlass.Float32,
+    use_focus_norm: cutlass.Constexpr[bool],
+):
+    edges, _ = focus.shape
+    nodes, _, _ = x_l0.shape
+    neo_attention_prelude_forward_kernel(
+        focus,
+        x_l0,
+        focus_weight,
+        focus_scale,
+        q_weight,
+        k_weight,
+        qk_scale,
+        focus_alpha,
+        q_node,
+        k_node,
+        focus_eps,
+        qk_eps,
+        tau,
+        label_smoothing,
+        use_focus_norm,
+    ).launch(
+        # The launch covers two independent domains.  Using their sum avoids a
+        # host-side branch on symbolic E/N while guaranteeing coverage of both.
+        grid=[cute.ceil_div(edges + nodes, 8), 1, 1],
+        block=[256, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_attention_prelude_forward_kernel(
+    focus: cute.Tensor,
+    x_l0: cute.Tensor,
+    focus_weight: cute.Tensor,
+    focus_scale: cute.Tensor,
+    q_weight: cute.Tensor,
+    k_weight: cute.Tensor,
+    qk_scale: cute.Tensor,
+    focus_alpha: cute.Tensor,
+    q_node: cute.Tensor,
+    k_node: cute.Tensor,
+    focus_eps: cutlass.Float32,
+    qk_eps: cutlass.Float32,
+    tau: cutlass.Float32,
+    label_smoothing: cutlass.Float32,
+    use_focus_norm: cutlass.Constexpr[bool],
+):
+    tid, _, _ = cute.arch.thread_idx()
+    block, _, _ = cute.arch.block_idx()
+    lane = tid % 32
+    edge = block * 8 + tid // 32
+    edges, _ = focus.shape
+
+    if edge < edges:
+        x0 = focus[edge, lane].to(cutlass.Float32)
+        x1 = focus[edge, 32 + lane].to(cutlass.Float32)
+        if cutlass.const_expr(use_focus_norm):
+            inv0 = cute.rsqrt(_warp_sum(x0 * x0) / cutlass.Float32(32.0) + focus_eps)
+            inv1 = cute.rsqrt(_warp_sum(x1 * x1) / cutlass.Float32(32.0) + focus_eps)
+            norm0 = x0 * inv0 * focus_scale[0, lane].to(cutlass.Float32)
+            norm1 = x1 * inv1 * focus_scale[1, lane].to(cutlass.Float32)
+        else:
+            norm0 = x0
+            norm1 = x1
+        logit0 = _warp_sum(norm0 * focus_weight[lane, 0].to(cutlass.Float32))
+        logit1 = _warp_sum(norm1 * focus_weight[lane, 1].to(cutlass.Float32))
+
+        if lane == 0:
+            z0 = logit0 / tau
+            z1 = logit1 / tau
+            zmax = z0
+            if z1 > zmax:
+                zmax = z1
+            e0 = cute.exp(z0 - zmax)
+            e1 = cute.exp(z1 - zmax)
+            denom = e0 + e1
+            keep = cutlass.Float32(1.0) - label_smoothing
+            smooth = label_smoothing / cutlass.Float32(2.0)
+            focus_alpha[edge, 0] = (e0 / denom * keep + smooth).to(
+                focus_alpha.element_type
+            )
+            focus_alpha[edge, 1] = (e1 / denom * keep + smooth).to(
+                focus_alpha.element_type
+            )
+
+    # Q/K is a per-node chain, not a child of the per-edge focus chain.  Keep
+    # this guard independent so every node row is initialized even when E < N.
+    nodes, _, _ = x_l0.shape
+    node = block * 8 + tid // 32
+    if node < nodes:
+        qk_x0 = x_l0[node, 0, lane].to(cutlass.Float32)
+        qk_x1 = x_l0[node, 1, lane].to(cutlass.Float32)
+        qk_norm0 = (
+            qk_x0
+            * cute.rsqrt(_warp_sum(qk_x0 * qk_x0) / cutlass.Float32(32.0) + qk_eps)
+            * qk_scale[0, lane].to(cutlass.Float32)
+        )
+        qk_norm1 = (
+            qk_x1
+            * cute.rsqrt(_warp_sum(qk_x1 * qk_x1) / cutlass.Float32(32.0) + qk_eps)
+            * qk_scale[1, lane].to(cutlass.Float32)
+        )
+        q0 = cutlass.Float32(0.0)
+        k0 = cutlass.Float32(0.0)
+        q1 = cutlass.Float32(0.0)
+        so2 = cutlass.Float32(0.0)
+        for input_channel in cutlass.range_constexpr(32):
+            value0 = cute.arch.shuffle_sync(qk_norm0, input_channel)
+            value1 = cute.arch.shuffle_sync(qk_norm1, input_channel)
+            q0 += value0 * q_weight[input_channel, 0, lane].to(cutlass.Float32)
+            k0 += value0 * k_weight[input_channel, 0, lane].to(cutlass.Float32)
+            q1 += value1 * q_weight[input_channel, 1, lane].to(cutlass.Float32)
+            so2 += value1 * k_weight[input_channel, 1, lane].to(cutlass.Float32)
+        q_node[node, 0, lane] = q0.to(q_node.element_type)
+        k_node[node, 0, lane] = k0.to(k_node.element_type)
+        q_node[node, 1, lane] = q1.to(q_node.element_type)
+        k_node[node, 1, lane] = so2.to(k_node.element_type)
+
+
+@lru_cache(maxsize=8)
+def compile_neo_attention_prelude_forward(
+    focus_eps: float,
+    qk_eps: float,
+    tau: float,
+    label_smoothing: float,
+    compile_identity: tuple[int, int, int] | None = None,
+    *,
+    use_focus_norm: bool = True,
+) -> Callable:
+    # The identity keeps independently compiled device/architecture binaries in
+    # distinct cache entries. Compilation itself runs under the runner's device.
+    del compile_identity
+    edges = cute.sym_int64()
+    nodes = cute.sym_int64()
+    fake_focus = make_fake_compact_tensor(
+        cutlass.Float32, (edges, 64), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    fake_x_l0 = make_fake_compact_tensor(
+        cutlass.Float32, (nodes, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    fake_focus_weight = make_fake_compact_tensor(
+        cutlass.Float32, (32, 2), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    fake_focus_scale = make_fake_compact_tensor(
+        cutlass.Float32, (2, 32), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    fake_q_weight = make_fake_compact_tensor(
+        cutlass.Float32, (32, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    fake_k_weight = make_fake_compact_tensor(
+        cutlass.Float32, (32, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    fake_qk_scale = make_fake_compact_tensor(
+        cutlass.Float32, (2, 32), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    fake_focus_alpha = make_fake_compact_tensor(
+        cutlass.Float32, (edges, 2), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    fake_q_node = make_fake_compact_tensor(
+        cutlass.Float32, (nodes, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    fake_k_node = make_fake_compact_tensor(
+        cutlass.Float32, (nodes, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    compiled = cute.compile(
+        neo_attention_prelude_forward_jit,
+        fake_focus,
+        fake_x_l0,
+        fake_focus_weight,
+        fake_focus_scale,
+        fake_q_weight,
+        fake_k_weight,
+        fake_qk_scale,
+        fake_focus_alpha,
+        fake_q_node,
+        fake_k_node,
+        fake_stream,
+        cutlass.Float32(focus_eps),
+        cutlass.Float32(qk_eps),
+        cutlass.Float32(tau),
+        cutlass.Float32(label_smoothing),
+        bool(use_focus_norm),
+        options="--enable-tvm-ffi",
+    )
+
+    def run(
+        focus,
+        x_l0,
+        focus_weight,
+        focus_scale,
+        q_weight,
+        k_weight,
+        qk_scale,
+        focus_alpha,
+        q_node,
+        k_node,
+    ):
+        return compiled(
+            focus,
+            x_l0,
+            focus_weight,
+            focus_scale,
+            q_weight,
+            k_weight,
+            qk_scale,
+            focus_alpha,
+            q_node,
+            k_node,
+            cutlass.Float32(focus_eps),
+            cutlass.Float32(qk_eps),
+            cutlass.Float32(tau),
+            cutlass.Float32(label_smoothing),
+        )
+
+    return run
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/gate_linear_residual_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/gate_linear_residual_backward.py
new file mode 100644
index 0000000000..f0b21e1c18
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/gate_linear_residual_backward.py
@@ -0,0 +1,169 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""CuTe Neo gate-linear + gate/residual backward without saved logits."""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN201, ANN202, TC002
+
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+ROWS_PER_BLOCK = 2
+
+
+@cute.jit
+def _sigmoid(value):
+    return cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-value))
+
+
+@cute.jit
+def neo_gate_linear_residual_backward_fused_jit(
+    grad_out: cute.Tensor,
+    y: cute.Tensor,
+    gate_weight: cute.Tensor,
+    grad_y: cute.Tensor,
+    stream: CUstream,
+    rows_per_block: cutlass.Constexpr[int],
+):
+    rows, _ = y.shape
+    neo_gate_linear_residual_backward_fused_kernel(
+        grad_out,
+        y,
+        gate_weight,
+        grad_y,
+        rows_per_block,
+    ).launch(
+        grid=[cute.ceil_div(rows, rows_per_block), 1, 1],
+        block=[32 * rows_per_block, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_gate_linear_residual_backward_fused_kernel(
+    grad_out: cute.Tensor,
+    y: cute.Tensor,
+    gate_weight: cute.Tensor,
+    grad_y: cute.Tensor,
+    rows_per_block: cutlass.Constexpr[int],
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    block_row, _, _ = cute.arch.block_idx()
+    row_slot = tidx // 32
+    channel = tidx - row_slot * 32
+    row = block_row * rows_per_block + row_slot
+    rows, _ = y.shape
+
+    smem = cutlass.utils.SmemAllocator()
+    grad_logits = smem.allocate_tensor(cutlass.Float32, rows_per_block * 3 * 32)
+    smem_base = row_slot * 3 * 32
+
+    for gate_degree in cutlass.range_constexpr(3):
+        grad_logits[smem_base + gate_degree * 32 + channel] = cutlass.Float32(0.0)
+    cute.arch.sync_threads()
+
+    grad_l0 = cutlass.Float32(0.0)
+    if row < rows:
+        focus = row - (row // 2) * 2
+        y0 = y[row, channel].to(cutlass.Float32)
+        sig0 = _sigmoid(y0)
+        grad0 = grad_out[row, channel].to(cutlass.Float32)
+        grad_l0 = (
+            grad0 * sig0 * (cutlass.Float32(1.0) + y0 * (cutlass.Float32(1.0) - sig0))
+        )
+
+        gate0_logit = cutlass.Float32(0.0)
+        gate1_logit = cutlass.Float32(0.0)
+        gate2_logit = cutlass.Float32(0.0)
+        for k in cutlass.range_constexpr(32):
+            src = y[row, k].to(cutlass.Float32)
+            gate0_logit += src * gate_weight[k, focus, channel].to(cutlass.Float32)
+            gate1_logit += src * gate_weight[k, focus, 32 + channel].to(cutlass.Float32)
+            gate2_logit += src * gate_weight[k, focus, 64 + channel].to(cutlass.Float32)
+
+        gate0 = _sigmoid(gate0_logit)
+        gate1 = _sigmoid(gate1_logit)
+        gate2 = _sigmoid(gate2_logit)
+        for d in cutlass.range_constexpr(1, 10, 1):
+            gate_idx = channel
+            gate = gate0
+            if cutlass.const_expr((d - 1) % 3 == 1):
+                gate_idx = 32 + channel
+                gate = gate1
+            if cutlass.const_expr((d - 1) % 3 == 2):
+                gate_idx = 64 + channel
+                gate = gate2
+            idx = d * 32 + channel
+            gout = grad_out[row, idx].to(cutlass.Float32)
+            yv = y[row, idx].to(cutlass.Float32)
+            grad_y[row, idx] = (gout * gate).to(grad_y.element_type)
+            old = grad_logits[smem_base + gate_idx]
+            grad_logits[smem_base + gate_idx] = old + gout * yv * gate * (
+                cutlass.Float32(1.0) - gate
+            )
+    cute.arch.sync_threads()
+
+    if row < rows:
+        focus = row - (row // 2) * 2
+        grad_gate_src = cutlass.Float32(0.0)
+        for out_idx in cutlass.range_constexpr(3 * 32):
+            grad_gate_src += grad_logits[smem_base + out_idx] * gate_weight[
+                channel, focus, out_idx
+            ].to(cutlass.Float32)
+        grad_y[row, channel] = (grad_l0 + grad_gate_src).to(grad_y.element_type)
+
+
+@device_aware_lru_cache(maxsize=2)
+def compile_neo_gate_linear_residual_backward_fused() -> Callable:
+    rows = cute.sym_int64()
+    fake_grad_out = make_fake_compact_tensor(
+        cutlass.Float32, (rows, 10 * 32), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    fake_y = make_fake_compact_tensor(
+        cutlass.Float32, (rows, 10 * 32), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    fake_gate_weight = make_fake_compact_tensor(
+        cutlass.Float32, (32, 2, 3 * 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    fake_grad_y = make_fake_compact_tensor(
+        cutlass.Float32, (rows, 10 * 32), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        neo_gate_linear_residual_backward_fused_jit,
+        fake_grad_out,
+        fake_y,
+        fake_gate_weight,
+        fake_grad_y,
+        fake_stream,
+        ROWS_PER_BLOCK,
+        options="--enable-tvm-ffi",
+    )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/message_grid_product.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/message_grid_product.py
new file mode 100644
index 0000000000..72cc3bbe3d
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/message_grid_product.py
@@ -0,0 +1,275 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Packed strict-FP32 CuTe grid product for Neo's F=2 SO2 branch."""
+
+from __future__ import (
+    annotations,
+)
+
+from collections.abc import (
+    Callable,
+)
+from functools import (
+    lru_cache,
+)
+from typing import (
+    Any,
+)
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ... import (
+    runtime_policy,
+)
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN201, TC003
+
+
+PACKED_COEFF_DIM = 48
+HIDDEN_CHANNELS = 64
+GRID_SIZE = 152
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+_SUPPORTED_CAPABILITIES = runtime_policy.SM80_PROFILE_CAPABILITIES | {
+    runtime_policy.SM90_CAPABILITY
+}
+
+
+def _make_grid_operation(
+    operation_type: type[Any],
+    *,
+    panel_adjoint: bool = False,
+) -> Any:
+    from ...output_grid.kernels import tiled_product as tiled
+
+    # F=2 and C=32 form one complete 64-channel panel. Bypass only the
+    # shared readout policy; the tiled implementation itself is unchanged.
+    operation = operation_type.__new__(operation_type)
+    operation.hidden_channels = HIDDEN_CHANNELS
+    operation.tile_k = tiled.TILE_K
+    operation.sm80_c96_n48_panel = bool(panel_adjoint)
+    operation.cta_tiler = (tiled.TILE_M, tiled.TILE_N, operation.tile_k)
+    operation.channel_tile_start = 0
+    operation.channel_tiles = 1
+    operation.has_channel_residue = False
+    operation.cta_sync_barrier = pipeline.NamedBarrier(
+        barrier_id=1,
+        num_threads=tiled.THREADS,
+    )
+    return operation
+
+
+def _validate_compile_target(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> None:
+    import torch
+
+    actual = tuple(torch.cuda.get_device_capability(device_index))
+    if actual != tuple(compute_capability):
+        raise ValueError("compile target does not match the selected CUDA device")
+    if actual not in _SUPPORTED_CAPABILITIES:
+        raise ValueError(
+            "packed Neo message-grid product requires the SM80-family profile or sm90"
+        )
+
+
+def _fake_inputs() -> tuple[Any, Any, Any]:
+    nodes = cute.sym_int64()
+    coeff = make_fake_compact_tensor(
+        cutlass.Float32,
+        (nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    to_grid = make_fake_compact_tensor(
+        cutlass.Float32,
+        (GRID_SIZE, PACKED_COEFF_DIM),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    from_grid = make_fake_compact_tensor(
+        cutlass.Float32,
+        (PACKED_COEFF_DIM, GRID_SIZE),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    return coeff, to_grid, from_grid
+
+
+def compile_message_grid_product_forward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> Callable:
+    """Compile the F=2 packed product with a symbolic node count."""
+    import torch
+
+    _validate_compile_target(device_index, compute_capability)
+    from ...output_grid.kernels.tiled_product import (
+        TiledOutputGridProductForward,
+    )
+
+    coeff, to_grid, from_grid = _fake_inputs()
+    stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    with torch.cuda.device(device_index):
+        return cute.compile(
+            _make_grid_operation(TiledOutputGridProductForward),
+            coeff,
+            coeff,
+            to_grid,
+            from_grid,
+            coeff,
+            stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+def compile_message_grid_product_backward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> Callable:
+    """Compile both packed input adjoints with a symbolic node count."""
+    import torch
+
+    _validate_compile_target(device_index, compute_capability)
+    from ...output_grid.kernels.tiled_product import (
+        TiledOutputGridProductBackward,
+    )
+
+    coeff, to_grid, from_grid = _fake_inputs()
+    stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    with torch.cuda.device(device_index):
+        return cute.compile(
+            _make_grid_operation(
+                TiledOutputGridProductBackward,
+                panel_adjoint=(
+                    compute_capability in runtime_policy.SM80_PROFILE_CAPABILITIES
+                ),
+            ),
+            coeff,
+            coeff,
+            coeff,
+            to_grid,
+            from_grid,
+            coeff,
+            coeff,
+            stream,
+            options="--enable-tvm-ffi",
+        )
+
+
+@lru_cache(maxsize=8)
+def _compiled_forward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> Callable:
+    return compile_message_grid_product_forward(device_index, compute_capability)
+
+
+@lru_cache(maxsize=8)
+def _compiled_backward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> Callable:
+    return compile_message_grid_product_backward(device_index, compute_capability)
+
+
+def _compile_identity(tensor) -> tuple[int, tuple[int, int]]:
+    import torch
+
+    device_index = tensor.device.index
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    return int(device_index), tuple(torch.cuda.get_device_capability(device_index))
+
+
+def _validate_tensors(left, right, to_grid, from_grid, grad_out=None) -> None:
+    import torch
+
+    floating = (left, right, to_grid, from_grid)
+    if (
+        tuple(left.shape[1:]) != (PACKED_COEFF_DIM, HIDDEN_CHANNELS)
+        or left.shape[0] <= 0
+        or right.shape != left.shape
+        or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM)
+        or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE)
+        or any(not tensor.is_cuda for tensor in floating)
+        or any(tensor.device != left.device for tensor in floating)
+        or any(tensor.dtype != torch.float32 for tensor in floating)
+        or any(not tensor.is_contiguous() for tensor in floating)
+        or any(tensor.data_ptr() % 16 != 0 for tensor in floating)
+        or tuple(torch.cuda.get_device_capability(left.device))
+        not in _SUPPORTED_CAPABILITIES
+        or not runtime_policy.uses_strict_fp32_matmul()
+    ):
+        raise ValueError(
+            "packed message-grid product requires contiguous SM80-family/SM90 FP32 "
+            "left/right=(N,48,64), to_grid=(152,48), and from_grid=(48,152)"
+        )
+    if grad_out is not None and (
+        grad_out.shape != left.shape
+        or grad_out.device != left.device
+        or grad_out.dtype != torch.float32
+        or not grad_out.is_contiguous()
+        or grad_out.data_ptr() % 16 != 0
+    ):
+        raise ValueError(
+            "packed message-grid grad_out must be contiguous and match left"
+        )
+
+
+def run_message_grid_product(left, right, to_grid, from_grid):
+    """Run the F=2 product without materializing projection-layout clones."""
+    import torch
+
+    _validate_tensors(left, right, to_grid, from_grid)
+    out = torch.empty_like(left)
+    with torch.cuda.device(left.device):
+        _compiled_forward(*_compile_identity(left))(
+            left,
+            right,
+            to_grid,
+            from_grid,
+            out,
+        )
+    return out
+
+
+def run_message_grid_product_backward(
+    grad_out,
+    left,
+    right,
+    to_grid,
+    from_grid,
+):
+    """Run both F=2 product input adjoints in the packed layout."""
+    import torch
+
+    _validate_tensors(left, right, to_grid, from_grid, grad_out)
+    grad_left = torch.empty_like(left)
+    grad_right = torch.empty_like(right)
+    with torch.cuda.device(left.device):
+        _compiled_backward(*_compile_identity(left))(
+            grad_out,
+            left,
+            right,
+            to_grid,
+            from_grid,
+            grad_left,
+            grad_right,
+        )
+    return grad_left, grad_right
+
+
+__all__ = [
+    "compile_message_grid_product_backward",
+    "compile_message_grid_product_forward",
+    "run_message_grid_product",
+    "run_message_grid_product_backward",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/output_gate_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/output_gate_backward.py
new file mode 100644
index 0000000000..e368e32d40
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/output_gate_backward.py
@@ -0,0 +1,204 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Manual backward for the fused Neo attention output gate.
+
+Only input gradients are produced.  The gate and RMSNorm parameter gradients
+are intentionally omitted for the E/F/S path.  ``grad_phase`` may alias
+``grad_gated``; ``grad_x_wide`` is an existing accumulation buffer and only its
+first 64 values per node are updated.
+
+The forward gate is recomputed.  Its logit gradient uses
+``sum(grad_gated * gated_out) * (1 - gate)``, so backward does not need either
+an ungated Phase-C aggregate or a saved gate tensor.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    Callable,
+)
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035
+
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+DEGREE_COUNT = 16
+FOCUS_COUNT = 2
+CHANNELS = 32
+HIDDEN = FOCUS_COUNT * CHANNELS
+OUTPUT_WIDTH = DEGREE_COUNT * HIDDEN
+
+
+@cute.jit
+def _warp_sum(value):
+    return cute.arch.warp_reduction_sum(value)
+
+
+@cute.jit
+def _sigmoid(value):
+    one = cutlass.Float32(1.0)
+    return one / (one + cute.exp(-value))
+
+
+@cute.jit
+def compute_neo_inv_rms(x, eps: cutlass.Constexpr[float]):
+    square_sum = _warp_sum(x * x)
+    return cute.rsqrt(square_sum / cutlass.Float32(CHANNELS) + cutlass.Float32(eps))
+
+
+@cute.jit
+def compute_neo_output_gate(
+    x,
+    norm_scale: cute.Tensor,
+    gate_weight: cute.Tensor,
+    inv_rms,
+    focus,
+    channel,
+):
+    logit_part = (
+        x
+        * inv_rms
+        * norm_scale[focus, channel].to(cutlass.Float32)
+        * gate_weight[channel, focus, 0].to(cutlass.Float32)
+    )
+    return _sigmoid(_warp_sum(logit_part))
+
+
+@cute.jit
+def neo_output_gate_backward_jit(
+    grad_gated: cute.Tensor,
+    gated_out: cute.Tensor,
+    x_wide: cute.Tensor,
+    norm_scale: cute.Tensor,
+    gate_weight: cute.Tensor,
+    grad_phase: cute.Tensor,
+    grad_x_wide: cute.Tensor,
+    stream: CUstream,
+    eps: cutlass.Constexpr[float],
+):
+    nodes, _ = gated_out.shape
+    neo_output_gate_backward_kernel(
+        grad_gated,
+        gated_out,
+        x_wide,
+        norm_scale,
+        gate_weight,
+        grad_phase,
+        grad_x_wide,
+        eps,
+    ).launch(
+        grid=[nodes, 1, 1],
+        block=[HIDDEN, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_output_gate_backward_kernel(
+    grad_gated: cute.Tensor,
+    gated_out: cute.Tensor,
+    x_wide: cute.Tensor,
+    norm_scale: cute.Tensor,
+    gate_weight: cute.Tensor,
+    grad_phase: cute.Tensor,
+    grad_x_wide: cute.Tensor,
+    eps: cutlass.Constexpr[float],
+):
+    tid, _, _ = cute.arch.thread_idx()
+    node, _, _ = cute.arch.block_idx()
+    focus = tid // CHANNELS
+    channel = tid - focus * CHANNELS
+
+    x = x_wide[node, tid].to(cutlass.Float32)
+    inv_rms = compute_neo_inv_rms(x, eps)
+    gate_value = compute_neo_output_gate(
+        x,
+        norm_scale,
+        gate_weight,
+        inv_rms,
+        focus,
+        channel,
+    )
+
+    gate_dot = cutlass.Float32(0.0)
+    for degree in cutlass.range_constexpr(DEGREE_COUNT):
+        idx = degree * HIDDEN + tid
+        grad = grad_gated[node, idx].to(cutlass.Float32)
+        gated = gated_out[node, idx].to(cutlass.Float32)
+        grad_phase[node, idx] = (grad * gate_value).to(grad_phase.element_type)
+        gate_dot += grad * gated
+
+    gate_dot = _warp_sum(gate_dot)
+    grad_logit = gate_dot * (cutlass.Float32(1.0) - gate_value)
+
+    scale = norm_scale[focus, channel].to(cutlass.Float32)
+    weight = gate_weight[channel, focus, 0].to(cutlass.Float32)
+    grad_scaled = grad_logit * weight * scale
+    rms_coeff = _warp_sum(grad_scaled * x) / cutlass.Float32(CHANNELS)
+
+    grad_x = grad_scaled * inv_rms
+    grad_x -= x * inv_rms * inv_rms * inv_rms * rms_coeff
+    previous = grad_x_wide[node, tid].to(cutlass.Float32)
+    grad_x_wide[node, tid] = (previous + grad_x).to(grad_x_wide.element_type)
+
+
+def compile_neo_output_gate_backward(eps: float) -> Callable:
+    """Compile input-only backward.
+
+    Runtime order is ``grad_gated, gated_out, x_wide, norm_scale, gate_weight,
+    grad_phase, grad_x_wide``.  ``grad_phase`` may alias ``grad_gated`` and the
+    kernel adds only the scalar row into the preinitialized ``grad_x_wide``.
+    """
+    nodes = cute.sym_int64()
+
+    def fake_node_tensor():
+        return make_fake_compact_tensor(
+            cutlass.Float32,
+            (nodes, OUTPUT_WIDTH),
+            stride_order=(1, 0),
+            **FAKE_TENSOR_KW,
+        )
+
+    grad_gated = fake_node_tensor()
+    gated_out = fake_node_tensor()
+    x_wide = fake_node_tensor()
+    norm_scale = make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    gate_weight = make_fake_compact_tensor(
+        cutlass.Float32,
+        (CHANNELS, FOCUS_COUNT, 1),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    grad_phase = fake_node_tensor()
+    grad_x_wide = fake_node_tensor()
+    stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        neo_output_gate_backward_jit,
+        grad_gated,
+        gated_out,
+        x_wide,
+        norm_scale,
+        gate_weight,
+        grad_phase,
+        grad_x_wide,
+        stream,
+        eps,
+        options="--enable-tvm-ffi",
+    )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_a_radial_forward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_a_radial_forward.py
new file mode 100644
index 0000000000..09fdf3b137
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_a_radial_forward.py
@@ -0,0 +1,394 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Packed-direct Phase-A/radial forward without an ``x_rot`` boundary."""
+
+# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    Callable,
+)
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+from ..wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES
+
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+
+@dataclass(frozen=True)
+class NeoPhaseARadialForwardParams:
+    x_wide: cute.Tensor
+    src: cute.Tensor
+    d_full: cute.Tensor
+    radial_m0: cute.Tensor
+    combined_weight: cute.Tensor
+    hidden_weight: cute.Tensor
+    channel_basis: cute.Tensor
+    out: cute.Tensor
+    rad_l0: cute.Tensor
+
+
+@cute.jit
+def _store_packed_phase_a_value(
+    params: NeoPhaseARadialForwardParams,
+    x_local: cute.Tensor,
+    edge,
+    src_node,
+    channel,
+    reduced: cutlass.Constexpr[int],
+    panel_start: cutlass.Constexpr[int],
+    full_start: cutlass.Constexpr[int],
+    width: cutlass.Constexpr[int],
+):
+    acc = cutlass.Float32(0.0)
+    for local_col in cutlass.range_constexpr(width):
+        d_val = params.d_full[edge, panel_start + local_col].to(cutlass.Float32)
+        x_val = params.x_wide[
+            src_node,
+            (full_start + local_col) * 64 + channel,
+        ].to(cutlass.Float32)
+        acc += d_val * x_val
+    x_local[reduced * 64 + channel] = acc
+
+
+@cute.jit
+def neo_phase_a_radial_forward_packed_direct_saved_jit(
+    x_wide: cute.Tensor,
+    src: cute.Tensor,
+    d_full: cute.Tensor,
+    radial_m0: cute.Tensor,
+    combined_weight: cute.Tensor,
+    hidden_weight: cute.Tensor,
+    channel_basis: cute.Tensor,
+    out: cute.Tensor,
+    rad_l0: cute.Tensor,
+    compact_out: cute.Tensor,
+    stream: CUstream,
+):
+    params = NeoPhaseARadialForwardParams(
+        x_wide=x_wide,
+        src=src,
+        d_full=d_full,
+        radial_m0=radial_m0,
+        combined_weight=combined_weight,
+        hidden_weight=hidden_weight,
+        channel_basis=channel_basis,
+        out=out,
+        rad_l0=rad_l0,
+    )
+    edges, _ = out.shape
+    neo_phase_a_radial_forward_packed_direct_saved_kernel(
+        params,
+        compact_out,
+    ).launch(
+        grid=[edges, 1, 1],
+        block=[64, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_phase_a_radial_forward_packed_direct_saved_kernel(
+    params: NeoPhaseARadialForwardParams,
+    compact_out: cute.Tensor,
+):
+    channel, _, _ = cute.arch.thread_idx()
+    edge, _, _ = cute.arch.block_idx()
+
+    smem = cutlass.utils.SmemAllocator()
+    x_local = smem.allocate_tensor(cutlass.Float32, 10 * 64)
+    compact = smem.allocate_tensor(cutlass.Float32, 25)
+    src_node = params.src[edge]
+
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 0, 0, 0, 1)
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 1, 1, 1, 3)
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 2, 10, 4, 5)
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 3, 25, 9, 7)
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 4, 4, 1, 3)
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 5, 15, 4, 5)
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 6, 32, 9, 7)
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 7, 7, 1, 3)
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 8, 20, 4, 5)
+    _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 9, 39, 9, 7)
+
+    if channel < 25:
+        acc = cutlass.Float32(0.0)
+        for radial_idx in cutlass.range_constexpr(4 * 32):
+            radial_value = params.radial_m0[edge, radial_idx].to(cutlass.Float32)
+            weight = params.combined_weight[radial_idx, channel].to(cutlass.Float32)
+            acc += radial_value * weight
+        compact[channel] = acc
+        compact_out[edge, channel] = acc
+
+    acc_l0 = cutlass.Float32(0.0)
+    for radial_channel in cutlass.range_constexpr(32):
+        radial_value = params.radial_m0[edge, radial_channel].to(cutlass.Float32)
+        weight = params.hidden_weight[radial_channel, channel].to(cutlass.Float32)
+        acc_l0 += radial_value * weight
+    params.rad_l0[edge, channel] = acc_l0
+
+    cute.arch.sync_threads()
+
+    for coeff in cutlass.range_constexpr(10):
+        acc = cutlass.Float32(0.0)
+        if coeff < 4:
+            out_coeff = coeff
+            for in_coeff in cutlass.range_constexpr(4):
+                kval = compact[in_coeff * 4 + out_coeff]
+                acc += kval * x_local[in_coeff * 64 + channel]
+        elif coeff < 7:
+            out_coeff = coeff - 4
+            for in_coeff in cutlass.range_constexpr(3):
+                kval = compact[16 + in_coeff * 3 + out_coeff]
+                acc += kval * x_local[(4 + in_coeff) * 64 + channel]
+        else:
+            out_coeff = coeff - 7
+            for in_coeff in cutlass.range_constexpr(3):
+                kval = compact[16 + in_coeff * 3 + out_coeff]
+                acc += kval * x_local[(7 + in_coeff) * 64 + channel]
+        acc *= params.channel_basis[channel].to(cutlass.Float32)
+        focus = channel // 32
+        focus_channel = channel - focus * 32
+        out_idx = focus * 10 * 32 + coeff * 32 + focus_channel
+        params.out[edge, out_idx] = acc
+
+
+def compile_neo_phase_a_radial_forward_packed_direct() -> Callable:
+    edge_count = cute.sym_int64()
+    node_count = cute.sym_int64()
+    fake_x_wide = make_fake_compact_tensor(
+        cutlass.Float32,
+        (node_count, 16 * 64),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_src = make_fake_compact_tensor(
+        cutlass.Int32,
+        (edge_count,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_d = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, PACKED_WIGNER_VALUES),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_radial = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, 4 * 32),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_combined = make_fake_compact_tensor(
+        cutlass.Float32,
+        (4 * 32, 25),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_hidden = make_fake_compact_tensor(
+        cutlass.Float32,
+        (32, 64),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_basis = make_fake_compact_tensor(
+        cutlass.Float32,
+        (64,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_out = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, 10 * 64),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_rad_l0 = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, 64),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    fake_compact = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, 25),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    return cute.compile(
+        neo_phase_a_radial_forward_packed_direct_saved_jit,
+        fake_x_wide,
+        fake_src,
+        fake_d,
+        fake_radial,
+        fake_combined,
+        fake_hidden,
+        fake_basis,
+        fake_out,
+        fake_rad_l0,
+        fake_compact,
+        fake_stream,
+        options="--enable-tvm-ffi",
+    )
+
+
+@device_aware_lru_cache(maxsize=4)
+def _compiled_neo_phase_a_radial_forward_packed_direct() -> Callable:
+    return compile_neo_phase_a_radial_forward_packed_direct()
+
+
+def _combined_radial_weight(radial_hidden_proj, radial_degree_mixer):
+    import torch
+
+    hidden_weight = radial_hidden_proj.weight
+    mixer_weight = radial_degree_mixer.weight
+    cache = getattr(radial_hidden_proj, "_deepmd_cute_neo_radial_combined", None)
+    key = (
+        hidden_weight.data_ptr(),
+        hidden_weight._version,
+        hidden_weight.dtype,
+        hidden_weight.device,
+        tuple(hidden_weight.shape),
+        tuple(hidden_weight.stride()),
+        hidden_weight.storage_offset(),
+        mixer_weight.data_ptr(),
+        mixer_weight._version,
+        mixer_weight.dtype,
+        mixer_weight.device,
+        tuple(mixer_weight.shape),
+        tuple(mixer_weight.stride()),
+        mixer_weight.storage_offset(),
+    )
+    if cache is not None and cache[0] == key:
+        return cache[1]
+
+    blocks = []
+    for degree in range(4):
+        mixer_block = mixer_weight.detach()[degree * 64 : (degree + 1) * 64, :]
+        blocks.append(torch.mm(hidden_weight.detach(), mixer_block))
+    combined = torch.cat(blocks, dim=0).contiguous()
+    radial_hidden_proj._deepmd_cute_neo_radial_combined = (key, combined)
+    return combined
+
+
+def run_neo_phase_a_radial_forward_packed_direct(
+    *,
+    radial_hidden_proj,
+    radial_degree_mixer,
+    x_wide,
+    src,
+    D_full,
+    radial_feat_m0,
+):
+    """Validate and launch the packed Phase-A/radial forward kernel."""
+    import torch
+
+    if x_wide.shape[1:] != (16, 64):
+        raise ValueError(f"expected x_wide shape (N,16,64), got {x_wide.shape}")
+    edge_count = src.numel()
+    if tuple(D_full.shape) != (edge_count, PACKED_WIGNER_VALUES):
+        raise ValueError(
+            "expected packed Wigner shape "
+            f"{(edge_count, PACKED_WIGNER_VALUES)}, got {tuple(D_full.shape)}"
+        )
+    if radial_feat_m0.shape != (edge_count, 4, 32):
+        raise ValueError(
+            f"expected radial_feat_m0 shape {(edge_count, 4, 32)}, "
+            f"got {tuple(radial_feat_m0.shape)}"
+        )
+    device = x_wide.device
+    if device.type != "cuda":
+        raise ValueError("packed Phase-A/radial forward requires CUDA tensors")
+    if src.device != device or src.dtype not in (torch.int32, torch.int64):
+        raise ValueError("src must be an int32 or int64 tensor on the input device")
+    if src.data_ptr() % 16:
+        raise ValueError("src must be 16-byte aligned")
+    source_tensors = (
+        ("x_wide", x_wide),
+        ("D_full", D_full),
+        ("radial_feat_m0", radial_feat_m0),
+        ("radial_hidden_proj.weight", radial_hidden_proj.weight),
+        ("radial_degree_mixer.weight", radial_degree_mixer.weight),
+        ("radial_degree_mixer.channel_basis", radial_degree_mixer.channel_basis),
+    )
+    for name, tensor in source_tensors:
+        if tensor.device != device or tensor.dtype != torch.float32:
+            raise ValueError(f"{name} must be FP32 on {device}")
+        if tensor.data_ptr() % 16:
+            raise ValueError(f"{name} must be 16-byte aligned")
+    if radial_hidden_proj.bias is not None:
+        raise NotImplementedError("collapsed radial mixer expects no hidden bias")
+    if tuple(radial_hidden_proj.weight.shape) != (32, 64):
+        raise NotImplementedError("collapsed radial mixer expects a (32,64) projection")
+    if radial_degree_mixer.mode != "degree_channel" or radial_degree_mixer.rank != 1:
+        raise NotImplementedError(
+            "collapsed radial mixer expects degree_channel rank=1"
+        )
+    if tuple(radial_degree_mixer.weight.shape) != (4 * 64, 25):
+        raise NotImplementedError("collapsed radial mixer expects lmax=3,mmax=1,C=64")
+    if tuple(radial_degree_mixer.channel_basis.shape) != (1, 64):
+        raise NotImplementedError(
+            "collapsed radial mixer expects a rank-1, 64-channel basis"
+        )
+
+    combined_weight = _combined_radial_weight(
+        radial_hidden_proj,
+        radial_degree_mixer,
+    )
+    if not combined_weight.is_contiguous() or combined_weight.data_ptr() % 16:
+        raise ValueError("combined radial weight must be contiguous and aligned")
+    kernel = _compiled_neo_phase_a_radial_forward_packed_direct()
+    out = torch.empty(
+        edge_count,
+        10 * 64,
+        device=x_wide.device,
+        dtype=x_wide.dtype,
+    )
+    rad_l0 = torch.empty(
+        edge_count,
+        64,
+        device=radial_feat_m0.device,
+        dtype=radial_feat_m0.dtype,
+    )
+    compact_out = torch.empty(
+        (edge_count, 25),
+        device=radial_feat_m0.device,
+        dtype=torch.float32,
+    )
+    kernel(
+        x_wide.contiguous().view(x_wide.shape[0], 16 * 64),
+        src.to(torch.int32).contiguous(),
+        D_full.contiguous(),
+        radial_feat_m0.contiguous().view(edge_count, 4 * 32),
+        combined_weight,
+        radial_hidden_proj.weight.detach().contiguous(),
+        radial_degree_mixer.channel_basis.detach().view(64).contiguous(),
+        out,
+        rad_l0,
+        compact_out,
+    )
+    return (
+        out.view(edge_count, 2, 10, 32),
+        rad_l0.view(edge_count, 2, 32),
+        compact_out,
+    )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py
new file mode 100644
index 0000000000..b14b4c3402
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py
@@ -0,0 +1,702 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Exact-shape Neo Phase-C backward with fused layout-boundary reductions.
+
+This kernel specializes ``D=16``, ``Dm=10``, ``F=2``, ``C=32`` and the
+46-value packed Wigner panel. It keeps the destination adjoint in shared
+memory once per node, but streams each degree once per edge into a lane-local
+ten-value input-adjoint fragment. ``grad_Dt`` is reduced by the two focus
+warps, so the effective stack input does not need a shared-memory slab.
+
+The kernel also owns the two reductions immediately downstream of Phase C:
+
+* attention ``grad_alpha`` is consumed in-place to produce envelope-softmax
+  ``grad_logits``, ``grad_edge`` and ``grad_z``;
+* focus ``grad_alpha`` is consumed in-place by the two-focus softmax/RMSNorm
+  backward to produce ``grad_focus_src``.
+
+The stack adjoint is written edge-major into the fully consumed final saved
+activation, avoiding a separate edge-sized allocation. The focus-source
+gradient remains focus-major.
+"""
+
+# ruff: noqa: ANN001, ANN201, ANN202, TC002
+
+from __future__ import (
+    annotations,
+)
+
+import operator
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ..wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+DEGREE_COUNT = 16
+REDUCED_COUNT = 10
+N_FOCUS = 2
+FOCUS_CHANNELS = 32
+HIDDEN = N_FOCUS * FOCUS_CHANNELS
+
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+
+@dataclass(frozen=True)
+class NeoPhaseCBackwardLayoutParams:
+    grad_out: cute.Tensor
+    stack: cute.Tensor
+    wigner_dt: cute.Tensor
+    alpha: cute.Tensor
+    focus_alpha: cute.Tensor
+    dst_ptr: cute.Tensor
+    rotate_inv_rescale: cute.Tensor
+    edge_gate: cute.Tensor
+    z_bias_raw: cute.Tensor
+    group_max: cute.Tensor
+    denom: cute.Tensor
+    focus_src: cute.Tensor
+    focus_weight: cute.Tensor
+    focus_scale: cute.Tensor
+    grad_stack: cute.Tensor
+    grad_wigner_dt: cute.Tensor
+    grad_logits: cute.Tensor
+    grad_edge: cute.Tensor
+    grad_z_partial: cute.Tensor
+    grad_z: cute.Tensor
+    grad_focus_src: cute.Tensor
+
+
+@cute.jit
+def _warp_sum(value):
+    return cute.arch.warp_reduction(value, operator.add)
+
+
+@cute.jit
+def _sigmoid(value):
+    one = cutlass.Float32(1.0)
+    return one / (one + cute.exp(-value))
+
+
+@cute.jit
+def _record_panel_term(
+    panel_values,
+    panel_offset,
+    dt_partial,
+    raw_fragment,
+    grad_fragment,
+    transformed,
+    head,
+    lane,
+    reduced: cutlass.Constexpr[int],
+    panel_index: cutlass.Constexpr[int],
+    alpha_value,
+    focus_value,
+):
+    """Consume one structural Wigner entry for grad-stack and grad-Dt."""
+    panel = panel_values[panel_offset + panel_index]
+    grad_fragment[reduced] = (
+        grad_fragment[reduced].to(cutlass.Float32) + panel * transformed
+    )
+
+    effective_stack = raw_fragment[reduced].to(cutlass.Float32) * focus_value
+    partial = _warp_sum(transformed * effective_stack * alpha_value)
+    if lane == 0:
+        dt_partial[head * PACKED_WIGNER_VALUES + panel_index] = partial
+
+
+@cute.jit
+def _focus_source_backward(
+    params: NeoPhaseCBackwardLayoutParams,
+    focus_grad,
+    focus_inv,
+    focus_grad_logits,
+    focus_coeff,
+    edge,
+    head,
+    lane,
+    tidx,
+    eps: cutlass.Constexpr[float],
+    tau: cutlass.Constexpr[float],
+    label_smoothing: cutlass.Constexpr[float],
+    use_focus_norm: cutlass.Constexpr[bool],
+):
+    """Consume the Phase-C focus-alpha reduction without a global temporary."""
+    weight = params.focus_weight[lane, head].to(cutlass.Float32)
+
+    if tidx == 0:
+        probability_keep = cutlass.Float32(1.0 - label_smoothing)
+        smooth = cutlass.Float32(label_smoothing / N_FOCUS)
+        probability0 = (
+            params.focus_alpha[edge, 0].to(cutlass.Float32) - smooth
+        ) / probability_keep
+        probability1 = (
+            params.focus_alpha[edge, 1].to(cutlass.Float32) - smooth
+        ) / probability_keep
+        keep = cutlass.Float32(1.0 - label_smoothing)
+        grad0 = focus_grad[0] * keep
+        grad1 = focus_grad[1] * keep
+        dot = grad0 * probability0 + grad1 * probability1
+        inv_tau = cutlass.Float32(1.0 / tau)
+        focus_grad_logits[0] = probability0 * (grad0 - dot) * inv_tau
+        focus_grad_logits[1] = probability1 * (grad1 - dot) * inv_tau
+    cute.arch.sync_threads()
+
+    if cutlass.const_expr(use_focus_norm):
+        value = params.focus_src[edge, head, lane].to(cutlass.Float32)
+        scale = params.focus_scale[head, lane].to(cutlass.Float32)
+        square_sum = _warp_sum(value * value)
+        if lane == 0:
+            focus_inv[head] = cute.rsqrt(
+                square_sum / cutlass.Float32(FOCUS_CHANNELS) + cutlass.Float32(eps)
+            )
+        grad_scaled = focus_grad_logits[head] * weight * scale
+        coeff_sum = _warp_sum(grad_scaled * value)
+        if lane == 0:
+            focus_coeff[head] = coeff_sum / cutlass.Float32(FOCUS_CHANNELS)
+        cute.arch.sync_warp()
+
+        inv = focus_inv[head]
+        grad_value = grad_scaled * inv
+        grad_value -= value * inv * inv * inv * focus_coeff[head]
+    else:
+        grad_value = focus_grad_logits[head] * weight
+    params.grad_focus_src[head, edge, lane] = grad_value.to(
+        params.grad_focus_src.element_type
+    )
+
+
+@cute.kernel
+def neo_phase_c_backward_layout_kernel(
+    params: NeoPhaseCBackwardLayoutParams,
+    raw_layout: cute.Layout,
+    raw_tiled_copy: cute.TiledCopy,
+    focus_eps: cutlass.Constexpr[float],
+    focus_tau: cutlass.Constexpr[float],
+    focus_label_smoothing: cutlass.Constexpr[float],
+    use_focus_norm: cutlass.Constexpr[bool],
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    node, _, _ = cute.arch.block_idx()
+    head = tidx // FOCUS_CHANNELS
+    lane = tidx - head * FOCUS_CHANNELS
+    lo = params.dst_ptr[node]
+    hi = params.dst_ptr[node + 1]
+
+    smem = cutlass.utils.SmemAllocator()
+    t_values = smem.allocate_tensor(cutlass.Float32, DEGREE_COUNT * HIDDEN)
+    panel_values = smem.allocate_tensor(
+        cutlass.Float32,
+        PACKED_WIGNER_VALUES,
+    )
+    dt_partial = smem.allocate_tensor(cutlass.Float32, N_FOCUS * PACKED_WIGNER_VALUES)
+    focus_grad = smem.allocate_tensor(cutlass.Float32, N_FOCUS)
+    focus_inv = smem.allocate_tensor(cutlass.Float32, N_FOCUS)
+    focus_grad_logits = smem.allocate_tensor(cutlass.Float32, N_FOCUS)
+    focus_coeff = smem.allocate_tensor(cutlass.Float32, N_FOCUS)
+    softmax_dot_by_focus = smem.allocate_tensor(cutlass.Float32, N_FOCUS)
+    gate_tile = smem.allocate_tensor(cutlass.Float32, HIDDEN)
+
+    for degree in cutlass.range_constexpr(DEGREE_COUNT):
+        index = degree * HIDDEN + tidx
+        upstream = params.grad_out[node, degree, tidx].to(cutlass.Float32)
+        rotate = params.rotate_inv_rescale[degree].to(cutlass.Float32)
+        t_values[index] = upstream * rotate
+    cute.arch.sync_threads()
+
+    raw_thread_copy = raw_tiled_copy.get_slice(lane)
+    softmax_dot = cutlass.Float32(0.0)
+
+    for edge in cutlass.range(lo, hi, 1, unroll=1):
+        stack_tile = cute.local_tile(
+            params.stack,
+            tiler=(1, 1, REDUCED_COUNT, FOCUS_CHANNELS),
+            coord=(edge, head, 0, 0),
+        )
+        stack_head = cute.make_tensor(stack_tile.iterator, raw_layout)
+        thread_stack = raw_thread_copy.partition_S(stack_head)
+        raw_fragment = cute.make_fragment_like(thread_stack, cutlass.Float32)
+        cute.copy(raw_tiled_copy, thread_stack, raw_fragment)
+        if tidx < PACKED_WIGNER_VALUES:
+            panel_values[tidx] = params.wigner_dt[edge, tidx].to(cutlass.Float32)
+        cute.arch.sync_threads()
+
+        alpha_value = params.alpha[edge, head].to(cutlass.Float32)
+        focus_value = params.focus_alpha[edge, head].to(cutlass.Float32)
+        grad_fragment = cute.make_fragment_like(raw_fragment, cutlass.Float32)
+        grad_fragment.fill(0.0)
+
+        transformed = t_values[tidx]
+        _record_panel_term(
+            panel_values,
+            0,
+            dt_partial,
+            raw_fragment,
+            grad_fragment,
+            transformed,
+            head,
+            lane,
+            0,
+            0,
+            alpha_value,
+            focus_value,
+        )
+        for local_col in cutlass.range_constexpr(3):
+            transformed = t_values[(1 + local_col) * HIDDEN + tidx]
+            for row_slot in cutlass.range_constexpr(3):
+                _record_panel_term(
+                    panel_values,
+                    0,
+                    dt_partial,
+                    raw_fragment,
+                    grad_fragment,
+                    transformed,
+                    head,
+                    lane,
+                    1 + row_slot * 3,
+                    1 + row_slot * 3 + local_col,
+                    alpha_value,
+                    focus_value,
+                )
+        for local_col in cutlass.range_constexpr(5):
+            transformed = t_values[(4 + local_col) * HIDDEN + tidx]
+            for row_slot in cutlass.range_constexpr(3):
+                _record_panel_term(
+                    panel_values,
+                    0,
+                    dt_partial,
+                    raw_fragment,
+                    grad_fragment,
+                    transformed,
+                    head,
+                    lane,
+                    2 + row_slot * 3,
+                    10 + row_slot * 5 + local_col,
+                    alpha_value,
+                    focus_value,
+                )
+        for local_col in cutlass.range_constexpr(7):
+            transformed = t_values[(9 + local_col) * HIDDEN + tidx]
+            for row_slot in cutlass.range_constexpr(3):
+                _record_panel_term(
+                    panel_values,
+                    0,
+                    dt_partial,
+                    raw_fragment,
+                    grad_fragment,
+                    transformed,
+                    head,
+                    lane,
+                    3 + row_slot * 3,
+                    25 + row_slot * 7 + local_col,
+                    alpha_value,
+                    focus_value,
+                )
+
+        grad_focus_part = cutlass.Float32(0.0)
+        grad_alpha_part = cutlass.Float32(0.0)
+        for reduced in cutlass.range_constexpr(REDUCED_COUNT):
+            raw = raw_fragment[reduced].to(cutlass.Float32)
+            grad_raw = grad_fragment[reduced].to(cutlass.Float32)
+            grad_focus_part += grad_raw * raw * alpha_value
+            grad_alpha_part += grad_raw * raw * focus_value
+            grad_fragment[reduced] = grad_raw * focus_value * alpha_value
+
+        # Every edge belongs to exactly one node CTA. Both focus warps have
+        # loaded their complete source fragment, so this exact-address store
+        # may reuse the fully consumed stack allocation.
+        grad_stack_tile = cute.local_tile(
+            params.grad_stack,
+            tiler=(1, 1, REDUCED_COUNT, FOCUS_CHANNELS),
+            coord=(edge, head, 0, 0),
+        )
+        grad_stack_head = cute.make_tensor(grad_stack_tile.iterator, raw_layout)
+        thread_grad_stack = raw_thread_copy.partition_D(grad_stack_head)
+        cute.copy(raw_tiled_copy, grad_fragment, thread_grad_stack)
+
+        grad_focus_value = _warp_sum(grad_focus_part)
+        grad_alpha_value = _warp_sum(grad_alpha_part)
+        if lane == 0:
+            focus_grad[head] = grad_focus_value
+            # grad_logits is the dead grad-alpha slab during the first node pass.
+            params.grad_logits[edge, head] = grad_alpha_value.to(
+                params.grad_logits.element_type
+            )
+            softmax_dot += grad_alpha_value * alpha_value
+
+        cute.arch.sync_threads()
+        if tidx < PACKED_WIGNER_VALUES:
+            value = dt_partial[tidx] + dt_partial[PACKED_WIGNER_VALUES + tidx]
+            params.grad_wigner_dt[edge, tidx] = value.to(
+                params.grad_wigner_dt.element_type
+            )
+
+        _focus_source_backward(
+            params,
+            focus_grad,
+            focus_inv,
+            focus_grad_logits,
+            focus_coeff,
+            edge,
+            head,
+            lane,
+            tidx,
+            focus_eps,
+            focus_tau,
+            focus_label_smoothing,
+            use_focus_norm,
+        )
+        cute.arch.sync_threads()
+
+    if lane == 0:
+        softmax_dot_by_focus[head] = softmax_dot
+        max_value = params.group_max[node, head].to(cutlass.Float32)
+        denom = params.denom[node, head].to(cutlass.Float32)
+        z_sigmoid = _sigmoid(params.z_bias_raw[head].to(cutlass.Float32))
+        params.grad_z_partial[node, head] = (
+            -softmax_dot * cute.exp(-max_value) / denom * z_sigmoid
+        ).to(params.grad_z_partial.element_type)
+    cute.arch.sync_threads()
+    # The edge loop below is lane-striped, so every lane needs the node dot.
+    softmax_dot = softmax_dot_by_focus[head]
+
+    for edge_base in cutlass.range(lo, hi, FOCUS_CHANNELS, unroll=1):
+        edge = edge_base + lane
+        gate_contribution = cutlass.Float32(0.0)
+        if edge < hi:
+            upstream = params.grad_logits[edge, head].to(cutlass.Float32)
+            centered = upstream - softmax_dot
+            alpha_value = params.alpha[edge, head].to(cutlass.Float32)
+            params.grad_logits[edge, head] = (alpha_value * centered).to(
+                params.grad_logits.element_type
+            )
+
+            gate = params.edge_gate[edge].to(cutlass.Float32)
+            if gate < cutlass.Float32(0.0):
+                gate = cutlass.Float32(0.0)
+            if gate > cutlass.Float32(0.0):
+                gate_contribution = alpha_value * centered * cutlass.Float32(2.0) / gate
+        gate_tile[tidx] = gate_contribution
+        cute.arch.sync_threads()
+        if head == 0 and edge < hi:
+            params.grad_edge[edge] = (gate_tile[lane] + gate_tile[lane + 32]).to(
+                params.grad_edge.element_type
+            )
+        cute.arch.sync_threads()
+
+
+@cute.jit
+def _cta_sum(value, scratch, threads: cutlass.Constexpr[int]):
+    lane = cute.arch.lane_idx()
+    warp = cute.arch.warp_idx()
+    warps = threads // 32
+    value = _warp_sum(value)
+    if lane == 0:
+        scratch[warp] = value
+    cute.arch.barrier()
+
+    total = cutlass.Float32(0.0)
+    if lane < warps:
+        total = scratch[lane]
+    return _warp_sum(total)
+
+
+@cute.kernel
+def neo_phase_c_backward_z_reduce_kernel(
+    params: NeoPhaseCBackwardLayoutParams,
+    threads: cutlass.Constexpr[int],
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    head, _, _ = cute.arch.block_idx()
+    node_count, _ = params.grad_z_partial.shape
+    warps = threads // 32
+
+    smem = cutlass.utils.SmemAllocator()
+    scratch = smem.allocate_tensor(cutlass.Float32, warps)
+    local = cutlass.Float32(0.0)
+    for node in cutlass.range(tidx, node_count, threads, unroll=1):
+        local += params.grad_z_partial[node, head].to(cutlass.Float32)
+    total = _cta_sum(local, scratch, threads)
+    if tidx == 0:
+        params.grad_z[head] = total.to(params.grad_z.element_type)
+
+
+@cute.jit
+def neo_phase_c_backward_layout_jit(
+    grad_out: cute.Tensor,
+    stack: cute.Tensor,
+    wigner_dt: cute.Tensor,
+    alpha: cute.Tensor,
+    focus_alpha: cute.Tensor,
+    dst_ptr: cute.Tensor,
+    rotate_inv_rescale: cute.Tensor,
+    edge_gate: cute.Tensor,
+    z_bias_raw: cute.Tensor,
+    group_max: cute.Tensor,
+    denom: cute.Tensor,
+    focus_src: cute.Tensor,
+    focus_weight: cute.Tensor,
+    focus_scale: cute.Tensor,
+    grad_stack: cute.Tensor,
+    grad_wigner_dt: cute.Tensor,
+    grad_logits: cute.Tensor,
+    grad_edge: cute.Tensor,
+    grad_z_partial: cute.Tensor,
+    grad_z: cute.Tensor,
+    grad_focus_src: cute.Tensor,
+    stream: CUstream,
+    focus_eps: cutlass.Constexpr[float],
+    focus_tau: cutlass.Constexpr[float],
+    focus_label_smoothing: cutlass.Constexpr[float],
+    use_focus_norm: cutlass.Constexpr[bool],
+):
+    params = NeoPhaseCBackwardLayoutParams(
+        grad_out=grad_out,
+        stack=stack,
+        wigner_dt=wigner_dt,
+        alpha=alpha,
+        focus_alpha=focus_alpha,
+        dst_ptr=dst_ptr,
+        rotate_inv_rescale=rotate_inv_rescale,
+        edge_gate=edge_gate,
+        z_bias_raw=z_bias_raw,
+        group_max=group_max,
+        denom=denom,
+        focus_src=focus_src,
+        focus_weight=focus_weight,
+        focus_scale=focus_scale,
+        grad_stack=grad_stack,
+        grad_wigner_dt=grad_wigner_dt,
+        grad_logits=grad_logits,
+        grad_edge=grad_edge,
+        grad_z_partial=grad_z_partial,
+        grad_z=grad_z,
+        grad_focus_src=grad_focus_src,
+    )
+    copy_atom = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        stack.element_type,
+        num_bits_per_copy=32,
+    )
+    channel_thread_layout = cute.make_ordered_layout((1, 32), order=(1, 0))
+    reduced_value_layout = cute.make_ordered_layout((10, 1), order=(1, 0))
+    raw_tiled_copy = cute.make_tiled_copy_tv(
+        copy_atom,
+        channel_thread_layout,
+        reduced_value_layout,
+    )
+    raw_layout = cute.make_layout((REDUCED_COUNT, FOCUS_CHANNELS), stride=(32, 1))
+    node_count, _, _ = grad_out.shape
+    neo_phase_c_backward_layout_kernel(
+        params,
+        raw_layout,
+        raw_tiled_copy,
+        focus_eps,
+        focus_tau,
+        focus_label_smoothing,
+        use_focus_norm,
+    ).launch(
+        grid=[node_count, 1, 1],
+        block=[HIDDEN, 1, 1],
+        stream=stream,
+    )
+    neo_phase_c_backward_z_reduce_kernel(params, 128).launch(
+        grid=[N_FOCUS, 1, 1],
+        block=[128, 1, 1],
+        stream=stream,
+    )
+
+
+def compile_neo_phase_c_backward_layout(
+    *,
+    focus_eps: float,
+    focus_tau: float,
+    focus_label_smoothing: float,
+    use_focus_norm: bool = True,
+) -> Callable:
+    """Compile the exact Neo Phase-C layout-boundary backward callable."""
+    if use_focus_norm and focus_eps <= 0.0:
+        raise ValueError("focus_eps must be positive")
+    if focus_tau <= 0.0:
+        raise ValueError("focus_tau must be positive")
+    if not 0.0 <= focus_label_smoothing < 1.0:
+        raise ValueError("focus_label_smoothing must be in [0, 1)")
+
+    edge_count = cute.sym_int64()
+    node_count = cute.sym_int64()
+    fake_grad_out = make_fake_compact_tensor(
+        cutlass.Float32,
+        (node_count, DEGREE_COUNT, HIDDEN),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_stack = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, N_FOCUS, REDUCED_COUNT, FOCUS_CHANNELS),
+        stride_order=(3, 2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_wigner_dt = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, PACKED_WIGNER_VALUES),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_alpha = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, N_FOCUS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_focus_alpha = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, N_FOCUS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_dst_ptr = make_fake_compact_tensor(
+        cutlass.Int32,
+        (cute.sym_int64(),),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_rotate = make_fake_compact_tensor(
+        cutlass.Float32,
+        (DEGREE_COUNT,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_edge_gate = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_z_bias = make_fake_compact_tensor(
+        cutlass.Float32,
+        (N_FOCUS,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_group_max = make_fake_compact_tensor(
+        cutlass.Float32,
+        (node_count, N_FOCUS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_denom = make_fake_compact_tensor(
+        cutlass.Float32,
+        (node_count, N_FOCUS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_focus_src = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, N_FOCUS, FOCUS_CHANNELS),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_focus_weight = make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_CHANNELS, N_FOCUS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_focus_scale = make_fake_compact_tensor(
+        cutlass.Float32,
+        (N_FOCUS, FOCUS_CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_stack = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, N_FOCUS, REDUCED_COUNT, FOCUS_CHANNELS),
+        stride_order=(3, 2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_wigner_dt = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, PACKED_WIGNER_VALUES),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_logits = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, N_FOCUS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_edge = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_z_partial = make_fake_compact_tensor(
+        cutlass.Float32,
+        (node_count, N_FOCUS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_z = make_fake_compact_tensor(
+        cutlass.Float32,
+        (N_FOCUS,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_focus_src = make_fake_compact_tensor(
+        cutlass.Float32,
+        (N_FOCUS, edge_count, FOCUS_CHANNELS),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        neo_phase_c_backward_layout_jit,
+        fake_grad_out,
+        fake_stack,
+        fake_wigner_dt,
+        fake_alpha,
+        fake_focus_alpha,
+        fake_dst_ptr,
+        fake_rotate,
+        fake_edge_gate,
+        fake_z_bias,
+        fake_group_max,
+        fake_denom,
+        fake_focus_src,
+        fake_focus_weight,
+        fake_focus_scale,
+        fake_grad_stack,
+        fake_grad_wigner_dt,
+        fake_grad_logits,
+        fake_grad_edge,
+        fake_grad_z_partial,
+        fake_grad_z,
+        fake_grad_focus_src,
+        fake_stream,
+        focus_eps,
+        focus_tau,
+        focus_label_smoothing,
+        bool(use_focus_norm),
+        options="--enable-tvm-ffi",
+    )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_forward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_forward.py
new file mode 100644
index 0000000000..634d7cde8c
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_forward.py
@@ -0,0 +1,565 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Strict-FP32 one-pass Neo Phase-C forward with output gating.
+
+This module specializes the supported Neo shape ``D=16, Dm=10, F=2,
+C=32``.  One 64-thread CTA owns one destination CSR row.  Lane
+``focus * 32 + channel`` sweeps every edge in that row and retains its 16
+output-degree accumulators in a CuTe register fragment.  Wigner values are
+cooperatively staged through a double-buffered shared-memory panel, so no
+``node * chunks`` partial tensor or second reduction launch is required.
+
+The final stores fuse the output-side attention gate:
+
+``sigmoid(project(RMSNorm(x_wide[:, 0]))) * rotate_inv_rescale * aggregate``.
+
+Dense Wigner, packed Wigner through the generic loader, and the 46-value
+packed-direct path share the same runtime tensor signature.  Every floating
+tensor and every arithmetic operation is FP32 by construction.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+from ..wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN201, ANN202, TC002
+
+DEGREE_COUNT = 16
+REDUCED_COUNT = 10
+FOCUS_COUNT = 2
+CHANNELS = 32
+
+if PACKED_WIGNER_VALUES != 46 or PACKED_WIGNER_VALUES > 2 * CHANNELS:
+    raise RuntimeError("one-pass Neo Phase-C requires the 46-value Wigner layout")
+HIDDEN = FOCUS_COUNT * CHANNELS
+PHASE_WIDTH = REDUCED_COUNT * HIDDEN
+OUTPUT_WIDTH = DEGREE_COUNT * HIDDEN
+THREADS = HIDDEN
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+
+@dataclass(frozen=True)
+class NeoPhaseCOnePassParams:
+    """Runtime tensors for the warp-private packed kernel."""
+
+    x_local: cute.Tensor
+    wigner_dt: cute.Tensor
+    alpha: cute.Tensor
+    focus_alpha: cute.Tensor
+    out: cute.Tensor
+    x_wide: cute.Tensor
+    norm_scale: cute.Tensor
+    gate_weight: cute.Tensor
+    dst_ptr: cute.Tensor
+    rotate_inv_rescale: cute.Tensor
+
+
+@cute.jit
+def _sigmoid(value):
+    one = cutlass.Float32(1.0)
+    return one / (one + cute.exp(-value))
+
+
+@cute.jit
+def _weighted_input(
+    params: NeoPhaseCOnePassParams,
+    edge,
+    tid,
+    focus,
+    channel,
+    reduced: cutlass.Constexpr[int],
+    alpha,
+    focus_scale,
+):
+    """Load one reduced coefficient and apply both edge attention weights."""
+    load_idx = focus * REDUCED_COUNT * CHANNELS + reduced * CHANNELS + channel
+    value = params.x_local[edge, load_idx].to(cutlass.Float32)
+    # Match the two-launch path's FP32 association exactly.
+    return value * alpha * focus_scale
+
+
+@cute.jit
+def _focus_scale(
+    params: NeoPhaseCOnePassParams,
+    edge,
+    focus,
+):
+    return params.focus_alpha[edge, focus].to(cutlass.Float32)
+
+
+@cute.jit
+def _output_gate(
+    params: NeoPhaseCOnePassParams,
+    node,
+    tid,
+    focus,
+    channel,
+    eps: cutlass.Constexpr[float],
+):
+    """Compute one gate per focus; each warp owns one 32-channel focus."""
+    x = params.x_wide[node, tid].to(cutlass.Float32)
+    square_sum = cute.arch.warp_reduction_sum(x * x)
+    inv_rms = cute.rsqrt(square_sum / cutlass.Float32(CHANNELS) + cutlass.Float32(eps))
+    logit_part = (
+        x
+        * inv_rms
+        * params.norm_scale[focus, channel].to(cutlass.Float32)
+        * params.gate_weight[channel, focus, 0].to(cutlass.Float32)
+    )
+    logit = cute.arch.warp_reduction_sum(logit_part)
+    return _sigmoid(logit)
+
+
+@cute.jit
+def _store_gated_output(
+    params: NeoPhaseCOnePassParams,
+    accumulator: cute.Tensor,
+    node,
+    tid,
+    gate,
+):
+    for degree in cutlass.range_constexpr(DEGREE_COUNT):
+        value = accumulator[degree]
+        value *= params.rotate_inv_rescale[degree].to(cutlass.Float32)
+        value *= gate
+        params.out[node, degree * HIDDEN + tid] = value
+
+
+@cute.jit
+def _warp_private_packed_wigner(
+    panel,
+    focus,
+    panel_index: cutlass.Constexpr[int],
+):
+    """Load one scalar from a focus warp's private shared Wigner panel."""
+    return panel[focus, panel_index]
+
+
+@cute.jit
+def neo_phase_c_onepass_output_gate_packed_direct_warp_private_jit(
+    x_local: cute.Tensor,
+    wigner_dt: cute.Tensor,
+    alpha: cute.Tensor,
+    focus_alpha: cute.Tensor,
+    out: cute.Tensor,
+    x_wide: cute.Tensor,
+    norm_scale: cute.Tensor,
+    gate_weight: cute.Tensor,
+    dst_ptr: cute.Tensor,
+    rotate_inv_rescale: cute.Tensor,
+    eps: cutlass.Constexpr[float],
+    stream: CUstream,
+):
+    """Launch the packed-direct kernel with a warp-local Wigner epilogue."""
+    params = NeoPhaseCOnePassParams(
+        x_local=x_local,
+        wigner_dt=wigner_dt,
+        alpha=alpha,
+        focus_alpha=focus_alpha,
+        out=out,
+        x_wide=x_wide,
+        norm_scale=norm_scale,
+        gate_weight=gate_weight,
+        dst_ptr=dst_ptr,
+        rotate_inv_rescale=rotate_inv_rescale,
+    )
+    accumulator_layout = cute.make_layout((DEGREE_COUNT,), stride=(1,))
+    node_count, _ = out.shape
+    neo_phase_c_onepass_output_gate_packed_direct_warp_private_kernel(
+        params,
+        accumulator_layout,
+        eps,
+    ).launch(
+        grid=[node_count, 1, 1],
+        block=[THREADS, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_phase_c_onepass_output_gate_packed_direct_warp_private_kernel(
+    params: NeoPhaseCOnePassParams,
+    accumulator_layout: cute.Layout,
+    eps: cutlass.Constexpr[float],
+):
+    """Store the gated output without a Wigner shared-memory round trip."""
+    tid, _, _ = cute.arch.thread_idx()
+    node, _, _ = cute.arch.block_idx()
+    focus = tid // CHANNELS
+    channel = tid - focus * CHANNELS
+    lane = cute.arch.lane_idx()
+
+    # Each focus is exactly one warp.  The gate reduction already broadcasts
+    # its result within that warp, so it does not need shared memory either.
+    smem = cutlass.utils.SmemAllocator()
+    panel_storage = smem.allocate_tensor(
+        cutlass.Float32,
+        FOCUS_COUNT * PACKED_WIGNER_VALUES,
+    )
+    panel_layout = cute.make_layout(
+        (FOCUS_COUNT, PACKED_WIGNER_VALUES),
+        stride=(PACKED_WIGNER_VALUES, 1),
+    )
+    panel = cute.make_tensor(panel_storage.iterator, panel_layout)
+    gate = _output_gate(params, node, tid, focus, channel, eps)
+    accumulator = cute.make_rmem_tensor(accumulator_layout, cutlass.Float32)
+    accumulator.fill(0.0)
+
+    lo = params.dst_ptr[node]
+    hi = params.dst_ptr[node + 1]
+    for edge in cutlass.range(lo, hi, 1, unroll=1):
+        panel[focus, lane] = params.wigner_dt[edge, lane].to(cutlass.Float32)
+        if lane < PACKED_WIGNER_VALUES - CHANNELS:
+            panel[focus, lane + CHANNELS] = params.wigner_dt[edge, lane + CHANNELS].to(
+                cutlass.Float32
+            )
+        cute.arch.sync_warp()
+
+        alpha = params.alpha[edge, focus].to(cutlass.Float32)
+        focus_scale = _focus_scale(params, edge, focus)
+        value0 = _weighted_input(
+            params,
+            edge,
+            tid,
+            focus,
+            channel,
+            0,
+            alpha,
+            focus_scale,
+        )
+        accumulator[0] += _warp_private_packed_wigner(panel, focus, 0) * value0
+
+        for row_slot in cutlass.range_constexpr(3):
+            value1 = _weighted_input(
+                params,
+                edge,
+                tid,
+                focus,
+                channel,
+                1 + row_slot * 3,
+                alpha,
+                focus_scale,
+            )
+            panel_start1 = 1 + row_slot * 3
+            for local_row in cutlass.range_constexpr(3):
+                accumulator[1 + local_row] += (
+                    _warp_private_packed_wigner(
+                        panel,
+                        focus,
+                        panel_start1 + local_row,
+                    )
+                    * value1
+                )
+
+        for row_slot in cutlass.range_constexpr(3):
+            value2 = _weighted_input(
+                params,
+                edge,
+                tid,
+                focus,
+                channel,
+                2 + row_slot * 3,
+                alpha,
+                focus_scale,
+            )
+            panel_start2 = 10 + row_slot * 5
+            for local_row in cutlass.range_constexpr(5):
+                accumulator[4 + local_row] += (
+                    _warp_private_packed_wigner(
+                        panel,
+                        focus,
+                        panel_start2 + local_row,
+                    )
+                    * value2
+                )
+
+        for row_slot in cutlass.range_constexpr(3):
+            value3 = _weighted_input(
+                params,
+                edge,
+                tid,
+                focus,
+                channel,
+                3 + row_slot * 3,
+                alpha,
+                focus_scale,
+            )
+            panel_start3 = 25 + row_slot * 7
+            for local_row in cutlass.range_constexpr(7):
+                accumulator[9 + local_row] += (
+                    _warp_private_packed_wigner(
+                        panel,
+                        focus,
+                        panel_start3 + local_row,
+                    )
+                    * value3
+                )
+
+        # Do not overwrite this warp's panel until every lane has consumed it.
+        cute.arch.sync_warp()
+
+    _store_gated_output(params, accumulator, node, tid, gate)
+
+
+def _fake_common_tensors():
+    edges = cute.sym_int64()
+    nodes = cute.sym_int64()
+    x_local = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edges, PHASE_WIDTH),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    wigner_dt = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edges, PACKED_WIGNER_VALUES),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    alpha = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edges, FOCUS_COUNT),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    focus_alpha = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edges, FOCUS_COUNT),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    out = make_fake_compact_tensor(
+        cutlass.Float32,
+        (nodes, OUTPUT_WIDTH),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    x_wide = make_fake_compact_tensor(
+        cutlass.Float32,
+        (nodes, OUTPUT_WIDTH),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    norm_scale = make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    gate_weight = make_fake_compact_tensor(
+        cutlass.Float32,
+        (CHANNELS, FOCUS_COUNT, 1),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    dst_ptr = make_fake_compact_tensor(
+        cutlass.Int32,
+        (cute.sym_int64(),),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    rotate_inv_rescale = make_fake_compact_tensor(
+        cutlass.Float32,
+        (DEGREE_COUNT,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    return (
+        x_local,
+        wigner_dt,
+        alpha,
+        focus_alpha,
+        out,
+        x_wide,
+        norm_scale,
+        gate_weight,
+        dst_ptr,
+        rotate_inv_rescale,
+    )
+
+
+def compile_neo_phase_c_onepass_output_gate(
+    eps: float,
+) -> Callable:
+    """Compile the packed, focus-major, warp-private specialization."""
+    if not (eps > 0.0 and eps < float("inf")):
+        raise ValueError("output-gate RMSNorm eps must be finite and positive")
+    common_args = _fake_common_tensors()
+    stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        neo_phase_c_onepass_output_gate_packed_direct_warp_private_jit,
+        *common_args,
+        eps,
+        stream,
+        options="--enable-tvm-ffi",
+    )
+
+
+@device_aware_lru_cache(maxsize=8)
+def _compiled_neo_phase_c_onepass_output_gate(eps: float) -> Callable:
+    return compile_neo_phase_c_onepass_output_gate(eps)
+
+
+def _expect_shape(name: str, tensor: Any, expected: tuple[int, ...]) -> None:
+    actual = tuple(tensor.shape)
+    if actual != expected:
+        raise ValueError(f"expected {name} shape {expected}, got {actual}")
+
+
+def _expect_fp32_cuda(name: str, tensor: Any, *, torch: Any, device: Any) -> None:
+    if tensor.dtype != torch.float32:
+        raise TypeError(f"{name} must be strict float32, got {tensor.dtype}")
+    if not tensor.is_cuda:
+        raise ValueError(f"{name} must be a CUDA tensor")
+    if tensor.device != device:
+        raise ValueError(f"{name} must be on {device}, got {tensor.device}")
+
+
+def run_neo_phase_c_onepass_output_gate(
+    *,
+    x_local_flat: Any,
+    Dt_full: Any,
+    alpha_focus: Any,
+    focus_compete_alpha: Any,
+    dst_ptr: Any,
+    rotate_inv_rescale: Any,
+    x_wide: Any,
+    output_gate_norm_scale: Any,
+    output_gate_weight: Any,
+    output_gate_eps: float,
+    out: Any | None = None,
+) -> Any:
+    """Validate and launch fused Phase C and output gating."""
+    import torch
+
+    if not x_local_flat.is_cuda:
+        raise ValueError("x_local_flat must be a CUDA tensor")
+
+    device = x_local_flat.device
+    edge_count = x_local_flat.shape[0]
+    node_count = x_wide.shape[0]
+    _expect_shape(
+        "x_local_flat",
+        x_local_flat,
+        (edge_count, FOCUS_COUNT, REDUCED_COUNT, CHANNELS),
+    )
+    _expect_shape("x_wide", x_wide, (node_count, DEGREE_COUNT, HIDDEN))
+    _expect_shape("alpha_focus", alpha_focus, (edge_count, FOCUS_COUNT))
+    _expect_shape(
+        "output_gate_norm_scale",
+        output_gate_norm_scale,
+        (FOCUS_COUNT, CHANNELS),
+    )
+    _expect_shape(
+        "output_gate_weight",
+        output_gate_weight,
+        (CHANNELS, FOCUS_COUNT, 1),
+    )
+    _expect_shape("rotate_inv_rescale", rotate_inv_rescale, (DEGREE_COUNT,))
+    _expect_shape("dst_ptr", dst_ptr, (node_count + 1,))
+    _expect_shape(
+        "focus_compete_alpha",
+        focus_compete_alpha,
+        (edge_count, FOCUS_COUNT),
+    )
+    _expect_shape("Dt_full", Dt_full, (edge_count, PACKED_WIGNER_VALUES))
+
+    floating_tensors = {
+        "x_local_flat": x_local_flat,
+        "Dt_full": Dt_full,
+        "alpha_focus": alpha_focus,
+        "focus_compete_alpha": focus_compete_alpha,
+        "rotate_inv_rescale": rotate_inv_rescale,
+        "x_wide": x_wide,
+        "output_gate_norm_scale": output_gate_norm_scale,
+        "output_gate_weight": output_gate_weight,
+    }
+    for name, tensor in floating_tensors.items():
+        _expect_fp32_cuda(name, tensor, torch=torch, device=device)
+
+    if dst_ptr.device != device:
+        raise ValueError("dst_ptr must be on the input CUDA device")
+    if dst_ptr.dtype not in (torch.int32, torch.int64):
+        raise TypeError(f"dst_ptr must be int32 or int64, got {dst_ptr.dtype}")
+    torch._assert_async(
+        dst_ptr[0] == 0,
+        "Phase-C forward requires dst_ptr[0] == 0",
+    )
+    torch._assert_async(
+        dst_ptr[-1] == edge_count,
+        "Phase-C forward requires dst_ptr[-1] == edge_count",
+    )
+    if dst_ptr.numel() > 1:
+        torch._assert_async(
+            torch.all(dst_ptr[1:] >= dst_ptr[:-1]),
+            "Phase-C forward requires nondecreasing dst_ptr",
+        )
+
+    if out is None:
+        out = torch.empty(
+            node_count,
+            DEGREE_COUNT,
+            HIDDEN,
+            device=device,
+            dtype=torch.float32,
+        )
+    else:
+        _expect_shape("out", out, (node_count, DEGREE_COUNT, HIDDEN))
+        _expect_fp32_cuda("out", out, torch=torch, device=device)
+        if not out.is_contiguous():
+            raise ValueError("out must be contiguous")
+
+    with torch.cuda.device(device):
+        kernel = _compiled_neo_phase_c_onepass_output_gate(float(output_gate_eps))
+        kernel(
+            x_local_flat.contiguous().view(edge_count, REDUCED_COUNT * HIDDEN),
+            Dt_full.contiguous(),
+            alpha_focus.contiguous(),
+            focus_compete_alpha.contiguous(),
+            out.view(node_count, DEGREE_COUNT * HIDDEN),
+            x_wide.contiguous().view(node_count, DEGREE_COUNT * HIDDEN),
+            output_gate_norm_scale.contiguous(),
+            output_gate_weight.contiguous(),
+            dst_ptr.to(dtype=torch.int32).contiguous(),
+            rotate_inv_rescale.contiguous(),
+        )
+    return out
+
+
+__all__ = [
+    "compile_neo_phase_c_onepass_output_gate",
+    "neo_phase_c_onepass_output_gate_packed_direct_warp_private_jit",
+    "neo_phase_c_onepass_output_gate_packed_direct_warp_private_kernel",
+    "run_neo_phase_c_onepass_output_gate",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/qk_edge.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/qk_edge.py
new file mode 100644
index 0000000000..d3ccbe8ad7
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/qk_edge.py
@@ -0,0 +1,364 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035
+"""Fused Neo Q/K edge logits and first-backward input adjoints."""
+
+from __future__ import (
+    annotations,
+)
+
+from functools import (
+    lru_cache,
+)
+from typing import (
+    Callable,
+)
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+
+@cute.kernel
+def neo_qk_edge_forward_kernel(
+    q_node: cute.Tensor,
+    k_node: cute.Tensor,
+    radial_l0: cute.Tensor,
+    attention_weight: cute.Tensor,
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    logits: cute.Tensor,
+    scale: cutlass.Constexpr[float],
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    block, focus, _ = cute.arch.block_idx()
+    edge = block * 128 + tidx
+    edges, _ = logits.shape
+    if edge < edges:
+        src_node = src[edge]
+        dst_node = dst[edge]
+        qk_acc = cutlass.Float32(0.0)
+        radial_acc = cutlass.Float32(0.0)
+        for channel in cutlass.range_constexpr(32):
+            qk_acc += q_node[dst_node, focus, channel].to(cutlass.Float32) * k_node[
+                src_node, focus, channel
+            ].to(cutlass.Float32)
+            radial_acc += radial_l0[edge, focus, channel].to(
+                cutlass.Float32
+            ) * attention_weight[channel, focus, 0].to(cutlass.Float32)
+        logits[edge, focus] = (qk_acc * cutlass.Float32(scale) + radial_acc).to(
+            logits.element_type
+        )
+
+
+@cute.jit
+def neo_qk_edge_forward_jit(
+    q_node: cute.Tensor,
+    k_node: cute.Tensor,
+    radial_l0: cute.Tensor,
+    attention_weight: cute.Tensor,
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    logits: cute.Tensor,
+    stream: CUstream,
+    scale: cutlass.Constexpr[float],
+):
+    edges, _ = logits.shape
+    neo_qk_edge_forward_kernel(
+        q_node,
+        k_node,
+        radial_l0,
+        attention_weight,
+        src,
+        dst,
+        logits,
+        scale,
+    ).launch(
+        grid=[cute.ceil_div(edges, 128), 2, 1],
+        block=[128, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_qk_edge_backward_kernel(
+    grad_logits: cute.Tensor,
+    q_node: cute.Tensor,
+    k_node: cute.Tensor,
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    grad_q_node: cute.Tensor,
+    grad_k_node: cute.Tensor,
+    scale: cutlass.Constexpr[float],
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    block, focus, _ = cute.arch.block_idx()
+    edge = block * 8 + tidx // 32
+    channel = tidx % 32
+    edges, _ = grad_logits.shape
+    if edge < edges:
+        src_node = src[edge]
+        dst_node = dst[edge]
+        grad = grad_logits[edge, focus].to(cutlass.Float32) * cutlass.Float32(scale)
+        grad_q = grad * k_node[src_node, focus, channel].to(cutlass.Float32)
+        grad_k = grad * q_node[dst_node, focus, channel].to(cutlass.Float32)
+        q_offset = (dst_node * 2 + focus) * 32 + channel
+        k_offset = (src_node * 2 + focus) * 32 + channel
+        q_ptr = grad_q_node.iterator + q_offset
+        k_ptr = grad_k_node.iterator + k_offset
+        cute.arch.atomic_add(q_ptr.llvm_ptr, grad_q, sem="relaxed", scope="gpu")
+        cute.arch.atomic_add(k_ptr.llvm_ptr, grad_k, sem="relaxed", scope="gpu")
+
+
+@cute.jit
+def neo_qk_edge_backward_jit(
+    grad_logits: cute.Tensor,
+    q_node: cute.Tensor,
+    k_node: cute.Tensor,
+    src: cute.Tensor,
+    dst: cute.Tensor,
+    grad_q_node: cute.Tensor,
+    grad_k_node: cute.Tensor,
+    stream: CUstream,
+    scale: cutlass.Constexpr[float],
+):
+    edges, _ = grad_logits.shape
+    neo_qk_edge_backward_kernel(
+        grad_logits,
+        q_node,
+        k_node,
+        src,
+        dst,
+        grad_q_node,
+        grad_k_node,
+        scale,
+    ).launch(
+        grid=[cute.ceil_div(edges, 8), 2, 1],
+        block=[256, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_qk_node_input_adjoint_kernel(
+    x_l0: cute.Tensor,
+    grad_q_node: cute.Tensor,
+    grad_k_node: cute.Tensor,
+    q_weight: cute.Tensor,
+    k_weight: cute.Tensor,
+    norm_scale: cute.Tensor,
+    grad_x_wide: cute.Tensor,
+    eps: cutlass.Float32,
+):
+    tid, _, _ = cute.arch.thread_idx()
+    block, _, _ = cute.arch.block_idx()
+    node_in_block = tid // 64
+    local_tid = tid % 64
+    focus = local_tid // 32
+    channel = local_tid % 32
+    node = block * 4 + node_in_block
+    nodes, _, _ = x_l0.shape
+
+    if node < nodes:
+        for flat_index in cutlass.range(local_tid, 16 * 64, 64):
+            grad_x_wide[node, flat_index] = cutlass.Float32(0.0)
+
+        grad_norm = cutlass.Float32(0.0)
+        for output_channel in cutlass.range_constexpr(32):
+            grad_norm += grad_q_node[node, focus, output_channel].to(
+                cutlass.Float32
+            ) * q_weight[channel, focus, output_channel].to(cutlass.Float32)
+            grad_norm += grad_k_node[node, focus, output_channel].to(
+                cutlass.Float32
+            ) * k_weight[channel, focus, output_channel].to(cutlass.Float32)
+
+        x = x_l0[node, focus, channel].to(cutlass.Float32)
+        grad_scaled = grad_norm * norm_scale[focus, channel].to(cutlass.Float32)
+        inv = cute.rsqrt(
+            cute.arch.warp_reduction_sum(x * x) / cutlass.Float32(32.0) + eps
+        )
+        coeff = cute.arch.warp_reduction_sum(grad_scaled * x) / cutlass.Float32(32.0)
+        grad_x = grad_scaled * inv - x * inv * inv * inv * coeff
+        grad_x_wide[node, focus * 32 + channel] = grad_x.to(grad_x_wide.element_type)
+
+
+@cute.jit
+def neo_qk_node_input_adjoint_jit(
+    x_l0: cute.Tensor,
+    grad_q_node: cute.Tensor,
+    grad_k_node: cute.Tensor,
+    q_weight: cute.Tensor,
+    k_weight: cute.Tensor,
+    norm_scale: cute.Tensor,
+    grad_x_wide: cute.Tensor,
+    stream: CUstream,
+    eps: cutlass.Float32,
+):
+    nodes, _, _ = x_l0.shape
+    neo_qk_node_input_adjoint_kernel(
+        x_l0,
+        grad_q_node,
+        grad_k_node,
+        q_weight,
+        k_weight,
+        norm_scale,
+        grad_x_wide,
+        eps,
+    ).launch(
+        grid=[cute.ceil_div(nodes, 4), 1, 1],
+        block=[256, 1, 1],
+        stream=stream,
+    )
+
+
+def _fake_inputs():
+    edge_count = cute.sym_int64()
+    node_count = cute.sym_int64()
+    q_node = make_fake_compact_tensor(
+        cutlass.Float32, (node_count, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    k_node = make_fake_compact_tensor(
+        cutlass.Float32, (node_count, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    radial = make_fake_compact_tensor(
+        cutlass.Float32, (edge_count, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    weight = make_fake_compact_tensor(
+        cutlass.Float32, (32, 2, 1), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    src = make_fake_compact_tensor(
+        cutlass.Int32, (edge_count,), stride_order=(0,), **FAKE_TENSOR_KW
+    )
+    dst = make_fake_compact_tensor(
+        cutlass.Int32, (edge_count,), stride_order=(0,), **FAKE_TENSOR_KW
+    )
+    logits = make_fake_compact_tensor(
+        cutlass.Float32, (edge_count, 2), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    return q_node, k_node, radial, weight, src, dst, logits
+
+
+@lru_cache(maxsize=8)
+def compile_neo_qk_edge_forward(
+    scale: float,
+    compile_identity: tuple[int, int, int] | None = None,
+) -> Callable:
+    del compile_identity
+    q_node, k_node, radial, weight, src, dst, logits = _fake_inputs()
+    stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        neo_qk_edge_forward_jit,
+        q_node,
+        k_node,
+        radial,
+        weight,
+        src,
+        dst,
+        logits,
+        stream,
+        scale,
+        options="--enable-tvm-ffi",
+    )
+
+
+@lru_cache(maxsize=8)
+def compile_neo_qk_edge_backward(
+    scale: float,
+    compile_identity: tuple[int, int, int] | None = None,
+) -> Callable:
+    del compile_identity
+    q_node, k_node, _radial, _weight, src, dst, logits = _fake_inputs()
+    grad_q = make_fake_compact_tensor(
+        cutlass.Float32, q_node.shape, stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    grad_k = make_fake_compact_tensor(
+        cutlass.Float32, k_node.shape, stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        neo_qk_edge_backward_jit,
+        logits,
+        q_node,
+        k_node,
+        src,
+        dst,
+        grad_q,
+        grad_k,
+        stream,
+        scale,
+        options="--enable-tvm-ffi",
+    )
+
+
+@lru_cache(maxsize=8)
+def compile_neo_qk_node_input_adjoint(
+    eps: float,
+    compile_identity: tuple[int, int, int] | None = None,
+) -> Callable:
+    del compile_identity
+    node_count = cute.sym_int64()
+    x_l0 = make_fake_compact_tensor(
+        cutlass.Float32, (node_count, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    grad_q = make_fake_compact_tensor(
+        cutlass.Float32, x_l0.shape, stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    grad_k = make_fake_compact_tensor(
+        cutlass.Float32, x_l0.shape, stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    q_weight = make_fake_compact_tensor(
+        cutlass.Float32, (32, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    k_weight = make_fake_compact_tensor(
+        cutlass.Float32, (32, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW
+    )
+    norm_scale = make_fake_compact_tensor(
+        cutlass.Float32, (2, 32), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    grad_x_wide = make_fake_compact_tensor(
+        cutlass.Float32, (node_count, 16 * 64), stride_order=(1, 0), **FAKE_TENSOR_KW
+    )
+    stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    compiled = cute.compile(
+        neo_qk_node_input_adjoint_jit,
+        x_l0,
+        grad_q,
+        grad_k,
+        q_weight,
+        k_weight,
+        norm_scale,
+        grad_x_wide,
+        stream,
+        cutlass.Float32(eps),
+        options="--enable-tvm-ffi",
+    )
+
+    def run(
+        x_l0_tensor,
+        grad_q_tensor,
+        grad_k_tensor,
+        q_weight_tensor,
+        k_weight_tensor,
+        norm_scale_tensor,
+        grad_x_wide_tensor,
+    ):
+        return compiled(
+            x_l0_tensor,
+            grad_q_tensor,
+            grad_k_tensor,
+            q_weight_tensor,
+            k_weight_tensor,
+            norm_scale_tensor,
+            grad_x_wide_tensor,
+            cutlass.Float32(eps),
+        )
+
+    return run
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/radial_phase_a_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/radial_phase_a_backward.py
new file mode 100644
index 0000000000..6c7f5254a6
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/radial_phase_a_backward.py
@@ -0,0 +1,718 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Source-CSR node-tiled Neo radial and Phase-A backward.
+
+This is the contention-free node-tiled implementation for the exact Neo SO2 shape
+``lmax=3, D=16, Dm=10, Cwide=64, F=2``. Edge tensors retain their physical
+destination-sorted order. ``source_ptr`` delimits intervals in
+``source_order``, whose slots hold the corresponding physical edge ids.
+
+One 64-thread CTA owns one source node. It keeps the node feature row in shared
+memory and 16 adjoint values per thread in registers, processes all incident
+edges, writes the per-edge radial/Wigner adjoints directly, then writes the
+node adjoint once. There are no global atomics and no ``(E, 16, 64)``
+reduction intermediate.
+"""
+
+# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    Callable,
+)
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ..wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES
+
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+DEGREE_COUNT = 16
+REDUCED_COUNT = 10
+HIDDEN = 64
+FOCUS_COUNT = 2
+FOCUS_HIDDEN = 32
+FOCUS_ROW = REDUCED_COUNT * FOCUS_HIDDEN
+RADIAL_WIDTH = 4 * FOCUS_HIDDEN
+COMPACT_WIDTH = 25
+SHARED_ROW_PITCH = HIDDEN + 4
+WARP_REDUCTION_GROUP = 4
+GROUPS_PER_WARP = 32 // WARP_REDUCTION_GROUP
+GROUPS_PER_CTA = 2 * GROUPS_PER_WARP
+CHANNELS_PER_SUBGROUP_LANE = HIDDEN // WARP_REDUCTION_GROUP
+
+
+@dataclass(frozen=True)
+class NeoRadialPhaseABackwardNodeParams:
+    grad_out_focus: cute.Tensor
+    grad_focus_src: cute.Tensor
+    grad_logits: cute.Tensor
+    radial_state: cute.Tensor
+    channel_basis: cute.Tensor
+    x_wide: cute.Tensor
+    source_order: cute.Tensor
+    source_ptr: cute.Tensor
+    d_full: cute.Tensor
+    grad_x_wide: cute.Tensor
+    grad_d_full: cute.Tensor
+
+
+@cute.jit
+def _focus_grad_value(
+    grad_out_focus: cute.Tensor,
+    grad_focus_src: cute.Tensor,
+    edge,
+    coeff: cutlass.Constexpr[int],
+    channel,
+):
+    focus = channel // FOCUS_HIDDEN
+    focus_channel = channel - focus * FOCUS_HIDDEN
+    offset = focus * FOCUS_ROW + coeff * FOCUS_HIDDEN + focus_channel
+    value = grad_out_focus[edge, offset].to(cutlass.Float32)
+    if cutlass.const_expr(coeff == 0):
+        value += grad_focus_src[focus, edge, focus_channel].to(cutlass.Float32)
+    return value
+
+
+@cute.jit
+def _recompute_local_value(
+    local_values,
+    x_values,
+    d_values,
+    channel,
+    reduced: cutlass.Constexpr[int],
+    panel_start: cutlass.Constexpr[int],
+    full_start: cutlass.Constexpr[int],
+    width: cutlass.Constexpr[int],
+    shared_row_pitch: cutlass.Constexpr[int],
+    x_row_pitch: cutlass.Constexpr[int],
+):
+    acc = cutlass.Float32(0.0)
+    for local_col in cutlass.range_constexpr(width):
+        acc += (
+            d_values[panel_start + local_col]
+            * x_values[(full_start + local_col) * x_row_pitch + channel]
+        )
+    local_values[reduced * shared_row_pitch + channel] = acc
+
+
+@cute.jit
+def _warp_owned_grad_compact(
+    focus_grad,
+    local_values,
+    channel_basis: cute.Tensor,
+    compact_idx,
+    subgroup_lane,
+    shared_row_pitch: cutlass.Constexpr[int],
+):
+    """Reduce one compact-kernel gradient inside a four-lane subgroup."""
+    value = cutlass.Float32(0.0)
+    if compact_idx < 16:
+        in_coeff = compact_idx // 4
+        out_coeff = compact_idx - in_coeff * 4
+        for channel_step in cutlass.range_constexpr(CHANNELS_PER_SUBGROUP_LANE):
+            hidden_channel = subgroup_lane + channel_step * WARP_REDUCTION_GROUP
+            basis = channel_basis[hidden_channel].to(cutlass.Float32)
+            value += (
+                focus_grad[out_coeff * shared_row_pitch + hidden_channel]
+                * local_values[in_coeff * shared_row_pitch + hidden_channel]
+                * basis
+            )
+    else:
+        pair = compact_idx - 16
+        in_coeff = pair // 3
+        out_coeff = pair - in_coeff * 3
+        for channel_step in cutlass.range_constexpr(CHANNELS_PER_SUBGROUP_LANE):
+            hidden_channel = subgroup_lane + channel_step * WARP_REDUCTION_GROUP
+            basis = channel_basis[hidden_channel].to(cutlass.Float32)
+            value += basis * (
+                focus_grad[(4 + out_coeff) * shared_row_pitch + hidden_channel]
+                * local_values[(4 + in_coeff) * shared_row_pitch + hidden_channel]
+                + focus_grad[(7 + out_coeff) * shared_row_pitch + hidden_channel]
+                * local_values[(7 + in_coeff) * shared_row_pitch + hidden_channel]
+            )
+    return cute.arch.warp_reduction_sum(
+        value,
+        threads_in_group=WARP_REDUCTION_GROUP,
+    )
+
+
+@cute.jit
+def _warp_owned_grad_d(
+    local_values,
+    x_values,
+    panel_idx,
+    subgroup_lane,
+    shared_row_pitch: cutlass.Constexpr[int],
+    x_row_pitch: cutlass.Constexpr[int],
+):
+    """Reduce one packed Wigner adjoint inside a four-lane subgroup."""
+    reduced = cutlass.Int32(0)
+    full_col = cutlass.Int32(0)
+    if panel_idx >= 25:
+        local_idx = panel_idx - 25
+        row_slot = local_idx // 7
+        reduced = 3 + row_slot * 3
+        full_col = 9 + local_idx - row_slot * 7
+    elif panel_idx >= 10:
+        local_idx = panel_idx - 10
+        row_slot = local_idx // 5
+        reduced = 2 + row_slot * 3
+        full_col = 4 + local_idx - row_slot * 5
+    elif panel_idx >= 1:
+        local_idx = panel_idx - 1
+        row_slot = local_idx // 3
+        reduced = 1 + row_slot * 3
+        full_col = 1 + local_idx - row_slot * 3
+
+    value = cutlass.Float32(0.0)
+    for channel_step in cutlass.range_constexpr(CHANNELS_PER_SUBGROUP_LANE):
+        hidden_channel = subgroup_lane + channel_step * WARP_REDUCTION_GROUP
+        value += (
+            local_values[reduced * shared_row_pitch + hidden_channel]
+            * x_values[full_col * x_row_pitch + hidden_channel]
+        )
+    return cute.arch.warp_reduction_sum(
+        value,
+        threads_in_group=WARP_REDUCTION_GROUP,
+    )
+
+
+@cute.jit
+def _grad_x_value(
+    local_values,
+    d_values,
+    channel,
+    degree: cutlass.Constexpr[int],
+    local_col: cutlass.Constexpr[int],
+    panel_start: cutlass.Constexpr[int],
+    width: cutlass.Constexpr[int],
+    rows: cutlass.Constexpr[int],
+    shared_row_pitch: cutlass.Constexpr[int],
+):
+    acc = cutlass.Float32(0.0)
+    for row_slot in cutlass.range_constexpr(rows):
+        reduced = degree + row_slot * 3
+        panel_offset = panel_start + row_slot * width + local_col
+        acc += (
+            d_values[panel_offset] * local_values[reduced * shared_row_pitch + channel]
+        )
+    return acc
+
+
+@cute.jit
+def neo_radial_phase_a_backward_node_jit(
+    grad_out_focus: cute.Tensor,
+    grad_focus_src: cute.Tensor,
+    grad_logits: cute.Tensor,
+    radial_state: cute.Tensor,
+    channel_basis: cute.Tensor,
+    x_wide: cute.Tensor,
+    source_order: cute.Tensor,
+    source_ptr: cute.Tensor,
+    d_full: cute.Tensor,
+    grad_x_wide: cute.Tensor,
+    grad_d_full: cute.Tensor,
+    stream: CUstream,
+):
+    params = NeoRadialPhaseABackwardNodeParams(
+        grad_out_focus=grad_out_focus,
+        grad_focus_src=grad_focus_src,
+        grad_logits=grad_logits,
+        radial_state=radial_state,
+        channel_basis=channel_basis,
+        x_wide=x_wide,
+        source_order=source_order,
+        source_ptr=source_ptr,
+        d_full=d_full,
+        grad_x_wide=grad_x_wide,
+        grad_d_full=grad_d_full,
+    )
+    node_count, _ = grad_x_wide.shape
+    neo_radial_phase_a_backward_node_kernel(params).launch(
+        grid=[node_count, 1, 1],
+        block=[HIDDEN, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_radial_phase_a_backward_node_kernel(
+    params: NeoRadialPhaseABackwardNodeParams,
+):
+    channel, _, _ = cute.arch.thread_idx()
+    node, _, _ = cute.arch.block_idx()
+    shared_row_pitch = SHARED_ROW_PITCH
+
+    smem = cutlass.utils.SmemAllocator()
+    x_row_pitch = SHARED_ROW_PITCH
+    x_values = smem.allocate_tensor(
+        cutlass.Float32,
+        DEGREE_COUNT * x_row_pitch,
+    )
+    focus_grad = smem.allocate_tensor(
+        cutlass.Float32,
+        REDUCED_COUNT * SHARED_ROW_PITCH,
+    )
+    # The primal local rows are dead after grad_compact. Reuse this panel for
+    # their adjoints instead of reserving another 2.5 KiB per CTA.
+    local_values = smem.allocate_tensor(
+        cutlass.Float32,
+        REDUCED_COUNT * SHARED_ROW_PITCH,
+    )
+    d_values = smem.allocate_tensor(cutlass.Float32, PACKED_WIGNER_VALUES)
+    compact = smem.allocate_tensor(cutlass.Float32, COMPACT_WIDTH)
+    grad_compact = smem.allocate_tensor(cutlass.Float32, COMPACT_WIDTH)
+
+    for full_row in cutlass.range_constexpr(DEGREE_COUNT):
+        x_values[full_row * x_row_pitch + channel] = params.x_wide[
+            node,
+            full_row * HIDDEN + channel,
+        ].to(cutlass.Float32)
+    cute.arch.sync_threads()
+
+    grad_x_0 = cutlass.Float32(0.0)
+    grad_x_1 = cutlass.Float32(0.0)
+    grad_x_2 = cutlass.Float32(0.0)
+    grad_x_3 = cutlass.Float32(0.0)
+    grad_x_4 = cutlass.Float32(0.0)
+    grad_x_5 = cutlass.Float32(0.0)
+    grad_x_6 = cutlass.Float32(0.0)
+    grad_x_7 = cutlass.Float32(0.0)
+    grad_x_8 = cutlass.Float32(0.0)
+    grad_x_9 = cutlass.Float32(0.0)
+    grad_x_10 = cutlass.Float32(0.0)
+    grad_x_11 = cutlass.Float32(0.0)
+    grad_x_12 = cutlass.Float32(0.0)
+    grad_x_13 = cutlass.Float32(0.0)
+    grad_x_14 = cutlass.Float32(0.0)
+    grad_x_15 = cutlass.Float32(0.0)
+
+    lo = params.source_ptr[node]
+    hi = params.source_ptr[node + 1]
+    for slot in cutlass.range(lo, hi, 1, unroll=1):
+        edge = params.source_order[slot]
+        for coeff in cutlass.range_constexpr(REDUCED_COUNT):
+            row_offset = coeff * shared_row_pitch + channel
+            focus_grad[row_offset] = _focus_grad_value(
+                params.grad_out_focus,
+                params.grad_focus_src,
+                edge,
+                coeff,
+                channel,
+            )
+        if channel < COMPACT_WIDTH:
+            compact[channel] = params.radial_state[edge, channel].to(cutlass.Float32)
+
+        if channel < PACKED_WIGNER_VALUES:
+            d_values[channel] = params.d_full[edge, channel].to(cutlass.Float32)
+        cute.arch.sync_threads()
+
+        if cutlass.const_expr(True):
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                0,
+                0,
+                0,
+                1,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                1,
+                1,
+                1,
+                3,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                2,
+                10,
+                4,
+                5,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                3,
+                25,
+                9,
+                7,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                4,
+                4,
+                1,
+                3,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                5,
+                15,
+                4,
+                5,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                6,
+                32,
+                9,
+                7,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                7,
+                7,
+                1,
+                3,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                8,
+                20,
+                4,
+                5,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            _recompute_local_value(
+                local_values,
+                x_values,
+                d_values,
+                channel,
+                9,
+                39,
+                9,
+                7,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            cute.arch.sync_threads()
+
+        lane = channel % 32
+        warp = channel // 32
+        subgroup = lane // WARP_REDUCTION_GROUP
+        subgroup_lane = lane % WARP_REDUCTION_GROUP
+        group = warp * GROUPS_PER_WARP + subgroup
+        for batch in cutlass.range_constexpr(
+            (COMPACT_WIDTH + GROUPS_PER_CTA - 1) // GROUPS_PER_CTA
+        ):
+            compact_idx = batch * GROUPS_PER_CTA + group
+            safe_compact_idx = compact_idx
+            if compact_idx >= COMPACT_WIDTH:
+                safe_compact_idx = cutlass.Int32(0)
+            grad_value = _warp_owned_grad_compact(
+                focus_grad,
+                local_values,
+                params.channel_basis,
+                safe_compact_idx,
+                subgroup_lane,
+                shared_row_pitch,
+            )
+            if compact_idx < COMPACT_WIDTH:
+                if subgroup_lane == 0:
+                    grad_compact[compact_idx] = grad_value
+        cute.arch.sync_threads()
+
+        basis = params.channel_basis[channel].to(cutlass.Float32)
+        for coeff in cutlass.range_constexpr(REDUCED_COUNT):
+            grad_value = cutlass.Float32(0.0)
+            if coeff < 4:
+                for out_coeff in cutlass.range_constexpr(4):
+                    grad_value += (
+                        focus_grad[out_coeff * shared_row_pitch + channel]
+                        * compact[coeff * 4 + out_coeff]
+                    )
+            elif coeff < 7:
+                in_coeff = coeff - 4
+                for out_coeff in cutlass.range_constexpr(3):
+                    grad_value += (
+                        focus_grad[(4 + out_coeff) * shared_row_pitch + channel]
+                        * compact[16 + in_coeff * 3 + out_coeff]
+                    )
+            else:
+                in_coeff = coeff - 7
+                for out_coeff in cutlass.range_constexpr(3):
+                    grad_value += (
+                        focus_grad[(7 + out_coeff) * shared_row_pitch + channel]
+                        * compact[16 + in_coeff * 3 + out_coeff]
+                    )
+            local_values[coeff * shared_row_pitch + channel] = grad_value * basis
+
+        if channel < COMPACT_WIDTH:
+            # grad_out_focus is dead once this edge has been reduced. Pack
+            # the 27-column GEMM operand in-place to avoid a separate cat.
+            params.grad_out_focus[edge, channel] = grad_compact[channel].to(
+                params.grad_out_focus.element_type
+            )
+        if channel < FOCUS_COUNT:
+            params.grad_out_focus[edge, COMPACT_WIDTH + channel] = params.grad_logits[
+                edge, channel
+            ].to(params.grad_out_focus.element_type)
+        cute.arch.sync_threads()
+
+        for batch in cutlass.range_constexpr(
+            (PACKED_WIGNER_VALUES + GROUPS_PER_CTA - 1) // GROUPS_PER_CTA
+        ):
+            panel_idx = batch * GROUPS_PER_CTA + group
+            safe_panel_idx = panel_idx
+            if panel_idx >= PACKED_WIGNER_VALUES:
+                safe_panel_idx = cutlass.Int32(0)
+            grad_d_value = _warp_owned_grad_d(
+                local_values,
+                x_values,
+                safe_panel_idx,
+                subgroup_lane,
+                shared_row_pitch,
+                x_row_pitch,
+            )
+            if panel_idx < PACKED_WIGNER_VALUES:
+                if subgroup_lane == 0:
+                    params.grad_d_full[edge, panel_idx] = grad_d_value.to(
+                        params.grad_d_full.element_type
+                    )
+
+        grad_x_0 += _grad_x_value(
+            local_values, d_values, channel, 0, 0, 0, 1, 1, shared_row_pitch
+        )
+        grad_x_1 += _grad_x_value(
+            local_values, d_values, channel, 1, 0, 1, 3, 3, shared_row_pitch
+        )
+        grad_x_2 += _grad_x_value(
+            local_values, d_values, channel, 1, 1, 1, 3, 3, shared_row_pitch
+        )
+        grad_x_3 += _grad_x_value(
+            local_values, d_values, channel, 1, 2, 1, 3, 3, shared_row_pitch
+        )
+        grad_x_4 += _grad_x_value(
+            local_values, d_values, channel, 2, 0, 10, 5, 3, shared_row_pitch
+        )
+        grad_x_5 += _grad_x_value(
+            local_values, d_values, channel, 2, 1, 10, 5, 3, shared_row_pitch
+        )
+        grad_x_6 += _grad_x_value(
+            local_values, d_values, channel, 2, 2, 10, 5, 3, shared_row_pitch
+        )
+        grad_x_7 += _grad_x_value(
+            local_values, d_values, channel, 2, 3, 10, 5, 3, shared_row_pitch
+        )
+        grad_x_8 += _grad_x_value(
+            local_values, d_values, channel, 2, 4, 10, 5, 3, shared_row_pitch
+        )
+        grad_x_9 += _grad_x_value(
+            local_values, d_values, channel, 3, 0, 25, 7, 3, shared_row_pitch
+        )
+        grad_x_10 += _grad_x_value(
+            local_values, d_values, channel, 3, 1, 25, 7, 3, shared_row_pitch
+        )
+        grad_x_11 += _grad_x_value(
+            local_values, d_values, channel, 3, 2, 25, 7, 3, shared_row_pitch
+        )
+        grad_x_12 += _grad_x_value(
+            local_values, d_values, channel, 3, 3, 25, 7, 3, shared_row_pitch
+        )
+        grad_x_13 += _grad_x_value(
+            local_values, d_values, channel, 3, 4, 25, 7, 3, shared_row_pitch
+        )
+        grad_x_14 += _grad_x_value(
+            local_values, d_values, channel, 3, 5, 25, 7, 3, shared_row_pitch
+        )
+        grad_x_15 += _grad_x_value(
+            local_values, d_values, channel, 3, 6, 25, 7, 3, shared_row_pitch
+        )
+        cute.arch.sync_threads()
+
+    params.grad_x_wide[node, 0 * HIDDEN + channel] = grad_x_0.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 1 * HIDDEN + channel] = grad_x_1.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 2 * HIDDEN + channel] = grad_x_2.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 3 * HIDDEN + channel] = grad_x_3.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 4 * HIDDEN + channel] = grad_x_4.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 5 * HIDDEN + channel] = grad_x_5.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 6 * HIDDEN + channel] = grad_x_6.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 7 * HIDDEN + channel] = grad_x_7.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 8 * HIDDEN + channel] = grad_x_8.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 9 * HIDDEN + channel] = grad_x_9.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 10 * HIDDEN + channel] = grad_x_10.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 11 * HIDDEN + channel] = grad_x_11.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 12 * HIDDEN + channel] = grad_x_12.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 13 * HIDDEN + channel] = grad_x_13.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 14 * HIDDEN + channel] = grad_x_14.to(
+        params.grad_x_wide.element_type
+    )
+    params.grad_x_wide[node, 15 * HIDDEN + channel] = grad_x_15.to(
+        params.grad_x_wide.element_type
+    )
+
+
+def compile_neo_radial_phase_a_backward_node_tiled() -> Callable:
+    """Compile the exact-shape source-CSR node-tiled backward specialization."""
+    edge_count = cute.sym_int64()
+    node_count = cute.sym_int64()
+    source_ptr_count = cute.sym_int64()
+    fake_grad_out = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, FOCUS_COUNT * REDUCED_COUNT * FOCUS_HIDDEN),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_focus_src = make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, edge_count, FOCUS_HIDDEN),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_logits = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, FOCUS_COUNT),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_radial_state = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, COMPACT_WIDTH),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_basis = make_fake_compact_tensor(
+        cutlass.Float32,
+        (HIDDEN,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_x_wide = make_fake_compact_tensor(
+        cutlass.Float32,
+        (node_count, DEGREE_COUNT * HIDDEN),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_source_order = make_fake_compact_tensor(
+        cutlass.Int32,
+        (edge_count,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_source_ptr = make_fake_compact_tensor(
+        cutlass.Int32,
+        (source_ptr_count,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+    fake_d = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, PACKED_WIGNER_VALUES),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_x = make_fake_compact_tensor(
+        cutlass.Float32,
+        (node_count, DEGREE_COUNT * HIDDEN),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_d = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, PACKED_WIGNER_VALUES),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        neo_radial_phase_a_backward_node_jit,
+        fake_grad_out,
+        fake_grad_focus_src,
+        fake_grad_logits,
+        fake_radial_state,
+        fake_basis,
+        fake_x_wide,
+        fake_source_order,
+        fake_source_ptr,
+        fake_d,
+        fake_grad_x,
+        fake_grad_d,
+        fake_stream,
+        options="--enable-tvm-ffi",
+    )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/structural_gate_sm80.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/structural_gate_sm80.py
new file mode 100644
index 0000000000..7ad1616496
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/structural_gate_sm80.py
@@ -0,0 +1,497 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+# ruff: noqa: ANN001, ANN201, ANN202, TC002, TC003
+"""Strict-FP32 vectorized structural-gate forward/backward for Neo SO2.
+
+The split path keeps both neighboring SO2 products and the two gate projections
+in cuBLAS. Its elementwise consumer assigns eight threads to each 32-channel
+row, with each thread moving an aligned float4. Both directions preserve the
+split-gate tensor contract and leave dense projections in PyTorch/cuBLAS.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from collections.abc import (
+    Callable,
+)
+from functools import (
+    lru_cache,
+)
+
+import cutlass
+import cutlass.cute as cute
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ... import (
+    runtime_policy,
+)
+
+FOCUS_COUNT = 2
+REDUCED_COUNT = 10
+CHANNELS = 32
+GATE_COUNT = 3
+VECTOR_WIDTH = 4
+CHANNEL_GROUPS = CHANNELS // VECTOR_WIDTH
+ROWS_PER_BLOCK = 16
+THREADS = ROWS_PER_BLOCK * CHANNEL_GROUPS
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+FORWARD_TENSOR_NAMES = ("residual", "y", "logits", "out")
+BACKWARD_TENSOR_NAMES = ("grad_out", "y", "logits", "grad_y", "grad_logits")
+
+
+def _guard_vec4_dispatch(
+    kernel: Callable,
+    tensor_names: tuple[str, ...],
+) -> Callable:
+    from ..structural_gate import (
+        _dispatch_aligned_vec4_kernel,
+    )
+
+    def dispatch(*tensors: object):
+        return _dispatch_aligned_vec4_kernel(kernel, tensor_names, *tensors)
+
+    return dispatch
+
+
+@cute.jit
+def _sigmoid(value):
+    return cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-value))
+
+
+@cute.jit
+def neo_gate_split_structural_vec4_sm80_forward_jit(
+    residual: cute.Tensor,
+    y: cute.Tensor,
+    logits: cute.Tensor,
+    out: cute.Tensor,
+    stream: CUstream,
+):
+    copy_atom = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        residual.element_type,
+        num_bits_per_copy=residual.element_type.width * VECTOR_WIDTH,
+    )
+    channel_thread_layout = cute.make_ordered_layout(
+        (1, CHANNEL_GROUPS),
+        order=(1, 0),
+    )
+    channel_value_layout = cute.make_ordered_layout(
+        (1, VECTOR_WIDTH),
+        order=(1, 0),
+    )
+    vector_copy = cute.make_tiled_copy_tv(
+        copy_atom,
+        channel_thread_layout,
+        channel_value_layout,
+    )
+    channel_layout = cute.make_layout((1, CHANNELS), stride=(CHANNELS, 1))
+    rows, _ = y.shape
+    neo_gate_split_structural_vec4_sm80_forward_kernel(
+        residual,
+        y,
+        logits,
+        out,
+        channel_layout,
+        vector_copy,
+    ).launch(
+        grid=[cute.ceil_div(rows, ROWS_PER_BLOCK), 1, 1],
+        block=[THREADS, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_gate_split_structural_vec4_sm80_forward_kernel(
+    residual: cute.Tensor,
+    y: cute.Tensor,
+    logits: cute.Tensor,
+    out: cute.Tensor,
+    channel_layout: cute.Layout,
+    vector_copy: cute.TiledCopy,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    block_row, _, _ = cute.arch.block_idx()
+    row_slot = tidx // CHANNEL_GROUPS
+    channel_group = tidx - row_slot * CHANNEL_GROUPS
+    row = block_row * ROWS_PER_BLOCK + row_slot
+    rows, _ = y.shape
+
+    if row < rows:
+        edge = row // FOCUS_COUNT
+        focus = row - edge * FOCUS_COUNT
+        thread_copy = vector_copy.get_slice(channel_group)
+
+        y0_tile = cute.local_tile(
+            y,
+            tiler=(1, CHANNELS),
+            coord=(row, 0),
+        )
+        residual0_tile = cute.local_tile(
+            residual,
+            tiler=(1, CHANNELS),
+            coord=(row, 0),
+        )
+        out0_tile = cute.local_tile(
+            out,
+            tiler=(1, CHANNELS),
+            coord=(row, 0),
+        )
+        thread_y = thread_copy.partition_S(y0_tile)
+        thread_residual = thread_copy.partition_S(residual0_tile)
+        thread_out = thread_copy.partition_D(out0_tile)
+        y_fragment = cute.make_fragment_like(thread_y, cutlass.Float32)
+        residual_fragment = cute.make_fragment_like(
+            thread_residual,
+            cutlass.Float32,
+        )
+        cute.copy(vector_copy, thread_y, y_fragment)
+        cute.copy(vector_copy, thread_residual, residual_fragment)
+        for value_idx in cutlass.range_constexpr(VECTOR_WIDTH):
+            y0 = y_fragment[value_idx].to(cutlass.Float32)
+            value = y0 * _sigmoid(y0)
+            value += residual_fragment[value_idx].to(cutlass.Float32)
+            residual_fragment[value_idx] = value
+        cute.copy(vector_copy, residual_fragment, thread_out)
+
+        for gate_index in cutlass.range_constexpr(GATE_COUNT):
+            logits_tile = cute.local_tile(
+                logits,
+                tiler=(1, 1, CHANNELS),
+                coord=(focus, edge, gate_index),
+            )
+            logits_panel = cute.make_tensor(logits_tile.iterator, channel_layout)
+            thread_logits = thread_copy.partition_S(logits_panel)
+            gate_fragment = cute.make_fragment_like(
+                thread_logits,
+                cutlass.Float32,
+            )
+            cute.copy(vector_copy, thread_logits, gate_fragment)
+            for value_idx in cutlass.range_constexpr(VECTOR_WIDTH):
+                gate_fragment[value_idx] = _sigmoid(
+                    gate_fragment[value_idx].to(cutlass.Float32)
+                )
+
+            for repeat in cutlass.range_constexpr(3):
+                degree = 1 + gate_index + repeat * GATE_COUNT
+                y_tile = cute.local_tile(
+                    y,
+                    tiler=(1, CHANNELS),
+                    coord=(row, degree),
+                )
+                residual_tile = cute.local_tile(
+                    residual,
+                    tiler=(1, CHANNELS),
+                    coord=(row, degree),
+                )
+                out_tile = cute.local_tile(
+                    out,
+                    tiler=(1, CHANNELS),
+                    coord=(row, degree),
+                )
+                thread_y = thread_copy.partition_S(y_tile)
+                thread_residual = thread_copy.partition_S(residual_tile)
+                thread_out = thread_copy.partition_D(out_tile)
+                y_fragment = cute.make_fragment_like(thread_y, cutlass.Float32)
+                residual_fragment = cute.make_fragment_like(
+                    thread_residual,
+                    cutlass.Float32,
+                )
+                cute.copy(vector_copy, thread_y, y_fragment)
+                cute.copy(vector_copy, thread_residual, residual_fragment)
+                for value_idx in cutlass.range_constexpr(VECTOR_WIDTH):
+                    value = y_fragment[value_idx].to(cutlass.Float32) * gate_fragment[
+                        value_idx
+                    ].to(cutlass.Float32)
+                    value += residual_fragment[value_idx].to(cutlass.Float32)
+                    residual_fragment[value_idx] = value
+                cute.copy(vector_copy, residual_fragment, thread_out)
+
+
+@cute.jit
+def neo_gate_split_structural_vec4_sm80_backward_jit(
+    grad_out: cute.Tensor,
+    y: cute.Tensor,
+    logits: cute.Tensor,
+    grad_y: cute.Tensor,
+    grad_logits: cute.Tensor,
+    stream: CUstream,
+):
+    copy_atom = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        y.element_type,
+        num_bits_per_copy=y.element_type.width * VECTOR_WIDTH,
+    )
+    channel_thread_layout = cute.make_ordered_layout(
+        (1, CHANNEL_GROUPS),
+        order=(1, 0),
+    )
+    channel_value_layout = cute.make_ordered_layout(
+        (1, VECTOR_WIDTH),
+        order=(1, 0),
+    )
+    vector_copy = cute.make_tiled_copy_tv(
+        copy_atom,
+        channel_thread_layout,
+        channel_value_layout,
+    )
+    channel_layout = cute.make_layout((1, CHANNELS), stride=(CHANNELS, 1))
+    rows, _ = y.shape
+    neo_gate_split_structural_vec4_sm80_backward_kernel(
+        grad_out,
+        y,
+        logits,
+        grad_y,
+        grad_logits,
+        channel_layout,
+        vector_copy,
+    ).launch(
+        grid=[cute.ceil_div(rows, ROWS_PER_BLOCK), 1, 1],
+        block=[THREADS, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def neo_gate_split_structural_vec4_sm80_backward_kernel(
+    grad_out: cute.Tensor,
+    y: cute.Tensor,
+    logits: cute.Tensor,
+    grad_y: cute.Tensor,
+    grad_logits: cute.Tensor,
+    channel_layout: cute.Layout,
+    vector_copy: cute.TiledCopy,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    block_row, _, _ = cute.arch.block_idx()
+    row_slot = tidx // CHANNEL_GROUPS
+    channel_group = tidx - row_slot * CHANNEL_GROUPS
+    row = block_row * ROWS_PER_BLOCK + row_slot
+    rows, _ = y.shape
+
+    if row < rows:
+        edge = row // FOCUS_COUNT
+        focus = row - edge * FOCUS_COUNT
+        thread_copy = vector_copy.get_slice(channel_group)
+
+        y0_tile = cute.local_tile(y, tiler=(1, CHANNELS), coord=(row, 0))
+        grad_y0_tile = cute.local_tile(
+            grad_y,
+            tiler=(1, CHANNELS),
+            coord=(row, 0),
+        )
+        grad_out0_panel = cute.local_tile(
+            grad_out,
+            tiler=(1, CHANNELS),
+            coord=(row, 0),
+        )
+        thread_y0 = thread_copy.partition_S(y0_tile)
+        thread_grad_out0 = thread_copy.partition_S(grad_out0_panel)
+        thread_grad_y0 = thread_copy.partition_D(grad_y0_tile)
+        y0_fragment = cute.make_fragment_like(thread_y0, cutlass.Float32)
+        grad_out0_fragment = cute.make_fragment_like(
+            thread_grad_out0,
+            cutlass.Float32,
+        )
+        cute.copy(vector_copy, thread_y0, y0_fragment)
+        cute.copy(vector_copy, thread_grad_out0, grad_out0_fragment)
+        for value_idx in cutlass.range_constexpr(VECTOR_WIDTH):
+            y0 = y0_fragment[value_idx].to(cutlass.Float32)
+            sig0 = _sigmoid(y0)
+            grad0 = grad_out0_fragment[value_idx].to(cutlass.Float32)
+            y0_fragment[value_idx] = (
+                grad0
+                * sig0
+                * (cutlass.Float32(1.0) + y0 * (cutlass.Float32(1.0) - sig0))
+            )
+        cute.copy(vector_copy, y0_fragment, thread_grad_y0)
+
+        for gate_index in cutlass.range_constexpr(GATE_COUNT):
+            logits_tile = cute.local_tile(
+                logits,
+                tiler=(1, 1, CHANNELS),
+                coord=(focus, edge, gate_index),
+            )
+            logits_panel = cute.make_tensor(logits_tile.iterator, channel_layout)
+            thread_logits = thread_copy.partition_S(logits_panel)
+            gate_fragment = cute.make_fragment_like(
+                thread_logits,
+                cutlass.Float32,
+            )
+            grad_logit_fragment = cute.make_fragment_like(
+                thread_logits,
+                cutlass.Float32,
+            )
+            cute.copy(vector_copy, thread_logits, gate_fragment)
+            for value_idx in cutlass.range_constexpr(VECTOR_WIDTH):
+                gate_fragment[value_idx] = _sigmoid(
+                    gate_fragment[value_idx].to(cutlass.Float32)
+                )
+                grad_logit_fragment[value_idx] = cutlass.Float32(0.0)
+
+            for repeat in cutlass.range_constexpr(3):
+                degree = 1 + gate_index + repeat * GATE_COUNT
+                y_tile = cute.local_tile(
+                    y,
+                    tiler=(1, CHANNELS),
+                    coord=(row, degree),
+                )
+                grad_y_tile = cute.local_tile(
+                    grad_y,
+                    tiler=(1, CHANNELS),
+                    coord=(row, degree),
+                )
+                grad_out_panel = cute.local_tile(
+                    grad_out,
+                    tiler=(1, CHANNELS),
+                    coord=(row, degree),
+                )
+                thread_y = thread_copy.partition_S(y_tile)
+                thread_grad_out = thread_copy.partition_S(grad_out_panel)
+                thread_grad_y = thread_copy.partition_D(grad_y_tile)
+                y_fragment = cute.make_fragment_like(thread_y, cutlass.Float32)
+                grad_out_fragment = cute.make_fragment_like(
+                    thread_grad_out,
+                    cutlass.Float32,
+                )
+                cute.copy(vector_copy, thread_y, y_fragment)
+                cute.copy(vector_copy, thread_grad_out, grad_out_fragment)
+                for value_idx in cutlass.range_constexpr(VECTOR_WIDTH):
+                    gate = gate_fragment[value_idx].to(cutlass.Float32)
+                    gout = grad_out_fragment[value_idx].to(cutlass.Float32)
+                    y_value = y_fragment[value_idx].to(cutlass.Float32)
+                    y_fragment[value_idx] = gout * gate
+                    grad_logit_fragment[value_idx] += (
+                        gout * y_value * gate * (cutlass.Float32(1.0) - gate)
+                    )
+                cute.copy(vector_copy, y_fragment, thread_grad_y)
+
+            grad_logits_tile = cute.local_tile(
+                grad_logits,
+                tiler=(1, 1, CHANNELS),
+                coord=(focus, edge, gate_index),
+            )
+            grad_logits_panel = cute.make_tensor(
+                grad_logits_tile.iterator,
+                channel_layout,
+            )
+            thread_grad_logits = thread_copy.partition_D(grad_logits_panel)
+            cute.copy(vector_copy, grad_logit_fragment, thread_grad_logits)
+
+
+@lru_cache(maxsize=8)
+def compile_neo_gate_split_structural_vec4_sm80_forward(
+    compile_identity: tuple[int, int, int] | None = None,
+) -> Callable:
+    if (
+        compile_identity is not None
+        and compile_identity[1:] not in runtime_policy.SUPPORTED_SO2_CAPABILITIES
+    ):
+        raise ValueError(
+            "vectorized structural gate forward requires a supported SO2 device"
+        )
+    rows = cute.sym_int64()
+    edges = cute.sym_int64()
+    fake_residual = make_fake_compact_tensor(
+        cutlass.Float32,
+        (rows, REDUCED_COUNT * CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_y = make_fake_compact_tensor(
+        cutlass.Float32,
+        (rows, REDUCED_COUNT * CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_logits = make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, edges, GATE_COUNT * CHANNELS),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_out = make_fake_compact_tensor(
+        cutlass.Float32,
+        (rows, REDUCED_COUNT * CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return _guard_vec4_dispatch(
+        cute.compile(
+            neo_gate_split_structural_vec4_sm80_forward_jit,
+            fake_residual,
+            fake_y,
+            fake_logits,
+            fake_out,
+            fake_stream,
+            options="--enable-tvm-ffi",
+        ),
+        FORWARD_TENSOR_NAMES,
+    )
+
+
+@lru_cache(maxsize=8)
+def compile_neo_gate_split_structural_vec4_sm80_backward(
+    compile_identity: tuple[int, int, int] | None = None,
+) -> Callable:
+    if (
+        compile_identity is not None
+        and compile_identity[1:] not in runtime_policy.SUPPORTED_SO2_CAPABILITIES
+    ):
+        raise ValueError(
+            "vectorized structural gate backward requires a supported SO2 device"
+        )
+    rows = cute.sym_int64()
+    edges = cute.sym_int64()
+    fake_grad_out = make_fake_compact_tensor(
+        cutlass.Float32,
+        (rows, REDUCED_COUNT * CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_y = make_fake_compact_tensor(
+        cutlass.Float32,
+        (rows, REDUCED_COUNT * CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_logits = make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, edges, GATE_COUNT * CHANNELS),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_y = make_fake_compact_tensor(
+        cutlass.Float32,
+        (rows, REDUCED_COUNT * CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_grad_logits = make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, edges, GATE_COUNT * CHANNELS),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return _guard_vec4_dispatch(
+        cute.compile(
+            neo_gate_split_structural_vec4_sm80_backward_jit,
+            fake_grad_out,
+            fake_y,
+            fake_logits,
+            fake_grad_y,
+            fake_grad_logits,
+            fake_stream,
+            options="--enable-tvm-ffi",
+        ),
+        BACKWARD_TENSOR_NAMES,
+    )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/linear.py b/deepmd/pt_expt/kernels/cute/sezm/so2/linear.py
new file mode 100644
index 0000000000..8912823e9a
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/linear.py
@@ -0,0 +1,259 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""SO2Linear helpers for the Neo CuTe SO2 path.
+
+The in-place residual adjoint enforces FP32 operand dtypes. Strict FP32 GEMM
+also requires the caller to select highest float32 matmul precision and disable
+TF32; these helpers do not mutate process-wide backend settings.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    Any,
+)
+
+import torch
+from torch import (
+    Tensor,
+)
+
+FOCUS_COUNT = 2
+REDUCED_COUNT = 10
+CHANNELS = 32
+M0_WIDTH = 4 * 32
+PAIR_WIDTH = 6 * CHANNELS
+FULL_WIDTH = REDUCED_COUNT * CHANNELS
+
+
+def _validate_neo_so2_linear(so2_linear: Any) -> None:
+    if (
+        so2_linear.lmax != 3
+        or so2_linear.mmax != 1
+        or so2_linear.in_channels != 32
+        or so2_linear.out_channels != 32
+        or so2_linear.n_focus != 2
+        or so2_linear.mlp_bias
+    ):
+        raise NotImplementedError("Neo CuTe SO2Linear expects lmax=3,mmax=1,F=2,C=32")
+
+
+def run_neo_so2_linear_manual(
+    so2_linear: Any,
+    x_local: Any,
+    *,
+    add_residual: bool = False,
+    per_focus_pair: bool = False,
+) -> Any:
+    """Run the fixed Neo SO2Linear block with two dense focus batched GEMMs."""
+    import torch
+
+    _validate_neo_so2_linear(so2_linear)
+    if add_residual:
+        w0, wpair = cached_neo_so2_linear_residual_weights(so2_linear)
+    else:
+        w0, wpair = cached_neo_so2_linear_weights(so2_linear)
+    x_flat = x_local.reshape(x_local.shape[0], 2, 10 * 32).transpose(0, 1)
+    out = x_local.new_empty(x_local.shape[0], 2, 10 * 32)
+    out_t = out.transpose(0, 1)
+    torch.bmm(x_flat[:, :, : 4 * 32], w0, out=out_t[:, :, : 4 * 32])
+    if per_focus_pair:
+        for focus in range(FOCUS_COUNT):
+            torch.mm(
+                x_flat[focus, :, M0_WIDTH:],
+                wpair[focus],
+                out=out_t[focus, :, M0_WIDTH:],
+            )
+    else:
+        torch.bmm(x_flat[:, :, M0_WIDTH:], wpair, out=out_t[:, :, M0_WIDTH:])
+    return out.reshape(x_local.shape[0], 2, 10, 32)
+
+
+def cached_neo_so2_linear_residual_weights(so2_linear: Any) -> tuple[Any, Any]:
+    """Return dense weights with the fixed SO2 residual folded in."""
+    w0, wpair = cached_neo_so2_linear_weights(so2_linear)
+    cache = getattr(so2_linear, "_deepmd_cute_neo_manual_residual_weights", None)
+    cache_key = (w0.data_ptr(), wpair.data_ptr(), w0.dtype, w0.device)
+    if (
+        cache is not None
+        and cache[0] is w0
+        and cache[1] is wpair
+        and cache[2] == cache_key
+    ):
+        return cache[3], cache[4]
+
+    w0_residual = w0.clone()
+    wpair_residual = wpair.clone()
+    w0_residual.diagonal(dim1=-2, dim2=-1).add_(1.0)
+    wpair_residual.diagonal(dim1=-2, dim2=-1).add_(1.0)
+    so2_linear._deepmd_cute_neo_manual_residual_weights = (
+        w0,
+        wpair,
+        cache_key,
+        w0_residual,
+        wpair_residual,
+    )
+    return w0_residual, wpair_residual
+
+
+def cached_neo_so2_linear_weights(so2_linear: Any) -> tuple[Any, Any]:
+    """Return cached dense block weights for Neo's fixed SO2Linear layout."""
+    import torch
+
+    cache = getattr(so2_linear, "_deepmd_cute_neo_manual_weights", None)
+    cache_key = (
+        so2_linear.weight_m0.data_ptr(),
+        so2_linear.weight_m[0].data_ptr(),
+        so2_linear.weight_m0._version,
+        so2_linear.weight_m[0]._version,
+        so2_linear.weight_m0.dtype,
+        so2_linear.weight_m0.device,
+    )
+    if (
+        cache is not None
+        and cache[0] is so2_linear.weight_m0
+        and cache[1] is so2_linear.weight_m[0]
+        and cache[2] == cache_key
+    ):
+        return cache[3], cache[4]
+
+    w0 = so2_linear.weight_m0.detach().view(4 * 32, 2, 4 * 32)
+    w0 = w0.permute(1, 0, 2).contiguous()
+    raw_pair = so2_linear.weight_m[0].detach().view(3 * 32, 2, 2 * 3 * 32)
+    w_u = raw_pair[:, :, : 3 * 32]
+    w_v = raw_pair[:, :, 3 * 32 :]
+    wpair = torch.empty(
+        2,
+        2 * 3 * 32,
+        2 * 3 * 32,
+        device=raw_pair.device,
+        dtype=raw_pair.dtype,
+    )
+    wpair[:, : 3 * 32, : 3 * 32] = w_u.permute(1, 0, 2)
+    wpair[:, : 3 * 32, 3 * 32 :] = w_v.permute(1, 0, 2)
+    wpair[:, 3 * 32 :, : 3 * 32] = -w_v.permute(1, 0, 2)
+    wpair[:, 3 * 32 :, 3 * 32 :] = w_u.permute(1, 0, 2)
+    wpair = wpair.contiguous()
+    so2_linear._deepmd_cute_neo_manual_weights = (
+        so2_linear.weight_m0,
+        so2_linear.weight_m[0],
+        cache_key,
+        w0,
+        wpair,
+    )
+    return w0, wpair
+
+
+def _has_direct_cublas_layout(tensor: Tensor) -> bool:
+    """Check only the direct-cuBLAS stride layouts used by PyTorch 2.10.
+
+    This is a stride-layout predicate, not a complete cuBLAS eligibility test;
+    dtype/device requirements are validated separately.
+    """
+    if tensor.ndim != 2:
+        return False
+    rows, columns = tensor.shape
+    stride0, stride1 = tensor.stride()
+    return (stride0 == 1 and stride1 >= max(1, rows)) or (
+        stride1 == 1 and stride0 >= max(1, columns)
+    )
+
+
+def _contiguous_tensors_overlap(lhs: Tensor, rhs: Tensor) -> bool:
+    """Return exact byte overlap for validated, nonempty contiguous tensors."""
+    if lhs.device != rhs.device:
+        return False
+    if lhs.device.type == "meta":
+        return torch._C._overlaps(lhs, rhs)
+
+    lhs_start = lhs.data_ptr()
+    rhs_start = rhs.data_ptr()
+    lhs_stop = lhs_start + lhs.numel() * lhs.element_size()
+    rhs_stop = rhs_start + rhs.numel() * rhs.element_size()
+    return lhs_start < rhs_stop and rhs_start < lhs_stop
+
+
+def _validate_edge_focus(tensor: Tensor, name: str) -> None:
+    if tensor.dtype != torch.float32:
+        raise TypeError(f"{name} must be float32, got {tensor.dtype}")
+    if tensor.ndim != 4 or tuple(tensor.shape[1:]) != (
+        FOCUS_COUNT,
+        REDUCED_COUNT,
+        CHANNELS,
+    ):
+        raise ValueError(
+            f"{name} must have shape (E,2,10,32), got {tuple(tensor.shape)}"
+        )
+    if tensor.shape[0] <= 0:
+        raise ValueError(f"{name} requires E > 0")
+    if not tensor.is_contiguous():
+        raise ValueError(f"{name} must use the canonical contiguous SO2 layout")
+
+
+def _validate_weight(
+    weight: Tensor,
+    name: str,
+    width: int,
+    device: torch.device,
+) -> None:
+    if weight.dtype != torch.float32:
+        raise TypeError(f"{name} must be float32, got {weight.dtype}")
+    expected_shape = (FOCUS_COUNT, width, width)
+    if tuple(weight.shape) != expected_shape:
+        raise ValueError(
+            f"{name} must have shape {expected_shape}, got {tuple(weight.shape)}"
+        )
+    if weight.device != device:
+        raise ValueError(f"{name} must be on {device}, got {weight.device}")
+    if not weight.is_contiguous():
+        raise ValueError(f"{name} must be contiguous")
+
+
+def neo_so2_linear_backward_residual_inplace(
+    residual: Tensor,
+    grad_out: Tensor,
+    w0_t: Tensor,
+    wpair_t: Tensor,
+) -> Tensor:
+    """Overwrite a dead residual with ``grad_out @ W.T + residual``."""
+    _validate_edge_focus(residual, "residual")
+    _validate_edge_focus(grad_out, "grad_out")
+    if residual.shape != grad_out.shape:
+        raise ValueError("residual and grad_out shapes must match")
+    if residual.device != grad_out.device:
+        raise ValueError("residual and grad_out devices must match")
+    if _contiguous_tensors_overlap(residual, grad_out):
+        raise ValueError(
+            "residual and grad_out must not alias; keep the final layer out-of-place"
+        )
+
+    _validate_weight(w0_t, "w0_t", M0_WIDTH, residual.device)
+    _validate_weight(wpair_t, "wpair_t", PAIR_WIDTH, residual.device)
+    for name, weight in (("w0_t", w0_t), ("wpair_t", wpair_t)):
+        if _contiguous_tensors_overlap(residual, weight):
+            raise ValueError(f"residual and {name} must not alias")
+
+    edge_count = residual.shape[0]
+    residual_flat = residual.view(edge_count, FOCUS_COUNT, FULL_WIDTH)
+    grad_flat = grad_out.view(edge_count, FOCUS_COUNT, FULL_WIDTH)
+    for focus in range(FOCUS_COUNT):
+        for start, stop, weight in (
+            (0, M0_WIDTH, w0_t[focus]),
+            (M0_WIDTH, FULL_WIDTH, wpair_t[focus]),
+        ):
+            residual_block = residual_flat[:, focus, start:stop]
+            grad_block = grad_flat[:, focus, start:stop]
+            if not _has_direct_cublas_layout(residual_block) or not (
+                _has_direct_cublas_layout(grad_block)
+                and _has_direct_cublas_layout(weight)
+            ):
+                raise ValueError("SO2 block layout would require cuBLAS staging")
+            residual_block.addmm_(
+                grad_block,
+                weight,
+                beta=1.0,
+                alpha=1.0,
+            )
+    return residual
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/message_grid.py b/deepmd/pt_expt/kernels/cute/sezm/so2/message_grid.py
new file mode 100644
index 0000000000..0a272a235c
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/message_grid.py
@@ -0,0 +1,365 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Packed-layout Neo message-grid forward and first input adjoint."""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    Any,
+    Protocol,
+)
+
+import torch
+from torch import (
+    Tensor,
+)
+
+COEFF_DIM = 16
+N_FOCUS = 2
+N_FRAMES = 3
+
+
+class _Sm90MessageGridState(Protocol):
+    """State subset consumed by the packed SM90 forward path."""
+
+    schedule: Tensor
+
+
+CHANNELS = 32
+HIDDEN_CHANNELS = N_FOCUS * CHANNELS
+
+
+def _validate_module_contract(net: Any) -> None:
+    expected = {
+        "layout": "flat",
+        "mode": "cross",
+        "op_type": "glu",
+        "n_focus": N_FOCUS,
+        "n_frames": N_FRAMES,
+        "channels": CHANNELS,
+        "dtype": torch.float32,
+    }
+    mismatches = {
+        name: (getattr(net, name, None), value)
+        for name, value in expected.items()
+        if getattr(net, name, None) != value
+    }
+    frame_expand = getattr(net, "frame_expand", None)
+    frame_contract = getattr(net, "frame_contract", None)
+    if frame_expand is None or frame_contract is None:
+        mismatches["frame_modules"] = (
+            (frame_expand is not None, frame_contract is not None),
+            (True, True),
+        )
+    else:
+        expected_frame = ("packed", N_FRAMES, CHANNELS)
+        for name, module in (
+            ("frame_expand", frame_expand),
+            ("frame_contract", frame_contract),
+        ):
+            actual_frame = (
+                getattr(module, "coefficient_layout", None),
+                getattr(module, "n_frames", None),
+                getattr(module, "channels", None),
+            )
+            if actual_frame != expected_frame:
+                mismatches[name] = (actual_frame, expected_frame)
+    if mismatches:
+        details = ", ".join(
+            f"{name}={actual!r} (expected {wanted!r})"
+            for name, (actual, wanted) in mismatches.items()
+        )
+        raise ValueError(f"packed message-grid module contract mismatch: {details}")
+
+
+def is_supported_message_grid(net: Any) -> bool:
+    """Return whether forward and manual backward share the same contract."""
+    try:
+        _validate_module_contract(net)
+    except ValueError:
+        return False
+    return True
+
+
+def _as_flat_ndfc(
+    net: Any,
+    name: str,
+    value: Tensor,
+    *,
+    like: Tensor | None = None,
+) -> Tensor:
+    """Adapt GridNet's flat layout without materializing a contiguous copy."""
+    expected_shape = (COEFF_DIM, HIDDEN_CHANNELS)
+    if value.ndim != 3 or tuple(value.shape[1:]) != expected_shape:
+        raise ValueError(
+            f"packed message-grid {name} must have shape (N, {COEFF_DIM}, "
+            f"{HIDDEN_CHANNELS}), got {tuple(value.shape)}"
+        )
+    if value.dtype != torch.float32:
+        raise ValueError(f"packed message-grid {name} must be FP32, got {value.dtype}")
+    if like is not None and (value.shape != like.shape or value.device != like.device):
+        raise ValueError(
+            f"packed message-grid {name} must match query shape/device; got "
+            f"shape={tuple(value.shape)}, device={value.device}"
+        )
+    # SO3Linear's einsum returns the valid flat GridNet layout with stride
+    # (64, N*64, 1), while Phase C returns compact (1024, 64, 1). Both split
+    # their unit-stride final axis into (F=2, C=32) without a copy.
+    if value.stride(-1) != 1:
+        raise ValueError(
+            f"packed message-grid {name} requires a unit-stride folded F*C "
+            f"axis, got stride={tuple(value.stride())}"
+        )
+    value_ndfc, shape_info = net._to_ndfc(value)
+    expected_ndfc = (value.shape[0], COEFF_DIM, N_FOCUS, CHANNELS)
+    if (
+        tuple(shape_info) != tuple(value.shape)
+        or tuple(value_ndfc.shape) != expected_ndfc
+    ):
+        raise ValueError(
+            f"packed message-grid {name} did not adapt to {expected_ndfc}; got "
+            f"shape={tuple(value_ndfc.shape)}, stride={tuple(value_ndfc.stride())}"
+        )
+    return value_ndfc
+
+
+def _validate_contract(
+    net: Any,
+    query: Tensor,
+    context: Tensor,
+) -> tuple[Tensor, Tensor]:
+    _validate_module_contract(net)
+    query_ndfc = _as_flat_ndfc(net, "query", query)
+    context_ndfc = _as_flat_ndfc(net, "context", context, like=query)
+    return query_ndfc, context_ndfc
+
+
+def _expanded_frame_weight(module: Any) -> Tensor:
+    return module.weight.index_select(0, module.degree_index)
+
+
+def _frame_expand_packed(module: Any, coeff: Tensor) -> Tensor:
+    weight = _expanded_frame_weight(module).view(
+        COEFF_DIM,
+        CHANNELS,
+        N_FRAMES,
+        CHANNELS,
+    )
+    # Output order D,K,F,C is the native CuTe grid-product contract. The
+    # trailing F,C panel remains unit-stride and therefore coalesced.
+    # PyTorch's degree-batched einsum naturally returns a degree-major stride.
+    # Normalize once at this producer because the CuTe consumer's packed
+    # contract is compact (N,D,K,F,C), not because the flat input was strided.
+    return torch.einsum("ndfi,dikc->ndkfc", coeff, weight).contiguous()
+
+
+def _frame_expand_packed_backward(
+    module: Any,
+    grad_packed: Tensor,
+) -> Tensor:
+    weight = _expanded_frame_weight(module).view(
+        COEFF_DIM,
+        CHANNELS,
+        N_FRAMES,
+        CHANNELS,
+    )
+    return torch.einsum("ndkfc,dikc->ndfi", grad_packed, weight)
+
+
+def _frame_contract_packed(module: Any, coeff_packed: Tensor) -> Tensor:
+    weight = _expanded_frame_weight(module).view(
+        COEFF_DIM,
+        N_FRAMES,
+        CHANNELS,
+        CHANNELS,
+    )
+    return torch.einsum("ndkfc,dkco->ndfo", coeff_packed, weight)
+
+
+def _frame_contract_packed_backward(module: Any, grad_out: Tensor) -> Tensor:
+    weight = _expanded_frame_weight(module).view(
+        COEFF_DIM,
+        N_FRAMES,
+        CHANNELS,
+        CHANNELS,
+    )
+    return torch.einsum("ndfo,dkco->ndkfc", grad_out, weight)
+
+
+def _focus_linear_backward_input(linear: Any, grad_out: Tensor) -> Tensor:
+    weight = linear.weight.view(linear.in_channels, linear.n_focus, linear.out_channels)
+    return torch.einsum("bfo,ifo->bfi", grad_out, weight)
+
+
+def _swiglu_backward_input(x: Tensor, grad_out: Tensor) -> Tensor:
+    gate, value = torch.chunk(x, chunks=2, dim=-1)
+    sigmoid = torch.sigmoid(gate)
+    grad_gate = grad_out * value * (sigmoid + gate * sigmoid * (1.0 - sigmoid))
+    grad_value = grad_out * gate * sigmoid
+    return torch.cat([grad_gate, grad_value], dim=-1)
+
+
+def run_packed_message_grid_forward(
+    net: Any,
+    query_flat: Tensor,
+    context_flat: Tensor,
+    *,
+    return_product: bool = False,
+    sm90_state: _Sm90MessageGridState | None = None,
+) -> Tensor | tuple[Tensor, Tensor]:
+    """Run only the message-grid module and return its canonical flat output."""
+    query, context = _validate_contract(net, query_flat, context_flat)
+    from .kernels.message_grid_product import (
+        run_message_grid_product,
+    )
+
+    nodes = query_flat.shape[0]
+    scalar_pair = torch.cat([query[:, 0], context[:, 0]], dim=-1).to(net.dtype)
+
+    left_packed = _frame_expand_packed(net.frame_expand, query)
+    right_packed = _frame_expand_packed(net.frame_expand, context)
+    left = left_packed.view(nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS)
+    right = right_packed.view(nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS)
+    if sm90_state is None:
+        product_flat = run_message_grid_product(
+            left,
+            right,
+            net.projector.to_grid_mat,
+            net.projector.from_grid_mat,
+        )
+    else:
+        from .sm90.message_grid_gaunt import (
+            run_sm90_gaunt_forward,
+        )
+
+        product_flat = run_sm90_gaunt_forward(
+            left,
+            right,
+            sm90_state.schedule,
+        )
+    product = product_flat.view(
+        nodes,
+        COEFF_DIM,
+        N_FRAMES,
+        N_FOCUS,
+        CHANNELS,
+    )
+
+    scalar_out = net.scalar_act(scalar_pair)
+    scalar_gate = torch.sigmoid(net.scalar_gate(scalar_pair))
+    coeff_packed = product * scalar_gate[:, None, None, :, :]
+    coeff_packed[:, 0, net.frame_zero_index].add_(scalar_out)
+    coeff = _frame_contract_packed(net.frame_contract, coeff_packed)
+    if net.residual_scale is not None:
+        coeff = coeff * net.residual_scale.view(1, 1, N_FOCUS, CHANNELS)
+    output = coeff.reshape_as(query_flat)
+    if return_product:
+        return output, product_flat
+    return output
+
+
+def run_packed_message_grid_backward(
+    net: Any,
+    query_flat: Tensor,
+    context_flat: Tensor,
+    grad_out_flat: Tensor,
+    *,
+    product_flat: Tensor | None = None,
+) -> tuple[Tensor, Tensor]:
+    """Return query/context adjoints while retaining packed grid intermediates."""
+    query, context = _validate_contract(net, query_flat, context_flat)
+    from .kernels.message_grid_product import (
+        run_message_grid_product,
+        run_message_grid_product_backward,
+    )
+
+    nodes = query_flat.shape[0]
+    scalar_pair = torch.cat([query[:, 0], context[:, 0]], dim=-1).to(net.dtype)
+
+    left_packed = _frame_expand_packed(net.frame_expand, query)
+    right_packed = _frame_expand_packed(net.frame_expand, context)
+    left = left_packed.view(nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS)
+    right = right_packed.view(nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS)
+    if product_flat is None:
+        product_flat = run_message_grid_product(
+            left,
+            right,
+            net.projector.to_grid_mat,
+            net.projector.from_grid_mat,
+        )
+    elif (
+        tuple(product_flat.shape) != (nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS)
+        or product_flat.dtype != torch.float32
+        or product_flat.device != query_flat.device
+        or not product_flat.is_contiguous()
+    ):
+        raise ValueError(
+            "saved packed message-grid product must be contiguous FP32 with shape "
+            f"({nodes}, {COEFF_DIM * N_FRAMES}, {HIDDEN_CHANNELS}) on "
+            f"{query_flat.device}"
+        )
+    product = product_flat.view(
+        nodes,
+        COEFF_DIM,
+        N_FRAMES,
+        N_FOCUS,
+        CHANNELS,
+    )
+
+    scalar_gate = torch.sigmoid(net.scalar_gate(scalar_pair))
+
+    grad = _as_flat_ndfc(
+        net,
+        "grad_out",
+        grad_out_flat,
+        like=query_flat,
+    ).to(net.dtype)
+    if net.residual_scale is not None:
+        grad = grad * net.residual_scale.view(1, 1, N_FOCUS, CHANNELS)
+    grad_scalar_packed = _frame_contract_packed_backward(net.frame_contract, grad)
+
+    grad_product = grad_scalar_packed * scalar_gate[:, None, None, :, :]
+    grad_scalar_gate = (grad_scalar_packed * product).sum(dim=(1, 2))
+    grad_scalar_out = grad_scalar_packed[:, 0, net.frame_zero_index]
+    grad_scalar_logits = grad_scalar_gate * scalar_gate * (1.0 - scalar_gate)
+    grad_scalar_pair = _focus_linear_backward_input(
+        net.scalar_gate,
+        grad_scalar_logits,
+    ) + _swiglu_backward_input(scalar_pair, grad_scalar_out)
+
+    # Broadcast multiplication preserves the degree-major einsum layout. The
+    # packed CuTe adjoint requires coefficient-major compact storage.
+    grad_product_flat = grad_product.contiguous().view(
+        nodes,
+        COEFF_DIM * N_FRAMES,
+        HIDDEN_CHANNELS,
+    )
+    grad_left, grad_right = run_message_grid_product_backward(
+        grad_product_flat,
+        left,
+        right,
+        net.projector.to_grid_mat,
+        net.projector.from_grid_mat,
+    )
+    grad_query = _frame_expand_packed_backward(
+        net.frame_expand,
+        grad_left.view(nodes, COEFF_DIM, N_FRAMES, N_FOCUS, CHANNELS),
+    )
+    grad_context = _frame_expand_packed_backward(
+        net.frame_expand,
+        grad_right.view(nodes, COEFF_DIM, N_FRAMES, N_FOCUS, CHANNELS),
+    )
+    grad_query[:, 0].add_(grad_scalar_pair[:, :, :CHANNELS])
+    grad_context[:, 0].add_(grad_scalar_pair[:, :, CHANNELS:])
+    return (
+        grad_query.reshape_as(query_flat).to(dtype=query_flat.dtype),
+        grad_context.reshape_as(context_flat).to(dtype=context_flat.dtype),
+    )
+
+
+__all__ = [
+    "run_packed_message_grid_backward",
+    "run_packed_message_grid_forward",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py b/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py
new file mode 100644
index 0000000000..99cb83220e
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py
@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Build CuTe SO2 metadata when the input does not carry graph-owned CSR."""
+
+from __future__ import (
+    annotations,
+)
+
+import torch
+
+
+def _destination_row_ptr_impl(
+    dst: torch.Tensor,
+    n_nodes: int,
+) -> torch.Tensor:
+    """Build destination CSR where real storage offsets remain visible."""
+    boundaries = torch.arange(
+        n_nodes + 1,
+        device=dst.device,
+        dtype=dst.dtype,
+    )
+    return torch.searchsorted(
+        dst.contiguous(),
+        boundaries,
+        out_int32=True,
+    ).contiguous()
+
+
+_destination_row_ptr_op = torch.library.custom_op(
+    "sezm_cute::destination_row_ptr",
+    mutates_args=(),
+)(_destination_row_ptr_impl)
+
+
+@_destination_row_ptr_op.register_fake
+def _(
+    dst: torch.Tensor,
+    n_nodes: int,
+) -> torch.Tensor:
+    return torch.empty(
+        (n_nodes + 1,),
+        device=dst.device,
+        dtype=torch.int32,
+    )
+
+
+def build_destination_row_ptr(dst: torch.Tensor, n_nodes: int) -> torch.Tensor:
+    """Build integer CSR for sorted destinations through an opaque runtime op.
+
+    Parameters
+    ----------
+    dst : torch.Tensor
+        Nondecreasing int32 or int64 destinations with shape ``(E,)``.
+    n_nodes : int
+        Number of nodes addressed by the edge list.
+
+    Returns
+    -------
+    torch.Tensor
+        Contiguous int32 row pointers with shape ``(n_nodes + 1,)`` and no
+        autograd history. The runtime op preserves nonzero storage offsets
+        when ``dst`` is a view such as ``edge_index[1]`` under Inductor.
+    """
+    return _destination_row_ptr_op(dst, n_nodes)
+
+
+def build_sorted_edge_index_metadata(
+    src: torch.Tensor,
+    dst: torch.Tensor,
+    n_nodes: int,
+    *,
+    validate_sorted: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Build destination and source CSR metadata for one sorted edge list.
+
+    Parameters
+    ----------
+    src : torch.Tensor
+        Source indices with shape ``(E,)``.
+    dst : torch.Tensor
+        Nondecreasing destination indices with shape ``(E,)``.
+    n_nodes : int
+        Number of nodes addressed by the edge list.
+    validate_sorted : bool, default: False
+        Whether to assert the destination-order contract on the device.
+
+    Returns
+    -------
+    tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        Destination row pointers, the source-major edge permutation, and source
+        row pointers. All three tensors use int32 storage.
+
+    Raises
+    ------
+    ValueError
+        If the edge arrays have incompatible shapes or devices, or if their
+        edge count exceeds int32 indexing.
+    TypeError
+        If either edge-index tensor is not int32 or int64.
+    """
+    if src.dim() != 1 or dst.dim() != 1:
+        raise ValueError("src and dst must be one-dimensional")
+    if src.shape != dst.shape:
+        raise ValueError("src and dst must have the same shape")
+    if src.device != dst.device:
+        raise ValueError("src and dst must be on the same device")
+    if src.dtype not in (torch.int32, torch.int64):
+        raise TypeError("src must have dtype int32 or int64")
+    if dst.dtype not in (torch.int32, torch.int64):
+        raise TypeError("dst must have dtype int32 or int64")
+    if n_nodes < 0:
+        raise ValueError("n_nodes must be non-negative")
+    if src.numel() > 2**31 - 1:
+        raise ValueError("sorted edge metadata requires E <= 2**31 - 1")
+
+    src = src.contiguous()
+    dst = dst.contiguous()
+    if validate_sorted and dst.numel() > 1:
+        torch._assert_async(
+            torch.all(dst[1:] >= dst[:-1]),
+            "Neo SO2 destinations_sorted=True requires monotonically "
+            "nondecreasing destination indices",
+        )
+    destination_row_ptr = _destination_row_ptr_op(dst, n_nodes)
+
+    source_order_i64 = torch.argsort(src, stable=True)
+    sorted_src = src.index_select(0, source_order_i64)
+    source_boundaries = torch.arange(
+        n_nodes + 1,
+        device=src.device,
+        dtype=src.dtype,
+    )
+    source_row_ptr = torch.searchsorted(
+        sorted_src,
+        source_boundaries,
+        out_int32=True,
+    ).contiguous()
+    source_order = source_order_i64.to(dtype=torch.int32).contiguous()
+    return destination_row_ptr, source_order, source_row_ptr
+
+
+__all__ = ["build_destination_row_ptr", "build_sorted_edge_index_metadata"]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py
new file mode 100644
index 0000000000..50c8d73927
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py
@@ -0,0 +1,2292 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Opt-in CuTe Neo SO2 custom op for SeZM/DPA4 inference."""
+
+from __future__ import (
+    annotations,
+)
+
+import threading
+import weakref
+from dataclasses import (
+    dataclass,
+    field,
+    replace,
+)
+from functools import (
+    lru_cache,
+)
+from itertools import (
+    count,
+)
+from types import (
+    SimpleNamespace,
+)
+from typing import (
+    Any,
+)
+
+import torch
+from torch import (
+    Tensor,
+)
+
+from .. import (
+    runtime_policy,
+)
+from .message_grid import (
+    is_supported_message_grid,
+)
+
+
+@dataclass(frozen=True)
+class NeoSO2RuntimeConfig:
+    """Architecture-specific choices for the compact Neo SO(2) path."""
+
+    native_sm90_path: bool = False
+    per_focus_so2_fwd_pair: bool = False
+    combined_so2_gate: bool = False
+
+
+@dataclass(frozen=True)
+class NeoSO2Spec:
+    """Shape and branch contract for one Neo SO(2) unit replacement."""
+
+    lmax: int
+    node_lmax: int
+    mmax: int
+    full_dim: int
+    reduced_dim: int
+    channels: int
+    n_focus: int
+    focus_dim: int
+    hidden_channels: int
+    so2_layers: int
+    n_atten_head: int
+    radial_so2_mode: str
+    radial_so2_rank: int
+    so2_norm: bool
+    focus_compete: bool
+    message_node_so3: bool
+    atten_f_mix: bool
+    atten_v_proj: bool
+    atten_o_proj: bool
+    mlp_bias: bool
+    layer_scale: bool
+    use_so2_attn_res: bool
+    has_pre_so2_norm: bool
+    has_post_so2_norm: bool
+
+    @property
+    def is_current_neo_target(self) -> bool:
+        return (
+            self.lmax == 3
+            and self.node_lmax == self.lmax
+            and self.mmax == 1
+            and self.full_dim == 16
+            and self.reduced_dim == 10
+            and self.channels == 32
+            and self.n_focus == 2
+            and self.focus_dim == 32
+            and self.hidden_channels == 64
+            and self.so2_layers == 3
+            and self.n_atten_head == 1
+            and self.radial_so2_mode == "degree_channel"
+            and self.radial_so2_rank == 1
+            and not self.so2_norm
+            and self.focus_compete
+            and self.message_node_so3
+            and not self.atten_f_mix
+            and not self.atten_v_proj
+            and not self.atten_o_proj
+            and not self.mlp_bias
+            and not self.layer_scale
+            and not self.use_so2_attn_res
+            and not self.has_pre_so2_norm
+            and self.has_post_so2_norm
+        )
+
+
+def _is_identity(module: Any) -> bool:
+    return (
+        module.__class__.__name__ == "Identity"
+        or module.__class__.__name__ == "_Identity"
+    )
+
+
+def get_neo_so2_spec(block: Any) -> NeoSO2Spec:
+    """Extract the SO2 contract from a SeZM interaction block."""
+    so2 = block.so2_conv
+    return NeoSO2Spec(
+        lmax=int(so2.lmax),
+        node_lmax=int(block.node_lmax),
+        mmax=int(so2.mmax),
+        full_dim=int(so2.ebed_dim_full),
+        reduced_dim=int(so2.reduced_dim),
+        channels=int(so2.channels),
+        n_focus=int(so2.n_focus),
+        focus_dim=int(so2.so2_focus_dim),
+        hidden_channels=int(so2.hidden_channels),
+        so2_layers=int(so2.mixing_layers),
+        n_atten_head=int(so2.n_atten_head),
+        radial_so2_mode=str(so2.radial_so2_mode),
+        radial_so2_rank=int(so2.radial_so2_rank),
+        so2_norm=bool(so2.so2_norm),
+        focus_compete=bool(so2.focus_compete),
+        message_node_so3=getattr(so2, "message_node_grid_product", None) is not None,
+        atten_f_mix=bool(so2.atten_f_mix),
+        atten_v_proj=getattr(so2, "attn_v_proj", None) is not None,
+        atten_o_proj=getattr(so2, "attn_o_proj", None) is not None,
+        mlp_bias=bool(so2.mlp_bias),
+        layer_scale=bool(so2.layer_scale),
+        use_so2_attn_res=bool(so2.use_so2_attn_res),
+        has_pre_so2_norm=not _is_identity(block.pre_so2_norm),
+        has_post_so2_norm=not _is_identity(block.post_so2_norm),
+    )
+
+
+def has_equivariant_rms_norm_contract(module: Any) -> bool:
+    """Return whether a norm exposes the exact state consumed by Neo SO2.
+
+    Parameters
+    ----------
+    module : Any
+        Candidate equivariant normalization module.
+
+    Returns
+    -------
+    bool
+        Whether the module's layout and state match the fused kernel contract.
+    """
+    return (
+        type(module).__name__ == "EquivariantRMSNorm"
+        and int(getattr(module, "lmax", -1)) == 3
+        and int(getattr(module, "channels", -1)) == 32
+        and int(getattr(module, "n_focus", -1)) == 1
+        and tuple(getattr(getattr(module, "adam_scale", None), "shape", ()))
+        == (4, 1, 32)
+        and tuple(getattr(getattr(module, "bias", None), "shape", ())) == (1, 32)
+        and tuple(getattr(getattr(module, "balance_weight", None), "shape", ()))
+        == (16,)
+        and tuple(getattr(getattr(module, "expand_index", None), "shape", ())) == (16,)
+        and hasattr(module, "eps")
+    )
+
+
+def has_focus_compete_norm_contract(so2: Any) -> bool:
+    """Return whether focus competition matches the optional norm contract."""
+    module = getattr(so2, "focus_compete_norm", None)
+    if not bool(getattr(so2, "focus_norm", False)):
+        return module is None or _is_identity(module)
+    return (
+        type(module).__name__ == "ScalarRMSNorm"
+        and int(getattr(module, "channels", -1)) == 32
+        and int(getattr(module, "n_focus", -1)) == 2
+        and tuple(getattr(getattr(module, "adam_scale", None), "shape", ())) == (2, 32)
+        and hasattr(module, "eps")
+    )
+
+
+def is_supported_neo_so2_block(block: Any) -> bool:
+    """Return whether this block can use the current Neo CuTe SO2 path."""
+    so2 = block.so2_conv
+    return (
+        get_neo_so2_spec(block).is_current_neo_target
+        and has_focus_compete_norm_contract(so2)
+        and not bool(so2.edge_cartesian)
+        and getattr(so2, "node_cartesian_tp", None) is None
+        and has_equivariant_rms_norm_contract(block.post_so2_norm)
+        and is_supported_message_grid(so2.message_node_grid_product)
+    )
+
+
+def _module_floating_state_uses_strict_fp32(
+    module: Any,
+    *,
+    require_floating_tensor: bool,
+) -> bool:
+    """Check live floating parameters and buffers without reading device data."""
+    saw_floating_tensor = False
+    for getter_name in ("parameters", "buffers"):
+        getter = getattr(module, getter_name, None)
+        if getter is None:
+            return False
+        for tensor in getter():
+            if tensor.is_floating_point():
+                saw_floating_tensor = True
+                if tensor.dtype != torch.float32:
+                    return False
+    return saw_floating_tensor or not require_floating_tensor
+
+
+def _module_uses_strict_fp32(module: Any) -> bool:
+    return _module_floating_state_uses_strict_fp32(
+        module,
+        require_floating_tensor=True,
+    )
+
+
+def _module_is_frozen(module: Any) -> bool:
+    """Return whether autograd cannot request gradients for module state."""
+    parameters = getattr(module, "parameters", None)
+    return parameters is not None and not any(
+        parameter.requires_grad for parameter in parameters()
+    )
+
+
+def _tensor_is_aligned(tensor: Tensor, alignment: int = 16) -> bool:
+    return tensor.data_ptr() % alignment == 0
+
+
+def _module_state_is_aligned(module: Any, alignment: int = 16) -> bool:
+    """Check CuTe's declared alignment contract for parameters and buffers."""
+    return all(
+        not tensor.is_floating_point() or _tensor_is_aligned(tensor, alignment)
+        for getter_name in ("parameters", "buffers")
+        for tensor in getattr(module, getter_name)()
+    )
+
+
+def _gate_expand_index_is_supported(block: Any) -> bool:
+    """Check the degree-to-gate map assumed by fused Neo gate kernels."""
+    non_linearities = getattr(
+        getattr(block, "so2_conv", None),
+        "non_linearities",
+        None,
+    )
+    if non_linearities is None:
+        return True
+    buffers = tuple(
+        expand_index
+        for non_linear in non_linearities
+        if (expand_index := getattr(non_linear, "expand_index", None)) is not None
+        and expand_index.numel() > 0
+    )
+    signature = tuple(
+        (
+            tensor.data_ptr(),
+            tensor._version,
+            tensor.dtype,
+            tensor.device,
+            tuple(tensor.shape),
+        )
+        for tensor in buffers
+    )
+    cached = getattr(block, "_deepmd_cute_gate_expand_contract", None)
+    if cached is not None and cached[0] == signature:
+        return bool(cached[1])
+    expected = torch.tensor(
+        [0, 1, 2, 0, 1, 2, 0, 1, 2],
+        dtype=torch.long,
+        device="cpu",
+    )
+    supported = all(
+        torch.equal(
+            expand_index.detach().to(device="cpu", dtype=torch.long),
+            expected,
+        )
+        for expand_index in buffers
+    )
+    block._deepmd_cute_gate_expand_contract = (signature, supported)
+    return supported
+
+
+def _gate_expand_index_structure_is_supported(block: Any) -> bool:
+    """Check graph-visible gate-index metadata without reading tensor values."""
+    non_linearities = getattr(
+        getattr(block, "so2_conv", None),
+        "non_linearities",
+        None,
+    )
+    if non_linearities is None:
+        return True
+    return all(
+        expand_index.dtype == torch.long and tuple(expand_index.shape) == (9,)
+        for non_linear in non_linearities
+        if (expand_index := getattr(non_linear, "expand_index", None)) is not None
+        and expand_index.numel() > 0
+    )
+
+
+def _aligned_contiguous(tensor: Tensor, alignment: int = 16) -> Tensor:
+    """Return canonical storage satisfying CuTe's assumed alignment."""
+    if tensor.is_contiguous() and _tensor_is_aligned(tensor, alignment):
+        return tensor
+    return tensor.clone(memory_format=torch.contiguous_format)
+
+
+def _producer_modules_use_strict_fp32(modules: Any) -> bool:
+    return all(
+        _module_floating_state_uses_strict_fp32(
+            module,
+            require_floating_tensor=False,
+        )
+        for module in modules
+    )
+
+
+def _dtypes_use_strict_fp32(dtypes: Any) -> bool:
+    return all(dtype == torch.float32 for dtype in dtypes)
+
+
+def is_supported_so2_compute_capability(
+    compute_capability: tuple[int, int],
+) -> bool:
+    """Return whether SO2 supports this compute capability."""
+    return runtime_policy.is_supported_so2_capability(compute_capability)
+
+
+def _device_compute_capability(device: torch.device) -> tuple[int, int]:
+    device_index = device.index
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    return _cuda_compute_capability(device_index)
+
+
+def _tensor_compute_capability(tensor: Any) -> tuple[int, int] | None:
+    """Resolve dispatch capability from an operand instead of global CUDA state."""
+    device = getattr(tensor, "device", None)
+    if device is None or device.type != "cuda":
+        return None
+    return _device_compute_capability(device)
+
+
+def _device_is_supported_for_so2(device: torch.device) -> bool:
+    # Metadata-only eligibility checks may use a CUDA device on a host without
+    # a CUDA runtime. Concrete CUDA dispatch always validates the capability.
+    return not torch.cuda.is_available() or is_supported_so2_compute_capability(
+        _device_compute_capability(device)
+    )
+
+
+def packed_wigner_edges_eligible(
+    *,
+    candidate: bool,
+    edge_count: int,
+    node_count: int,
+    destinations_sorted: bool,
+    runtime_dtypes: Any = (),
+) -> bool:
+    """Finish packed eligibility from host-side edge-order provenance."""
+    return (
+        candidate
+        and edge_count > 0
+        and destinations_sorted
+        and _dtypes_use_strict_fp32(runtime_dtypes)
+        and runtime_policy.so2_int32_indexing_is_safe(
+            edge_count=edge_count,
+            node_count=node_count,
+        )
+    )
+
+
+def is_neo_so2_static_eligible(
+    block: Any,
+    *,
+    training: bool,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> bool:
+    """Check SO2 conditions that are stable during one descriptor forward."""
+    return (
+        runtime_policy.is_cute_infer_enabled()
+        and not training
+        and not bool(getattr(block, "training", False))
+        and device.type == "cuda"
+        and _device_is_supported_for_so2(device)
+        and dtype == torch.float32
+        and not torch.is_autocast_enabled(device.type)
+        and _has_frozen_fp32_contract(block)
+        and _module_state_is_aligned(block)
+        and _gate_expand_index_structure_is_supported(block)
+        and getattr(block, "_deepmd_cute_so2_state", None) is not False
+        and is_supported_neo_so2_block(block)
+    )
+
+
+def _has_frozen_fp32_contract(block: Any) -> bool:
+    """Recheck mutable precision and parameter state before inference dispatch."""
+    return (
+        not block.training
+        and runtime_policy.uses_strict_fp32_matmul()
+        and _module_uses_strict_fp32(block)
+        and _module_is_frozen(block)
+    )
+
+
+def is_neo_so2_runtime_eligible(
+    block: Any,
+    *,
+    training: bool,
+    device: torch.device,
+    dtype: torch.dtype,
+    edge_count: int,
+    node_count: int,
+    destinations_sorted: bool,
+) -> bool:
+    """Return the exact strict-FP32 inference contract for SO2 dispatch."""
+    return (
+        edge_count > 0
+        and destinations_sorted
+        and is_neo_so2_static_eligible(
+            block,
+            training=training,
+            device=device,
+            dtype=dtype,
+        )
+    )
+
+
+def is_packed_wigner_candidate(
+    *,
+    blocks: Any,
+    training: bool,
+    device: torch.device,
+    dtype: torch.dtype,
+    producer_modules: Any = (),
+    producer_dtypes: Any = (),
+    has_edge_src_gate: bool = False,
+) -> bool:
+    """Check packed-Wigner conditions known before edge construction."""
+    if has_edge_src_gate:
+        # SO2 currently falls back for SFPG so eager must receive dense Wigner data.
+        return False
+    block_tuple = tuple(blocks)
+    if device.type == "cuda":
+        try:
+            compute_capability = _device_compute_capability(device)
+        except (AssertionError, RuntimeError):
+            packed_wigner_enabled = False
+        else:
+            packed_wigner_enabled = (
+                runtime_policy.is_cute_infer_enabled()
+                and runtime_policy.is_supported_so2_capability(compute_capability)
+            )
+    else:
+        packed_wigner_enabled = False
+    if (
+        not block_tuple
+        or not packed_wigner_enabled
+        or not _producer_modules_use_strict_fp32(producer_modules)
+        or not _dtypes_use_strict_fp32(producer_dtypes)
+        or torch.is_autocast_enabled(device.type)
+    ):
+        return False
+    return all(
+        is_neo_so2_static_eligible(
+            block,
+            training=training,
+            device=device,
+            dtype=dtype,
+        )
+        for block in block_tuple
+    )
+
+
+class _RegistryEntry:
+    """Weakly retain a block so discarded models do not leak the registry."""
+
+    def __init__(
+        self,
+        block: Any,
+        config: Any,
+        *,
+        on_collect: Any | None = None,
+    ) -> None:
+        try:
+            self._block_ref = weakref.ref(block, on_collect)
+        except TypeError:
+            self._block_ref = lambda: block
+        self.config = config
+
+    @property
+    def block(self) -> Any:
+        block = self._block_ref()
+        if block is None:
+            raise RuntimeError("the registered Neo SO2 block has been released")
+        return block
+
+
+@dataclass(frozen=True)
+class _RegisteredSO2State:
+    device_index: int
+    handle: int
+    config: NeoSO2RuntimeConfig
+
+
+@dataclass
+class _RunnerState:
+    runner: Any | None
+    backward_calls: int = 0
+    reservation_lock: threading.Lock = field(
+        default_factory=threading.Lock,
+        repr=False,
+        compare=False,
+    )
+
+
+_REGISTRY: dict[int, _RegistryEntry] = {}
+_HANDLE_COUNTER = count(1)
+_REGISTRY_LOCK = threading.Lock()
+_PACKED_RUNNER_CACHE: dict[tuple[str, int | None, int], _RunnerState] = {}
+_PACKED_RUNNER_CACHE_LOCK = threading.Lock()
+
+
+def _runner_compile_identity(runner: Any) -> tuple[int, int, int]:
+    """Return the runner's device-specific CuTe compilation identity."""
+    return getattr(runner, "compile_identity", (-1, 0, 0))
+
+
+def _compile_on_runner_device(
+    compile_identity: tuple[int, int, int],
+    compiler: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    """Compile under the CUDA device represented by the cache key."""
+    device_index = int(compile_identity[0])
+    if device_index < 0:
+        return compiler(*args, **kwargs)
+    with torch.cuda.device(device_index):
+        return compiler(*args, **kwargs)
+
+
+@lru_cache(maxsize=8)
+def _compile_output_gate_backward(
+    compile_identity: tuple[int, int, int], eps: float
+) -> Any:
+    from .kernels.output_gate_backward import (
+        compile_neo_output_gate_backward,
+    )
+
+    return _compile_on_runner_device(
+        compile_identity,
+        compile_neo_output_gate_backward,
+        eps,
+    )
+
+
+def register_cute_so2_block(block: Any, config: Any) -> int:
+    """Register a DeePMD block/config pair and return a stable integer handle."""
+
+    def remove_collected_entry(block_ref: weakref.ReferenceType[Any]) -> None:
+        with _REGISTRY_LOCK:
+            entry = _REGISTRY.get(handle)
+            if entry is not None and entry._block_ref is block_ref:
+                _REGISTRY.pop(handle, None)
+
+    with _REGISTRY_LOCK:
+        handle = next(_HANDLE_COUNTER)
+        _REGISTRY[handle] = _RegistryEntry(
+            block=block,
+            config=config,
+            on_collect=remove_collected_entry,
+        )
+    return handle
+
+
+def invalidate_cute_so2_state(block: Any) -> None:
+    """Release a block's registered CuTe state after its modules change."""
+    state = getattr(block, "_deepmd_cute_so2_state", None)
+    if isinstance(state, _RegisteredSO2State):
+        with _REGISTRY_LOCK:
+            _REGISTRY.pop(state.handle, None)
+    if hasattr(block, "_deepmd_cute_so2_state"):
+        delattr(block, "_deepmd_cute_so2_state")
+    if hasattr(block, "_deepmd_cute_gate_expand_contract"):
+        delattr(block, "_deepmd_cute_gate_expand_contract")
+
+
+def _validate_gate_expand_index(block: Any) -> None:
+    """Pin the degree-to-gate map assumed by fused Neo gate kernels."""
+    if not _gate_expand_index_is_supported(block):
+        raise ValueError(
+            "Neo SO2 fused gate kernels require expand_index=[0,1,2,0,1,2,0,1,2]"
+        )
+
+
+@torch.compiler.disable
+def _register_cute_so2_state(
+    block: Any,
+    device_index: int,
+    config: NeoSO2RuntimeConfig,
+) -> _RegisteredSO2State | None:
+    """Publish cold SO2 state eagerly before its handle enters a custom op."""
+    _validate_gate_expand_index(block)
+    old_state = getattr(block, "_deepmd_cute_so2_state", None)
+    if isinstance(old_state, _RegisteredSO2State):
+        with _REGISTRY_LOCK:
+            old_entry = _REGISTRY.get(old_state.handle)
+            if (
+                old_state.device_index == device_index
+                and old_state.config == config
+                and old_entry is not None
+                and old_entry.block is block
+            ):
+                return old_state
+            _REGISTRY.pop(old_state.handle, None)
+    if not _module_state_is_aligned(block):
+        # Module parameters and buffers are frozen for this inference path, so
+        # cache a failed static contract until explicit state invalidation.
+        block._deepmd_cute_so2_state = False
+        return None
+    state = _RegisteredSO2State(
+        device_index=device_index,
+        handle=register_cute_so2_block(block, config),
+        config=config,
+    )
+    block._deepmd_cute_so2_state = state
+    return state
+
+
+@torch.compiler.disable
+def prepare_cute_so2_blocks(
+    blocks: Any,
+    *,
+    training: bool,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> bool:
+    """Validate and register SO2 state before model graph capture begins."""
+    if training or device.type != "cuda" or dtype != torch.float32:
+        return False
+    block_tuple = tuple(blocks)
+    if not block_tuple:
+        return False
+    device_index = device.index
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    compute_capability = _cuda_compute_capability(device_index)
+    if not is_supported_so2_compute_capability(compute_capability):
+        return False
+    config = _architecture_default_config(compute_capability)
+    if not all(
+        is_neo_so2_static_eligible(
+            block,
+            training=training,
+            device=device,
+            dtype=dtype,
+        )
+        for block in block_tuple
+    ):
+        return False
+    return all(
+        _register_cute_so2_state(block, device_index, config) is not None
+        for block in block_tuple
+    )
+
+
+def _runner_token_key(
+    runner_token: Tensor, *, path: str
+) -> tuple[str, int | None, int]:
+    if runner_token.dtype != torch.uint8 or runner_token.numel() != 1:
+        raise ValueError(f"{path} Neo SO2 runner token must be one uint8 value")
+    return (
+        runner_token.device.type,
+        runner_token.device.index,
+        int(runner_token.data_ptr()),
+    )
+
+
+def _packed_runner_key(runner_token: Tensor) -> tuple[str, int | None, int]:
+    return _runner_token_key(runner_token, path="packed")
+
+
+def _release_packed_runner(
+    key: tuple[str, int | None, int],
+    state: _RunnerState,
+) -> None:
+    with _PACKED_RUNNER_CACHE_LOCK:
+        if _PACKED_RUNNER_CACHE.get(key) is state:
+            _PACKED_RUNNER_CACHE.pop(key, None)
+
+
+def _store_packed_runner(runner_token: Tensor, runner: Any) -> None:
+    key = _packed_runner_key(runner_token)
+    state = _RunnerState(runner=runner)
+    with _PACKED_RUNNER_CACHE_LOCK:
+        old_state = _PACKED_RUNNER_CACHE.get(key)
+        if old_state is not None:
+            raise RuntimeError("packed Neo SO2 runner token is already outstanding")
+        else:
+            _PACKED_RUNNER_CACHE[key] = state
+    weakref.finalize(runner_token, _release_packed_runner, key, state)
+
+
+def _borrow_packed_runner(runner_token: Tensor) -> Any | None:
+    key = _packed_runner_key(runner_token)
+    with _PACKED_RUNNER_CACHE_LOCK:
+        state = _PACKED_RUNNER_CACHE.get(key)
+    if state is None:
+        raise RuntimeError("packed Neo SO2 runner token is not outstanding")
+    with state.reservation_lock:
+        # Transfer the forward runner once; retained VJPs rebuild isolated workspaces.
+        runner = state.runner
+        state.runner = None
+        state.backward_calls += 1
+    return runner
+
+
+def _edge_src_gate_arg(edge_src_gate: Tensor) -> Tensor | None:
+    return None if edge_src_gate.numel() == 0 else edge_src_gate
+
+
+def _layout_like(tensor: Tensor, like: Tensor) -> Tensor:
+    if tensor.shape == like.shape and tensor.stride() == like.stride():
+        return tensor
+    out = torch.empty_strided(
+        like.shape,
+        like.stride(),
+        device=like.device,
+        dtype=like.dtype,
+    )
+    out.copy_(tensor)
+    return out
+
+
+def _grad_layout_like(tensor: Tensor, like: Tensor, *, skip: bool) -> Tensor:
+    if skip:
+        return tensor
+    return _layout_like(tensor, like)
+
+
+def _assert_grad_meta_contract(
+    actual: Tensor,
+    expected_like: Tensor,
+    *,
+    name: str,
+    expected_stride: tuple[int, ...] | None = None,
+) -> None:
+    """Fail before AOT consumes a gradient that contradicts register_fake."""
+    stride = expected_like.stride() if expected_stride is None else expected_stride
+    if (
+        actual.shape != expected_like.shape
+        or actual.dtype != expected_like.dtype
+        or actual.device != expected_like.device
+        or actual.stride() != stride
+    ):
+        raise RuntimeError(
+            f"Neo SO2 {name} gradient violates the custom-op meta contract: "
+            f"got shape={tuple(actual.shape)} stride={actual.stride()}, expected "
+            f"shape={tuple(expected_like.shape)} stride={stride}"
+        )
+
+
+def _assert_not_cuda_graph_capturing(path: str, tensor: Tensor) -> None:
+    """Reject stateful paths that cannot be safely captured and replayed."""
+    if tensor.is_cuda and torch.cuda.is_current_stream_capturing():
+        raise RuntimeError(
+            f"Neo SO2 {path} cannot run inside direct CUDA graph capture; "
+            "use torch.compile without wrapping this stateful custom op in "
+            "torch.cuda.graph"
+        )
+
+
+def _fake_x_wide_grad_like(x: Tensor, *, skip: bool) -> Tensor:
+    """Fake the native SO2 grad-x layout when layout restore is skipped.
+
+    The manual SO3 pre-mix backward returns a degree-major view with shape
+    ``(N, D, 1, C)`` and stride ``(C, N*C, C, 1)``.  If the fake kernel
+    promises ``empty_like(x)`` instead, AOTAutograd emits stride assertions for
+    a contiguous ``(N, D, 1, C)`` gradient and rejects the runtime result.
+    """
+    if not skip:
+        return torch.empty_like(x)
+    if x.ndim != 4:
+        return torch.empty_like(x)
+    n_node = x.shape[0]
+    channels = x.shape[-1]
+    return torch.empty_strided(
+        x.shape,
+        (channels, n_node * channels, channels, 1),
+        device=x.device,
+        dtype=x.dtype,
+    )
+
+
+def _equivariant_rmsnorm_backward(norm: Any, x: Tensor, grad_out: Tensor) -> Tensor:
+    # Matches EquivariantRMSNorm.forward for the inference case. Parameters are
+    # frozen; only the input gradient is needed for forces/stress.
+    in_dtype = x.dtype
+    norm_dtype = norm.adam_scale.dtype
+    xf = x.to(dtype=norm_dtype)
+    gf = grad_out.to(dtype=norm_dtype)
+    x0_in = xf[:, :1, :, :]
+    xt = xf[:, 1:, :, :]
+    x0 = x0_in - x0_in.mean(dim=-1, keepdim=True)
+
+    mean_variance = x0.square().sum(dim=(1, 3)) * norm.balance_weight[0]
+    if xt.numel() > 0:
+        mean_variance = mean_variance + torch.einsum(
+            "ndfc,d->nf", xt * xt, norm.balance_weight[1:]
+        )
+    eps = getattr(norm, "eps_tensor", norm.eps)
+    inv = torch.rsqrt(mean_variance + eps).unsqueeze(1).unsqueeze(-1)
+    expanded_scale = torch.index_select(
+        norm.adam_scale, dim=0, index=norm.expand_index
+    ).unsqueeze(0)
+
+    grad_pre = gf * expanded_scale
+    grad_x0 = grad_pre[:, :1, :, :]
+    grad_xt = grad_pre[:, 1:, :, :]
+
+    dinv = (grad_x0 * x0).sum(dim=(1, 3), keepdim=True)
+    if xt.numel() > 0:
+        dinv = dinv + (grad_xt * xt).sum(dim=(1, 3), keepdim=True)
+    dvar = -0.5 * dinv * inv.pow(3)
+
+    grad_centered = grad_x0 * inv + dvar * (2.0 * norm.balance_weight[0] * x0)
+    grad_x0_in = grad_centered - grad_centered.mean(dim=-1, keepdim=True)
+    if xt.numel() == 0:
+        return grad_x0_in.to(dtype=in_dtype)
+    grad_xt_in = grad_xt * inv + dvar * (
+        2.0 * norm.balance_weight[1:].view(1, -1, 1, 1) * xt
+    )
+    return torch.cat([grad_x0_in, grad_xt_in], dim=1).to(dtype=in_dtype)
+
+
+def _so3_linear_backward_input(linear: Any, grad_out: Tensor) -> Tensor:
+    weight = linear.weight.view(
+        linear.lmax + 1,
+        linear.in_channels,
+        linear.n_focus,
+        linear.out_channels,
+    )
+    weight_expanded = torch.index_select(weight, dim=0, index=linear.expand_index)
+    return torch.einsum("ndfo,difo->ndfi", grad_out, weight_expanded)
+
+
+def _focus_linear_forward(linear: Any, x: Tensor) -> Tensor:
+    weight = linear.weight.view(linear.in_channels, linear.n_focus, linear.out_channels)
+    out = torch.einsum("bfi,ifo->bfo", x, weight)
+    if linear.use_bias:
+        out = out + linear.bias.view(linear.n_focus, linear.out_channels).unsqueeze(0)
+    return out
+
+
+def _focus_linear_backward_input(linear: Any, grad_out: Tensor) -> Tensor:
+    weight = linear.weight.view(linear.in_channels, linear.n_focus, linear.out_channels)
+    return torch.einsum("bfo,ifo->bfi", grad_out, weight)
+
+
+def _swiglu_forward(x: Tensor) -> Tensor:
+    gate, value = torch.chunk(x, chunks=2, dim=-1)
+    return gate * torch.sigmoid(gate) * value
+
+
+def _swiglu_backward_input(x: Tensor, grad_out: Tensor) -> Tensor:
+    gate, value = torch.chunk(x, chunks=2, dim=-1)
+    sig = torch.sigmoid(gate)
+    grad_gate = grad_out * value * (sig + gate * sig * (1.0 - sig))
+    grad_value = grad_out * gate * sig
+    return torch.cat([grad_gate, grad_value], dim=-1)
+
+
+def _frame_expand_forward(module: Any, coeff: Tensor) -> Tensor:
+    weight = module.weight.index_select(0, module.degree_index)
+    return torch.einsum("ndfi,dio->ndfo", coeff, weight)
+
+
+def _frame_expand_backward_input(module: Any, grad_out: Tensor) -> Tensor:
+    weight = module.weight.index_select(0, module.degree_index)
+    return torch.einsum("ndfo,dio->ndfi", grad_out, weight)
+
+
+def _frame_contract_backward_input(module: Any, grad_out: Tensor) -> Tensor:
+    weight = module.weight.index_select(0, module.degree_index)
+    return torch.einsum("ndfo,dio->ndfi", grad_out, weight)
+
+
+def _neo_so2_linear_backward_input_with_residual(
+    so2_linear: Any,
+    grad_out: Tensor,
+    residual: Tensor,
+    *,
+    inplace_residual: bool = False,
+    out: Tensor | None = None,
+) -> Tensor:
+    from .linear import (
+        cached_neo_so2_linear_weights,
+    )
+
+    w0, wpair = cached_neo_so2_linear_weights(so2_linear)
+    cache = getattr(so2_linear, "_deepmd_cute_neo_manual_weights_t", None)
+    cache_key = (w0.data_ptr(), wpair.data_ptr(), w0.dtype, w0.device)
+    if (
+        cache is None
+        or cache[0] is not w0
+        or cache[1] is not wpair
+        or cache[2] != cache_key
+    ):
+        w0_t = w0.transpose(1, 2).contiguous()
+        wpair_t = wpair.transpose(1, 2).contiguous()
+        so2_linear._deepmd_cute_neo_manual_weights_t = (
+            w0,
+            wpair,
+            cache_key,
+            w0_t,
+            wpair_t,
+        )
+    else:
+        w0_t, wpair_t = cache[3], cache[4]
+    if inplace_residual:
+        if out is not None:
+            raise ValueError("SO2 backward cannot select both in-place and out storage")
+        from .linear import (
+            neo_so2_linear_backward_residual_inplace,
+        )
+
+        return neo_so2_linear_backward_residual_inplace(
+            residual,
+            grad_out,
+            w0_t,
+            wpair_t,
+        )
+    if grad_out is residual:
+        from .structural_gate import (
+            focus_major_so2_backward_with_folded_residual_out,
+        )
+
+        folded_cache = getattr(
+            so2_linear,
+            "_deepmd_cute_neo_manual_folded_weights_t",
+            None,
+        )
+        folded_key = (
+            w0_t.data_ptr(),
+            wpair_t.data_ptr(),
+            w0_t.dtype,
+            w0_t.device,
+        )
+        if (
+            folded_cache is None
+            or folded_cache[0] is not w0_t
+            or folded_cache[1] is not wpair_t
+            or folded_cache[2] != folded_key
+        ):
+            w0_folded_t = w0_t.clone()
+            wpair_folded_t = wpair_t.clone()
+            w0_folded_t.diagonal(dim1=-2, dim2=-1).add_(1.0)
+            wpair_folded_t.diagonal(dim1=-2, dim2=-1).add_(1.0)
+            folded_cache = (
+                w0_t,
+                wpair_t,
+                folded_key,
+                w0_folded_t,
+                wpair_folded_t,
+            )
+            so2_linear._deepmd_cute_neo_manual_folded_weights_t = folded_cache
+        if out is None:
+            out = grad_out.new_empty(grad_out.shape)
+        return focus_major_so2_backward_with_folded_residual_out(
+            grad_out,
+            folded_cache[3],
+            folded_cache[4],
+            out=out,
+        )
+
+    raise RuntimeError("Neo SO2 backward reached an unsupported storage pattern")
+
+
+def _so3_grid_cross_glu_flat_backward(
+    net: Any,
+    query_flat: Tensor,
+    context_flat: Tensor,
+    grad_out_flat: Tensor,
+) -> tuple[Tensor, Tensor]:
+    """Input-gradient for Neo's flat cross SO3GridNet GLU branch."""
+    if (
+        net.layout != "flat"
+        or net.mode != "cross"
+        or net.op_type != "glu"
+        or net.frame_expand is None
+        or net.frame_contract is None
+    ):
+        raise NotImplementedError("manual grid backward supports Neo cross/flat/glu")
+
+    q_dtype = query_flat.dtype
+    c_dtype = context_flat.dtype
+    n_batch, coeff_dim, _ = query_flat.shape
+    n_focus = net.n_focus
+    channels = net.channels
+    n_frames = net.n_frames
+    expanded = net.expanded_channels
+
+    query = query_flat.reshape(n_batch, coeff_dim, n_focus, channels)
+    context = context_flat.reshape(n_batch, coeff_dim, n_focus, channels)
+    scalar_pair = torch.cat(
+        [query[:, 0, :, :], context[:, 0, :, :]],
+        dim=-1,
+    ).to(dtype=net.dtype)
+
+    left = _frame_expand_forward(net.frame_expand, query).to(dtype=net.dtype)
+    right = _frame_expand_forward(net.frame_expand, context).to(dtype=net.dtype)
+    left_view = left.reshape(n_batch, coeff_dim, n_focus, n_frames, channels)
+    right_view = right.reshape(n_batch, coeff_dim, n_focus, n_frames, channels)
+    to_grid = net.projector.to_grid_mat.reshape(
+        net.projector.grid_size,
+        coeff_dim,
+        n_frames,
+    )
+    from_grid = net.projector.from_grid_mat.reshape(
+        coeff_dim,
+        n_frames,
+        net.projector.grid_size,
+    )
+    left_grid = torch.einsum("gdk,ndfkc->ngfc", to_grid, left_view)
+    right_grid = torch.einsum("gdk,ndfkc->ngfc", to_grid, right_view)
+    coeff = torch.einsum("dkg,ngfc->ndfkc", from_grid, left_grid * right_grid)
+    coeff_flat = coeff.reshape(n_batch, coeff_dim, n_focus, expanded)
+
+    scalar_logits = _focus_linear_forward(net.scalar_gate, scalar_pair)
+    scalar_gate = torch.sigmoid(scalar_logits)
+    coeff_view = coeff_flat.reshape(n_batch, coeff_dim, n_focus, n_frames, channels)
+
+    grad = grad_out_flat.reshape(n_batch, coeff_dim, n_focus, channels).to(
+        dtype=net.dtype
+    )
+    if net.residual_scale is not None:
+        grad = grad * net.residual_scale.reshape(1, 1, n_focus, channels)
+    grad_scalar_flat = _frame_contract_backward_input(net.frame_contract, grad)
+    grad_scalar_view = grad_scalar_flat.reshape(
+        n_batch,
+        coeff_dim,
+        n_focus,
+        n_frames,
+        channels,
+    )
+
+    grad_coeff = grad_scalar_view * scalar_gate[:, None, :, None, :]
+    grad_scalar_gate = (grad_scalar_view * coeff_view).sum(dim=(1, 3))
+    grad_scalar_out = grad_scalar_view[:, 0, :, net.frame_zero_index, :]
+    grad_scalar_logits = grad_scalar_gate * scalar_gate * (1.0 - scalar_gate)
+    grad_scalar_pair = _focus_linear_backward_input(net.scalar_gate, grad_scalar_logits)
+    grad_scalar_pair = grad_scalar_pair + _swiglu_backward_input(
+        scalar_pair,
+        grad_scalar_out,
+    )
+
+    grad_grid = torch.einsum("dkg,ndfkc->ngfc", from_grid, grad_coeff)
+    grad_left_grid = grad_grid * right_grid
+    grad_right_grid = grad_grid * left_grid
+    grad_left = torch.einsum("gdk,ngfc->ndfkc", to_grid, grad_left_grid).reshape(
+        n_batch,
+        coeff_dim,
+        n_focus,
+        expanded,
+    )
+    grad_right = torch.einsum("gdk,ngfc->ndfkc", to_grid, grad_right_grid).reshape(
+        n_batch,
+        coeff_dim,
+        n_focus,
+        expanded,
+    )
+    grad_query = _frame_expand_backward_input(net.frame_expand, grad_left)
+    grad_context = _frame_expand_backward_input(net.frame_expand, grad_right)
+    grad_query[:, 0, :, :].add_(grad_scalar_pair[:, :, :channels])
+    grad_context[:, 0, :, :].add_(grad_scalar_pair[:, :, channels:])
+    return (
+        grad_query.reshape_as(query_flat).to(dtype=q_dtype),
+        grad_context.reshape_as(context_flat).to(dtype=c_dtype),
+    )
+
+
+def _final_manual_backward(runner: Any, grad_out: Tensor) -> tuple[Tensor, Tensor]:
+    so2 = runner.so2
+    block = runner.block
+    n_node = runner.node_count
+    if runner.use_full_node:
+        grad_so2_out = grad_out
+    else:
+        grad_so2_out = grad_out[:, : block.mp_ebed_dim, :, :]
+
+    phase = runner.phase_c_out.detach()
+    x_wide = runner.x_wide.detach()
+    out_gate_flat = runner.out_gate_flat
+    post_norm_in = runner.post_norm_input
+
+    grad_post_norm_in = _equivariant_rmsnorm_backward(
+        block.post_so2_norm,
+        post_norm_in,
+        grad_so2_out,
+    )
+    grad_post_mix = _so3_linear_backward_input(
+        so2.post_focus_mix,
+        grad_post_norm_in.squeeze(2).unsqueeze(2),
+    ).squeeze(2)
+
+    if so2.message_node_grid_product is not None:
+        if runner.packed_message_grid:
+            message_grid_product = runner.message_grid_product
+            runner.message_grid_product = None
+            from .message_grid import (
+                run_packed_message_grid_backward,
+            )
+
+            grad_out_gate_flat, grad_grid_context = run_packed_message_grid_backward(
+                so2.message_node_grid_product,
+                out_gate_flat,
+                x_wide,
+                grad_post_mix,
+                product_flat=message_grid_product,
+            )
+            del message_grid_product
+        else:
+            grad_out_gate_flat, grad_grid_context = _so3_grid_cross_glu_flat_backward(
+                so2.message_node_grid_product,
+                out_gate_flat,
+                x_wide,
+                grad_post_mix,
+            )
+        grad_out_gate_flat.add_(grad_post_mix)
+        grad_x_wide_down = torch.zeros(
+            n_node,
+            16 * 64,
+            device=x_wide.device,
+            dtype=x_wide.dtype,
+        ).view(n_node, 16, 64)
+        grad_x_wide_down.add_(grad_grid_context)
+    else:
+        grad_out_gate_flat = grad_post_mix
+        grad_x_wide_down = torch.zeros(
+            n_node,
+            16 * 64,
+            device=x_wide.device,
+            dtype=x_wide.dtype,
+        ).view(n_node, 16, 64)
+
+    grad_phase = grad_out_gate_flat.contiguous()
+    output_gate_backward = _compile_output_gate_backward(
+        _runner_compile_identity(runner),
+        float(so2.attn_output_gate_norm.eps),
+    )
+    output_gate_backward(
+        grad_phase.view(n_node, 16 * 64),
+        phase.contiguous().view(n_node, 16 * 64),
+        x_wide.contiguous().view(n_node, 16 * 64),
+        so2.attn_output_gate_norm.adam_scale.detach()
+        .float()
+        .reshape(2, 32)
+        .contiguous(),
+        so2.adamw_attn_gate_w.detach().float().reshape(32, 2, 1).contiguous(),
+        grad_phase.view(n_node, 16 * 64),
+        grad_x_wide_down.view(n_node, 16 * 64),
+    )
+    return grad_phase.reshape_as(phase), grad_x_wide_down
+
+
+def _qk_manual_backward(
+    runner: Any,
+    grad_logits: Tensor,
+) -> Tensor:
+    so2 = runner.so2
+    n_node = runner.node_count
+    x_wide = runner.x_wide.detach()
+    x_l0 = x_wide[:, 0, :].reshape(n_node, 2, 32)
+    q_node = runner.q_node
+    k_node = runner.k_node
+    grad_q_node = getattr(runner, "grad_q_node", None)
+    grad_k_node = getattr(runner, "grad_k_node", None)
+    if grad_q_node is None:
+        grad_q_node = torch.empty_like(
+            q_node,
+            memory_format=torch.contiguous_format,
+        )
+        grad_k_node = torch.empty_like(
+            k_node,
+            memory_format=torch.contiguous_format,
+        )
+        runner.grad_q_node = grad_q_node
+        runner.grad_k_node = grad_k_node
+    if not grad_q_node.is_contiguous() or not grad_k_node.is_contiguous():
+        raise RuntimeError("Neo Q/K backward buffers must be compact N x 2 x 32")
+    grad_q_node.zero_()
+    grad_k_node.zero_()
+    runner.qk_edge_backward(
+        grad_logits.contiguous(),
+        q_node,
+        k_node,
+        runner.src_i32,
+        runner.dst_i32,
+        grad_q_node,
+        grad_k_node,
+    )
+    grad_x_wide = getattr(runner, "grad_x_wide_qk", None)
+    if grad_x_wide is None or grad_x_wide.shape != x_wide.shape:
+        grad_x_wide = torch.empty(
+            x_wide.shape,
+            device=x_wide.device,
+            dtype=x_wide.dtype,
+        )
+        runner.grad_x_wide_qk = grad_x_wide
+    runner.qk_node_input_adjoint(
+        x_l0.contiguous(),
+        grad_q_node,
+        grad_k_node,
+        so2.attn_q_proj.weight.detach().float().view(32, 2, 32).contiguous(),
+        so2.attn_k_proj.weight.detach().float().view(32, 2, 32).contiguous(),
+        so2.attn_qk_norm.adam_scale.detach().float().contiguous(),
+        grad_x_wide.view(n_node, 16 * 64),
+    )
+    return grad_x_wide
+
+
+def _native_edge_major_stack_grad(grad_stack_out: Tensor) -> Tensor:
+    """Validate the exact in-place Phase-C adjoint consumed by final SO2."""
+    edge_count = grad_stack_out.shape[0]
+    expected_shape = (edge_count, 2, 10, 32)
+    expected_stride = (2 * 10 * 32, 10 * 32, 32, 1)
+    if (
+        grad_stack_out.shape != expected_shape
+        or grad_stack_out.stride() != expected_stride
+        or grad_stack_out.dtype != torch.float32
+    ):
+        raise RuntimeError(
+            "in-place Phase-C adjoint requires compact edge-major storage: "
+            f"got shape={tuple(grad_stack_out.shape)} "
+            f"stride={grad_stack_out.stride()} dtype={grad_stack_out.dtype}, "
+            f"expected shape={expected_shape} stride={expected_stride}"
+        )
+    return grad_stack_out
+
+
+def _phase_c_layout_grad_stack(runner: Any) -> Tensor:
+    """Select aliased edge-major or ordinary Phase-C output storage."""
+    if runner.phase_c_y is not None:
+        raise RuntimeError(
+            "in-place Phase-C adjoint requires the folded single-input stack"
+        )
+    grad_mixed = runner.grad_mixed_slab
+    if grad_mixed is None or (
+        grad_mixed.untyped_storage()._cdata
+        == runner.phase_c_stack.untyped_storage()._cdata
+    ):
+        raise RuntimeError(
+            "in-place Phase-C adjoint must retain a distinct grad_mixed_slab"
+        )
+    # Phase C owns every destination edge and stores G at the same
+    # edge/focus address only after its complete 10x32 input fragment is
+    # register-resident. Final SO2 consumes G before gate scratch reuses y2.
+    return _native_edge_major_stack_grad(runner.phase_c_stack)
+
+
+def _stack_backward_manual(runner: Any, grad_stack_out: Tensor) -> Tensor:
+    cur_grad = _native_edge_major_stack_grad(grad_stack_out)
+    for cache_index in range(len(runner.stack_caches) - 1, -1, -1):
+        cache = runner.stack_caches[cache_index]
+        residual_grad = cur_grad
+        if cache.final:
+            grad_y = cur_grad
+        else:
+            if runner.config.combined_so2_gate:
+                assert runner.combined_gate_backward is not None
+                runner.combined_gate_backward(
+                    cur_grad.view(-1, 10 * 32),
+                    cache.y.detach().view(-1, 10 * 32),
+                    cache.non_linear.gate_linear.weight.detach()
+                    .view(32, 2, 3 * 32)
+                    .contiguous(),
+                    runner.grad_y,
+                )
+                grad_y = runner.grad_y.view_as(cache.y)
+            else:
+                grad_gate_logits = runner._run_structural_gate_backward(
+                    runner.structural_gate_backward,
+                    cur_grad,
+                    cache.y,
+                    cache.logits,
+                    runner.grad_y.view_as(cache.y),
+                    grad_logits=runner.grad_gate_logits,
+                    overwrite_logits=True,
+                )
+                grad_y = runner.grad_y.view_as(cache.y)
+                runner._focus_major_gate_linear_backward_add(
+                    grad_y,
+                    grad_gate_logits,
+                    cache.non_linear.gate_linear.weight.detach(),
+                )
+        linear = runner.so2.so2_linears[cache_index]
+        so2_out = runner.grad_mixed_slab if cache.final else None
+        cur_grad = _neo_so2_linear_backward_input_with_residual(
+            linear,
+            grad_y,
+            residual_grad,
+            inplace_residual=not cache.final,
+            out=so2_out,
+        )
+    return cur_grad
+
+
+def _x_wide_manual_backward(runner: Any, grad_x_wide_total: Tensor) -> Tensor:
+    block = runner.block
+    so2 = runner.so2
+    x_so2 = runner.x if runner.use_full_node else runner.x[:, : block.mp_ebed_dim, :, :]
+    if not _is_identity(block.pre_so2_norm):
+        raise NotImplementedError(
+            "manual x-wide backward currently expects Identity pre norm"
+        )
+    grad_x_pre_flat = _so3_linear_backward_input(
+        so2.pre_focus_mix,
+        grad_x_wide_total.unsqueeze(2),
+    )
+    grad_x_pre = grad_x_pre_flat.squeeze(2).reshape_as(x_so2)
+    if runner.use_full_node:
+        grad_x = grad_x_pre
+    else:
+        grad_x = torch.zeros_like(runner.x)
+        grad_x[:, : block.mp_ebed_dim, :, :] = grad_x_pre
+    expected_stride = (
+        grad_x.shape[-1],
+        grad_x.shape[0] * grad_x.shape[-1],
+        grad_x.shape[-1],
+        1,
+    )
+    if grad_x.stride() != expected_stride:
+        grad_x = grad_x.permute(1, 0, 2, 3).contiguous().permute(1, 0, 2, 3)
+    return grad_x
+
+
+def _make_edge_cache(
+    *,
+    src: Tensor,
+    dst: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    edge_env: Tensor,
+    edge_src_gate: Tensor,
+) -> Any:
+    return SimpleNamespace(
+        src=src,
+        dst=dst,
+        D_full=d_full,
+        Dt_full=dt_full,
+        edge_env=edge_env,
+        edge_src_gate=_edge_src_gate_arg(edge_src_gate),
+        D_to_m_cache={},
+        Dt_from_m_cache={},
+    )
+
+
+def _destination_degrees_fit_limit(
+    destination_row_ptr: Tensor,
+    max_edges_per_node: int,
+) -> bool:
+    """Return whether every destination row fits a bounded native kernel."""
+    destination_degrees = destination_row_ptr[1:] - destination_row_ptr[:-1]
+    return bool(
+        torch.all(
+            (destination_degrees >= 0) & (destination_degrees <= max_edges_per_node)
+        ).item()
+    )
+
+
+def _build_runner(
+    handle: int,
+    x: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    radial_feat: Tensor,
+    edge_env: Tensor,
+    src: Tensor,
+    dst: Tensor,
+    dst_ptr: Tensor,
+    source_order: Tensor,
+    source_ptr: Tensor,
+    edge_src_gate: Tensor,
+) -> Any:
+    """Build a runner after enforcing CuTe's runtime pointer contract.
+
+    This function executes below the custom-op boundary, including for the
+    compile-visible thin path.  Exact ``data_ptr`` checks are therefore safe
+    here and repair contiguous offset views without introducing a Dynamo graph
+    break above the op.
+    """
+    with _REGISTRY_LOCK:
+        entry = _REGISTRY[int(handle)]
+    # Compiled graphs keep the handle, not guards on the registered parameters.
+    # Enforce the contract here too if state changes after graph capture.
+    if not _has_frozen_fp32_contract(entry.block):
+        raise RuntimeError(
+            "CuTe SO2 requires frozen FP32 parameters and strict FP32 matmul; "
+            "model or precision state changed after preparation"
+        )
+    config = entry.config
+    if config.native_sm90_path:
+        from .sm90.phase_c_attention_backward import (
+            MAX_EDGES_PER_NODE,
+        )
+
+        native_degree_eligible = _destination_degrees_fit_limit(
+            dst_ptr,
+            MAX_EDGES_PER_NODE,
+        )
+        if native_degree_eligible:
+            from .sm90.runner import NeoSm90SO2Runner as Runner
+        else:
+            from .runner import NeoFullCuteBackward as Runner
+
+            config = replace(config, native_sm90_path=False)
+    else:
+        from .runner import NeoFullCuteBackward as Runner
+
+    x = _aligned_contiguous(x)
+    shared_wigner_storage = d_full.data_ptr() == dt_full.data_ptr()
+    d_full = _aligned_contiguous(d_full)
+    dt_full = d_full if shared_wigner_storage else _aligned_contiguous(dt_full)
+    radial_feat = _aligned_contiguous(radial_feat)
+    edge_env = _aligned_contiguous(edge_env)
+    src = _aligned_contiguous(src)
+    dst = _aligned_contiguous(dst)
+    dst_ptr = _aligned_contiguous(dst_ptr)
+    source_order = _aligned_contiguous(source_order)
+    source_ptr = _aligned_contiguous(source_ptr)
+    edge_src_gate = _aligned_contiguous(edge_src_gate)
+    edge_cache = _make_edge_cache(
+        src=src,
+        dst=dst,
+        d_full=d_full,
+        dt_full=dt_full,
+        edge_env=edge_env,
+        edge_src_gate=edge_src_gate,
+    )
+    record = SimpleNamespace(edge_cache=edge_cache)
+    with torch.cuda.device(x.device), torch.no_grad():
+        runner = Runner(
+            torch,
+            entry.block,
+            record,
+            x,
+            d_full,
+            dt_full,
+            radial_feat,
+            dst_ptr,
+            source_order,
+            source_ptr,
+            runtime_config=config,
+        )
+    return runner
+
+
+def _so2_packed_direct_forward_impl(
+    handle: int,
+    x: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    radial_feat: Tensor,
+    edge_env: Tensor,
+    src: Tensor,
+    dst: Tensor,
+    dst_ptr: Tensor,
+    source_order: Tensor,
+    source_ptr: Tensor,
+    edge_src_gate: Tensor,
+) -> tuple[Tensor, Tensor]:
+    _assert_not_cuda_graph_capturing("packed forward", x)
+    runner = _build_runner(
+        handle,
+        x,
+        d_full,
+        dt_full,
+        radial_feat,
+        edge_env,
+        src,
+        dst,
+        dst_ptr,
+        source_order,
+        source_ptr,
+        edge_src_gate,
+    )
+    runner_token = torch.empty(
+        (1,),
+        device=x.device,
+        dtype=torch.uint8,
+    )
+    _store_packed_runner(runner_token, runner)
+    return _layout_like(runner.final.detach(), x), runner_token
+
+
+def _so2_backward_from_runner_current_device(
+    handle: int,
+    grad_out: Tensor,
+    x: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    radial_feat: Tensor,
+    edge_env: Tensor,
+    edge_src_gate: Tensor,
+    runner: Any,
+) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+    del handle
+    grad_x, grad_d, grad_dt, grad_radial = _runner_backward_manual(
+        runner,
+        grad_out,
+    )
+
+    grad_edge_env = runner.grad_edge.view_as(edge_env)
+    if edge_src_gate.numel() != 0:
+        factor = edge_src_gate.reshape_as(edge_env).float().clamp_min(0.0).sqrt()
+        grad_edge_env = grad_edge_env * factor.to(dtype=grad_edge_env.dtype)
+    grad_edge_env = (
+        grad_edge_env.clone() if grad_edge_env._base is not None else grad_edge_env
+    )
+    grad_edge_env.masked_fill_(edge_env <= 0, 0)
+    grad_x_out = _grad_layout_like(grad_x, x, skip=True)
+    grad_d_out = _grad_layout_like(grad_d, d_full, skip=True)
+    grad_dt_out = _grad_layout_like(grad_dt, dt_full, skip=True)
+    grad_radial_out = _grad_layout_like(grad_radial, radial_feat, skip=True)
+    grad_edge_env_out = _grad_layout_like(grad_edge_env, edge_env, skip=True)
+    if x.ndim == 4:
+        n_node = x.shape[0]
+        channels = x.shape[-1]
+        grad_x_stride = (channels, n_node * channels, channels, 1)
+    else:
+        grad_x_stride = x.stride()
+    for name, actual, expected_like, expected_stride in (
+        ("x", grad_x_out, x, grad_x_stride),
+        ("D", grad_d_out, d_full, d_full.stride()),
+        ("Dt", grad_dt_out, dt_full, dt_full.stride()),
+        ("radial", grad_radial_out, radial_feat, radial_feat.stride()),
+        ("edge_env", grad_edge_env_out, edge_env, edge_env.stride()),
+    ):
+        _assert_grad_meta_contract(
+            actual,
+            expected_like,
+            name=name,
+            expected_stride=expected_stride,
+        )
+    return (
+        grad_x_out,
+        grad_d_out,
+        grad_dt_out,
+        grad_radial_out,
+        grad_edge_env_out,
+    )
+
+
+def _so2_backward_from_runner(
+    handle: int,
+    grad_out: Tensor,
+    x: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    radial_feat: Tensor,
+    edge_env: Tensor,
+    edge_src_gate: Tensor,
+    runner: Any,
+) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+    """Run all compilation and launches on the operand's CUDA device."""
+    with torch.cuda.device(grad_out.device):
+        return _so2_backward_from_runner_current_device(
+            handle,
+            grad_out,
+            x,
+            d_full,
+            dt_full,
+            radial_feat,
+            edge_env,
+            edge_src_gate,
+            runner,
+        )
+
+
+def _so2_packed_direct_backward_impl(
+    handle: int,
+    grad_out: Tensor,
+    x: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    radial_feat: Tensor,
+    edge_env: Tensor,
+    src: Tensor,
+    dst: Tensor,
+    dst_ptr: Tensor,
+    source_order: Tensor,
+    source_ptr: Tensor,
+    edge_src_gate: Tensor,
+    runner_token: Tensor,
+) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+    _assert_not_cuda_graph_capturing("packed backward", grad_out)
+    runner = _borrow_packed_runner(runner_token)
+    if runner is None:
+        runner = _build_runner(
+            handle,
+            x,
+            d_full,
+            dt_full,
+            radial_feat,
+            edge_env,
+            src,
+            dst,
+            dst_ptr,
+            source_order,
+            source_ptr,
+            edge_src_gate,
+        )
+    return _so2_backward_from_runner(
+        handle,
+        grad_out,
+        x,
+        d_full,
+        dt_full,
+        radial_feat,
+        edge_env,
+        edge_src_gate,
+        runner,
+    )
+
+
+def _runner_backward_manual(
+    runner: Any,
+    grad_out: Tensor,
+) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+    if getattr(runner, "uses_native_sm90_path", False):
+        return runner.input_adjoint(grad_out)
+
+    so2 = runner.so2
+    n_edge = runner.edge_count
+
+    grad_phase_c_out, grad_x_wide_down = _final_manual_backward(runner, grad_out)
+    # Keep the large SO2 slabs out of the readout-backward allocation crest.
+    runner.ensure_backward_workspace()
+
+    from .phase_c import (
+        NeoPhaseCBackwardLayoutOutputs,
+    )
+
+    layout_outputs = getattr(runner, "phase_c_layout_outputs", None)
+    if layout_outputs is None:
+        layout_outputs = NeoPhaseCBackwardLayoutOutputs(
+            grad_stack=_phase_c_layout_grad_stack(runner),
+            grad_wigner_dt=runner.grad_dt,
+            grad_logits=runner.grad_logits,
+            grad_edge=runner.grad_edge,
+            grad_z_partial=runner.grad_z_partial,
+            grad_z=runner.grad_z,
+            grad_focus_src=torch.empty(
+                2,
+                runner.edge_count,
+                32,
+                device=runner.x.device,
+                dtype=torch.float32,
+            ),
+        )
+        runner.phase_c_layout_outputs = layout_outputs
+    runner.phase_c_layout_backward(
+        grad_phase_c_out.contiguous(),
+        runner.phase_c_stack.detach().contiguous(),
+        runner.dt.detach().contiguous(),
+        runner.alpha,
+        runner.focus_alpha,
+        runner.dst_ptr_i32,
+        runner.rotate,
+        runner.edge_gate,
+        so2.adamw_attn_z_bias_raw.detach().reshape(2).float().contiguous(),
+        runner.group_max,
+        runner.denom,
+        runner.focus_gate_src.detach().contiguous(),
+        so2.adamw_focus_compete_w.detach().float().contiguous(),
+        runner.focus_norm_scale,
+        layout_outputs,
+    )
+    grad_stack_out = _native_edge_major_stack_grad(layout_outputs.grad_stack)
+    grad_focus_src_focus = layout_outputs.grad_focus_src
+
+    grad_x_wide_qk = _qk_manual_backward(
+        runner,
+        runner.grad_logits.view(n_edge, 2),
+    )
+    grad_mixed = _stack_backward_manual(runner, grad_stack_out)
+
+    grad_mixed_focus = grad_mixed
+    if not grad_mixed_focus.is_contiguous():
+        grad_mixed_focus = grad_mixed_focus.contiguous()
+    x_wide_flat = runner.x_wide.detach().contiguous().view(runner.node_count, 16 * 64)
+    radial_state = runner.radial_compact
+    from .radial_phase_a import (
+        run_neo_radial_phase_a_backward_node_tiled,
+    )
+
+    run_neo_radial_phase_a_backward_node_tiled(
+        grad_mixed_focus.view(runner.edge_count, 2 * 10 * 32),
+        runner.grad_logits,
+        radial_state,
+        runner.so2.radial_degree_mixer.channel_basis.detach().view(64).contiguous(),
+        x_wide_flat,
+        runner.source_order_i32,
+        runner.source_ptr_i32,
+        runner.d.detach(),
+        grad_focus_src_focus=grad_focus_src_focus,
+        batched_radial_projection_weight=runner.batched_radial_projection_weight,
+        grad_x_wide=runner.grad_x_wide_phase_a,
+        grad_d_full=runner.grad_d,
+        grad_radial_m0=runner.grad_radial_flat,
+        validate_csr=runtime_policy.is_cute_strict_enabled(),
+    )
+    grad_radial = runner.grad_radial_flat.view_as(runner.radial)
+    grad_x_wide_phase_a = runner.grad_x_wide_phase_a.view_as(runner.x_wide)
+    grad_d = runner.grad_d
+
+    grad_x_wide_total = grad_x_wide_phase_a
+    grad_x_wide_total.add_(grad_x_wide_qk)
+    grad_x_wide_total.add_(grad_x_wide_down)
+    grad_x = _x_wide_manual_backward(runner, grad_x_wide_total)
+    return grad_x, grad_d, runner.grad_dt, grad_radial
+
+
+def _stateful_custom_op_tags() -> tuple[Any, ...] | None:
+    """Tag hidden-state runner ops as unsafe for direct CUDA graph capture."""
+    tag_type = getattr(getattr(torch, "_C", None), "Tag", None)
+    cudagraph_unsafe = getattr(tag_type, "cudagraph_unsafe", None)
+    if cudagraph_unsafe is None:
+        return None
+    return (cudagraph_unsafe,)
+
+
+_SO2_CUSTOM_OP_TAGS = _stateful_custom_op_tags()
+_so2_packed_direct_op = torch.library.custom_op(
+    "sezm_cute::so2_packed_direct", mutates_args=(), tags=_SO2_CUSTOM_OP_TAGS
+)(_so2_packed_direct_forward_impl)
+_so2_packed_direct_bwd_op = torch.library.custom_op(
+    "sezm_cute::so2_packed_direct_bwd", mutates_args=(), tags=_SO2_CUSTOM_OP_TAGS
+)(_so2_packed_direct_backward_impl)
+
+
+@_so2_packed_direct_op.register_fake
+def _(
+    handle: int,
+    x: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    radial_feat: Tensor,
+    edge_env: Tensor,
+    src: Tensor,
+    dst: Tensor,
+    dst_ptr: Tensor,
+    source_order: Tensor,
+    source_ptr: Tensor,
+    edge_src_gate: Tensor,
+) -> tuple[Tensor, Tensor]:
+    del handle
+    del d_full, dt_full, radial_feat, edge_env, src, dst, dst_ptr
+    del source_order, source_ptr, edge_src_gate
+    return (
+        torch.empty_like(x),
+        torch.empty((1,), device=x.device, dtype=torch.uint8),
+    )
+
+
+@_so2_packed_direct_bwd_op.register_fake
+def _(
+    handle: int,
+    grad_out: Tensor,
+    x: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    radial_feat: Tensor,
+    edge_env: Tensor,
+    src: Tensor,
+    dst: Tensor,
+    dst_ptr: Tensor,
+    source_order: Tensor,
+    source_ptr: Tensor,
+    edge_src_gate: Tensor,
+    runner_token: Tensor,
+) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+    del handle, grad_out, src, dst, dst_ptr, source_order, source_ptr
+    del edge_src_gate, runner_token
+    return (
+        _fake_x_wide_grad_like(x, skip=True),
+        torch.empty_like(d_full),
+        torch.empty_like(dt_full),
+        torch.empty_like(radial_feat),
+        torch.empty_like(edge_env),
+    )
+
+
+def _so2_packed_direct_setup_context(
+    ctx: Any,
+    inputs: tuple[Any, ...],
+    output: tuple[Tensor, Tensor],
+) -> None:
+    _, runner_token = output
+    (
+        handle,
+        x,
+        d_full,
+        dt_full,
+        radial_feat,
+        edge_env,
+        src,
+        dst,
+        dst_ptr,
+        source_order,
+        source_ptr,
+        edge_src_gate,
+    ) = inputs
+    ctx.handle = int(handle)
+    ctx.save_for_backward(
+        x,
+        d_full,
+        dt_full,
+        radial_feat,
+        edge_env,
+        src,
+        dst,
+        dst_ptr,
+        source_order,
+        source_ptr,
+        edge_src_gate,
+        runner_token,
+    )
+
+
+def _so2_packed_direct_registered_backward_impl(
+    ctx: Any,
+    grad_out: Tensor,
+) -> tuple[Any, ...]:
+    (
+        x,
+        d_full,
+        dt_full,
+        radial_feat,
+        edge_env,
+        src,
+        dst,
+        dst_ptr,
+        source_order,
+        source_ptr,
+        edge_src_gate,
+        runner_token,
+    ) = ctx.saved_tensors
+    grad_x, grad_d, grad_dt, grad_radial, grad_edge_env = _so2_packed_direct_bwd_op(
+        ctx.handle,
+        grad_out,
+        x,
+        d_full,
+        dt_full,
+        radial_feat,
+        edge_env,
+        src,
+        dst,
+        dst_ptr,
+        source_order,
+        source_ptr,
+        edge_src_gate,
+        runner_token,
+    )
+    return (
+        None,
+        grad_x,
+        grad_d,
+        grad_dt,
+        grad_radial,
+        grad_edge_env,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+    )
+
+
+def _so2_packed_direct_backward(
+    ctx: Any,
+    grad_out: Tensor,
+    grad_runner_token: Tensor | None,
+) -> tuple[Any, ...]:
+    """Run packed-direct backward with its custom op visible to compilation."""
+    del grad_runner_token
+    return _so2_packed_direct_registered_backward_impl(ctx, grad_out)
+
+
+_so2_packed_direct_op.register_autograd(
+    _so2_packed_direct_backward,
+    setup_context=_so2_packed_direct_setup_context,
+)
+
+
+def _cute_so2_impl(
+    handle: int,
+    x: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    radial_feat: Tensor,
+    edge_env: Tensor,
+    src: Tensor,
+    dst: Tensor,
+    dst_ptr: Tensor,
+    source_order: Tensor,
+    source_ptr: Tensor,
+    edge_src_gate: Tensor,
+) -> Tensor:
+    output, _runner_token = _so2_packed_direct_op(
+        int(handle),
+        x,
+        d_full,
+        dt_full,
+        radial_feat,
+        edge_env,
+        src,
+        dst,
+        dst_ptr,
+        source_order,
+        source_ptr,
+        edge_src_gate,
+    )
+    return output
+
+
+def cute_so2(
+    handle: int,
+    x: Tensor,
+    d_full: Tensor,
+    dt_full: Tensor,
+    radial_feat: Tensor,
+    edge_env: Tensor,
+    src: Tensor,
+    dst: Tensor,
+    dst_ptr: Tensor,
+    edge_src_gate: Tensor,
+    source_order: Tensor | None = None,
+    source_ptr: Tensor | None = None,
+) -> Tensor:
+    """Run SO2 through its registered custom-op boundary."""
+    if source_order is None:
+        source_order = src.new_empty((0,), dtype=torch.int32)
+    if source_ptr is None:
+        source_ptr = src.new_empty((0,), dtype=torch.int32)
+    return _cute_so2_impl(
+        handle,
+        x,
+        d_full,
+        dt_full,
+        radial_feat,
+        edge_env,
+        src,
+        dst,
+        dst_ptr,
+        source_order,
+        source_ptr,
+        edge_src_gate,
+    )
+
+
+def _dst_ptr_from_sorted(
+    torch_module: Any,
+    dst: Tensor,
+    n_node: int,
+    *,
+    destinations_sorted: bool,
+) -> Tensor | None:
+    if not destinations_sorted:
+        return None
+    if runtime_policy.is_cute_strict_enabled() and dst.numel() > 1:
+        torch_module._assert_async(
+            torch_module.all(dst[1:] >= dst[:-1]),
+            "Neo SO2 destinations_sorted=True requires monotonically "
+            "nondecreasing destination indices",
+        )
+    boundaries = torch_module.arange(
+        n_node + 1,
+        device=dst.device,
+        dtype=torch_module.int64,
+    )
+    return torch_module.searchsorted(dst.contiguous(), boundaries)
+
+
+def _validated_sorted_edge_metadata_args(
+    edge_cache: Any,
+    *,
+    node_count: int,
+    dst_ptr: Tensor | None,
+    source_order: Tensor | None,
+    source_ptr: Tensor | None,
+) -> tuple[Tensor, Tensor, Tensor] | None:
+    """Return validated invocation-local CSR tensors for this edge cache."""
+    if not getattr(edge_cache, "destinations_sorted", False):
+        return None
+    if dst_ptr is None or source_order is None or source_ptr is None:
+        return None
+    device = edge_cache.src.device
+    if (
+        dst_ptr.device != device
+        or source_order.device != device
+        or source_ptr.device != device
+        or dst_ptr.dtype != torch.int32
+        or source_order.dtype != torch.int32
+        or source_ptr.dtype != torch.int32
+        or dst_ptr.numel() != node_count + 1
+        or source_order.numel() != edge_cache.src.numel()
+        or source_ptr.numel() != node_count + 1
+    ):
+        return None
+    return (
+        dst_ptr.contiguous(),
+        source_order.contiguous(),
+        source_ptr.contiguous(),
+    )
+
+
+def _cuda_compute_capability(device_index: int) -> tuple[int, int]:
+    return tuple(torch.cuda.get_device_capability(device_index))
+
+
+def _architecture_default_config(
+    compute_capability: tuple[int, int],
+) -> NeoSO2RuntimeConfig:
+    return NeoSO2RuntimeConfig(
+        native_sm90_path=compute_capability == runtime_policy.SM90_CAPABILITY,
+        per_focus_so2_fwd_pair=(
+            compute_capability in runtime_policy.SM80_PROFILE_CAPABILITIES
+        ),
+        combined_so2_gate=(
+            compute_capability in runtime_policy.FUSED_SO2_GATE_CAPABILITIES
+        ),
+    )
+
+
+def _maybe_run_prepared_cute_so2(
+    block: Any,
+    x: Tensor,
+    edge_cache: Any,
+    radial_feat: Tensor,
+    dst_ptr: Tensor | None = None,
+    source_order: Tensor | None = None,
+    source_ptr: Tensor | None = None,
+) -> Tensor | None:
+    """Dispatch prevalidated packed SO2 state without a Python graph break."""
+    state = getattr(block, "_deepmd_cute_so2_state", None)
+    if not isinstance(state, _RegisteredSO2State):
+        return None
+    if not runtime_policy.is_cute_infer_enabled() or not _has_frozen_fp32_contract(
+        block
+    ):
+        return None
+    with _REGISTRY_LOCK:
+        entry = _REGISTRY.get(state.handle)
+        if entry is None or entry.block is not block:
+            return None
+    d_full = edge_cache.D_packed
+    dt_full = d_full
+    destinations_sorted = bool(getattr(edge_cache, "destinations_sorted", False))
+    device_index = x.device.index
+    if (
+        block.training
+        or x.device.type != "cuda"
+        or device_index is None
+        or device_index != state.device_index
+        or x.dtype != torch.float32
+        or torch.is_autocast_enabled(x.device.type)
+        or d_full is None
+        or dt_full is None
+        or d_full is not dt_full
+        or d_full.dim() != 2
+        or dt_full.dim() != 2
+        or edge_cache.edge_src_gate is not None
+        or not destinations_sorted
+        or not _dtypes_use_strict_fp32(
+            (
+                d_full.dtype,
+                dt_full.dtype,
+                radial_feat.dtype,
+                edge_cache.edge_env.dtype,
+            )
+        )
+    ):
+        return None
+    edge_count = edge_cache.src.numel()
+    if not runtime_policy.so2_int32_indexing_is_safe(
+        edge_count=edge_count,
+        node_count=x.shape[0],
+    ):
+        return None
+    metadata_args = _validated_sorted_edge_metadata_args(
+        edge_cache,
+        node_count=x.shape[0],
+        dst_ptr=dst_ptr,
+        source_order=source_order,
+        source_ptr=source_ptr,
+    )
+    if metadata_args is None:
+        dst_ptr = _dst_ptr_from_sorted(
+            torch,
+            edge_cache.dst,
+            x.shape[0],
+            destinations_sorted=destinations_sorted,
+        )
+        if dst_ptr is None:
+            return None
+        source_order = edge_cache.src.new_empty((0,), dtype=torch.int32)
+        source_ptr = edge_cache.src.new_empty((0,), dtype=torch.int32)
+    else:
+        dst_ptr, source_order, source_ptr = metadata_args
+    # The opaque custom-op implementation performs exact pointer-alignment
+    # canonicalization in ``_build_runner``. Keep this wrapper graph-visible.
+    empty_edge_src_gate = edge_cache.edge_env.new_empty((0,))
+    d_arg = d_full.contiguous()
+    output = cute_so2(
+        state.handle,
+        x.contiguous(),
+        d_arg,
+        d_arg,
+        radial_feat.contiguous(),
+        edge_cache.edge_env.contiguous(),
+        edge_cache.src.contiguous(),
+        edge_cache.dst.contiguous(),
+        dst_ptr.contiguous(),
+        empty_edge_src_gate,
+        source_order=source_order.contiguous(),
+        source_ptr=source_ptr.contiguous(),
+    )
+    return _layout_like(output, x)
+
+
+@torch.compiler.disable
+def _maybe_run_cute_so2_fallback(
+    block: Any,
+    x: Tensor,
+    edge_cache: Any,
+    radial_feat: Tensor,
+    dst_ptr: Tensor | None = None,
+    source_order: Tensor | None = None,
+    source_ptr: Tensor | None = None,
+) -> Tensor | None:
+    """Run the opt-in Neo CuTe SO2 path, or return ``None`` for fallback."""
+    if edge_cache.D_packed is None:
+        return None
+    # The optimized backward does not expose the differentiable SFPG
+    # source-gate adjoint. Preserve force/stress correctness via eager fallback.
+    if edge_cache.edge_src_gate is not None:
+        return None
+    destinations_sorted = bool(getattr(edge_cache, "destinations_sorted", False))
+    if not is_neo_so2_runtime_eligible(
+        block,
+        training=bool(block.training),
+        device=x.device,
+        dtype=x.dtype,
+        edge_count=edge_cache.src.numel(),
+        node_count=x.shape[0],
+        destinations_sorted=destinations_sorted,
+    ):
+        return None
+    if not _dtypes_use_strict_fp32(
+        (
+            edge_cache.D_packed.dtype,
+            radial_feat.dtype,
+            edge_cache.edge_env.dtype,
+        )
+    ):
+        return None
+
+    state = getattr(block, "_deepmd_cute_so2_state", None)
+    if state is False:
+        return None
+    if state is not None and not isinstance(state, _RegisteredSO2State):
+        state = None
+    device_index = x.device.index
+    if device_index is None:
+        device_index = torch.cuda.current_device()
+    compute_capability = _cuda_compute_capability(device_index)
+    if not is_supported_so2_compute_capability(compute_capability):
+        return None
+    config = _architecture_default_config(compute_capability)
+    if not runtime_policy.so2_int32_indexing_is_safe(
+        edge_count=edge_cache.src.numel(),
+        node_count=x.shape[0],
+    ):
+        return None
+    if edge_cache.D_packed.dim() != 2:
+        return None
+    if state is None or state.device_index != device_index or state.config != config:
+        state = _register_cute_so2_state(block, device_index, config)
+        if state is None:
+            return None
+
+    metadata_args = _validated_sorted_edge_metadata_args(
+        edge_cache,
+        node_count=x.shape[0],
+        dst_ptr=dst_ptr,
+        source_order=source_order,
+        source_ptr=source_ptr,
+    )
+    if metadata_args is None:
+        dst_ptr = _dst_ptr_from_sorted(
+            torch,
+            edge_cache.dst,
+            x.shape[0],
+            destinations_sorted=destinations_sorted,
+        )
+        if dst_ptr is None:
+            return None
+        source_order = edge_cache.src.new_empty((0,), dtype=torch.int32)
+        source_ptr = edge_cache.src.new_empty((0,), dtype=torch.int32)
+    else:
+        dst_ptr, source_order, source_ptr = metadata_args
+    x_arg = _aligned_contiguous(x)
+    d_arg = _aligned_contiguous(edge_cache.D_packed)
+    dt_arg = d_arg
+    radial_arg = _aligned_contiguous(radial_feat)
+    edge_env_arg = _aligned_contiguous(edge_cache.edge_env)
+    src_arg = _aligned_contiguous(edge_cache.src)
+    dst_arg = _aligned_contiguous(edge_cache.dst)
+    dst_ptr_arg = _aligned_contiguous(dst_ptr)
+    source_order_arg = _aligned_contiguous(source_order)
+    source_ptr_arg = _aligned_contiguous(source_ptr)
+    edge_src_gate = edge_cache.edge_src_gate
+    if edge_src_gate is None:
+        edge_src_gate = edge_cache.edge_env.new_empty((0,))
+    edge_src_gate_arg = _aligned_contiguous(edge_src_gate)
+    output = cute_so2(
+        state.handle,
+        x_arg,
+        d_arg,
+        dt_arg,
+        radial_arg,
+        edge_env_arg,
+        src_arg,
+        dst_arg,
+        dst_ptr_arg,
+        edge_src_gate_arg,
+        source_order=source_order_arg,
+        source_ptr=source_ptr_arg,
+    )
+    return _layout_like(output, x)
+
+
+def maybe_run_cute_so2(
+    block: Any,
+    x: Tensor,
+    edge_cache: Any,
+    radial_feat: Tensor,
+    dst_ptr: Tensor | None = None,
+    source_order: Tensor | None = None,
+    source_ptr: Tensor | None = None,
+) -> Tensor | None:
+    """Use prevalidated opaque dispatch, or the conservative eager fallback."""
+    state = getattr(block, "_deepmd_cute_so2_state", None)
+    use_prepared = isinstance(
+        state,
+        _RegisteredSO2State,
+    ) or runtime_policy.is_so2_thin_wrapper_enabled(_tensor_compute_capability(x))
+    if use_prepared:
+        output = _maybe_run_prepared_cute_so2(
+            block,
+            x,
+            edge_cache,
+            radial_feat,
+            dst_ptr,
+            source_order,
+            source_ptr,
+        )
+        if output is not None:
+            return output
+    return _maybe_run_cute_so2_fallback(
+        block,
+        x,
+        edge_cache,
+        radial_feat,
+        dst_ptr,
+        source_order,
+        source_ptr,
+    )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py b/deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py
new file mode 100644
index 0000000000..471ec21d5a
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py
@@ -0,0 +1,380 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Runtime contract for exact-shape Neo Phase-C backward."""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    Any,
+)
+
+import torch
+
+from ..compile_cache import (
+    device_aware_lru_cache,
+)
+from .kernels.phase_c_backward import (
+    DEGREE_COUNT,
+    FOCUS_CHANNELS,
+    HIDDEN,
+    N_FOCUS,
+    REDUCED_COUNT,
+    compile_neo_phase_c_backward_layout,
+)
+from .wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES
+
+REQUIRED_ALIGNMENT = 16
+
+
+@dataclass(frozen=True)
+class NeoPhaseCBackwardLayoutOutputs:
+    """Caller-owned outputs and scratch for one fused Phase-C invocation."""
+
+    grad_stack: torch.Tensor
+    grad_wigner_dt: torch.Tensor
+    grad_logits: torch.Tensor
+    grad_edge: torch.Tensor
+    grad_z_partial: torch.Tensor
+    grad_z: torch.Tensor
+    grad_focus_src: torch.Tensor
+
+
+def _storage_id(tensor: torch.Tensor) -> int:
+    return tensor.untyped_storage()._cdata
+
+
+def _is_exact_view(lhs: torch.Tensor, rhs: torch.Tensor) -> bool:
+    """Return whether two tensors name the same logical and physical view."""
+    return (
+        _storage_id(lhs) == _storage_id(rhs)
+        and lhs.data_ptr() == rhs.data_ptr()
+        and lhs.storage_offset() == rhs.storage_offset()
+        and lhs.shape == rhs.shape
+        and lhs.stride() == rhs.stride()
+        and lhs.dtype == rhs.dtype
+        and lhs.device == rhs.device
+    )
+
+
+def _tensor_byte_region(tensor: torch.Tensor) -> tuple[int, int] | None:
+    """Return the physical byte range of a compact runtime tensor."""
+    if tensor.numel() == 0 or tensor.device.type == "meta":
+        return None
+    start = tensor.data_ptr()
+    return start, start + tensor.numel() * tensor.element_size()
+
+
+def _tensor_regions_overlap(
+    lhs: torch.Tensor,
+    lhs_region: tuple[int, int] | None,
+    rhs: torch.Tensor,
+    rhs_region: tuple[int, int] | None,
+) -> bool:
+    """Compare precomputed physical regions, including external storage views."""
+    if lhs.device != rhs.device:
+        return False
+    if lhs.device.type == "meta":
+        return torch._C._overlaps(lhs, rhs)
+    if lhs_region is None or rhs_region is None:
+        return False
+    lhs_start, lhs_stop = lhs_region
+    rhs_start, rhs_stop = rhs_region
+    return lhs_start < rhs_stop and rhs_start < lhs_stop
+
+
+def _require_alignment(
+    name: str,
+    tensor: torch.Tensor,
+    alignment: int = REQUIRED_ALIGNMENT,
+) -> None:
+    """Enforce the alignment promised to CuTe by ``assumed_align``."""
+    if tensor.device.type == "meta":
+        return
+
+    byte_offset = tensor.storage_offset() * tensor.element_size()
+    pointer_remainder = tensor.data_ptr() % alignment
+    storage_remainder = tensor.untyped_storage().data_ptr() % alignment
+    offset_remainder = byte_offset % alignment
+    if pointer_remainder or storage_remainder or offset_remainder:
+        raise ValueError(
+            f"{name} must be {alignment}-byte aligned; got data pointer "
+            f"remainder {pointer_remainder}, storage pointer remainder "
+            f"{storage_remainder}, and byte storage-offset remainder "
+            f"{offset_remainder}"
+        )
+
+
+def _require_tensor(
+    name: str,
+    tensor: torch.Tensor,
+    shape: tuple[int, ...],
+    *,
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+) -> None:
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+    if tensor.device != device:
+        raise ValueError(f"{name} must be on {device}, got {tensor.device}")
+    if tensor.dtype != dtype:
+        raise ValueError(f"{name} must have dtype {dtype}, got {tensor.dtype}")
+    if not tensor.is_contiguous():
+        raise ValueError(f"{name} must be compact")
+    _require_alignment(name, tensor)
+
+
+@device_aware_lru_cache(maxsize=32)
+def _compile_layout_boundary(
+    focus_eps: float,
+    focus_tau: float,
+    focus_label_smoothing: float,
+    use_focus_norm: bool,
+) -> Any:
+    return compile_neo_phase_c_backward_layout(
+        focus_eps=focus_eps,
+        focus_tau=focus_tau,
+        focus_label_smoothing=focus_label_smoothing,
+        use_focus_norm=use_focus_norm,
+    )
+
+
+class CuteNeoPhaseCBackwardLayout:
+    """Callable for the fused node-owned Phase-C boundary.
+
+    One invocation replaces Phase-C backward, envelope-softmax backward plus
+    its z reduction, and focus-source backward.  The Phase-C residual adjoint
+    is the unmodified ``grad_out`` input and remains owned by the caller.
+
+    Inputs retain their model-native ranks. ``grad_stack`` aliases compact
+    edge-major ``stack`` storage and is written only after the complete source
+    fragment has been consumed. ``grad_focus_src`` is compact ``(F,E,C)``.
+    """
+
+    def __init__(
+        self,
+        *,
+        focus_eps: float,
+        focus_tau: float,
+        focus_label_smoothing: float,
+        use_focus_norm: bool = True,
+    ) -> None:
+        self._compiled = _compile_layout_boundary(
+            float(focus_eps),
+            float(focus_tau),
+            float(focus_label_smoothing),
+            bool(use_focus_norm),
+        )
+
+    def __call__(
+        self,
+        grad_out: torch.Tensor,
+        stack: torch.Tensor,
+        wigner_dt: torch.Tensor,
+        alpha: torch.Tensor,
+        focus_alpha: torch.Tensor,
+        dst_ptr: torch.Tensor,
+        rotate_inv_rescale: torch.Tensor,
+        edge_gate: torch.Tensor,
+        z_bias_raw: torch.Tensor,
+        group_max: torch.Tensor,
+        denom: torch.Tensor,
+        focus_src: torch.Tensor,
+        focus_weight: torch.Tensor,
+        focus_scale: torch.Tensor,
+        outputs: NeoPhaseCBackwardLayoutOutputs,
+    ) -> NeoPhaseCBackwardLayoutOutputs:
+        edge_count = stack.shape[0]
+        node_count = grad_out.shape[0]
+        device = stack.device
+        stack_shape = (
+            edge_count,
+            N_FOCUS,
+            REDUCED_COUNT,
+            FOCUS_CHANNELS,
+        )
+
+        _require_tensor(
+            "grad_out",
+            grad_out,
+            (node_count, DEGREE_COUNT, HIDDEN),
+            device=device,
+        )
+        _require_tensor("stack", stack, stack_shape, device=device)
+        _require_tensor(
+            "wigner_dt",
+            wigner_dt,
+            (edge_count, PACKED_WIGNER_VALUES),
+            device=device,
+        )
+        _require_tensor("alpha", alpha, (edge_count, N_FOCUS), device=device)
+        _require_tensor(
+            "focus_alpha", focus_alpha, (edge_count, N_FOCUS), device=device
+        )
+        _require_tensor(
+            "dst_ptr",
+            dst_ptr,
+            (node_count + 1,),
+            device=device,
+            dtype=torch.int32,
+        )
+        _require_tensor(
+            "rotate_inv_rescale",
+            rotate_inv_rescale,
+            (DEGREE_COUNT,),
+            device=device,
+        )
+        _require_tensor("edge_gate", edge_gate, (edge_count,), device=device)
+        _require_tensor("z_bias_raw", z_bias_raw, (N_FOCUS,), device=device)
+        _require_tensor("group_max", group_max, (node_count, N_FOCUS), device=device)
+        _require_tensor("denom", denom, (node_count, N_FOCUS), device=device)
+        _require_tensor(
+            "focus_src",
+            focus_src,
+            (edge_count, N_FOCUS, FOCUS_CHANNELS),
+            device=device,
+        )
+        _require_tensor(
+            "focus_weight",
+            focus_weight,
+            (FOCUS_CHANNELS, N_FOCUS),
+            device=device,
+        )
+        _require_tensor(
+            "focus_scale",
+            focus_scale,
+            (N_FOCUS, FOCUS_CHANNELS),
+            device=device,
+        )
+
+        _require_tensor(
+            "outputs.grad_stack",
+            outputs.grad_stack,
+            stack_shape,
+            device=device,
+        )
+        if not _is_exact_view(outputs.grad_stack, stack):
+            raise ValueError("outputs.grad_stack must be the exact in-place stack view")
+        _require_tensor(
+            "outputs.grad_wigner_dt",
+            outputs.grad_wigner_dt,
+            (edge_count, PACKED_WIGNER_VALUES),
+            device=device,
+        )
+        _require_tensor(
+            "outputs.grad_logits",
+            outputs.grad_logits,
+            (edge_count, N_FOCUS),
+            device=device,
+        )
+        _require_tensor(
+            "outputs.grad_edge", outputs.grad_edge, (edge_count,), device=device
+        )
+        _require_tensor(
+            "outputs.grad_z_partial",
+            outputs.grad_z_partial,
+            (node_count, N_FOCUS),
+            device=device,
+        )
+        _require_tensor("outputs.grad_z", outputs.grad_z, (N_FOCUS,), device=device)
+        _require_tensor(
+            "outputs.grad_focus_src",
+            outputs.grad_focus_src,
+            (N_FOCUS, edge_count, FOCUS_CHANNELS),
+            device=device,
+        )
+
+        input_tensors = tuple(
+            (name, tensor, _tensor_byte_region(tensor))
+            for name, tensor in (
+                ("grad_out", grad_out),
+                ("stack", stack),
+                ("wigner_dt", wigner_dt),
+                ("alpha", alpha),
+                ("focus_alpha", focus_alpha),
+                ("dst_ptr", dst_ptr),
+                ("rotate_inv_rescale", rotate_inv_rescale),
+                ("edge_gate", edge_gate),
+                ("z_bias_raw", z_bias_raw),
+                ("group_max", group_max),
+                ("denom", denom),
+                ("focus_src", focus_src),
+                ("focus_weight", focus_weight),
+                ("focus_scale", focus_scale),
+            )
+        )
+        output_tensors = tuple(
+            (
+                f"outputs.{field_name}",
+                getattr(outputs, field_name),
+                _tensor_byte_region(getattr(outputs, field_name)),
+            )
+            for field_name in (
+                "grad_stack",
+                "grad_wigner_dt",
+                "grad_logits",
+                "grad_edge",
+                "grad_z_partial",
+                "grad_z",
+                "grad_focus_src",
+            )
+        )
+
+        for output_index, (
+            output_name,
+            output,
+            output_region,
+        ) in enumerate(output_tensors):
+            for other_name, other, other_region in output_tensors[output_index + 1 :]:
+                if _tensor_regions_overlap(
+                    output,
+                    output_region,
+                    other,
+                    other_region,
+                ):
+                    raise ValueError(
+                        f"{output_name} must not overlap output {other_name}"
+                    )
+            for input_name, input_tensor, input_region in input_tensors:
+                if not _tensor_regions_overlap(
+                    output,
+                    output_region,
+                    input_tensor,
+                    input_region,
+                ):
+                    continue
+                is_stack_adjoint = (
+                    output_name == "outputs.grad_stack" and input_name == "stack"
+                )
+                if is_stack_adjoint:
+                    continue
+                raise ValueError(f"{output_name} must not overlap input {input_name}")
+
+        self._compiled(
+            grad_out,
+            stack,
+            wigner_dt,
+            alpha,
+            focus_alpha,
+            dst_ptr,
+            rotate_inv_rescale,
+            edge_gate,
+            z_bias_raw,
+            group_max,
+            denom,
+            focus_src,
+            focus_weight,
+            focus_scale,
+            outputs.grad_stack,
+            outputs.grad_wigner_dt,
+            outputs.grad_logits,
+            outputs.grad_edge,
+            outputs.grad_z_partial,
+            outputs.grad_z,
+            outputs.grad_focus_src,
+        )
+        return outputs
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/radial_phase_a.py b/deepmd/pt_expt/kernels/cute/sezm/so2/radial_phase_a.py
new file mode 100644
index 0000000000..1b63f3e640
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/radial_phase_a.py
@@ -0,0 +1,430 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Runtime wrapper for the source-CSR node-tiled radial/Phase-A backward."""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+from ..compile_cache import (
+    device_aware_lru_cache,
+)
+from .wigner_layout import (
+    PACKED_VALUE_COUNT,
+)
+
+if TYPE_CHECKING:
+    from torch import (
+        Tensor,
+    )
+
+
+DEGREE_COUNT = 16
+REDUCED_COUNT = 10
+HIDDEN = 64
+FOCUS_COUNT = 2
+FOCUS_HIDDEN = 32
+PACKED_WIGNER_VALUES = PACKED_VALUE_COUNT
+RADIAL_WIDTH = 4 * FOCUS_HIDDEN
+COMPACT_WIDTH = 25
+PROJECTION_INPUT_WIDTH = COMPACT_WIDTH + FOCUS_COUNT
+
+
+@dataclass(frozen=True)
+class NeoSourceCSR:
+    """Indirect source CSR over the unchanged physical edge order."""
+
+    source_order: Tensor
+    source_ptr: Tensor
+
+
+@dataclass(frozen=True)
+class NeoRadialPhaseABackwardNodeResult:
+    """Output buffers populated by the node-tiled backward."""
+
+    grad_x_wide: Tensor
+    grad_d_full: Tensor
+    grad_radial_m0: Tensor
+
+
+def build_source_csr(
+    src: Tensor,
+    node_count: int,
+    *,
+    validate_sources: bool = False,
+) -> NeoSourceCSR:
+    """Build indirect source CSR without changing the physical edge order.
+
+    Source bounds are always checked before constructing the CSR. Setting
+    ``validate_sources=True`` reports an eager ``ValueError`` and therefore
+    synchronizes when ``src`` is CUDA; the default uses an asynchronous device
+    assertion. Callers should build this once with the edge cache and retain
+    both tensors.
+    """
+    import torch
+
+    if src.dim() != 1:
+        raise ValueError("src must be one-dimensional")
+    if src.dtype not in (torch.int32, torch.int64):
+        raise TypeError("src must have dtype int32 or int64")
+    if node_count < 0:
+        raise ValueError("node_count must be non-negative")
+    if src.numel() > 2**31 - 1:
+        raise ValueError("source CSR int32 indexing requires E <= 2**31 - 1")
+
+    src = src.contiguous()
+    if src.numel() != 0:
+        valid = torch.all((src >= 0) & (src < node_count))
+        message = "source-CSR backward requires source indices in [0, node_count)"
+        if validate_sources:
+            if not bool(valid):
+                raise ValueError(message)
+        else:
+            torch._assert_async(
+                valid,
+                message,
+            )
+
+    source_order_i64 = torch.argsort(src, stable=True)
+    sorted_src = src.index_select(0, source_order_i64)
+    boundaries = torch.arange(
+        node_count + 1,
+        device=src.device,
+        dtype=src.dtype,
+    )
+    source_ptr = torch.searchsorted(
+        sorted_src,
+        boundaries,
+        out_int32=True,
+    ).contiguous()
+    source_order = source_order_i64.to(dtype=torch.int32).contiguous()
+    return NeoSourceCSR(source_order=source_order, source_ptr=source_ptr)
+
+
+@device_aware_lru_cache(maxsize=4)
+def _compile_node_tiled() -> Any:
+    from .kernels.radial_phase_a_backward import (
+        compile_neo_radial_phase_a_backward_node_tiled,
+    )
+
+    return compile_neo_radial_phase_a_backward_node_tiled()
+
+
+def _expect_tensor(
+    name: str,
+    tensor: Tensor,
+    shape: tuple[int, ...],
+    *,
+    device: Any,
+    dtype: Any,
+) -> None:
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+    if tensor.device != device:
+        raise ValueError(f"{name} must be on {device}, got {tensor.device}")
+    if tensor.dtype != dtype:
+        raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}")
+    if not tensor.is_contiguous():
+        raise ValueError(f"{name} must be contiguous")
+
+
+def prepare_batched_radial_projection_weight(
+    combined_weight: Tensor,
+    attention_radial_weight: Tensor,
+) -> Tensor:
+    """Precombine the compact and attention adjoints for one FP32 GEMM."""
+    import torch
+
+    device = combined_weight.device
+    _expect_tensor(
+        "combined_weight",
+        combined_weight,
+        (RADIAL_WIDTH, COMPACT_WIDTH),
+        device=device,
+        dtype=torch.float32,
+    )
+    _expect_tensor(
+        "attention_radial_weight",
+        attention_radial_weight,
+        (FOCUS_HIDDEN, FOCUS_COUNT),
+        device=device,
+        dtype=torch.float32,
+    )
+
+    projection_weight = combined_weight.new_zeros(
+        (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH)
+    )
+    projection_weight[:COMPACT_WIDTH].copy_(combined_weight.transpose(0, 1))
+    projection_weight[COMPACT_WIDTH:, :FOCUS_HIDDEN].copy_(
+        attention_radial_weight.transpose(0, 1)
+    )
+    return projection_weight
+
+
+def _project_batched_radial_adjoint(
+    projection_weight: Tensor,
+    grad_radial_m0: Tensor,
+    consumed_workspace: Tensor,
+) -> None:
+    """Project the 27-column adjoint packed in consumed edge scratch."""
+    import torch
+
+    if torch.backends.cuda.matmul.allow_tf32:
+        raise RuntimeError("strict FP32 requires allow_tf32=False")
+    if torch.get_float32_matmul_precision() != "highest":
+        raise RuntimeError("strict FP32 requires float32 matmul precision 'highest'")
+    edge_count = consumed_workspace.shape[0]
+    device = consumed_workspace.device
+    tensors = (
+        (
+            "projection_weight",
+            projection_weight,
+            (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH),
+        ),
+        ("grad_radial_m0", grad_radial_m0, (edge_count, RADIAL_WIDTH)),
+        (
+            "consumed_workspace",
+            consumed_workspace,
+            (edge_count, FOCUS_COUNT * REDUCED_COUNT * FOCUS_HIDDEN),
+        ),
+    )
+    for name, tensor, shape in tensors:
+        _expect_tensor(
+            name,
+            tensor,
+            shape,
+            device=device,
+            dtype=torch.float32,
+        )
+
+    projection_input = consumed_workspace[:, :PROJECTION_INPUT_WIDTH]
+    torch.mm(projection_input, projection_weight, out=grad_radial_m0)
+
+
+def _validate_csr_values(
+    source_order: Tensor,
+    source_ptr: Tensor,
+    edge_count: int,
+) -> None:
+    import torch
+
+    valid = (source_ptr[0] == 0) & (source_ptr[-1] == edge_count)
+    if source_ptr.numel() > 1:
+        valid = valid & torch.all(source_ptr[1:] >= source_ptr[:-1])
+    if not bool(valid):
+        raise ValueError(
+            "source_ptr must be nondecreasing, begin at zero, and end at E"
+        )
+    expected = torch.arange(
+        edge_count,
+        device=source_order.device,
+        dtype=source_order.dtype,
+    )
+    if not torch.equal(torch.sort(source_order).values, expected):
+        raise ValueError("source_order must be a permutation of [0, E)")
+
+
+def run_neo_radial_phase_a_backward_node_tiled(
+    grad_out_focus: Tensor,
+    grad_logits: Tensor,
+    radial_state: Tensor,
+    channel_basis: Tensor,
+    x_wide: Tensor,
+    source_order: Tensor,
+    source_ptr: Tensor,
+    d_full: Tensor,
+    *,
+    grad_focus_src_focus: Tensor,
+    batched_radial_projection_weight: Tensor,
+    grad_x_wide: Tensor | None = None,
+    grad_d_full: Tensor | None = None,
+    grad_radial_m0: Tensor | None = None,
+    validate_csr: bool = False,
+) -> NeoRadialPhaseABackwardNodeResult:
+    """Run the node-owned backward over indirect source CSR.
+
+    The kernel recomputes Phase A, fuses the focus-source adjoint, uses
+    four-lane warp reductions with a 68-float shared row pitch, and packs the
+    27-column radial projection input for one strict-FP32 matrix call.
+    ``grad_out_focus`` is repacked in place as projection workspace and must
+    not be reused after this function returns.
+    """
+    import torch
+
+    if not x_wide.is_cuda:
+        raise ValueError("node-tiled radial Phase-A backward requires CUDA tensors")
+    if x_wide.dtype != torch.float32:
+        raise TypeError("node-tiled radial Phase-A backward specializes float32")
+    if not x_wide.is_contiguous() or x_wide.dim() != 2:
+        raise ValueError("x_wide must be a contiguous two-dimensional tensor")
+
+    device = x_wide.device
+    dtype = x_wide.dtype
+    node_count = x_wide.shape[0]
+    edge_count = grad_out_focus.shape[0]
+    _expect_tensor(
+        "x_wide",
+        x_wide,
+        (node_count, DEGREE_COUNT * HIDDEN),
+        device=device,
+        dtype=dtype,
+    )
+    _expect_tensor(
+        "grad_out_focus",
+        grad_out_focus,
+        (edge_count, FOCUS_COUNT * REDUCED_COUNT * FOCUS_HIDDEN),
+        device=device,
+        dtype=dtype,
+    )
+    _expect_tensor(
+        "grad_focus_src_focus",
+        grad_focus_src_focus,
+        (FOCUS_COUNT, edge_count, FOCUS_HIDDEN),
+        device=device,
+        dtype=dtype,
+    )
+    _expect_tensor(
+        "grad_logits",
+        grad_logits,
+        (edge_count, FOCUS_COUNT),
+        device=device,
+        dtype=dtype,
+    )
+    _expect_tensor(
+        "radial_state",
+        radial_state,
+        (edge_count, COMPACT_WIDTH),
+        device=device,
+        dtype=dtype,
+    )
+    _expect_tensor(
+        "batched_radial_projection_weight",
+        batched_radial_projection_weight,
+        (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH),
+        device=device,
+        dtype=dtype,
+    )
+    _expect_tensor(
+        "channel_basis",
+        channel_basis,
+        (HIDDEN,),
+        device=device,
+        dtype=dtype,
+    )
+    _expect_tensor(
+        "source_order",
+        source_order,
+        (edge_count,),
+        device=device,
+        dtype=torch.int32,
+    )
+    _expect_tensor(
+        "source_ptr",
+        source_ptr,
+        (node_count + 1,),
+        device=device,
+        dtype=torch.int32,
+    )
+    _expect_tensor(
+        "d_full",
+        d_full,
+        (edge_count, PACKED_WIGNER_VALUES),
+        device=device,
+        dtype=dtype,
+    )
+    if node_count == 0 and edge_count != 0:
+        raise ValueError("a non-empty edge list requires at least one source node")
+    torch._assert_async(
+        source_ptr[0] == 0,
+        "source-CSR backward requires source_ptr[0] == 0",
+    )
+    torch._assert_async(
+        source_ptr[-1] == edge_count,
+        "source-CSR backward requires source_ptr[-1] == edge_count",
+    )
+    if source_ptr.numel() > 1:
+        torch._assert_async(
+            torch.all(source_ptr[1:] >= source_ptr[:-1]),
+            "source-CSR backward requires nondecreasing source_ptr",
+        )
+    if source_order.numel() != 0:
+        torch._assert_async(
+            torch.all((source_order >= 0) & (source_order < edge_count)),
+            "source-CSR backward requires source_order entries in [0, E)",
+        )
+    if validate_csr:
+        _validate_csr_values(source_order, source_ptr, edge_count)
+
+    if grad_x_wide is None:
+        grad_x_wide = torch.empty_like(x_wide)
+    else:
+        _expect_tensor(
+            "grad_x_wide",
+            grad_x_wide,
+            (node_count, DEGREE_COUNT * HIDDEN),
+            device=device,
+            dtype=dtype,
+        )
+    if grad_d_full is None:
+        grad_d_full = torch.empty_like(d_full)
+    else:
+        _expect_tensor(
+            "grad_d_full",
+            grad_d_full,
+            (edge_count, PACKED_WIGNER_VALUES),
+            device=device,
+            dtype=dtype,
+        )
+    if grad_radial_m0 is None:
+        grad_radial_m0 = torch.empty(
+            (edge_count, RADIAL_WIDTH),
+            device=device,
+            dtype=dtype,
+        )
+    else:
+        _expect_tensor(
+            "grad_radial_m0",
+            grad_radial_m0,
+            (edge_count, RADIAL_WIDTH),
+            device=device,
+            dtype=dtype,
+        )
+
+    result = NeoRadialPhaseABackwardNodeResult(
+        grad_x_wide=grad_x_wide,
+        grad_d_full=grad_d_full,
+        grad_radial_m0=grad_radial_m0,
+    )
+    if edge_count == 0:
+        grad_x_wide.zero_()
+        return result
+
+    with torch.cuda.device(device):
+        kernel = _compile_node_tiled()
+        kernel(
+            grad_out_focus,
+            grad_focus_src_focus,
+            grad_logits,
+            radial_state,
+            channel_basis,
+            x_wide,
+            source_order,
+            source_ptr,
+            d_full,
+            grad_x_wide,
+            grad_d_full,
+        )
+        _project_batched_radial_adjoint(
+            batched_radial_projection_weight,
+            grad_radial_m0,
+            grad_out_focus,
+        )
+    return result
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py b/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py
new file mode 100644
index 0000000000..8cf05fe469
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py
@@ -0,0 +1,918 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Runtime runner for the Neo CuTe SO2 unit."""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    Any,
+)
+
+from ..runtime_policy import (
+    FUSED_SO2_GATE_CAPABILITIES,
+    SM80_PROFILE_CAPABILITIES,
+    SM90_CAPABILITY,
+    SUPPORTED_SO2_CAPABILITIES,
+)
+
+
+def _uses_packed_message_grid(compute_capability: tuple[int, int]) -> bool:
+    """Return whether the packed message-grid kernels support this GPU."""
+    return (
+        compute_capability in SM80_PROFILE_CAPABILITIES
+        or compute_capability == SM90_CAPABILITY
+    )
+
+
+def _validate_runtime_config(
+    runtime_config: Any,
+    *,
+    compute_capability: tuple[int, int] | None = None,
+) -> None:
+    if compute_capability not in SUPPORTED_SO2_CAPABILITIES:
+        raise RuntimeError("Neo SO2 requires a supported compute capability")
+    if runtime_config.native_sm90_path and compute_capability != SM90_CAPABILITY:
+        raise RuntimeError("the native SM90 SO2 path must be selected only on sm_90")
+    if runtime_config.per_focus_so2_fwd_pair != (
+        compute_capability in SM80_PROFILE_CAPABILITIES
+    ):
+        raise RuntimeError("the per-focus SO2 path must be selected on sm_80/sm_86")
+    if runtime_config.combined_so2_gate != (
+        compute_capability in FUSED_SO2_GATE_CAPABILITIES
+    ):
+        raise RuntimeError(
+            "the combined SO2/gate path must be selected on sm_89/sm_120"
+        )
+
+
+def _combined_radial_weight(
+    torch: Any,
+    radial_hidden_proj: Any,
+    radial_degree_mixer: Any,
+) -> Any:
+    hidden_weight = radial_hidden_proj.weight.detach()
+    mixer_weight = radial_degree_mixer.weight.detach()
+    cache_key = (
+        hidden_weight.data_ptr(),
+        hidden_weight._version,
+        mixer_weight.data_ptr(),
+        mixer_weight._version,
+        hidden_weight.dtype,
+        hidden_weight.device,
+    )
+    cache = getattr(radial_degree_mixer, "_deepmd_cute_combined_weight", None)
+    if cache is not None and cache[0] == cache_key:
+        return cache[1]
+    blocks = []
+    for degree in range(4):
+        mixer_block = mixer_weight[degree * 64 : (degree + 1) * 64, :]
+        blocks.append(torch.mm(hidden_weight, mixer_block))
+    combined = torch.cat(blocks, dim=0).contiguous()
+    radial_degree_mixer._deepmd_cute_combined_weight = (cache_key, combined)
+    return combined
+
+
+def _combined_attention_radial_weight(
+    torch: Any,
+    radial_hidden_proj: Any,
+    attention_weight: Any,
+) -> Any:
+    hidden_weight = radial_hidden_proj.weight.detach()
+    attention_weight = attention_weight.detach()
+    cache_key = (
+        hidden_weight.data_ptr(),
+        hidden_weight._version,
+        attention_weight.data_ptr(),
+        attention_weight._version,
+        hidden_weight.dtype,
+        hidden_weight.device,
+    )
+    cache = getattr(radial_hidden_proj, "_deepmd_cute_attention_radial_weight", None)
+    if cache is not None and cache[0] == cache_key:
+        return cache[1]
+    blocks = []
+    for focus in range(2):
+        hidden_block = hidden_weight[:, focus * 32 : (focus + 1) * 32]
+        blocks.append(torch.mv(hidden_block, attention_weight[:, focus, 0]))
+    combined = torch.stack(blocks, dim=1).contiguous()
+    radial_hidden_proj._deepmd_cute_attention_radial_weight = (cache_key, combined)
+    return combined
+
+
+def _batched_radial_projection_weight(
+    radial_degree_mixer: Any,
+    combined_weight: Any,
+    attention_radial_weight: Any,
+) -> Any:
+    cache_key = (
+        combined_weight.data_ptr(),
+        combined_weight._version,
+        attention_radial_weight.data_ptr(),
+        attention_radial_weight._version,
+        combined_weight.dtype,
+        combined_weight.device,
+    )
+    cache = getattr(
+        radial_degree_mixer,
+        "_deepmd_cute_batched_radial_projection_weight",
+        None,
+    )
+    if cache is not None and cache[0] == cache_key:
+        return cache[1]
+
+    from .radial_phase_a import (
+        prepare_batched_radial_projection_weight,
+    )
+
+    projection_weight = prepare_batched_radial_projection_weight(
+        combined_weight,
+        attention_radial_weight,
+    )
+    # Keep the derived operands alive so allocator pointer reuse cannot spoof
+    # the versioned cache key after a parameter update.
+    radial_degree_mixer._deepmd_cute_batched_radial_projection_weight = (
+        cache_key,
+        projection_weight,
+        combined_weight,
+        attention_radial_weight,
+    )
+    return projection_weight
+
+
+def _edge_gate(torch: Any, edge_cache: Any) -> Any:
+    gate = edge_cache.edge_env.reshape(-1).float().clamp_min(0.0)
+    if edge_cache.edge_src_gate is not None:
+        gate = gate * edge_cache.edge_src_gate.reshape(-1).float().clamp_min(0.0).sqrt()
+    return gate.contiguous()
+
+
+@dataclass
+class StackCache:
+    y: Any
+    logits: Any | None
+    non_linear: Any
+    final: bool
+
+    def __post_init__(self) -> None:
+        self.y = self.y.detach()
+        if self.logits is not None:
+            self.logits = self.logits.detach()
+
+
+class _NeoSO2BackwardWorkspaceBase:
+    _EXPORTED_NAMES = (
+        "grad_stack_focus",
+        "grad_focus_alpha",
+        "grad_dt",
+        "grad_alpha",
+        "grad_logits",
+        "grad_edge",
+        "grad_z_partial",
+        "grad_z",
+        "grad_x_rot",
+        "grad_radial_flat",
+        "grad_x_wide_phase_a",
+        "grad_d",
+        "grad_y",
+        "grad_gate_logits",
+        "grad_mixed_slab",
+    )
+
+    def attach_to(self, runner: Any) -> None:
+        for name in self._EXPORTED_NAMES:
+            setattr(runner, name, getattr(self, name))
+
+
+def _structural_memory_views(
+    torch: Any,
+    *,
+    edge_count: int,
+    like: Any,
+    phase_c_stack: Any,
+    phase_c_y: Any | None,
+    radial_scratch: Any,
+    radial_values_per_edge: int,
+    phase_c_single_input_reuse: bool,
+) -> tuple[Any, Any | None, Any]:
+    """Validate and expose saved stack buffers that become backward scratch."""
+    expected_phase_shape = (edge_count, 2, 10, 32)
+    if phase_c_single_input_reuse:
+        if phase_c_y is not None:
+            raise ValueError(
+                "structural memory reuse requires phase_c_y to be absent in "
+                "single-input mode"
+            )
+        phase_tensors = (("phase_c_stack", phase_c_stack, expected_phase_shape),)
+    else:
+        phase_tensors = (
+            ("phase_c_stack", phase_c_stack, expected_phase_shape),
+            ("phase_c_y", phase_c_y, expected_phase_shape),
+        )
+    for name, tensor, expected_shape in phase_tensors:
+        if tensor is None or tuple(tensor.shape) != expected_shape:
+            raise ValueError(
+                f"structural memory reuse requires {name} shape {expected_shape}"
+            )
+        if tensor.dtype != torch.float32 or tensor.dtype != like.dtype:
+            raise ValueError(f"structural memory reuse requires FP32 {name} storage")
+        if tensor.device != like.device or not tensor.is_contiguous():
+            raise ValueError(
+                f"structural memory reuse requires contiguous {name} on {like.device}"
+            )
+        if tensor.storage_offset() != 0:
+            raise ValueError(
+                f"structural memory reuse requires zero-offset {name} storage"
+            )
+    required_radial_values = edge_count * radial_values_per_edge
+    if radial_scratch is None or radial_scratch.numel() < required_radial_values:
+        raise ValueError(
+            "structural memory reuse requires radial_scratch capacity of at least "
+            f"{required_radial_values} values"
+        )
+    if radial_scratch.dtype != torch.float32 or radial_scratch.dtype != like.dtype:
+        raise ValueError("structural memory reuse requires FP32 radial_scratch storage")
+    if radial_scratch.device != like.device or not radial_scratch.is_contiguous():
+        raise ValueError(
+            "structural memory reuse requires contiguous radial_scratch on "
+            f"{like.device}"
+        )
+    if radial_scratch.storage_offset() != 0:
+        raise ValueError(
+            "structural memory reuse requires zero-offset radial_scratch storage"
+        )
+    tensors = (*[tensor for _, tensor, _ in phase_tensors], radial_scratch)
+    storages = {tensor.untyped_storage()._cdata for tensor in tensors}
+    if len(storages) != len(tensors):
+        raise ValueError("structural memory reuse requires distinct storages")
+    return (
+        phase_c_stack.view(edge_count, 10 * 64),
+        None if phase_c_y is None else phase_c_y.view(edge_count, 10 * 64),
+        radial_scratch.flatten(),
+    )
+
+
+class NeoSO2BackwardWorkspace(_NeoSO2BackwardWorkspaceBase):
+    """Lazily allocated SO2 backward scratch with lifetime-based slab reuse."""
+
+    def __init__(
+        self,
+        torch: Any,
+        *,
+        edge_count: int,
+        node_count: int,
+        d_full: Any,
+        dt_full: Any,
+        radial: Any,
+        phase_c_stack: Any | None = None,
+        phase_c_y: Any | None = None,
+        radial_scratch: Any | None = None,
+        structural_memory_reuse: bool = False,
+        phase_c_single_input_reuse: bool = False,
+    ) -> None:
+        opts = {"device": d_full.device, "dtype": d_full.dtype}
+        d_values_per_edge = 1
+        for size in d_full.shape[1:]:
+            d_values_per_edge *= size
+        radial_values_per_edge = 1
+        for size in radial.shape[1:]:
+            radial_values_per_edge *= size
+        if d_values_per_edge > 10 * 64:
+            raise ValueError("Neo SO2 grad_D does not fit the Phase-C scratch slab")
+
+        if phase_c_single_input_reuse and not structural_memory_reuse:
+            raise ValueError(
+                "Phase-C single-input reuse requires structural memory reuse"
+            )
+
+        phase_c_stack_flat = None
+        phase_c_y_flat = None
+        radial_scratch_flat = None
+        if structural_memory_reuse:
+            phase_c_stack_flat, phase_c_y_flat, radial_scratch_flat = (
+                _structural_memory_views(
+                    torch,
+                    edge_count=edge_count,
+                    like=d_full,
+                    phase_c_stack=phase_c_stack,
+                    phase_c_y=phase_c_y,
+                    radial_scratch=radial_scratch,
+                    radial_values_per_edge=radial_values_per_edge,
+                    phase_c_single_input_reuse=phase_c_single_input_reuse,
+                )
+            )
+
+        self._phase_c_or_d = (
+            phase_c_stack_flat
+            if phase_c_stack_flat is not None
+            else torch.empty(
+                edge_count,
+                10 * 64,
+                device=opts["device"],
+                dtype=opts["dtype"],
+            )
+        )
+        self.grad_stack_focus = self._phase_c_or_d.view(edge_count, 2, 10, 32)
+        self.grad_d = self._phase_c_or_d.flatten()[
+            : edge_count * d_values_per_edge
+        ].view_as(d_full)
+        self.grad_dt = torch.empty_like(dt_full)
+
+        # Gate backward completes before radial backward writes the final radial grad.
+        if radial_scratch_flat is not None:
+            self._gate_or_radial = radial_scratch_flat
+        else:
+            self._gate_or_radial = torch.empty(
+                edge_count,
+                radial_values_per_edge,
+                device=opts["device"],
+                dtype=opts["dtype"],
+            )
+        self.grad_gate_logits = None
+        self.grad_radial_flat = self._gate_or_radial.flatten()[
+            : edge_count * radial_values_per_edge
+        ].view(edge_count, radial_values_per_edge)
+
+        # The per-layer gate output is dead before radial backward writes grad_x_rot.
+        self._gate_or_x_rot = (
+            phase_c_stack_flat
+            if phase_c_stack_flat is not None
+            else torch.empty(
+                edge_count,
+                10 * 64,
+                device=opts["device"],
+                dtype=opts["dtype"],
+            )
+        )
+        self.grad_y = self._gate_or_x_rot.view(edge_count * 2, 10 * 32)
+        self.grad_x_rot = self._gate_or_x_rot
+        if phase_c_stack_flat is None:
+            self.grad_mixed_slab = None
+        elif phase_c_y_flat is not None:
+            self.grad_mixed_slab = phase_c_y_flat.view(edge_count, 2, 10, 32)
+        else:
+            self.grad_mixed_slab = torch.empty(
+                edge_count,
+                2,
+                10,
+                32,
+                device=opts["device"],
+                dtype=opts["dtype"],
+            )
+
+        self.grad_focus_alpha = torch.empty(
+            edge_count, 2, device=opts["device"], dtype=opts["dtype"]
+        )
+        self.grad_alpha = torch.empty(
+            edge_count, 2, device=opts["device"], dtype=opts["dtype"]
+        )
+        self.grad_logits = torch.empty(
+            edge_count, 2, device=opts["device"], dtype=opts["dtype"]
+        )
+        self.grad_edge = torch.empty(
+            edge_count, device=opts["device"], dtype=opts["dtype"]
+        )
+        self.grad_z_partial = torch.empty(
+            node_count, 2, device=opts["device"], dtype=opts["dtype"]
+        )
+        self.grad_z = torch.empty(2, device=opts["device"], dtype=opts["dtype"])
+        self.grad_x_wide_phase_a = torch.empty(
+            node_count,
+            16 * 64,
+            device=opts["device"],
+            dtype=opts["dtype"],
+        )
+
+
+class NeoFullCuteBackward:
+    def __init__(
+        self,
+        torch: Any,
+        block: Any,
+        record: Any,
+        x: Any,
+        d_full: Any,
+        dt_full: Any,
+        radial_feat: Any,
+        dst_ptr: Any,
+        source_order: Any,
+        source_ptr: Any,
+        *,
+        runtime_config: Any,
+    ) -> None:
+        from .kernels.envelope_softmax import (
+            compile_envelope_softmax_forward,
+        )
+        from .kernels.phase_a_radial_forward import (
+            run_neo_phase_a_radial_forward_packed_direct,
+        )
+        from .linear import (
+            run_neo_so2_linear_manual,
+        )
+        from .wigner_layout import (
+            PACKED_VALUE_COUNT,
+        )
+
+        self.torch = torch
+        self.block = block
+        self.so2 = block.so2_conv
+        self.record = record
+        self.config = runtime_config
+        device_index = x.device.index
+        if device_index is None:
+            device_index = torch.cuda.current_device()
+        compute_capability = tuple(torch.cuda.get_device_capability(device_index))
+        self.device_index = device_index
+        self.compute_capability = compute_capability
+        self.packed_message_grid = _uses_packed_message_grid(compute_capability)
+        self.compile_identity = (device_index, *compute_capability)
+        _validate_runtime_config(
+            runtime_config,
+            compute_capability=compute_capability,
+        )
+        self.x = x
+        self.d = d_full
+        self.dt = dt_full
+        expected_shape = (
+            record.edge_cache.src.numel(),
+            PACKED_VALUE_COUNT,
+        )
+        if (
+            tuple(d_full.shape) != expected_shape
+            or tuple(dt_full.shape) != expected_shape
+        ):
+            raise ValueError(
+                f"packed Neo SO2 Wigner tensors must have shape {expected_shape}"
+            )
+        if d_full.data_ptr() != dt_full.data_ptr():
+            raise ValueError("packed Neo SO2 D and Dt must share one storage")
+        self.radial = radial_feat
+        self.dst_ptr_i32 = dst_ptr.to(device=x.device, dtype=torch.int32).contiguous()
+        self.edge_count = record.edge_cache.src.numel()
+        self.node_count = x.shape[0]
+        self.src_i32 = record.edge_cache.src.to(torch.int32).contiguous()
+        self.src_i64 = record.edge_cache.src.contiguous()
+        self.dst_i32 = record.edge_cache.dst.to(torch.int32).contiguous()
+        self.dst_i64 = record.edge_cache.dst.contiguous()
+        if (
+            source_order.numel() == self.edge_count
+            and source_ptr.numel() == self.node_count + 1
+        ):
+            self.source_order_i32 = source_order.to(
+                device=x.device,
+                dtype=torch.int32,
+            ).contiguous()
+            self.source_ptr_i32 = source_ptr.to(
+                device=x.device,
+                dtype=torch.int32,
+            ).contiguous()
+        else:
+            from .radial_phase_a import (
+                build_source_csr,
+            )
+
+            source_csr = build_source_csr(self.src_i32, self.node_count)
+            self.source_order_i32 = source_csr.source_order
+            self.source_ptr_i32 = source_csr.source_ptr
+        self.rotate = self.so2.rotate_inv_rescale_full.contiguous()
+        self.edge_gate = _edge_gate(torch, record.edge_cache)
+        self.focus_norm_enabled = bool(self.so2.focus_norm)
+        if self.focus_norm_enabled:
+            focus_norm = self.so2.focus_compete_norm
+            self.focus_norm_eps = float(focus_norm.eps)
+            self.focus_norm_scale = focus_norm.adam_scale.detach().float().contiguous()
+        else:
+            self.focus_norm_eps = 0.0
+            # The no-norm specialization does not read this fixed-ABI argument;
+            # reuse an existing shape-compatible tensor instead of allocating one.
+            self.focus_norm_scale = (
+                self.so2.attn_qk_norm.adam_scale.detach().float().contiguous()
+            )
+        from .kernels.focus_source_backward import (
+            compile_neo_attention_prelude_forward,
+        )
+
+        with torch.cuda.device(self.device_index):
+            self.attention_prelude_forward = compile_neo_attention_prelude_forward(
+                self.focus_norm_eps,
+                float(self.so2.attn_qk_norm.eps),
+                float(self.so2.focus_softmax_tau),
+                float(self.so2.focus_label_smoothing),
+                self.compile_identity,
+                use_focus_norm=self.focus_norm_enabled,
+            )
+        from .kernels.qk_edge import (
+            compile_neo_qk_edge_backward,
+            compile_neo_qk_edge_forward,
+        )
+
+        qk_scale = float(self.so2.head_dim**-0.5)
+        with torch.cuda.device(self.device_index):
+            self.qk_edge_forward = compile_neo_qk_edge_forward(
+                qk_scale,
+                self.compile_identity,
+            )
+            self.qk_edge_backward = compile_neo_qk_edge_backward(
+                qk_scale,
+                self.compile_identity,
+            )
+        from .kernels.qk_edge import (
+            compile_neo_qk_node_input_adjoint,
+        )
+
+        with torch.cuda.device(self.device_index):
+            self.qk_node_input_adjoint = compile_neo_qk_node_input_adjoint(
+                float(self.so2.attn_qk_norm.eps),
+                self.compile_identity,
+            )
+        from .phase_c import (
+            CuteNeoPhaseCBackwardLayout,
+        )
+
+        self.phase_c_layout_backward = CuteNeoPhaseCBackwardLayout(
+            focus_eps=self.focus_norm_eps,
+            focus_tau=float(self.so2.focus_softmax_tau),
+            focus_label_smoothing=float(self.so2.focus_label_smoothing),
+            use_focus_norm=self.focus_norm_enabled,
+        )
+        with torch.cuda.device(self.device_index):
+            self.softmax_fwd = compile_envelope_softmax_forward(
+                128,
+                float(self.so2.eps),
+            )
+        self.structural_gate_forward = None
+        self.structural_gate_backward = None
+        self.combined_gate_backward = None
+        self._focus_major_gate_linear_forward = None
+        self._focus_major_gate_linear_backward_add = None
+        self._run_structural_gate_forward = None
+        self._run_structural_gate_backward = None
+        if runtime_config.combined_so2_gate:
+            from .kernels.gate_linear_residual_backward import (
+                compile_neo_gate_linear_residual_backward_fused,
+            )
+
+            with torch.cuda.device(self.device_index):
+                self.combined_gate_backward = (
+                    compile_neo_gate_linear_residual_backward_fused()
+                )
+        elif not runtime_config.native_sm90_path:
+            from .kernels.structural_gate_sm80 import (
+                compile_neo_gate_split_structural_vec4_sm80_backward,
+                compile_neo_gate_split_structural_vec4_sm80_forward,
+            )
+            from .structural_gate import (
+                focus_major_gate_linear_backward_add_,
+                focus_major_gate_linear_forward,
+                run_structural_gate_backward,
+                run_structural_gate_forward,
+            )
+
+            self.structural_gate_forward = (
+                compile_neo_gate_split_structural_vec4_sm80_forward(
+                    self.compile_identity,
+                )
+            )
+            self.structural_gate_backward = (
+                compile_neo_gate_split_structural_vec4_sm80_backward(
+                    self.compile_identity,
+                )
+            )
+            self._focus_major_gate_linear_forward = focus_major_gate_linear_forward
+            self._focus_major_gate_linear_backward_add = (
+                focus_major_gate_linear_backward_add_
+            )
+            self._run_structural_gate_forward = run_structural_gate_forward
+            self._run_structural_gate_backward = run_structural_gate_backward
+        opts = {"device": x.device, "dtype": x.dtype}
+        self._backward_workspace = None
+        self.alpha = torch.empty(
+            self.edge_count, 2, device=opts["device"], dtype=opts["dtype"]
+        )
+        self.group_max = torch.empty(
+            self.node_count, 2, device=opts["device"], dtype=opts["dtype"]
+        )
+        self.denom = torch.empty(
+            self.node_count, 2, device=opts["device"], dtype=opts["dtype"]
+        )
+
+        self._run_cute_phase_a_radial_forward = (
+            run_neo_phase_a_radial_forward_packed_direct
+        )
+        self._run_neo_so2_linear_manual = run_neo_so2_linear_manual
+
+        self.combined_radial = _combined_radial_weight(
+            torch, self.so2.radial_hidden_proj, self.so2.radial_degree_mixer
+        )
+        self.combined_attention_radial = _combined_attention_radial_weight(
+            torch,
+            self.so2.radial_hidden_proj,
+            self.so2.adamw_attn_logit_w,
+        )
+        self.batched_radial_projection_weight = _batched_radial_projection_weight(
+            self.so2.radial_degree_mixer,
+            self.combined_radial,
+            self.combined_attention_radial,
+        )
+
+        self._build_forward_graph()
+
+    def ensure_backward_workspace(
+        self,
+    ) -> NeoSO2BackwardWorkspace:
+        if self._backward_workspace is None:
+            radial_scratch = getattr(self, "structural_scratch", None)
+            if radial_scratch is None:
+                radial_scratch = next(
+                    (
+                        cache.logits
+                        for cache in self.stack_caches
+                        if cache.logits is not None
+                    ),
+                    None,
+                )
+            if radial_scratch is None:
+                raise RuntimeError(
+                    "Neo SO2 backward requires a reusable radial scratch buffer; "
+                    "no stack layer stored gate logits"
+                )
+            if self.phase_c_stack.device.type == "cuda":
+                stream = self.torch.cuda.current_stream(self.phase_c_stack.device)
+                self.phase_c_stack.record_stream(stream)
+                radial_scratch.record_stream(stream)
+            workspace = NeoSO2BackwardWorkspace(
+                self.torch,
+                edge_count=self.edge_count,
+                node_count=self.node_count,
+                d_full=self.d,
+                dt_full=self.dt,
+                radial=self.radial,
+                phase_c_stack=self.phase_c_stack,
+                phase_c_y=None,
+                radial_scratch=radial_scratch,
+                structural_memory_reuse=True,
+                phase_c_single_input_reuse=True,
+            )
+            workspace.attach_to(self)
+            self._backward_workspace = workspace
+        return self._backward_workspace
+
+    def _build_forward_graph(self) -> None:
+        torch = self.torch
+        so2 = self.so2
+        block = self.block
+        n_node = self.node_count
+        n_edge = self.edge_count
+        self.structural_scratch = None
+        if self.config.combined_so2_gate:
+            self.structural_scratch = self.x.new_empty(self.radial.shape)
+
+        use_full_node = block.node_lmax == block.lmax
+        self.use_full_node = use_full_node
+        x_so2 = self.x if use_full_node else self.x[:, : block.mp_ebed_dim, :, :]
+        x_pre = block.pre_so2_norm(x_so2)
+        self.x_wide = so2.pre_focus_mix(
+            x_pre.reshape(n_node, x_so2.shape[1], block.channels).unsqueeze(2)
+        ).squeeze(2)
+
+        cur, rad_l0, radial_compact = self._run_cute_phase_a_radial_forward(
+            radial_hidden_proj=so2.radial_hidden_proj,
+            radial_degree_mixer=so2.radial_degree_mixer,
+            x_wide=self.x_wide.detach(),
+            src=self.src_i32,
+            D_full=self.d.detach(),
+            radial_feat_m0=self.radial.detach(),
+        )
+        self.radial_compact = radial_compact.detach()
+
+        self.focus_gate_src = cur[:, :, 0, :].detach().contiguous()
+        self.stack_caches: list[StackCache] = []
+
+        if self.config.combined_so2_gate:
+            from .kernels.combined_gate_forward import (
+                CuteNeoSO2GateCombinedFwdRunner,
+                prepare_neo_so2_gate_combined_weights,
+            )
+            from .linear import (
+                cached_neo_so2_linear_weights,
+            )
+
+        stack_layers = zip(
+            so2.so2_linears,
+            so2.so2_inter_norms,
+            so2.non_linearities,
+            strict=True,
+        )
+        for layer_idx, (so2_linear, _inter_norm, non_linear) in enumerate(stack_layers):
+            x_layer = cur.detach()
+            final = layer_idx == so2.mixing_layers - 1
+            logits = None
+            if not final and self.config.combined_so2_gate:
+                w0, wpair = cached_neo_so2_linear_weights(so2_linear)
+                gate_parameter = non_linear.gate_linear.weight
+                gate_weight = gate_parameter.detach().view(32, 2, 3 * 32).contiguous()
+                pack_key = (
+                    so2_linear.weight_m0.data_ptr(),
+                    so2_linear.weight_m0._version,
+                    so2_linear.weight_m[0].data_ptr(),
+                    so2_linear.weight_m[0]._version,
+                    gate_parameter.data_ptr(),
+                    gate_parameter._version,
+                    x_layer.device,
+                )
+                pack_cache = getattr(
+                    so2_linear,
+                    "_deepmd_cute_neo_combined_gate_weights",
+                    None,
+                )
+                if (
+                    not isinstance(pack_cache, tuple)
+                    or len(pack_cache) != 6
+                    or pack_cache[0] != pack_key
+                ):
+                    with torch.cuda.device(x_layer.device):
+                        pack_stream = torch.cuda.current_stream(x_layer.device)
+                        packed_weights = prepare_neo_so2_gate_combined_weights(
+                            w0,
+                            wpair,
+                            gate_weight,
+                        )
+                        ready_event = torch.cuda.Event()
+                        ready_event.record(pack_stream)
+                    packed_weights_ready = (
+                        ready_event,
+                        pack_stream.cuda_stream,
+                    )
+                    pack_cache = (
+                        pack_key,
+                        packed_weights,
+                        packed_weights_ready,
+                        # Retain source parameters so allocator pointer reuse
+                        # cannot spoof the versioned cache key.
+                        so2_linear.weight_m0,
+                        so2_linear.weight_m[0],
+                        gate_parameter,
+                    )
+                    so2_linear._deepmd_cute_neo_combined_gate_weights = pack_cache
+                else:
+                    packed_weights = pack_cache[1]
+                    packed_weights_ready = pack_cache[2]
+                y = torch.empty_like(x_layer)
+                # Each CTA reads its residual tile before storing it, and the
+                # m=0 and m>0 regions are disjoint. Reuse x_layer for output to
+                # retain the optimized two-buffer stack footprint.
+                combined_forward = CuteNeoSO2GateCombinedFwdRunner(
+                    x_layer,
+                    x_layer,
+                    y,
+                    x_layer,
+                    packed_weights=packed_weights,
+                    packed_weights_ready=packed_weights_ready,
+                )
+                cur = combined_forward()
+            else:
+                y = self._run_neo_so2_linear_manual(
+                    so2_linear,
+                    x_layer,
+                    add_residual=final,
+                    per_focus_pair=self.config.per_focus_so2_fwd_pair,
+                )
+            if not final and not self.config.combined_so2_gate:
+                gate_src = y[:, :, 0, :]
+                logits = self._focus_major_gate_linear_forward(
+                    gate_src,
+                    non_linear.gate_linear.weight.detach(),
+                )
+                self._run_structural_gate_forward(
+                    self.structural_gate_forward,
+                    x_layer,
+                    y,
+                    logits,
+                    out=x_layer,
+                )
+                cur = x_layer
+            elif final:
+                self.phase_c_stack = y.detach()
+                self.phase_c_y = None
+                cur = None
+            self.stack_caches.append(
+                StackCache(
+                    y=y,
+                    logits=logits,
+                    non_linear=non_linear,
+                    final=final,
+                )
+            )
+
+        x_wide_qk = self.x_wide.detach()
+        rad_l0_qk = rad_l0.detach().view(n_edge, 2, 32)
+        x_l0_node = x_wide_qk[:, 0, :].reshape(n_node, 2, 32)
+        focus_alpha = torch.empty(
+            n_edge,
+            2,
+            device=self.focus_gate_src.device,
+            dtype=torch.float32,
+        )
+        q_node = torch.empty_like(
+            x_l0_node,
+            memory_format=torch.contiguous_format,
+        )
+        k_node = torch.empty_like(
+            x_l0_node,
+            memory_format=torch.contiguous_format,
+        )
+        self.attention_prelude_forward(
+            self.focus_gate_src.view(n_edge, 64),
+            x_l0_node.contiguous(),
+            so2.adamw_focus_compete_w.detach().float().contiguous(),
+            self.focus_norm_scale,
+            so2.attn_q_proj.weight.detach().float().view(32, 2, 32).contiguous(),
+            so2.attn_k_proj.weight.detach().float().view(32, 2, 32).contiguous(),
+            so2.attn_qk_norm.adam_scale.detach().float().contiguous(),
+            focus_alpha,
+            q_node,
+            k_node,
+        )
+        self.focus_alpha = focus_alpha.detach()
+        self.q_node = q_node.detach()
+        self.k_node = k_node.detach()
+        self.attn_logits = torch.empty(
+            n_edge,
+            2,
+            device=self.q_node.device,
+            dtype=self.q_node.dtype,
+        )
+        self.qk_edge_forward(
+            self.q_node,
+            self.k_node,
+            rad_l0_qk.contiguous(),
+            so2.adamw_attn_logit_w.detach().contiguous(),
+            self.src_i32,
+            self.dst_i32,
+            self.attn_logits,
+        )
+
+        self.softmax_fwd(
+            self.attn_logits.detach().contiguous(),
+            self.edge_gate,
+            self.dst_ptr_i32,
+            so2.adamw_attn_z_bias_raw.detach().reshape(2).float().contiguous(),
+            self.alpha,
+            self.group_max,
+            self.denom,
+        )
+        self.attn_logits = None
+        x_wide_down = self.x_wide.detach()
+        from .kernels.phase_c_forward import (
+            run_neo_phase_c_onepass_output_gate,
+        )
+
+        self.phase_c_out = run_neo_phase_c_onepass_output_gate(
+            x_local_flat=self.phase_c_stack,
+            Dt_full=self.dt.detach(),
+            alpha_focus=self.alpha,
+            focus_compete_alpha=self.focus_alpha,
+            dst_ptr=self.dst_ptr_i32,
+            rotate_inv_rescale=so2.rotate_inv_rescale_full,
+            x_wide=x_wide_down,
+            output_gate_norm_scale=so2.attn_output_gate_norm.adam_scale.detach()
+            .float()
+            .reshape(2, 32)
+            .contiguous(),
+            output_gate_weight=so2.adamw_attn_gate_w.detach()
+            .float()
+            .reshape(32, 2, 1)
+            .contiguous(),
+            output_gate_eps=float(so2.attn_output_gate_norm.eps),
+        ).to(dtype=so2.compute_dtype)
+        out = self.phase_c_out.detach().to(dtype=so2.dtype)
+        self.out_gate_flat = out.detach()
+        self.message_grid_product = None
+        if so2.message_node_grid_product is not None:
+            if self.packed_message_grid:
+                from .message_grid import (
+                    run_packed_message_grid_forward,
+                )
+
+                grid_out = run_packed_message_grid_forward(
+                    so2.message_node_grid_product,
+                    out,
+                    x_wide_down,
+                )
+            else:
+                grid_out = so2.message_node_grid_product(out, x_wide_down)
+            out = out + grid_out
+        self.post_mix_input = out.detach()
+        out = so2.post_focus_mix(out.unsqueeze(2)).squeeze(2)
+        self.post_norm_input = out.unsqueeze(2).detach()
+        so2_out = block.post_so2_norm(self.post_norm_input)
+        if use_full_node:
+            self.final = so2_out
+        else:
+            final = self.x.new_zeros(self.x.shape)
+            final[:, : block.mp_ebed_dim, :, :] = so2_out
+            self.final = final
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/__init__.py
new file mode 100644
index 0000000000..74481d7e7e
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Native SM90 split-complex Neo SO(2) implementation."""
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/final_phase_c.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/final_phase_c.py
new file mode 100644
index 0000000000..2d432b1294
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/final_phase_c.py
@@ -0,0 +1,486 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Direct destination statistics for the SM90 final SO2Linear/Phase-C boundary.
+
+One CTA owns one ``(node, focus)`` segment.  Its feature-owning threads keep
+all output-degree statistics in registers while walking the destination's
+edges in 64-edge chunks.  The CTA writes the final node-scale ``a0`` and
+``a1`` statistics directly, eliminating the global chunk partials and their
+second reduction launch.
+
+The strict-FP32 node GEMMs are applied only after edge values have been
+reduced to node-scale sufficient statistics.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+from ..wigner_layout import (
+    PACKED_VALUE_COUNT,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN202, ANN204, TC002
+
+EDGE_CHUNK = 64
+FOCUS_COUNT = 2
+CHANNELS = 32
+DEGREE_COUNT = 16
+M0_WIDTH = 128
+M1_WIDTH = 96
+PACKED_WIGNER_VALUES = PACKED_VALUE_COUNT
+M0_THREADS = M0_WIDTH
+M1_THREADS = M1_WIDTH * 2
+THREADS = M0_THREADS + M1_THREADS
+M1_SCALE_VALUES = EDGE_CHUNK * (DEGREE_COUNT - 1) * 2
+M0_SCALE_VALUES = EDGE_CHUNK * DEGREE_COUNT
+TOTAL_SCALE_VALUES = M0_SCALE_VALUES + M1_SCALE_VALUES
+LOADS_PER_THREAD = (TOTAL_SCALE_VALUES + THREADS - 1) // THREADS
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+
+@dataclass(frozen=True)
+class ExpandedFinalWeights:
+    """Dense final SO2Linear blocks selected for each output degree."""
+
+    w0: torch.Tensor
+    wc: torch.Tensor
+
+
+@dataclass(frozen=True)
+class ExpandedComplexWorkspace:
+    """Node-scale real and complex sufficient statistics."""
+
+    m0: torch.Tensor
+    m1: torch.Tensor
+
+    @property
+    def storage_bytes(self) -> int:
+        return sum(
+            tensor.numel() * tensor.element_size() for tensor in (self.m0, self.m1)
+        )
+
+
+def prepare_expanded_final_weights(
+    w0: torch.Tensor,
+    wc: torch.Tensor,
+) -> ExpandedFinalWeights:
+    """Select the dense input block needed by each full output degree."""
+    if tuple(w0.shape) != (FOCUS_COUNT, M0_WIDTH, M0_WIDTH):
+        raise ValueError("w0 must have shape (2,128,128)")
+    if tuple(wc.shape) != (FOCUS_COUNT, M1_WIDTH, M1_WIDTH):
+        raise ValueError("wc must have shape (2,96,96)")
+    if w0.dtype != torch.float32 or wc.dtype != torch.complex64:
+        raise TypeError("w0/wc must be float32/complex64")
+    degree_by_q = (0, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3)
+    blocks0 = torch.stack(
+        [
+            w0[:, :, degree * CHANNELS : (degree + 1) * CHANNELS]
+            for degree in degree_by_q
+        ],
+        dim=1,
+    ).contiguous()
+    blocks1 = torch.stack(
+        [
+            wc[:, :, (degree - 1) * CHANNELS : degree * CHANNELS]
+            for degree in degree_by_q[1:]
+        ],
+        dim=1,
+    ).contiguous()
+    return ExpandedFinalWeights(w0=blocks0, wc=blocks1)
+
+
+def _require_sm90_strict_fp32(device: torch.device) -> None:
+    if device.type != "cuda":
+        raise ValueError("SM90 final Phase C requires CUDA")
+    if tuple(torch.cuda.get_device_capability(device)) != (9, 0):
+        raise RuntimeError("SM90 final Phase C requires compute capability 9.0")
+    if torch.backends.cuda.matmul.allow_tf32:
+        raise RuntimeError("strict FP32 requires allow_tf32=False")
+    if torch.get_float32_matmul_precision() != "highest":
+        raise RuntimeError("strict FP32 requires float32 matmul precision 'highest'")
+
+
+@dataclass(frozen=True)
+class DirectStatisticsForwardResult:
+    """Forward output and the final node-scale sufficient statistics."""
+
+    output: torch.Tensor
+    workspace: ExpandedComplexWorkspace
+
+    @property
+    def statistics_storage_bytes(self) -> int:
+        return self.workspace.storage_bytes
+
+
+@cute.jit
+def _packed_index_runtime(q, row_slot):
+    base = cutlass.Int32(0)
+    width = cutlass.Int32(1)
+    local = cutlass.Int32(0)
+    if q >= 9:
+        base = cutlass.Int32(25)
+        width = cutlass.Int32(7)
+        local = q - 9
+    elif q >= 4:
+        base = cutlass.Int32(10)
+        width = cutlass.Int32(5)
+        local = q - 4
+    elif q >= 1:
+        base = cutlass.Int32(1)
+        width = cutlass.Int32(3)
+        local = q - 1
+    return base + row_slot * width + local
+
+
+class CuteDirectNodeStatistics:
+    """Accumulate all real and complex statistics in one destination CTA."""
+
+    @cute.jit
+    def __call__(
+        self,
+        m0: cute.Tensor,
+        m1_ri: cute.Tensor,
+        dt_packed: cute.Tensor,
+        beta: cute.Tensor,
+        dst_ptr: cute.Tensor,
+        a0: cute.Tensor,
+        a1_ri: cute.Tensor,
+        stream: CUstream,
+    ):
+        self.kernel(m0, m1_ri, dt_packed, beta, dst_ptr, a0, a1_ri).launch(
+            grid=[a0.shape[2], FOCUS_COUNT, 1],
+            block=[THREADS, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        m0: cute.Tensor,
+        m1_ri: cute.Tensor,
+        dt_packed: cute.Tensor,
+        beta: cute.Tensor,
+        dst_ptr: cute.Tensor,
+        a0: cute.Tensor,
+        a1_ri: cute.Tensor,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        node, focus, _ = cute.arch.block_idx()
+        node_lo = dst_ptr[node]
+        node_hi = dst_ptr[node + 1]
+        edge_count = node_hi - node_lo
+        chunk_count = (edge_count + EDGE_CHUNK - 1) // EDGE_CHUNK
+
+        smem = cutlass.utils.SmemAllocator()
+        m0_scale_storage = smem.allocate_tensor(cutlass.Float32, M0_SCALE_VALUES)
+        m0_scales = cute.make_tensor(
+            m0_scale_storage.iterator,
+            cute.make_layout(
+                (EDGE_CHUNK, DEGREE_COUNT),
+                stride=(DEGREE_COUNT, 1),
+            ),
+        )
+        m1_scale_storage = smem.allocate_tensor(cutlass.Float32, M1_SCALE_VALUES)
+        m1_scales = cute.make_tensor(
+            m1_scale_storage.iterator,
+            cute.make_layout(
+                (EDGE_CHUNK, DEGREE_COUNT - 1, 2),
+                stride=((DEGREE_COUNT - 1) * 2, 2, 1),
+            ),
+        )
+
+        # A single 16-value register fragment serves either a real feature or
+        # one component of a complex feature.  Complex threads use entries
+        # [0, 15), avoiding two live accumulator arrays in generated code.
+        accumulators = cute.make_rmem_tensor(
+            cute.make_layout((DEGREE_COUNT,), stride=(1,)),
+            cutlass.Float32,
+        )
+        accumulators.fill(0.0)
+
+        for chunk_slot in cutlass.range(chunk_count, unroll=1):
+            lo = node_lo + chunk_slot * EDGE_CHUNK
+            hi = lo + EDGE_CHUNK
+            if node_hi < hi:
+                hi = node_hi
+
+            for load_slot in cutlass.range_constexpr(LOADS_PER_THREAD):
+                linear = tidx + load_slot * THREADS
+                if linear < M0_SCALE_VALUES:
+                    edge_slot = linear // DEGREE_COUNT
+                    q = linear - edge_slot * DEGREE_COUNT
+                    edge = lo + edge_slot
+                    value = cutlass.Float32(0.0)
+                    if edge < hi:
+                        panel = _packed_index_runtime(q, cutlass.Int32(0))
+                        value = beta[edge, focus].to(cutlass.Float32) * dt_packed[
+                            edge, panel
+                        ].to(cutlass.Float32)
+                    m0_scales[edge_slot, q] = value
+                elif linear < TOTAL_SCALE_VALUES:
+                    item = linear - M0_SCALE_VALUES
+                    edge_slot = item // ((DEGREE_COUNT - 1) * 2)
+                    remainder = item - edge_slot * (DEGREE_COUNT - 1) * 2
+                    q1 = remainder // 2
+                    component = remainder - q1 * 2
+                    edge = lo + edge_slot
+                    value = cutlass.Float32(0.0)
+                    if edge < hi:
+                        panel = _packed_index_runtime(q1 + 1, component + 1)
+                        value = beta[edge, focus].to(cutlass.Float32) * dt_packed[
+                            edge, panel
+                        ].to(cutlass.Float32)
+                    m1_scales[edge_slot, q1, component] = value
+            cute.arch.sync_threads()
+
+            if tidx < M0_THREADS:
+                feature = tidx
+                for edge_slot in cutlass.range_constexpr(EDGE_CHUNK):
+                    edge = lo + edge_slot
+                    if edge < hi:
+                        x = m0[focus, edge, feature].to(cutlass.Float32)
+                        for q in cutlass.range_constexpr(DEGREE_COUNT):
+                            value = accumulators[q].to(cutlass.Float32)
+                            value += m0_scales[edge_slot, q] * x
+                            accumulators[q] = value
+            elif tidx < THREADS:
+                complex_thread = tidx - M0_THREADS
+                feature = complex_thread // 2
+                component = complex_thread - feature * 2
+                for edge_slot in cutlass.range_constexpr(EDGE_CHUNK):
+                    edge = lo + edge_slot
+                    if edge < hi:
+                        xr = m1_ri[focus, edge, feature, 0].to(cutlass.Float32)
+                        xi = m1_ri[focus, edge, feature, 1].to(cutlass.Float32)
+                        for q1 in cutlass.range_constexpr(DEGREE_COUNT - 1):
+                            dr = m1_scales[edge_slot, q1, 0]
+                            di = m1_scales[edge_slot, q1, 1]
+                            value = accumulators[q1].to(cutlass.Float32)
+                            if component == 0:
+                                value += dr * xr + di * xi
+                            else:
+                                value += dr * xi - di * xr
+                            accumulators[q1] = value
+
+            # Every thread must finish reading shared scales before the next
+            # 64-edge chunk overwrites them.
+            cute.arch.sync_threads()
+
+        if tidx < M0_THREADS:
+            feature = tidx
+            for q in cutlass.range_constexpr(DEGREE_COUNT):
+                a0[focus, q, node, feature] = accumulators[q]
+        elif tidx < THREADS:
+            complex_thread = tidx - M0_THREADS
+            feature = complex_thread // 2
+            component = complex_thread - feature * 2
+            for q1 in cutlass.range_constexpr(DEGREE_COUNT - 1):
+                a1_ri[focus, q1, node, feature, component] = accumulators[q1]
+
+
+def _fake_m0_edges():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, cute.sym_int64(), M0_WIDTH),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_m1_edges_ri():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, cute.sym_int64(), M1_WIDTH, 2),
+        stride_order=(3, 2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_dt():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (cute.sym_int64(), PACKED_WIGNER_VALUES),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_beta():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (cute.sym_int64(), FOCUS_COUNT),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_index():
+    return make_fake_compact_tensor(
+        cutlass.Int32,
+        (cute.sym_int64(),),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_a0():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, DEGREE_COUNT, cute.sym_int64(), M0_WIDTH),
+        stride_order=(3, 2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_a1_ri():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, DEGREE_COUNT - 1, cute.sym_int64(), M1_WIDTH, 2),
+        stride_order=(4, 3, 2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+@device_aware_lru_cache(maxsize=2)
+def _compiled_direct_statistics() -> Callable:
+    return cute.compile(
+        CuteDirectNodeStatistics(),
+        _fake_m0_edges(),
+        _fake_m1_edges_ri(),
+        _fake_dt(),
+        _fake_beta(),
+        _fake_index(),
+        _fake_a0(),
+        _fake_a1_ri(),
+        stream=make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _validate_forward_inputs(
+    m0: torch.Tensor,
+    m1: torch.Tensor,
+    dt_packed: torch.Tensor,
+    beta: torch.Tensor,
+    dst_ptr: torch.Tensor,
+) -> tuple[torch.device, int]:
+    device = m0.device
+    _require_sm90_strict_fp32(device)
+    edge_count = int(m0.shape[1])
+    node_count = int(dst_ptr.numel() - 1)
+    expected = (
+        ("m0", m0, (FOCUS_COUNT, edge_count, M0_WIDTH), torch.float32),
+        ("m1", m1, (FOCUS_COUNT, edge_count, M1_WIDTH), torch.complex64),
+        (
+            "dt_packed",
+            dt_packed,
+            (edge_count, PACKED_WIGNER_VALUES),
+            torch.float32,
+        ),
+        ("beta", beta, (edge_count, FOCUS_COUNT), torch.float32),
+        ("dst_ptr", dst_ptr, (node_count + 1,), torch.int32),
+    )
+    for name, tensor, shape, dtype in expected:
+        if tuple(tensor.shape) != shape or tensor.dtype != dtype:
+            raise ValueError(f"{name} must have shape {shape} and dtype {dtype}")
+        if tensor.device != device or not tensor.is_contiguous():
+            raise ValueError(f"{name} must be contiguous on {device}")
+        if tensor.data_ptr() % 16:
+            raise ValueError(f"{name} must be at least 16-byte aligned")
+    return device, node_count
+
+
+def run_direct_statistics_forward(
+    *,
+    m0: torch.Tensor,
+    m1: torch.Tensor,
+    dt_packed: torch.Tensor,
+    beta: torch.Tensor,
+    dst_ptr: torch.Tensor,
+    weights: ExpandedFinalWeights,
+    edge_chunk: int = EDGE_CHUNK,
+    chunk_slots: int | None = None,
+) -> DirectStatisticsForwardResult:
+    """Build final node statistics directly and apply node-scale SO2Linear.
+
+    ``edge_chunk`` and ``chunk_slots`` retain the chunked-forward argument
+    contract used by the SM90 SO2 runner.
+    Only the 64-edge internal schedule is supported; ``chunk_slots`` is not an
+    allocation dimension in this implementation.
+    """
+    if edge_chunk != EDGE_CHUNK:
+        raise ValueError(f"direct statistics requires edge_chunk={EDGE_CHUNK}")
+    if chunk_slots is not None and chunk_slots <= 0:
+        raise ValueError("chunk_slots must be positive when provided")
+    device, node_count = _validate_forward_inputs(m0, m1, dt_packed, beta, dst_ptr)
+    a0 = torch.empty(
+        (FOCUS_COUNT, DEGREE_COUNT, node_count, M0_WIDTH),
+        device=device,
+        dtype=torch.float32,
+    )
+    a1 = torch.empty(
+        (FOCUS_COUNT, DEGREE_COUNT - 1, node_count, M1_WIDTH),
+        device=device,
+        dtype=torch.complex64,
+    )
+    with torch.cuda.device(device):
+        _compiled_direct_statistics()(
+            m0,
+            torch.view_as_real(m1),
+            dt_packed,
+            beta,
+            dst_ptr,
+            a0,
+            torch.view_as_real(a1),
+        )
+        out0 = torch.bmm(a0.flatten(0, 1), weights.w0.flatten(0, 1)).view(
+            FOCUS_COUNT, DEGREE_COUNT, node_count, CHANNELS
+        )
+        out1 = torch.bmm(a1.flatten(0, 1), weights.wc.flatten(0, 1)).view(
+            FOCUS_COUNT, DEGREE_COUNT - 1, node_count, CHANNELS
+        )
+        output = out0.permute(2, 0, 1, 3).contiguous()
+        output[:, :, 1:] += out1.real.permute(2, 0, 1, 3)
+    return DirectStatisticsForwardResult(
+        output=output,
+        workspace=ExpandedComplexWorkspace(m0=a0, m1=a1),
+    )
+
+
+__all__ = [
+    "EDGE_CHUNK",
+    "THREADS",
+    "DirectStatisticsForwardResult",
+    "ExpandedFinalWeights",
+    "prepare_expanded_final_weights",
+    "run_direct_statistics_forward",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_gaunt.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_gaunt.py
new file mode 100644
index 0000000000..b4ab9da8c2
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_gaunt.py
@@ -0,0 +1,323 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Exact normalized-Gaunt message-grid product for the Neo SM90 path."""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as cute_utils
+import torch
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN202, ANN204, TC002
+
+COEFF_DIM = 48
+CHANNELS = 64
+THREADS = 256
+GROUPS = 4
+VALUES_PER_NODE = COEFF_DIM * CHANNELS
+EXPECTED_COMPACT_PATHS = 1968
+EXPECTED_ORDERED_PATHS = 3833
+SUPPORT_GAP_THRESHOLD = 1.0e-4
+SM90_CAPABILITY = (9, 0)
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+# (left row, right row, C[p,i,j], d[i]G[p,i,j], d[j]G[p,i,j]).
+# The last value is zero on a diagonal path.
+GauntTerm = tuple[int, int, float, float, float]
+
+
+@dataclass(frozen=True, eq=False)
+class Sm90GauntSchedule:
+    """Normalized-Gaunt rows captured as CuTe compile-time constants."""
+
+    rows: tuple[tuple[GauntTerm, ...], ...]
+    output_groups: tuple[tuple[int, ...], ...]
+
+    def __post_init__(self) -> None:
+        if len(self.rows) != COEFF_DIM:
+            raise ValueError("Gaunt schedule must contain 48 rows")
+        if len(self.output_groups) != GROUPS:
+            raise ValueError("Gaunt schedule must contain four groups")
+        assigned = tuple(row for group in self.output_groups for row in group)
+        if tuple(sorted(assigned)) != tuple(range(COEFF_DIM)):
+            raise ValueError("each coefficient row must occur in exactly one group")
+        if sum(len(row) for row in self.rows) != EXPECTED_COMPACT_PATHS:
+            raise ValueError("Gaunt schedule must contain 1,968 paths")
+        for output_row, row in enumerate(self.rows):
+            for left_row, right_row, forward, adj_left, adj_right in row:
+                if not (0 <= left_row <= right_row < COEFF_DIM):
+                    raise ValueError(
+                        f"invalid compact path in row {output_row}: "
+                        f"({left_row}, {right_row})"
+                    )
+                if forward == 0.0 or adj_left == 0.0:
+                    raise ValueError("Gaunt paths must not contain zero weights")
+                if (left_row == right_row) != (adj_right == 0.0):
+                    raise ValueError(
+                        "only diagonal Gaunt paths may omit the mirrored adjoint"
+                    )
+
+
+def build_sm90_gaunt_schedule(
+    to_grid: torch.Tensor,
+    from_grid: torch.Tensor,
+) -> Sm90GauntSchedule:
+    """Recover and certify the fixed Neo Gaunt support from its projectors."""
+    if (
+        tuple(to_grid.shape) != (152, COEFF_DIM)
+        or tuple(from_grid.shape) != (COEFF_DIM, 152)
+        or to_grid.dtype != torch.float32
+        or from_grid.dtype != torch.float32
+    ):
+        raise ValueError("Neo Gaunt requires FP32 (152,48)/(48,152) projectors")
+
+    to_cpu = to_grid.detach().to(device="cpu", dtype=torch.float64)
+    from_cpu = from_grid.detach().to(device="cpu", dtype=torch.float64)
+    tensor = torch.einsum("pg,gi,gj->pij", from_cpu, to_cpu, to_cpu)
+    support = tensor.abs() > SUPPORT_GAP_THRESHOLD
+    ordered_paths = int(support.sum())
+    if ordered_paths != EXPECTED_ORDERED_PATHS:
+        raise ValueError(
+            "Neo projector support changed: expected "
+            f"{EXPECTED_ORDERED_PATHS} paths, got {ordered_paths}"
+        )
+    minimum_signal = float(tensor.abs()[support].min())
+    maximum_residual = float(tensor.abs()[~support].max())
+    if minimum_signal <= 10.0 * SUPPORT_GAP_THRESHOLD:
+        raise ValueError("Neo Gaunt support no longer has a certified magnitude gap")
+    if maximum_residual >= SUPPORT_GAP_THRESHOLD:
+        raise ValueError("Neo Gaunt structural-zero residual exceeds its certificate")
+
+    degree_weight = tuple(
+        2 * degree + 1
+        for degree in range(4)
+        for _order in range(-degree, degree + 1)
+        for _frame in range(3)
+    )
+    weights = torch.tensor(degree_weight, device="cpu", dtype=torch.float64)
+    normalized = tensor / weights[:, None, None]
+    permutation_error = max(
+        float((normalized - normalized.permute(permutation)).abs().max())
+        for permutation in (
+            (0, 1, 2),
+            (0, 2, 1),
+            (1, 0, 2),
+            (1, 2, 0),
+            (2, 0, 1),
+            (2, 1, 0),
+        )
+    )
+    if permutation_error > 2.0e-6:
+        raise ValueError(
+            f"Neo normalized-Gaunt symmetry changed: max error {permutation_error:.3e}"
+        )
+
+    rows: list[tuple[GauntTerm, ...]] = []
+    for output_row in range(COEFF_DIM):
+        terms: list[GauntTerm] = []
+        for left_row in range(COEFF_DIM):
+            for right_row in range(left_row, COEFF_DIM):
+                if not bool(support[output_row, left_row, right_row]):
+                    continue
+                forward = float(tensor[output_row, left_row, right_row].float())
+                symmetric = normalized[output_row, left_row, right_row]
+                adj_left = float((symmetric * degree_weight[left_row]).float())
+                adj_right = (
+                    0.0
+                    if left_row == right_row
+                    else float((symmetric * degree_weight[right_row]).float())
+                )
+                terms.append((left_row, right_row, forward, adj_left, adj_right))
+        rows.append(tuple(terms))
+
+    group_rows: list[list[int]] = [[] for _ in range(GROUPS)]
+    group_loads = [0] * GROUPS
+    for output_row in sorted(
+        range(COEFF_DIM),
+        key=lambda row: (-len(rows[row]), row),
+    ):
+        group = min(range(GROUPS), key=lambda item: (group_loads[item], item))
+        group_rows[group].append(output_row)
+        group_loads[group] += len(rows[output_row])
+    if max(group_loads) - min(group_loads) > 32:
+        raise ValueError(f"Neo Gaunt static groups are imbalanced: {group_loads}")
+    return Sm90GauntSchedule(
+        tuple(rows),
+        tuple(tuple(sorted(group)) for group in group_rows),
+    )
+
+
+class _Sm90GauntForward:
+    def __init__(self, schedule: Sm90GauntSchedule) -> None:
+        self.schedule = schedule
+
+    @cute.jit
+    def __call__(
+        self,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        output: cute.Tensor,
+        stream: CUstream,
+    ):
+        operand_layout = cute.make_layout(
+            (2, COEFF_DIM, CHANNELS),
+            stride=(VALUES_PER_NODE, CHANNELS, 1),
+        )
+        self.kernel(left, right, output, operand_layout).launch(
+            grid=(left.shape[0], 1, 1),
+            block=[THREADS, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        output: cute.Tensor,
+        operand_layout: cute.Layout,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        node, _, _ = cute.arch.block_idx()
+        group = tidx >> 6
+        channel = tidx & (CHANNELS - 1)
+
+        smem = cute_utils.SmemAllocator()
+        operands = smem.allocate_tensor(cutlass.Float32, operand_layout, 16)
+        for side in cutlass.range_constexpr(2):
+            for slot in cutlass.range_constexpr(VALUES_PER_NODE // THREADS):
+                linear = tidx + slot * THREADS
+                row = linear >> 6
+                scalar_channel = linear & (CHANNELS - 1)
+                if cutlass.const_expr(side == 1):
+                    value = right[node, row, scalar_channel].to(cutlass.Float32)
+                else:
+                    value = left[node, row, scalar_channel].to(cutlass.Float32)
+                operands[side, row, scalar_channel] = value
+        cute.arch.sync_threads()
+
+        if group == 0:
+            self._accumulate_group(operands, output, node, channel, 0)
+        elif group == 1:
+            self._accumulate_group(operands, output, node, channel, 1)
+        elif group == 2:
+            self._accumulate_group(operands, output, node, channel, 2)
+        else:
+            self._accumulate_group(operands, output, node, channel, 3)
+
+    @cute.jit
+    def _accumulate_group(
+        self,
+        operands: cute.Tensor,
+        output: cute.Tensor,
+        node: cutlass.Int32,
+        channel: cutlass.Int32,
+        group: cutlass.Constexpr,
+    ):
+        rows = self.schedule.output_groups[group]
+        for row_slot in cutlass.range_constexpr(len(rows)):
+            output_row = rows[row_slot]
+            accumulator = cutlass.Float32(0.0)
+            terms = self.schedule.rows[output_row]
+            for path in cutlass.range_constexpr(len(terms)):
+                left_row, right_row, coefficient_value, _, _ = terms[path]
+                product = (
+                    operands[0, left_row, channel] * operands[1, right_row, channel]
+                )
+                if cutlass.const_expr(left_row != right_row):
+                    product = product + (
+                        operands[0, right_row, channel] * operands[1, left_row, channel]
+                    )
+                accumulator = accumulator + cutlass.Float32(coefficient_value) * product
+            output[node, output_row, channel] = accumulator
+
+
+def _fake_dense() -> cute.Tensor:
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (cute.sym_int64(), COEFF_DIM, CHANNELS),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+@device_aware_lru_cache(maxsize=8)
+def _compiled_sm90_gaunt_forward(schedule: Sm90GauntSchedule) -> Callable:
+    dense = _fake_dense()
+    return cute.compile(
+        _Sm90GauntForward(schedule),
+        dense,
+        dense,
+        dense,
+        make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _validate_dense(*tensors: torch.Tensor) -> None:
+    reference = tensors[0]
+    if (
+        tuple(reference.shape[1:]) != (COEFF_DIM, CHANNELS)
+        or reference.shape[0] <= 0
+        or not reference.is_cuda
+        or reference.dtype != torch.float32
+        or tuple(torch.cuda.get_device_capability(reference.device)) != SM90_CAPABILITY
+    ):
+        raise ValueError("Gaunt operands must be SM90 FP32 (N,48,64)")
+    if any(
+        tensor.shape != reference.shape
+        or tensor.device != reference.device
+        or tensor.dtype != torch.float32
+        or not tensor.is_contiguous()
+        or tensor.data_ptr() % 16 != 0
+        for tensor in tensors
+    ):
+        raise ValueError("Gaunt operands must match and be 16-byte aligned")
+
+
+def run_sm90_gaunt_forward(
+    left: torch.Tensor,
+    right: torch.Tensor,
+    schedule: Sm90GauntSchedule,
+) -> torch.Tensor:
+    """Evaluate the exact normalized-Gaunt product on SM90."""
+    output = torch.empty_like(left)
+    _validate_dense(left, right, output)
+    with torch.cuda.device(left.device):
+        _compiled_sm90_gaunt_forward(schedule)(left, right, output)
+    return output
+
+
+__all__ = [
+    "Sm90GauntSchedule",
+    "build_sm90_gaunt_schedule",
+    "run_sm90_gaunt_forward",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_readout.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_readout.py
new file mode 100644
index 0000000000..abf3cb2774
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_readout.py
@@ -0,0 +1,849 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Fused strict-FP32 Neo message-grid input adjoint for SM90."""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as cute_utils
+import torch
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+from ...runtime_policy import (
+    SM90_CAPABILITY,
+    uses_strict_fp32_matmul,
+)
+from .message_grid_gaunt import (
+    Sm90GauntSchedule,
+    build_sm90_gaunt_schedule,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN202, ANN204, TC002
+
+DEGREE_COUNT = 16
+FRAME_COUNT = 3
+FOCUS_COUNT = 2
+CHANNELS = 32
+FOLDED_CHANNELS = FOCUS_COUNT * CHANNELS
+PACKED_COEFF_DIM = DEGREE_COUNT * FRAME_COUNT
+
+FRAME_CONTRACT_THREADS = 128
+FRAME_CONTRACT_M_TILE = 128
+FRAME_CONTRACT_N_TILE = 64
+FRAME_CONTRACT_K_TILE = 16
+FRAME_CONTRACT_N_TILES = 2
+FRAME_CONTRACT_WIDTH = FRAME_COUNT * CHANNELS
+FRAME_CONTRACT_SMEM_PADDING = 4
+
+READOUT_THREADS = 256
+READOUT_GROUPS = 4
+VALUES_PER_PANEL = PACKED_COEFF_DIM * FOLDED_CHANNELS
+ROWS_PER_GROUP = PACKED_COEFF_DIM // READOUT_GROUPS
+DEGREE_CHANNEL_VALUES_PER_NODE = DEGREE_COUNT * CHANNELS
+DEGREES_PER_THREAD = DEGREE_CHANNEL_VALUES_PER_NODE // READOUT_THREADS
+LEFT_PANEL = 0
+RIGHT_PANEL = 1
+GATED_COEFFICIENT_PANEL = 2
+GRAD_LEFT_PANEL = 3
+GRAD_RIGHT_PANEL = 4
+WORKSPACE_PANELS = 5
+
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+
+@dataclass(frozen=True)
+class Sm90MessageGridState:
+    """Certified Gaunt schedule and packed immutable readout weights."""
+
+    schedule: Sm90GauntSchedule
+    frame_contract: torch.Tensor
+    residual_scale: torch.Tensor
+    frame_expand_t: torch.Tensor
+
+
+def _tensor_cache_key(tensor: torch.Tensor | None) -> tuple[Any, ...] | None:
+    if tensor is None:
+        return None
+    return (
+        tensor.data_ptr(),
+        tensor._version,
+        tensor.dtype,
+        tensor.device,
+        tuple(tensor.shape),
+        tuple(tensor.stride()),
+    )
+
+
+def prepare_sm90_message_grid_state(net: Any) -> Sm90MessageGridState:
+    """Prepare and cache the fixed Neo readout contract outside the hot path."""
+    projector = net.projector
+    key = (
+        _tensor_cache_key(projector.to_grid_mat),
+        _tensor_cache_key(projector.from_grid_mat),
+        _tensor_cache_key(net.frame_contract.weight),
+        _tensor_cache_key(net.frame_contract.degree_index),
+        _tensor_cache_key(net.frame_expand.weight),
+        _tensor_cache_key(net.frame_expand.degree_index),
+        _tensor_cache_key(net.residual_scale),
+    )
+    cached = getattr(net, "_deepmd_cute_sm90_message_grid_state", None)
+    if cached is not None and cached[0] == key:
+        return cached[1]
+
+    schedule = build_sm90_gaunt_schedule(
+        projector.to_grid_mat,
+        projector.from_grid_mat,
+    )
+    frame_contract = net.frame_contract.weight.index_select(
+        0,
+        net.frame_contract.degree_index,
+    ).view(DEGREE_COUNT, FRAME_COUNT, CHANNELS, CHANNELS)
+    frame_expand = net.frame_expand.weight.index_select(
+        0,
+        net.frame_expand.degree_index,
+    ).view(DEGREE_COUNT, CHANNELS, FRAME_COUNT, CHANNELS)
+    residual_scale = net.residual_scale
+    if residual_scale is None:
+        residual_scale = torch.ones(
+            (FOCUS_COUNT, CHANNELS),
+            device=frame_contract.device,
+            dtype=torch.float32,
+        )
+    state = Sm90MessageGridState(
+        schedule=schedule,
+        frame_contract=frame_contract.contiguous(),
+        residual_scale=residual_scale.view(FOCUS_COUNT, CHANNELS).contiguous(),
+        frame_expand_t=frame_expand.permute(0, 2, 3, 1).contiguous(),
+    )
+    net._deepmd_cute_sm90_message_grid_state = (key, state)
+    return state
+
+
+@cute.jit
+def _make_frame_contract_mma():
+    atoms_layout = cute.make_layout(
+        (FRAME_CONTRACT_THREADS // 16, 16, 1),
+        stride=(16, 1, 0),
+    )
+    permutation_m = cute.make_layout(
+        (atoms_layout.shape[0], 4),
+        stride=(4, 1),
+    )
+    permutation_n = cute.make_layout(
+        (atoms_layout.shape[1], 4),
+        stride=(4, 1),
+    )
+    return cute.make_tiled_mma(
+        cute.nvgpu.MmaUniversalOp(cutlass.Float32),
+        atoms_layout,
+        permutation_mnk=(permutation_m, permutation_n, None),
+    )
+
+
+class _FrameContractAdjoint:
+    """Batched ``N x 32 @ 32 x 96`` strict-FP32 FrameContract adjoint."""
+
+    @cute.jit
+    def __call__(
+        self,
+        grad_out: cute.Tensor,
+        frame_contract: cute.Tensor,
+        residual_scale: cute.Tensor,
+        coefficient_slab: cute.Tensor,
+        stream: CUstream,
+    ):
+        s_a_layout = cute.make_layout(
+            (FRAME_CONTRACT_M_TILE, FRAME_CONTRACT_K_TILE),
+            stride=(1, FRAME_CONTRACT_M_TILE + FRAME_CONTRACT_SMEM_PADDING),
+        )
+        s_b_layout = cute.make_layout(
+            (FRAME_CONTRACT_N_TILE, FRAME_CONTRACT_K_TILE),
+            stride=(1, FRAME_CONTRACT_N_TILE + FRAME_CONTRACT_SMEM_PADDING),
+        )
+        output_reference_layout = cute.make_layout(
+            (FRAME_CONTRACT_M_TILE, FRAME_CONTRACT_N_TILE),
+            stride=(FRAME_CONTRACT_N_TILE, 1),
+        )
+        self.kernel(
+            grad_out,
+            frame_contract,
+            residual_scale,
+            coefficient_slab,
+            s_a_layout,
+            s_b_layout,
+            output_reference_layout,
+            _make_frame_contract_mma(),
+        ).launch(
+            grid=(
+                cute.ceil_div(grad_out.shape[0], FRAME_CONTRACT_M_TILE),
+                FRAME_CONTRACT_N_TILES,
+                DEGREE_COUNT * FOCUS_COUNT,
+            ),
+            block=[FRAME_CONTRACT_THREADS, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        grad_out: cute.Tensor,
+        frame_contract: cute.Tensor,
+        residual_scale: cute.Tensor,
+        coefficient_slab: cute.Tensor,
+        s_a_layout: cute.Layout,
+        s_b_layout: cute.Layout,
+        output_reference_layout: cute.Layout,
+        tiled_mma: cute.TiledMma,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        m_tile, n_tile, batch = cute.arch.block_idx()
+        m_base = m_tile * FRAME_CONTRACT_M_TILE
+        n_base = n_tile * FRAME_CONTRACT_N_TILE
+        degree = batch // FOCUS_COUNT
+        focus = batch - degree * FOCUS_COUNT
+
+        smem = cute_utils.SmemAllocator()
+        s_a = smem.allocate_tensor(cutlass.Float32, s_a_layout, 16)
+        s_b = smem.allocate_tensor(cutlass.Float32, s_b_layout, 16)
+        thr_mma = tiled_mma.get_slice(tidx)
+        t_s_a = thr_mma.partition_A(s_a)
+        t_s_b = thr_mma.partition_B(s_b)
+        r_a = tiled_mma.make_fragment_A(t_s_a)
+        r_b = tiled_mma.make_fragment_B(t_s_b)
+
+        # The reference supplies the C-fragment partition. The epilogue maps
+        # logical columns directly into the packed coefficient slab.
+        output_reference = cute.make_tensor(
+            coefficient_slab.iterator,
+            output_reference_layout,
+        )
+        t_c_reference = thr_mma.partition_C(output_reference)
+        accumulator = tiled_mma.make_fragment_C(t_c_reference)
+        accumulator.fill(0.0)
+        k_blocks = cute.size(r_a, mode=[2])
+
+        for k_tile in cutlass.range_constexpr(CHANNELS // FRAME_CONTRACT_K_TILE):
+            self._stage_operands(
+                grad_out,
+                frame_contract,
+                residual_scale,
+                s_a,
+                s_b,
+                tidx,
+                m_base,
+                n_base,
+                degree,
+                focus,
+                k_tile * FRAME_CONTRACT_K_TILE,
+            )
+            for k_block in cutlass.range(k_blocks, unroll_full=True):
+                cute.autovec_copy(
+                    t_s_a[None, None, k_block],
+                    r_a[None, None, k_block],
+                )
+                cute.autovec_copy(
+                    t_s_b[None, None, k_block],
+                    r_b[None, None, k_block],
+                )
+                cute.gemm(
+                    tiled_mma,
+                    accumulator,
+                    r_a[None, None, k_block],
+                    r_b[None, None, k_block],
+                    accumulator,
+                )
+            cute.arch.sync_threads()
+
+        self._store_compact_epilogue(
+            coefficient_slab,
+            accumulator,
+            thr_mma,
+            output_reference,
+            m_base,
+            n_base,
+            degree,
+            focus,
+        )
+
+    @cute.jit
+    def _stage_operands(
+        self,
+        grad_out: cute.Tensor,
+        frame_contract: cute.Tensor,
+        residual_scale: cute.Tensor,
+        s_a: cute.Tensor,
+        s_b: cute.Tensor,
+        tidx: cutlass.Int32,
+        m_base: cutlass.Int32,
+        n_base: cutlass.Int32,
+        degree: cutlass.Int32,
+        focus: cutlass.Int32,
+        k_base: cutlass.Constexpr[int],
+    ):
+        a_slots = (
+            FRAME_CONTRACT_M_TILE * FRAME_CONTRACT_K_TILE + FRAME_CONTRACT_THREADS - 1
+        ) // FRAME_CONTRACT_THREADS
+        for slot in cutlass.range_constexpr(a_slots):
+            linear = tidx + slot * FRAME_CONTRACT_THREADS
+            row = linear // FRAME_CONTRACT_K_TILE
+            k = linear - row * FRAME_CONTRACT_K_TILE
+            node = m_base + row
+            value = cutlass.Float32(0.0)
+            if node < grad_out.shape[0]:
+                output_channel = k_base + k
+                value = grad_out[node, degree, focus, output_channel].to(
+                    cutlass.Float32
+                ) * residual_scale[focus, output_channel].to(cutlass.Float32)
+            s_a[row, k] = value
+
+        b_slots = (
+            FRAME_CONTRACT_N_TILE * FRAME_CONTRACT_K_TILE + FRAME_CONTRACT_THREADS - 1
+        ) // FRAME_CONTRACT_THREADS
+        for slot in cutlass.range_constexpr(b_slots):
+            linear = tidx + slot * FRAME_CONTRACT_THREADS
+            output_column = linear // FRAME_CONTRACT_K_TILE
+            k = linear - output_column * FRAME_CONTRACT_K_TILE
+            logical_column = n_base + output_column
+            value = cutlass.Float32(0.0)
+            if logical_column < FRAME_CONTRACT_WIDTH:
+                frame = logical_column // CHANNELS
+                channel = logical_column - frame * CHANNELS
+                value = frame_contract[
+                    degree,
+                    frame,
+                    channel,
+                    k_base + k,
+                ].to(cutlass.Float32)
+            s_b[output_column, k] = value
+        cute.arch.sync_threads()
+
+    @cute.jit
+    def _store_compact_epilogue(
+        self,
+        coefficient_slab: cute.Tensor,
+        accumulator: cute.Tensor,
+        thr_mma,
+        output_reference: cute.Tensor,
+        m_base: cutlass.Int32,
+        n_base: cutlass.Int32,
+        degree: cutlass.Int32,
+        focus: cutlass.Int32,
+    ):
+        accumulator.store(accumulator.load())
+        identity = cute.make_identity_tensor(output_reference.shape)
+        coordinates = thr_mma.partition_C(identity)
+        for value_idx in range(cute.size(accumulator.shape)):
+            coordinate = coordinates[value_idx]
+            node = m_base + coordinate[0]
+            logical_column = n_base + coordinate[1]
+            if (
+                node < coefficient_slab.shape[0]
+                and logical_column < FRAME_CONTRACT_WIDTH
+            ):
+                frame = logical_column // CHANNELS
+                channel = logical_column - frame * CHANNELS
+                packed = degree * FRAME_COUNT + frame
+                hidden = focus * CHANNELS + channel
+                coefficient_slab[node, packed, hidden] = accumulator[value_idx].to(
+                    cutlass.Float32
+                )
+
+
+class _FusedReadoutAdjoint:
+    """Fuse gate, normalized-Gaunt adjoint, and FrameExpand transpose."""
+
+    def __init__(self, schedule: Sm90GauntSchedule) -> None:
+        self.schedule = schedule
+
+    @cute.jit
+    def __call__(
+        self,
+        coefficient_slab: cute.Tensor,
+        scalar_gate: cute.Tensor,
+        product: cute.Tensor,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        frame_expand_t: cute.Tensor,
+        grad_query: cute.Tensor,
+        grad_context: cute.Tensor,
+        grad_scalar_gate: cute.Tensor,
+        grad_scalar_out: cute.Tensor,
+        stream: CUstream,
+    ):
+        workspace_layout = cute.make_layout(
+            (WORKSPACE_PANELS, PACKED_COEFF_DIM, FOLDED_CHANNELS),
+            stride=(VALUES_PER_PANEL, FOLDED_CHANNELS, 1),
+        )
+        statistic_layout = cute.make_layout(
+            (READOUT_GROUPS, FOLDED_CHANNELS),
+            stride=(FOLDED_CHANNELS, 1),
+        )
+        self.kernel(
+            coefficient_slab,
+            scalar_gate,
+            product,
+            left,
+            right,
+            frame_expand_t,
+            grad_query,
+            grad_context,
+            grad_scalar_gate,
+            grad_scalar_out,
+            workspace_layout,
+            statistic_layout,
+        ).launch(
+            grid=(left.shape[0], 1, 1),
+            block=[READOUT_THREADS, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        coefficient_slab: cute.Tensor,
+        scalar_gate: cute.Tensor,
+        product: cute.Tensor,
+        left: cute.Tensor,
+        right: cute.Tensor,
+        frame_expand_t: cute.Tensor,
+        grad_query: cute.Tensor,
+        grad_context: cute.Tensor,
+        grad_scalar_gate: cute.Tensor,
+        grad_scalar_out: cute.Tensor,
+        workspace_layout: cute.Layout,
+        statistic_layout: cute.Layout,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        node, _, _ = cute.arch.block_idx()
+        group = tidx >> 6
+        folded_channel = tidx & (FOLDED_CHANNELS - 1)
+        gate = scalar_gate[node, folded_channel].to(cutlass.Float32)
+
+        smem = cute_utils.SmemAllocator()
+        workspace = smem.allocate_tensor(cutlass.Float32, workspace_layout, 16)
+        statistic_partials = smem.allocate_tensor(
+            cutlass.Float32,
+            statistic_layout,
+            16,
+        )
+        statistic = cutlass.Float32(0.0)
+        scalar_out = cutlass.Float32(0.0)
+
+        for operand in cutlass.range_constexpr(3):
+            for slot in cutlass.range_constexpr(ROWS_PER_GROUP):
+                row = group * ROWS_PER_GROUP + slot
+                value = left[node, row, folded_channel].to(cutlass.Float32)
+                if cutlass.const_expr(operand == RIGHT_PANEL):
+                    value = right[node, row, folded_channel].to(cutlass.Float32)
+                if cutlass.const_expr(operand == GATED_COEFFICIENT_PANEL):
+                    value = coefficient_slab[node, row, folded_channel].to(
+                        cutlass.Float32
+                    )
+                    statistic = statistic + value * product[
+                        node,
+                        row,
+                        folded_channel,
+                    ].to(cutlass.Float32)
+                    if row == 0:
+                        scalar_out = value
+                    value = value * gate
+                workspace[operand, row, folded_channel] = value
+
+        statistic_partials[group, folded_channel] = statistic
+        cute.arch.sync_threads()
+
+        if group == 0:
+            total = statistic_partials[0, folded_channel]
+            for statistic_group in cutlass.range_constexpr(
+                1,
+                READOUT_GROUPS,
+                1,
+            ):
+                total = (
+                    total
+                    + statistic_partials[
+                        statistic_group,
+                        folded_channel,
+                    ]
+                )
+            grad_scalar_gate[node, folded_channel] = total
+            grad_scalar_out[node, folded_channel] = scalar_out
+
+        if group == 0:
+            self._accumulate_group(workspace, folded_channel, 0)
+        elif group == 1:
+            self._accumulate_group(workspace, folded_channel, 1)
+        elif group == 2:
+            self._accumulate_group(workspace, folded_channel, 2)
+        else:
+            self._accumulate_group(workspace, folded_channel, 3)
+        cute.arch.sync_threads()
+
+        for slot in cutlass.range_constexpr(DEGREES_PER_THREAD):
+            linear = tidx + slot * READOUT_THREADS
+            input_channel = linear & (CHANNELS - 1)
+            degree = linear >> 5
+            accumulator_left_0 = cutlass.Float32(0.0)
+            accumulator_left_1 = cutlass.Float32(0.0)
+            accumulator_right_0 = cutlass.Float32(0.0)
+            accumulator_right_1 = cutlass.Float32(0.0)
+            for frame in cutlass.range_constexpr(FRAME_COUNT):
+                packed_row = degree * FRAME_COUNT + frame
+                for output_channel in cutlass.range_constexpr(CHANNELS):
+                    weight = frame_expand_t[
+                        degree,
+                        frame,
+                        output_channel,
+                        input_channel,
+                    ].to(cutlass.Float32)
+                    accumulator_left_0 = accumulator_left_0 + (
+                        workspace[GRAD_LEFT_PANEL, packed_row, output_channel] * weight
+                    )
+                    accumulator_right_0 = accumulator_right_0 + (
+                        workspace[GRAD_RIGHT_PANEL, packed_row, output_channel] * weight
+                    )
+                    accumulator_left_1 = accumulator_left_1 + (
+                        workspace[
+                            GRAD_LEFT_PANEL,
+                            packed_row,
+                            CHANNELS + output_channel,
+                        ]
+                        * weight
+                    )
+                    accumulator_right_1 = accumulator_right_1 + (
+                        workspace[
+                            GRAD_RIGHT_PANEL,
+                            packed_row,
+                            CHANNELS + output_channel,
+                        ]
+                        * weight
+                    )
+            grad_query[node, degree, 0, input_channel] = accumulator_left_0
+            grad_query[node, degree, 1, input_channel] = accumulator_left_1
+            grad_context[node, degree, 0, input_channel] = accumulator_right_0
+            grad_context[node, degree, 1, input_channel] = accumulator_right_1
+
+    @cute.jit
+    def _accumulate_group(
+        self,
+        workspace: cute.Tensor,
+        folded_channel: cutlass.Int32,
+        group: cutlass.Constexpr,
+    ):
+        rows = self.schedule.output_groups[group]
+        for row_slot in cutlass.range_constexpr(len(rows)):
+            input_row = rows[row_slot]
+            accumulator_left = cutlass.Float32(0.0)
+            accumulator_right = cutlass.Float32(0.0)
+            terms = self.schedule.rows[input_row]
+            for path in cutlass.range_constexpr(len(terms)):
+                (
+                    left_row,
+                    right_row,
+                    _,
+                    left_weight_value,
+                    right_weight_value,
+                ) = terms[path]
+                common = (
+                    cutlass.Float32(left_weight_value)
+                    * workspace[GATED_COEFFICIENT_PANEL, left_row, folded_channel]
+                )
+                accumulator_left = accumulator_left + (
+                    common * workspace[RIGHT_PANEL, right_row, folded_channel]
+                )
+                accumulator_right = accumulator_right + (
+                    common * workspace[LEFT_PANEL, right_row, folded_channel]
+                )
+                if cutlass.const_expr(left_row != right_row):
+                    mirrored_common = (
+                        cutlass.Float32(right_weight_value)
+                        * workspace[
+                            GATED_COEFFICIENT_PANEL,
+                            right_row,
+                            folded_channel,
+                        ]
+                    )
+                    accumulator_left = accumulator_left + (
+                        mirrored_common
+                        * workspace[RIGHT_PANEL, left_row, folded_channel]
+                    )
+                    accumulator_right = accumulator_right + (
+                        mirrored_common
+                        * workspace[LEFT_PANEL, left_row, folded_channel]
+                    )
+            workspace[GRAD_LEFT_PANEL, input_row, folded_channel] = accumulator_left
+            workspace[GRAD_RIGHT_PANEL, input_row, folded_channel] = accumulator_right
+
+
+def _fake_state() -> cute.Tensor:
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (cute.sym_int64(), DEGREE_COUNT, FOCUS_COUNT, CHANNELS),
+        stride_order=(3, 2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_frame_weight() -> cute.Tensor:
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (DEGREE_COUNT, FRAME_COUNT, CHANNELS, CHANNELS),
+        stride_order=(3, 2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_focus_channel() -> cute.Tensor:
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_packed() -> cute.Tensor:
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (cute.sym_int64(), PACKED_COEFF_DIM, FOLDED_CHANNELS),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_node_channel() -> cute.Tensor:
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (cute.sym_int64(), FOLDED_CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+@device_aware_lru_cache(maxsize=4)
+def _compiled_frame_contract_adjoint() -> Callable:
+    return cute.compile(
+        _FrameContractAdjoint(),
+        _fake_state(),
+        _fake_frame_weight(),
+        _fake_focus_channel(),
+        _fake_packed(),
+        make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+@device_aware_lru_cache(maxsize=8)
+def _compiled_fused_readout_adjoint(schedule: Sm90GauntSchedule) -> Callable:
+    packed = _fake_packed()
+    node_channel = _fake_node_channel()
+    state = _fake_state()
+    return cute.compile(
+        _FusedReadoutAdjoint(schedule),
+        packed,
+        node_channel,
+        packed,
+        packed,
+        packed,
+        _fake_frame_weight(),
+        state,
+        state,
+        node_channel,
+        node_channel,
+        make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _validate_tensor(
+    name: str,
+    tensor: torch.Tensor,
+    shape: tuple[int, ...],
+    device: torch.device,
+) -> None:
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+    if tensor.dtype != torch.float32:
+        raise TypeError(f"{name} must be FP32")
+    if tensor.device != device or not tensor.is_contiguous():
+        raise ValueError(f"{name} must be contiguous on {device}")
+    if tensor.data_ptr() % 16:
+        raise ValueError(f"{name} must be at least 16-byte aligned")
+
+
+def _validate_runtime(
+    grad_out: torch.Tensor,
+    scalar_gate: torch.Tensor,
+    product: torch.Tensor,
+    left: torch.Tensor,
+    right: torch.Tensor,
+    state: Sm90MessageGridState,
+) -> int:
+    device = left.device
+    if (
+        device.type != "cuda"
+        or tuple(torch.cuda.get_device_capability(device)) != SM90_CAPABILITY
+    ):
+        raise RuntimeError("the fused message-grid readout requires SM90")
+    if not uses_strict_fp32_matmul():
+        raise RuntimeError("the fused message-grid readout requires strict FP32")
+    nodes = int(left.shape[0])
+    if nodes <= 0:
+        raise ValueError("the fused message-grid readout requires N > 0")
+    packed_shape = (nodes, PACKED_COEFF_DIM, FOLDED_CHANNELS)
+    state_shape = (nodes, DEGREE_COUNT, FOCUS_COUNT, CHANNELS)
+    node_channel_shape = (nodes, FOLDED_CHANNELS)
+    for name, tensor, shape in (
+        ("grad_out", grad_out, state_shape),
+        ("scalar_gate", scalar_gate, node_channel_shape),
+        ("product", product, packed_shape),
+        ("left", left, packed_shape),
+        ("right", right, packed_shape),
+        (
+            "frame_contract",
+            state.frame_contract,
+            (DEGREE_COUNT, FRAME_COUNT, CHANNELS, CHANNELS),
+        ),
+        ("residual_scale", state.residual_scale, (FOCUS_COUNT, CHANNELS)),
+        (
+            "frame_expand_t",
+            state.frame_expand_t,
+            (DEGREE_COUNT, FRAME_COUNT, CHANNELS, CHANNELS),
+        ),
+    ):
+        _validate_tensor(name, tensor, shape, device)
+    return nodes
+
+
+def run_sm90_message_grid_backward(
+    net: Any,
+    query_flat: torch.Tensor,
+    context_flat: torch.Tensor,
+    grad_out_flat: torch.Tensor,
+    product_flat: torch.Tensor,
+    state: Sm90MessageGridState,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return query/context adjoints without expanded global adjoint slabs."""
+    from ..message_grid import (
+        _as_flat_ndfc,
+        _focus_linear_backward_input,
+        _frame_expand_packed,
+        _swiglu_backward_input,
+        _validate_contract,
+    )
+
+    query, context = _validate_contract(net, query_flat, context_flat)
+    nodes = int(query_flat.shape[0])
+    scalar_pair = torch.cat([query[:, 0], context[:, 0]], dim=-1).to(net.dtype)
+    scalar_gate = (
+        torch.sigmoid(net.scalar_gate(scalar_pair))
+        .reshape(
+            nodes,
+            FOLDED_CHANNELS,
+        )
+        .contiguous()
+    )
+    left = _frame_expand_packed(net.frame_expand, query).view(
+        nodes,
+        PACKED_COEFF_DIM,
+        FOLDED_CHANNELS,
+    )
+    right = _frame_expand_packed(net.frame_expand, context).view(
+        nodes,
+        PACKED_COEFF_DIM,
+        FOLDED_CHANNELS,
+    )
+    grad_out = (
+        _as_flat_ndfc(
+            net,
+            "grad_out",
+            grad_out_flat,
+            like=query_flat,
+        )
+        .to(net.dtype)
+        .contiguous()
+    )
+    _validate_runtime(grad_out, scalar_gate, product_flat, left, right, state)
+
+    coefficient_slab = torch.empty_like(left)
+    grad_query = torch.empty(
+        query.shape,
+        device=query.device,
+        dtype=query.dtype,
+    )
+    grad_context = torch.empty(
+        context.shape,
+        device=context.device,
+        dtype=context.dtype,
+    )
+    grad_scalar_gate = torch.empty_like(scalar_gate)
+    grad_scalar_out = torch.empty_like(scalar_gate)
+    with torch.cuda.device(left.device):
+        _compiled_frame_contract_adjoint()(
+            grad_out,
+            state.frame_contract,
+            state.residual_scale,
+            coefficient_slab,
+        )
+        _compiled_fused_readout_adjoint(state.schedule)(
+            coefficient_slab,
+            scalar_gate,
+            product_flat,
+            left,
+            right,
+            state.frame_expand_t,
+            grad_query,
+            grad_context,
+            grad_scalar_gate,
+            grad_scalar_out,
+        )
+
+    gate = scalar_gate.view(nodes, FOCUS_COUNT, CHANNELS)
+    grad_scalar_logits = grad_scalar_gate.view_as(gate) * gate * (1.0 - gate)
+    grad_scalar_pair = _focus_linear_backward_input(
+        net.scalar_gate,
+        grad_scalar_logits,
+    ) + _swiglu_backward_input(
+        scalar_pair,
+        grad_scalar_out.view(nodes, FOCUS_COUNT, CHANNELS),
+    )
+    grad_query[:, 0].add_(grad_scalar_pair[:, :, :CHANNELS])
+    grad_context[:, 0].add_(grad_scalar_pair[:, :, CHANNELS:])
+    return (
+        grad_query.reshape_as(query_flat).to(dtype=query_flat.dtype),
+        grad_context.reshape_as(context_flat).to(dtype=context_flat.dtype),
+    )
+
+
+__all__ = [
+    "Sm90MessageGridState",
+    "prepare_sm90_message_grid_state",
+    "run_sm90_message_grid_backward",
+]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/output_gate.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/output_gate.py
new file mode 100644
index 0000000000..4020acb4a6
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/output_gate.py
@@ -0,0 +1,204 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""SM90 CuTe epilogue for chunked final-SO2Linear sufficient statistics."""
+
+from __future__ import (
+    annotations,
+)
+
+import math
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN202, TC002
+
+FOCUS_COUNT = 2
+DEGREE_COUNT = 16
+CHANNELS = 32
+HIDDEN = FOCUS_COUNT * CHANNELS
+THREADS = HIDDEN
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+
+@cute.jit
+def _sigmoid(value):
+    one = cutlass.Float32(1.0)
+    return one / (one + cute.exp(-value))
+
+
+@cute.jit
+def _chunked_final_output_gate_jit(
+    raw: cute.Tensor,
+    x_wide: cute.Tensor,
+    norm_scale: cute.Tensor,
+    gate_weight: cute.Tensor,
+    rotate_inv_rescale: cute.Tensor,
+    out: cute.Tensor,
+    eps: cutlass.Constexpr[float],
+    stream: CUstream,
+):
+    _chunked_final_output_gate_kernel(
+        raw,
+        x_wide,
+        norm_scale,
+        gate_weight,
+        rotate_inv_rescale,
+        out,
+        eps,
+    ).launch(
+        grid=[raw.shape[0], 1, 1],
+        block=[THREADS, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def _chunked_final_output_gate_kernel(
+    raw: cute.Tensor,
+    x_wide: cute.Tensor,
+    norm_scale: cute.Tensor,
+    gate_weight: cute.Tensor,
+    rotate_inv_rescale: cute.Tensor,
+    out: cute.Tensor,
+    eps: cutlass.Constexpr[float],
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    node, _, _ = cute.arch.block_idx()
+    focus = tidx // CHANNELS
+    channel = tidx - focus * CHANNELS
+
+    x = x_wide[node, 0, tidx].to(cutlass.Float32)
+    square_sum = cute.arch.warp_reduction_sum(x * x)
+    inv_rms = cute.rsqrt(square_sum / cutlass.Float32(CHANNELS) + cutlass.Float32(eps))
+    logit_part = (
+        x
+        * inv_rms
+        * norm_scale[focus, channel].to(cutlass.Float32)
+        * gate_weight[channel, focus, 0].to(cutlass.Float32)
+    )
+    gate = _sigmoid(cute.arch.warp_reduction_sum(logit_part))
+    for degree in cutlass.range_constexpr(DEGREE_COUNT):
+        value = raw[node, focus, degree, channel].to(cutlass.Float32)
+        value *= rotate_inv_rescale[degree].to(cutlass.Float32)
+        out[node, degree, tidx] = value * gate
+
+
+def _fake_float(shape: tuple[object, ...], stride_order: tuple[int, ...]):
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        shape,
+        stride_order=stride_order,
+        **FAKE_TENSOR_KW,
+    )
+
+
+@device_aware_lru_cache(maxsize=8)
+def _compiled_chunked_final_output_gate(eps: float) -> Callable:
+    nodes = cute.sym_int64()
+    return cute.compile(
+        _chunked_final_output_gate_jit,
+        _fake_float(
+            (nodes, FOCUS_COUNT, DEGREE_COUNT, CHANNELS),
+            (3, 2, 1, 0),
+        ),
+        _fake_float((nodes, DEGREE_COUNT, HIDDEN), (2, 1, 0)),
+        _fake_float((FOCUS_COUNT, CHANNELS), (1, 0)),
+        _fake_float((CHANNELS, FOCUS_COUNT, 1), (2, 1, 0)),
+        _fake_float((DEGREE_COUNT,), (0,)),
+        _fake_float((nodes, DEGREE_COUNT, HIDDEN), (2, 1, 0)),
+        eps,
+        make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _require_tensor(
+    name: str,
+    tensor: torch.Tensor,
+    shape: tuple[int, ...],
+    device: torch.device,
+) -> None:
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+    if tensor.dtype != torch.float32 or tensor.device != device:
+        raise ValueError(f"{name} must be FP32 on {device}")
+    if not tensor.is_contiguous() or tensor.data_ptr() % 16:
+        raise ValueError(f"{name} must be contiguous and 16-byte aligned")
+
+
+def run_chunked_final_output_gate(
+    *,
+    raw: torch.Tensor,
+    x_wide: torch.Tensor,
+    norm_scale: torch.Tensor,
+    gate_weight: torch.Tensor,
+    rotate_inv_rescale: torch.Tensor,
+    eps: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Apply the real Neo rotate-rescale and output gate to node statistics."""
+    if not math.isfinite(eps) or eps <= 0.0:
+        raise ValueError("output-gate epsilon must be finite and positive")
+    device = raw.device
+    if device.type != "cuda" or tuple(torch.cuda.get_device_capability(device)) != (
+        9,
+        0,
+    ):
+        raise RuntimeError("final output gate requires SM90")
+    if torch.backends.cuda.matmul.allow_tf32:
+        raise RuntimeError("strict FP32 requires allow_tf32=False")
+    node_count = int(raw.shape[0])
+    _require_tensor(
+        "raw",
+        raw,
+        (node_count, FOCUS_COUNT, DEGREE_COUNT, CHANNELS),
+        device,
+    )
+    _require_tensor("x_wide", x_wide, (node_count, DEGREE_COUNT, HIDDEN), device)
+    _require_tensor("norm_scale", norm_scale, (FOCUS_COUNT, CHANNELS), device)
+    _require_tensor("gate_weight", gate_weight, (CHANNELS, FOCUS_COUNT, 1), device)
+    _require_tensor("rotate_inv_rescale", rotate_inv_rescale, (DEGREE_COUNT,), device)
+    if out is None:
+        out = torch.empty(
+            (node_count, DEGREE_COUNT, HIDDEN),
+            device=device,
+            dtype=torch.float32,
+        )
+    else:
+        _require_tensor("out", out, (node_count, DEGREE_COUNT, HIDDEN), device)
+    with torch.cuda.device(device):
+        _compiled_chunked_final_output_gate(float(eps))(
+            raw,
+            x_wide,
+            norm_scale,
+            gate_weight,
+            rotate_inv_rescale,
+            out,
+        )
+    return out
+
+
+__all__ = ["run_chunked_final_output_gate"]
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/persistent.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/persistent.py
new file mode 100644
index 0000000000..921d403399
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/persistent.py
@@ -0,0 +1,593 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Persistent-complex strict-FP32 Neo SO2 stack for SM90.
+
+The fixed Neo ``m=1`` block is a complex 96-wide representation of the dense
+real block ``[[U,V],[-V,U]]``. The first two frozen SO2Linear layers and gated
+residuals remain in this representation without intermediate packing.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+from torch import (
+    Tensor,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN202, TC002
+
+FOCUS_COUNT = 2
+CHANNELS = 32
+M0_ROWS = 4
+M1_ROWS = 3
+M0_WIDTH = M0_ROWS * CHANNELS
+M1_WIDTH = M1_ROWS * CHANNELS
+PAIR_WIDTH = 2 * M1_WIDTH
+GATE_GROUPS = 3
+GATE_WIDTH = GATE_GROUPS * CHANNELS
+GATED_LAYERS = 2
+STACK_LAYERS = 3
+
+ROWS_PER_BLOCK = 8
+THREADS = ROWS_PER_BLOCK * CHANNELS
+WEIGHT_SMEM_STRIDE = GATE_WIDTH + 1
+WEIGHT_VALUES = CHANNELS * GATE_WIDTH
+WEIGHT_LOADS_PER_THREAD = WEIGHT_VALUES // THREADS
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+__all__ = [
+    "NeoPersistentComplexSaved",
+    "NeoPersistentComplexState",
+    "NeoPersistentComplexWeights",
+    "prepare_neo_persistent_complex_weights",
+    "validate_neo_persistent_complex_state",
+]
+
+
+@dataclass(frozen=True)
+class NeoPersistentComplexState:
+    """Focus-major state consumed directly by the persistent stack.
+
+    ``m0`` has shape ``(2,E,128)`` and dtype ``float32``.  ``m1`` has shape
+    ``(2,E,96)`` and dtype ``complex64``.  Both tensors are contiguous; no
+    packing or transposition is needed by the real or complex batched GEMMs.
+    """
+
+    m0: Tensor
+    m1: Tensor
+
+    @property
+    def edge_count(self) -> int:
+        return int(self.m0.shape[1])
+
+    @property
+    def storage_bytes(self) -> int:
+        return sum(
+            tensor.numel() * tensor.element_size() for tensor in (self.m0, self.m1)
+        )
+
+
+@dataclass(frozen=True)
+class NeoPersistentComplexWeights:
+    """Frozen strict-FP32 operands in forward and input-adjoint orientation."""
+
+    w0: Tensor
+    wc: Tensor
+    w0_h: Tensor
+    wc_h: Tensor
+    gate: Tensor
+
+
+@dataclass(frozen=True)
+class NeoPersistentComplexSaved:
+    """Minimal exact gate state for the two nonlinear layers."""
+
+    z0: tuple[Tensor, Tensor]
+    z1: tuple[Tensor, Tensor]
+
+    @property
+    def storage_bytes(self) -> int:
+        return sum(
+            tensor.numel() * tensor.element_size() for tensor in (*self.z0, *self.z1)
+        )
+
+
+def _require_shape(name: str, tensor: Tensor, shape: tuple[int, ...]) -> None:
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+
+
+def _require_frozen_cuda_tensor(
+    name: str,
+    tensor: Tensor,
+    *,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> None:
+    if tensor.dtype != dtype:
+        raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}")
+    if tensor.device != device:
+        raise ValueError(f"{name} must be on {device}, got {tensor.device}")
+    if not tensor.is_contiguous():
+        raise ValueError(f"{name} must be contiguous")
+    if tensor.requires_grad:
+        raise ValueError(f"{name} must be frozen; parameter gradients are out of scope")
+
+
+def validate_neo_persistent_complex_state(
+    state: NeoPersistentComplexState,
+    *,
+    name: str = "state",
+) -> None:
+    """Validate the direct Phase-A/Phase-C split interface."""
+    if (
+        state.m0.ndim != 3
+        or state.m0.shape[0] != FOCUS_COUNT
+        or state.m0.shape[2] != M0_WIDTH
+    ):
+        raise ValueError(
+            f"{name}.m0 must have shape (2,E,128), got {tuple(state.m0.shape)}"
+        )
+    _require_shape(
+        f"{name}.m1",
+        state.m1,
+        (FOCUS_COUNT, state.edge_count, M1_WIDTH),
+    )
+    if state.edge_count <= 0:
+        raise ValueError(f"{name} requires E > 0")
+    if state.m0.dtype != torch.float32 or state.m1.dtype != torch.complex64:
+        raise TypeError(f"{name} requires float32 m0 and complex64 m1")
+    if state.m0.device != state.m1.device:
+        raise ValueError(f"{name}.m0 and {name}.m1 must share a device")
+    if not state.m0.is_contiguous() or not state.m1.is_contiguous():
+        raise ValueError(f"{name} tensors must be focus-major contiguous")
+
+
+def prepare_neo_persistent_complex_weights(
+    w0: Tensor,
+    wp: Tensor,
+    gate: Tensor,
+) -> NeoPersistentComplexWeights:
+    """Convert exact block-real frozen weights to persistent complex weights.
+
+    ``w0`` and ``wp`` use the live ``(input, output)`` orientation consumed by
+    ``torch.bmm``.  The pair block must be exactly ``[[U,V],[-V,U]]``.
+    """
+    _require_shape("w0", w0, (STACK_LAYERS, FOCUS_COUNT, M0_WIDTH, M0_WIDTH))
+    _require_shape(
+        "wp",
+        wp,
+        (STACK_LAYERS, FOCUS_COUNT, PAIR_WIDTH, PAIR_WIDTH),
+    )
+    _require_shape(
+        "gate",
+        gate,
+        (GATED_LAYERS, FOCUS_COUNT, CHANNELS, GATE_WIDTH),
+    )
+    if not w0.is_cuda:
+        raise ValueError("weights must be CUDA tensors")
+    device = w0.device
+    for name, tensor in (("w0", w0), ("wp", wp), ("gate", gate)):
+        _require_frozen_cuda_tensor(
+            name,
+            tensor,
+            dtype=torch.float32,
+            device=device,
+        )
+
+    u = wp[:, :, :M1_WIDTH, :M1_WIDTH]
+    v = wp[:, :, :M1_WIDTH, M1_WIDTH:]
+    if not torch.equal(wp[:, :, M1_WIDTH:, :M1_WIDTH], -v):
+        raise ValueError("wp lower-left block must equal -V exactly")
+    if not torch.equal(wp[:, :, M1_WIDTH:, M1_WIDTH:], u):
+        raise ValueError("wp lower-right block must equal U exactly")
+
+    w0_live = w0.detach().contiguous()
+    wc_live = torch.complex(u, v).contiguous()
+    w0_h = w0_live.transpose(-2, -1).contiguous()
+    wc_h = wc_live.conj().transpose(-2, -1).contiguous()
+
+    return NeoPersistentComplexWeights(
+        w0=w0_live,
+        wc=wc_live,
+        w0_h=w0_h,
+        wc_h=wc_h,
+        gate=gate.detach().contiguous(),
+    )
+
+
+@cute.jit
+def _sigmoid(value):
+    return cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-value))
+
+
+@cute.jit
+def _stage_gate_weight(gate_weight, shared_weight, focus, tidx):
+    for load_slot in cutlass.range_constexpr(WEIGHT_LOADS_PER_THREAD):
+        linear = tidx + load_slot * THREADS
+        source_channel = linear // GATE_WIDTH
+        gate_channel = linear - source_channel * GATE_WIDTH
+        shared_weight[source_channel * WEIGHT_SMEM_STRIDE + gate_channel] = gate_weight[
+            focus, source_channel, gate_channel
+        ].to(cutlass.Float32)
+
+
+@cute.jit
+def _load_gate_values(shared_weight, scalar_rows, row_slot, channel):
+    gate0_logit = cutlass.Float32(0.0)
+    gate1_logit = cutlass.Float32(0.0)
+    gate2_logit = cutlass.Float32(0.0)
+    scalar_base = row_slot * CHANNELS
+    for source_channel in cutlass.range_constexpr(CHANNELS):
+        source = scalar_rows[scalar_base + source_channel]
+        weight_base = source_channel * WEIGHT_SMEM_STRIDE + channel
+        gate0_logit += source * shared_weight[weight_base]
+        gate1_logit += source * shared_weight[weight_base + CHANNELS]
+        gate2_logit += source * shared_weight[weight_base + 2 * CHANNELS]
+    return _sigmoid(gate0_logit), _sigmoid(gate1_logit), _sigmoid(gate2_logit)
+
+
+@cute.jit
+def _select_gate(gate0, gate1, gate2, group):
+    gate = gate0
+    if cutlass.const_expr(group == 1):
+        gate = gate1
+    if cutlass.const_expr(group == 2):
+        gate = gate2
+    return gate
+
+
+@cute.jit
+def _persistent_gate_forward_jit(
+    residual0: cute.Tensor,
+    residual1_ri: cute.Tensor,
+    z0: cute.Tensor,
+    z1_ri: cute.Tensor,
+    gate_weight: cute.Tensor,
+    out0: cute.Tensor,
+    out1_ri: cute.Tensor,
+    stream: CUstream,
+):
+    edge_count = z0.shape[1]
+    _persistent_gate_forward_kernel(
+        residual0,
+        residual1_ri,
+        z0,
+        z1_ri,
+        gate_weight,
+        out0,
+        out1_ri,
+    ).launch(
+        grid=[cute.ceil_div(edge_count, ROWS_PER_BLOCK), FOCUS_COUNT, 1],
+        block=[THREADS, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def _persistent_gate_forward_kernel(
+    residual0: cute.Tensor,
+    residual1_ri: cute.Tensor,
+    z0: cute.Tensor,
+    z1_ri: cute.Tensor,
+    gate_weight: cute.Tensor,
+    out0: cute.Tensor,
+    out1_ri: cute.Tensor,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    edge_block, focus, _ = cute.arch.block_idx()
+    row_slot = tidx // CHANNELS
+    channel = tidx - row_slot * CHANNELS
+    edge = edge_block * ROWS_PER_BLOCK + row_slot
+    edge_count = z0.shape[1]
+
+    smem = cutlass.utils.SmemAllocator()
+    shared_weight = smem.allocate_tensor(
+        cutlass.Float32,
+        CHANNELS * WEIGHT_SMEM_STRIDE,
+    )
+    scalar_rows = smem.allocate_tensor(
+        cutlass.Float32,
+        ROWS_PER_BLOCK * CHANNELS,
+    )
+    _stage_gate_weight(gate_weight, shared_weight, focus, tidx)
+    scalar = cutlass.Float32(0.0)
+    if edge < edge_count:
+        scalar = z0[focus, edge, channel].to(cutlass.Float32)
+    scalar_rows[row_slot * CHANNELS + channel] = scalar
+    cute.arch.sync_threads()
+
+    if edge < edge_count:
+        gate0, gate1, gate2 = _load_gate_values(
+            shared_weight,
+            scalar_rows,
+            row_slot,
+            channel,
+        )
+        out0[focus, edge, channel] = (
+            residual0[focus, edge, channel].to(cutlass.Float32)
+            + scalar * _sigmoid(scalar)
+        ).to(out0.element_type)
+        for group in cutlass.range_constexpr(GATE_GROUPS):
+            gate = _select_gate(gate0, gate1, gate2, group)
+            m0_col = (group + 1) * CHANNELS + channel
+            out0[focus, edge, m0_col] = (
+                residual0[focus, edge, m0_col].to(cutlass.Float32)
+                + z0[focus, edge, m0_col].to(cutlass.Float32) * gate
+            ).to(out0.element_type)
+            m1_col = group * CHANNELS + channel
+            for component in cutlass.range_constexpr(2):
+                out1_ri[focus, edge, m1_col, component] = (
+                    residual1_ri[focus, edge, m1_col, component].to(cutlass.Float32)
+                    + z1_ri[focus, edge, m1_col, component].to(cutlass.Float32) * gate
+                ).to(out1_ri.element_type)
+
+
+@cute.jit
+def _persistent_gate_adjoint_jit(
+    grad0: cute.Tensor,
+    grad1_ri: cute.Tensor,
+    z0: cute.Tensor,
+    z1_ri: cute.Tensor,
+    gate_weight: cute.Tensor,
+    grad_z0: cute.Tensor,
+    grad_z1_ri: cute.Tensor,
+    stream: CUstream,
+):
+    edge_count = z0.shape[1]
+    _persistent_gate_adjoint_kernel(
+        grad0,
+        grad1_ri,
+        z0,
+        z1_ri,
+        gate_weight,
+        grad_z0,
+        grad_z1_ri,
+    ).launch(
+        grid=[cute.ceil_div(edge_count, ROWS_PER_BLOCK), FOCUS_COUNT, 1],
+        block=[THREADS, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def _persistent_gate_adjoint_kernel(
+    grad0: cute.Tensor,
+    grad1_ri: cute.Tensor,
+    z0: cute.Tensor,
+    z1_ri: cute.Tensor,
+    gate_weight: cute.Tensor,
+    grad_z0: cute.Tensor,
+    grad_z1_ri: cute.Tensor,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    edge_block, focus, _ = cute.arch.block_idx()
+    row_slot = tidx // CHANNELS
+    channel = tidx - row_slot * CHANNELS
+    edge = edge_block * ROWS_PER_BLOCK + row_slot
+    edge_count = z0.shape[1]
+
+    smem = cutlass.utils.SmemAllocator()
+    shared_weight = smem.allocate_tensor(
+        cutlass.Float32,
+        CHANNELS * WEIGHT_SMEM_STRIDE,
+    )
+    scalar_rows = smem.allocate_tensor(
+        cutlass.Float32,
+        ROWS_PER_BLOCK * CHANNELS,
+    )
+    grad_logits = smem.allocate_tensor(
+        cutlass.Float32,
+        ROWS_PER_BLOCK * GATE_WIDTH,
+    )
+    _stage_gate_weight(gate_weight, shared_weight, focus, tidx)
+    scalar = cutlass.Float32(0.0)
+    if edge < edge_count:
+        scalar = z0[focus, edge, channel].to(cutlass.Float32)
+    scalar_rows[row_slot * CHANNELS + channel] = scalar
+    cute.arch.sync_threads()
+
+    grad_scalar = cutlass.Float32(0.0)
+    grad_logit0 = cutlass.Float32(0.0)
+    grad_logit1 = cutlass.Float32(0.0)
+    grad_logit2 = cutlass.Float32(0.0)
+    if edge < edge_count:
+        gate0, gate1, gate2 = _load_gate_values(
+            shared_weight,
+            scalar_rows,
+            row_slot,
+            channel,
+        )
+        scalar_sigmoid = _sigmoid(scalar)
+        grad_scalar = (
+            grad0[focus, edge, channel].to(cutlass.Float32)
+            * scalar_sigmoid
+            * (cutlass.Float32(1.0) + scalar * (cutlass.Float32(1.0) - scalar_sigmoid))
+        )
+
+        for group in cutlass.range_constexpr(GATE_GROUPS):
+            gate = _select_gate(gate0, gate1, gate2, group)
+            m0_col = (group + 1) * CHANNELS + channel
+            upstream0 = grad0[focus, edge, m0_col].to(cutlass.Float32)
+            value0 = z0[focus, edge, m0_col].to(cutlass.Float32)
+            grad_z0[focus, edge, m0_col] = (upstream0 * gate).to(grad_z0.element_type)
+            contribution = upstream0 * value0
+
+            m1_col = group * CHANNELS + channel
+            for component in cutlass.range_constexpr(2):
+                upstream1 = grad1_ri[focus, edge, m1_col, component].to(cutlass.Float32)
+                value1 = z1_ri[focus, edge, m1_col, component].to(cutlass.Float32)
+                grad_z1_ri[focus, edge, m1_col, component] = (upstream1 * gate).to(
+                    grad_z1_ri.element_type
+                )
+                contribution += upstream1 * value1
+
+            gate_derivative = gate * (cutlass.Float32(1.0) - gate)
+            if cutlass.const_expr(group == 0):
+                grad_logit0 = contribution * gate_derivative
+            if cutlass.const_expr(group == 1):
+                grad_logit1 = contribution * gate_derivative
+            if cutlass.const_expr(group == 2):
+                grad_logit2 = contribution * gate_derivative
+
+    grad_base = row_slot * GATE_WIDTH + channel
+    grad_logits[grad_base] = grad_logit0
+    grad_logits[grad_base + CHANNELS] = grad_logit1
+    grad_logits[grad_base + 2 * CHANNELS] = grad_logit2
+    cute.arch.sync_threads()
+
+    if edge < edge_count:
+        for gate_channel in cutlass.range_constexpr(GATE_WIDTH):
+            grad_scalar += (
+                grad_logits[row_slot * GATE_WIDTH + gate_channel]
+                * shared_weight[channel * WEIGHT_SMEM_STRIDE + gate_channel]
+            )
+        grad_z0[focus, edge, channel] = grad_scalar.to(grad_z0.element_type)
+
+
+def _fake_m0():
+    edge_count = cute.sym_int64()
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, edge_count, M0_WIDTH),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_m1_ri():
+    edge_count = cute.sym_int64()
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, edge_count, M1_WIDTH, 2),
+        stride_order=(3, 2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_gate_weight():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, CHANNELS, GATE_WIDTH),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _compile_forward() -> Callable:
+    return cute.compile(
+        _persistent_gate_forward_jit,
+        _fake_m0(),
+        _fake_m1_ri(),
+        _fake_m0(),
+        _fake_m1_ri(),
+        _fake_gate_weight(),
+        _fake_m0(),
+        _fake_m1_ri(),
+        make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _compile_adjoint() -> Callable:
+    return cute.compile(
+        _persistent_gate_adjoint_jit,
+        _fake_m0(),
+        _fake_m1_ri(),
+        _fake_m0(),
+        _fake_m1_ri(),
+        _fake_gate_weight(),
+        _fake_m0(),
+        _fake_m1_ri(),
+        make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+@device_aware_lru_cache(maxsize=2)
+def _compiled_forward() -> Callable:
+    return _compile_forward()
+
+
+@device_aware_lru_cache(maxsize=2)
+def _compiled_adjoint() -> Callable:
+    return _compile_adjoint()
+
+
+def _m1_real_view(m1: Tensor) -> Tensor:
+    view = torch.view_as_real(m1)
+    if not view.is_contiguous():
+        raise ValueError("complex state must expose a contiguous real/imag view")
+    return view
+
+
+def _run_gate_forward(
+    residual: NeoPersistentComplexState,
+    z: NeoPersistentComplexState,
+    gate_weight: Tensor,
+    out: NeoPersistentComplexState,
+) -> None:
+    with torch.cuda.device(z.m0.device):
+        _compiled_forward()(
+            residual.m0,
+            _m1_real_view(residual.m1),
+            z.m0,
+            _m1_real_view(z.m1),
+            gate_weight,
+            out.m0,
+            _m1_real_view(out.m1),
+        )
+
+
+def _run_gate_adjoint(
+    grad: NeoPersistentComplexState,
+    z: NeoPersistentComplexState,
+    gate_weight: Tensor,
+    out: NeoPersistentComplexState,
+) -> None:
+    with torch.cuda.device(z.m0.device):
+        _compiled_adjoint()(
+            grad.m0,
+            _m1_real_view(grad.m1),
+            z.m0,
+            _m1_real_view(z.m1),
+            gate_weight,
+            out.m0,
+            _m1_real_view(out.m1),
+        )
+
+
+def _empty_state_like(state: NeoPersistentComplexState) -> NeoPersistentComplexState:
+    return NeoPersistentComplexState(
+        m0=torch.empty_like(state.m0),
+        m1=torch.empty_like(state.m1),
+    )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a.py
new file mode 100644
index 0000000000..5341d1f9fb
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a.py
@@ -0,0 +1,555 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Direct strict-FP32 Phase-A producer for the persistent-complex SO2 stack.
+
+The generic boundary materializes Neo's reduced SO2 state as block-real
+``(E,2,10,32)`` and then launches a second kernel to transpose/split it into
+focus-major ``m=0`` real and interleaved ``m=1`` complex panels. This producer
+applies the same packed-Wigner rotation, compact radial
+maps, and rank-1 channel basis, but writes the persistent representation
+directly:
+
+* reduced rows 0..3 -> ``m0[focus, edge, 4 * channel]``;
+* reduced rows 4..6 -> the real component of ``m1``;
+* reduced rows 7..9 -> the imaginary component of ``m1``.
+
+No full-edge block-real slab exists on this path. ``N`` and ``E``
+remain runtime dimensions; only the Neo representation contract is static.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as cute_utils
+import torch
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+from ..wigner_layout import (
+    PACKED_VALUE_COUNT,
+)
+from .persistent import (
+    NeoPersistentComplexState,
+    validate_neo_persistent_complex_state,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN202, ANN204
+
+EDGE_TILE = 32
+THREADS = 256
+FOCUS_COUNT = 2
+FOCUS_DIM = 32
+FULL_CHANNELS = FOCUS_COUNT * FOCUS_DIM
+M0_ROWS = 4
+M1_ROWS = 3
+M0_WIDTH = M0_ROWS * FOCUS_DIM
+M1_WIDTH = M1_ROWS * FOCUS_DIM
+RADIAL_COMPACT = 25
+
+D_CACHE_BYTES = EDGE_TILE * PACKED_VALUE_COUNT * 4
+RADIAL_CACHE_BYTES = EDGE_TILE * RADIAL_COMPACT * 4
+SRC_CACHE_BYTES = EDGE_TILE * 4
+BASIS_CACHE_BYTES = FULL_CHANNELS * 4
+CTA_SHARED_BYTES = (
+    D_CACHE_BYTES + RADIAL_CACHE_BYTES + SRC_CACHE_BYTES + BASIS_CACHE_BYTES
+)
+
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+DEFAULT_STREAM = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT)
+
+__all__ = [
+    "CTA_SHARED_BYTES",
+    "run_neo_phase_a_persistent_complex_fp32",
+]
+
+
+def _fake_x_wide():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (cute.sym_int64(), 16 * FULL_CHANNELS),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_src():
+    return make_fake_compact_tensor(
+        cutlass.Int32,
+        (cute.sym_int64(),),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_edge_matrix(columns: int):
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (cute.sym_int64(), columns),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_channel_basis():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FULL_CHANNELS,),
+        stride_order=(0,),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_m0():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, cute.sym_int64(), M0_WIDTH),
+        stride_order=(2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _fake_m1_ri():
+    return make_fake_compact_tensor(
+        cutlass.Float32,
+        (FOCUS_COUNT, cute.sym_int64(), M1_WIDTH, 2),
+        stride_order=(3, 2, 1, 0),
+        **FAKE_TENSOR_KW,
+    )
+
+
+@cute.jit
+def _rotation_row_cached(
+    x_wide,
+    d_cache,
+    edge_row,
+    src_node,
+    channel,
+    panel_start: cutlass.Constexpr[int],
+    full_start: cutlass.Constexpr[int],
+    width: cutlass.Constexpr[int],
+):
+    """Evaluate one packed-Wigner row with the Phase-A reduction order."""
+    value = cutlass.Float32(0.0)
+    for local_col in cutlass.range_constexpr(width):
+        value += d_cache[edge_row, panel_start + local_col].to(
+            cutlass.Float32
+        ) * x_wide[
+            src_node,
+            (full_start + local_col) * FULL_CHANNELS + channel,
+        ].to(cutlass.Float32)
+    return value
+
+
+class CuteNeoPhaseAPersistentComplexFP32:
+    """Produce native persistent-complex panels without a dense boundary."""
+
+    @cute.jit
+    def __call__(
+        self,
+        x_wide,
+        src,
+        d_full,
+        radial_compact,
+        channel_basis,
+        m0,
+        m1_ri,
+        stream: cuda.CUstream = DEFAULT_STREAM,
+    ):
+        d_layout = cute.make_layout(
+            (EDGE_TILE, PACKED_VALUE_COUNT),
+            stride=(PACKED_VALUE_COUNT, 1),
+        )
+        radial_layout = cute.make_layout(
+            (EDGE_TILE, RADIAL_COMPACT),
+            stride=(RADIAL_COMPACT, 1),
+        )
+        edge_layout = cute.make_layout((EDGE_TILE,), stride=(1,))
+        basis_layout = cute.make_layout((FULL_CHANNELS,), stride=(1,))
+        self.kernel(
+            x_wide,
+            src,
+            d_full,
+            radial_compact,
+            channel_basis,
+            m0,
+            m1_ri,
+            d_layout,
+            radial_layout,
+            edge_layout,
+            basis_layout,
+        ).launch(
+            grid=(cute.ceil_div(m0.shape[1], EDGE_TILE), 1, 1),
+            block=[THREADS, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        x_wide,
+        src,
+        d_full,
+        radial_compact,
+        channel_basis,
+        m0,
+        m1_ri,
+        d_layout,
+        radial_layout,
+        edge_layout,
+        basis_layout,
+    ):
+        tidx, _, _ = cute.arch.thread_idx()
+        edge_tile, _, _ = cute.arch.block_idx()
+        edge_count = m0.shape[1]
+
+        smem = cute_utils.SmemAllocator()
+        d_cache = smem.allocate_tensor(cutlass.Float32, d_layout, 16)
+        radial_cache = smem.allocate_tensor(cutlass.Float32, radial_layout, 16)
+        src_cache = smem.allocate_tensor(cutlass.Int32, edge_layout, 16)
+        basis_cache = smem.allocate_tensor(cutlass.Float32, basis_layout, 16)
+
+        self._load_edge_state(
+            src,
+            d_full,
+            radial_compact,
+            channel_basis,
+            d_cache,
+            radial_cache,
+            src_cache,
+            basis_cache,
+            edge_count,
+            tidx,
+            edge_tile,
+        )
+        self._produce_split_panels(
+            x_wide,
+            d_cache,
+            radial_cache,
+            src_cache,
+            basis_cache,
+            m0,
+            m1_ri,
+            edge_count,
+            tidx,
+            edge_tile,
+        )
+
+    @cute.jit
+    def _load_edge_state(
+        self,
+        src,
+        d_full,
+        radial_compact,
+        channel_basis,
+        d_cache,
+        radial_cache,
+        src_cache,
+        basis_cache,
+        edge_count,
+        tidx,
+        edge_tile,
+    ):
+        d_slots = (EDGE_TILE * PACKED_VALUE_COUNT + THREADS - 1) // THREADS
+        for slot in cutlass.range_constexpr(d_slots):
+            linear = tidx + slot * THREADS
+            if linear < EDGE_TILE * PACKED_VALUE_COUNT:
+                edge_row = linear // PACKED_VALUE_COUNT
+                column = linear - edge_row * PACKED_VALUE_COUNT
+                edge = edge_tile * EDGE_TILE + edge_row
+                value = cutlass.Float32(0.0)
+                if edge < edge_count:
+                    value = d_full[edge, column].to(cutlass.Float32)
+                d_cache[edge_row, column] = value
+
+        radial_slots = (EDGE_TILE * RADIAL_COMPACT + THREADS - 1) // THREADS
+        for slot in cutlass.range_constexpr(radial_slots):
+            linear = tidx + slot * THREADS
+            if linear < EDGE_TILE * RADIAL_COMPACT:
+                edge_row = linear // RADIAL_COMPACT
+                column = linear - edge_row * RADIAL_COMPACT
+                edge = edge_tile * EDGE_TILE + edge_row
+                value = cutlass.Float32(0.0)
+                if edge < edge_count:
+                    value = radial_compact[edge, column].to(cutlass.Float32)
+                radial_cache[edge_row, column] = value
+
+        if tidx < EDGE_TILE:
+            edge = edge_tile * EDGE_TILE + tidx
+            src_node = cutlass.Int32(0)
+            if edge < edge_count:
+                src_node = src[edge]
+            src_cache[tidx] = src_node
+        if tidx < FULL_CHANNELS:
+            basis_cache[tidx] = channel_basis[tidx].to(cutlass.Float32)
+        cute.arch.sync_threads()
+
+    @cute.jit
+    def _produce_split_panels(
+        self,
+        x_wide,
+        d_cache,
+        radial_cache,
+        src_cache,
+        basis_cache,
+        m0,
+        m1_ri,
+        edge_count,
+        tidx,
+        edge_tile,
+    ):
+        tasks = (EDGE_TILE * FULL_CHANNELS) // THREADS
+        for task in cutlass.range_constexpr(tasks):
+            linear = tidx + task * THREADS
+            edge_row = linear // FULL_CHANNELS
+            channel = linear - edge_row * FULL_CHANNELS
+            edge = edge_tile * EDGE_TILE + edge_row
+
+            if edge < edge_count:
+                focus = channel // FOCUS_DIM
+                focus_channel = channel - focus * FOCUS_DIM
+                src_node = src_cache[edge_row]
+                x0 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 0, 0, 1
+                )
+                x1 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 1, 1, 3
+                )
+                x2 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 10, 4, 5
+                )
+                x3 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 25, 9, 7
+                )
+                x4 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 4, 1, 3
+                )
+                x5 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 15, 4, 5
+                )
+                x6 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 32, 9, 7
+                )
+                x7 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 7, 1, 3
+                )
+                x8 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 20, 4, 5
+                )
+                x9 = _rotation_row_cached(
+                    x_wide, d_cache, edge_row, src_node, channel, 39, 9, 7
+                )
+
+                basis = basis_cache[channel].to(cutlass.Float32)
+                y0 = (
+                    radial_cache[edge_row, 0] * x0
+                    + radial_cache[edge_row, 4] * x1
+                    + radial_cache[edge_row, 8] * x2
+                    + radial_cache[edge_row, 12] * x3
+                ) * basis
+                y1 = (
+                    radial_cache[edge_row, 1] * x0
+                    + radial_cache[edge_row, 5] * x1
+                    + radial_cache[edge_row, 9] * x2
+                    + radial_cache[edge_row, 13] * x3
+                ) * basis
+                y2 = (
+                    radial_cache[edge_row, 2] * x0
+                    + radial_cache[edge_row, 6] * x1
+                    + radial_cache[edge_row, 10] * x2
+                    + radial_cache[edge_row, 14] * x3
+                ) * basis
+                y3 = (
+                    radial_cache[edge_row, 3] * x0
+                    + radial_cache[edge_row, 7] * x1
+                    + radial_cache[edge_row, 11] * x2
+                    + radial_cache[edge_row, 15] * x3
+                ) * basis
+                y4 = (
+                    radial_cache[edge_row, 16] * x4
+                    + radial_cache[edge_row, 19] * x5
+                    + radial_cache[edge_row, 22] * x6
+                ) * basis
+                y5 = (
+                    radial_cache[edge_row, 17] * x4
+                    + radial_cache[edge_row, 20] * x5
+                    + radial_cache[edge_row, 23] * x6
+                ) * basis
+                y6 = (
+                    radial_cache[edge_row, 18] * x4
+                    + radial_cache[edge_row, 21] * x5
+                    + radial_cache[edge_row, 24] * x6
+                ) * basis
+                y7 = (
+                    radial_cache[edge_row, 16] * x7
+                    + radial_cache[edge_row, 19] * x8
+                    + radial_cache[edge_row, 22] * x9
+                ) * basis
+                y8 = (
+                    radial_cache[edge_row, 17] * x7
+                    + radial_cache[edge_row, 20] * x8
+                    + radial_cache[edge_row, 23] * x9
+                ) * basis
+                y9 = (
+                    radial_cache[edge_row, 18] * x7
+                    + radial_cache[edge_row, 21] * x8
+                    + radial_cache[edge_row, 24] * x9
+                ) * basis
+
+                m0[focus, edge, focus_channel] = y0
+                m0[focus, edge, FOCUS_DIM + focus_channel] = y1
+                m0[focus, edge, 2 * FOCUS_DIM + focus_channel] = y2
+                m0[focus, edge, 3 * FOCUS_DIM + focus_channel] = y3
+                m1_ri[focus, edge, focus_channel, 0] = y4
+                m1_ri[focus, edge, focus_channel, 1] = y7
+                m1_ri[focus, edge, FOCUS_DIM + focus_channel, 0] = y5
+                m1_ri[focus, edge, FOCUS_DIM + focus_channel, 1] = y8
+                m1_ri[focus, edge, 2 * FOCUS_DIM + focus_channel, 0] = y6
+                m1_ri[focus, edge, 2 * FOCUS_DIM + focus_channel, 1] = y9
+
+
+@device_aware_lru_cache(maxsize=8)
+def _compiled_producer(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> Callable:
+    if compute_capability != (9, 0):
+        raise RuntimeError("direct split Phase A requires SM90")
+    with torch.cuda.device(device_index):
+        return cute.compile(
+            CuteNeoPhaseAPersistentComplexFP32(),
+            _fake_x_wide(),
+            _fake_src(),
+            _fake_edge_matrix(PACKED_VALUE_COUNT),
+            _fake_edge_matrix(RADIAL_COMPACT),
+            _fake_channel_basis(),
+            _fake_m0(),
+            _fake_m1_ri(),
+            stream=make_fake_stream(use_tvm_ffi_env_stream=False),
+            options="--enable-tvm-ffi",
+        )
+
+
+def _validate_inputs(
+    x_wide: torch.Tensor,
+    src: torch.Tensor,
+    d_full: torch.Tensor,
+    radial_compact: torch.Tensor,
+    channel_basis: torch.Tensor,
+) -> tuple[int, torch.Tensor]:
+    if x_wide.ndim != 3 or tuple(x_wide.shape[1:]) != (16, FULL_CHANNELS):
+        raise ValueError(f"x_wide must have shape (N,16,64), got {x_wide.shape}")
+    if x_wide.shape[0] <= 0:
+        raise ValueError("direct split Phase A requires N > 0")
+    if src.ndim != 1 or src.dtype not in (torch.int32, torch.int64):
+        raise TypeError("src must be a one-dimensional int32 or int64 tensor")
+    edge_count = src.numel()
+    if edge_count <= 0:
+        raise ValueError("direct split Phase A requires E > 0")
+    if tuple(d_full.shape) != (edge_count, PACKED_VALUE_COUNT):
+        raise ValueError(f"d_full must have shape {(edge_count, PACKED_VALUE_COUNT)}")
+    if tuple(radial_compact.shape) != (edge_count, RADIAL_COMPACT):
+        raise ValueError(
+            f"radial_compact must have shape {(edge_count, RADIAL_COMPACT)}"
+        )
+    if tuple(channel_basis.shape) != (FULL_CHANNELS,):
+        raise ValueError("channel_basis must have shape (64,)")
+
+    float_tensors = (x_wide, d_full, radial_compact, channel_basis)
+    if any(t.dtype != torch.float32 or not t.is_cuda for t in float_tensors):
+        raise TypeError("all Phase-A floating-point operands must be CUDA float32")
+    if not src.is_cuda:
+        raise TypeError("src must be a CUDA tensor")
+    if any(t.device != x_wide.device for t in (*float_tensors[1:], src)):
+        raise ValueError("all Phase-A operands must share x_wide.device")
+    if any(not t.is_contiguous() for t in float_tensors):
+        raise ValueError("all Phase-A floating-point operands must be contiguous")
+    src_i32 = (
+        src
+        if src.dtype == torch.int32 and src.is_contiguous()
+        else src.to(dtype=torch.int32).contiguous()
+    )
+    return edge_count, src_i32
+
+
+def _allocate_state(edge_count: int, device: torch.device) -> NeoPersistentComplexState:
+    return NeoPersistentComplexState(
+        m0=torch.empty(
+            (FOCUS_COUNT, edge_count, M0_WIDTH),
+            dtype=torch.float32,
+            device=device,
+        ),
+        m1=torch.empty(
+            (FOCUS_COUNT, edge_count, M1_WIDTH),
+            dtype=torch.complex64,
+            device=device,
+        ),
+    )
+
+
+def run_neo_phase_a_persistent_complex_fp32(
+    *,
+    x_wide: torch.Tensor,
+    src: torch.Tensor,
+    d_full: torch.Tensor,
+    radial_compact: torch.Tensor,
+    channel_basis: torch.Tensor,
+    out: NeoPersistentComplexState | None = None,
+) -> NeoPersistentComplexState:
+    """Write Phase A directly into the persistent stack's native split state."""
+    edge_count, src_i32 = _validate_inputs(
+        x_wide,
+        src,
+        d_full,
+        radial_compact,
+        channel_basis,
+    )
+    if out is None:
+        out = _allocate_state(edge_count, x_wide.device)
+    validate_neo_persistent_complex_state(out, name="out")
+    if out.edge_count != edge_count or out.m0.device != x_wide.device:
+        raise ValueError("out must have matching E and share x_wide.device")
+
+    if torch.backends.cuda.matmul.allow_tf32:
+        raise RuntimeError("strict FP32 requires allow_tf32=False")
+    if torch.get_float32_matmul_precision() != "highest":
+        raise RuntimeError("strict FP32 requires float32 matmul precision 'highest'")
+    device_index = x_wide.device.index
+    if device_index is None:
+        raise RuntimeError("direct split Phase A requires CUDA")
+    compute_capability = tuple(torch.cuda.get_device_capability(device_index))
+    compiled = _compiled_producer(device_index, compute_capability)
+    m1_ri = torch.view_as_real(out.m1)
+    if not m1_ri.is_contiguous():
+        raise ValueError("out.m1 must expose a contiguous interleaved real/imag view")
+    stream = cuda.CUstream(torch.cuda.current_stream(x_wide.device).cuda_stream)
+    compiled(
+        x_wide.view(x_wide.shape[0], 16 * FULL_CHANNELS),
+        src_i32,
+        d_full,
+        radial_compact,
+        channel_basis,
+        out.m0,
+        m1_ri,
+        stream=stream,
+    )
+    return out
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a_backward.py
new file mode 100644
index 0000000000..1116f95a2b
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a_backward.py
@@ -0,0 +1,936 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Exact split-complex input adjoint for the SM90 Neo Phase-A boundary.
+
+The incoming adjoint remains in the persistent stack's focus-major contract:
+``m0`` is float32 ``(2,E,128)`` and ``m1`` is complex64 ``(2,E,96)``.  One
+CTA owns one source node and reads those panels directly while recomputing the
+packed-Wigner rotation.  It emits the exact input adjoints for node features,
+packed Wigner values, compact radial maps, and the rank-1 channel basis without
+ever reconstructing an ``(E,2,10,32)`` block-real gradient slab.
+
+Source CSR is a caller-owned edge-cache property.  It preserves physical edge
+order while giving each node exclusive ownership of its feature adjoint, so no
+atomics or ``(E,16,64)`` edge-local reduction tensor is required.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils as cute_utils
+import torch
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+from torch import (
+    Tensor,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+from ..kernels.radial_phase_a_backward import (
+    COMPACT_WIDTH,
+    DEGREE_COUNT,
+    FOCUS_COUNT,
+    FOCUS_HIDDEN,
+    GROUPS_PER_CTA,
+    GROUPS_PER_WARP,
+    HIDDEN,
+    PACKED_WIGNER_VALUES,
+    REDUCED_COUNT,
+    SHARED_ROW_PITCH,
+    WARP_REDUCTION_GROUP,
+    _grad_x_value,
+    _recompute_local_value,
+    _warp_owned_grad_compact,
+    _warp_owned_grad_d,
+)
+from .persistent import (
+    M0_WIDTH,
+    M1_WIDTH,
+    NeoPersistentComplexState,
+    validate_neo_persistent_complex_state,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN202, TC002
+
+THREADS = HIDDEN
+REDUCE_THREADS = 32
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+__all__ = [
+    "NeoPhaseAPersistentComplexAdjoints",
+    "NeoPhaseAPersistentComplexBackwardWorkspace",
+    "allocate_neo_phase_a_persistent_complex_backward",
+    "run_neo_phase_a_persistent_complex_backward_fp32",
+]
+
+
+@dataclass(frozen=True)
+class NeoPhaseAPersistentComplexAdjoints:
+    """Differentiable input adjoints for the strict-FP32 Phase-A producer."""
+
+    grad_x_wide: Tensor
+    grad_d_full: Tensor
+    grad_radial_compact: Tensor
+    grad_channel_basis: Tensor
+
+
+@dataclass(frozen=True)
+class NeoPhaseAPersistentComplexBackwardWorkspace:
+    """Small deterministic channel-basis reduction workspace."""
+
+    grad_basis_by_node: Tensor
+
+    @property
+    def storage_bytes(self) -> int:
+        return self.grad_basis_by_node.numel() * self.grad_basis_by_node.element_size()
+
+
+@dataclass(frozen=True)
+class _BackwardParams:
+    grad_m0: cute.Tensor
+    grad_m1_ri: cute.Tensor
+    radial_compact: cute.Tensor
+    channel_basis: cute.Tensor
+    x_wide: cute.Tensor
+    source_order: cute.Tensor
+    source_ptr: cute.Tensor
+    d_full: cute.Tensor
+    grad_x_wide: cute.Tensor
+    grad_d_full: cute.Tensor
+    grad_radial_compact: cute.Tensor
+    grad_basis_by_node: cute.Tensor
+
+
+@cute.jit
+def _split_grad_value(
+    grad_m0,
+    grad_m1_ri,
+    edge,
+    reduced: cutlass.Constexpr[int],
+    channel,
+):
+    focus = channel // FOCUS_HIDDEN
+    focus_channel = channel - focus * FOCUS_HIDDEN
+    if cutlass.const_expr(reduced < 4):
+        return grad_m0[
+            focus,
+            edge,
+            reduced * FOCUS_HIDDEN + focus_channel,
+        ].to(cutlass.Float32)
+    if cutlass.const_expr(reduced < 7):
+        return grad_m1_ri[
+            focus,
+            edge,
+            (reduced - 4) * FOCUS_HIDDEN + focus_channel,
+            0,
+        ].to(cutlass.Float32)
+    return grad_m1_ri[
+        focus,
+        edge,
+        (reduced - 7) * FOCUS_HIDDEN + focus_channel,
+        1,
+    ].to(cutlass.Float32)
+
+
+@cute.jit
+def _radial_output_value(
+    local_values,
+    compact,
+    reduced: cutlass.Constexpr[int],
+    channel,
+):
+    """Recompute the pre-basis radial output for one reduced row/channel."""
+    value = cutlass.Float32(0.0)
+    if cutlass.const_expr(reduced < 4):
+        for input_row in cutlass.range_constexpr(4):
+            value += (
+                compact[input_row * 4 + reduced]
+                * local_values[input_row * SHARED_ROW_PITCH + channel]
+            )
+    elif cutlass.const_expr(reduced < 7):
+        output_row = reduced - 4
+        for input_row in cutlass.range_constexpr(3):
+            value += (
+                compact[16 + input_row * 3 + output_row]
+                * local_values[(4 + input_row) * SHARED_ROW_PITCH + channel]
+            )
+    else:
+        output_row = reduced - 7
+        for input_row in cutlass.range_constexpr(3):
+            value += (
+                compact[16 + input_row * 3 + output_row]
+                * local_values[(7 + input_row) * SHARED_ROW_PITCH + channel]
+            )
+    return value
+
+
+@cute.jit
+def _phase_a_split_backward_jit(
+    grad_m0: cute.Tensor,
+    grad_m1_ri: cute.Tensor,
+    radial_compact: cute.Tensor,
+    channel_basis: cute.Tensor,
+    x_wide: cute.Tensor,
+    source_order: cute.Tensor,
+    source_ptr: cute.Tensor,
+    d_full: cute.Tensor,
+    grad_x_wide: cute.Tensor,
+    grad_d_full: cute.Tensor,
+    grad_radial_compact: cute.Tensor,
+    grad_basis_by_node: cute.Tensor,
+    grad_channel_basis: cute.Tensor,
+    stream: CUstream,
+):
+    params = _BackwardParams(
+        grad_m0=grad_m0,
+        grad_m1_ri=grad_m1_ri,
+        radial_compact=radial_compact,
+        channel_basis=channel_basis,
+        x_wide=x_wide,
+        source_order=source_order,
+        source_ptr=source_ptr,
+        d_full=d_full,
+        grad_x_wide=grad_x_wide,
+        grad_d_full=grad_d_full,
+        grad_radial_compact=grad_radial_compact,
+        grad_basis_by_node=grad_basis_by_node,
+    )
+    _phase_a_split_backward_kernel(params).launch(
+        grid=[x_wide.shape[0], 1, 1],
+        block=[THREADS, 1, 1],
+        stream=stream,
+    )
+    _reduce_channel_basis_kernel(grad_basis_by_node, grad_channel_basis).launch(
+        grid=[HIDDEN, 1, 1],
+        block=[REDUCE_THREADS, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def _phase_a_split_backward_kernel(params: _BackwardParams):
+    channel, _, _ = cute.arch.thread_idx()
+    node, _, _ = cute.arch.block_idx()
+    x_row_pitch = SHARED_ROW_PITCH
+
+    smem = cute_utils.SmemAllocator()
+    x_values = smem.allocate_tensor(
+        cutlass.Float32,
+        DEGREE_COUNT * x_row_pitch,
+    )
+    focus_grad = smem.allocate_tensor(
+        cutlass.Float32,
+        REDUCED_COUNT * SHARED_ROW_PITCH,
+    )
+    # Primal local rows are overwritten by their adjoints after grad-radial and
+    # grad-basis have consumed them.
+    local_values = smem.allocate_tensor(
+        cutlass.Float32,
+        REDUCED_COUNT * SHARED_ROW_PITCH,
+    )
+    d_values = smem.allocate_tensor(cutlass.Float32, PACKED_WIGNER_VALUES)
+    compact = smem.allocate_tensor(cutlass.Float32, COMPACT_WIDTH)
+
+    for full_row in cutlass.range_constexpr(DEGREE_COUNT):
+        x_values[full_row * x_row_pitch + channel] = params.x_wide[
+            node,
+            full_row * HIDDEN + channel,
+        ].to(cutlass.Float32)
+    cute.arch.sync_threads()
+
+    grad_x_0 = cutlass.Float32(0.0)
+    grad_x_1 = cutlass.Float32(0.0)
+    grad_x_2 = cutlass.Float32(0.0)
+    grad_x_3 = cutlass.Float32(0.0)
+    grad_x_4 = cutlass.Float32(0.0)
+    grad_x_5 = cutlass.Float32(0.0)
+    grad_x_6 = cutlass.Float32(0.0)
+    grad_x_7 = cutlass.Float32(0.0)
+    grad_x_8 = cutlass.Float32(0.0)
+    grad_x_9 = cutlass.Float32(0.0)
+    grad_x_10 = cutlass.Float32(0.0)
+    grad_x_11 = cutlass.Float32(0.0)
+    grad_x_12 = cutlass.Float32(0.0)
+    grad_x_13 = cutlass.Float32(0.0)
+    grad_x_14 = cutlass.Float32(0.0)
+    grad_x_15 = cutlass.Float32(0.0)
+    grad_basis = cutlass.Float32(0.0)
+
+    lo = params.source_ptr[node]
+    hi = params.source_ptr[node + 1]
+    for slot in cutlass.range(lo, hi, 1, unroll=1):
+        edge = params.source_order[slot]
+        for reduced in cutlass.range_constexpr(REDUCED_COUNT):
+            focus_grad[reduced * SHARED_ROW_PITCH + channel] = _split_grad_value(
+                params.grad_m0,
+                params.grad_m1_ri,
+                edge,
+                reduced,
+                channel,
+            )
+        if channel < COMPACT_WIDTH:
+            compact[channel] = params.radial_compact[edge, channel].to(cutlass.Float32)
+        if channel < PACKED_WIGNER_VALUES:
+            d_values[channel] = params.d_full[edge, channel].to(cutlass.Float32)
+        cute.arch.sync_threads()
+
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            0,
+            0,
+            0,
+            1,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            1,
+            1,
+            1,
+            3,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            2,
+            10,
+            4,
+            5,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            3,
+            25,
+            9,
+            7,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            4,
+            4,
+            1,
+            3,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            5,
+            15,
+            4,
+            5,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            6,
+            32,
+            9,
+            7,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            7,
+            7,
+            1,
+            3,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            8,
+            20,
+            4,
+            5,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        _recompute_local_value(
+            local_values,
+            x_values,
+            d_values,
+            channel,
+            9,
+            39,
+            9,
+            7,
+            SHARED_ROW_PITCH,
+            x_row_pitch,
+        )
+        cute.arch.sync_threads()
+
+        for reduced in cutlass.range_constexpr(REDUCED_COUNT):
+            grad_basis += focus_grad[
+                reduced * SHARED_ROW_PITCH + channel
+            ] * _radial_output_value(
+                local_values,
+                compact,
+                reduced,
+                channel,
+            )
+
+        lane = channel % 32
+        warp = channel // 32
+        subgroup = lane // WARP_REDUCTION_GROUP
+        subgroup_lane = lane % WARP_REDUCTION_GROUP
+        group = warp * GROUPS_PER_WARP + subgroup
+        for batch in cutlass.range_constexpr(
+            (COMPACT_WIDTH + GROUPS_PER_CTA - 1) // GROUPS_PER_CTA
+        ):
+            compact_idx = batch * GROUPS_PER_CTA + group
+            safe_compact_idx = compact_idx
+            if compact_idx >= COMPACT_WIDTH:
+                safe_compact_idx = cutlass.Int32(0)
+            grad_value = _warp_owned_grad_compact(
+                focus_grad,
+                local_values,
+                params.channel_basis,
+                safe_compact_idx,
+                subgroup_lane,
+                SHARED_ROW_PITCH,
+            )
+            if compact_idx < COMPACT_WIDTH:
+                if subgroup_lane == 0:
+                    params.grad_radial_compact[edge, compact_idx] = grad_value.to(
+                        params.grad_radial_compact.element_type
+                    )
+        # Every warp reads the full shared primal panel while reducing its
+        # assigned compact columns.  Do not let an early warp reuse that panel
+        # for input adjoints until all compact reductions have finished.
+        cute.arch.sync_threads()
+
+        basis = params.channel_basis[channel].to(cutlass.Float32)
+        for reduced in cutlass.range_constexpr(REDUCED_COUNT):
+            grad_value = cutlass.Float32(0.0)
+            if reduced < 4:
+                for output_row in cutlass.range_constexpr(4):
+                    grad_value += (
+                        focus_grad[output_row * SHARED_ROW_PITCH + channel]
+                        * compact[reduced * 4 + output_row]
+                    )
+            elif reduced < 7:
+                input_row = reduced - 4
+                for output_row in cutlass.range_constexpr(3):
+                    grad_value += (
+                        focus_grad[(4 + output_row) * SHARED_ROW_PITCH + channel]
+                        * compact[16 + input_row * 3 + output_row]
+                    )
+            else:
+                input_row = reduced - 7
+                for output_row in cutlass.range_constexpr(3):
+                    grad_value += (
+                        focus_grad[(7 + output_row) * SHARED_ROW_PITCH + channel]
+                        * compact[16 + input_row * 3 + output_row]
+                    )
+            local_values[reduced * SHARED_ROW_PITCH + channel] = grad_value * basis
+        cute.arch.sync_threads()
+
+        for batch in cutlass.range_constexpr(
+            (PACKED_WIGNER_VALUES + GROUPS_PER_CTA - 1) // GROUPS_PER_CTA
+        ):
+            panel_idx = batch * GROUPS_PER_CTA + group
+            safe_panel_idx = panel_idx
+            if panel_idx >= PACKED_WIGNER_VALUES:
+                safe_panel_idx = cutlass.Int32(0)
+            grad_d_value = _warp_owned_grad_d(
+                local_values,
+                x_values,
+                safe_panel_idx,
+                subgroup_lane,
+                SHARED_ROW_PITCH,
+                x_row_pitch,
+            )
+            if panel_idx < PACKED_WIGNER_VALUES:
+                if subgroup_lane == 0:
+                    params.grad_d_full[edge, panel_idx] = grad_d_value.to(
+                        params.grad_d_full.element_type
+                    )
+
+        grad_x_0 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            0,
+            0,
+            0,
+            1,
+            1,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_1 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            1,
+            0,
+            1,
+            3,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_2 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            1,
+            1,
+            1,
+            3,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_3 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            1,
+            2,
+            1,
+            3,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_4 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            2,
+            0,
+            10,
+            5,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_5 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            2,
+            1,
+            10,
+            5,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_6 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            2,
+            2,
+            10,
+            5,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_7 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            2,
+            3,
+            10,
+            5,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_8 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            2,
+            4,
+            10,
+            5,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_9 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            3,
+            0,
+            25,
+            7,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_10 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            3,
+            1,
+            25,
+            7,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_11 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            3,
+            2,
+            25,
+            7,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_12 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            3,
+            3,
+            25,
+            7,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_13 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            3,
+            4,
+            25,
+            7,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_14 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            3,
+            5,
+            25,
+            7,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        grad_x_15 += _grad_x_value(
+            local_values,
+            d_values,
+            channel,
+            3,
+            6,
+            25,
+            7,
+            3,
+            SHARED_ROW_PITCH,
+        )
+        cute.arch.sync_threads()
+
+    params.grad_x_wide[node, 0 * HIDDEN + channel] = grad_x_0
+    params.grad_x_wide[node, 1 * HIDDEN + channel] = grad_x_1
+    params.grad_x_wide[node, 2 * HIDDEN + channel] = grad_x_2
+    params.grad_x_wide[node, 3 * HIDDEN + channel] = grad_x_3
+    params.grad_x_wide[node, 4 * HIDDEN + channel] = grad_x_4
+    params.grad_x_wide[node, 5 * HIDDEN + channel] = grad_x_5
+    params.grad_x_wide[node, 6 * HIDDEN + channel] = grad_x_6
+    params.grad_x_wide[node, 7 * HIDDEN + channel] = grad_x_7
+    params.grad_x_wide[node, 8 * HIDDEN + channel] = grad_x_8
+    params.grad_x_wide[node, 9 * HIDDEN + channel] = grad_x_9
+    params.grad_x_wide[node, 10 * HIDDEN + channel] = grad_x_10
+    params.grad_x_wide[node, 11 * HIDDEN + channel] = grad_x_11
+    params.grad_x_wide[node, 12 * HIDDEN + channel] = grad_x_12
+    params.grad_x_wide[node, 13 * HIDDEN + channel] = grad_x_13
+    params.grad_x_wide[node, 14 * HIDDEN + channel] = grad_x_14
+    params.grad_x_wide[node, 15 * HIDDEN + channel] = grad_x_15
+    params.grad_basis_by_node[node, channel] = grad_basis
+
+
+@cute.kernel
+def _reduce_channel_basis_kernel(
+    grad_basis_by_node: cute.Tensor,
+    grad_channel_basis: cute.Tensor,
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    channel, _, _ = cute.arch.block_idx()
+    value = cutlass.Float32(0.0)
+    for node in cutlass.range(
+        tidx,
+        grad_basis_by_node.shape[0],
+        REDUCE_THREADS,
+        unroll=1,
+    ):
+        value += grad_basis_by_node[node, channel].to(cutlass.Float32)
+    value = cute.arch.warp_reduction_sum(value)
+    if tidx == 0:
+        grad_channel_basis[channel] = value
+
+
+def _fake(dtype, shape: tuple, stride_order: tuple[int, ...]):
+    return make_fake_compact_tensor(
+        dtype,
+        shape,
+        stride_order=stride_order,
+        **FAKE_TENSOR_KW,
+    )
+
+
+@device_aware_lru_cache(maxsize=8)
+def _compiled_backward(
+    device_index: int,
+    compute_capability: tuple[int, int],
+) -> Callable:
+    if compute_capability != (9, 0):
+        raise RuntimeError("the direct split Phase-A adjoint is sm_90-only")
+    edge_count = cute.sym_int64()
+    node_count = cute.sym_int64()
+    source_ptr_count = cute.sym_int64()
+    with torch.cuda.device(device_index):
+        return cute.compile(
+            _phase_a_split_backward_jit,
+            _fake(cutlass.Float32, (FOCUS_COUNT, edge_count, M0_WIDTH), (2, 1, 0)),
+            _fake(
+                cutlass.Float32,
+                (FOCUS_COUNT, edge_count, M1_WIDTH, 2),
+                (3, 2, 1, 0),
+            ),
+            _fake(cutlass.Float32, (edge_count, COMPACT_WIDTH), (1, 0)),
+            _fake(cutlass.Float32, (HIDDEN,), (0,)),
+            _fake(cutlass.Float32, (node_count, DEGREE_COUNT * HIDDEN), (1, 0)),
+            _fake(cutlass.Int32, (edge_count,), (0,)),
+            _fake(cutlass.Int32, (source_ptr_count,), (0,)),
+            _fake(cutlass.Float32, (edge_count, PACKED_WIGNER_VALUES), (1, 0)),
+            _fake(cutlass.Float32, (node_count, DEGREE_COUNT * HIDDEN), (1, 0)),
+            _fake(cutlass.Float32, (edge_count, PACKED_WIGNER_VALUES), (1, 0)),
+            _fake(cutlass.Float32, (edge_count, COMPACT_WIDTH), (1, 0)),
+            _fake(cutlass.Float32, (node_count, HIDDEN), (1, 0)),
+            _fake(cutlass.Float32, (HIDDEN,), (0,)),
+            make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+
+
+def allocate_neo_phase_a_persistent_complex_backward(
+    *,
+    edge_count: int,
+    node_count: int,
+    device: torch.device,
+) -> tuple[
+    NeoPhaseAPersistentComplexAdjoints,
+    NeoPhaseAPersistentComplexBackwardWorkspace,
+]:
+    """Allocate caller-owned outputs and the bounded basis-reduction workspace."""
+    if edge_count <= 0 or node_count <= 0:
+        raise ValueError("the Phase-A adjoint requires N > 0 and E > 0")
+    opts = {"device": device, "dtype": torch.float32}
+    outputs = NeoPhaseAPersistentComplexAdjoints(
+        grad_x_wide=torch.empty(
+            (node_count, DEGREE_COUNT, HIDDEN),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+        grad_d_full=torch.empty(
+            (edge_count, PACKED_WIGNER_VALUES),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+        grad_radial_compact=torch.empty(
+            (edge_count, COMPACT_WIDTH),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+        grad_channel_basis=torch.empty(
+            (HIDDEN,),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+    )
+    workspace = NeoPhaseAPersistentComplexBackwardWorkspace(
+        grad_basis_by_node=torch.empty(
+            (node_count, HIDDEN),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        )
+    )
+    return outputs, workspace
+
+
+def _expect(
+    name: str,
+    tensor: Tensor,
+    shape: tuple[int, ...],
+    device: torch.device,
+    dtype: torch.dtype,
+) -> None:
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+    if tensor.device != device or tensor.dtype != dtype or not tensor.is_contiguous():
+        raise ValueError(f"{name} must be contiguous {dtype} on {device}")
+
+
+def run_neo_phase_a_persistent_complex_backward_fp32(
+    *,
+    grad_state: NeoPersistentComplexState,
+    radial_compact: Tensor,
+    channel_basis: Tensor,
+    x_wide: Tensor,
+    source_order: Tensor,
+    source_ptr: Tensor,
+    d_full: Tensor,
+    outputs: NeoPhaseAPersistentComplexAdjoints | None = None,
+    workspace: NeoPhaseAPersistentComplexBackwardWorkspace | None = None,
+) -> NeoPhaseAPersistentComplexAdjoints:
+    """Run the direct split-gradient Phase-A input adjoint."""
+    if torch.is_grad_enabled():
+        raise RuntimeError("the explicit Phase-A adjoint must run under no_grad")
+    validate_neo_persistent_complex_state(grad_state, name="grad_state")
+    device = grad_state.m0.device
+    edge_count = grad_state.edge_count
+    if x_wide.ndim != 3 or tuple(x_wide.shape[1:]) != (DEGREE_COUNT, HIDDEN):
+        raise ValueError("x_wide must have shape (N,16,64)")
+    node_count = x_wide.shape[0]
+    if node_count <= 0:
+        raise ValueError("the Phase-A adjoint requires N > 0")
+
+    specs = (
+        ("radial_compact", radial_compact, (edge_count, COMPACT_WIDTH), torch.float32),
+        ("channel_basis", channel_basis, (HIDDEN,), torch.float32),
+        ("x_wide", x_wide, (node_count, DEGREE_COUNT, HIDDEN), torch.float32),
+        ("source_order", source_order, (edge_count,), torch.int32),
+        ("source_ptr", source_ptr, (node_count + 1,), torch.int32),
+        ("d_full", d_full, (edge_count, PACKED_WIGNER_VALUES), torch.float32),
+    )
+    for name, tensor, shape, dtype in specs:
+        _expect(name, tensor, shape, device, dtype)
+    torch._assert_async(
+        source_ptr[0] == 0,
+        "SM90 Phase-A backward requires source_ptr[0] == 0",
+    )
+    torch._assert_async(
+        source_ptr[-1] == edge_count,
+        "SM90 Phase-A backward requires source_ptr[-1] == edge_count",
+    )
+    if source_ptr.numel() > 1:
+        torch._assert_async(
+            torch.all(source_ptr[1:] >= source_ptr[:-1]),
+            "SM90 Phase-A backward requires nondecreasing source_ptr",
+        )
+    if source_order.numel() != 0:
+        torch._assert_async(
+            torch.all((source_order >= 0) & (source_order < edge_count)),
+            "SM90 Phase-A backward requires source_order entries in [0, E)",
+        )
+
+    if outputs is None or workspace is None:
+        if outputs is not None or workspace is not None:
+            raise ValueError("outputs and workspace must be supplied together")
+        outputs, workspace = allocate_neo_phase_a_persistent_complex_backward(
+            edge_count=edge_count,
+            node_count=node_count,
+            device=device,
+        )
+    output_specs = (
+        ("grad_x_wide", outputs.grad_x_wide, (node_count, DEGREE_COUNT, HIDDEN)),
+        ("grad_d_full", outputs.grad_d_full, (edge_count, PACKED_WIGNER_VALUES)),
+        (
+            "grad_radial_compact",
+            outputs.grad_radial_compact,
+            (edge_count, COMPACT_WIDTH),
+        ),
+        ("grad_channel_basis", outputs.grad_channel_basis, (HIDDEN,)),
+        (
+            "grad_basis_by_node",
+            workspace.grad_basis_by_node,
+            (node_count, HIDDEN),
+        ),
+    )
+    for name, tensor, shape in output_specs:
+        _expect(name, tensor, shape, device, torch.float32)
+
+    if torch.backends.cuda.matmul.allow_tf32:
+        raise RuntimeError("strict FP32 requires allow_tf32=False")
+    if torch.get_float32_matmul_precision() != "highest":
+        raise RuntimeError("strict FP32 requires float32 matmul precision 'highest'")
+    device_index = device.index
+    if device_index is None:
+        raise RuntimeError("the direct split Phase-A adjoint requires CUDA")
+    compute_capability = tuple(torch.cuda.get_device_capability(device_index))
+    kernel = _compiled_backward(device_index, compute_capability)
+    grad_m1_ri = torch.view_as_real(grad_state.m1)
+    if not grad_m1_ri.is_contiguous():
+        raise ValueError("grad_state.m1 must expose interleaved real/imag storage")
+    with torch.cuda.device(device):
+        kernel(
+            grad_state.m0,
+            grad_m1_ri,
+            radial_compact,
+            channel_basis,
+            x_wide.view(node_count, DEGREE_COUNT * HIDDEN),
+            source_order,
+            source_ptr,
+            d_full,
+            outputs.grad_x_wide.view(node_count, DEGREE_COUNT * HIDDEN),
+            outputs.grad_d_full,
+            outputs.grad_radial_compact,
+            workspace.grad_basis_by_node,
+            outputs.grad_channel_basis,
+        )
+    return outputs
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py
new file mode 100644
index 0000000000..bb8f93eda2
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py
@@ -0,0 +1,863 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Grouped SM90 final Phase-C/attention adjoint with split-state gather.
+
+Four independent 64-thread groups process four CSR edges
+per round while sharing one destination node's reverse-linear ``b0/b1``
+panels. Each ``(edge, focus)`` warp also retains the expanded
+``grad_m0`` and complex ``grad_m1`` accumulators in registers. The same panel
+loads therefore serve both the Phase-C scale contraction and split-state
+adjoint, replacing the separate ``_expanded_adjoint_gather`` launch.
+
+The focus competition, segmented attention softmax, envelope, and Q/K
+adjoints remain fused. No edge-sized temporary is added.
+Source-K is atomically accumulated, so callers must clear ``grad_k_node``
+before every invocation.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+from ..wigner_layout import (
+    PACKED_VALUE_COUNT,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN202, TC002
+
+FOCUS_COUNT = 2
+CHANNELS = 32
+DEGREE_COUNT = 16
+M0_WIDTH = 128
+M1_WIDTH = 96
+PACKED_WIGNER_VALUES = PACKED_VALUE_COUNT
+MAX_EDGES_PER_NODE = 256
+QK_SCALE = CHANNELS**-0.5
+B0_NODE_VALUES = FOCUS_COUNT * DEGREE_COUNT * M0_WIDTH
+B1_NODE_VALUES = FOCUS_COUNT * (DEGREE_COUNT - 1) * M1_WIDTH * 2
+PACKED_M0 = (0, 1, 2, 3, 10, 11, 12, 13, 14, 25, 26, 27, 28, 29, 30, 31)
+PACKED_RE = (4, 5, 6, 15, 16, 17, 18, 19, 32, 33, 34, 35, 36, 37, 38)
+PACKED_IM = (7, 8, 9, 20, 21, 22, 23, 24, 39, 40, 41, 42, 43, 44, 45)
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+EDGE_GROUPS = 4
+GROUP_WIDTH = FOCUS_COUNT * CHANNELS
+THREADS = EDGE_GROUPS * GROUP_WIDTH
+M0_ROWS = M0_WIDTH // CHANNELS
+M1_ROWS = M1_WIDTH // CHANNELS
+
+
+@cute.jit
+def _warp_sum(value):
+    return cute.arch.warp_reduction_sum(value)
+
+
+def _fake(dtype, shape: tuple[object, ...], stride_order: tuple[int, ...]):
+    return make_fake_compact_tensor(
+        dtype,
+        shape,
+        stride_order=stride_order,
+        **FAKE_TENSOR_KW,
+    )
+
+
+def _require_tensor(
+    name: str,
+    tensor: torch.Tensor,
+    shape: tuple[int, ...],
+    dtype: torch.dtype,
+    device: torch.device,
+) -> None:
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+    if (
+        tensor.dtype != dtype
+        or tensor.device != device
+        or not tensor.is_cuda
+        or not tensor.is_contiguous()
+        or tensor.data_ptr() % 16
+    ):
+        raise ValueError(
+            f"{name} must be contiguous, 16-byte-aligned {dtype} on {device}"
+        )
+
+
+__all__ = [
+    "GroupedExpandedFinalPhaseCAttentionAdjointOutputs",
+    "allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs",
+    "compile_grouped_expanded_final_phase_c_attention_adjoint",
+    "run_grouped_expanded_final_phase_c_attention_adjoint",
+]
+
+
+@dataclass(frozen=True)
+class GroupedExpandedFinalPhaseCAttentionAdjointOutputs:
+    """True adjoints emitted by the one-launch boundary."""
+
+    grad_m0: torch.Tensor
+    grad_m1: torch.Tensor
+    grad_dt: torch.Tensor
+    grad_logits: torch.Tensor
+    grad_edge: torch.Tensor
+    grad_focus_src: torch.Tensor
+    grad_q_node: torch.Tensor
+    grad_k_node: torch.Tensor
+
+
+@cute.jit
+def _grouped_expanded_adjoint_jit(
+    b0: cute.Tensor,
+    b1_ri: cute.Tensor,
+    m0: cute.Tensor,
+    m1_ri: cute.Tensor,
+    dt_packed: cute.Tensor,
+    beta: cute.Tensor,
+    alpha: cute.Tensor,
+    focus_alpha: cute.Tensor,
+    focus_src: cute.Tensor,
+    focus_weight: cute.Tensor,
+    focus_scale: cute.Tensor,
+    q_node: cute.Tensor,
+    k_node: cute.Tensor,
+    edge_gate: cute.Tensor,
+    src: cute.Tensor,
+    dst_ptr: cute.Tensor,
+    grad_m0: cute.Tensor,
+    grad_m1_ri: cute.Tensor,
+    grad_dt: cute.Tensor,
+    grad_logits: cute.Tensor,
+    grad_edge: cute.Tensor,
+    grad_focus_src: cute.Tensor,
+    grad_q_node: cute.Tensor,
+    grad_k_node: cute.Tensor,
+    focus_eps: cutlass.Constexpr[float],
+    focus_tau: cutlass.Constexpr[float],
+    label_smoothing: cutlass.Constexpr[float],
+    qk_scale: cutlass.Constexpr[float],
+    use_focus_norm: cutlass.Constexpr[bool],
+    stream: CUstream,
+):
+    _grouped_expanded_adjoint_kernel(
+        b0,
+        b1_ri,
+        m0,
+        m1_ri,
+        dt_packed,
+        beta,
+        alpha,
+        focus_alpha,
+        focus_src,
+        focus_weight,
+        focus_scale,
+        q_node,
+        k_node,
+        edge_gate,
+        src,
+        dst_ptr,
+        grad_m0,
+        grad_m1_ri,
+        grad_dt,
+        grad_logits,
+        grad_edge,
+        grad_focus_src,
+        grad_q_node,
+        grad_k_node,
+        focus_eps,
+        focus_tau,
+        label_smoothing,
+        qk_scale,
+        use_focus_norm,
+    ).launch(
+        grid=[dst_ptr.shape[0] - 1, 1, 1],
+        block=[THREADS, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def _grouped_expanded_adjoint_kernel(
+    b0: cute.Tensor,
+    b1_ri: cute.Tensor,
+    m0: cute.Tensor,
+    m1_ri: cute.Tensor,
+    dt_packed: cute.Tensor,
+    beta: cute.Tensor,
+    alpha: cute.Tensor,
+    focus_alpha: cute.Tensor,
+    focus_src: cute.Tensor,
+    focus_weight: cute.Tensor,
+    focus_scale: cute.Tensor,
+    q_node: cute.Tensor,
+    k_node: cute.Tensor,
+    edge_gate: cute.Tensor,
+    src: cute.Tensor,
+    dst_ptr: cute.Tensor,
+    grad_m0: cute.Tensor,
+    grad_m1_ri: cute.Tensor,
+    grad_dt: cute.Tensor,
+    grad_logits: cute.Tensor,
+    grad_edge: cute.Tensor,
+    grad_focus_src: cute.Tensor,
+    grad_q_node: cute.Tensor,
+    grad_k_node: cute.Tensor,
+    focus_eps: cutlass.Constexpr[float],
+    focus_tau: cutlass.Constexpr[float],
+    label_smoothing: cutlass.Constexpr[float],
+    qk_scale: cutlass.Constexpr[float],
+    use_focus_norm: cutlass.Constexpr[bool],
+):
+    tidx, _, _ = cute.arch.thread_idx()
+    node, _, _ = cute.arch.block_idx()
+    group = tidx // GROUP_WIDTH
+    local = tidx - group * GROUP_WIDTH
+    focus = local // CHANNELS
+    lane = local - focus * CHANNELS
+    node_lo = dst_ptr[node]
+    node_hi = dst_ptr[node + 1]
+    degree = node_hi - node_lo
+    rounds = (degree + EDGE_GROUPS - 1) // EDGE_GROUPS
+
+    smem = cutlass.utils.SmemAllocator()
+    b0_node = smem.allocate_tensor(cutlass.Float32, B0_NODE_VALUES)
+    b1_node = smem.allocate_tensor(cutlass.Float32, B1_NODE_VALUES)
+    beta_adjoint = smem.allocate_tensor(
+        cutlass.Float32,
+        MAX_EDGES_PER_NODE * FOCUS_COUNT,
+    )
+    dt_partial = smem.allocate_tensor(
+        cutlass.Float32,
+        EDGE_GROUPS * FOCUS_COUNT * PACKED_WIGNER_VALUES,
+    )
+    softmax_dot = smem.allocate_tensor(cutlass.Float32, FOCUS_COUNT)
+    q_partial = smem.allocate_tensor(
+        cutlass.Float32,
+        EDGE_GROUPS * FOCUS_COUNT * CHANNELS,
+    )
+
+    for linear in cutlass.range(tidx, B0_NODE_VALUES, THREADS, unroll=1):
+        quotient = linear // M0_WIDTH
+        feature = linear - quotient * M0_WIDTH
+        panel_focus = quotient // DEGREE_COUNT
+        q = quotient - panel_focus * DEGREE_COUNT
+        b0_node[linear] = b0[panel_focus, q, node, feature].to(cutlass.Float32)
+    for linear in cutlass.range(tidx, B1_NODE_VALUES, THREADS, unroll=1):
+        quotient = linear // 2
+        component = linear - quotient * 2
+        feature_quotient = quotient // M1_WIDTH
+        feature = quotient - feature_quotient * M1_WIDTH
+        panel_focus = feature_quotient // (DEGREE_COUNT - 1)
+        q1 = feature_quotient - panel_focus * (DEGREE_COUNT - 1)
+        b1_node[linear] = b1_ri[
+            panel_focus,
+            q1,
+            node,
+            feature,
+            component,
+        ].to(cutlass.Float32)
+    cute.arch.sync_threads()
+
+    m0_fragment = cute.make_rmem_tensor(
+        cute.make_layout((M0_ROWS,), stride=(1,)),
+        cutlass.Float32,
+    )
+    m1_real_fragment = cute.make_rmem_tensor(
+        cute.make_layout((M1_ROWS,), stride=(1,)),
+        cutlass.Float32,
+    )
+    m1_imag_fragment = cute.make_rmem_tensor(
+        cute.make_layout((M1_ROWS,), stride=(1,)),
+        cutlass.Float32,
+    )
+    grad_m0_fragment = cute.make_rmem_tensor(
+        cute.make_layout((M0_ROWS,), stride=(1,)),
+        cutlass.Float32,
+    )
+    grad_m1_real_fragment = cute.make_rmem_tensor(
+        cute.make_layout((M1_ROWS,), stride=(1,)),
+        cutlass.Float32,
+    )
+    grad_m1_imag_fragment = cute.make_rmem_tensor(
+        cute.make_layout((M1_ROWS,), stride=(1,)),
+        cutlass.Float32,
+    )
+
+    # All four groups execute the same round count, making CTA barriers safe.
+    for edge_round in cutlass.range(rounds, unroll=1):
+        edge_slot = edge_round * EDGE_GROUPS + group
+        edge = node_lo + edge_slot
+        edge_active = edge_slot < degree
+
+        for row in cutlass.range_constexpr(M0_ROWS):
+            feature = row * CHANNELS + lane
+            value = cutlass.Float32(0.0)
+            if edge_active:
+                value = m0[focus, edge, feature].to(cutlass.Float32)
+            m0_fragment[row] = value
+            grad_m0_fragment[row] = cutlass.Float32(0.0)
+        for row in cutlass.range_constexpr(M1_ROWS):
+            feature = row * CHANNELS + lane
+            real = cutlass.Float32(0.0)
+            imag = cutlass.Float32(0.0)
+            if edge_active:
+                real = m1_ri[focus, edge, feature, 0].to(cutlass.Float32)
+                imag = m1_ri[focus, edge, feature, 1].to(cutlass.Float32)
+            m1_real_fragment[row] = real
+            m1_imag_fragment[row] = imag
+            grad_m1_real_fragment[row] = cutlass.Float32(0.0)
+            grad_m1_imag_fragment[row] = cutlass.Float32(0.0)
+
+        beta_value = cutlass.Float32(0.0)
+        if edge_active:
+            beta_value = beta[edge, focus].to(cutlass.Float32)
+        grad_beta_value = cutlass.Float32(0.0)
+        dt_base = (group * FOCUS_COUNT + focus) * PACKED_WIGNER_VALUES
+
+        # Ascending q order matches the existing expanded gather exactly.
+        for q in cutlass.range_constexpr(DEGREE_COUNT):
+            panel0 = PACKED_M0[q]
+            dt0 = cutlass.Float32(0.0)
+            if edge_active:
+                dt0 = dt_packed[edge, panel0].to(cutlass.Float32)
+            scalar0_lane = cutlass.Float32(0.0)
+            for row in cutlass.range_constexpr(M0_ROWS):
+                feature = row * CHANNELS + lane
+                b0_offset = (focus * DEGREE_COUNT + q) * M0_WIDTH + feature
+                b_value = b0_node[b0_offset]
+                grad_m0_fragment[row] += beta_value * dt0 * b_value
+                scalar0_lane += b_value * m0_fragment[row]
+            scalar0 = _warp_sum(scalar0_lane)
+            if lane == 0:
+                dt_partial[dt_base + panel0] = beta_value * scalar0
+                grad_beta_value += dt0 * scalar0
+
+            if cutlass.const_expr(q > 0):
+                panel_re = PACKED_RE[q - 1]
+                panel_im = PACKED_IM[q - 1]
+                dt_re = cutlass.Float32(0.0)
+                dt_im = cutlass.Float32(0.0)
+                if edge_active:
+                    dt_re = dt_packed[edge, panel_re].to(cutlass.Float32)
+                    dt_im = dt_packed[edge, panel_im].to(cutlass.Float32)
+                scalar1_re_lane = cutlass.Float32(0.0)
+                scalar1_im_lane = cutlass.Float32(0.0)
+                for row in cutlass.range_constexpr(M1_ROWS):
+                    feature = row * CHANNELS + lane
+                    b1_offset = (
+                        (focus * (DEGREE_COUNT - 1) + q - 1) * M1_WIDTH + feature
+                    ) * 2
+                    br = b1_node[b1_offset]
+                    bi = b1_node[b1_offset + 1]
+                    grad_m1_real_fragment[row] += beta_value * (dt_re * br - dt_im * bi)
+                    grad_m1_imag_fragment[row] += beta_value * (dt_re * bi + dt_im * br)
+                    xr = m1_real_fragment[row]
+                    xi = m1_imag_fragment[row]
+                    scalar1_re_lane += br * xr + bi * xi
+                    scalar1_im_lane += br * xi - bi * xr
+                scalar1_re = _warp_sum(scalar1_re_lane)
+                scalar1_im = _warp_sum(scalar1_im_lane)
+                if lane == 0:
+                    dt_partial[dt_base + panel_re] = beta_value * scalar1_re
+                    dt_partial[dt_base + panel_im] = beta_value * scalar1_im
+                    grad_beta_value += dt_re * scalar1_re + dt_im * scalar1_im
+
+        if edge_active:
+            for row in cutlass.range_constexpr(M0_ROWS):
+                feature = row * CHANNELS + lane
+                grad_m0[focus, edge, feature] = grad_m0_fragment[row]
+            for row in cutlass.range_constexpr(M1_ROWS):
+                feature = row * CHANNELS + lane
+                grad_m1_ri[focus, edge, feature, 0] = grad_m1_real_fragment[row]
+                grad_m1_ri[focus, edge, feature, 1] = grad_m1_imag_fragment[row]
+        if lane == 0:
+            beta_adjoint[edge_slot * FOCUS_COUNT + focus] = grad_beta_value
+        cute.arch.sync_threads()
+
+        if tidx < EDGE_GROUPS * PACKED_WIGNER_VALUES:
+            output_group = tidx // PACKED_WIGNER_VALUES
+            panel = tidx - output_group * PACKED_WIGNER_VALUES
+            output_slot = edge_round * EDGE_GROUPS + output_group
+            if output_slot < degree:
+                output_edge = node_lo + output_slot
+                focus0 = output_group * FOCUS_COUNT * PACKED_WIGNER_VALUES
+                grad_dt[output_edge, panel] = (
+                    dt_partial[focus0 + panel]
+                    + dt_partial[focus0 + PACKED_WIGNER_VALUES + panel]
+                )
+        cute.arch.sync_threads()
+
+    # Match the forward path's ascending-edge segmented-softmax reduction order.
+    if group == 0 and lane == 0:
+        value = cutlass.Float32(0.0)
+        for edge_slot in cutlass.range(degree, unroll=1):
+            edge = node_lo + edge_slot
+            grad_attention = beta_adjoint[
+                edge_slot * FOCUS_COUNT + focus
+            ] * focus_alpha[edge, focus].to(cutlass.Float32)
+            value += alpha[edge, focus].to(cutlass.Float32) * grad_attention
+        softmax_dot[focus] = value
+    cute.arch.sync_threads()
+
+    keep = cutlass.Float32(1.0 - label_smoothing)
+    smooth = cutlass.Float32(label_smoothing / FOCUS_COUNT)
+    inv_tau = cutlass.Float32(1.0 / focus_tau)
+    grad_q = cutlass.Float32(0.0)
+
+    for edge_slot in cutlass.range(group, degree, EDGE_GROUPS, unroll=1):
+        edge = node_lo + edge_slot
+        source = src[edge]
+        grad_beta_value = beta_adjoint[edge_slot * FOCUS_COUNT + focus]
+        alpha_value = alpha[edge, focus].to(cutlass.Float32)
+        focus_value = focus_alpha[edge, focus].to(cutlass.Float32)
+        grad_attention = grad_beta_value * focus_value
+        grad_logit = alpha_value * (grad_attention - softmax_dot[focus])
+        if lane == 0:
+            grad_logits[edge, focus] = grad_logit
+
+        scaled_grad_logit = grad_logit * cutlass.Float32(qk_scale)
+        grad_q += scaled_grad_logit * k_node[source, focus, lane].to(cutlass.Float32)
+        grad_k = scaled_grad_logit * q_node[node, focus, lane].to(cutlass.Float32)
+        k_offset = (source * FOCUS_COUNT + focus) * CHANNELS + lane
+        k_ptr = grad_k_node.iterator + k_offset
+        cute.arch.atomic_add(
+            k_ptr.llvm_ptr,
+            grad_k,
+            sem="relaxed",
+            scope="gpu",
+        )
+
+        probability0 = (focus_alpha[edge, 0].to(cutlass.Float32) - smooth) / keep
+        probability1 = (focus_alpha[edge, 1].to(cutlass.Float32) - smooth) / keep
+        focus_grad0 = (
+            beta_adjoint[edge_slot * FOCUS_COUNT]
+            * alpha[edge, 0].to(cutlass.Float32)
+            * keep
+        )
+        focus_grad1 = (
+            beta_adjoint[edge_slot * FOCUS_COUNT + 1]
+            * alpha[edge, 1].to(cutlass.Float32)
+            * keep
+        )
+        focus_dot = focus_grad0 * probability0 + focus_grad1 * probability1
+        probability = probability0
+        focus_grad_probability = focus_grad0
+        if focus == 1:
+            probability = probability1
+            focus_grad_probability = focus_grad1
+        focus_grad_logit = probability * (focus_grad_probability - focus_dot) * inv_tau
+
+        weight = focus_weight[lane, focus].to(cutlass.Float32)
+        if cutlass.const_expr(use_focus_norm):
+            focus_value_raw = focus_src[edge, focus, lane].to(cutlass.Float32)
+            scale = focus_scale[focus, lane].to(cutlass.Float32)
+            inv_rms = cute.rsqrt(
+                _warp_sum(focus_value_raw * focus_value_raw) / cutlass.Float32(CHANNELS)
+                + cutlass.Float32(focus_eps)
+            )
+            grad_scaled = focus_grad_logit * weight * scale
+            coeff = _warp_sum(grad_scaled * focus_value_raw) / cutlass.Float32(CHANNELS)
+            grad_focus_src[focus, edge, lane] = (
+                grad_scaled * inv_rms
+                - focus_value_raw * inv_rms * inv_rms * inv_rms * coeff
+            )
+        else:
+            grad_focus_src[focus, edge, lane] = focus_grad_logit * weight
+
+        if local == 0:
+            grad_attention0 = beta_adjoint[edge_slot * FOCUS_COUNT] * focus_alpha[
+                edge, 0
+            ].to(cutlass.Float32)
+            grad_attention1 = beta_adjoint[edge_slot * FOCUS_COUNT + 1] * focus_alpha[
+                edge, 1
+            ].to(cutlass.Float32)
+            grad_logit0 = alpha[edge, 0].to(cutlass.Float32) * (
+                grad_attention0 - softmax_dot[0]
+            )
+            grad_logit1 = alpha[edge, 1].to(cutlass.Float32) * (
+                grad_attention1 - softmax_dot[1]
+            )
+            gate = edge_gate[edge].to(cutlass.Float32)
+            value = cutlass.Float32(0.0)
+            if gate > cutlass.Float32(0.0):
+                value = cutlass.Float32(2.0) * (grad_logit0 + grad_logit1) / gate
+            grad_edge[edge] = value
+
+    q_partial[group * GROUP_WIDTH + local] = grad_q
+    cute.arch.sync_threads()
+    if group == 0:
+        total = cutlass.Float32(0.0)
+        for partial_group in cutlass.range_constexpr(EDGE_GROUPS):
+            total += q_partial[partial_group * GROUP_WIDTH + local]
+        grad_q_node[node, focus, lane] = total
+
+
+@device_aware_lru_cache(maxsize=8)
+def compile_grouped_expanded_final_phase_c_attention_adjoint(
+    focus_eps: float,
+    focus_tau: float,
+    label_smoothing: float,
+    qk_scale: float = QK_SCALE,
+    use_focus_norm: bool = True,
+) -> Callable:
+    """Compile the fixed four-group SM90 schedule."""
+    edges = cute.sym_int64()
+    nodes = cute.sym_int64()
+    b0 = _fake(
+        cutlass.Float32,
+        (FOCUS_COUNT, DEGREE_COUNT, nodes, M0_WIDTH),
+        (3, 2, 1, 0),
+    )
+    b1 = _fake(
+        cutlass.Float32,
+        (FOCUS_COUNT, DEGREE_COUNT - 1, nodes, M1_WIDTH, 2),
+        (4, 3, 2, 1, 0),
+    )
+    m0 = _fake(cutlass.Float32, (FOCUS_COUNT, edges, M0_WIDTH), (2, 1, 0))
+    m1 = _fake(
+        cutlass.Float32,
+        (FOCUS_COUNT, edges, M1_WIDTH, 2),
+        (3, 2, 1, 0),
+    )
+    dt = _fake(cutlass.Float32, (edges, PACKED_WIGNER_VALUES), (1, 0))
+    edge_focus = _fake(cutlass.Float32, (edges, FOCUS_COUNT), (1, 0))
+    focus_src = _fake(
+        cutlass.Float32,
+        (edges, FOCUS_COUNT, CHANNELS),
+        (2, 1, 0),
+    )
+    focus_weight = _fake(cutlass.Float32, (CHANNELS, FOCUS_COUNT), (1, 0))
+    focus_scale = _fake(cutlass.Float32, (FOCUS_COUNT, CHANNELS), (1, 0))
+    node_focus = _fake(
+        cutlass.Float32,
+        (nodes, FOCUS_COUNT, CHANNELS),
+        (2, 1, 0),
+    )
+    edge_scalar = _fake(cutlass.Float32, (edges,), (0,))
+    edge_index = _fake(cutlass.Int32, (edges,), (0,))
+    dst_ptr = _fake(cutlass.Int32, (cute.sym_int64(),), (0,))
+    grad_focus_src = _fake(
+        cutlass.Float32,
+        (FOCUS_COUNT, edges, CHANNELS),
+        (2, 1, 0),
+    )
+    return cute.compile(
+        _grouped_expanded_adjoint_jit,
+        b0,
+        b1,
+        m0,
+        m1,
+        dt,
+        edge_focus,
+        edge_focus,
+        edge_focus,
+        focus_src,
+        focus_weight,
+        focus_scale,
+        node_focus,
+        node_focus,
+        edge_scalar,
+        edge_index,
+        dst_ptr,
+        _fake(cutlass.Float32, (FOCUS_COUNT, edges, M0_WIDTH), (2, 1, 0)),
+        _fake(
+            cutlass.Float32,
+            (FOCUS_COUNT, edges, M1_WIDTH, 2),
+            (3, 2, 1, 0),
+        ),
+        _fake(cutlass.Float32, (edges, PACKED_WIGNER_VALUES), (1, 0)),
+        _fake(cutlass.Float32, (edges, FOCUS_COUNT), (1, 0)),
+        _fake(cutlass.Float32, (edges,), (0,)),
+        grad_focus_src,
+        _fake(cutlass.Float32, (nodes, FOCUS_COUNT, CHANNELS), (2, 1, 0)),
+        _fake(cutlass.Float32, (nodes, FOCUS_COUNT, CHANNELS), (2, 1, 0)),
+        float(focus_eps),
+        float(focus_tau),
+        float(label_smoothing),
+        float(qk_scale),
+        bool(use_focus_norm),
+        make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs(
+    *,
+    edge_count: int,
+    node_count: int,
+    device: torch.device,
+    grad_m0: torch.Tensor | None = None,
+    grad_m1: torch.Tensor | None = None,
+) -> GroupedExpandedFinalPhaseCAttentionAdjointOutputs:
+    """Allocate outputs, optionally reusing dead split-state storage."""
+    if (grad_m0 is None) != (grad_m1 is None):
+        raise ValueError("grad_m0 and grad_m1 must be supplied together")
+    opts = {"device": device, "dtype": torch.float32}
+    return GroupedExpandedFinalPhaseCAttentionAdjointOutputs(
+        grad_m0=(
+            grad_m0
+            if grad_m0 is not None
+            else torch.empty(
+                (FOCUS_COUNT, edge_count, M0_WIDTH),
+                device=opts["device"],
+                dtype=opts["dtype"],
+            )
+        ),
+        grad_m1=(
+            grad_m1
+            if grad_m1 is not None
+            else torch.empty(
+                (FOCUS_COUNT, edge_count, M1_WIDTH),
+                device=device,
+                dtype=torch.complex64,
+            )
+        ),
+        grad_dt=torch.empty(
+            (edge_count, PACKED_WIGNER_VALUES),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+        grad_logits=torch.empty(
+            (edge_count, FOCUS_COUNT),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+        grad_edge=torch.empty(
+            (edge_count,),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+        grad_focus_src=torch.empty(
+            (FOCUS_COUNT, edge_count, CHANNELS),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+        grad_q_node=torch.empty(
+            (node_count, FOCUS_COUNT, CHANNELS),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+        grad_k_node=torch.zeros(
+            (node_count, FOCUS_COUNT, CHANNELS),
+            device=opts["device"],
+            dtype=opts["dtype"],
+        ),
+    )
+
+
+def run_grouped_expanded_final_phase_c_attention_adjoint(
+    *,
+    b0: torch.Tensor,
+    b1: torch.Tensor,
+    m0: torch.Tensor,
+    m1: torch.Tensor,
+    dt_packed: torch.Tensor,
+    beta: torch.Tensor,
+    alpha: torch.Tensor,
+    focus_alpha: torch.Tensor,
+    focus_src: torch.Tensor,
+    focus_weight: torch.Tensor,
+    focus_scale: torch.Tensor,
+    q_node: torch.Tensor,
+    k_node: torch.Tensor,
+    edge_gate: torch.Tensor,
+    src: torch.Tensor,
+    dst_ptr: torch.Tensor,
+    focus_eps: float,
+    focus_tau: float,
+    label_smoothing: float,
+    qk_scale: float = QK_SCALE,
+    use_focus_norm: bool = True,
+    outputs: GroupedExpandedFinalPhaseCAttentionAdjointOutputs | None = None,
+) -> GroupedExpandedFinalPhaseCAttentionAdjointOutputs:
+    """Run the adjoint; reusable outputs require clearing ``grad_k_node``."""
+    device = b0.device
+    if device.type != "cuda" or tuple(torch.cuda.get_device_capability(device)) != (
+        9,
+        0,
+    ):
+        raise RuntimeError("grouped expanded Phase-C adjoint requires SM90")
+    if torch.backends.cuda.matmul.allow_tf32:
+        raise RuntimeError("strict FP32 requires allow_tf32=False")
+    if torch.get_float32_matmul_precision() != "highest":
+        raise RuntimeError("strict FP32 requires float32 matmul precision 'highest'")
+
+    node_count = int(dst_ptr.numel() - 1)
+    edge_count = int(src.numel())
+    expected_inputs = (
+        ("b0", b0, (FOCUS_COUNT, DEGREE_COUNT, node_count, M0_WIDTH), torch.float32),
+        (
+            "b1",
+            b1,
+            (FOCUS_COUNT, DEGREE_COUNT - 1, node_count, M1_WIDTH),
+            torch.complex64,
+        ),
+        ("m0", m0, (FOCUS_COUNT, edge_count, M0_WIDTH), torch.float32),
+        ("m1", m1, (FOCUS_COUNT, edge_count, M1_WIDTH), torch.complex64),
+        (
+            "dt_packed",
+            dt_packed,
+            (edge_count, PACKED_WIGNER_VALUES),
+            torch.float32,
+        ),
+        ("beta", beta, (edge_count, FOCUS_COUNT), torch.float32),
+        ("alpha", alpha, (edge_count, FOCUS_COUNT), torch.float32),
+        ("focus_alpha", focus_alpha, (edge_count, FOCUS_COUNT), torch.float32),
+        (
+            "focus_src",
+            focus_src,
+            (edge_count, FOCUS_COUNT, CHANNELS),
+            torch.float32,
+        ),
+        (
+            "focus_weight",
+            focus_weight,
+            (CHANNELS, FOCUS_COUNT),
+            torch.float32,
+        ),
+        (
+            "focus_scale",
+            focus_scale,
+            (FOCUS_COUNT, CHANNELS),
+            torch.float32,
+        ),
+        (
+            "q_node",
+            q_node,
+            (node_count, FOCUS_COUNT, CHANNELS),
+            torch.float32,
+        ),
+        (
+            "k_node",
+            k_node,
+            (node_count, FOCUS_COUNT, CHANNELS),
+            torch.float32,
+        ),
+        ("edge_gate", edge_gate, (edge_count,), torch.float32),
+        ("src", src, (edge_count,), torch.int32),
+        ("dst_ptr", dst_ptr, (node_count + 1,), torch.int32),
+    )
+    for name, tensor, shape, dtype in expected_inputs:
+        _require_tensor(name, tensor, shape, dtype, device)
+    torch._assert_async(
+        dst_ptr[0] == 0,
+        "SM90 Phase-C backward requires dst_ptr[0] == 0",
+    )
+    torch._assert_async(
+        dst_ptr[-1] == edge_count,
+        "SM90 Phase-C backward requires dst_ptr[-1] == edge_count",
+    )
+    if dst_ptr.numel() > 1:
+        destination_degrees = dst_ptr[1:] - dst_ptr[:-1]
+        torch._assert_async(
+            torch.all(
+                (destination_degrees >= 0) & (destination_degrees <= MAX_EDGES_PER_NODE)
+            ),
+            "SM90 Phase-C backward requires destination degrees in "
+            f"[0, {MAX_EDGES_PER_NODE}]",
+        )
+
+    if outputs is None:
+        outputs = allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs(
+            edge_count=edge_count,
+            node_count=node_count,
+            device=device,
+        )
+    expected_outputs = (
+        (
+            "grad_m0",
+            outputs.grad_m0,
+            (FOCUS_COUNT, edge_count, M0_WIDTH),
+            torch.float32,
+        ),
+        (
+            "grad_m1",
+            outputs.grad_m1,
+            (FOCUS_COUNT, edge_count, M1_WIDTH),
+            torch.complex64,
+        ),
+        ("grad_dt", outputs.grad_dt, (edge_count, PACKED_WIGNER_VALUES), torch.float32),
+        ("grad_logits", outputs.grad_logits, (edge_count, FOCUS_COUNT), torch.float32),
+        ("grad_edge", outputs.grad_edge, (edge_count,), torch.float32),
+        (
+            "grad_focus_src",
+            outputs.grad_focus_src,
+            (FOCUS_COUNT, edge_count, CHANNELS),
+            torch.float32,
+        ),
+        (
+            "grad_q_node",
+            outputs.grad_q_node,
+            (node_count, FOCUS_COUNT, CHANNELS),
+            torch.float32,
+        ),
+        (
+            "grad_k_node",
+            outputs.grad_k_node,
+            (node_count, FOCUS_COUNT, CHANNELS),
+            torch.float32,
+        ),
+    )
+    for name, tensor, shape, dtype in expected_outputs:
+        _require_tensor(name, tensor, shape, dtype, device)
+
+    with torch.cuda.device(device):
+        compile_grouped_expanded_final_phase_c_attention_adjoint(
+            float(focus_eps),
+            float(focus_tau),
+            float(label_smoothing),
+            float(qk_scale),
+            bool(use_focus_norm),
+        )(
+            b0,
+            torch.view_as_real(b1),
+            m0,
+            torch.view_as_real(m1),
+            dt_packed,
+            beta,
+            alpha,
+            focus_alpha,
+            focus_src,
+            focus_weight,
+            focus_scale,
+            q_node,
+            k_node,
+            edge_gate,
+            src,
+            dst_ptr,
+            outputs.grad_m0,
+            torch.view_as_real(outputs.grad_m1),
+            outputs.grad_dt,
+            outputs.grad_logits,
+            outputs.grad_edge,
+            outputs.grad_focus_src,
+            outputs.grad_q_node,
+            outputs.grad_k_node,
+        )
+    return outputs
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/prefix.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/prefix.py
new file mode 100644
index 0000000000..c678069083
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/prefix.py
@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""In-place forward and reverse recurrences for the SM90 persistent prefix.
+
+The CuTe gate is elementwise alias-safe for ``out=residual``: every output
+element reads only the corresponding residual element, while all gate values
+are derived from the separate ``z`` tensors.  The input state may therefore
+become the running state without changing the recurrence.
+
+The reverse recurrence overwrites dead saved preactivations and aliases the
+running residual through ``torch.baddbmm`` to minimize edge-sized storage.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+import torch
+
+from .persistent import (
+    NeoPersistentComplexSaved,
+    NeoPersistentComplexState,
+    NeoPersistentComplexWeights,
+    _empty_state_like,
+    _run_gate_adjoint,
+    _run_gate_forward,
+    validate_neo_persistent_complex_state,
+)
+
+GATED_LAYERS = 2
+
+__all__ = [
+    "run_persistent_prefix_forward_inplace",
+    "run_persistent_prefix_input_adjoint_destructive_saved",
+]
+
+
+def _saved(z0: list[torch.Tensor], z1: list[torch.Tensor]) -> NeoPersistentComplexSaved:
+    return NeoPersistentComplexSaved(
+        z0=(z0[0], z0[1]),
+        z1=(z1[0], z1[1]),
+    )
+
+
+def run_persistent_prefix_forward_inplace(
+    state: NeoPersistentComplexState,
+    weights: NeoPersistentComplexWeights,
+) -> tuple[NeoPersistentComplexState, NeoPersistentComplexSaved]:
+    """Overwrite the caller-owned running state after each separate GEMM."""
+    validate_neo_persistent_complex_state(state)
+    saved_m0: list[torch.Tensor] = []
+    saved_m1: list[torch.Tensor] = []
+    for layer in range(GATED_LAYERS):
+        z = _empty_state_like(state)
+        torch.bmm(state.m0, weights.w0[layer], out=z.m0)
+        torch.bmm(state.m1, weights.wc[layer], out=z.m1)
+        saved_m0.append(z.m0)
+        saved_m1.append(z.m1)
+        _run_gate_forward(state, z, weights.gate[layer], state)
+    return state, _saved(saved_m0, saved_m1)
+
+
+def run_persistent_prefix_input_adjoint_destructive_saved(
+    grad_out: NeoPersistentComplexState,
+    saved: NeoPersistentComplexSaved,
+    weights: NeoPersistentComplexWeights,
+) -> NeoPersistentComplexState:
+    """Overwrite each dead saved preactivation with its exact gate adjoint.
+
+    The gate kernel stages the scalar row before writing any output.  Every
+    remaining preactivation element is read and then replaced by its own
+    adjoint, so input/output aliasing is safe.  Backward visits layer 1 before
+    layer 0; each saved state is therefore dead when it becomes ``grad_z``.
+    """
+    validate_neo_persistent_complex_state(grad_out, name="grad_out")
+    running = grad_out
+    for layer in range(GATED_LAYERS - 1, -1, -1):
+        grad_z = NeoPersistentComplexState(saved.z0[layer], saved.z1[layer])
+        _run_gate_adjoint(running, grad_z, weights.gate[layer], grad_z)
+        torch.baddbmm(
+            running.m0,
+            grad_z.m0,
+            weights.w0_h[layer],
+            out=running.m0,
+        )
+        torch.baddbmm(
+            running.m1,
+            grad_z.m1,
+            weights.wc_h[layer],
+            out=running.m1,
+        )
+    return running
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/radial.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/radial.py
new file mode 100644
index 0000000000..d8c256f474
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/radial.py
@@ -0,0 +1,310 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Compact radial state for the SM90 persistent-complex Neo SO2 path.
+
+The fused Phase-A kernel computes three logically separate values in one
+edge CTA: the packed-Wigner rotation, the 25-value compact radial map, and the
+64-value degree-zero radial attention feature.  A persistent-complex Phase A
+must not call that kernel merely to recover the latter two values because doing
+so would also allocate and write the discarded ``(E,2,10,32)`` real stack.
+
+This module retains the same FP32 reduction order for the two
+radial projections while omitting the dense SO2 output.  The resulting compact
+state is consumed directly by the split-complex Phase A and its adjoint.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    TYPE_CHECKING,
+)
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from ...compile_cache import (
+    device_aware_lru_cache,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Callable,
+    )
+
+
+# CuTe JIT functions use DSL-inferred argument and return types.
+# ruff: noqa: ANN001, ANN202, TC002
+
+RADIAL_WIDTH = 4 * 32
+COMPACT_WIDTH = 25
+ATTENTION_WIDTH = 64
+THREADS = 64
+FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True}
+
+__all__ = [
+    "project_neo_radial_input_adjoint_fp32",
+    "run_neo_radial_state_forward_fp32",
+]
+
+
+@cute.jit
+def _radial_state_forward_jit(
+    radial,
+    combined_weight,
+    hidden_weight,
+    compact_out,
+    attention_out,
+    stream: CUstream,
+):
+    edge_count, _ = radial.shape
+    _radial_state_forward_kernel(
+        radial,
+        combined_weight,
+        hidden_weight,
+        compact_out,
+        attention_out,
+    ).launch(
+        grid=[edge_count, 1, 1],
+        block=[THREADS, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.kernel
+def _radial_state_forward_kernel(
+    radial,
+    combined_weight,
+    hidden_weight,
+    compact_out,
+    attention_out,
+):
+    channel, _, _ = cute.arch.thread_idx()
+    edge, _, _ = cute.arch.block_idx()
+
+    if channel < COMPACT_WIDTH:
+        compact = cutlass.Float32(0.0)
+        for radial_channel in cutlass.range_constexpr(RADIAL_WIDTH):
+            compact += radial[edge, radial_channel].to(
+                cutlass.Float32
+            ) * combined_weight[radial_channel, channel].to(cutlass.Float32)
+        compact_out[edge, channel] = compact.to(compact_out.element_type)
+
+    attention = cutlass.Float32(0.0)
+    for radial_channel in cutlass.range_constexpr(32):
+        attention += radial[edge, radial_channel].to(cutlass.Float32) * hidden_weight[
+            radial_channel, channel
+        ].to(cutlass.Float32)
+    attention_out[edge, channel] = attention.to(attention_out.element_type)
+
+
+@device_aware_lru_cache(maxsize=4)
+def _compiled_radial_state_forward() -> Callable:
+    edge_count = cute.sym_int64()
+    fake_radial = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, RADIAL_WIDTH),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_combined = make_fake_compact_tensor(
+        cutlass.Float32,
+        (RADIAL_WIDTH, COMPACT_WIDTH),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_hidden = make_fake_compact_tensor(
+        cutlass.Float32,
+        (32, ATTENTION_WIDTH),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_compact = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, COMPACT_WIDTH),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    fake_attention = make_fake_compact_tensor(
+        cutlass.Float32,
+        (edge_count, ATTENTION_WIDTH),
+        stride_order=(1, 0),
+        **FAKE_TENSOR_KW,
+    )
+    return cute.compile(
+        _radial_state_forward_jit,
+        fake_radial,
+        fake_combined,
+        fake_hidden,
+        fake_compact,
+        fake_attention,
+        make_fake_stream(use_tvm_ffi_env_stream=True),
+        options="--enable-tvm-ffi",
+    )
+
+
+def _expect_fp32_cuda(
+    name: str,
+    tensor: torch.Tensor,
+    shape: tuple[int, ...],
+    device: torch.device,
+) -> None:
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+    if (
+        tensor.dtype != torch.float32
+        or tensor.device != device
+        or not tensor.is_cuda
+        or not tensor.is_contiguous()
+    ):
+        raise ValueError(f"{name} must be contiguous CUDA float32 on {device}")
+
+
+def run_neo_radial_state_forward_fp32(
+    *,
+    radial_feat: torch.Tensor,
+    combined_weight: torch.Tensor,
+    hidden_weight: torch.Tensor,
+    compact_out: torch.Tensor | None = None,
+    attention_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Produce the 25-value Phase-A map and 64-value attention feature."""
+    if radial_feat.ndim != 3 or tuple(radial_feat.shape[1:]) != (4, 32):
+        raise ValueError(
+            f"radial_feat must have shape (E,4,32), got {tuple(radial_feat.shape)}"
+        )
+    if not radial_feat.is_cuda:
+        raise ValueError("radial_feat must be a CUDA tensor")
+    device = radial_feat.device
+    edge_count = radial_feat.shape[0]
+    _expect_fp32_cuda(
+        "radial_feat",
+        radial_feat,
+        (edge_count, 4, 32),
+        device,
+    )
+    _expect_fp32_cuda(
+        "combined_weight",
+        combined_weight,
+        (RADIAL_WIDTH, COMPACT_WIDTH),
+        device,
+    )
+    _expect_fp32_cuda(
+        "hidden_weight",
+        hidden_weight,
+        (32, ATTENTION_WIDTH),
+        device,
+    )
+    if edge_count <= 0:
+        raise ValueError("persistent-complex SO2 requires E > 0")
+    if tuple(torch.cuda.get_device_capability(device)) != (9, 0):
+        raise RuntimeError("persistent-complex SO2 requires SM90")
+    if torch.backends.cuda.matmul.allow_tf32:
+        raise RuntimeError("strict FP32 requires allow_tf32=False")
+    if torch.get_float32_matmul_precision() != "highest":
+        raise RuntimeError("strict FP32 requires float32 matmul precision 'highest'")
+
+    if compact_out is None:
+        compact_out = torch.empty(
+            (edge_count, COMPACT_WIDTH),
+            dtype=torch.float32,
+            device=device,
+        )
+    if attention_out is None:
+        attention_out = torch.empty(
+            (edge_count, ATTENTION_WIDTH),
+            dtype=torch.float32,
+            device=device,
+        )
+    _expect_fp32_cuda(
+        "compact_out",
+        compact_out,
+        (edge_count, COMPACT_WIDTH),
+        device,
+    )
+    _expect_fp32_cuda(
+        "attention_out",
+        attention_out,
+        (edge_count, ATTENTION_WIDTH),
+        device,
+    )
+
+    with torch.cuda.device(device):
+        _compiled_radial_state_forward()(
+            radial_feat.view(edge_count, RADIAL_WIDTH),
+            combined_weight,
+            hidden_weight,
+            compact_out,
+            attention_out,
+        )
+    return compact_out, attention_out
+
+
+def project_neo_radial_input_adjoint_fp32(
+    *,
+    grad_compact: torch.Tensor,
+    grad_logits: torch.Tensor,
+    combined_weight: torch.Tensor,
+    combined_attention_weight: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Project compact Phase-A and attention adjoints to ``(E,4,32)``.
+
+    This intentionally follows the strict-FP32 cuBLAS route. No approximation
+    is introduced: the first matrix product is the compact radial adjoint and
+    the degree-zero slice receives the independent attention-logit adjoint.
+    """
+    if torch.backends.cuda.matmul.allow_tf32:
+        raise RuntimeError("strict FP32 requires allow_tf32=False")
+    if torch.get_float32_matmul_precision() != "highest":
+        raise RuntimeError("strict FP32 requires float32 matmul precision 'highest'")
+    edge_count = grad_compact.shape[0]
+    device = grad_compact.device
+    _expect_fp32_cuda(
+        "grad_compact",
+        grad_compact,
+        (edge_count, COMPACT_WIDTH),
+        device,
+    )
+    _expect_fp32_cuda(
+        "grad_logits",
+        grad_logits,
+        (edge_count, 2),
+        device,
+    )
+    _expect_fp32_cuda(
+        "combined_weight",
+        combined_weight,
+        (RADIAL_WIDTH, COMPACT_WIDTH),
+        device,
+    )
+    _expect_fp32_cuda(
+        "combined_attention_weight",
+        combined_attention_weight,
+        (32, 2),
+        device,
+    )
+    if out is None:
+        out = torch.empty(
+            (edge_count, 4, 32),
+            dtype=torch.float32,
+            device=device,
+        )
+    _expect_fp32_cuda("out", out, (edge_count, 4, 32), device)
+
+    out_flat = out.view(edge_count, RADIAL_WIDTH)
+    torch.mm(grad_compact, combined_weight.transpose(0, 1), out=out_flat)
+    out_flat[:, :32].addmm_(
+        grad_logits,
+        combined_attention_weight.transpose(0, 1),
+    )
+    return out
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py
new file mode 100644
index 0000000000..a937402c65
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py
@@ -0,0 +1,658 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Split-complex strict-FP32 SM90 implementation of the complete Neo SO2.
+
+The value path uses one split-complex edge representation throughout:
+
+* compact radial projection and direct split-complex Phase A;
+* two persistent split-complex gated residual layers;
+* the third residual SO2Linear commuted through Phase C and evaluated from
+  64-edge destination-CSR sufficient statistics;
+* the Neo output gate followed by the message-grid/readout.
+
+No ``(E,2,10,32)`` block-real SO2 slab is constructed, and no split/block-real
+pack or unpack kernel is used.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    Any,
+)
+
+import torch
+from torch import (
+    Tensor,
+)
+
+from ..linear import (
+    cached_neo_so2_linear_weights,
+)
+from ..operation import (
+    _compile_output_gate_backward,
+    _equivariant_rmsnorm_backward,
+    _runner_compile_identity,
+    _so3_linear_backward_input,
+    _x_wide_manual_backward,
+)
+from ..runner import (
+    NeoFullCuteBackward,
+)
+from .final_phase_c import (
+    ExpandedFinalWeights,
+    prepare_expanded_final_weights,
+    run_direct_statistics_forward,
+)
+from .message_grid_readout import (
+    prepare_sm90_message_grid_state,
+    run_sm90_message_grid_backward,
+)
+from .output_gate import (
+    run_chunked_final_output_gate,
+)
+from .persistent import (
+    NeoPersistentComplexSaved,
+    NeoPersistentComplexState,
+    NeoPersistentComplexWeights,
+    prepare_neo_persistent_complex_weights,
+)
+from .phase_a import (
+    run_neo_phase_a_persistent_complex_fp32,
+)
+from .phase_a_backward import (
+    run_neo_phase_a_persistent_complex_backward_fp32,
+)
+from .phase_c_attention_backward import (
+    allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs,
+    run_grouped_expanded_final_phase_c_attention_adjoint,
+)
+from .prefix import (
+    run_persistent_prefix_forward_inplace,
+    run_persistent_prefix_input_adjoint_destructive_saved,
+)
+from .radial import (
+    project_neo_radial_input_adjoint_fp32,
+    run_neo_radial_state_forward_fp32,
+)
+
+FOCUS_COUNT = 2
+CHANNELS = 32
+HIDDEN = FOCUS_COUNT * CHANNELS
+DEGREE_COUNT = 16
+M0_WIDTH = 128
+M1_WIDTH = 96
+GATED_LAYERS = 2
+_PERSISTENT_WEIGHT_CACHE = "_deepmd_cute_neo_sm90_persistent_weights"
+_FINAL_WEIGHT_CACHE = "_deepmd_cute_neo_sm90_final_weights"
+
+
+def _tensor_version_key(tensor: Tensor) -> tuple[Any, ...]:
+    return (
+        tensor.data_ptr(),
+        tensor._version,
+        tuple(tensor.shape),
+        tuple(tensor.stride()),
+        tensor.dtype,
+        tensor.device,
+    )
+
+
+def _prepare_persistent_weights(so2: Any) -> NeoPersistentComplexWeights:
+    """Pack the three SO2Linear blocks and two scalar gates once."""
+    linears = tuple(so2.so2_linears)
+    nonlinearities = tuple(so2.non_linearities)
+    if len(linears) != 3 or len(nonlinearities) != 3:
+        raise NotImplementedError("SM90 SO2 requires three SO2 layers")
+    if any(type(norm).__name__ != "Identity" for norm in so2.so2_inter_norms):
+        raise NotImplementedError("SM90 SO2 requires disabled inter-layer norms")
+    for nonlinearity in nonlinearities[:GATED_LAYERS]:
+        if getattr(nonlinearity, "layout", None) != "fndc":
+            raise NotImplementedError("SM90 SO2 requires fndc gates")
+        activation = getattr(nonlinearity, "activation_function", None)
+        if activation is None:
+            activation = getattr(
+                getattr(nonlinearity, "scalar_act", None),
+                "activation",
+                None,
+            )
+        if str(activation).lower() != "silu":
+            raise NotImplementedError("SM90 SO2 requires SiLU gates")
+        if getattr(nonlinearity.gate_linear, "bias", None) is not None:
+            raise NotImplementedError("SM90 SO2 does not support gate bias")
+
+    sources = tuple(
+        tensor
+        for linear in linears
+        for tensor in (linear.weight_m0, linear.weight_m[0])
+    ) + tuple(
+        nonlinearity.gate_linear.weight
+        for nonlinearity in nonlinearities[:GATED_LAYERS]
+    )
+    cache_key = tuple(_tensor_version_key(tensor) for tensor in sources)
+    cached = getattr(so2, _PERSISTENT_WEIGHT_CACHE, None)
+    if isinstance(cached, tuple) and len(cached) == 3 and cached[0] == cache_key:
+        return cached[1]
+
+    w0_layers: list[Tensor] = []
+    wp_layers: list[Tensor] = []
+    for linear in linears:
+        w0, wp = cached_neo_so2_linear_weights(linear)
+        w0_layers.append(w0)
+        wp_layers.append(wp)
+    gate_layers = [
+        nonlinearity.gate_linear.weight.detach()
+        .view(CHANNELS, FOCUS_COUNT, 3 * CHANNELS)
+        .permute(1, 0, 2)
+        .contiguous()
+        for nonlinearity in nonlinearities[:GATED_LAYERS]
+    ]
+    packed = prepare_neo_persistent_complex_weights(
+        torch.stack(w0_layers, dim=0).contiguous(),
+        torch.stack(wp_layers, dim=0).contiguous(),
+        torch.stack(gate_layers, dim=0).contiguous(),
+    )
+    setattr(so2, _PERSISTENT_WEIGHT_CACHE, (cache_key, packed, sources))
+    return packed
+
+
+def _prepare_final_weights(so2: Any) -> ExpandedFinalWeights:
+    """Cache the residual-folded final SO2Linear in node-GEMM form."""
+    linear = tuple(so2.so2_linears)[-1]
+    sources = (linear.weight_m0, linear.weight_m[0])
+    key = tuple(_tensor_version_key(tensor) for tensor in sources)
+    cached = getattr(so2, _FINAL_WEIGHT_CACHE, None)
+    if isinstance(cached, tuple) and len(cached) == 3 and cached[0] == key:
+        return cached[1]
+
+    w0, wp = cached_neo_so2_linear_weights(linear)
+    if tuple(w0.shape) != (FOCUS_COUNT, M0_WIDTH, M0_WIDTH):
+        raise ValueError("SM90 SO2 requires final m=0 weights (2,128,128)")
+    if tuple(wp.shape) != (FOCUS_COUNT, 2 * M1_WIDTH, 2 * M1_WIDTH):
+        raise ValueError("SM90 SO2 requires final pair weights (2,192,192)")
+    u = wp[:, :M1_WIDTH, :M1_WIDTH]
+    v = wp[:, :M1_WIDTH, M1_WIDTH:]
+    if not torch.equal(wp[:, M1_WIDTH:, :M1_WIDTH], -v):
+        raise ValueError("final pair weight lower-left block must be -V")
+    if not torch.equal(wp[:, M1_WIDTH:, M1_WIDTH:], u):
+        raise ValueError("final pair weight lower-right block must be U")
+    weights = prepare_expanded_final_weights(
+        (w0 + torch.eye(M0_WIDTH, device=w0.device, dtype=torch.float32)).contiguous(),
+        (
+            torch.complex(u, v)
+            + torch.eye(M1_WIDTH, device=w0.device, dtype=torch.complex64)
+        ).contiguous(),
+    )
+    setattr(so2, _FINAL_WEIGHT_CACHE, (key, weights, sources))
+    return weights
+
+
+__all__ = ["NeoSm90SO2Runner"]
+
+
+@dataclass(frozen=True)
+class _PersistentPrefixForward:
+    """Pre-final split state and exact gate preactivations for its adjoint."""
+
+    state: NeoPersistentComplexState
+    saved: NeoPersistentComplexSaved
+
+
+def _run_persistent_prefix_forward(
+    state: NeoPersistentComplexState,
+    weights: NeoPersistentComplexWeights,
+) -> _PersistentPrefixForward:
+    """Run two gated layers while reusing the direct Phase-A state buffer."""
+    current, saved = run_persistent_prefix_forward_inplace(state, weights)
+    return _PersistentPrefixForward(
+        state=current,
+        saved=saved,
+    )
+
+
+def _run_persistent_prefix_input_adjoint(
+    grad_out: NeoPersistentComplexState,
+    saved: NeoPersistentComplexSaved,
+    weights: NeoPersistentComplexWeights,
+) -> NeoPersistentComplexState:
+    """Overwrite dead gate checkpoints and the running residual in-place."""
+    return run_persistent_prefix_input_adjoint_destructive_saved(
+        grad_out,
+        saved,
+        weights,
+    )
+
+
+def _run_final_reverse_panels(
+    *,
+    grad_out: Tensor,
+    weights: Any,
+) -> tuple[Tensor, Tensor]:
+    """Form the two strict-FP32 reverse node panels once."""
+    node_count = int(grad_out.shape[0])
+    grad_node = grad_out.permute(1, 2, 0, 3).contiguous()
+    b0 = torch.bmm(
+        grad_node.flatten(0, 1),
+        weights.w0.flatten(0, 1).transpose(-1, -2),
+    ).view(FOCUS_COUNT, DEGREE_COUNT, node_count, M0_WIDTH)
+    grad_node1 = grad_node[:, 1:].to(torch.complex64).contiguous()
+    b1 = torch.bmm(
+        grad_node1.flatten(0, 1),
+        weights.wc.flatten(0, 1).conj().transpose(-1, -2),
+    ).view(FOCUS_COUNT, DEGREE_COUNT - 1, node_count, M1_WIDTH)
+    return b0, b1
+
+
+def _qk_node_input_adjoint(
+    runner: NeoSm90SO2Runner,
+    grad_q_node: Tensor,
+    grad_k_node: Tensor,
+) -> Tensor:
+    """Map fused Q/K node adjoints into the wide SO2 input."""
+    so2 = runner.so2
+    x_wide = runner.x_wide.detach()
+    x_l0 = x_wide[:, 0, :].reshape(runner.node_count, FOCUS_COUNT, CHANNELS)
+    grad_x_wide = torch.empty_like(x_wide, memory_format=torch.contiguous_format)
+    runner.qk_node_input_adjoint(
+        x_l0.contiguous(),
+        grad_q_node,
+        grad_k_node,
+        so2.attn_q_proj.weight.detach()
+        .float()
+        .view(CHANNELS, FOCUS_COUNT, CHANNELS)
+        .contiguous(),
+        so2.attn_k_proj.weight.detach()
+        .float()
+        .view(CHANNELS, FOCUS_COUNT, CHANNELS)
+        .contiguous(),
+        so2.attn_qk_norm.adam_scale.detach().float().contiguous(),
+        grad_x_wide.view(runner.node_count, DEGREE_COUNT * HIDDEN),
+    )
+    return grad_x_wide
+
+
+class NeoSm90SO2Runner(NeoFullCuteBackward):
+    """Complete Neo SO2 runner with one native split representation."""
+
+    uses_native_sm90_path = True
+
+    def _build_forward_graph(self) -> None:
+        torch_module = self.torch
+        so2 = self.so2
+        block = self.block
+        node_count = self.node_count
+        edge_count = self.edge_count
+        if self.compute_capability != (9, 0):
+            raise RuntimeError("split-complex SO2 requires SM90")
+
+        self.structural_scratch = None
+        self.use_full_node = block.node_lmax == block.lmax
+        x_so2 = self.x if self.use_full_node else self.x[:, : block.mp_ebed_dim]
+        x_pre = block.pre_so2_norm(x_so2)
+        self.x_wide = (
+            so2.pre_focus_mix(
+                x_pre.reshape(
+                    node_count,
+                    x_so2.shape[1],
+                    block.channels,
+                ).unsqueeze(2)
+            )
+            .squeeze(2)
+            .contiguous()
+        )
+
+        self.radial_compact, radial_l0 = run_neo_radial_state_forward_fp32(
+            radial_feat=self.radial.detach().contiguous(),
+            combined_weight=self.combined_radial,
+            hidden_weight=so2.radial_hidden_proj.weight.detach().contiguous(),
+        )
+        phase_a_state = run_neo_phase_a_persistent_complex_fp32(
+            x_wide=self.x_wide.detach(),
+            src=self.src_i32,
+            d_full=self.d.detach(),
+            radial_compact=self.radial_compact,
+            channel_basis=so2.radial_degree_mixer.channel_basis.detach()
+            .view(HIDDEN)
+            .contiguous(),
+        )
+        self.focus_gate_src = (
+            phase_a_state.m0[:, :, :CHANNELS].permute(1, 0, 2).contiguous()
+        )
+
+        self.persistent_weights = _prepare_persistent_weights(so2)
+        prefix = _run_persistent_prefix_forward(
+            phase_a_state,
+            self.persistent_weights,
+        )
+        self.phase_c_state = prefix.state
+        self.persistent_saved = prefix.saved
+        del phase_a_state, prefix
+
+        x_l0_node = self.x_wide[:, 0, :].reshape(
+            node_count,
+            FOCUS_COUNT,
+            CHANNELS,
+        )
+        self.focus_alpha = torch_module.empty(
+            edge_count,
+            FOCUS_COUNT,
+            device=self.x.device,
+            dtype=torch_module.float32,
+        )
+        self.q_node = torch_module.empty_like(
+            x_l0_node,
+            memory_format=torch_module.contiguous_format,
+        )
+        self.k_node = torch_module.empty_like(
+            x_l0_node,
+            memory_format=torch_module.contiguous_format,
+        )
+        self.attention_prelude_forward(
+            self.focus_gate_src.view(edge_count, HIDDEN),
+            x_l0_node.contiguous(),
+            so2.adamw_focus_compete_w.detach().float().contiguous(),
+            self.focus_norm_scale,
+            so2.attn_q_proj.weight.detach()
+            .float()
+            .view(CHANNELS, FOCUS_COUNT, CHANNELS)
+            .contiguous(),
+            so2.attn_k_proj.weight.detach()
+            .float()
+            .view(CHANNELS, FOCUS_COUNT, CHANNELS)
+            .contiguous(),
+            so2.attn_qk_norm.adam_scale.detach().float().contiguous(),
+            self.focus_alpha,
+            self.q_node,
+            self.k_node,
+        )
+
+        self.attn_logits = torch_module.empty(
+            edge_count,
+            FOCUS_COUNT,
+            device=self.x.device,
+            dtype=torch_module.float32,
+        )
+        self.qk_edge_forward(
+            self.q_node,
+            self.k_node,
+            radial_l0.view(edge_count, FOCUS_COUNT, CHANNELS),
+            so2.adamw_attn_logit_w.detach().contiguous(),
+            self.src_i32,
+            self.dst_i32,
+            self.attn_logits,
+        )
+        del radial_l0
+        self.softmax_fwd(
+            self.attn_logits,
+            self.edge_gate,
+            self.dst_ptr_i32,
+            so2.adamw_attn_z_bias_raw.detach()
+            .reshape(FOCUS_COUNT)
+            .float()
+            .contiguous(),
+            self.alpha,
+            self.group_max,
+            self.denom,
+        )
+        # First input adjoints need alpha, Q, K, and edge metadata, but not the
+        # materialized logits or the null-mass parameter-gradient statistics.
+        self.attn_logits = None
+        self.group_max = None
+        self.denom = None
+
+        self.beta = (self.alpha * self.focus_alpha).contiguous()
+        self.final_weights = _prepare_final_weights(so2)
+        raw_phase_c = run_direct_statistics_forward(
+            m0=self.phase_c_state.m0,
+            m1=self.phase_c_state.m1,
+            dt_packed=self.dt.detach(),
+            beta=self.beta,
+            dst_ptr=self.dst_ptr_i32,
+            weights=self.final_weights,
+        ).output
+        self.phase_c_out = run_chunked_final_output_gate(
+            raw=raw_phase_c,
+            x_wide=self.x_wide.detach(),
+            norm_scale=so2.attn_output_gate_norm.adam_scale.detach()
+            .float()
+            .reshape(FOCUS_COUNT, CHANNELS)
+            .contiguous(),
+            gate_weight=so2.adamw_attn_gate_w.detach()
+            .float()
+            .reshape(CHANNELS, FOCUS_COUNT, 1)
+            .contiguous(),
+            rotate_inv_rescale=so2.rotate_inv_rescale_full.detach().contiguous(),
+            eps=float(so2.attn_output_gate_norm.eps),
+        ).to(dtype=so2.compute_dtype)
+        # The partial/node statistics and ungated output have no backward role:
+        # the exact adjoint recomputes from grad_out and the pre-final state.
+        del raw_phase_c
+
+        out = self.phase_c_out.detach().to(dtype=so2.dtype)
+        self.out_gate_flat = out.detach()
+        self.message_grid_product = None
+        self.message_grid_sm90_state = None
+        if so2.message_node_grid_product is not None:
+            if self.packed_message_grid:
+                from ..message_grid import (
+                    run_packed_message_grid_forward,
+                )
+
+                self.message_grid_sm90_state = prepare_sm90_message_grid_state(
+                    so2.message_node_grid_product
+                )
+                grid_out, product = run_packed_message_grid_forward(
+                    so2.message_node_grid_product,
+                    out,
+                    self.x_wide,
+                    return_product=True,
+                    sm90_state=self.message_grid_sm90_state,
+                )
+                self.message_grid_product = product.detach()
+            else:
+                grid_out = so2.message_node_grid_product(out, self.x_wide)
+            out = out + grid_out
+
+        self.post_mix_input = out.detach()
+        out = so2.post_focus_mix(out.unsqueeze(2)).squeeze(2)
+        self.post_norm_input = out.unsqueeze(2).detach()
+        so2_out = block.post_so2_norm(self.post_norm_input)
+        if self.use_full_node:
+            self.final = so2_out
+        else:
+            final = self.x.new_zeros(self.x.shape)
+            final[:, : block.mp_ebed_dim] = so2_out
+            self.final = final
+
+    def input_adjoint(self, grad_out: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        """Return SO2 input adjoints through the split-complex reverse path."""
+        return _runner_backward(self, grad_out)
+
+
+def _final_manual_backward_sm90(
+    runner: NeoSm90SO2Runner,
+    grad_out: Tensor,
+) -> tuple[Tensor, Tensor]:
+    """Use the one-slab tiled message-grid adjoint in the final SO2 boundary."""
+    if not (
+        runner.packed_message_grid
+        and runner.so2.message_node_grid_product is not None
+        and runner.message_grid_sm90_state is not None
+    ):
+        raise RuntimeError("SM90 SO2 requires the packed message-grid path")
+
+    so2 = runner.so2
+    block = runner.block
+    if runner.use_full_node:
+        grad_so2_out = grad_out
+    else:
+        grad_so2_out = grad_out[:, : block.mp_ebed_dim, :, :]
+
+    phase = runner.phase_c_out.detach()
+    x_wide = runner.x_wide.detach()
+    grad_post_norm_in = _equivariant_rmsnorm_backward(
+        block.post_so2_norm,
+        runner.post_norm_input,
+        grad_so2_out,
+    )
+    grad_post_mix = _so3_linear_backward_input(
+        so2.post_focus_mix,
+        grad_post_norm_in.squeeze(2).unsqueeze(2),
+    ).squeeze(2)
+
+    message_grid_product = runner.message_grid_product
+    runner.message_grid_product = None
+    grad_out_gate_flat, grad_grid_context = run_sm90_message_grid_backward(
+        so2.message_node_grid_product,
+        runner.out_gate_flat,
+        x_wide,
+        grad_post_mix,
+        message_grid_product,
+        runner.message_grid_sm90_state,
+    )
+    del message_grid_product
+    runner.message_grid_sm90_state = None
+
+    grad_out_gate_flat.add_(grad_post_mix)
+    # FrameExpand's input adjoint preserves its degree-major einsum stride.
+    # The output-gate kernel updates one flat node panel in place, so establish
+    # that writable contract once at this consumer boundary.
+    grad_x_wide_down = grad_grid_context.contiguous()
+    grad_phase = grad_out_gate_flat.contiguous()
+    output_gate_backward = _compile_output_gate_backward(
+        _runner_compile_identity(runner),
+        float(so2.attn_output_gate_norm.eps),
+    )
+    output_gate_backward(
+        grad_phase.view(runner.node_count, DEGREE_COUNT * HIDDEN),
+        phase.contiguous().view(runner.node_count, DEGREE_COUNT * HIDDEN),
+        x_wide.contiguous().view(runner.node_count, DEGREE_COUNT * HIDDEN),
+        so2.attn_output_gate_norm.adam_scale.detach()
+        .float()
+        .reshape(FOCUS_COUNT, CHANNELS)
+        .contiguous(),
+        so2.adamw_attn_gate_w.detach()
+        .float()
+        .reshape(CHANNELS, FOCUS_COUNT, 1)
+        .contiguous(),
+        grad_phase.view(runner.node_count, DEGREE_COUNT * HIDDEN),
+        grad_x_wide_down.view(runner.node_count, DEGREE_COUNT * HIDDEN),
+    )
+    return grad_phase.reshape_as(phase), grad_x_wide_down
+
+
+def _runner_backward(
+    runner: NeoSm90SO2Runner,
+    grad_out: Tensor,
+) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+    so2 = runner.so2
+    grad_phase, grad_x_wide_down = _final_manual_backward_sm90(runner, grad_out)
+    runner.phase_c_out = None
+    runner.out_gate_flat = None
+    runner.post_mix_input = None
+    runner.post_norm_input = None
+    grad_node = (
+        grad_phase.view(
+            runner.node_count,
+            DEGREE_COUNT,
+            FOCUS_COUNT,
+            CHANNELS,
+        )
+        .permute(0, 2, 1, 3)
+        .contiguous()
+    )
+    grad_node.mul_(runner.rotate.view(1, 1, DEGREE_COUNT, 1))
+    b0, b1 = _run_final_reverse_panels(
+        grad_out=grad_node,
+        weights=runner.final_weights,
+    )
+    fused_outputs = allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs(
+        edge_count=runner.edge_count,
+        node_count=runner.node_count,
+        device=grad_node.device,
+        grad_m0=runner.phase_c_state.m0,
+        grad_m1=runner.phase_c_state.m1,
+    )
+    fused_adjoint = run_grouped_expanded_final_phase_c_attention_adjoint(
+        b0=b0,
+        b1=b1,
+        m0=runner.phase_c_state.m0,
+        m1=runner.phase_c_state.m1,
+        dt_packed=runner.dt.detach(),
+        beta=runner.beta,
+        alpha=runner.alpha,
+        focus_alpha=runner.focus_alpha,
+        focus_src=runner.focus_gate_src,
+        focus_weight=so2.adamw_focus_compete_w.detach().float().contiguous(),
+        focus_scale=runner.focus_norm_scale,
+        q_node=runner.q_node,
+        k_node=runner.k_node,
+        edge_gate=runner.edge_gate,
+        src=runner.src_i32,
+        dst_ptr=runner.dst_ptr_i32,
+        focus_eps=runner.focus_norm_eps,
+        focus_tau=float(so2.focus_softmax_tau),
+        label_smoothing=float(so2.focus_label_smoothing),
+        qk_scale=CHANNELS**-0.5,
+        use_focus_norm=runner.focus_norm_enabled,
+        outputs=fused_outputs,
+    )
+    grad_stack = NeoPersistentComplexState(
+        fused_adjoint.grad_m0,
+        fused_adjoint.grad_m1,
+    )
+    grad_dt = fused_adjoint.grad_dt
+    grad_logits = fused_adjoint.grad_logits
+    grad_edge = fused_adjoint.grad_edge
+    grad_focus_src = fused_adjoint.grad_focus_src
+    grad_q_node = fused_adjoint.grad_q_node
+    grad_k_node = fused_adjoint.grad_k_node
+    del fused_adjoint, b0, b1, grad_node
+    runner.phase_c_state = None
+    runner.beta = None
+    runner.final_weights = None
+    runner.focus_gate_src = None
+    runner.alpha = None
+    runner.focus_alpha = None
+
+    grad_phase_a = _run_persistent_prefix_input_adjoint(
+        grad_stack,
+        runner.persistent_saved,
+        runner.persistent_weights,
+    )
+    runner.persistent_saved = None
+    # focus_gate_src aliases the first m=0 row of the direct Phase-A result.
+    grad_phase_a.m0[:, :, :CHANNELS].add_(grad_focus_src)
+    del grad_stack, grad_focus_src
+    phase_a = run_neo_phase_a_persistent_complex_backward_fp32(
+        grad_state=grad_phase_a,
+        radial_compact=runner.radial_compact,
+        channel_basis=so2.radial_degree_mixer.channel_basis.detach()
+        .view(HIDDEN)
+        .contiguous(),
+        x_wide=runner.x_wide.detach(),
+        source_order=runner.source_order_i32,
+        source_ptr=runner.source_ptr_i32,
+        d_full=runner.d.detach(),
+    )
+    runner.radial_compact = None
+    runner.persistent_weights = None
+    grad_radial = project_neo_radial_input_adjoint_fp32(
+        grad_compact=phase_a.grad_radial_compact,
+        grad_logits=grad_logits,
+        combined_weight=runner.combined_radial,
+        combined_attention_weight=runner.combined_attention_radial,
+    )
+    grad_x_wide = phase_a.grad_x_wide.view_as(runner.x_wide)
+    grad_x_wide.add_(_qk_node_input_adjoint(runner, grad_q_node, grad_k_node))
+    runner.q_node = None
+    runner.k_node = None
+    grad_x_wide.add_(grad_x_wide_down)
+    grad_x = _x_wide_manual_backward(runner, grad_x_wide)
+    runner.grad_edge = grad_edge
+
+    return grad_x, phase_a.grad_d_full, grad_dt, grad_radial
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/structural_gate.py b/deepmd/pt_expt/kernels/cute/sezm/so2/structural_gate.py
new file mode 100644
index 0000000000..4ea1c7cda9
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/structural_gate.py
@@ -0,0 +1,175 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""PyTorch/cuBLAS glue for the structural Neo SO2 split-gate path.
+
+The two focus panels stay contiguous across the gate-linear boundary. Forward
+uses direct cuBLAS matrix products, while backward accumulates into the scalar
+slice of ``grad_y`` through the GEMM beta epilogue.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from typing import (
+    Any,
+)
+
+import torch
+from torch import (
+    Tensor,
+)
+
+FOCUS_COUNT = 2
+CHANNELS = 32
+GATE_WIDTH = 3 * CHANNELS
+VEC4_ALIGNMENT_BYTES = 16
+VEC4_STORAGE_OFFSET_MULTIPLE = 4
+
+
+def _dispatch_aligned_vec4_kernel(
+    kernel: Any,
+    tensor_names: tuple[str, ...],
+    *tensors: Tensor,
+) -> Any:
+    """Validate the float4 load/store contract, then dispatch the kernel."""
+    if len(tensor_names) != len(tensors):
+        raise ValueError("vec4 dispatch tensor names and arguments must match")
+    if not tensors or tensors[0].numel() == 0:
+        return None
+
+    for name, tensor in zip(tensor_names, tensors, strict=True):
+        if tensor.dtype != torch.float32:
+            raise TypeError(
+                "SM80 vec4 structural gate requires float32 tensors; "
+                f"{name} has dtype={tensor.dtype}"
+            )
+        if not tensor.is_contiguous():
+            raise ValueError(
+                "SM80 vec4 structural gate requires compact tensors; "
+                f"{name} has shape={tuple(tensor.shape)} and stride={tensor.stride()}"
+            )
+        if tensor.storage_offset() % VEC4_STORAGE_OFFSET_MULTIPLE:
+            raise ValueError(
+                "SM80 vec4 structural gate requires storage offsets divisible "
+                f"by {VEC4_STORAGE_OFFSET_MULTIPLE} float32 elements; "
+                f"{name} has storage_offset={tensor.storage_offset()}"
+            )
+        pointer_remainder = tensor.data_ptr() % VEC4_ALIGNMENT_BYTES
+        if pointer_remainder:
+            raise ValueError(
+                "SM80 vec4 structural gate requires 16-byte-aligned tensors; "
+                f"{name} has data_ptr modulo 16={pointer_remainder}"
+            )
+
+    return kernel(*tensors)
+
+
+def focus_major_gate_linear_forward(
+    gate_src: Tensor,
+    gate_weight: Tensor,
+) -> Tensor:
+    """Project ``(E, 2, 32)`` into contiguous ``(2, E, 96)`` panels."""
+    edge_count = gate_src.shape[0]
+    weight = gate_weight.view(CHANNELS, FOCUS_COUNT, GATE_WIDTH)
+    logits = torch.empty(
+        FOCUS_COUNT,
+        edge_count,
+        GATE_WIDTH,
+        dtype=gate_src.dtype,
+        device=gate_src.device,
+    )
+    for focus in range(FOCUS_COUNT):
+        torch.mm(gate_src[:, focus, :], weight[:, focus, :], out=logits[focus])
+    return logits
+
+
+def focus_major_gate_linear_backward_add_(
+    grad_y: Tensor,
+    grad_logits: Tensor,
+    gate_weight: Tensor,
+) -> Tensor:
+    """Accumulate the gate-linear adjoint into ``grad_y`` in place."""
+    weight = gate_weight.view(CHANNELS, FOCUS_COUNT, GATE_WIDTH)
+    for focus in range(FOCUS_COUNT):
+        grad_y[:, focus, 0, :].addmm_(
+            grad_logits[focus],
+            weight[:, focus, :].T,
+        )
+    return grad_y
+
+
+def focus_major_so2_backward_with_folded_residual_out(
+    grad_out: Tensor,
+    w0_folded_t: Tensor,
+    wpair_folded_t: Tensor,
+    *,
+    out: Tensor,
+) -> Tensor:
+    """Write ``grad_out @ (W.T + I)`` without seeding output copies."""
+    if out.shape != grad_out.shape:
+        raise ValueError("SO2 backward tensors must have identical shapes")
+    if not out.is_contiguous():
+        raise ValueError("SO2 backward output storage must be contiguous")
+    if out.untyped_storage()._cdata == grad_out.untyped_storage()._cdata:
+        raise ValueError("SO2 backward output must not alias its input")
+    edge_count = grad_out.shape[0]
+    grad_flat = grad_out.view(edge_count, FOCUS_COUNT, 10 * CHANNELS)
+    out_flat = out.view(edge_count, FOCUS_COUNT, 10 * CHANNELS)
+    split = 4 * CHANNELS
+    for focus in range(FOCUS_COUNT):
+        torch.mm(
+            grad_flat[:, focus, :split],
+            w0_folded_t[focus],
+            out=out_flat[:, focus, :split],
+        )
+        torch.mm(
+            grad_flat[:, focus, split:],
+            wpair_folded_t[focus],
+            out=out_flat[:, focus, split:],
+        )
+    return out
+
+
+def run_structural_gate_forward(
+    kernel: Any,
+    residual: Tensor,
+    y: Tensor,
+    logits: Tensor,
+    *,
+    out: Tensor,
+) -> Tensor:
+    """Run the alias-safe CuTe forward into caller-owned storage."""
+    rows = residual.shape[0] * residual.shape[1]
+    kernel(
+        residual.view(rows, 10 * CHANNELS),
+        y.view(rows, 10 * CHANNELS),
+        logits,
+        out.view(rows, 10 * CHANNELS),
+    )
+    return out
+
+
+def run_structural_gate_backward(
+    kernel: Any,
+    grad_out: Tensor,
+    y: Tensor,
+    logits: Tensor,
+    grad_y: Tensor,
+    *,
+    grad_logits: Tensor | None,
+    overwrite_logits: bool,
+) -> Tensor:
+    """Run gate backward, optionally replacing consumed logits with adjoints."""
+    grad_logits_out = logits if overwrite_logits else grad_logits
+    if grad_logits_out is None:
+        raise ValueError("grad_logits storage is required when logits are preserved")
+    rows = y.shape[0] * y.shape[1]
+    kernel(
+        grad_out.view(rows, 10 * CHANNELS),
+        y.view(rows, 10 * CHANNELS),
+        logits,
+        grad_y.view(rows, 10 * CHANNELS),
+        grad_logits_out,
+    )
+    return grad_logits_out
diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py b/deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py
new file mode 100644
index 0000000000..f6d977ae18
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py
@@ -0,0 +1,137 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Fixed packed Wigner layout for the Neo ``lmax=3, mmax=1`` SO2 path.
+
+The panel stores only rows selected by ``coeff_index_m`` and only columns from
+the matching Wigner block. Phase A reads ``D[coeff, degree]`` while Phase C
+reads ``Dt[degree, coeff]``; both expressions therefore address the same slot.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from dataclasses import (
+    dataclass,
+)
+from functools import (
+    lru_cache,
+)
+from typing import (
+    TYPE_CHECKING,
+)
+
+import torch
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Iterator,
+    )
+
+    from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+        WignerDCalculator,
+    )
+
+
+BLOCK_WIDTHS = (1, 3, 5, 7)
+FULL_BLOCK_OFFSETS = (0, 1, 4, 9, 16)
+
+# Rows are ordered m=0, m=-1, m=+1 within each non-scalar block.
+SELECTED_LOCAL_ROWS = ((0,), (1, 0, 2), (2, 1, 3), (3, 2, 4))
+PANEL_BLOCK_OFFSETS = (0, 1, 10, 25, 46)
+PACKED_VALUE_COUNT = PANEL_BLOCK_OFFSETS[-1]
+
+
+@lru_cache(maxsize=8)
+def _reference_wigner_calculator(
+    lmax: int, eps: float, dtype: torch.dtype, device: torch.device
+) -> WignerDCalculator:
+    from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+        WignerDCalculator,
+    )
+
+    calculator = WignerDCalculator(lmax, eps=eps, dtype=dtype).to(device)
+    # The recovery path must not depend on an optional accelerated backend.
+    calculator.cutile_infer_monomials = False
+    calculator.triton_infer_l_1_monomials = False
+    calculator.triton_train_l_1_monomials = False
+    return calculator.eval()
+
+
+@torch.compiler.disable
+def dense_wigner_for_fallback(
+    quaternion: torch.Tensor,
+    *,
+    lmax: int,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Rebuild dense rotations only when a packed SO2 invocation declines.
+
+    The quaternion retains its geometry autograd history. Only constant
+    coefficient tables are cached, never edge tensors or their gradients.
+    """
+    calculator = _reference_wigner_calculator(
+        lmax, eps, quaternion.dtype, quaternion.device
+    )
+    return calculator(quaternion)
+
+
+# DeePMD's m-major reduced ordering for lmax=3, mmax=1.
+COEFF_INDEX_M = (0, 2, 6, 12, 1, 5, 11, 3, 7, 13)
+REDUCED_DEGREES = (0, 1, 2, 3, 1, 2, 3, 1, 2, 3)
+REDUCED_PANEL_ROW_OFFSETS = (0, 1, 10, 25, 4, 15, 32, 7, 20, 39)
+ZONAL_PANEL_OFFSETS = tuple(range(1, 4)) + tuple(range(10, 15)) + tuple(range(25, 32))
+
+
+@dataclass(frozen=True)
+class PackedWignerEntry:
+    offset: int
+    degree: int
+    reduced: int
+    full_row: int
+    full_col: int
+
+
+_REDUCED_BY_FULL_ROW = {
+    full_row: reduced for reduced, full_row in enumerate(COEFF_INDEX_M)
+}
+
+
+def d_offset(reduced: int, full_col: int) -> int | None:
+    """Map ``D[coeff_index_m[reduced], full_col]`` into the packed panel."""
+    degree = REDUCED_DEGREES[reduced]
+    block_start = FULL_BLOCK_OFFSETS[degree]
+    block_stop = FULL_BLOCK_OFFSETS[degree + 1]
+    if full_col < block_start or full_col >= block_stop:
+        return None
+    return REDUCED_PANEL_ROW_OFFSETS[reduced] + full_col - block_start
+
+
+def dt_offset(full_row: int, reduced: int) -> int | None:
+    """Map ``Dt[full_row, coeff_index_m[reduced]]`` to its shared D slot."""
+    return d_offset(reduced, full_row)
+
+
+def iter_packed_entries() -> Iterator[PackedWignerEntry]:
+    """Yield the 46 stored entries in contiguous panel order."""
+    for degree, (width, block_start, panel_start, local_rows) in enumerate(
+        zip(
+            BLOCK_WIDTHS,
+            FULL_BLOCK_OFFSETS[:-1],
+            PANEL_BLOCK_OFFSETS[:-1],
+            SELECTED_LOCAL_ROWS,
+            strict=True,
+        )
+    ):
+        for row_slot, local_row in enumerate(local_rows):
+            full_row = block_start + local_row
+            reduced = _REDUCED_BY_FULL_ROW[full_row]
+            row_start = panel_start + row_slot * width
+            for local_col in range(width):
+                yield PackedWignerEntry(
+                    offset=row_start + local_col,
+                    degree=degree,
+                    reduced=reduced,
+                    full_row=full_row,
+                    full_col=block_start + local_col,
+                )
diff --git a/deepmd/pt_expt/kernels/cute/sezm/wignerd.py b/deepmd/pt_expt/kernels/cute/sezm/wignerd.py
new file mode 100644
index 0000000000..cb70a10573
--- /dev/null
+++ b/deepmd/pt_expt/kernels/cute/sezm/wignerd.py
@@ -0,0 +1,1009 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+# pyright: reportMissingImports=false
+# ruff: noqa: ANN001, ANN201, ANN202, ANN204, TC002, UP035
+"""Packed CuTe Wigner-D panel for the Neo SO2 inference path."""
+
+from __future__ import (
+    annotations,
+)
+
+import threading
+from dataclasses import (
+    dataclass,
+)
+from typing import (
+    Any,
+    Callable,
+)
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings.driver import (
+    CUstream,
+)
+from cutlass.cute.runtime import (
+    make_fake_compact_tensor,
+    make_fake_stream,
+)
+
+from . import (
+    runtime_policy,
+)
+from .compile_cache import (
+    device_aware_lru_cache,
+)
+from .so2.wigner_layout import PACKED_VALUE_COUNT as SO2_PANEL_VALUES
+
+L2_SPARSE_TERMS = 10
+L3_SPARSE_TERMS = 20
+
+
+def _load_dpa4_wignerd_calculator():
+    """Load DeePMD's WignerDCalculator to reuse its coefficient tables."""
+    from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+        WignerDCalculator,
+    )
+
+    return WignerDCalculator
+
+
+def _sparsify_rows(
+    coeffs: torch.Tensor, max_terms: int, threshold: float = 1.0e-12
+) -> tuple[torch.Tensor, torch.Tensor]:
+    values = torch.zeros(
+        coeffs.shape[0], max_terms, device=coeffs.device, dtype=coeffs.dtype
+    )
+    indices = torch.zeros(
+        coeffs.shape[0], max_terms, device=coeffs.device, dtype=torch.int32
+    )
+    for row in range(coeffs.shape[0]):
+        nz = torch.nonzero(coeffs[row].abs() > threshold, as_tuple=False).flatten()
+        if nz.numel() > max_terms:
+            raise RuntimeError(
+                f"sparse row {row} has {nz.numel()} terms, max_terms={max_terms}"
+            )
+        values[row, : nz.numel()] = coeffs[row, nz]
+        indices[row, : nz.numel()] = nz.to(torch.int32)
+    return values.contiguous(), indices.contiguous()
+
+
+def _build_l2_l3_tables(
+    dtype: torch.dtype, device: torch.device
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+]:
+    calc_cls = _load_dpa4_wignerd_calculator()
+    try:
+        cache = calc_cls._get_small_order_cache_cpu_fp64(3)
+    except TypeError:
+        cache = calc_cls._get_small_order_cache_cpu_fp64()
+    c_l2 = cache["C_l2"]
+    monomials_l2 = calc_cls._generate_monomials(4, 4)
+    c_l2_flat = torch.zeros(
+        25,
+        len(monomials_l2),
+        device=c_l2.device,
+        dtype=c_l2.dtype,
+    )
+    for mono_idx, exponents in enumerate(monomials_l2):
+        for a in range(4):
+            for b in range(4):
+                for c in range(4):
+                    for d in range(4):
+                        counts = (
+                            int(a == 0) + int(b == 0) + int(c == 0) + int(d == 0),
+                            int(a == 1) + int(b == 1) + int(c == 1) + int(d == 1),
+                            int(a == 2) + int(b == 2) + int(c == 2) + int(d == 2),
+                            int(a == 3) + int(b == 3) + int(c == 3) + int(d == 3),
+                        )
+                        if counts == exponents:
+                            c_l2_flat[:, mono_idx] += c_l2[:, :, a, b, c, d].reshape(25)
+    exp_l2 = torch.tensor(
+        monomials_l2,
+        device=c_l2.device,
+        dtype=torch.int32,
+    )
+    c_l2_sparse, c_l2_sparse_idx = _sparsify_rows(c_l2_flat, L2_SPARSE_TERMS)
+    c_l3_sparse, c_l3_sparse_idx = _sparsify_rows(cache["C_l3"], L3_SPARSE_TERMS)
+    return (
+        exp_l2.to(device=device, dtype=torch.int32).contiguous(),
+        c_l2_sparse.to(device=device, dtype=dtype).contiguous(),
+        c_l2_sparse_idx.to(device=device, dtype=torch.int32).contiguous(),
+        cache["exp_l3"].to(device=device, dtype=torch.int32).contiguous(),
+        c_l3_sparse.to(device=device, dtype=dtype).contiguous(),
+        c_l3_sparse_idx.to(device=device, dtype=torch.int32).contiguous(),
+    )
+
+
+@dataclass(frozen=True)
+class WignerDParams:
+    q: cute.Tensor
+    panel: cute.Tensor
+    exp_l2: cute.Tensor
+    c_l2_sparse: cute.Tensor
+    c_l2_sparse_idx: cute.Tensor
+    exp_l3: cute.Tensor
+    c_l3_sparse: cute.Tensor
+    c_l3_sparse_idx: cute.Tensor
+
+
+class WignerDForward:
+    def __init__(self, threads: int, dtype):
+        if threads % 32 != 0:
+            raise ValueError("threads must be a multiple of 32")
+        self.threads = int(threads)
+        self.warps = threads // 32
+        self.dtype = dtype
+
+    @cute.jit
+    def rotmat(self, w, x, y, z, row, col):
+        two = self.dtype(2.0)
+        one = self.dtype(1.0)
+        value = self.dtype(0.0)
+        if row == 0 and col == 0:
+            value = one - two * (y * y + z * z)
+        elif row == 0 and col == 1:
+            value = two * (x * y - w * z)
+        elif row == 0 and col == 2:
+            value = two * (x * z + w * y)
+        elif row == 1 and col == 0:
+            value = two * (x * y + w * z)
+        elif row == 1 and col == 1:
+            value = one - two * (x * x + z * z)
+        elif row == 1 and col == 2:
+            value = two * (y * z - w * x)
+        elif row == 2 and col == 0:
+            value = two * (x * z - w * y)
+        elif row == 2 and col == 1:
+            value = two * (y * z + w * x)
+        elif row == 2 and col == 2:
+            value = one - two * (x * x + y * y)
+        return value
+
+    @cute.jit
+    def perm(self, idx):
+        out = idx
+        if idx == 0:
+            out = 1
+        elif idx == 1:
+            out = 2
+        elif idx == 2:
+            out = 0
+        return out
+
+    @cute.jit
+    def sign(self, idx):
+        value = self.dtype(-1.0)
+        if idx == 2:
+            value = self.dtype(1.0)
+        return value
+
+    @cute.jit
+    def d1(self, w, x, y, z, row, col):
+        r = self.rotmat(w, x, y, z, self.perm(row), self.perm(col))
+        return r * self.sign(row) * self.sign(col)
+
+    @cute.jit
+    def component(self, w, x, y, z, comp):
+        value = w
+        if comp == 1:
+            value = x
+        elif comp == 2:
+            value = y
+        elif comp == 3:
+            value = z
+        return value
+
+    @cute.jit
+    def pow_small(self, base, exponent):
+        value = self.dtype(1.0)
+        if exponent >= 1:
+            value *= base
+        if exponent >= 2:
+            value *= base
+        if exponent >= 3:
+            value *= base
+        if exponent >= 4:
+            value *= base
+        if exponent >= 5:
+            value *= base
+        if exponent >= 6:
+            value *= base
+        return value
+
+    @cute.jit
+    def monomial_l3(self, exp_l3, mono, w, x, y, z):
+        value = self.dtype(1.0)
+        for comp in cutlass.range_constexpr(4):
+            value *= self.pow_small(
+                self.component(w, x, y, z, comp),
+                exp_l3[mono, comp],
+            )
+        return value
+
+    @cute.kernel
+    def kernel_warp_edges_panel(self, params: WignerDParams):
+        tidx, _, _ = cute.arch.thread_idx()
+        edge_block, _, _ = cute.arch.block_idx()
+        lane = tidx % 32
+        warp = tidx // 32
+        edge_count, _ = params.q.shape
+        edge = edge_block * self.warps + warp
+
+        smem = cutlass.utils.SmemAllocator()
+        l2_monomials = smem.allocate_tensor(self.dtype, self.warps * 35)
+        l3_monomials = smem.allocate_tensor(self.dtype, self.warps * 84)
+
+        if edge < edge_count:
+            qw = params.q[edge, 0].to(self.dtype)
+            qx = params.q[edge, 1].to(self.dtype)
+            qy = params.q[edge, 2].to(self.dtype)
+            qz = params.q[edge, 3].to(self.dtype)
+            inv_norm = cute.rsqrt(
+                qw * qw + qx * qx + qy * qy + qz * qz + self.dtype(1.0e-14)
+            )
+            qw = qw * inv_norm
+            qx = qx * inv_norm
+            qy = qy * inv_norm
+            qz = qz * inv_norm
+
+            if lane == 0:
+                params.panel[edge, 0] = self.dtype(1.0).to(params.panel.element_type)
+
+            for flat in cutlass.range(lane, 9, 32, unroll=1):
+                row_slot = flat // 3
+                col = flat - row_slot * 3
+                row = cutlass.Int32(1)
+                if row_slot == 1:
+                    row = cutlass.Int32(0)
+                elif row_slot == 2:
+                    row = cutlass.Int32(2)
+                value = self.d1(qw, qx, qy, qz, row, col)
+                params.panel[edge, 1 + flat] = value.to(params.panel.element_type)
+
+            l2_base = warp * 35
+            for mono in cutlass.range(lane, 35, 32, unroll=1):
+                l2_monomials[l2_base + mono] = self.monomial_l3(
+                    params.exp_l2, mono, qw, qx, qy, qz
+                )
+            cute.arch.sync_warp()
+            for flat in cutlass.range(lane, 15, 32, unroll=1):
+                row_slot = flat // 5
+                col = flat - row_slot * 5
+                row = cutlass.Int32(2)
+                if row_slot == 1:
+                    row = cutlass.Int32(1)
+                elif row_slot == 2:
+                    row = cutlass.Int32(3)
+                block_flat = row * 5 + col
+                value = self.dtype(0.0)
+                for term in cutlass.range_constexpr(L2_SPARSE_TERMS):
+                    mono = params.c_l2_sparse_idx[block_flat, term]
+                    value += (
+                        params.c_l2_sparse[block_flat, term].to(self.dtype)
+                        * l2_monomials[l2_base + mono]
+                    )
+                params.panel[edge, 10 + flat] = value.to(params.panel.element_type)
+
+            l3_base = warp * 84
+            for mono in cutlass.range(lane, 84, 32, unroll=1):
+                l3_monomials[l3_base + mono] = self.monomial_l3(
+                    params.exp_l3, mono, qw, qx, qy, qz
+                )
+            cute.arch.sync_warp()
+            for flat in cutlass.range(lane, 21, 32, unroll=1):
+                row_slot = flat // 7
+                col = flat - row_slot * 7
+                row = cutlass.Int32(3)
+                if row_slot == 1:
+                    row = cutlass.Int32(2)
+                elif row_slot == 2:
+                    row = cutlass.Int32(4)
+                block_flat = row * 7 + col
+                value = self.dtype(0.0)
+                for term in cutlass.range_constexpr(L3_SPARSE_TERMS):
+                    mono = params.c_l3_sparse_idx[block_flat, term]
+                    value += (
+                        params.c_l3_sparse[block_flat, term].to(self.dtype)
+                        * l3_monomials[l3_base + mono]
+                    )
+                params.panel[edge, 25 + flat] = value.to(params.panel.element_type)
+
+
+@dataclass(frozen=True)
+class WignerDBwdParams:
+    q: cute.Tensor
+    grad_panel: cute.Tensor
+    grad_q: cute.Tensor
+    exp_l2: cute.Tensor
+    c_l2_sparse: cute.Tensor
+    c_l2_sparse_idx: cute.Tensor
+    exp_l3: cute.Tensor
+    c_l3_sparse: cute.Tensor
+    c_l3_sparse_idx: cute.Tensor
+
+
+class WignerDBackward:
+    def __init__(self, threads: int, dtype):
+        if threads % 32 != 0:
+            raise ValueError("threads must be a multiple of 32")
+        self.threads = int(threads)
+        self.warps = threads // 32
+        self.dtype = dtype
+
+    @cute.jit
+    def warp_sum(self, value):
+        return cute.arch.warp_reduction_sum(value)
+
+    @cute.jit
+    def cta_sum(self, value, scratch, tidx):
+        lane = tidx % 32
+        warp = tidx // 32
+        value = self.warp_sum(value)
+        if lane == 0:
+            scratch[warp] = value
+        cute.arch.sync_threads()
+
+        total = self.dtype(0.0)
+        if tidx < self.warps:
+            total = scratch[tidx]
+        total = self.warp_sum(total)
+        if tidx == 0:
+            scratch[0] = total
+        cute.arch.sync_threads()
+        return scratch[0]
+
+    @cute.jit
+    def rotmat_grad(self, w, x, y, z, row, col, comp):
+        two = self.dtype(2.0)
+        four = self.dtype(4.0)
+        value = self.dtype(0.0)
+        if row == 0 and col == 0:
+            if comp == 2:
+                value = -four * y
+            elif comp == 3:
+                value = -four * z
+        elif row == 0 and col == 1:
+            if comp == 0:
+                value = -two * z
+            elif comp == 1:
+                value = two * y
+            elif comp == 2:
+                value = two * x
+            elif comp == 3:
+                value = -two * w
+        elif row == 0 and col == 2:
+            if comp == 0:
+                value = two * y
+            elif comp == 1:
+                value = two * z
+            elif comp == 2:
+                value = two * w
+            elif comp == 3:
+                value = two * x
+        elif row == 1 and col == 0:
+            if comp == 0:
+                value = two * z
+            elif comp == 1:
+                value = two * y
+            elif comp == 2:
+                value = two * x
+            elif comp == 3:
+                value = two * w
+        elif row == 1 and col == 1:
+            if comp == 1:
+                value = -four * x
+            elif comp == 3:
+                value = -four * z
+        elif row == 1 and col == 2:
+            if comp == 0:
+                value = -two * x
+            elif comp == 1:
+                value = -two * w
+            elif comp == 2:
+                value = two * z
+            elif comp == 3:
+                value = two * y
+        elif row == 2 and col == 0:
+            if comp == 0:
+                value = -two * y
+            elif comp == 1:
+                value = two * z
+            elif comp == 2:
+                value = -two * w
+            elif comp == 3:
+                value = two * x
+        elif row == 2 and col == 1:
+            if comp == 0:
+                value = two * x
+            elif comp == 1:
+                value = two * w
+            elif comp == 2:
+                value = two * z
+            elif comp == 3:
+                value = two * y
+        elif row == 2 and col == 2:
+            if comp == 1:
+                value = -four * x
+            elif comp == 2:
+                value = -four * y
+        return value
+
+    @cute.jit
+    def perm(self, idx):
+        out = idx
+        if idx == 0:
+            out = 1
+        elif idx == 1:
+            out = 2
+        elif idx == 2:
+            out = 0
+        return out
+
+    @cute.jit
+    def sign(self, idx):
+        value = self.dtype(-1.0)
+        if idx == 2:
+            value = self.dtype(1.0)
+        return value
+
+    @cute.jit
+    def d1_grad(self, w, x, y, z, row, col, comp):
+        r = self.rotmat_grad(w, x, y, z, self.perm(row), self.perm(col), comp)
+        return r * self.sign(row) * self.sign(col)
+
+    @cute.jit
+    def component(self, w, x, y, z, comp):
+        value = w
+        if comp == 1:
+            value = x
+        elif comp == 2:
+            value = y
+        elif comp == 3:
+            value = z
+        return value
+
+    @cute.jit
+    def pow_small(self, base, exponent):
+        value = self.dtype(1.0)
+        if exponent >= 1:
+            value *= base
+        if exponent >= 2:
+            value *= base
+        if exponent >= 3:
+            value *= base
+        if exponent >= 4:
+            value *= base
+        if exponent >= 5:
+            value *= base
+        if exponent >= 6:
+            value *= base
+        return value
+
+    @cute.jit
+    def monomial_l3_grad(self, exp_l3, mono, w, x, y, z, comp):
+        exp_comp = exp_l3[mono, comp]
+        value = self.dtype(0.0)
+        if exp_comp > 0:
+            value = exp_comp.to(self.dtype)
+            for c in cutlass.range_constexpr(4):
+                exp_c = exp_l3[mono, c]
+                if c == comp:
+                    exp_c = exp_c - 1
+                value *= self.pow_small(self.component(w, x, y, z, c), exp_c)
+        return value
+
+    @cute.kernel
+    def kernel_sparse_panel(self, params: WignerDBwdParams):
+        tidx, _, _ = cute.arch.thread_idx()
+        edge, _, _ = cute.arch.block_idx()
+
+        smem = cutlass.utils.SmemAllocator()
+        scratch = smem.allocate_tensor(self.dtype, self.warps)
+        dmono_l2 = smem.allocate_tensor(self.dtype, 4 * 35)
+        dmono_l3 = smem.allocate_tensor(self.dtype, 4 * 84)
+
+        raw_w = params.q[edge, 0].to(self.dtype)
+        raw_x = params.q[edge, 1].to(self.dtype)
+        raw_y = params.q[edge, 2].to(self.dtype)
+        raw_z = params.q[edge, 3].to(self.dtype)
+        inv_norm = cute.rsqrt(
+            raw_w * raw_w
+            + raw_x * raw_x
+            + raw_y * raw_y
+            + raw_z * raw_z
+            + self.dtype(1.0e-14)
+        )
+        w = raw_w * inv_norm
+        x = raw_x * inv_norm
+        y = raw_y * inv_norm
+        z = raw_z * inv_norm
+
+        for idx in cutlass.range(tidx, 4 * 35, self.threads, unroll=1):
+            comp = idx // 35
+            mono = idx - comp * 35
+            dmono_l2[idx] = self.monomial_l3_grad(params.exp_l2, mono, w, x, y, z, comp)
+        for idx in cutlass.range(tidx, 4 * 84, self.threads, unroll=1):
+            comp = idx // 84
+            mono = idx - comp * 84
+            dmono_l3[idx] = self.monomial_l3_grad(params.exp_l3, mono, w, x, y, z, comp)
+        cute.arch.sync_threads()
+
+        gw_local = self.dtype(0.0)
+        gx_local = self.dtype(0.0)
+        gy_local = self.dtype(0.0)
+        gz_local = self.dtype(0.0)
+
+        for flat in cutlass.range(tidx, 9, self.threads, unroll=1):
+            row_slot = flat // 3
+            col = flat - row_slot * 3
+            row = cutlass.Int32(1)
+            if row_slot == 1:
+                row = cutlass.Int32(0)
+            elif row_slot == 2:
+                row = cutlass.Int32(2)
+            grad = params.grad_panel[edge, 1 + flat].to(self.dtype)
+            gw_local += grad * self.d1_grad(w, x, y, z, row, col, 0)
+            gx_local += grad * self.d1_grad(w, x, y, z, row, col, 1)
+            gy_local += grad * self.d1_grad(w, x, y, z, row, col, 2)
+            gz_local += grad * self.d1_grad(w, x, y, z, row, col, 3)
+
+        for flat in cutlass.range(tidx, 15, self.threads, unroll=1):
+            row_slot = flat // 5
+            col = flat - row_slot * 5
+            row = cutlass.Int32(2)
+            if row_slot == 1:
+                row = cutlass.Int32(1)
+            elif row_slot == 2:
+                row = cutlass.Int32(3)
+            block_flat = row * 5 + col
+            grad = params.grad_panel[edge, 10 + flat].to(self.dtype)
+            for term in cutlass.range_constexpr(L2_SPARSE_TERMS):
+                mono = params.c_l2_sparse_idx[block_flat, term]
+                coeff = params.c_l2_sparse[block_flat, term].to(self.dtype)
+                gw_local += grad * coeff * dmono_l2[mono]
+                gx_local += grad * coeff * dmono_l2[35 + mono]
+                gy_local += grad * coeff * dmono_l2[70 + mono]
+                gz_local += grad * coeff * dmono_l2[105 + mono]
+
+        for flat in cutlass.range(tidx, 21, self.threads, unroll=1):
+            row_slot = flat // 7
+            col = flat - row_slot * 7
+            row = cutlass.Int32(3)
+            if row_slot == 1:
+                row = cutlass.Int32(2)
+            elif row_slot == 2:
+                row = cutlass.Int32(4)
+            block_flat = row * 7 + col
+            grad = params.grad_panel[edge, 25 + flat].to(self.dtype)
+            for term in cutlass.range_constexpr(L3_SPARSE_TERMS):
+                mono = params.c_l3_sparse_idx[block_flat, term]
+                coeff = params.c_l3_sparse[block_flat, term].to(self.dtype)
+                gw_local += grad * coeff * dmono_l3[mono]
+                gx_local += grad * coeff * dmono_l3[84 + mono]
+                gy_local += grad * coeff * dmono_l3[168 + mono]
+                gz_local += grad * coeff * dmono_l3[252 + mono]
+
+        gw = self.cta_sum(gw_local, scratch, tidx)
+        cute.arch.sync_threads()
+        gx = self.cta_sum(gx_local, scratch, tidx)
+        cute.arch.sync_threads()
+        gy = self.cta_sum(gy_local, scratch, tidx)
+        cute.arch.sync_threads()
+        gz = self.cta_sum(gz_local, scratch, tidx)
+
+        if tidx == 0:
+            dot = gw * raw_w + gx * raw_x + gy * raw_y + gz * raw_z
+            inv3 = inv_norm * inv_norm * inv_norm
+            params.grad_q[edge, 0] = (inv_norm * gw - raw_w * inv3 * dot).to(
+                params.grad_q.element_type
+            )
+            params.grad_q[edge, 1] = (inv_norm * gx - raw_x * inv3 * dot).to(
+                params.grad_q.element_type
+            )
+            params.grad_q[edge, 2] = (inv_norm * gy - raw_y * inv3 * dot).to(
+                params.grad_q.element_type
+            )
+            params.grad_q[edge, 3] = (inv_norm * gz - raw_z * inv3 * dot).to(
+                params.grad_q.element_type
+            )
+
+
+@cute.jit
+def wignerd_panel_forward_warp_edges_jit(
+    q: cute.Tensor,
+    panel: cute.Tensor,
+    exp_l2: cute.Tensor,
+    c_l2_sparse: cute.Tensor,
+    c_l2_sparse_idx: cute.Tensor,
+    exp_l3: cute.Tensor,
+    c_l3_sparse: cute.Tensor,
+    c_l3_sparse_idx: cute.Tensor,
+    threads: cutlass.Constexpr[int],
+    stream: CUstream,
+):
+    params = WignerDParams(
+        q=q,
+        panel=panel,
+        exp_l2=exp_l2,
+        c_l2_sparse=c_l2_sparse,
+        c_l2_sparse_idx=c_l2_sparse_idx,
+        exp_l3=exp_l3,
+        c_l3_sparse=c_l3_sparse,
+        c_l3_sparse_idx=c_l3_sparse_idx,
+    )
+    edge_count, _ = q.shape
+    warps = threads // 32
+    edge_blocks = cute.ceil_div(edge_count, warps)
+    WignerDForward(threads, cutlass.Float32).kernel_warp_edges_panel(params).launch(
+        grid=[edge_blocks, 1, 1],
+        block=[threads, 1, 1],
+        stream=stream,
+    )
+
+
+@cute.jit
+def wignerd_panel_backward_jit(
+    q: cute.Tensor,
+    grad_panel: cute.Tensor,
+    grad_q: cute.Tensor,
+    exp_l2: cute.Tensor,
+    c_l2_sparse: cute.Tensor,
+    c_l2_sparse_idx: cute.Tensor,
+    exp_l3: cute.Tensor,
+    c_l3_sparse: cute.Tensor,
+    c_l3_sparse_idx: cute.Tensor,
+    threads: cutlass.Constexpr[int],
+    stream: CUstream,
+):
+    params = WignerDBwdParams(
+        q=q,
+        grad_panel=grad_panel,
+        grad_q=grad_q,
+        exp_l2=exp_l2,
+        c_l2_sparse=c_l2_sparse,
+        c_l2_sparse_idx=c_l2_sparse_idx,
+        exp_l3=exp_l3,
+        c_l3_sparse=c_l3_sparse,
+        c_l3_sparse_idx=c_l3_sparse_idx,
+    )
+    edge_count, _ = q.shape
+    WignerDBackward(threads, cutlass.Float32).kernel_sparse_panel(params).launch(
+        grid=[edge_count, 1, 1],
+        block=[threads, 1, 1],
+        stream=stream,
+    )
+
+
+def compile_wignerd_panel_forward(
+    threads: int,
+) -> Callable:
+    if threads != 32:
+        raise ValueError("packed Wigner forward requires one warp per edge")
+    e = cute.sym_int64()
+    fake_q = make_fake_compact_tensor(cutlass.Float32, (e, 4), stride_order=(1, 0))
+    fake_panel = make_fake_compact_tensor(
+        cutlass.Float32,
+        (e, SO2_PANEL_VALUES),
+        stride_order=(1, 0),
+    )
+    fake_exp_l2 = make_fake_compact_tensor(cutlass.Int32, (35, 4), stride_order=(1, 0))
+    fake_c_l2_sparse = make_fake_compact_tensor(
+        cutlass.Float32,
+        (25, L2_SPARSE_TERMS),
+        stride_order=(1, 0),
+    )
+    fake_c_l2_sparse_idx = make_fake_compact_tensor(
+        cutlass.Int32,
+        (25, L2_SPARSE_TERMS),
+        stride_order=(1, 0),
+    )
+    fake_exp_l3 = make_fake_compact_tensor(cutlass.Int32, (84, 4), stride_order=(1, 0))
+    fake_c_l3_sparse = make_fake_compact_tensor(
+        cutlass.Float32,
+        (49, L3_SPARSE_TERMS),
+        stride_order=(1, 0),
+    )
+    fake_c_l3_sparse_idx = make_fake_compact_tensor(
+        cutlass.Int32,
+        (49, L3_SPARSE_TERMS),
+        stride_order=(1, 0),
+    )
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        wignerd_panel_forward_warp_edges_jit,
+        fake_q,
+        fake_panel,
+        fake_exp_l2,
+        fake_c_l2_sparse,
+        fake_c_l2_sparse_idx,
+        fake_exp_l3,
+        fake_c_l3_sparse,
+        fake_c_l3_sparse_idx,
+        threads,
+        fake_stream,
+        options="--enable-tvm-ffi",
+    )
+
+
+def compile_wignerd_panel_backward(
+    threads: int,
+) -> Callable:
+    if threads % 32 != 0:
+        raise ValueError("packed Wigner backward threads must be warp-aligned")
+    e = cute.sym_int64()
+    fake_q = make_fake_compact_tensor(cutlass.Float32, (e, 4), stride_order=(1, 0))
+    fake_grad_panel = make_fake_compact_tensor(
+        cutlass.Float32,
+        (e, SO2_PANEL_VALUES),
+        stride_order=(1, 0),
+    )
+    fake_grad_q = make_fake_compact_tensor(cutlass.Float32, (e, 4), stride_order=(1, 0))
+    fake_exp_l2 = make_fake_compact_tensor(cutlass.Int32, (35, 4), stride_order=(1, 0))
+    fake_c_l2_sparse = make_fake_compact_tensor(
+        cutlass.Float32,
+        (25, L2_SPARSE_TERMS),
+        stride_order=(1, 0),
+    )
+    fake_c_l2_sparse_idx = make_fake_compact_tensor(
+        cutlass.Int32,
+        (25, L2_SPARSE_TERMS),
+        stride_order=(1, 0),
+    )
+    fake_exp_l3 = make_fake_compact_tensor(cutlass.Int32, (84, 4), stride_order=(1, 0))
+    fake_c_l3_sparse = make_fake_compact_tensor(
+        cutlass.Float32,
+        (49, L3_SPARSE_TERMS),
+        stride_order=(1, 0),
+    )
+    fake_c_l3_sparse_idx = make_fake_compact_tensor(
+        cutlass.Int32,
+        (49, L3_SPARSE_TERMS),
+        stride_order=(1, 0),
+    )
+    fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True)
+    return cute.compile(
+        wignerd_panel_backward_jit,
+        fake_q,
+        fake_grad_panel,
+        fake_grad_q,
+        fake_exp_l2,
+        fake_c_l2_sparse,
+        fake_c_l2_sparse_idx,
+        fake_exp_l3,
+        fake_c_l3_sparse,
+        fake_c_l3_sparse_idx,
+        threads,
+        fake_stream,
+        options="--enable-tvm-ffi",
+    )
+
+
+@device_aware_lru_cache(maxsize=16)
+def _cached_wignerd_panel_forward(threads: int) -> Callable:
+    return compile_wignerd_panel_forward(threads)
+
+
+@device_aware_lru_cache(maxsize=16)
+def _cached_wignerd_panel_backward(threads: int) -> Callable:
+    return compile_wignerd_panel_backward(threads)
+
+
+_TABLE_CACHE: dict[tuple[str, int], tuple[torch.Tensor, ...]] = {}
+_TABLE_CACHE_LOCK = threading.Lock()
+
+
+def _device_cache_key(device: torch.device) -> tuple[str, int]:
+    index = -1 if device.index is None else int(device.index)
+    return (device.type, index)
+
+
+def _get_lmax3_tables(device: torch.device) -> tuple[torch.Tensor, ...]:
+    key = _device_cache_key(device)
+    tables = _TABLE_CACHE.get(key)
+    if tables is None:
+        with _TABLE_CACHE_LOCK:
+            tables = _TABLE_CACHE.get(key)
+            if tables is None:
+                tables = _build_l2_l3_tables(torch.float32, device)
+                _TABLE_CACHE[key] = tables
+    return tables
+
+
+def _wignerd_panel_impl(edge_quat: torch.Tensor) -> torch.Tensor:
+    q = edge_quat.detach().contiguous()
+    if q.dtype != torch.float32:
+        raise TypeError(f"packed WignerD requires float32, got {q.dtype}")
+
+    panel = torch.empty(
+        q.shape[0],
+        SO2_PANEL_VALUES,
+        device=q.device,
+        dtype=q.dtype,
+    )
+    if q.shape[0] == 0:
+        return panel
+
+    (
+        exp_l2,
+        c_l2_sparse,
+        c_l2_sparse_idx,
+        exp_l3,
+        c_l3_sparse,
+        c_l3_sparse_idx,
+    ) = _get_lmax3_tables(q.device)
+    with torch.cuda.device(q.device):
+        compiled = _cached_wignerd_panel_forward(32)
+        compiled(
+            q,
+            panel,
+            exp_l2,
+            c_l2_sparse,
+            c_l2_sparse_idx,
+            exp_l3,
+            c_l3_sparse,
+            c_l3_sparse_idx,
+        )
+    return panel
+
+
+def _wignerd_panel_bwd_impl(
+    grad_panel: torch.Tensor,
+    edge_quat: torch.Tensor,
+) -> torch.Tensor:
+    q = edge_quat.detach().contiguous()
+    if tuple(grad_panel.shape) != (q.shape[0], SO2_PANEL_VALUES):
+        raise ValueError(
+            f"packed Wigner gradient must have shape ({q.shape[0]}, {SO2_PANEL_VALUES})"
+        )
+    if q.dtype != torch.float32:
+        raise TypeError(f"packed WignerD requires float32, got {q.dtype}")
+    if q.shape[0] == 0:
+        return torch.empty_like(q)
+
+    (
+        exp_l2,
+        c_l2_sparse,
+        c_l2_sparse_idx,
+        exp_l3,
+        c_l3_sparse,
+        c_l3_sparse_idx,
+    ) = _get_lmax3_tables(q.device)
+    grad_q = torch.empty_like(q)
+    with torch.cuda.device(q.device):
+        compiled = _cached_wignerd_panel_backward(128)
+        compiled(
+            q,
+            grad_panel.detach().contiguous(),
+            grad_q,
+            exp_l2,
+            c_l2_sparse,
+            c_l2_sparse_idx,
+            exp_l3,
+            c_l3_sparse,
+            c_l3_sparse_idx,
+        )
+    return grad_q
+
+
+def _stateful_custom_op_tags() -> tuple[Any, ...] | None:
+    """Tag hidden-state runner ops as unsafe for direct CUDA graph capture."""
+    tag_type = getattr(getattr(torch, "_C", None), "Tag", None)
+    cudagraph_unsafe = getattr(tag_type, "cudagraph_unsafe", None)
+    if cudagraph_unsafe is None:
+        return None
+    return (cudagraph_unsafe,)
+
+
+_WIGNERD_CUSTOM_OP_TAGS = _stateful_custom_op_tags()
+_wignerd_panel_op = torch.library.custom_op(
+    "sezm_cute::wignerd_so2_panel", mutates_args=(), tags=_WIGNERD_CUSTOM_OP_TAGS
+)(_wignerd_panel_impl)
+_wignerd_panel_bwd_op = torch.library.custom_op(
+    "sezm_cute::wignerd_so2_panel_bwd", mutates_args=(), tags=_WIGNERD_CUSTOM_OP_TAGS
+)(_wignerd_panel_bwd_impl)
+
+
+@_wignerd_panel_op.register_fake
+def _(edge_quat: torch.Tensor) -> torch.Tensor:
+    return edge_quat.new_empty((edge_quat.shape[0], SO2_PANEL_VALUES))
+
+
+@_wignerd_panel_bwd_op.register_fake
+def _(grad_panel: torch.Tensor, edge_quat: torch.Tensor) -> torch.Tensor:
+    del grad_panel
+    return torch.empty_like(edge_quat)
+
+
+def _tensor_compute_capability(
+    tensor: torch.Tensor,
+) -> tuple[int, int] | None:
+    """Resolve dispatch capability from the Wigner operand's CUDA device."""
+    if tensor.device.type != "cuda":
+        return None
+    return tuple(torch.cuda.get_device_capability(tensor.device))
+
+
+def _wignerd_panel_setup_context(
+    ctx: Any,
+    inputs: tuple,
+    output: torch.Tensor,
+) -> None:
+    del output
+    (edge_quat,) = inputs
+    ctx.save_for_backward(edge_quat)
+
+
+def _wignerd_panel_registered_backward_impl(
+    ctx: Any,
+    grad_panel: torch.Tensor,
+) -> tuple[torch.Tensor]:
+    (edge_quat,) = ctx.saved_tensors
+    return (_wignerd_panel_bwd_op(grad_panel, edge_quat),)
+
+
+def _wignerd_panel_backward(
+    ctx: Any,
+    grad_panel: torch.Tensor,
+) -> tuple[torch.Tensor]:
+    """Run packed Wigner-D backward with its custom op visible to compilation."""
+    return _wignerd_panel_registered_backward_impl(ctx, grad_panel)
+
+
+_wignerd_panel_op.register_autograd(
+    _wignerd_panel_backward,
+    setup_context=_wignerd_panel_setup_context,
+)
+
+
+def _run_cute_wignerd_impl(
+    edge_quat: torch.Tensor,
+    wigner_calc: Any,
+    *,
+    packed_wigner: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    """Return CuTe Wigner data for the Neo ``lmax=3`` CUDA path.
+
+    Unsupported shapes/devices return ``None`` so the caller can use DeePMD's
+    original implementation unchanged. A prevalidated strict-FP32 packed
+    request returns the same ``(E,46)`` panel object in both tuple positions.
+    """
+    lmax = int(getattr(wigner_calc, "lmax", -1))
+    if (
+        not packed_wigner
+        or lmax != 3
+        or edge_quat.dim() != 2
+        or edge_quat.shape[-1] != 4
+        or not edge_quat.is_cuda
+        or edge_quat.dtype != torch.float32
+        or edge_quat.shape[0] == 0
+    ):
+        return None
+    compute_capability = _tensor_compute_capability(edge_quat)
+    if compute_capability is None or not runtime_policy.is_packed_wigner_enabled(
+        compute_capability
+    ):
+        return None
+    panel = _wignerd_panel_op(edge_quat)
+    return panel, panel
+
+
+def run_cute_wignerd(
+    edge_quat: torch.Tensor,
+    wigner_calc: Any,
+    *,
+    packed_wigner: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    """Run packed Wigner-D through its registered custom-op boundary."""
+    return _run_cute_wignerd_impl(
+        edge_quat,
+        wigner_calc,
+        packed_wigner=packed_wigner,
+    )
diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
index 7a8ac3fd43..0dc70ee7ee 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -5377,7 +5377,7 @@ def __call__(
 def prepare_triton_value_path_weights(model: torch.nn.Module) -> None:
     """Prepare fixed SO(2) weight layouts for every bound Triton value path."""
     for module in model.modules():
-        value_path = getattr(module, "_triton_value_path", None)
+        value_path = getattr(module, "triton_infer_l_2_value", None)
         if isinstance(value_path, _TritonSO2ValuePath):
             value_path.prepare_inference_weights()
 
diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py
index c6e54bd17f..2387b139dd 100644
--- a/deepmd/pt_expt/kernels/utils.py
+++ b/deepmd/pt_expt/kernels/utils.py
@@ -329,12 +329,12 @@ def operator_available(name: str) -> bool:
 
 
 def use_cute_infer() -> bool:
-    """Return whether the opt-in CuTe inference operator is enabled.
+    """Return whether the opt-in SeZM CuTe inference path is enabled.
 
     The flag is controlled by the ``DP_CUTE_INFER`` environment variable and is
-    read at module construction time. It selects the fused CuTe SO(2) value-path
-    operator (an independent path from ``DP_TRITON_INFER``) and only takes effect
-    during inference; training always uses the dense reference path.
+    read before accelerated state is prepared. CuTe replaces eligible SeZM spans
+    and leaves unsupported shapes to the selected Triton, CUDA, or reference
+    path. It only takes effect during inference.
 
     Returns
     -------
@@ -352,11 +352,10 @@ def use_cutile_infer() -> bool:
     written in the ``cuda.tile`` DSL and only takes effect during inference;
     training always uses the dense reference path.
 
-    The path is mutually exclusive with ``DP_TRITON_INFER`` and
-    ``DP_CUTE_INFER``: when it is enabled no Triton kernel executes, and a
-    convolution whose layout it does not support falls back to the dense
-    reference rather than to another accelerated backend. Enabling more than one
-    of the three is rejected at construction.
+    The path is mutually exclusive with ``DP_TRITON_INFER``: when it is enabled
+    no Triton kernel executes, and a convolution whose layout it does not support
+    falls back to the dense reference. ``DP_CUTE_INFER`` remains an independent
+    exact-shape overlay around the enclosing SeZM block.
 
     Returns
     -------
diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py
index aeb11d798a..606904ef2d 100644
--- a/deepmd/pt_expt/utils/serialization.py
+++ b/deepmd/pt_expt/utils/serialization.py
@@ -1495,7 +1495,7 @@ def _prepare_dpa4_triton_value_path_weights(
     if not _uses_dpa4_kernel_defaults(model_data):
         return
     if not any(
-        getattr(module, "_triton_value_path", None) is not None
+        getattr(module, "triton_infer_l_2_value", None) is not None
         for module in model.modules()
     ):
         return
@@ -1854,9 +1854,9 @@ def _trace_and_export_impl(
     force_fused_scatter = target_device.type == "cuda"
     for module in model.modules():
         if isinstance(module, SO2Linear):
-            module._force_block_diag_matmul = force_block_diag
+            module.force_block_diag_matmul = force_block_diag
         if isinstance(module, GeometricInitialEmbedding):
-            module._force_fused_scatter = force_fused_scatter
+            module.force_cuda_infer_l_1_scatter = force_fused_scatter
 
     if lower_kind == "graph" and not _supports_graph_export(model):
         raise NotImplementedError(
diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md
index 23d4e82327..e8479a3ca3 100644
--- a/doc/install/easy-install.md
+++ b/doc/install/easy-install.md
@@ -191,6 +191,22 @@ pip install deepmd-kit
 
 ::::::
 
+Optional CuTe inference kernels are available on Linux with Python 3.11 or
+newer for the PT and PT-expt backends on supported NVIDIA Ampere, Ada, Hopper,
+and Blackwell GPUs:
+
+```bash
+pip install "deepmd-kit[torch,cute]"
+```
+
+The `cute` extra installs the CUTLASS CuTe DSL, TVM FFI, and NVIDIA Alchemi
+Toolkit-Ops runtimes used by the optional CuTe paths. Set `DP_CUTE_INFER=1` to
+enable the DPA4-Neo implementation in PT or PT-expt inference. It may coexist
+with `DP_TRITON_INFER`: CuTe replaces exact supported spans, while unsupported
+model shapes retain the selected Triton or reference path. CuTe is not captured
+in a frozen `.pt2`; requesting it without the runtime dependencies raises an error
+rather than silently selecting another implementation.
+
 The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64 and arm64, and Windows x86-64.
 
 > [!WARNING]
diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md
index 5bb6a8097a..bf1f73a72e 100644
--- a/doc/model/dpa4.md
+++ b/doc/model/dpa4.md
@@ -471,29 +471,167 @@ Three options control training precision and the compiled path:
   through the energy gradient) and can speed training markedly on supported
   setups.
 
+Two environment variables independently select the fused training kernels:
+
+| Environment variable | Default | Effect                                                                                                             |
+| -------------------- | ------- | ------------------------------------------------------------------------------------------------------------------ |
+| `DP_TRITON_TRAIN`    | `0`     | Triton training level `0` or `1`. Level `1` enables the eligible second-order-complete stage kernels.              |
+| `DP_CUDA_TRAIN`      | off     | Enable the fused CUDA SO(2) value path on supported blocks. Accepted true values are `1`, `true`, `yes`, and `on`. |
+
+The corresponding false values for `DP_CUDA_TRAIN` are `0`, `false`, `no`,
+and `off`.
+
+#### Training kernel composition
+
+Training and inference use separate gates because their derivative contracts
+are different. Inference holds model parameters fixed and requires the forward
+and coordinate gradient. A force-loss training step additionally requires
+gradients of every consumed parameter and differentiates the coordinate
+gradient once more. The training kernels therefore provide their own parameter
+backward and second-order formulas rather than reusing inference-only kernels.
+
+`DP_TRITON_TRAIN=1` enables a composition of eligible Triton operators,
+including Wigner monomials, the local-frame rotations, radial degree mixing,
+SO(2) block GEMMs, gated activations, segmented attention softmax, flash
+attention aggregation, and the wider SO(3) grid products. Unsupported stages
+remain ordinary PyTorch operations. There are no training levels `2` or `3`:
+the fused Triton value path and fp16x3 mixing selected by
+`DP_TRITON_INFER>=2` are inference-only.
+
+`DP_CUDA_TRAIN=1` enables one fused CUDA operator for each supported SO(2)
+block. It owns the value stream from the source-node rotation through radial
+degree mixing, focus competition, and the gated mixing stack, stopping before
+the attention aggregation. The operator supplies parameter gradients and the
+second derivative required by force-loss training. A block outside its
+supported layout or an installation without the CUDA operator falls through to
+the enabled Triton stages or to PyTorch.
+
+The CUDA and Triton training gates are complementary, not mutually exclusive.
+For a span covered by both, CUDA takes precedence only over the SO(2) value
+stream; Triton remains active for the attention span and other non-overlapping
+operators. Their combinations are:
+
+| `DP_TRITON_TRAIN` | `DP_CUDA_TRAIN` | Result                                                                                                                                    |
+| ----------------- | --------------- | ----------------------------------------------------------------------------------------------------------------------------------------- |
+| `0`               | `0`             | PyTorch implements the complete training graph.                                                                                           |
+| `1`               | `0`             | Eligible stages use the Triton composition.                                                                                               |
+| `0`               | `1`             | Eligible value streams use CUDA; attention and unsupported blocks use PyTorch.                                                            |
+| `1`               | `1`             | Eligible value streams use CUDA, Triton serves attention and the other supported stages, and unsupported CUDA blocks fall back to Triton. |
+
+When every block supports the CUDA value path and Triton flash attention, the
+descriptor reuses the CUDA operator's packed Wigner runs instead of
+materializing dense per-edge Wigner-D matrices. Enabling both gates can
+therefore reduce memory as well as compose the two faster execution paths.
+The current inference-only force and virial assembly kernels are not selected
+during training. The training path keeps this assembly in the differentiable
+PyTorch graph so the force loss can traverse it.
+
+The intended combined setting is:
+
+```bash
+export DP_TRITON_TRAIN=1
+export DP_CUDA_TRAIN=1
+```
+
+Both gates serve PT and PT-expt and are read when the model is constructed.
+Set them before loading or constructing the model and before `torch.compile`.
+They do not change inference dispatch; likewise, `DP_CUTE_INFER`,
+`DP_TRITON_INFER`, and `DP_CUDA_INFER` do not change training dispatch.
+
 ### Inference and deployment settings
 
 Inference behavior is controlled by environment variables, each with an
 equivalent input-file option used during training validation:
 
-| Environment variable | Input-file option           | Default       | Effect                                                                                                                                                                                                                                                                                                                             |
-| -------------------- | --------------------------- | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `DP_COMPILE_INFER`   | `validating.compiled_infer` | off           | Use the compile path for evaluation/inference. Same `torch==2.11` / CUDA ≥ 12.6 requirements as `model.use_compile`.                                                                                                                                                                                                               |
-| `DP_TF32_INFER`      | `validating.tf32_infer`     | `0` (highest) | float32 matmul precision for inference: `0` highest, `1` high, `2` medium. Higher values improve throughput but make the potential energy surface less smooth.                                                                                                                                                                     |
-| `DP_AMP_INFER`       | `validating.amp_infer`      | off           | bf16 autocast inside the descriptor interaction blocks for inference, independently of `descriptor.use_amp`. Training AMP remains controlled by `descriptor.use_amp`. Usually keeps aggregate MAE similar but can make the potential energy surface less smooth.                                                                   |
-| `DP_TRITON_INFER`    | —                           | `0`           | Triton inference kernel level `0`-`3` (CUDA eval only, compatible with `DP_COMPILE_INFER`). Levels `1` and `2` are exact float32; level `3` trades a small accuracy margin for a substantial speedup. Detailed below.                                                                                                              |
-| `DP_CUTILE_INFER`    | —                           | off           | cuTile inference path (CUDA eval only, compatible with `DP_COMPILE_INFER`, mutually exclusive with `DP_TRITON_INFER`). Python inference only, and **not** captured in a frozen `.pt2`. Detailed below.                                                                                                                             |
-| `DP_CUDA_INFER`      | —                           | `0`           | Hand-written CUDA operator level `0`-`2` (CUDA eval only, stacks on top of `DP_TRITON_INFER`). Level `1` is faster on every GPU and checkpoint measured; level `2` additionally offers the fused convolution, which routes itself per checkpoint and falls back to the level-`1` behaviour where it would not pay. Detailed below. |
-
-Accepted boolean values for the other switches are `1`/`true`/`yes`/`on` and
-`0`/`false`/`no`/`off`; `DP_TRITON_INFER` accepts only the numeric levels.
-`DP_TRITON_INFER`, `DP_CUTILE_INFER` and `DP_CUTE_INFER` each select a complete
-accelerated inference path and are mutually exclusive; enabling more than one is
-rejected when the model is constructed.
+| Environment variable | Input-file option           | Default       | Effect                                                                                                                                                                                                                                                                                                                                  |
+| -------------------- | --------------------------- | ------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `DP_COMPILE_INFER`   | `validating.compiled_infer` | off           | Use the compile path for evaluation/inference. Same `torch==2.11` / CUDA ≥ 12.6 requirements as `model.use_compile`.                                                                                                                                                                                                                    |
+| `DP_TF32_INFER`      | `validating.tf32_infer`     | `0` (highest) | float32 matmul precision for inference: `0` highest, `1` high, `2` medium. Higher values improve throughput but make the potential energy surface less smooth.                                                                                                                                                                          |
+| `DP_AMP_INFER`       | `validating.amp_infer`      | off           | bf16 autocast inside the descriptor interaction blocks for inference, independently of `descriptor.use_amp`. Training AMP remains controlled by `descriptor.use_amp`. Usually keeps aggregate MAE similar but can make the potential energy surface less smooth.                                                                        |
+| `DP_TRITON_INFER`    | —                           | `0`           | Triton inference kernel level `0`-`3` (CUDA eval only, compatible with `DP_COMPILE_INFER`). Levels `1` and `2` are exact float32; level `3` trades a small accuracy margin for a substantial speedup. Detailed below.                                                                                                                   |
+| `DP_CUTE_INFER`      | —                           | off           | CuTe inference kernels for supported SeZM shapes (PT and PT-expt CUDA eval only, compatible with `DP_TRITON_INFER`, `DP_CUDA_INFER`, and `DP_COMPILE_INFER`). Unsupported shapes retain the selected fallback. Python inference only, and **not** captured in a frozen `.pt2`. Detailed below.                                          |
+| `DP_CUTILE_INFER`    | —                           | off           | cuTile inference path (CUDA eval only, compatible with `DP_COMPILE_INFER`, mutually exclusive with `DP_TRITON_INFER`). Python inference only, and **not** captured in a frozen `.pt2`. Detailed below.                                                                                                                                  |
+| `DP_CUDA_INFER`      | —                           | `0`           | Hand-written CUDA operator level `0`-`2` (CUDA eval only, composed with the CuTe and Triton paths). Level `1` is faster on every GPU and checkpoint measured; level `2` additionally offers the fused convolution, which routes itself per checkpoint and falls back to the level-`1` behaviour where it would not pay. Detailed below. |
+
+Accepted boolean values are `1`/`true`/`yes`/`on` and
+`0`/`false`/`no`/`off`. `DP_TRITON_INFER` and `DP_CUDA_INFER` instead accept
+only their documented numeric levels.
 Shell exports take precedence over the input-file options and over values
 written in the input; they are read when the model is constructed and changing
 them afterward has no effect.
 
+#### How the inference kernels compose
+
+`DP_CUTE_INFER`, `DP_TRITON_INFER`, and `DP_CUDA_INFER` are permissions for
+different implementations, not mutually exclusive backend selectors. The
+model selects one implementation independently for each computational span;
+enabling several paths never evaluates the same span more than once. A larger
+eligible fused span takes precedence, and an unsupported span falls through to
+the next enabled implementation.
+
+For the main SO(2) interaction block, the dispatch order is:
+
+1. the exact-shape CuTe block when its complete eligibility contract holds;
+1. the level-2 CUDA fused convolution when its structure and profitability
+   checks hold;
+1. the selected Triton composition;
+1. the dense PyTorch implementation.
+
+The surrounding operations are selected separately:
+
+| Computational span                                 | First available implementation                                     |
+| -------------------------------------------------- | ------------------------------------------------------------------ |
+| cutoff envelope and radial basis                   | CUDA level 1, then PyTorch                                         |
+| Wigner-D construction                              | CuTe packed Wigner-D, CUDA level 1, Triton monomials, then PyTorch |
+| complete SO(2) interaction block                   | CuTe, CUDA level 2, Triton, then PyTorch                           |
+| geometric initial embedding and SO(3) grid product | CuTe, CUDA level 1, then PyTorch                                   |
+| force and virial assembly                          | CUDA level 1, Triton level 1, then PyTorch                         |
+
+The packed Wigner-D layout used by the CuTe SO(2) path is descriptor-wide.
+Consequently, every interaction block must satisfy the CuTe contract before
+any block enters that path. Other CuTe kernels, such as the grid product, apply
+their own local eligibility checks. CUDA level 1 remains useful around a CuTe
+block because it covers independent spans such as the radial basis and force
+assembly.
+
+With `DP_CUTILE_INFER=0`, the common combinations behave as follows, where `T`
+is a nonzero Triton level:
+
+| `DP_CUTE_INFER` | `DP_TRITON_INFER` | `DP_CUDA_INFER` | Result                                                                                                                                 |
+| --------------- | ----------------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------- |
+| `0`             | `0`               | `0`             | The dense PyTorch path is used.                                                                                                        |
+| `0`             | `T`               | `0`             | Triton serves its supported spans; other spans remain in PyTorch.                                                                      |
+| `0`             | `T`               | `1`             | CUDA level 1 serves its spans, while Triton serves the remaining accelerated spans.                                                    |
+| `0`             | `T`               | `2`             | Eligible SO(2) blocks use the CUDA fused convolution; ineligible blocks fall back to Triton. CUDA level 1 remains enabled.             |
+| `1`             | any               | any             | Eligible exact-shape spans use CuTe; the configured CUDA, Triton, and PyTorch paths remain as fallbacks and for non-overlapping spans. |
+
+When CUDA level 2 takes over an SO(2) block, its fused convolution also takes
+over the mixing stack. The fp16x3 GEMMs added by Triton level 3 therefore do not
+run inside that block, although they remain available to a block rejected by
+the CUDA gate. The same principle applies to a block selected by CuTe.
+
+Among these three variables there is no prohibited combination. Do not confuse
+`DP_CUTE_INFER` with `DP_CUTILE_INFER`: cuTile and Triton are mutually exclusive
+complete SO(2) paths, while CuTe is an exact-shape overlay that may coexist with
+either path and with the CUDA operators.
+
+With the default `DP_TF32_INFER=0` and `DP_AMP_INFER=0` precision policy, a
+conservative exact-float32 Python inference configuration enables CuTe with
+both general fallbacks:
+
+```bash
+export DP_CUTE_INFER=1
+export DP_TRITON_INFER=2
+export DP_CUDA_INFER=1
+```
+
+CUDA level 2 can replace level 1 for narrow checkpoints on GPUs where its fused
+SO(2) convolution is profitable. See
+[Training kernel composition](#training-kernel-composition) for the independent
+training gates.
+
+#### Triton levels
+
 `DP_TRITON_INFER` selects how much of the descriptor runs in fused Triton
 kernels. Level `1` adds universal fused kernels, numerically equivalent to the
 dense path with full float32 accumulation. Level `2` adds the table-configured
@@ -507,6 +645,20 @@ shipping built in. On other GPUs the kernels fall back to conservative
 configurations, and `dp --pt freeze` tunes the missing entries on the local GPU
 before exporting, a one-off sweep of a few minutes baked into the `.pt2`.
 
+#### CuTe overlay
+
+`DP_CUTE_INFER` enables CuTe kernels for the current DPA4-Neo inference
+contract in both PT and PT-expt. The path fuses the eligible SO(2) interaction
+block, packed Wigner-D construction, geometric initial embedding, output grid
+product, and scalar readout while retaining strict float32 arithmetic.
+Eligibility checks cover the model structure, tensor layout, dtype, device, and
+GPU capability. Any unsupported configuration falls through to the enabled
+Triton, CUDA, cuTile, or reference implementation rather than entering a
+partially compatible kernel. The kernels are JIT compiled at runtime and are
+not captured in a frozen `.pt2`.
+
+#### cuTile path
+
 `DP_CUTILE_INFER` replaces the whole SeZM edge pipeline — Wigner monomials,
 rotate-and-mix, the gated SO(2) mixing stack, the attention aggregation and the
 force / virial assembly — with kernels written in the `cuda.tile` DSL. On an
@@ -521,6 +673,8 @@ whose layout it does not support falls back to the dense reference rather than
 to Triton. The kernels are JIT compiled at runtime, so this path serves Python
 inference only and is not captured in a frozen `.pt2`.
 
+#### CUDA levels
+
 `DP_CUDA_INFER` enables hand-written CUDA operators that fuse spans of the SeZM
 descriptor. Unlike the paths above it is not an alternative backend: it stacks
 on top of `DP_TRITON_INFER`, taking over the spans it covers and leaving the
@@ -584,6 +738,8 @@ Both levels are precompiled custom operators that `make_fx` traces, so unlike
 `DP_CUTILE_INFER` they are baked into a frozen `.pt2` and keep their effect when
 it is later loaded by ASE or LAMMPS.
 
+#### Precision guidance
+
 For molecular dynamics and other workflows sensitive to the smoothness of the
 potential energy surface, keep `DP_TF32_INFER=0` and `DP_AMP_INFER=0`.
 `DP_AMP_INFER` can coexist with `DP_TF32_INFER`, but bf16 autocast dominates
@@ -596,6 +752,8 @@ for the target system. `DP_CUTILE_INFER` inherits that same rounding scale
 through its mixing stack and is the faster of the two on Blackwell, at the cost
 of being unavailable to the frozen `.pt2` route.
 
+#### Frozen `.pt2` kernel selection
+
 > [!IMPORTANT]
 > Set these variables **before** running `dp --pt freeze` or
 > `dp --pt-expt freeze`. The exported `.pt2` is an AOTInductor artifact, so the
@@ -610,9 +768,9 @@ of being unavailable to the frozen `.pt2` route.
 > chosen levels and whether each came from the environment or the default are
 > logged at export. A CPU-targeted archive disables GPU-only inference paths
 > and keeps the reference CPU implementation regardless of these settings.
-> `DP_CUTILE_INFER` is the exception:
-> its kernels are JIT compiled at runtime and do not bake into the artifact, so
-> it applies to Python inference only and has no effect on a frozen model. A frozen `.pt2` runs a forward-only
+> `DP_CUTILE_INFER` and `DP_CUTE_INFER` are the exceptions: their kernels are
+> JIT compiled at runtime and do not bake into the artifact, so they apply to
+> Python inference only and have no effect on a frozen model. A frozen `.pt2` runs a forward-only
 > package, so training-time memory-saving switches do not apply to it.
 
 ### Hardware selection
diff --git a/pyproject.toml b/pyproject.toml
index 5da7f9b1c6..b5e00d4a27 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -146,6 +146,11 @@ cu12 = [
     "nvidia-cudnn-cu12",
     "nvidia-cuda-nvcc-cu12",
 ]
+cute = [
+    'nvidia-cutlass-dsl[cu13]>=4.6.1,<5; python_version >= "3.11" and platform_system == "Linux"',
+    'apache-tvm-ffi; python_version >= "3.11" and platform_system == "Linux"',
+    'nvalchemi-toolkit-ops>=0.3.1; python_version >= "3.11" and platform_system == "Linux"',
+]
 jax = [
     # below is a funny workaround for
     # https://github.com/astral-sh/uv/issues/8601
diff --git a/source/tests/common/dpmodel/test_dpa4_edge_cache.py b/source/tests/common/dpmodel/test_dpa4_edge_cache.py
index 17f823f69f..0cdcae6856 100644
--- a/source/tests/common/dpmodel/test_dpa4_edge_cache.py
+++ b/source/tests/common/dpmodel/test_dpa4_edge_cache.py
@@ -2,13 +2,34 @@
 """Tests for the backend-neutral DPA4 edge-cache acceleration seams."""
 
 import numpy as np
+import pytest
 
 from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import (
     _edge_cache_from_arrays,
+    _finalize_edge_cache,
     edge_cache_to_dtype,
 )
 
 
+@pytest.mark.parametrize("floor", [0.0, -1.0])
+def test_degree_normalization_rejects_nonpositive_floor(floor: float) -> None:
+    with pytest.raises(ValueError, match="deg_norm_floor must be positive"):
+        _finalize_edge_cache(
+            n_nodes=1,
+            src=np.zeros(1, dtype=np.int64),
+            dst=np.zeros(1, dtype=np.int64),
+            edge_type_feat=np.ones((1, 1)),
+            edge_vec=np.ones((1, 3)),
+            edge_rbf=np.ones((1, 1)),
+            edge_env=np.ones((1, 1)),
+            D_full=None,
+            Dt_full=None,
+            D_packed=None,
+            edge_quat=np.ones((1, 4)),
+            deg_norm_floor=floor,
+        )
+
+
 def test_fused_builders_replace_reference_and_initialize_step_cache() -> None:
     calls: dict[str, int] = {"radial": 0, "wigner": 0}
     keep_seen: list[np.ndarray] = []
@@ -66,3 +87,42 @@ def fused_wigner(quaternion: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
 
     cache.csr_cache = None
     assert edge_cache_to_dtype(cache, np.float32).csr_cache is None
+
+
+def test_edge_cache_normalizes_graph_csr_to_int64() -> None:
+    def wigner_calc(quaternion: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        block = np.ones((quaternion.shape[0], 1, 1), dtype=quaternion.dtype)
+        return block, block
+
+    cache = _edge_cache_from_arrays(
+        type_ebed=np.array([[1.0], [2.0]]),
+        edge_index=np.array([[1, 0], [0, 1]], dtype=np.int32),
+        edge_vec=np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]),
+        edge_mask=np.array([True, True]),
+        compute_dtype=np.float64,
+        eps=1.0e-12,
+        deg_norm_floor=1.0,
+        inner_clamp=None,
+        bridging_switch=None,
+        edge_envelope=np.ones_like,
+        radial_basis=np.ones_like,
+        random_gamma=False,
+        wigner_calc=wigner_calc,
+        csr_cache={
+            "dst": (
+                np.array([0, 1], dtype=np.int32),
+                np.array([0, 1, 2], dtype=np.int32),
+            ),
+            "src": (
+                np.array([1, 0], dtype=np.int32),
+                np.array([0, 1, 2], dtype=np.int32),
+            ),
+        },
+    )
+
+    assert cache.src.dtype == np.int64
+    assert cache.dst.dtype == np.int64
+    assert cache.csr_cache is not None
+    for order, row_ptr in cache.csr_cache.values():
+        assert order.dtype == np.int64
+        assert row_ptr.dtype == np.int64
diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py
index d147ab3dd3..5ba771be86 100644
--- a/source/tests/common/dpmodel/test_neighbor_graph.py
+++ b/source/tests/common/dpmodel/test_neighbor_graph.py
@@ -8,6 +8,7 @@
     NeighborGraph,
     build_edge_csr,
     canonicalize_neighbor_graph,
+    compact_edges,
 )
 
 
@@ -141,6 +142,79 @@ def test_canonicalization_rebuilds_an_untrusted_canonical_claim(self) -> None:
             np.array([[1, 0], [0, 1]], dtype=np.int64),
         )
 
+    def test_compaction_preserves_canonical_csr(self) -> None:
+        graph = NeighborGraph(
+            n_node=np.array([3], dtype=np.int64),
+            edge_index=np.array([[2, 0, 0, 0], [0, 1, 0, 0]], dtype=np.int64),
+            edge_vec=np.arange(12, dtype=np.float64).reshape(4, 3),
+            edge_mask=np.array([True, True, False, False]),
+            destination_order=np.arange(4, dtype=np.int64),
+            destination_row_ptr=np.array([0, 1, 2, 2], dtype=np.int64),
+            source_order=np.array([1, 0, 2, 3], dtype=np.int64),
+            source_row_ptr=np.array([0, 1, 1, 2], dtype=np.int64),
+            destination_sorted=True,
+        )
+
+        result = compact_edges(graph)
+
+        np.testing.assert_array_equal(result.edge_index, graph.edge_index[:, :2])
+        np.testing.assert_array_equal(result.edge_vec, graph.edge_vec[:2])
+        np.testing.assert_array_equal(result.edge_mask, [True, True])
+        np.testing.assert_array_equal(result.destination_order, [0, 1])
+        np.testing.assert_array_equal(result.destination_row_ptr, [0, 1, 2, 2])
+        np.testing.assert_array_equal(result.source_order, [1, 0])
+        np.testing.assert_array_equal(result.source_row_ptr, [0, 1, 1, 2])
+        self.assertTrue(result.destination_sorted)
+
+    def test_compaction_preserves_csr_with_internal_mask(self) -> None:
+        graph = NeighborGraph(
+            n_node=np.array([2], dtype=np.int64),
+            edge_index=np.array([[1, 0, 0], [0, 0, 1]], dtype=np.int64),
+            edge_vec=np.arange(9, dtype=np.float64).reshape(3, 3),
+            edge_mask=np.array([True, False, True]),
+            destination_order=np.array([0, 1, 2], dtype=np.int64),
+            destination_row_ptr=np.array([0, 2, 3], dtype=np.int64),
+            source_order=np.array([1, 2, 0], dtype=np.int64),
+            source_row_ptr=np.array([0, 2, 3], dtype=np.int64),
+            destination_sorted=True,
+        )
+
+        result = compact_edges(graph)
+
+        np.testing.assert_array_equal(result.edge_index, graph.edge_index[:, [0, 2]])
+        np.testing.assert_array_equal(result.edge_vec, graph.edge_vec[[0, 2]])
+        np.testing.assert_array_equal(result.edge_mask, [True, True])
+        np.testing.assert_array_equal(result.destination_order, [0, 1])
+        np.testing.assert_array_equal(result.destination_row_ptr, [0, 1, 2])
+        np.testing.assert_array_equal(result.source_order, [1, 0])
+        np.testing.assert_array_equal(result.source_row_ptr, [0, 1, 2])
+        self.assertTrue(result.destination_sorted)
+
+    def test_compaction_preserves_permuted_csr_with_internal_mask(self) -> None:
+        graph = NeighborGraph(
+            n_node=np.array([3], dtype=np.int64),
+            edge_index=np.array(
+                [[2, 0, 1, 0], [1, 2, 0, 1]],
+                dtype=np.int32,
+            ),
+            edge_vec=np.arange(12, dtype=np.float64).reshape(4, 3),
+            edge_mask=np.array([True, False, True, True]),
+            destination_order=np.array([2, 0, 3, 1], dtype=np.int32),
+            destination_row_ptr=np.array([0, 1, 3, 4], dtype=np.int64),
+            source_order=np.array([1, 3, 2, 0], dtype=np.int32),
+            source_row_ptr=np.array([0, 2, 3, 4], dtype=np.int64),
+            destination_sorted=False,
+        )
+
+        result = compact_edges(graph)
+
+        np.testing.assert_array_equal(result.edge_index, graph.edge_index[:, [0, 2, 3]])
+        np.testing.assert_array_equal(result.destination_order, [1, 0, 2])
+        np.testing.assert_array_equal(result.destination_row_ptr, [0, 1, 3, 3])
+        np.testing.assert_array_equal(result.source_order, [2, 1, 0])
+        np.testing.assert_array_equal(result.source_row_ptr, [0, 1, 2, 3])
+        self.assertFalse(result.destination_sorted)
+
 
 from deepmd.dpmodel.utils.neighbor_graph import (
     node_ownership_mask,
diff --git a/source/tests/pt/model/sezm_cute/__init__.py b/source/tests/pt/model/sezm_cute/__init__.py
new file mode 100644
index 0000000000..6ceb116d85
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
diff --git a/source/tests/pt/model/sezm_cute/_paths.py b/source/tests/pt/model/sezm_cute/_paths.py
new file mode 100644
index 0000000000..c635e9f780
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/_paths.py
@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Source paths for CuTe kernels that cannot be imported without CUDA."""
+
+from pathlib import (
+    Path,
+)
+
+from deepmd.pt_expt.kernels.cute import (
+    sezm,
+)
+
+CUTE_ROOT = Path(sezm.__path__[0])
diff --git a/source/tests/pt/model/sezm_cute/test_compile_cache.py b/source/tests/pt/model/sezm_cute/test_compile_cache.py
new file mode 100644
index 0000000000..5c62379106
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_compile_cache.py
@@ -0,0 +1,94 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Behavioral tests for device-aware CuTe compile caching."""
+
+from __future__ import (
+    annotations,
+)
+
+import sys
+from types import (
+    SimpleNamespace,
+)
+from unittest import (
+    mock,
+)
+
+from deepmd.pt_expt.kernels.cute.sezm import (
+    compile_cache,
+)
+
+
+def test_cache_separates_device_and_compute_capability() -> None:
+    identity = [(0, 8, 0)]
+    calls: list[tuple[str, tuple[int, int, int]]] = []
+
+    @compile_cache.device_aware_lru_cache(
+        maxsize=8,
+        identity_getter=lambda: identity[0],
+    )
+    def compile_kernel(mode: str):
+        calls.append((mode, identity[0]))
+        return object()
+
+    sm80_first = compile_kernel("strict-fp32")
+    assert compile_kernel("strict-fp32") is sm80_first
+
+    identity[0] = (1, 9, 0)
+    sm90 = compile_kernel("strict-fp32")
+    assert sm90 is not sm80_first
+
+    identity[0] = (0, 8, 0)
+    assert compile_kernel("strict-fp32") is sm80_first
+    assert calls == [
+        ("strict-fp32", (0, 8, 0)),
+        ("strict-fp32", (1, 9, 0)),
+    ]
+
+
+def test_cache_exposes_standard_cache_controls() -> None:
+    @compile_cache.device_aware_lru_cache(
+        maxsize=2,
+        identity_getter=lambda: (0, 8, 6),
+    )
+    def compile_kernel(rows: int):
+        return object()
+
+    compile_kernel(4)
+    assert compile_kernel.cache_info().currsize == 1
+    compile_kernel.cache_clear()
+    assert compile_kernel.cache_info().currsize == 0
+
+
+def test_cache_compiles_inside_the_keyed_cuda_device() -> None:
+    entered: list[tuple[str, int]] = []
+
+    class DeviceContext:
+        def __init__(self, index: int) -> None:
+            self.index = index
+
+        def __enter__(self) -> None:
+            entered.append(("enter", self.index))
+
+        def __exit__(self, *_args) -> None:
+            entered.append(("exit", self.index))
+
+    fake_torch = SimpleNamespace(
+        cuda=SimpleNamespace(
+            is_available=lambda: True,
+            device=lambda index: DeviceContext(index),
+        )
+    )
+
+    @compile_cache.device_aware_lru_cache(
+        maxsize=2,
+        identity_getter=lambda: (3, 9, 0),
+    )
+    def compile_kernel() -> object:
+        entered.append(("compile", 3))
+        return object()
+
+    with mock.patch.dict(sys.modules, {"torch": fake_torch}):
+        compile_kernel()
+
+    assert entered == [("enter", 3), ("compile", 3), ("exit", 3)]
diff --git a/source/tests/pt/model/sezm_cute/test_efs.py b/source/tests/pt/model/sezm_cute/test_efs.py
new file mode 100644
index 0000000000..a729ba5fd1
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_efs.py
@@ -0,0 +1,518 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""End-to-end strict-FP32 acceptance tests for the Neo CuTe path."""
+
+from __future__ import (
+    annotations,
+)
+
+import contextlib
+import importlib
+import os
+import subprocess
+import sys
+from pathlib import (
+    Path,
+)
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
+
+import pytest
+import torch
+
+if TYPE_CHECKING:
+    from collections.abc import (
+        Iterator,
+    )
+
+
+TOL = 5.0e-5
+OUTPUT_KEYS = ("energy", "force", "atom_energy", "virial")
+SM80_CAPABILITIES = frozenset({(8, 0), (8, 6)})
+SUPPORTED_CAPABILITIES = SM80_CAPABILITIES | frozenset(
+    {(8, 9), (9, 0), (10, 0), (12, 0)}
+)
+_CHILD_MODE = "--neo-cute-efs-child"
+
+
+def _cute_runtime_skip_reason() -> str | None:
+    if not torch.cuda.is_available():
+        return "Neo CuTe E/F/S parity requires CUDA"
+    capability = tuple(torch.cuda.get_device_capability())
+    if capability not in SUPPORTED_CAPABILITIES:
+        return f"Neo CuTe E/F/S parity does not support {capability}"
+    try:
+        importlib.import_module("cutlass.cute")
+        importlib.import_module("cuda.bindings.driver")
+        importlib.import_module("tvm_ffi")
+    except Exception as exc:  # pragma: no cover - runtime dependent
+        return f"Neo CuTe E/F/S acceptance requires the CuTe DSL runtime: {exc}"
+    return None
+
+
+_CUTE_SKIP_REASON = _cute_runtime_skip_reason()
+
+
+@contextlib.contextmanager
+def _strict_fp32() -> Iterator[None]:
+    matmul = torch.backends.cuda.matmul
+    cudnn = torch.backends.cudnn
+    prior_matmul_tf32 = matmul.allow_tf32
+    prior_cudnn_tf32 = cudnn.allow_tf32
+    prior_precision = torch.get_float32_matmul_precision()
+    try:
+        matmul.allow_tf32 = False
+        cudnn.allow_tf32 = False
+        torch.set_float32_matmul_precision("highest")
+        yield
+    finally:
+        matmul.allow_tf32 = prior_matmul_tf32
+        cudnn.allow_tf32 = prior_cudnn_tf32
+        torch.set_float32_matmul_precision(prior_precision)
+
+
+def _neo_model(*, use_compile: bool, sel: int = 32) -> torch.nn.Module:
+    from deepmd.pt.model.model import (
+        get_sezm_model,
+    )
+
+    model = get_sezm_model(
+        {
+            "type": "SeZM",
+            "type_map": ["O", "H"],
+            "descriptor": {
+                "type": "SeZM",
+                "sel": sel,
+                "rcut": 3.0,
+                "channels": 32,
+                "n_radial": 16,
+                "use_env_seed": True,
+                "lmax": 3,
+                "mmax": 1,
+                "n_blocks": 2,
+                "so2_layers": 3,
+                "radial_so2_mode": "degree_channel",
+                "radial_so2_rank": 1,
+                "n_focus": 2,
+                "focus_dim": 0,
+                "n_atten_head": 1,
+                "message_node_so3": True,
+                "ffn_neurons": 0,
+                "ffn_so3_grid": True,
+                "grid_mlp": False,
+                "grid_branch": [0, 0, 1],
+                "ffn_blocks": 1,
+                "so3_readout": "mlp",
+                "use_amp": False,
+                "precision": "float32",
+                "seed": 42,
+            },
+            "fitting_net": {
+                "neuron": [0],
+                "precision": "float32",
+                "seed": 42,
+            },
+            "use_compile": use_compile,
+            "enable_tf32": False,
+        }
+    ).to(device="cuda", dtype=torch.float32)
+    model.eval()
+    for parameter in model.parameters():
+        parameter.requires_grad_(False)
+    return model
+
+
+def _water_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    coord = torch.tensor(
+        [
+            [
+                [0.70, 0.80, 0.90],
+                [1.58, 0.80, 0.90],
+                [0.43, 1.63, 0.90],
+                [3.20, 3.00, 2.80],
+                [4.08, 3.00, 2.80],
+                [2.93, 3.83, 2.80],
+            ]
+        ],
+        device="cuda",
+        dtype=torch.float64,
+    )
+    atype = torch.tensor(
+        [[0, 1, 1, 0, 1, 1]],
+        device="cuda",
+        dtype=torch.int32,
+    )
+    box = torch.tensor(
+        [[5.4, 0.0, 0.0, 0.0, 5.2, 0.0, 0.0, 0.0, 5.0]],
+        device="cuda",
+        dtype=torch.float64,
+    )
+    return coord, atype, box
+
+
+def _triclinic_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    box_matrix = torch.tensor(
+        [
+            [5.2, 0.0, 0.0],
+            [0.7, 4.8, 0.0],
+            [0.3, 0.5, 5.0],
+        ],
+        device="cuda",
+        dtype=torch.float32,
+    )
+    fractional = torch.tensor(
+        [
+            [0.08, 0.11, 0.14],
+            [0.25, 0.14, 0.18],
+            [0.12, 0.31, 0.22],
+            [0.54, 0.57, 0.49],
+            [0.72, 0.59, 0.51],
+            [0.57, 0.76, 0.55],
+            [0.88, 0.16, 0.81],
+            [0.06, 0.22, 0.84],
+        ],
+        device="cuda",
+        dtype=torch.float32,
+    )
+    coord = (fractional @ box_matrix).unsqueeze(0)
+    atype = torch.tensor(
+        [[0, 1, 1, 0, 1, 1, 0, 1]],
+        device="cuda",
+        dtype=torch.int32,
+    )
+    return coord, atype, box_matrix.reshape(1, 9)
+
+
+def _high_degree_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    node_count = 258
+    coord = torch.zeros(
+        (1, node_count, 3),
+        device="cuda",
+        dtype=torch.float32,
+    )
+    coord[0, :, 0] = (
+        torch.arange(
+            node_count,
+            device="cuda",
+            dtype=torch.float32,
+        )
+        * 0.01
+    )
+    atype = (torch.arange(node_count, device="cuda") % 2).to(dtype=torch.int32)[None, :]
+    return coord, atype, None
+
+
+def _detached_outputs(outputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    return {name: outputs[name].detach().cpu().clone() for name in OUTPUT_KEYS}
+
+
+_FRAME_FACTORIES = {
+    "water": _water_frame,
+    "triclinic": _triclinic_frame,
+    "high_degree": _high_degree_frame,
+}
+_STANDARD_FRAME_NAMES = ("water", "triclinic")
+
+
+def _assert_runtime_profile() -> dict[str, bool]:
+    from deepmd.pt_expt.kernels.cute.sezm import (
+        runtime_policy,
+    )
+    from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2
+
+    capability = tuple(torch.cuda.get_device_capability())
+    if capability not in SUPPORTED_CAPABILITIES:
+        raise AssertionError(f"unsupported SO2 capability {capability}")
+    policy_checks = {
+        "master": runtime_policy.is_cute_infer_enabled(),
+        "triton_infer_2": os.environ.get("DP_TRITON_INFER") == "2",
+        "supported": runtime_policy.is_supported_so2_capability(capability),
+        "packed_wigner": runtime_policy.is_packed_wigner_enabled(capability),
+    }
+    config = so2._architecture_default_config(capability)
+    if capability in SM80_CAPABILITIES:
+        architecture_checks = {
+            "sm80_profile": runtime_policy.is_sm80_profile_enabled(capability),
+            "gie": runtime_policy.is_gie_enabled(capability),
+            "thin_wrapper": runtime_policy.is_so2_thin_wrapper_enabled(capability),
+            "output_grid_bwd_panel": (
+                runtime_policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability)
+            ),
+            "output_grid_fwd": (
+                runtime_policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability)
+            ),
+            "readout_fold": runtime_policy.is_readout_input_fold_sm80_enabled(
+                capability
+            ),
+            "per_focus_so2_forward": config.per_focus_so2_fwd_pair,
+            "no_native_sm90_path": not config.native_sm90_path,
+        }
+    elif capability == (9, 0):
+        architecture_checks = {
+            "native_sm90_path": config.native_sm90_path,
+            "shared_so2_forward": not config.per_focus_so2_fwd_pair,
+            "sm90_output_grid": (
+                runtime_policy.is_output_grid_sm90_c96_asymmetric_panels_enabled(
+                    capability
+                )
+            ),
+            "sm90_readout_fold": runtime_policy.is_readout_input_fold_sm90_enabled(
+                capability
+            ),
+        }
+    elif capability in runtime_policy.FUSED_SO2_GATE_CAPABILITIES:
+        architecture_checks = {
+            "combined_so2_gate": config.combined_so2_gate,
+            "shared_so2_forward": not config.per_focus_so2_fwd_pair,
+            "no_native_sm90_path": not config.native_sm90_path,
+        }
+    else:
+        architecture_checks = {
+            "no_native_sm90_path": not config.native_sm90_path,
+            "shared_so2_forward": not config.per_focus_so2_fwd_pair,
+        }
+    selected = {**policy_checks, **architecture_checks}
+    missing = sorted(name for name, enabled in selected.items() if not enabled)
+    if missing:
+        raise AssertionError(
+            "DP_CUTE_INFER did not select the expected runtime profile: "
+            + ", ".join(missing)
+        )
+    return selected
+
+
+def _run_efs_child(
+    *,
+    frame_name: str,
+    use_compile: bool,
+    output_path: Path,
+) -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2
+
+    torch.manual_seed(20260726)
+    torch.cuda.manual_seed_all(20260726)
+    runner_calls = 0
+    runner_types: set[str] = set()
+    max_destination_degree = 0
+    original_build_runner = so2._build_runner
+
+    def counted_build_runner(*args: Any, **kwargs: Any) -> Any:
+        nonlocal runner_calls, max_destination_degree
+        runner = original_build_runner(*args, **kwargs)
+        runner_calls += 1
+        runner_types.add(f"{type(runner).__module__}.{type(runner).__qualname__}")
+        if runner_calls == 1:
+            destination_row_ptr = args[8]
+            degrees = destination_row_ptr[1:] - destination_row_ptr[:-1]
+            max_destination_degree = int(degrees.max().item())
+        return runner
+
+    if use_compile:
+        so2._build_runner = counted_build_runner
+    try:
+        model = _neo_model(
+            use_compile=use_compile,
+            sel=288 if frame_name == "high_degree" else 32,
+        )
+        state_keys_before = tuple(model.state_dict())
+        coord, atype, box = _FRAME_FACTORIES[frame_name]()
+        with _strict_fp32():
+            profile = _assert_runtime_profile() if use_compile else {}
+            run_count = 2 if use_compile else 1
+            runs = [
+                _detached_outputs(model(coord, atype, box=box))
+                for _ in range(run_count)
+            ]
+            torch.cuda.synchronize()
+        state_keys_after = tuple(model.state_dict())
+    finally:
+        if use_compile:
+            so2._build_runner = original_build_runner
+
+    if state_keys_after != state_keys_before:
+        raise AssertionError("CuTe warmup changed the model state_dict contract")
+    if use_compile and runner_calls == 0:
+        raise AssertionError("compiled CuTe run did not instantiate the SO2 runner")
+    torch.save(
+        {
+            "outputs": runs,
+            "profile": profile,
+            "runner_calls": runner_calls,
+            "runner_types": sorted(runner_types),
+            "max_destination_degree": max_destination_degree,
+            "use_compile": use_compile,
+        },
+        output_path,
+    )
+
+
+def _clean_child_environment() -> dict[str, str]:
+    environment = os.environ.copy()
+    for name in tuple(environment):
+        if name.startswith("DP_CUTE_"):
+            environment.pop(name)
+    for name in (
+        "DP_COMPILE_INFER",
+        "DP_INTERFACE_PREC",
+        "DP_TF32_INFER",
+        "DP_TRITON_INFER",
+        "NVIDIA_TF32_OVERRIDE",
+    ):
+        environment.pop(name, None)
+    environment.update(
+        {
+            "DP_INTERFACE_PREC": "low",
+            "DP_TF32_INFER": "0",
+            "NVIDIA_TF32_OVERRIDE": "0",
+            "PYTHONDONTWRITEBYTECODE": "1",
+        }
+    )
+    return environment
+
+
+def _child_environment(*, use_compile: bool) -> dict[str, str]:
+    environment = _clean_child_environment()
+    if use_compile:
+        environment.update(
+            {
+                "DP_COMPILE_INFER": "1",
+                "DP_CUTE_STRICT": "1",
+                "DP_CUTE_INFER": "1",
+                "DP_TRITON_INFER": "2",
+            }
+        )
+    else:
+        environment.update(
+            {
+                "DP_COMPILE_INFER": "0",
+                "DP_CUTE_INFER": "0",
+                "DP_TRITON_INFER": "0",
+            }
+        )
+    return environment
+
+
+def _run_child_process(
+    *,
+    frame_name: str,
+    use_compile: bool,
+    output_path: Path,
+) -> dict[str, Any]:
+    result = subprocess.run(
+        [
+            sys.executable,
+            str(Path(__file__).resolve()),
+            _CHILD_MODE,
+            frame_name,
+            "compiled" if use_compile else "eager",
+            str(output_path),
+        ],
+        env=_child_environment(use_compile=use_compile),
+        capture_output=True,
+        text=True,
+        timeout=900,
+        check=False,
+    )
+    if result.returncode:
+        pytest.fail(
+            "Neo CuTe E/F/S child failed.\n"
+            f"stdout:\n{result.stdout}\n"
+            f"stderr:\n{result.stderr}"
+        )
+    return torch.load(output_path, map_location="cpu", weights_only=True)
+
+
+@pytest.mark.skipif(
+    _CUTE_SKIP_REASON is not None,
+    reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable",
+)
+@pytest.mark.parametrize("frame_name", _STANDARD_FRAME_NAMES)
+def test_neo_cute_compiled_profile_matches_eager_at_shared_5e5(
+    tmp_path: Path,
+    frame_name: str,
+) -> None:
+    eager = _run_child_process(
+        frame_name=frame_name,
+        use_compile=False,
+        output_path=tmp_path / f"{frame_name}-eager.pt",
+    )
+    compiled = _run_child_process(
+        frame_name=frame_name,
+        use_compile=True,
+        output_path=tmp_path / f"{frame_name}-compiled.pt",
+    )
+
+    assert not eager["use_compile"]
+    assert compiled["use_compile"]
+    assert compiled["runner_calls"] > 0
+    assert compiled["profile"]
+    assert all(compiled["profile"].values())
+    expected = eager["outputs"][0]
+    for run_index, actual in enumerate(compiled["outputs"]):
+        for name in OUTPUT_KEYS:
+            torch.testing.assert_close(
+                actual[name],
+                expected[name],
+                atol=TOL,
+                rtol=TOL,
+                msg=(
+                    f"compiled Neo CuTe {frame_name} run {run_index} {name} "
+                    "differs from eager PyTorch"
+                ),
+            )
+
+
+@pytest.mark.skipif(
+    _CUTE_SKIP_REASON is not None,
+    reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable",
+)
+def test_sm90_high_degree_force_uses_portable_cute_fallback(
+    tmp_path: Path,
+) -> None:
+    if tuple(torch.cuda.get_device_capability()) != (9, 0):
+        pytest.skip("the bounded native SO2 runner exists only on SM90")
+    from deepmd.pt_expt.kernels.cute.sezm.so2.sm90.phase_c_attention_backward import (
+        MAX_EDGES_PER_NODE,
+    )
+
+    eager = _run_child_process(
+        frame_name="high_degree",
+        use_compile=False,
+        output_path=tmp_path / "high-degree-eager.pt",
+    )
+    compiled = _run_child_process(
+        frame_name="high_degree",
+        use_compile=True,
+        output_path=tmp_path / "high-degree-compiled.pt",
+    )
+
+    assert compiled["runner_calls"] > 0
+    assert compiled["max_destination_degree"] > MAX_EDGES_PER_NODE
+    assert any("NeoFullCuteBackward" in name for name in compiled["runner_types"])
+    assert all("NeoSm90SO2Runner" not in name for name in compiled["runner_types"])
+    expected = eager["outputs"][0]
+    for actual in compiled["outputs"]:
+        for name in OUTPUT_KEYS:
+            torch.testing.assert_close(
+                actual[name],
+                expected[name],
+                atol=TOL,
+                rtol=TOL,
+                msg=f"high-degree SM90 portable CuTe {name} differs from eager",
+            )
+
+
+def _main() -> None:
+    if len(sys.argv) != 5 or sys.argv[1] != _CHILD_MODE:
+        raise SystemExit("this test module is only executable in E/F/S child mode")
+    _run_efs_child(
+        frame_name=sys.argv[2],
+        use_compile=sys.argv[3] == "compiled",
+        output_path=Path(sys.argv[4]),
+    )
+
+
+if __name__ == "__main__":
+    _main()
diff --git a/source/tests/pt/model/sezm_cute/test_envelope_softmax.py b/source/tests/pt/model/sezm_cute/test_envelope_softmax.py
new file mode 100644
index 0000000000..e6c27b5d05
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_envelope_softmax.py
@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Numerical acceptance for the CuTe envelope-gated softmax."""
+
+from __future__ import (
+    annotations,
+)
+
+import importlib
+
+import pytest
+import torch
+
+from deepmd.pt.model.descriptor.sezm_nn.attention import (
+    segment_envelope_gated_softmax,
+)
+
+
+def _runtime_skip_reason() -> str | None:
+    if not torch.cuda.is_available():
+        return "CuTe envelope-softmax acceptance requires CUDA"
+    try:
+        importlib.import_module("cutlass.cute")
+        importlib.import_module("cuda.bindings.driver")
+    except Exception as exc:  # pragma: no cover - runtime dependent
+        return f"CuTe DSL runtime unavailable: {exc}"
+    return None
+
+
+_SKIP_REASON = _runtime_skip_reason()
+
+
+@pytest.mark.skipif(_SKIP_REASON is not None, reason=_SKIP_REASON or "unavailable")
+@pytest.mark.parametrize(
+    ("logit_scale", "edge_gate"),
+    [
+        (1.0, 1.0),
+        (15.0, 1.0e-2),
+        (110.0, 1.0e-23),
+        (-120.0, 1.0),
+    ],
+)
+def test_cute_envelope_softmax_matches_eager_across_extreme_frames(
+    logit_scale: float,
+    edge_gate: float,
+) -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.envelope_softmax import (
+        compile_envelope_softmax_forward,
+    )
+
+    logits = torch.tensor(
+        [
+            [logit_scale, logit_scale - 0.5],
+            [logit_scale - 1.0, logit_scale - 1.5],
+            [logit_scale - 0.25, logit_scale - 0.75],
+            [logit_scale - 2.0, logit_scale - 2.5],
+        ],
+        device="cuda",
+        dtype=torch.float32,
+    )
+    gate = torch.full((4,), edge_gate, device="cuda", dtype=torch.float32)
+    dst = torch.tensor([0, 0, 1, 1], device="cuda", dtype=torch.long)
+    dst_ptr = torch.tensor([0, 2, 4], device="cuda", dtype=torch.int32)
+    z_bias_raw = torch.tensor([0.1, -0.3], device="cuda", dtype=torch.float32)
+    eps = 1.0e-7
+
+    expected = segment_envelope_gated_softmax(
+        logits.view(4, 2, 1),
+        gate,
+        dst,
+        2,
+        z_bias_raw.view(2, 1),
+        eps,
+    ).view(4, 2)
+    actual = torch.empty_like(logits)
+    group_max = torch.empty((2, 2), device="cuda", dtype=torch.float32)
+    denom = torch.empty_like(group_max)
+    run = compile_envelope_softmax_forward(128, eps)
+    run(logits, gate, dst_ptr, z_bias_raw, actual, group_max, denom)
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(actual, expected, atol=5.0e-5, rtol=5.0e-5)
+    assert torch.isfinite(actual).all()
diff --git a/source/tests/pt/model/sezm_cute/test_gie.py b/source/tests/pt/model/sezm_cute/test_gie.py
new file mode 100644
index 0000000000..0dce88913a
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_gie.py
@@ -0,0 +1,330 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Differential tests for the opt-in CuTe geometric initial embedding."""
+
+from __future__ import (
+    annotations,
+)
+
+import unittest
+from types import (
+    SimpleNamespace,
+)
+
+import torch
+
+from deepmd.pt_expt.kernels.cute.sezm import gie as gie_module
+from deepmd.pt_expt.kernels.cute.sezm.so2.metadata import (
+    build_destination_row_ptr,
+)
+
+
+def _load_gie_module():
+    return gie_module
+
+
+def _degree_slots(lmax: int, *, device: torch.device) -> torch.Tensor:
+    degrees = torch.arange(1, lmax + 1, device=device, dtype=torch.long)
+    return torch.repeat_interleave(degrees - 1, 2 * degrees + 1)
+
+
+def _zonal_indices(
+    lmax: int, *, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    rows = torch.arange(1, (lmax + 1) ** 2, device=device, dtype=torch.long)
+    degrees = torch.arange(1, lmax + 1, device=device, dtype=torch.long)
+    degree_for_row = torch.repeat_interleave(degrees, 2 * degrees + 1)
+    return rows, degree_for_row * (degree_for_row + 1)
+
+
+def _materialized_reference(
+    radial: torch.Tensor,
+    zonal: torch.Tensor,
+    inv_sqrt_deg: torch.Tensor,
+    dst: torch.Tensor,
+    gate: torch.Tensor,
+    *,
+    n_nodes: int,
+    lmax: int,
+) -> torch.Tensor:
+    slots = _degree_slots(lmax, device=radial.device)
+    message = zonal.unsqueeze(-1) * radial.index_select(1, slots)
+    if gate.numel() != 0:
+        message = message * gate.reshape(-1, 1, 1)
+    non_scalar = radial.new_zeros(n_nodes, zonal.shape[1], radial.shape[2])
+    non_scalar.index_add_(0, dst, message)
+    out = radial.new_zeros(n_nodes, zonal.shape[1] + 1, radial.shape[2])
+    out[:, 1:, :] = non_scalar
+    return out * inv_sqrt_deg
+
+
+def _inputs(
+    *,
+    n_nodes: int,
+    dst_values: tuple[int, ...],
+    lmax: int,
+    channels: int,
+    device: torch.device,
+    with_gate: bool,
+) -> tuple[torch.Tensor, ...]:
+    edge_count = len(dst_values)
+    generator = torch.Generator(device=device).manual_seed(20260703 + edge_count)
+    radial = torch.randn(
+        edge_count,
+        lmax,
+        channels,
+        generator=generator,
+        device=device,
+        dtype=torch.float32,
+        requires_grad=True,
+    )
+    dense_dt = torch.randn(
+        edge_count,
+        (lmax + 1) ** 2,
+        (lmax + 1) ** 2,
+        generator=generator,
+        device=device,
+        dtype=torch.float32,
+        requires_grad=True,
+    )
+    rows, cols = _zonal_indices(lmax, device=device)
+    zonal = dense_dt[:, rows, cols]
+    inv_sqrt_deg = (
+        torch.rand(
+            n_nodes,
+            1,
+            1,
+            generator=generator,
+            device=device,
+            dtype=torch.float32,
+        )
+        + 0.25
+    ).requires_grad_(True)
+    dst = torch.tensor(dst_values, device=device, dtype=torch.long)
+    if with_gate:
+        gate = torch.rand(
+            edge_count,
+            1,
+            generator=generator,
+            device=device,
+            dtype=torch.float32,
+            requires_grad=True,
+        )
+    else:
+        gate = torch.empty(0, device=device, dtype=torch.float32)
+    return radial, dense_dt, zonal, inv_sqrt_deg, dst, gate
+
+
+class TestSeZMCuTeGIEContract(unittest.TestCase):
+    def test_backward_compile_key_separates_destination_dtypes(self):
+        gie = _load_gie_module()
+        common = ((0, 8, 0), 3, 32, True, (128, 32, 1))
+
+        int32_key = gie._backward_compile_key(*common, torch.int32)
+        int64_key = gie._backward_compile_key(*common, torch.int64)
+
+        self.assertNotEqual(int32_key, int64_key)
+
+    def test_contract_requires_sorted_dynamic_strict_fp32_inputs(self):
+        gie = _load_gie_module()
+        radial, _dense_dt, zonal, inv_sqrt_deg, dst, gate = _inputs(
+            n_nodes=5,
+            dst_values=(0, 0, 1, 3, 3, 4),
+            lmax=3,
+            channels=4,
+            device=torch.device("cpu"),
+            with_gate=True,
+        )
+        module = SimpleNamespace(
+            lmax=3,
+            channels=4,
+            training=False,
+            non_scalar_row_index=torch.arange(1, 16, device=torch.device("cpu")),
+            radial_slot_index_for_row=_degree_slots(3, device=torch.device("cpu")),
+        )
+        cache = SimpleNamespace(
+            dst=dst,
+            inv_sqrt_deg=inv_sqrt_deg,
+            edge_src_gate=gate,
+            destinations_sorted=True,
+        )
+        self.assertTrue(gie.validate_gie_contract(module, 5, cache, radial, zonal))
+
+        cache.destinations_sorted = False
+        self.assertFalse(gie.validate_gie_contract(module, 5, cache, radial, zonal))
+        cache.destinations_sorted = True
+        self.assertFalse(
+            gie.validate_gie_contract(module, 5, cache, radial.double(), zonal)
+        )
+        module.non_scalar_row_index = torch.arange(15, device=torch.device("cpu"))
+        self.assertFalse(gie.validate_gie_contract(module, 5, cache, radial, zonal))
+
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
+    def test_cuda_forward_and_wigner_radial_degree_gate_gradients(self):
+        gie = _load_gie_module()
+        if not gie.SEZM_CUTE_GIE_AVAILABLE:
+            self.skipTest("CuTe DSL is not available")
+        for n_nodes, dst_values, with_gate in (
+            (4, (0, 0, 2, 2, 2, 3), False),
+            (7, (0, 1, 1, 1, 4, 6, 6, 6, 6), True),
+        ):
+            with self.subTest(n_nodes=n_nodes, edges=len(dst_values), gate=with_gate):
+                expected_inputs = _inputs(
+                    n_nodes=n_nodes,
+                    dst_values=dst_values,
+                    lmax=3,
+                    channels=32,
+                    device=torch.device("cuda"),
+                    with_gate=with_gate,
+                )
+                radial, dense_dt, zonal, inv_sqrt_deg, dst, gate = expected_inputs
+                weight = torch.randn_like(
+                    radial.new_empty(n_nodes, 16, radial.shape[2])
+                )
+                expected = _materialized_reference(
+                    radial,
+                    zonal,
+                    inv_sqrt_deg,
+                    dst,
+                    gate,
+                    n_nodes=n_nodes,
+                    lmax=3,
+                )
+                expected_grads = torch.autograd.grad(
+                    (expected * weight).sum(),
+                    (radial, dense_dt, inv_sqrt_deg, *([gate] if with_gate else [])),
+                )
+
+                actual_inputs = _inputs(
+                    n_nodes=n_nodes,
+                    dst_values=dst_values,
+                    lmax=3,
+                    channels=32,
+                    device=torch.device("cuda"),
+                    with_gate=with_gate,
+                )
+                radial, dense_dt, zonal, inv_sqrt_deg, dst, gate = actual_inputs
+                actual = gie.gie_fused_cuda(
+                    radial,
+                    zonal,
+                    inv_sqrt_deg,
+                    dst,
+                    gate,
+                    n_nodes=n_nodes,
+                    lmax=3,
+                )
+                actual_grads = torch.autograd.grad(
+                    (actual * weight).sum(),
+                    (radial, dense_dt, inv_sqrt_deg, *([gate] if with_gate else [])),
+                )
+                torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+                for actual_grad, expected_grad in zip(
+                    actual_grads, expected_grads, strict=True
+                ):
+                    torch.testing.assert_close(
+                        actual_grad, expected_grad, rtol=2e-5, atol=2e-5
+                    )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
+    def test_inductor_gie_handles_offset_destination_view(self):
+        gie = _load_gie_module()
+        if not gie.SEZM_CUTE_GIE_AVAILABLE:
+            self.skipTest("CuTe DSL is not available")
+
+        def run_gie(edge_index, radial, zonal, inv, gate):
+            dst = edge_index[1]
+            n_nodes = inv.shape[0]
+            out = gie.gie_fused_cuda(
+                radial,
+                zonal,
+                inv,
+                dst,
+                gate,
+                n_nodes=n_nodes,
+                lmax=3,
+            )
+            return out, build_destination_row_ptr(dst, n_nodes)
+
+        compiled = torch.compile(
+            run_gie, backend="inductor", dynamic=True, fullgraph=True
+        )
+        for with_gate in (False, True):
+            for n_nodes, src_values, dst_values, row_ptr in (
+                (4, (2, 0, 3, 0, 1, 2), (0, 0, 1, 2, 2, 3), (0, 2, 3, 5, 6)),
+                (
+                    7,
+                    (6, 0, 5, 2, 0, 4, 1, 3, 2),
+                    (0, 1, 1, 1, 4, 6, 6, 6, 6),
+                    (0, 1, 4, 4, 4, 5, 5, 9),
+                ),
+            ):
+                with self.subTest(
+                    gate=with_gate, n_nodes=n_nodes, edges=len(dst_values)
+                ):
+                    expected_inputs = _inputs(
+                        n_nodes=n_nodes,
+                        dst_values=dst_values,
+                        lmax=3,
+                        channels=32,
+                        device=torch.device("cuda"),
+                        with_gate=with_gate,
+                    )
+                    radial, dense_dt, zonal, inv, dst, gate = expected_inputs
+                    weight = torch.randn_like(radial.new_empty(n_nodes, 16, 32))
+                    expected = _materialized_reference(
+                        radial,
+                        zonal,
+                        inv,
+                        dst,
+                        gate,
+                        n_nodes=n_nodes,
+                        lmax=3,
+                    )
+                    expected_grads = torch.autograd.grad(
+                        (expected * weight).sum(),
+                        (radial, dense_dt, zonal, inv, *([gate] if with_gate else [])),
+                    )
+
+                    actual_inputs = _inputs(
+                        n_nodes=n_nodes,
+                        dst_values=dst_values,
+                        lmax=3,
+                        channels=32,
+                        device=torch.device("cuda"),
+                        with_gate=with_gate,
+                    )
+                    radial, dense_dt, zonal, inv, _dst, gate = actual_inputs
+                    # The first row is deliberately not sorted: losing the second
+                    # row's offset makes searchsorted operate on unrelated sources.
+                    edge_index = torch.tensor(
+                        (src_values, dst_values), dtype=torch.int64, device="cuda"
+                    )
+                    self.assertTrue(edge_index[1].is_contiguous())
+                    self.assertEqual(edge_index[1].storage_offset(), len(dst_values))
+                    self.assertFalse(edge_index.requires_grad)
+                    actual, dst_ptr = compiled(edge_index, radial, zonal, inv, gate)
+                    self.assertEqual(dst_ptr.dtype, torch.int32)
+                    self.assertTrue(dst_ptr.is_contiguous())
+                    self.assertFalse(dst_ptr.requires_grad)
+                    self.assertIsNone(dst_ptr.grad_fn)
+                    torch.testing.assert_close(
+                        dst_ptr,
+                        torch.tensor(row_ptr, dtype=torch.int32, device="cuda"),
+                        rtol=0,
+                        atol=0,
+                    )
+                    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+                    actual_grads = torch.autograd.grad(
+                        (actual * weight).sum(),
+                        (radial, dense_dt, zonal, inv, *([gate] if with_gate else [])),
+                    )
+                    for actual_grad, expected_grad in zip(
+                        actual_grads, expected_grads, strict=True
+                    ):
+                        torch.testing.assert_close(
+                            actual_grad, expected_grad, rtol=2e-5, atol=2e-5
+                        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/model/sezm_cute/test_multigpu_dispatch.py b/source/tests/pt/model/sezm_cute/test_multigpu_dispatch.py
new file mode 100644
index 0000000000..f8863d777d
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_multigpu_dispatch.py
@@ -0,0 +1,279 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""CPU/mock coverage for operand-device-aware Neo CuTe dispatch."""
+
+from __future__ import (
+    annotations,
+)
+
+import importlib
+import os
+from types import (
+    SimpleNamespace,
+)
+from unittest import (
+    mock,
+)
+
+import pytest
+import torch
+
+from deepmd.pt.model.network import (
+    mlp,
+)
+from deepmd.pt_expt.kernels.cute.sezm import (
+    runtime_policy,
+)
+from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2
+
+
+def _cuda_operand(index: int = 1) -> SimpleNamespace:
+    return SimpleNamespace(
+        device=torch.device("cuda", index),
+        is_cuda=True,
+    )
+
+
+def test_so2_custom_ops_are_direct_and_thin_path_uses_operand_device() -> None:
+    operand = _cuda_operand()
+    sentinel = object()
+
+    with mock.patch.object(so2, "_cute_so2_impl", return_value=sentinel) as direct:
+        actual = so2.cute_so2(
+            1,
+            operand,
+            object(),
+            object(),
+            object(),
+            object(),
+            object(),
+            object(),
+            object(),
+            object(),
+            object(),
+            object(),
+        )
+
+    assert actual is sentinel
+    direct.assert_called_once()
+
+    context = object()
+    with mock.patch.object(
+        so2,
+        "_so2_packed_direct_registered_backward_impl",
+        return_value=sentinel,
+    ) as backward:
+        actual = so2._so2_packed_direct_backward(context, operand, None)
+
+    assert actual is sentinel
+    backward.assert_called_once_with(context, operand)
+
+    with (
+        mock.patch.object(
+            so2,
+            "_cuda_compute_capability",
+            return_value=(9, 0),
+        ) as capability,
+        mock.patch.object(
+            runtime_policy,
+            "is_so2_thin_wrapper_enabled",
+            return_value=False,
+        ) as thin_selector,
+        mock.patch.object(
+            so2,
+            "_maybe_run_cute_so2_fallback",
+            return_value=sentinel,
+        ) as fallback,
+    ):
+        actual = so2.maybe_run_cute_so2(object(), operand, object(), object())
+
+    assert actual is sentinel
+    capability.assert_called_once_with(1)
+    thin_selector.assert_called_once_with((9, 0))
+    fallback.assert_called_once()
+
+
+def test_so2_prepare_uses_requested_device_instead_of_current_device() -> None:
+    device = torch.device("cuda", 1)
+    with (
+        mock.patch.object(
+            so2,
+            "_cuda_compute_capability",
+            return_value=(8, 6),
+        ) as capability,
+        mock.patch.object(
+            so2,
+            "is_supported_so2_compute_capability",
+            return_value=False,
+        ) as supported,
+        mock.patch.object(
+            torch.cuda,
+            "current_device",
+            side_effect=AssertionError("must not read the current CUDA device"),
+        ),
+    ):
+        assert not so2.prepare_cute_so2_blocks(
+            [object()],
+            training=False,
+            device=device,
+            dtype=torch.float32,
+        )
+
+    capability.assert_called_once_with(1)
+    supported.assert_called_once_with((8, 6))
+
+
+def test_wignerd_custom_ops_are_called_directly() -> None:
+    pytest.importorskip("cutlass.cute")
+    wignerd = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.wignerd")
+    operand = _cuda_operand()
+    sentinel = object()
+
+    with mock.patch.object(
+        wignerd,
+        "_run_cute_wignerd_impl",
+        return_value=sentinel,
+    ) as forward:
+        actual = wignerd.run_cute_wignerd(operand, object())
+
+    assert actual is sentinel
+    forward.assert_called_once_with(operand, mock.ANY, packed_wigner=False)
+
+    context = SimpleNamespace(saved_tensors=(operand,))
+    grad_panel = object()
+    with mock.patch.object(
+        wignerd,
+        "_wignerd_panel_registered_backward_impl",
+        return_value=sentinel,
+    ) as backward:
+        actual = wignerd._wignerd_panel_backward(context, grad_panel)
+
+    assert actual is sentinel
+    backward.assert_called_once_with(context, grad_panel)
+
+
+def test_mlp_selector_and_forward_use_input_device() -> None:
+    cuda1 = torch.device("cuda", 1)
+
+    def capability(device: torch.device | None = None) -> tuple[int, int]:
+        return (8, 0) if device == cuda1 else (9, 0)
+
+    with (
+        mock.patch.dict(
+            os.environ,
+            {"DP_CUTE_INFER": "1"},
+            clear=True,
+        ),
+        mock.patch.object(torch.cuda, "is_available", return_value=True),
+        mock.patch.object(
+            torch.cuda,
+            "get_device_capability",
+            side_effect=capability,
+        ) as get_capability,
+    ):
+        assert mlp._use_so2_compile_visible_linear(cuda1)
+        assert not mlp._use_so2_compile_visible_linear(torch.device("cuda", 0))
+        assert not mlp._use_so2_compile_visible_linear(torch.device("cpu"))
+
+    assert get_capability.call_args_list == [
+        mock.call(cuda1),
+        mock.call(torch.device("cuda", 0)),
+    ]
+
+    layer = mlp.MLPLayer(
+        4,
+        4,
+        activation_function="none",
+        precision="float32",
+    ).to("cpu")
+    layer.eval()
+    mlp.enable_neo_cute_compile_visible_linears(layer)
+    value = torch.randn(3, 4, device="cpu")
+    with mock.patch.object(
+        mlp,
+        "_use_so2_compile_visible_linear",
+        return_value=False,
+    ) as selector:
+        layer(value)
+    selector.assert_called_once_with(value.device)
+
+
+def test_nv_neighbor_list_selector_and_build_use_coordinate_device() -> None:
+    sezm_model = importlib.import_module("deepmd.pt.model.model.sezm_model")
+    cuda1 = torch.device("cuda", 1)
+
+    def capability(device: torch.device | None = None) -> tuple[int, int]:
+        return (8, 0) if device == cuda1 else (9, 0)
+
+    with (
+        mock.patch.dict(
+            os.environ,
+            {"DP_CUTE_INFER": "1"},
+            clear=True,
+        ),
+        mock.patch.object(torch.cuda, "is_available", return_value=True),
+        mock.patch.object(
+            torch.cuda,
+            "get_device_capability",
+            side_effect=capability,
+        ) as get_capability,
+    ):
+        assert sezm_model._neo_cute_nlist_eager_island_enabled(cuda1)
+        assert not sezm_model._neo_cute_nlist_eager_island_enabled(
+            torch.device("cuda", 0)
+        )
+        assert not sezm_model._neo_cute_nlist_eager_island_enabled(torch.device("cpu"))
+
+    assert get_capability.call_args_list == [
+        mock.call(cuda1),
+        mock.call(torch.device("cuda", 0)),
+    ]
+
+    builder = sezm_model.NvNeighborList()
+    sentinel = object()
+    with (
+        mock.patch.object(
+            sezm_model,
+            "_neo_cute_nlist_eager_island_enabled",
+            return_value=False,
+        ) as selector,
+        mock.patch.object(builder, "build", return_value=sentinel) as build,
+    ):
+        actual = sezm_model._build_neo_neighbor_list(
+            builder,
+            mock.Mock(device=cuda1),
+            object(),
+            None,
+            4.0,
+            [32],
+            return_mode="edges",
+        )
+
+    assert actual is sentinel
+    selector.assert_called_once_with(cuda1)
+    build.assert_called_once()
+
+    with (
+        mock.patch.object(
+            sezm_model,
+            "_neo_cute_nlist_eager_island_enabled",
+            return_value=True,
+        ),
+        mock.patch.object(
+            sezm_model,
+            "_build_neo_neighbor_list_eager_island",
+            return_value=sentinel,
+        ) as eager,
+    ):
+        actual = sezm_model._build_neo_neighbor_list(
+            builder,
+            mock.Mock(device=cuda1),
+            object(),
+            None,
+            4.0,
+            [32],
+            return_mode="edges",
+        )
+
+    assert actual is sentinel
+    eager.assert_called_once()
diff --git a/source/tests/pt/model/sezm_cute/test_output_grid_product.py b/source/tests/pt/model/sezm_cute/test_output_grid_product.py
new file mode 100644
index 0000000000..f3531794d3
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_output_grid_product.py
@@ -0,0 +1,539 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Behavioral differentials for the fused Neo output-grid product."""
+
+from __future__ import (
+    annotations,
+)
+
+import importlib
+import sys
+import types
+
+import pytest
+import torch
+
+TOL = 5.0e-5
+N_FRAMES = 3
+COEFF_DIM = 16
+GRID_SIZE = 152
+SUPPORTED_HIDDEN_CHANNELS = (96, 192)
+
+
+def _cute_runtime_skip_reason() -> str | None:
+    if not torch.cuda.is_available():
+        return "output-grid differentials require CUDA"
+    if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6), (9, 0)}:
+        return "output-grid CuTe dispatch supports only sm80, sm86, and sm90"
+    try:
+        importlib.import_module("cutlass.cute")
+        importlib.import_module("cuda.bindings.driver")
+    except Exception as exc:  # pragma: no cover - runtime dependent
+        return f"output-grid differentials require the CuTe DSL runtime: {exc}"
+    return None
+
+
+_CUTE_SKIP_REASON = _cute_runtime_skip_reason()
+
+
+def _sm80_skip_reason() -> str | None:
+    if _CUTE_SKIP_REASON is not None:
+        return _CUTE_SKIP_REASON
+    if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6)}:
+        return "specialized output-grid differentials require sm80 or sm86"
+    return None
+
+
+_SM80_SKIP_REASON = _sm80_skip_reason()
+
+
+def _reference(left, right, to_grid, from_grid):
+    nodes = left.shape[0]
+    hidden_channels = left.shape[-1] // N_FRAMES
+    left_flat = left.reshape(nodes, COEFF_DIM * N_FRAMES, hidden_channels)
+    right_flat = right.reshape(nodes, COEFF_DIM * N_FRAMES, hidden_channels)
+    left_grid = torch.einsum("gj,njc->ngc", to_grid, left_flat)
+    right_grid = torch.einsum("gj,njc->ngc", to_grid, right_flat)
+    out = torch.einsum("jg,ngc->njc", from_grid, left_grid * right_grid)
+    return out.reshape_as(left)
+
+
+def test_grid_mlp_accepts_fused_pair_callback():
+    from deepmd.pt.model.descriptor.sezm_nn.grid_net import (
+        GridMLP,
+    )
+
+    module = GridMLP(
+        channels=2,
+        mode="self",
+        n_frames=3,
+        dtype=torch.float32,
+        trainable=False,
+        seed=7,
+    ).to("cpu")
+    left = torch.randn(2, 4, 1, 6, device="cpu")
+    right = torch.randn_like(left)
+    scalar_pair = torch.empty(2, 1, 4, device="cpu")
+    calls = 0
+
+    def fused_pair(projected_left, projected_right):
+        nonlocal calls
+        calls += 1
+        return projected_left * projected_right
+
+    out = module(
+        left,
+        right,
+        scalar_pair,
+        to_grid=lambda value: value,
+        from_grid=lambda value: value,
+        pair_grid=fused_pair,
+    )
+
+    assert calls == 1
+    assert out.shape == left.shape
+
+
+@pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+def test_dispatch_falls_back_without_exact_cuda_contract(
+    monkeypatch: pytest.MonkeyPatch,
+    hidden_channels: int,
+):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        product as output_grid_product,
+    )
+
+    monkeypatch.setattr(
+        output_grid_product.runtime_policy,
+        "is_cute_infer_enabled",
+        lambda: True,
+    )
+    left = torch.randn(2, COEFF_DIM, 1, N_FRAMES * hidden_channels, device="cpu")
+    right = torch.randn_like(left)
+    to_grid = torch.randn(GRID_SIZE, COEFF_DIM * N_FRAMES, device="cpu")
+    from_grid = torch.randn(COEFF_DIM * N_FRAMES, GRID_SIZE, device="cpu")
+
+    assert (
+        output_grid_product.maybe_run_cute_output_grid_product(
+            left,
+            right,
+            to_grid,
+            from_grid,
+            n_frames=N_FRAMES,
+        )
+        is None
+    )
+
+
+@pytest.mark.parametrize(
+    ("hidden_channels", "expected"),
+    [(96, 96), (192, 192), (128, None)],
+)
+def test_exact_shape_guard_accepts_only_validated_widths(
+    hidden_channels: int,
+    expected: int | None,
+) -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid.product import (
+        _exact_hidden_channels,
+    )
+
+    left = torch.empty(2, COEFF_DIM, 1, N_FRAMES * hidden_channels, device="cpu")
+    assert _exact_hidden_channels(left, N_FRAMES) == expected
+
+
+@pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+def test_fake_registrations_allocate_canonical_strides(hidden_channels: int) -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        product as output_grid_product,
+    )
+
+    width = N_FRAMES * hidden_channels
+    shape = (2, COEFF_DIM, 1, width)
+    canonical_stride = (COEFF_DIM * width, width, width, 1)
+    left = torch.empty_strided(
+        shape,
+        (canonical_stride[0], canonical_stride[1], 7, 1),
+        device="cpu",
+    )
+    right = torch.empty_strided(
+        shape,
+        (canonical_stride[0], canonical_stride[1], 11, 1),
+        device="cpu",
+    )
+    grad_out = torch.empty_strided(
+        shape,
+        (canonical_stride[0], canonical_stride[1], 13, 1),
+        device="cpu",
+    )
+    to_grid = torch.empty(GRID_SIZE, COEFF_DIM * N_FRAMES, device="cpu")
+    from_grid = torch.empty(COEFF_DIM * N_FRAMES, GRID_SIZE, device="cpu")
+
+    out = output_grid_product._output_grid_product_fake(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        N_FRAMES,
+    )
+    grad_left, grad_right = output_grid_product._output_grid_product_bwd_fake(
+        grad_out,
+        left,
+        right,
+        to_grid,
+        from_grid,
+        N_FRAMES,
+    )
+
+    assert out.stride() == canonical_stride
+    assert grad_left.stride() == canonical_stride
+    assert grad_right.stride() == canonical_stride
+
+
+@pytest.mark.parametrize(
+    ("compute_capability", "hidden_channels", "policy_enabled", "expected"),
+    [
+        ((9, 0), 96, True, True),
+        ((9, 0), 96, False, False),
+        ((9, 0), 192, True, False),
+        ((9, 1), 96, True, False),
+    ],
+)
+def test_sm90_c96_asymmetric_panel_dispatch_is_explicit(
+    monkeypatch: pytest.MonkeyPatch,
+    compute_capability: tuple[int, int],
+    hidden_channels: int,
+    policy_enabled: bool,
+    expected: bool,
+) -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        product as output_grid_product,
+    )
+
+    kernel_module_name = (
+        "deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product"
+    )
+    kernel_module = types.ModuleType(kernel_module_name)
+    calls: dict[str, dict[str, bool]] = {}
+
+    def run_forward(left, right, to_grid, from_grid, **kwargs):
+        del right, to_grid, from_grid
+        calls["forward"] = kwargs
+        return torch.empty_like(left)
+
+    def run_backward(grad_out, left, right, to_grid, from_grid, **kwargs):
+        del grad_out, to_grid, from_grid
+        calls["backward"] = kwargs
+        return torch.empty_like(left), torch.empty_like(right)
+
+    kernel_module.run_tiled_output_grid_product = run_forward
+    kernel_module.run_tiled_output_grid_product_backward = run_backward
+    monkeypatch.setitem(sys.modules, kernel_module_name, kernel_module)
+    monkeypatch.setattr(
+        output_grid_product,
+        "_validate_exact_contract",
+        lambda *_args: hidden_channels,
+    )
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_capability",
+        lambda _device=None: compute_capability,
+    )
+    monkeypatch.setattr(
+        output_grid_product.runtime_policy,
+        "is_output_grid_fwd_sm80_c96_n48_enabled",
+        lambda _compute_capability: False,
+    )
+    monkeypatch.setattr(
+        output_grid_product.runtime_policy,
+        "is_output_grid_bwd_sm80_c96_n48_panel_enabled",
+        lambda _compute_capability: False,
+    )
+    monkeypatch.setattr(
+        output_grid_product.runtime_policy,
+        "is_output_grid_sm90_c96_asymmetric_panels_enabled",
+        lambda _compute_capability: policy_enabled,
+    )
+
+    width = N_FRAMES * hidden_channels
+    left = torch.empty(2, COEFF_DIM, 1, width, device="cpu")
+    right = torch.empty_like(left)
+    grad_out = torch.empty_like(left)
+    to_grid = torch.empty(
+        GRID_SIZE,
+        COEFF_DIM * N_FRAMES,
+        device="cpu",
+    )
+    from_grid = torch.empty(
+        COEFF_DIM * N_FRAMES,
+        GRID_SIZE,
+        device="cpu",
+    )
+
+    output_grid_product._output_grid_product_impl(
+        left,
+        right,
+        to_grid,
+        from_grid,
+        N_FRAMES,
+    )
+    output_grid_product._output_grid_product_bwd_impl(
+        grad_out,
+        left,
+        right,
+        to_grid,
+        from_grid,
+        N_FRAMES,
+    )
+
+    assert calls["forward"] == {
+        "use_sm80_c96_n48": False,
+        "use_sm90_c96_asymmetric_panels": expected,
+    }
+    assert calls["backward"] == {
+        "use_sm80_c96_n48_panel": False,
+        "use_sm90_c96_asymmetric_panels": expected,
+    }
+
+
+@pytest.mark.skipif(
+    _CUTE_SKIP_REASON is not None,
+    reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable",
+)
+class TestOutputGridProductCuda:
+    @staticmethod
+    def _inputs(nodes: int, hidden_channels: int):
+        generator = torch.Generator(device="cuda").manual_seed(
+            20260703 + nodes + hidden_channels
+        )
+        left = 0.1 * torch.randn(
+            nodes,
+            COEFF_DIM,
+            1,
+            N_FRAMES * hidden_channels,
+            device="cuda",
+            generator=generator,
+        )
+        right = 0.1 * torch.randn(
+            left.shape,
+            device="cuda",
+            generator=generator,
+        )
+        to_grid = 0.1 * torch.randn(
+            GRID_SIZE,
+            COEFF_DIM * N_FRAMES,
+            device="cuda",
+            generator=generator,
+        )
+        from_grid = 0.1 * torch.randn(
+            COEFF_DIM * N_FRAMES,
+            GRID_SIZE,
+            device="cuda",
+            generator=generator,
+        )
+        return left, right, to_grid, from_grid
+
+    @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+    @pytest.mark.parametrize("nodes", [1, 7, 65])
+    def test_forward_and_first_backward_match_strict_fp32(
+        self,
+        nodes: int,
+        hidden_channels: int,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.product import (
+            output_grid_product_cute,
+        )
+
+        left, right, to_grid, from_grid = self._inputs(nodes, hidden_channels)
+        left_ref = left.detach().clone().requires_grad_(True)
+        right_ref = right.detach().clone().requires_grad_(True)
+        left_actual = left.detach().clone().requires_grad_(True)
+        right_actual = right.detach().clone().requires_grad_(True)
+        grad = torch.randn_like(left)
+
+        expected = _reference(left_ref, right_ref, to_grid, from_grid)
+        expected_grads = torch.autograd.grad(
+            expected,
+            (left_ref, right_ref),
+            grad,
+        )
+        actual = output_grid_product_cute(
+            left_actual,
+            right_actual,
+            to_grid,
+            from_grid,
+            n_frames=N_FRAMES,
+        )
+        actual_grads = torch.autograd.grad(
+            actual,
+            (left_actual, right_actual),
+            grad,
+        )
+
+        torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+        torch.testing.assert_close(
+            actual_grads[0], expected_grads[0], atol=TOL, rtol=TOL
+        )
+        torch.testing.assert_close(
+            actual_grads[1], expected_grads[1], atol=TOL, rtol=TOL
+        )
+
+    @pytest.mark.skipif(
+        _SM80_SKIP_REASON is not None,
+        reason=_SM80_SKIP_REASON or "sm80-family GPU is unavailable",
+    )
+    @pytest.mark.parametrize("nodes", [1, 7, 65])
+    def test_sm80_c96_n48_forward_custom_op_matches_strict_fp32(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        nodes: int,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+            product as output_grid_product,
+        )
+
+        monkeypatch.setattr(
+            output_grid_product.runtime_policy,
+            "is_output_grid_fwd_sm80_c96_n48_enabled",
+            lambda _compute_capability: True,
+        )
+        left, right, to_grid, from_grid = self._inputs(nodes, 96)
+        expected = _reference(left, right, to_grid, from_grid)
+        actual = output_grid_product.output_grid_product_cute(
+            left,
+            right,
+            to_grid,
+            from_grid,
+            n_frames=N_FRAMES,
+        )
+        torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+
+    @pytest.mark.skipif(
+        _SM80_SKIP_REASON is not None,
+        reason=_SM80_SKIP_REASON or "sm80-family GPU is unavailable",
+    )
+    def test_sm80_c96_n48_panel_custom_op_matches_strict_fp32(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+            product as output_grid_product,
+        )
+
+        monkeypatch.setattr(
+            output_grid_product.runtime_policy,
+            "is_output_grid_bwd_sm80_c96_n48_panel_enabled",
+            lambda _compute_capability: True,
+        )
+        left, right, to_grid, from_grid = self._inputs(7, 96)
+        left_ref = left.detach().clone().requires_grad_(True)
+        right_ref = right.detach().clone().requires_grad_(True)
+        left_actual = left.detach().clone().requires_grad_(True)
+        right_actual = right.detach().clone().requires_grad_(True)
+        grad = torch.randn_like(left)
+
+        expected = _reference(left_ref, right_ref, to_grid, from_grid)
+        expected_grads = torch.autograd.grad(
+            expected,
+            (left_ref, right_ref),
+            grad,
+        )
+        actual = output_grid_product.output_grid_product_cute(
+            left_actual,
+            right_actual,
+            to_grid,
+            from_grid,
+            n_frames=N_FRAMES,
+        )
+        actual_grads = torch.autograd.grad(
+            actual,
+            (left_actual, right_actual),
+            grad,
+        )
+
+        torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+        torch.testing.assert_close(
+            actual_grads[0], expected_grads[0], atol=TOL, rtol=TOL
+        )
+        torch.testing.assert_close(
+            actual_grads[1], expected_grads[1], atol=TOL, rtol=TOL
+        )
+
+    @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+    def test_dispatch_uses_only_the_master_gate(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        hidden_channels: int,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+            product as output_grid_product,
+        )
+
+        left, right, to_grid, from_grid = self._inputs(2, hidden_channels)
+        monkeypatch.setattr(
+            output_grid_product.runtime_policy,
+            "is_cute_infer_enabled",
+            lambda: False,
+        )
+        assert (
+            output_grid_product.maybe_run_cute_output_grid_product(
+                left,
+                right,
+                to_grid,
+                from_grid,
+                n_frames=N_FRAMES,
+            )
+            is None
+        )
+
+        monkeypatch.setattr(
+            output_grid_product.runtime_policy,
+            "is_cute_infer_enabled",
+            lambda: True,
+        )
+        actual = output_grid_product.maybe_run_cute_output_grid_product(
+            left,
+            right,
+            to_grid,
+            from_grid,
+            n_frames=N_FRAMES,
+        )
+        assert actual is not None
+        torch.testing.assert_close(
+            actual,
+            _reference(left, right, to_grid, from_grid),
+            atol=TOL,
+            rtol=TOL,
+        )
+
+    @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+    def test_dispatch_declines_non_strict_matmul_state(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        hidden_channels: int,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+            product as output_grid_product,
+        )
+
+        left, right, to_grid, from_grid = self._inputs(2, hidden_channels)
+        monkeypatch.setattr(
+            output_grid_product.runtime_policy,
+            "is_cute_infer_enabled",
+            lambda: True,
+        )
+        monkeypatch.setattr(
+            output_grid_product.runtime_policy,
+            "uses_strict_fp32_matmul",
+            lambda: False,
+        )
+
+        assert (
+            output_grid_product.maybe_run_cute_output_grid_product(
+                left,
+                right,
+                to_grid,
+                from_grid,
+                n_frames=N_FRAMES,
+            )
+            is None
+        )
diff --git a/source/tests/pt/model/sezm_cute/test_readout_l0.py b/source/tests/pt/model/sezm_cute/test_readout_l0.py
new file mode 100644
index 0000000000..8adaac0f81
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_readout_l0.py
@@ -0,0 +1,852 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Focused algebra, dispatch, custom-op, and fullgraph readout checks."""
+
+from __future__ import (
+    annotations,
+)
+
+import importlib
+
+import pytest
+import torch
+
+COEFF_DIM = 16
+N_FRAMES = 3
+PACKED_COEFF_DIM = COEFF_DIM * N_FRAMES
+GRID_SIZE = 152
+HIDDEN_CHANNELS = 192
+PACKED_WIDTH = N_FRAMES * HIDDEN_CHANNELS
+TOL = 5.0e-5
+
+
+def _cute_runtime_skip_reason() -> str | None:
+    if not torch.cuda.is_available():
+        return "Neo readout l=0 differentials require CUDA"
+    if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (9, 0)}:
+        return "Neo readout l=0 differentials require sm80 or sm90"
+    try:
+        importlib.import_module("cutlass.cute")
+        importlib.import_module("cuda.bindings.driver")
+    except Exception as exc:  # pragma: no cover - runtime dependent
+        return f"Neo readout l=0 differentials require CuTe DSL: {exc}"
+    return None
+
+
+_CUTE_SKIP_REASON = _cute_runtime_skip_reason()
+
+
+def _cpu_inputs(nodes: int = 2, hidden_channels: int = HIDDEN_CHANNELS):
+    width = N_FRAMES * hidden_channels
+    left = torch.randn(nodes, COEFF_DIM, 1, width, device="cpu")
+    right = torch.randn_like(left)
+    to_grid = torch.randn(GRID_SIZE, PACKED_COEFF_DIM, device="cpu")
+    from_grid = torch.randn(PACKED_COEFF_DIM, GRID_SIZE, device="cpu")
+    return left, right, to_grid, from_grid
+
+
+def _output_ffn(
+    *,
+    hidden_channels: int = 96,
+    trainable: bool = False,
+    device: str = "cpu",
+):
+    from deepmd.pt.model.descriptor.sezm_nn.ffn import (
+        EquivariantFFN,
+    )
+
+    return (
+        EquivariantFFN(
+            lmax=3,
+            channels=32,
+            hidden_channels=hidden_channels,
+            kmax=1,
+            grid_mlp=True,
+            grid_branch=0,
+            dtype=torch.float32,
+            s2_activation=False,
+            ffn_so3_grid=True,
+            activation_function="silu",
+            glu_activation=True,
+            mlp_bias=False,
+            trainable=trainable,
+            seed=29,
+        )
+        .to(device)
+        .eval()
+    )
+
+
+def _neo_descriptor():
+    from deepmd.pt.model.descriptor.sezm import (
+        DescrptSeZM,
+    )
+
+    return DescrptSeZM(
+        ntypes=2,
+        sel=4,
+        channels=32,
+        lmax=3,
+        mmax=1,
+        n_blocks=2,
+        so2_layers=3,
+        n_focus=2,
+        message_node_so3=True,
+        ffn_neurons=0,
+        ffn_so3_grid=True,
+        grid_branch=[0, 0, 1],
+        ffn_blocks=1,
+        so3_readout="mlp",
+        use_amp=False,
+        precision="float32",
+        trainable=False,
+        seed=42,
+    ).eval()
+
+
+@pytest.fixture
+def cpu_neo_descriptor(monkeypatch):
+    from deepmd.pt.model.network import (
+        mlp,
+    )
+    from deepmd.pt.utils import (
+        env,
+    )
+    from deepmd.pt.utils import utils as pt_utils
+
+    cpu = torch.device("cpu")
+    monkeypatch.setattr(env, "DEVICE", cpu)
+    monkeypatch.setattr(pt_utils, "DEVICE", cpu)
+    monkeypatch.setattr(mlp, "device", cpu)
+    return _neo_descriptor().to("cpu")
+
+
+def _capture_matmul_precision_state() -> tuple[str | None, str | None, str]:
+    matmul = torch.backends.cuda.matmul
+    try:
+        global_precision = torch.backends.fp32_precision
+    except AttributeError:
+        return None, None, torch.get_float32_matmul_precision()
+
+    torch.backends.fp32_precision = "none"
+    backend_precision = matmul.fp32_precision
+    matmul.fp32_precision = "none"
+    try:
+        legacy_precision = torch.get_float32_matmul_precision()
+    finally:
+        torch.backends.fp32_precision = global_precision
+        matmul.fp32_precision = backend_precision
+    return backend_precision, global_precision, legacy_precision
+
+
+def _restore_matmul_precision_state(
+    state: tuple[str | None, str | None, str],
+) -> None:
+    backend_precision, global_precision, legacy_precision = state
+    matmul = torch.backends.cuda.matmul
+    if backend_precision is None:
+        torch.set_float32_matmul_precision(legacy_precision)
+        return
+    matmul.fp32_precision = "none"
+    if global_precision is not None:
+        torch.backends.fp32_precision = "none"
+    torch.set_float32_matmul_precision(legacy_precision)
+    if global_precision is not None:
+        torch.backends.fp32_precision = global_precision
+    matmul.fp32_precision = backend_precision
+
+
+@pytest.fixture
+def matmul_precision_state():
+    state = _capture_matmul_precision_state()
+    try:
+        yield torch.backends.cuda.matmul
+    finally:
+        _restore_matmul_precision_state(state)
+
+
+def _set_new_matmul_precision(matmul, precision: str) -> None:
+    try:
+        matmul.fp32_precision = precision
+    except AttributeError:
+        pytest.skip("new CUDA matmul precision API is unavailable")
+
+
+def _reference_product(left, right, to_grid, from_grid):
+    nodes = left.shape[0]
+    left_flat = left.reshape(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS)
+    right_flat = right.reshape(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS)
+    left_grid = torch.einsum("gj,njh->ngh", to_grid, left_flat)
+    right_grid = torch.einsum("gj,njh->ngh", to_grid, right_flat)
+    return torch.einsum(
+        "g,ngh->nh",
+        from_grid[0],
+        left_grid * right_grid,
+    )
+
+
+def test_sm80_readout_input_fold_selector_uses_cute_gate(monkeypatch):
+    from deepmd.pt_expt.kernels.cute.sezm import (
+        runtime_policy,
+    )
+
+    monkeypatch.delenv("DP_CUTE_INFER", raising=False)
+    monkeypatch.delenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", raising=False)
+    assert not runtime_policy.is_cute_infer_enabled()
+    assert not runtime_policy.is_readout_input_fold_sm80_enabled((8, 0))
+
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    assert runtime_policy.is_cute_infer_enabled()
+    assert runtime_policy.is_readout_input_fold_sm80_enabled((8, 0))
+    monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", "0")
+    assert not runtime_policy.is_readout_input_fold_sm80_enabled((8, 0))
+    monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", "1")
+    assert runtime_policy.is_readout_input_fold_sm80_enabled((8, 0))
+    assert not runtime_policy.is_readout_input_fold_sm80_enabled((9, 0))
+
+
+def test_sm90_readout_input_fold_selector_uses_cute_gate(monkeypatch):
+    from deepmd.pt_expt.kernels.cute.sezm import (
+        runtime_policy,
+    )
+
+    monkeypatch.delenv("DP_CUTE_INFER", raising=False)
+    monkeypatch.delenv("DP_CUTE_READOUT_INPUT_FOLD_SM90", raising=False)
+    assert not runtime_policy.is_readout_input_fold_sm90_enabled((9, 0))
+
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    assert runtime_policy.is_readout_input_fold_sm90_enabled((9, 0))
+    assert runtime_policy.is_readout_input_fold_enabled((9, 0))
+    monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM90", "0")
+    assert not runtime_policy.is_readout_input_fold_sm90_enabled((9, 0))
+    monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM90", "1")
+    assert runtime_policy.is_readout_input_fold_sm90_enabled((9, 0))
+    assert not runtime_policy.is_readout_input_fold_sm90_enabled((8, 0))
+
+
+def test_readout_input_fold_guard_accepts_validated_architectures(monkeypatch):
+    from deepmd.pt_expt.kernels.cute.sezm import (
+        runtime_policy,
+    )
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu")
+    capabilities = []
+
+    def select_readout(capability):
+        capabilities.append(capability)
+        return capability in {(8, 0), (9, 0)}
+
+    monkeypatch.setattr(
+        runtime_policy,
+        "is_readout_input_fold_enabled",
+        select_readout,
+    )
+    monkeypatch.setattr(readout_l0, "_has_exact_neo_readout_contract", lambda *_: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: (8, 0))
+    assert readout_l0._can_use_sm80_readout_input_fold(module, value)
+
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: (9, 0))
+    assert readout_l0._can_use_sm80_readout_input_fold(module, value)
+    assert capabilities == [(8, 0), (9, 0)]
+
+
+def test_sm80_readout_input_fold_matches_forward_and_input_vjp(monkeypatch):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    with torch.no_grad():
+        module.so3_linear_2.weight.normal_(std=0.03)
+    value = torch.randn(3, COEFF_DIM, 1, 32, device="cpu")
+    cotangent = torch.randn(3, 32, device="cpu")
+
+    monkeypatch.setattr(
+        readout_l0,
+        "_can_use_sm80_readout_input_fold",
+        lambda *_: False,
+    )
+    value_ref = value.detach().clone().requires_grad_(True)
+    expected = readout_l0._run_neo_readout_l0(
+        module,
+        value_ref,
+        _reference_product,
+    )
+    expected_grad = torch.autograd.grad(expected, value_ref, cotangent)[0]
+
+    monkeypatch.setattr(
+        readout_l0,
+        "_can_use_sm80_readout_input_fold",
+        lambda *_: True,
+    )
+
+    def forbid_staged_projection(*args, **kwargs):
+        del args, kwargs
+        pytest.fail("folded readout must not execute a staged input projection")
+
+    monkeypatch.setattr(module.so3_linear_1, "forward", forbid_staged_projection)
+    monkeypatch.setattr(
+        module.act.grid_op.left_proj,
+        "forward",
+        forbid_staged_projection,
+    )
+    monkeypatch.setattr(
+        module.act.grid_op.right_proj,
+        "forward",
+        forbid_staged_projection,
+    )
+    monkeypatch.setattr(module.act.scalar_gate, "forward", forbid_staged_projection)
+
+    def compact_reference(left, right, to_grid, from_grid):
+        assert left.is_contiguous()
+        assert right.is_contiguous()
+        return _reference_product(left, right, to_grid, from_grid)
+
+    value_actual = value.detach().clone().requires_grad_(True)
+    actual = readout_l0._run_neo_readout_l0(
+        module,
+        value_actual,
+        compact_reference,
+    )
+    actual_grad = torch.autograd.grad(actual, value_actual, cotangent)[0]
+
+    torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+    torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL)
+
+
+def test_sm80_readout_input_fold_is_fullgraph_safe(monkeypatch):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    with torch.no_grad():
+        module.so3_linear_2.weight.normal_(std=0.03)
+    monkeypatch.setattr(
+        readout_l0,
+        "_can_use_sm80_readout_input_fold",
+        lambda *_: True,
+    )
+
+    def folded(value):
+        return readout_l0._run_neo_readout_l0(
+            module,
+            value,
+            _reference_product,
+        )
+
+    value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu")
+    readout_l0.prepare_sm80_readout_input_fold(module)
+    expected = folded(value)
+    compiled = torch.compile(folded, backend="eager", fullgraph=True)
+    actual = compiled(value)
+    torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+
+
+def test_sm80_readout_input_fold_cache_refreshes_after_inplace_change(monkeypatch):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    ready_calls = []
+    monkeypatch.setattr(
+        readout_l0,
+        "_synchronize_sm80_readout_input_fold_build",
+        lambda weights: ready_calls.append(weights),
+    )
+    first = readout_l0.prepare_sm80_readout_input_fold(module)
+    first_cache = getattr(module, readout_l0._READOUT_INPUT_FOLD_CACHE)
+    assert getattr(module, readout_l0._READOUT_INPUT_FOLD_LEFT) is first[0]
+    assert not any("readout_input_fold" in key for key in module.state_dict())
+
+    cached = readout_l0.prepare_sm80_readout_input_fold(module)
+    assert all(
+        actual is expected for actual, expected in zip(cached, first, strict=True)
+    )
+    assert len(ready_calls) == 1
+    with torch.no_grad():
+        module.act.grid_op.left_proj.weight.add_(0.01)
+    second = readout_l0.prepare_sm80_readout_input_fold(module)
+    second_cache = getattr(module, readout_l0._READOUT_INPUT_FOLD_CACHE)
+
+    assert second_cache is not first_cache
+    assert second[0] is not first[0]
+    assert not torch.equal(second[0], first[0])
+    assert len(ready_calls) == 2
+
+
+def test_sm80_readout_input_fold_cache_refreshes_after_parameter_replacement():
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    first = readout_l0.prepare_sm80_readout_input_fold(module)
+    old_weight = module.act.grid_op.left_proj.weight
+    module.act.grid_op.left_proj.weight = torch.nn.Parameter(
+        old_weight.detach().clone().add_(0.02),
+        requires_grad=False,
+    )
+    second = readout_l0.prepare_sm80_readout_input_fold(module)
+
+    assert second[0] is not first[0]
+    assert not torch.equal(second[0], first[0])
+
+
+@pytest.mark.parametrize("assign", [False, True])
+def test_sm80_readout_input_fold_load_state_dict_invalidates_cache(assign):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    first = readout_l0.prepare_sm80_readout_input_fold(module)
+    state = {key: value.clone() for key, value in module.state_dict().items()}
+    left_key = next(key for key in state if key.endswith("grid_op.left_proj.weight"))
+    state[left_key].add_(0.03)
+
+    module.load_state_dict(state, assign=assign)
+
+    assert getattr(module, readout_l0._READOUT_INPUT_FOLD_CACHE) is None
+    assert all(
+        getattr(module, name) is None for name in readout_l0._READOUT_INPUT_FOLD_BUFFERS
+    )
+    second = readout_l0.prepare_sm80_readout_input_fold(module)
+    assert second[0] is not first[0]
+    assert not torch.equal(second[0], first[0])
+    assert not any("readout_input_fold" in key for key in module.state_dict())
+
+
+def test_sm80_readout_input_fold_buffers_follow_device_moves():
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    readout_l0.prepare_sm80_readout_input_fold(module)
+
+    module.to("meta")
+
+    assert all(
+        getattr(module, name).device.type == "meta"
+        for name in readout_l0._READOUT_INPUT_FOLD_BUFFERS
+    )
+    refreshed = readout_l0.prepare_sm80_readout_input_fold(module)
+    assert all(weight.device.type == "meta" for weight in refreshed)
+
+
+def test_sm80_readout_input_fold_compile_requires_explicit_preparation(
+    monkeypatch,
+):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+    with pytest.raises(RuntimeError, match="prepare_sm80_readout_input_fold"):
+        readout_l0._get_sm80_readout_input_fold(module)
+
+
+def test_sm80_readout_input_fold_compiled_lookup_returns_prepared_buffers(
+    monkeypatch,
+):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    expected = readout_l0.prepare_sm80_readout_input_fold(module)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+    cached = readout_l0._get_sm80_readout_input_fold(module)
+    assert all(
+        actual is reference for actual, reference in zip(cached, expected, strict=True)
+    )
+
+
+@pytest.mark.parametrize(
+    ("hidden_channels", "expected"),
+    [(HIDDEN_CHANNELS, True), (96, False), (128, False)],
+)
+def test_shape_guard_accepts_only_exact_c192_readout(
+    hidden_channels: int,
+    expected: bool,
+):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    left, _, _, _ = _cpu_inputs(hidden_channels=hidden_channels)
+    assert readout_l0._has_exact_product_shape(left) is expected
+
+
+def test_fake_registrations_return_canonical_strides():
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    shape = (2, COEFF_DIM, 1, PACKED_WIDTH)
+    canonical = (COEFF_DIM * PACKED_WIDTH, PACKED_WIDTH, PACKED_WIDTH, 1)
+    left = torch.empty_strided(shape, (canonical[0], canonical[1], 7, 1), device="cpu")
+    right = torch.empty_strided(
+        shape, (canonical[0], canonical[1], 11, 1), device="cpu"
+    )
+    dq0 = torch.empty(2, HIDDEN_CHANNELS, device="cpu")
+    to_grid = torch.empty(GRID_SIZE, PACKED_COEFF_DIM, device="cpu")
+    from_grid = torch.empty(PACKED_COEFF_DIM, GRID_SIZE, device="cpu")
+
+    q0 = readout_l0._readout_l0_fake(left, right, to_grid, from_grid)
+    grad_left, grad_right = readout_l0._readout_l0_bwd_fake(
+        dq0, left, right, to_grid, from_grid
+    )
+
+    assert q0.shape == (2, HIDDEN_CHANNELS)
+    assert q0.stride() == (HIDDEN_CHANNELS, 1)
+    assert grad_left.stride() == canonical
+    assert grad_right.stride() == canonical
+
+
+def test_custom_ops_use_canonical_fake_metadata():
+    from torch._subclasses.fake_tensor import (
+        FakeTensorMode,
+    )
+
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    shape = (2, COEFF_DIM, 1, PACKED_WIDTH)
+    canonical = (COEFF_DIM * PACKED_WIDTH, PACKED_WIDTH, PACKED_WIDTH, 1)
+    with FakeTensorMode():
+        left = torch.empty_strided(
+            shape, (canonical[0], canonical[1], 7, 1), device="cuda"
+        )
+        right = torch.empty_strided(
+            shape, (canonical[0], canonical[1], 11, 1), device="cuda"
+        )
+        dq0 = torch.empty(2, HIDDEN_CHANNELS, device="cuda")
+        to_grid = torch.empty(GRID_SIZE, PACKED_COEFF_DIM, device="cuda")
+        from_grid = torch.empty(PACKED_COEFF_DIM, GRID_SIZE, device="cuda")
+        q0 = readout_l0._readout_l0_op(left, right, to_grid, from_grid)
+        grad_left, grad_right = readout_l0._readout_l0_bwd_op(
+            dq0, left, right, to_grid, from_grid
+        )
+
+    assert q0.stride() == (HIDDEN_CHANNELS, 1)
+    assert grad_left.stride() == canonical
+    assert grad_right.stride() == canonical
+
+
+def test_reference_completion_matches_module_output_and_input_vjp():
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    with torch.no_grad():
+        module.so3_linear_2.weight.normal_(std=0.05)
+    value = 0.1 * torch.randn(3, COEFF_DIM, 1, 32, device="cpu")
+    seed = torch.randn(3, 32, device="cpu")
+    value_ref = value.detach().clone().requires_grad_(True)
+    value_actual = value.detach().clone().requires_grad_(True)
+
+    expected = (value_ref + module(value_ref))[:, 0, 0, :]
+    expected_grad = torch.autograd.grad(expected, value_ref, seed)[0]
+    actual = readout_l0._run_neo_readout_l0(module, value_actual, _reference_product)
+    actual_grad = torch.autograd.grad(actual, value_actual, seed)[0]
+
+    torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+    torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL)
+
+
+def test_exact_structure_and_frozen_inference_guards():
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    assert readout_l0.has_neo_readout_contract(module)
+    assert readout_l0._inference_mode_is_frozen(module)
+    assert not readout_l0.has_neo_readout_contract(_output_ffn(hidden_channels=64))
+    module.train()
+    assert not readout_l0._inference_mode_is_frozen(module)
+    module.eval()
+    next(module.parameters()).requires_grad_(True)
+    assert not readout_l0._inference_mode_is_frozen(module)
+
+
+def test_module_boundary_falls_back_when_device_is_unsupported(monkeypatch):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu")
+    expected = (value + module(value))[:, 0, 0, :]
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+
+    assert readout_l0.maybe_run_neo_readout_l0(module, value) is None
+    torch.testing.assert_close(
+        readout_l0.run_neo_output_readout(module, value), expected
+    )
+
+
+def test_descriptor_readout_wiring(cpu_neo_descriptor, monkeypatch):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    descriptor = cpu_neo_descriptor
+    value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu")
+    expected = torch.randn(2, 32, device="cpu")
+    calls = []
+
+    def record_readout(output_ffn, ffn_in):
+        calls.append((output_ffn, ffn_in))
+        return expected
+
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    monkeypatch.setattr(readout_l0, "maybe_run_neo_readout_l0", record_readout)
+    assert descriptor.run_cute_infer_readout(value) is expected
+    assert calls[-1] == (descriptor.output_ffn, value)
+
+
+def test_trainable_descriptor_state_bypasses_candidate(monkeypatch):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    module = _output_ffn()
+    value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu")
+    expected = (value + module(value))[:, 0, 0, :]
+    monkeypatch.setattr(
+        readout_l0,
+        "maybe_run_neo_readout_l0",
+        lambda *args, **kwargs: pytest.fail("candidate must be bypassed"),
+    )
+
+    actual = readout_l0.run_neo_output_readout(module, value, parameters_frozen=False)
+    torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.parametrize(
+    ("precision", "expected"),
+    [("highest", True), ("high", False)],
+)
+def test_legacy_matmul_precision_guard_is_eager_and_fullgraph_safe(
+    matmul_precision_state,
+    precision: str,
+    expected: bool,
+):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    try:
+        matmul_precision_state.fp32_precision = "none"
+    except AttributeError:
+        pass
+    torch.set_float32_matmul_precision(precision)
+    value = torch.ones(1, device="cpu")
+
+    def guarded(tensor):
+        return tensor if readout_l0._uses_strict_fp32_matmul() else -tensor
+
+    eager = guarded(value)
+    fullgraph = torch.compile(guarded, backend="eager", fullgraph=True)(value)
+    assert readout_l0._uses_strict_fp32_matmul() is expected
+    assert torch.equal(fullgraph, eager)
+
+
+@pytest.mark.parametrize(
+    ("precision", "expected"),
+    [("ieee", True), ("tf32", False)],
+)
+def test_modern_matmul_precision_guard_rejects_tf32(
+    matmul_precision_state,
+    precision: str,
+    expected: bool,
+):
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+        readout_l0,
+    )
+
+    _set_new_matmul_precision(matmul_precision_state, precision)
+    assert readout_l0._uses_strict_fp32_matmul() is expected
+
+
+@pytest.mark.skipif(
+    _CUTE_SKIP_REASON is not None,
+    reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable",
+)
+class TestReadoutL0Cuda:
+    @staticmethod
+    def _inputs(nodes: int):
+        generator = torch.Generator(device="cuda").manual_seed(20260704 + nodes)
+        left = 0.1 * torch.randn(
+            nodes,
+            COEFF_DIM,
+            1,
+            PACKED_WIDTH,
+            device="cuda",
+            generator=generator,
+        )
+        right = 0.1 * torch.randn(left.shape, device="cuda", generator=generator)
+        to_grid = 0.1 * torch.randn(
+            GRID_SIZE,
+            PACKED_COEFF_DIM,
+            device="cuda",
+            generator=generator,
+        )
+        from_grid = 0.1 * torch.randn(
+            PACKED_COEFF_DIM,
+            GRID_SIZE,
+            device="cuda",
+            generator=generator,
+        )
+        return left, right, to_grid, from_grid
+
+    @pytest.mark.parametrize("nodes", [1, 7, 65])
+    def test_forward_and_input_vjp_match_strict_fp32(
+        self,
+        nodes: int,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import (
+            readout_l0_product_cute,
+        )
+
+        left, right, to_grid, from_grid = self._inputs(nodes)
+        left_ref = left.detach().clone().requires_grad_(True)
+        right_ref = right.detach().clone().requires_grad_(True)
+        left_actual = left.detach().clone().requires_grad_(True)
+        right_actual = right.detach().clone().requires_grad_(True)
+        dq0 = torch.randn(nodes, HIDDEN_CHANNELS, device="cuda")
+
+        expected = _reference_product(left_ref, right_ref, to_grid, from_grid)
+        expected_grads = torch.autograd.grad(expected, (left_ref, right_ref), dq0)
+        actual = readout_l0_product_cute(left_actual, right_actual, to_grid, from_grid)
+        actual_grads = torch.autograd.grad(actual, (left_actual, right_actual), dq0)
+
+        torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+        torch.testing.assert_close(
+            actual_grads[0], expected_grads[0], atol=TOL, rtol=TOL
+        )
+        torch.testing.assert_close(
+            actual_grads[1], expected_grads[1], atol=TOL, rtol=TOL
+        )
+
+    def test_opcheck_and_fullgraph_preserve_metadata_and_vjp(self):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+            readout_l0,
+        )
+
+        left, right, to_grid, from_grid = self._inputs(3)
+        dq0 = torch.randn(3, HIDDEN_CHANNELS, device="cuda")
+        torch.library.opcheck(
+            readout_l0._readout_l0_op,
+            (left, right, to_grid, from_grid),
+            test_utils=("test_schema", "test_faketensor"),
+        )
+        compiled = torch.compile(
+            readout_l0.readout_l0_product_cute,
+            dynamic=True,
+            fullgraph=True,
+        )
+        left.requires_grad_(True)
+        right.requires_grad_(True)
+        actual = compiled(left, right, to_grid, from_grid)
+        grads = torch.autograd.grad(actual, (left, right), dq0)
+        assert actual.shape == (3, HIDDEN_CHANNELS)
+        assert grads[0].shape == left.shape
+        assert grads[1].shape == right.shape
+
+    def test_exact_module_fullgraph_uses_full_neo_gate(self, monkeypatch):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+            readout_l0,
+        )
+
+        module = _output_ffn(device="cuda")
+        with torch.no_grad():
+            module.so3_linear_2.weight.normal_(std=0.05)
+        value = torch.randn(3, COEFF_DIM, 1, 32, device="cuda")
+        seed = torch.randn(3, 32, device="cuda")
+        value_ref = value.detach().clone().requires_grad_(True)
+        value_actual = value.detach().clone().requires_grad_(True)
+        monkeypatch.delenv("DP_CUTE_INFER", raising=False)
+        expected = (value_ref + module(value_ref))[:, 0, 0, :]
+        expected_grad = torch.autograd.grad(expected, value_ref, seed)[0]
+
+        monkeypatch.setenv("DP_CUTE_INFER", "1")
+        readout_l0.prepare_sm80_readout_input_fold(module)
+        compiled = torch.compile(
+            lambda tensor: readout_l0.run_neo_output_readout(module, tensor),
+            dynamic=True,
+            fullgraph=True,
+        )
+        actual = compiled(value_actual)
+        actual_grad = torch.autograd.grad(actual, value_actual, seed)[0]
+
+        torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+        torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL)
+
+    def test_sm80_input_fold_preparation_is_cross_stream_ready(self):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+            readout_l0,
+        )
+
+        if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6)}:
+            pytest.skip("readout input folding requires SM80 or SM86")
+
+        module = _output_ffn(device="cuda")
+        expected = readout_l0._build_sm80_readout_input_fold(module)
+        torch.cuda.synchronize()
+
+        producer = torch.cuda.Stream()
+        with torch.cuda.stream(producer):
+            prepared = readout_l0.prepare_sm80_readout_input_fold(module)
+        assert producer.query()
+
+        consumer = torch.cuda.Stream()
+        with torch.cuda.stream(consumer):
+            observed = tuple(weight.clone() for weight in prepared)
+        consumer.synchronize()
+        for actual, reference in zip(observed, expected, strict=True):
+            torch.testing.assert_close(actual, reference, atol=0.0, rtol=0.0)
+
+    def test_sm80_input_fold_fullgraph_matches_module_vjp(self, monkeypatch):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+            readout_l0,
+        )
+
+        if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6)}:
+            pytest.skip("readout input folding requires SM80 or SM86")
+
+        module = _output_ffn(device="cuda")
+        with torch.no_grad():
+            module.so3_linear_2.weight.normal_(std=0.05)
+        value = torch.randn(7, COEFF_DIM, 1, 32, device="cuda")
+        seed = torch.randn(7, 32, device="cuda")
+        value_ref = value.detach().clone().requires_grad_(True)
+        value_actual = value.detach().clone().requires_grad_(True)
+        expected = (value_ref + module(value_ref))[:, 0, 0, :]
+        expected_grad = torch.autograd.grad(expected, value_ref, seed)[0]
+
+        monkeypatch.setenv("DP_CUTE_INFER", "1")
+        monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", "1")
+        assert readout_l0._can_use_sm80_readout_input_fold(module, value_actual)
+        readout_l0.prepare_sm80_readout_input_fold(module)
+        compiled = torch.compile(
+            lambda tensor: readout_l0.run_neo_output_readout(module, tensor),
+            dynamic=True,
+            fullgraph=True,
+        )
+        actual = compiled(value_actual)
+        actual_grad = torch.autograd.grad(actual, value_actual, seed)[0]
+
+        torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+        torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL)
diff --git a/source/tests/pt/model/sezm_cute/test_runtime_policy.py b/source/tests/pt/model/sezm_cute/test_runtime_policy.py
new file mode 100644
index 0000000000..aca4b0938d
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_runtime_policy.py
@@ -0,0 +1,229 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Behavioral tests for CuTe inference feature policy."""
+
+from __future__ import (
+    annotations,
+)
+
+import os
+from unittest import (
+    mock,
+)
+
+import pytest
+
+from deepmd.pt_expt.kernels.cute.sezm import runtime_policy as policy
+
+
+@pytest.mark.parametrize("value", ["2", "unsupported", "tru"])
+@pytest.mark.parametrize("name", ["DP_CUTE_STRICT", "DP_CUTE_SO2_THIN_WRAPPER"])
+def test_invalid_boolean_override_raises(name: str, value: str) -> None:
+    with mock.patch.dict(os.environ, {name: value}, clear=True):
+        with pytest.raises(ValueError, match=name):
+            policy._env_override(name)
+
+
+def test_cute_master_gate_controls_sezm_path() -> None:
+    cases = (
+        ({}, False),
+        ({"DP_CUTE_INFER": "0"}, False),
+        ({"DP_CUTE_INFER": "1"}, True),
+        ({"DP_CUTE_INFER": "true"}, True),
+        ({"DP_CUTE_INFER": "unsupported"}, False),
+        ({"DP_CUTE_INFER": "1", "DP_TRITON_INFER": "2"}, True),
+    )
+    for environment, expected in cases:
+        with mock.patch.dict(os.environ, environment, clear=True):
+            assert policy.is_cute_infer_enabled() is expected
+
+
+def test_master_switch_controls_every_sm80_subfeature() -> None:
+    with mock.patch.dict(
+        os.environ,
+        {
+            "DP_CUTE_GIE": "1",
+            "DP_CUTE_SO2_PACKED_WIGNER": "1",
+            "DP_CUTE_SO2_THIN_WRAPPER": "1",
+            "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": "1",
+            "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48": "1",
+            "DP_CUTE_READOUT_INPUT_FOLD_SM80": "1",
+        },
+        clear=True,
+    ):
+        assert not policy.is_cute_infer_enabled()
+        for capability in policy.SM80_PROFILE_CAPABILITIES:
+            assert not policy.is_sm80_profile_enabled(capability)
+            assert not policy.is_gie_enabled(capability)
+            assert not policy.is_packed_wigner_enabled(capability)
+            assert not policy.is_so2_thin_wrapper_enabled(capability)
+            assert not policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability)
+            assert not policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability)
+            assert not policy.is_readout_input_fold_sm80_enabled(capability)
+
+
+def test_sm80_family_profile_defaults_from_master_only() -> None:
+    with mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=True):
+        assert policy.SM80_PROFILE_CAPABILITIES == frozenset({(8, 0), (8, 6)})
+        for capability in policy.SM80_PROFILE_CAPABILITIES:
+            assert policy.is_sm80_profile_enabled(capability)
+            assert policy.is_gie_enabled(capability)
+            assert policy.is_packed_wigner_enabled(capability)
+            assert policy.is_so2_thin_wrapper_enabled(capability)
+            assert policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability)
+            assert policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability)
+            assert policy.is_readout_input_fold_sm80_enabled(capability)
+
+
+def test_every_sm80_profile_feature_can_be_disabled() -> None:
+    checks = {
+        "DP_CUTE_GIE": lambda: policy.is_gie_enabled((8, 0)),
+        "DP_CUTE_SO2_PACKED_WIGNER": lambda: policy.is_packed_wigner_enabled((8, 0)),
+        "DP_CUTE_SO2_THIN_WRAPPER": lambda: policy.is_so2_thin_wrapper_enabled((8, 0)),
+        "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": lambda: (
+            policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled((8, 0))
+        ),
+        "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48": lambda: (
+            policy.is_output_grid_fwd_sm80_c96_n48_enabled((8, 0))
+        ),
+        "DP_CUTE_READOUT_INPUT_FOLD_SM80": lambda: (
+            policy.is_readout_input_fold_sm80_enabled((8, 0))
+        ),
+    }
+    for name, checker in checks.items():
+        with mock.patch.dict(
+            os.environ,
+            {"DP_CUTE_INFER": "1", name: "0"},
+            clear=True,
+        ):
+            assert not checker()
+
+
+def test_non_sm80_family_uses_architecture_defaults_without_sm80_features() -> None:
+    with mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=True):
+        capability = (8, 9)
+        assert not policy.is_sm80_profile_enabled(capability)
+        assert not policy.is_gie_enabled(capability)
+        assert policy.is_packed_wigner_enabled(capability)
+        assert not policy.is_so2_thin_wrapper_enabled(capability)
+        assert not policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability)
+        assert not policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability)
+        assert not policy.is_readout_input_fold_sm80_enabled(capability)
+
+
+def test_sm90_c96_asymmetric_panels_default_is_exact_arch_disable_only() -> None:
+    switch = policy.OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS_ENV
+
+    with mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=True):
+        assert policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0))
+        assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((8, 9))
+        assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 1))
+
+    with mock.patch.dict(
+        os.environ,
+        {"DP_CUTE_INFER": "1", switch: "0"},
+        clear=True,
+    ):
+        assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0))
+
+    with mock.patch.dict(os.environ, {switch: "1"}, clear=True):
+        assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0))
+
+    with mock.patch.dict(
+        os.environ,
+        {"DP_CUTE_INFER": "1", switch: "1"},
+        clear=True,
+    ):
+        assert policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0))
+        assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((8, 0))
+
+
+def test_overrides_remain_architecture_safe() -> None:
+    with mock.patch.dict(
+        os.environ,
+        {
+            "DP_CUTE_INFER": "1",
+            "DP_CUTE_SO2_THIN_WRAPPER": "1",
+            "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": "1",
+            "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48": "1",
+            "DP_CUTE_READOUT_INPUT_FOLD_SM80": "1",
+        },
+        clear=True,
+    ):
+        assert policy.is_so2_thin_wrapper_enabled((9, 0))
+        assert not policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled((9, 0))
+        assert not policy.is_output_grid_fwd_sm80_c96_n48_enabled((9, 0))
+        assert not policy.is_readout_input_fold_sm80_enabled((9, 0))
+
+    with mock.patch.dict(
+        os.environ,
+        {
+            "DP_CUTE_INFER": "1",
+            "DP_CUTE_GIE": "1",
+            "DP_CUTE_SO2_PACKED_WIGNER": "1",
+        },
+        clear=True,
+    ):
+        assert not policy.is_gie_enabled((9, 0))
+        assert not policy.is_packed_wigner_enabled((10, 1))
+
+
+def test_int32_so2_capacity_checks_every_flattened_axis() -> None:
+    max_edges = policy.INT32_MAX // policy.SO2_VALUES_PER_EDGE
+    max_nodes = policy.INT32_MAX // policy.SO2_VALUES_PER_NODE
+
+    assert policy.so2_int32_indexing_is_safe(
+        edge_count=max_edges,
+        node_count=max_nodes,
+    )
+    assert not policy.so2_int32_indexing_is_safe(
+        edge_count=max_edges + 1,
+        node_count=1,
+    )
+    assert not policy.so2_int32_indexing_is_safe(
+        edge_count=1,
+        node_count=max_nodes + 1,
+    )
+    assert not policy.so2_int32_indexing_is_safe(
+        edge_count=-1,
+        node_count=1,
+    )
+    assert not policy.so2_int32_indexing_is_safe(
+        edge_count=1,
+        node_count=-1,
+    )
+
+
+def test_strict_mode_is_explicitly_opt_in() -> None:
+    with mock.patch.dict(os.environ, {}, clear=True):
+        assert not policy.is_cute_strict_enabled()
+    with mock.patch.dict(os.environ, {"DP_CUTE_STRICT": "1"}, clear=True):
+        assert policy.is_cute_strict_enabled()
+
+
+def test_neighbor_list_eager_island_policy_is_owned_by_neo_runtime() -> None:
+    cases = (
+        ({}, (8, 0), False),
+        ({"DP_CUTE_INFER": "1"}, (8, 0), True),
+        ({"DP_CUTE_INFER": "1"}, (8, 6), False),
+        ({"DP_CUTE_INFER": "1"}, (9, 0), False),
+        (
+            {
+                "DP_CUTE_INFER": "1",
+                "DP_CUTE_SO2_EAGER_ISLANDS": "0",
+            },
+            (8, 0),
+            False,
+        ),
+        (
+            {
+                "DP_CUTE_INFER": "1",
+                "DP_CUTE_SO2_EAGER_ISLANDS": "1",
+            },
+            (9, 0),
+            True,
+        ),
+    )
+    for environment, capability, expected in cases:
+        with mock.patch.dict(os.environ, environment, clear=True):
+            assert policy.is_so2_eager_island_enabled(capability) is expected
diff --git a/source/tests/pt/model/sezm_cute/test_so2.py b/source/tests/pt/model/sezm_cute/test_so2.py
new file mode 100644
index 0000000000..69d898eb77
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2.py
@@ -0,0 +1,743 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Contracts for the opt-in CuTe Neo SO2 inference path."""
+
+from __future__ import (
+    annotations,
+)
+
+import importlib
+from dataclasses import (
+    fields,
+    replace,
+)
+from types import (
+    SimpleNamespace,
+)
+
+import pytest
+import torch
+
+from deepmd.pt.model.descriptor.sezm_nn.edge_cache import (
+    EdgeFeatureCache,
+    _build_edge_wigner,
+    _separate_packed_wigner,
+)
+from deepmd.pt.model.descriptor.sezm_nn.norm import (
+    EquivariantRMSNorm,
+    ScalarRMSNorm,
+)
+from deepmd.pt_expt.kernels.cute.sezm import (
+    runtime_policy,
+)
+from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as _SO2
+from deepmd.pt_expt.kernels.cute.sezm.so2 import runner as _SO2_RUNNER
+
+NeoSO2BackwardWorkspace = _SO2_RUNNER.NeoSO2BackwardWorkspace
+NeoSO2RuntimeConfig = _SO2.NeoSO2RuntimeConfig
+NeoFullCuteBackward = _SO2_RUNNER.NeoFullCuteBackward
+StackCache = _SO2_RUNNER.StackCache
+_validate_runtime_config = _SO2_RUNNER._validate_runtime_config
+_uses_packed_message_grid = _SO2_RUNNER._uses_packed_message_grid
+_destination_degrees_fit_limit = _SO2._destination_degrees_fit_limit
+
+
+class _Identity:
+    pass
+
+
+def _neo_like_block(**overrides):
+    frame_contract = {
+        "coefficient_layout": "packed",
+        "n_frames": 3,
+        "channels": 32,
+    }
+    message_node_grid_product = SimpleNamespace(
+        layout="flat",
+        mode="cross",
+        op_type="glu",
+        n_focus=2,
+        n_frames=3,
+        channels=32,
+        dtype=torch.float32,
+        frame_expand=SimpleNamespace(**frame_contract),
+        frame_contract=SimpleNamespace(**frame_contract),
+    )
+    so2 = SimpleNamespace(
+        lmax=3,
+        mmax=1,
+        ebed_dim_full=16,
+        reduced_dim=10,
+        channels=32,
+        n_focus=2,
+        so2_focus_dim=32,
+        hidden_channels=64,
+        mixing_layers=3,
+        n_atten_head=1,
+        head_dim=32,
+        radial_so2_mode="degree_channel",
+        radial_so2_rank=1,
+        so2_norm=False,
+        focus_compete=True,
+        focus_norm=True,
+        edge_cartesian=False,
+        node_cartesian_tp=None,
+        message_node_grid_product=message_node_grid_product,
+        atten_f_mix=False,
+        attn_v_proj=None,
+        attn_o_proj=None,
+        mlp_bias=False,
+        layer_scale=False,
+        use_so2_attn_res=False,
+    )
+    for key, value in overrides.items():
+        setattr(so2, key, value)
+    so2.focus_compete_norm = (
+        ScalarRMSNorm(
+            channels=32,
+            n_focus=2,
+            dtype=torch.float32,
+            trainable=False,
+        )
+        if so2.focus_norm
+        else None
+    )
+    block = torch.nn.Module()
+    block.so2_conv = so2
+    block.lmax = 3
+    block.node_lmax = 3
+    block.pre_so2_norm = _Identity()
+    block.post_so2_norm = EquivariantRMSNorm(
+        3,
+        32,
+        dtype=torch.float32,
+        trainable=False,
+    )
+    block.runtime_weight = torch.nn.Parameter(
+        torch.ones(1, dtype=torch.float32, device="cpu"),
+        requires_grad=False,
+    )
+    return block
+
+
+@pytest.mark.parametrize("focus_norm", [True, False])
+def test_exact_neo_contract_is_supported(focus_norm: bool) -> None:
+    block = _neo_like_block(focus_norm=focus_norm)
+
+    assert _SO2.get_neo_so2_spec(block).is_current_neo_target
+    assert _SO2.is_supported_neo_so2_block(block)
+
+
+@pytest.mark.parametrize(
+    ("field", "value"),
+    [
+        ("lmax", 4),
+        ("mmax", 2),
+        ("channels", 64),
+        ("n_focus", 1),
+        ("mixing_layers", 2),
+        ("focus_compete", False),
+        ("edge_cartesian", True),
+    ],
+)
+def test_non_neo_contracts_fall_back(field: str, value: object) -> None:
+    assert not _SO2.is_supported_neo_so2_block(_neo_like_block(**{field: value}))
+
+
+def test_unsupported_message_grid_contract_falls_back() -> None:
+    block = _neo_like_block()
+    block.so2_conv.message_node_grid_product.op_type = "mlp"
+
+    assert not _SO2.is_supported_neo_so2_block(block)
+
+
+def test_unsupported_post_norm_falls_back() -> None:
+    block = _neo_like_block()
+    block.post_so2_norm = torch.nn.LayerNorm(32, device="cpu")
+
+    assert not _SO2.is_supported_neo_so2_block(block)
+
+
+def test_gate_expand_contract_cache_tracks_buffer_versions() -> None:
+    block = _neo_like_block()
+    expand_index = torch.tensor(
+        [0, 1, 2, 0, 1, 2, 0, 1, 2],
+        dtype=torch.long,
+        device="cpu",
+    )
+    block.so2_conv.non_linearities = [SimpleNamespace(expand_index=expand_index)]
+
+    assert _SO2._gate_expand_index_structure_is_supported(block)
+    assert _SO2._gate_expand_index_is_supported(block)
+    expand_index[0] = 2
+    assert _SO2._gate_expand_index_structure_is_supported(block)
+    assert not _SO2._gate_expand_index_is_supported(block)
+
+
+@pytest.mark.parametrize(
+    "capability",
+    sorted(runtime_policy.SUPPORTED_SO2_CAPABILITIES),
+)
+def test_every_supported_architecture_has_a_validated_config(
+    capability: tuple[int, int],
+) -> None:
+    config = _SO2._architecture_default_config(capability)
+
+    assert _validate_runtime_config(config, compute_capability=capability) is None
+
+
+@pytest.mark.parametrize("capability", [(8, 0), (8, 6)])
+def test_sm80_family_uses_per_focus_so2_forward(
+    capability: tuple[int, int],
+) -> None:
+    config = _SO2._architecture_default_config(capability)
+
+    assert config.per_focus_so2_fwd_pair
+    assert not config.native_sm90_path
+    assert not config.combined_so2_gate
+
+
+def test_sm90_uses_native_split_complex_path() -> None:
+    config = _SO2._architecture_default_config((9, 0))
+
+    assert config.native_sm90_path
+    assert not config.per_focus_so2_fwd_pair
+    assert not config.combined_so2_gate
+
+
+def test_sm90_portable_fallback_is_a_validated_config() -> None:
+    config = replace(
+        _SO2._architecture_default_config((9, 0)),
+        native_sm90_path=False,
+    )
+
+    assert _validate_runtime_config(config, compute_capability=(9, 0)) is None
+
+
+@pytest.mark.parametrize(
+    ("row_ptr", "expected"),
+    [
+        ([0, 256, 512], True),
+        ([0, 257, 512], False),
+        ([0, 4, 3], False),
+    ],
+)
+def test_destination_degree_limit(
+    row_ptr: list[int],
+    expected: bool,
+) -> None:
+    destination_row_ptr = torch.tensor(row_ptr, dtype=torch.int32, device="cpu")
+
+    assert _destination_degrees_fit_limit(destination_row_ptr, 256) is expected
+
+
+def test_sm100_uses_shared_default_profile() -> None:
+    config = _SO2._architecture_default_config((10, 0))
+
+    assert not config.native_sm90_path
+    assert not config.per_focus_so2_fwd_pair
+    assert not config.combined_so2_gate
+
+
+@pytest.mark.parametrize("capability", [(8, 9), (12, 0)])
+def test_sm89_and_sm120_use_combined_so2_gate(
+    capability: tuple[int, int],
+) -> None:
+    config = _SO2._architecture_default_config(capability)
+
+    assert config.combined_so2_gate
+    assert not config.native_sm90_path
+    assert not config.per_focus_so2_fwd_pair
+
+
+def test_combined_so2_gate_rejects_misaligned_contiguous_tensor() -> None:
+    pytest.importorskip("cutlass.cute")
+    pytest.importorskip("cuda.bindings.driver")
+    from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.combined_gate_forward import (
+        _require_16_byte_alignment,
+    )
+
+    storage = torch.empty(9, dtype=torch.float32, device="cpu")
+    aligned = storage[:8]
+    misaligned = storage[1:9]
+    assert aligned.data_ptr() % 16 == 0
+    assert misaligned.is_contiguous()
+    _require_16_byte_alignment((aligned,))
+    with pytest.raises(ValueError, match="16-byte aligned"):
+        _require_16_byte_alignment((misaligned,))
+
+
+def test_runtime_config_contains_only_reached_selectors() -> None:
+    assert {field.name for field in fields(NeoSO2RuntimeConfig)} == {
+        "native_sm90_path",
+        "per_focus_so2_fwd_pair",
+        "combined_so2_gate",
+    }
+
+
+def test_runtime_config_rejects_incompatible_capability() -> None:
+    config = _SO2._architecture_default_config((8, 0))
+
+    with pytest.raises(RuntimeError, match="supported compute capability"):
+        _validate_runtime_config(config, compute_capability=(7, 5))
+
+
+def test_runtime_config_rejects_wrong_architecture_profile() -> None:
+    sm80_config = _SO2._architecture_default_config((8, 0))
+    with pytest.raises(RuntimeError, match="per-focus SO2"):
+        _validate_runtime_config(sm80_config, compute_capability=(12, 0))
+
+    sm90_config = _SO2._architecture_default_config((9, 0))
+    with pytest.raises(RuntimeError, match="native SM90 SO2"):
+        _validate_runtime_config(sm90_config, compute_capability=(8, 0))
+
+    sm120_config = _SO2._architecture_default_config((12, 0))
+    with pytest.raises(RuntimeError, match="combined SO2/gate"):
+        _validate_runtime_config(sm120_config, compute_capability=(10, 0))
+    with pytest.raises(RuntimeError, match="combined SO2/gate"):
+        _validate_runtime_config(NeoSO2RuntimeConfig(), compute_capability=(12, 0))
+
+
+@pytest.mark.parametrize(
+    ("capability", "expected"),
+    [
+        ((8, 0), True),
+        ((8, 6), True),
+        ((8, 9), False),
+        ((9, 0), True),
+        ((10, 0), False),
+        ((12, 0), False),
+    ],
+)
+def test_packed_message_grid_architecture_contract(
+    capability: tuple[int, int],
+    expected: bool,
+) -> None:
+    assert _uses_packed_message_grid(capability) is expected
+
+
+def test_packed_edge_eligibility_requires_sorted_strict_fp32() -> None:
+    kwargs = {
+        "candidate": True,
+        "edge_count": 12,
+        "node_count": 4,
+        "destinations_sorted": True,
+        "runtime_dtypes": (torch.float32, torch.float32),
+    }
+
+    assert _SO2.packed_wigner_edges_eligible(**kwargs)
+    assert not _SO2.packed_wigner_edges_eligible(
+        **{**kwargs, "destinations_sorted": False}
+    )
+    assert not _SO2.packed_wigner_edges_eligible(
+        **{**kwargs, "runtime_dtypes": (torch.float64,)}
+    )
+    assert _SO2.packed_wigner_edges_eligible(
+        **{**kwargs, "edge_count": 3, "node_count": 4}
+    )
+
+
+def test_packed_wigner_has_a_separate_backward_compatible_cache_field() -> None:
+    dense = torch.empty((2, 16, 16), device="cpu")
+    dense_t = torch.empty_like(dense)
+    actual_dense, actual_dense_t, packed = _separate_packed_wigner(dense, dense_t)
+    assert actual_dense is dense
+    assert actual_dense_t is dense_t
+    assert packed is None
+
+    panel = torch.empty((2, 46), device="cpu")
+    actual_dense, actual_dense_t, packed = _separate_packed_wigner(panel, panel)
+    assert actual_dense is None
+    assert actual_dense_t is None
+    assert packed is panel
+    assert EdgeFeatureCache._fields[-2:] == ("destinations_sorted", "D_packed")
+
+
+def test_ineligible_wigner_build_does_not_import_cute(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    edge_vec = torch.tensor([[1.0, 0.0, 0.0]], device="cpu")
+    edge_len = torch.linalg.vector_norm(edge_vec, dim=-1, keepdim=True)
+
+    def eager_wigner(edge_quat: torch.Tensor):
+        dense = edge_quat.new_zeros((edge_quat.shape[0], 1, 1))
+        return dense, dense.transpose(-1, -2)
+
+    dense, dense_t, edge_quat = _build_edge_wigner(
+        edge_vec=edge_vec,
+        edge_len=edge_len,
+        eps=1.0e-8,
+        random_gamma=False,
+        wigner_calc=eager_wigner,
+        packed_wigner=False,
+    )
+
+    assert dense is not None and tuple(dense.shape) == (1, 1, 1)
+    assert dense_t is not None and tuple(dense_t.shape) == (1, 1, 1)
+    assert tuple(edge_quat.shape) == (1, 4)
+
+
+@pytest.mark.parametrize(
+    ("training", "dtype", "edge_count", "node_count", "destinations_sorted"),
+    [
+        (True, torch.float32, 12, 4, True),
+        (False, torch.float64, 12, 4, True),
+        (False, torch.float32, 0, 4, True),
+        (False, torch.float32, 12, 4, False),
+    ],
+)
+def test_ineligible_runtime_contracts_fall_back(
+    monkeypatch: pytest.MonkeyPatch,
+    training: bool,
+    dtype: torch.dtype,
+    edge_count: int,
+    node_count: int,
+    destinations_sorted: bool,
+) -> None:
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    block = _neo_like_block().eval()
+
+    assert not _SO2.is_neo_so2_runtime_eligible(
+        block,
+        training=training,
+        device=torch.device("cuda", 0),
+        dtype=dtype,
+        edge_count=edge_count,
+        node_count=node_count,
+        destinations_sorted=destinations_sorted,
+    )
+
+
+def test_runtime_contract_allows_fewer_edges_than_nodes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    block = _neo_like_block().eval()
+
+    assert _SO2.is_neo_so2_runtime_eligible(
+        block,
+        training=False,
+        device=torch.device("cuda", 0),
+        dtype=torch.float32,
+        edge_count=3,
+        node_count=4,
+        destinations_sorted=True,
+    )
+
+
+def _cute_cuda_runtime_available() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        importlib.import_module("cutlass.cute")
+        importlib.import_module("cuda.bindings.driver")
+    except Exception:  # pragma: no cover - runtime dependent
+        return False
+    return True
+
+
+@pytest.mark.skipif(
+    not _cute_cuda_runtime_available(),
+    reason="CuTe attention-prelude regression requires CUDA and CuTe DSL",
+)
+@pytest.mark.parametrize("use_focus_norm", [True, False])
+def test_attention_prelude_initializes_all_qk_rows_when_edges_are_sparse(
+    use_focus_norm: bool,
+) -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.focus_source_backward import (
+        compile_neo_attention_prelude_forward,
+    )
+
+    torch.manual_seed(20260810)
+    device = torch.device("cuda")
+    edge_count = 2
+    node_count = 3
+    focus = torch.randn(edge_count, 64, device=device)
+    x_l0 = torch.randn(node_count, 2, 32, device=device)
+    focus_weight = torch.randn(32, 2, device=device)
+    focus_scale = torch.randn(2, 32, device=device)
+    q_weight = torch.randn(32, 2, 32, device=device)
+    k_weight = torch.randn_like(q_weight)
+    qk_scale = torch.randn(2, 32, device=device)
+    focus_alpha = torch.full((edge_count, 2), torch.nan, device=device)
+    q_node = torch.full((node_count, 2, 32), torch.nan, device=device)
+    k_node = torch.full_like(q_node, torch.nan)
+    focus_eps = 1.0e-5
+    qk_eps = 1.0e-5
+    tau = 0.7
+    label_smoothing = 0.1
+
+    focus_view = focus.view(edge_count, 2, 32)
+    focus_input = focus_view
+    if use_focus_norm:
+        focus_input = (
+            focus_view
+            * torch.rsqrt(focus_view.square().mean(dim=-1, keepdim=True) + focus_eps)
+            * focus_scale.unsqueeze(0)
+        )
+    focus_logits = torch.stack(
+        [
+            (focus_input[:, focus_idx] * focus_weight[:, focus_idx]).sum(dim=-1)
+            for focus_idx in range(2)
+        ],
+        dim=-1,
+    )
+    expected_alpha = torch.softmax(focus_logits / tau, dim=-1)
+    expected_alpha = expected_alpha * (1.0 - label_smoothing) + (label_smoothing / 2.0)
+    x_norm = (
+        x_l0
+        * torch.rsqrt(x_l0.square().mean(dim=-1, keepdim=True) + qk_eps)
+        * qk_scale.unsqueeze(0)
+    )
+    expected_q = torch.einsum("nfi,ifo->nfo", x_norm, q_weight)
+    expected_k = torch.einsum("nfi,ifo->nfo", x_norm, k_weight)
+
+    run = compile_neo_attention_prelude_forward(
+        focus_eps,
+        qk_eps,
+        tau,
+        label_smoothing,
+        use_focus_norm=use_focus_norm,
+    )
+    run(
+        focus,
+        x_l0,
+        focus_weight,
+        focus_scale,
+        q_weight,
+        k_weight,
+        qk_scale,
+        focus_alpha,
+        q_node,
+        k_node,
+    )
+    torch.cuda.synchronize()
+
+    assert torch.isfinite(q_node).all()
+    assert torch.isfinite(k_node).all()
+    torch.testing.assert_close(focus_alpha, expected_alpha, atol=5.0e-5, rtol=5.0e-5)
+    torch.testing.assert_close(q_node, expected_q, atol=5.0e-5, rtol=5.0e-5)
+    torch.testing.assert_close(k_node, expected_k, atol=5.0e-5, rtol=5.0e-5)
+
+
+def test_exact_runtime_contract_is_eligible_without_autocast(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: True)
+    block = _neo_like_block().eval()
+    kwargs = {
+        "training": False,
+        "device": torch.device("cuda", 0),
+        "dtype": torch.float32,
+        "edge_count": 12,
+        "node_count": 4,
+        "destinations_sorted": True,
+    }
+
+    assert _SO2.is_neo_so2_runtime_eligible(block, **kwargs)
+    monkeypatch.setattr(torch, "is_autocast_enabled", lambda device_type: True)
+    assert not _SO2.is_neo_so2_runtime_eligible(block, **kwargs)
+
+
+def test_runtime_contract_rejects_tf32_matmul_policy(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    monkeypatch.setattr(torch, "is_autocast_enabled", lambda device_type: False)
+    block = _neo_like_block().eval()
+    kwargs = {
+        "training": False,
+        "device": torch.device("cuda", 0),
+        "dtype": torch.float32,
+        "edge_count": 12,
+        "node_count": 4,
+        "destinations_sorted": True,
+    }
+
+    monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: True)
+    assert _SO2.is_neo_so2_runtime_eligible(block, **kwargs)
+
+    monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: False)
+    assert not _SO2.is_neo_so2_runtime_eligible(block, **kwargs)
+
+
+@pytest.mark.parametrize("change", ["tf32", "trainable", "dtype", "training"])
+def test_prepared_dispatch_rechecks_mutable_model_state(
+    monkeypatch: pytest.MonkeyPatch, change: str
+) -> None:
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: True)
+    block = _neo_like_block().eval()
+    config = NeoSO2RuntimeConfig()
+    handle = _SO2.register_cute_so2_block(block, config)
+    block._deepmd_cute_so2_state = _SO2._RegisteredSO2State(0, handle, config)
+    try:
+        assert _SO2._has_frozen_fp32_contract(block)
+        if change == "tf32":
+            monkeypatch.setattr(
+                runtime_policy, "uses_strict_fp32_matmul", lambda: False
+            )
+        elif change == "trainable":
+            block.runtime_weight.requires_grad_(True)
+        elif change == "dtype":
+            block.double()
+        else:
+            block.train()
+
+        # Rejection must happen before inspecting edge data or entering a kernel.
+        assert _SO2._maybe_run_prepared_cute_so2(block, None, None, None) is None
+        with pytest.raises(RuntimeError, match="model or precision state changed"):
+            _SO2._build_runner(handle, *([None] * 11))
+    finally:
+        _SO2.invalidate_cute_so2_state(block)
+
+
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_packed_dispatch_decline_recovers_outputs_and_gradients(
+    monkeypatch: pytest.MonkeyPatch, backend: str, dtype: torch.dtype
+) -> None:
+    from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import (
+        EdgeCache,
+    )
+    from deepmd.pt.model.descriptor.sezm_nn.block import (
+        SeZMInteractionBlock,
+    )
+    from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+        WignerDCalculator,
+    )
+    from deepmd.pt_expt.descriptor.dpa4_nn.block import (
+        SeZMInteractionBlock as ExptBlock,
+    )
+    from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import (
+        iter_packed_entries,
+    )
+
+    monkeypatch.setenv("DP_TRITON_INFER", "0")
+    monkeypatch.setenv("DP_CUTE_INFER", "0")
+    generator = torch.Generator(device="cpu").manual_seed(71)
+    block = (
+        SeZMInteractionBlock(
+            lmax=3,
+            mmax=1,
+            channels=4,
+            mixing_layers=3,
+            radial_so2_mode="degree_channel",
+            radial_so2_rank=1,
+            ffn_activation_function="silu",
+            dtype=dtype,
+            seed=11,
+            trainable=True,
+        )
+        .cpu()
+        .eval()
+    )
+    with torch.no_grad():
+        for parameter in block.parameters():
+            parameter.add_(
+                0.05
+                * torch.randn(
+                    parameter.shape, dtype=dtype, device="cpu", generator=generator
+                )
+            )
+    if backend == "pt_expt":
+        block = ExptBlock.deserialize(block.serialize()).cpu().eval()
+    calculator = WignerDCalculator(3, dtype=dtype).cpu()
+    x_base = torch.randn(3, 16, 1, 4, dtype=dtype, device="cpu", generator=generator)
+    radial_base = torch.randn(4, 4, 4, dtype=dtype, device="cpu", generator=generator)
+    quat_base = torch.randn(4, 4, dtype=dtype, device="cpu", generator=generator)
+    src = torch.tensor([1, 2, 0, 1], dtype=torch.int64, device="cpu")
+    dst = torch.tensor([0, 0, 1, 2], dtype=torch.int64, device="cpu")
+    parameter = block.so2_conv.so2_linears[0].weight_m0.requires_grad_(True)
+
+    def run(*, packed: bool):
+        x, radial, quaternion = (
+            value.clone().requires_grad_(True)
+            for value in (x_base, radial_base, quat_base)
+        )
+        d_full, dt_full = calculator(quaternion)
+        panel = torch.stack(
+            [
+                d_full[:, entry.full_row, entry.full_col]
+                for entry in iter_packed_entries()
+            ],
+            dim=1,
+        )
+        cache_type = EdgeFeatureCache if backend == "pt" else EdgeCache
+        cache = cache_type(
+            src=src,
+            dst=dst,
+            edge_type_feat=radial[:, 0],
+            edge_vec=quaternion[:, 1:],
+            edge_rbf=radial[:, 0],
+            edge_env=torch.ones(4, 1, dtype=dtype, device="cpu"),
+            deg=torch.tensor([2, 1, 1], dtype=dtype, device="cpu"),
+            inv_sqrt_deg=torch.ones(3, 1, 1, dtype=dtype, device="cpu"),
+            edge_quat=quaternion,
+            D_full=None if packed else d_full,
+            Dt_full=None if packed else dt_full,
+            D_packed=panel if packed else None,
+            destinations_sorted=True,
+        )
+        output = block._run_so2_unit_impl(x, cache, radial)
+        gradients = torch.autograd.grad(
+            output.square().sum(), (x, radial, quaternion, parameter)
+        )
+        if packed:
+            # Recovery is local to the declined block, not a mutation of shared state.
+            assert cache.D_full is None and cache.D_packed is panel
+        return output, gradients
+
+    expected, expected_grads = run(packed=False)
+    declined_calls = []
+
+    def decline(*args, **kwargs):
+        declined_calls.append(True)
+
+    monkeypatch.setattr(_SO2, "maybe_run_cute_so2", decline)
+    actual, actual_grads = run(packed=True)
+    assert declined_calls == [True]
+    tol = 5e-5 if dtype == torch.float32 else 1e-10
+    torch.testing.assert_close(actual, expected, atol=tol, rtol=tol)
+    for actual_grad, expected_grad in zip(actual_grads, expected_grads, strict=True):
+        torch.testing.assert_close(actual_grad, expected_grad, atol=tol, rtol=tol)
+
+
+def test_fake_native_gradient_layout_matches_runtime_contract() -> None:
+    x = torch.empty(5, 16, 1, 32, device="cpu")
+
+    grad = _SO2._fake_x_wide_grad_like(x, skip=True)
+
+    assert grad.shape == x.shape
+    assert grad.stride() == (32, 5 * 32, 32, 1)
+
+
+def test_custom_op_runtime_canonicalizes_misaligned_contiguous_view() -> None:
+    base = torch.arange(17, dtype=torch.float32, device="cpu")
+    offset_view = base[1:]
+
+    assert offset_view.is_contiguous()
+    assert offset_view.storage_offset() == 1
+    actual = _SO2._aligned_contiguous(offset_view)
+
+    torch.testing.assert_close(actual, offset_view)
+    assert actual.is_contiguous()
+    assert actual.storage_offset() == 0
+    assert actual.data_ptr() % 16 == 0
+
+
+def test_custom_op_runtime_preserves_aligned_compact_tensor() -> None:
+    tensor = torch.arange(16, dtype=torch.float32, device="cpu")
+
+    assert _SO2._aligned_contiguous(tensor) is tensor
+
+
+def test_stack_cache_retains_only_backward_state() -> None:
+    assert {field.name for field in fields(StackCache)} == {
+        "y",
+        "logits",
+        "non_linear",
+        "final",
+    }
diff --git a/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py
new file mode 100644
index 0000000000..e1646699d9
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py
@@ -0,0 +1,702 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Contracts for asynchronous SO2 destination-pointer construction."""
+
+from __future__ import (
+    annotations,
+)
+
+import ast
+import importlib
+import unittest
+from itertools import (
+    count,
+)
+from types import (
+    SimpleNamespace,
+)
+from typing import (
+    Any,
+)
+from unittest import (
+    mock,
+)
+
+from ._paths import (
+    CUTE_ROOT,
+)
+
+try:
+    import torch
+except ModuleNotFoundError:  # pragma: no cover - lightweight source-test host
+    torch = None
+
+
+SO2_PATH = CUTE_ROOT / "so2" / "operation.py"
+SO2_METADATA_PATH = CUTE_ROOT / "so2" / "metadata.py"
+
+
+class _Tensor:
+    pass
+
+
+def _function(tree: ast.AST, name: str) -> ast.FunctionDef:
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name == name:
+            return node
+    raise AssertionError(f"function {name!r} is missing")
+
+
+def _load_extracted_dst_ptr_function(*, strict: bool = False):
+    source = SO2_PATH.read_text(encoding="utf-8")
+    tree = ast.parse(source)
+    function = _function(tree, "_dst_ptr_from_sorted")
+    module = ast.Module(body=[function], type_ignores=[])
+    ast.fix_missing_locations(module)
+    namespace: dict[str, Any] = {
+        "Any": Any,
+        "Tensor": _Tensor,
+        "runtime_policy": SimpleNamespace(
+            is_cute_strict_enabled=lambda: strict,
+        ),
+    }
+    exec(compile(module, str(SO2_PATH), "exec"), namespace)
+    return namespace["_dst_ptr_from_sorted"]
+
+
+def _load_extracted_sorted_metadata_function():
+    source = SO2_METADATA_PATH.read_text(encoding="utf-8")
+    tree = ast.parse(source)
+    function = _function(tree, "build_sorted_edge_index_metadata")
+    module = ast.Module(body=[function], type_ignores=[])
+    ast.fix_missing_locations(module)
+
+    def destination_row_ptr(dst, n_nodes):
+        boundaries = torch.arange(
+            n_nodes + 1,
+            device=dst.device,
+            dtype=dst.dtype,
+        )
+        return torch.searchsorted(dst, boundaries, out_int32=True).contiguous()
+
+    namespace = {
+        "torch": torch,
+        "_destination_row_ptr_op": destination_row_ptr,
+    }
+    exec(compile(module, str(SO2_METADATA_PATH), "exec"), namespace)
+    return namespace["build_sorted_edge_index_metadata"]
+
+
+def _load_sorted_metadata_function() -> Any:
+    module = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.so2.metadata")
+    return module.build_sorted_edge_index_metadata
+
+
+class _FakeDst:
+    device = "cuda:0"
+    dtype = "int64"
+
+    def __init__(self):
+        self.contiguous_calls = 0
+
+    def contiguous(self):
+        self.contiguous_calls += 1
+        return self
+
+
+class _RecordingTorch:
+    int64 = "int64"
+
+    def __init__(self):
+        self.calls: list[tuple[Any, ...]] = []
+        self.boundaries = object()
+        self.result = object()
+
+    def arange(self, stop, *, device, dtype):
+        self.calls.append(("arange", stop, device, dtype))
+        return self.boundaries
+
+    def searchsorted(self, sorted_sequence, values):
+        self.calls.append(("searchsorted", sorted_sequence, values))
+        return self.result
+
+    def bincount(self, *args, **kwargs):
+        raise AssertionError("CUDA bincount must not construct sorted dst_ptr")
+
+
+class TestSO2DstPtrExtracted(unittest.TestCase):
+    def test_sorted_path_operator_and_output_contract(self):
+        helper = _load_extracted_dst_ptr_function()
+        torch_module = _RecordingTorch()
+        dst = _FakeDst()
+
+        result = helper(
+            torch_module,
+            dst,
+            8,
+            destinations_sorted=True,
+        )
+
+        self.assertIs(result, torch_module.result)
+        self.assertEqual(dst.contiguous_calls, 1)
+        self.assertEqual(
+            torch_module.calls,
+            [
+                ("arange", 9, dst.device, torch_module.int64),
+                ("searchsorted", dst, torch_module.boundaries),
+            ],
+        )
+
+    def test_unsorted_path_returns_before_tensor_work(self):
+        helper = _load_extracted_dst_ptr_function()
+        torch_module = _RecordingTorch()
+
+        self.assertIsNone(
+            helper(
+                torch_module,
+                _FakeDst(),
+                8,
+                destinations_sorted=False,
+            )
+        )
+        self.assertEqual(torch_module.calls, [])
+
+
+@unittest.skipIf(torch is None, "destination-pointer differential requires PyTorch")
+class TestSO2DstPtrTorch(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.helper = staticmethod(_load_extracted_dst_ptr_function())
+
+    def test_matches_bincount_reference_and_preserves_layout(self):
+        assert torch is not None
+        devices = [torch.device("cpu")]
+        if torch.cuda.is_available():
+            devices.append(torch.device("cuda"))
+
+        cases = (
+            (1, []),
+            (1, [0, 0, 0]),
+            (8, [0, 0, 2, 5, 5, 5, 7]),
+            (9, [1, 1, 1, 4, 8]),
+            (17, [0, 3, 3, 7, 7, 7, 15]),
+        )
+        for device in devices:
+            for n_node, values in cases:
+                with self.subTest(device=device, n_node=n_node, values=values):
+                    dst = torch.tensor(values, device=device, dtype=torch.int64)
+                    counts = torch.bincount(dst, minlength=n_node)
+                    expected = torch.empty(
+                        n_node + 1,
+                        device=device,
+                        dtype=torch.int64,
+                    )
+                    expected[0] = 0
+                    expected[1:] = counts.cumsum(0)
+
+                    actual = self.helper(
+                        torch,
+                        dst,
+                        n_node,
+                        destinations_sorted=True,
+                    )
+
+                    self.assertTrue(torch.equal(actual, expected))
+                    self.assertEqual(actual.dtype, torch.int64)
+                    self.assertEqual(actual.device, dst.device)
+                    self.assertTrue(actual.is_contiguous())
+                    self.assertEqual(actual.stride(), (1,))
+
+    def test_strict_mode_rejects_false_sorted_provenance(self):
+        assert torch is not None
+        helper = _load_extracted_dst_ptr_function(strict=True)
+        dst = torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cpu")
+
+        with self.assertRaisesRegex(RuntimeError, "monotonically nondecreasing"):
+            helper(
+                torch,
+                dst,
+                4,
+                destinations_sorted=True,
+            )
+
+    def test_strict_mode_accepts_duplicate_sorted_destinations(self):
+        assert torch is not None
+        helper = _load_extracted_dst_ptr_function(strict=True)
+        dst = torch.tensor([0, 0, 2, 3], dtype=torch.int64, device="cpu")
+
+        actual = helper(
+            torch,
+            dst,
+            4,
+            destinations_sorted=True,
+        )
+
+        self.assertTrue(
+            torch.equal(
+                actual,
+                torch.tensor([0, 2, 2, 3, 4], device="cpu"),
+            )
+        )
+
+    def test_fullgraph_dynamic_compile_preserves_pointer_contract(self):
+        assert torch is not None
+        helper = self.helper
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+            backend = "inductor"
+        else:
+            device = torch.device("cpu")
+            backend = "aot_eager"
+
+        def build_ptr(dst):
+            return helper(
+                torch,
+                dst,
+                8,
+                destinations_sorted=True,
+            )
+
+        compiled = torch.compile(
+            build_ptr,
+            backend=backend,
+            dynamic=True,
+            fullgraph=True,
+        )
+        dst = torch.tensor(
+            [0, 0, 2, 5, 5, 5, 7],
+            device=device,
+            dtype=torch.int64,
+        )
+        actual = compiled(dst)
+        expected = torch.tensor(
+            [0, 2, 2, 3, 3, 3, 6, 6, 7],
+            device=device,
+            dtype=torch.int64,
+        )
+
+        self.assertTrue(torch.equal(actual, expected))
+        self.assertEqual(actual.stride(), (1,))
+
+    @unittest.skipUnless(
+        torch is not None and torch.cuda.is_available(),
+        "SO2 custom-op registry lifetime regression requires CUDA",
+    )
+    def test_compile_cold_registration_survives_into_packed_custom_op_runtime(self):
+        assert torch is not None
+        so2 = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.so2.operation")
+        prior_registry = dict(so2._REGISTRY)
+        prior_next_handle = next(so2._HANDLE_COUNTER)
+        so2._HANDLE_COUNTER = count(prior_next_handle)
+
+        def cleanup():
+            torch._dynamo.reset()
+            so2._PACKED_RUNNER_CACHE.clear()
+            so2._REGISTRY.clear()
+            so2._REGISTRY.update(prior_registry)
+            so2._HANDLE_COUNTER = count(prior_next_handle)
+
+        self.addCleanup(cleanup)
+        so2._REGISTRY.clear()
+        so2._PACKED_RUNNER_CACHE.clear()
+        torch._dynamo.reset()
+        block = torch.nn.Module()
+        config = so2.NeoSO2RuntimeConfig()
+        runtime_handles = []
+
+        class FakeRunner:
+            pass
+
+        def fake_build(handle, x_arg, *args):
+            del args
+            entry = so2._REGISTRY[int(handle)]
+            self.assertIs(entry.block, block)
+            runtime_handles.append(int(handle))
+            runner = FakeRunner()
+            runner.final = x_arg.detach().clone()
+            return runner
+
+        def invoke(
+            x_arg,
+            d_arg,
+            dt_arg,
+            radial_arg,
+            edge_arg,
+            src_arg,
+            dst_arg,
+        ):
+            state = getattr(block, "_deepmd_cute_so2_state", None)
+            if state is None:
+                state = so2._register_cute_so2_state(
+                    block,
+                    x_arg.device.index,
+                    config,
+                )
+            dst_ptr = so2._dst_ptr_from_sorted(
+                torch,
+                dst_arg,
+                x_arg.shape[0],
+                destinations_sorted=True,
+            )
+            edge_src_gate = edge_arg.new_empty((0,))
+            return so2.cute_so2(
+                state.handle,
+                x_arg,
+                d_arg,
+                dt_arg,
+                radial_arg,
+                edge_arg,
+                src_arg,
+                dst_arg,
+                dst_ptr,
+                edge_src_gate,
+            )
+
+        disabled_invoke = torch.compiler.disable(invoke)
+
+        def compiled_entry(*args):
+            return disabled_invoke(*args)
+
+        device = torch.device("cuda:0")
+        x = torch.randn(2, 16, 1, 32, device=device)
+        d_full = torch.randn(3, 46, device=device)
+        dt_full = torch.randn_like(d_full)
+        radial = torch.randn(3, 4, 32, device=device)
+        edge_env = torch.ones(3, 1, device=device)
+        src = torch.tensor((0, 1, 0), dtype=torch.int64, device=device)
+        dst = torch.tensor((0, 0, 1), dtype=torch.int64, device=device)
+
+        with mock.patch.object(so2, "_build_runner", new=fake_build):
+            compiled = torch.compile(compiled_entry, backend="eager", dynamic=True)
+            first = compiled(x, d_full, dt_full, radial, edge_env, src, dst)
+            second = compiled(x, d_full, dt_full, radial, edge_env, src, dst)
+
+        state = block._deepmd_cute_so2_state
+        self.assertEqual(runtime_handles, [state.handle, state.handle])
+        self.assertEqual(list(so2._REGISTRY), [state.handle])
+        self.assertIs(so2._REGISTRY[state.handle].block, block)
+        torch.testing.assert_close(first, x, rtol=0.0, atol=0.0)
+        torch.testing.assert_close(second, x, rtol=0.0, atol=0.0)
+
+
+@unittest.skipIf(torch is None, "sorted edge metadata requires PyTorch")
+class TestSortedEdgeIndexMetadata(unittest.TestCase):
+    @staticmethod
+    def _cache(src, dst, node_count):
+        assert torch is not None
+        edge_cache = importlib.import_module(
+            "deepmd.pt.model.descriptor.sezm_nn.edge_cache"
+        )
+        device = torch.device("cpu")
+        src = src.to(device=device)
+        dst = dst.to(device=device)
+        edge_count = src.numel()
+        return edge_cache.EdgeFeatureCache(
+            src=src,
+            dst=dst,
+            edge_type_feat=torch.empty(edge_count, 1, device=device),
+            edge_vec=torch.empty(edge_count, 3, device=device),
+            edge_rbf=torch.empty(edge_count, 1, device=device),
+            edge_env=torch.ones(edge_count, 1, device=device),
+            deg=torch.zeros(node_count, device=device),
+            inv_sqrt_deg=torch.ones(node_count, 1, 1, device=device),
+            destinations_sorted=True,
+        )
+
+    def test_builds_destination_and_indirect_source_csr(self):
+        assert torch is not None
+        builder = _load_sorted_metadata_function()
+        src = torch.tensor(
+            [2, 0, 3, 0, 1, 2],
+            dtype=torch.int64,
+            device="cpu",
+        )
+        dst = torch.tensor(
+            [0, 0, 1, 2, 2, 3],
+            dtype=torch.int64,
+            device="cpu",
+        )
+        dst_ptr, source_order, source_ptr = builder(src, dst, 4)
+
+        self.assertEqual(dst_ptr.dtype, torch.int32)
+        self.assertEqual(source_order.dtype, torch.int32)
+        self.assertEqual(source_ptr.dtype, torch.int32)
+        torch.testing.assert_close(
+            dst_ptr,
+            torch.tensor(
+                [0, 2, 3, 5, 6],
+                dtype=torch.int32,
+                device="cpu",
+            ),
+            rtol=0.0,
+            atol=0.0,
+        )
+        torch.testing.assert_close(
+            src.index_select(0, source_order.to(torch.int64)),
+            torch.tensor(
+                [0, 0, 1, 2, 2, 3],
+                dtype=torch.int64,
+                device="cpu",
+            ),
+            rtol=0.0,
+            atol=0.0,
+        )
+        torch.testing.assert_close(
+            source_ptr,
+            torch.tensor(
+                [0, 2, 3, 5, 6],
+                dtype=torch.int32,
+                device="cpu",
+            ),
+            rtol=0.0,
+            atol=0.0,
+        )
+
+    def test_empty_edges_build_valid_zero_csr(self):
+        assert torch is not None
+        builder = _load_sorted_metadata_function()
+        dst_ptr, source_order, source_ptr = builder(
+            torch.empty(0, dtype=torch.int64, device="cpu"),
+            torch.empty(0, dtype=torch.int64, device="cpu"),
+            4,
+        )
+
+        torch.testing.assert_close(
+            dst_ptr,
+            torch.zeros(5, dtype=torch.int32, device="cpu"),
+            rtol=0.0,
+            atol=0.0,
+        )
+        self.assertEqual(source_order.numel(), 0)
+        torch.testing.assert_close(
+            source_ptr,
+            torch.zeros(5, dtype=torch.int32, device="cpu"),
+            rtol=0.0,
+            atol=0.0,
+        )
+
+    def test_strict_metadata_rejects_unsorted_destinations(self):
+        assert torch is not None
+        builder = _load_extracted_sorted_metadata_function()
+        src = torch.tensor([0, 1, 2, 3], dtype=torch.int64, device="cpu")
+        dst = torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cpu")
+
+        with self.assertRaisesRegex(RuntimeError, "monotonically nondecreasing"):
+            builder(
+                src,
+                dst,
+                4,
+                validate_sorted=True,
+            )
+
+    def test_strict_metadata_builder_is_fullgraph_compilable(self):
+        assert torch is not None
+        builder = _load_extracted_sorted_metadata_function()
+        dynamo = getattr(torch, "_dynamo", None)
+        if dynamo is None:
+            self.skipTest("torch._dynamo is unavailable")
+
+        def build_ptr(src, dst):
+            dst_ptr, _, _ = builder(
+                src,
+                dst,
+                4,
+                validate_sorted=True,
+            )
+            return dst_ptr
+
+        compiled = torch.compile(
+            build_ptr,
+            backend="eager",
+            dynamic=True,
+            fullgraph=True,
+        )
+        src = torch.tensor([2, 0, 1], dtype=torch.int64, device="cpu")
+        dst = torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu")
+        actual = compiled(src, dst)
+
+        torch.testing.assert_close(
+            actual,
+            torch.tensor([0, 1, 2, 3, 3], dtype=torch.int32, device="cpu"),
+            rtol=0.0,
+            atol=0.0,
+        )
+
+    @unittest.skipUnless(
+        torch is not None and torch.cuda.is_available(),
+        "offset destination compile regression requires CUDA",
+    )
+    def test_inductor_handles_offset_destination_view(self):
+        assert torch is not None
+        builder = _load_sorted_metadata_function()
+
+        def build_metadata(edge_index):
+            return builder(edge_index[0], edge_index[1], 4)
+
+        compiled = torch.compile(
+            build_metadata,
+            backend="inductor",
+            dynamic=True,
+            fullgraph=True,
+        )
+        edge_index = torch.tensor(
+            [
+                [2, 0, 3, 0, 1, 2],
+                [0, 0, 1, 2, 2, 3],
+            ],
+            dtype=torch.int64,
+            device="cuda",
+        )
+        self.assertGreater(edge_index[1].storage_offset(), 0)
+
+        dst_ptr, source_order, source_ptr = compiled(edge_index)
+        torch.testing.assert_close(
+            dst_ptr,
+            torch.tensor([0, 2, 3, 5, 6], dtype=torch.int32, device="cuda"),
+            rtol=0.0,
+            atol=0.0,
+        )
+        torch.testing.assert_close(
+            edge_index[0].index_select(0, source_order.to(torch.int64)),
+            torch.tensor([0, 0, 1, 2, 2, 3], device="cuda"),
+            rtol=0.0,
+            atol=0.0,
+        )
+        torch.testing.assert_close(
+            source_ptr,
+            torch.tensor([0, 2, 3, 5, 6], dtype=torch.int32, device="cuda"),
+            rtol=0.0,
+            atol=0.0,
+        )
+
+    def test_dynamic_node_and_edge_counts_build_independent_metadata(self):
+        assert torch is not None
+        builder = _load_sorted_metadata_function()
+        first = builder(
+            torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"),
+            torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"),
+            4,
+        )
+        second = builder(
+            torch.tensor(
+                [0, 2, 4, 1, 3],
+                dtype=torch.int64,
+                device="cpu",
+            ),
+            torch.tensor(
+                [0, 0, 2, 4, 5],
+                dtype=torch.int64,
+                device="cpu",
+            ),
+            6,
+        )
+
+        torch.testing.assert_close(
+            first[0],
+            torch.tensor([0, 1, 2, 3, 3], dtype=torch.int32, device="cpu"),
+            rtol=0.0,
+            atol=0.0,
+        )
+        torch.testing.assert_close(
+            second[0],
+            torch.tensor(
+                [0, 2, 2, 3, 3, 4, 5],
+                dtype=torch.int32,
+                device="cpu",
+            ),
+            rtol=0.0,
+            atol=0.0,
+        )
+        self.assertNotEqual(first[0].data_ptr(), second[0].data_ptr())
+        self.assertNotEqual(first[1].data_ptr(), second[1].data_ptr())
+        self.assertNotEqual(first[2].data_ptr(), second[2].data_ptr())
+
+    def test_explicit_tensor_metadata_survives_disabled_so2_boundary(self):
+        assert torch is not None
+        builder = _load_sorted_metadata_function()
+        dynamo = getattr(torch, "_dynamo", None)
+        if dynamo is None:
+            self.skipTest("torch._dynamo is unavailable")
+
+        def so2_boundary(
+            x,
+            cache,
+            radial,
+            dst_ptr,
+            source_order,
+            source_ptr,
+        ):
+            del radial
+            assert dst_ptr is not None
+            assert source_order is not None
+            assert source_ptr is not None
+            return (
+                x
+                + cache.edge_env.sum()
+                + dst_ptr.sum()
+                + source_order.sum()
+                + source_ptr.sum()
+            )
+
+        disabled_so2_boundary = torch.compiler.disable(so2_boundary)
+
+        def consume(cache, x, radial, dst_ptr, source_order, source_ptr):
+            before_break = x + cache.edge_env.sum()
+            return disabled_so2_boundary(
+                before_break,
+                cache,
+                radial,
+                dst_ptr,
+                source_order,
+                source_ptr,
+            )
+
+        compiled = torch.compile(consume, backend="eager", dynamic=True)
+        cases = (
+            (
+                torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"),
+                torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"),
+                4,
+            ),
+            (
+                torch.tensor(
+                    [0, 2, 4, 1, 3],
+                    dtype=torch.int64,
+                    device="cpu",
+                ),
+                torch.tensor(
+                    [0, 0, 2, 4, 5],
+                    dtype=torch.int64,
+                    device="cpu",
+                ),
+                6,
+            ),
+        )
+        for src, dst, node_count in cases:
+            cache = self._cache(src, dst, node_count)
+            dst_ptr, source_order, source_ptr = builder(
+                src,
+                dst,
+                node_count,
+            )
+            radial = torch.ones(src.numel(), 4, 1, device="cpu")
+            args = (
+                cache,
+                torch.ones(1, device="cpu"),
+                radial,
+                dst_ptr,
+                source_order,
+                source_ptr,
+            )
+            eager = consume(*args)
+            actual = compiled(*args)
+            torch.testing.assert_close(
+                actual,
+                eager,
+                rtol=0.0,
+                atol=0.0,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/model/sezm_cute/test_so2_gate_structural.py b/source/tests/pt/model/sezm_cute/test_so2_gate_structural.py
new file mode 100644
index 0000000000..d8ef04185f
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_gate_structural.py
@@ -0,0 +1,96 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Numerical differentials for the structural Neo SO2 gate."""
+
+from __future__ import (
+    annotations,
+)
+
+import unittest
+
+import torch
+
+from deepmd.pt_expt.kernels.cute.sezm.so2 import structural_gate as so2_gate_structural
+
+
+class TestStructuralGateHelpers(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.helper = so2_gate_structural
+
+    def setUp(self):
+        torch.manual_seed(20260703)
+        self.edge_count = 7
+        self.gate_src = torch.randn(
+            self.edge_count, 2, 32, dtype=torch.float32, device="cpu"
+        )
+        self.weight = torch.randn(32, 2 * 3 * 32, dtype=torch.float32, device="cpu")
+
+    def test_focus_major_forward_matches_focus_linear(self):
+        actual = self.helper.focus_major_gate_linear_forward(
+            self.gate_src,
+            self.weight,
+        )
+        expected = torch.einsum(
+            "efi,ifo->efo",
+            self.gate_src,
+            self.weight.view(32, 2, 3 * 32),
+        )
+        self.assertEqual(actual.shape, (2, self.edge_count, 3 * 32))
+        self.assertEqual(actual.stride(), (self.edge_count * 3 * 32, 3 * 32, 1))
+        torch.testing.assert_close(
+            actual.permute(1, 0, 2), expected, atol=5e-5, rtol=5e-5
+        )
+
+    def test_backward_addmm_accumulates_without_replacing_grad_y_storage(self):
+        grad_y = torch.randn(
+            self.edge_count, 2, 10, 32, dtype=torch.float32, device="cpu"
+        )
+        grad_logits = torch.randn(
+            2, self.edge_count, 3 * 32, dtype=torch.float32, device="cpu"
+        )
+        expected = grad_y.clone()
+        weight = self.weight.view(32, 2, 3 * 32)
+        for focus in range(2):
+            expected[:, focus, 0, :] += grad_logits[focus] @ weight[:, focus, :].T
+
+        pointer = grad_y.untyped_storage().data_ptr()
+        actual = self.helper.focus_major_gate_linear_backward_add_(
+            grad_y,
+            grad_logits,
+            self.weight,
+        )
+
+        self.assertIs(actual, grad_y)
+        self.assertEqual(actual.untyped_storage().data_ptr(), pointer)
+        torch.testing.assert_close(actual, expected, atol=5e-5, rtol=5e-5)
+
+    def test_forward_wrapper_preserves_caller_owned_storage(self):
+        residual = torch.randn(
+            self.edge_count, 2, 10, 32, dtype=torch.float32, device="cpu"
+        )
+        y = torch.randn_like(residual)
+        logits = torch.randn(
+            2, self.edge_count, 3 * 32, dtype=torch.float32, device="cpu"
+        )
+        pointer = residual.untyped_storage().data_ptr()
+
+        def fake_kernel(residual_flat, y_flat, logits_arg, out_flat):
+            self.assertEqual(residual_flat.data_ptr(), out_flat.data_ptr())
+            self.assertIs(logits_arg, logits)
+            out_flat.copy_(residual_flat + y_flat)
+
+        actual = self.helper.run_structural_gate_forward(
+            fake_kernel,
+            residual,
+            y,
+            logits,
+            out=residual,
+        )
+
+        self.assertIs(actual, residual)
+        self.assertEqual(actual.untyped_storage().data_ptr(), pointer)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/model/sezm_cute/test_so2_inplace_residual.py b/source/tests/pt/model/sezm_cute/test_so2_inplace_residual.py
new file mode 100644
index 0000000000..65e81cd758
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_inplace_residual.py
@@ -0,0 +1,364 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Contracts for the in-place CuTe SO2 residual/addmm helper."""
+
+from __future__ import (
+    annotations,
+)
+
+import unittest
+
+import torch
+from torch.utils._python_dispatch import (
+    TorchDispatchMode,
+)
+
+from deepmd.pt_expt.kernels.cute.sezm.so2 import linear as so2_linear
+
+
+def _load_candidate():
+    return so2_linear
+
+
+def _out_of_place_reference(grad_out, residual, w0_t, wpair_t):
+    edge_count = grad_out.shape[0]
+    grad_flat = grad_out.view(edge_count, 2, 10 * 32)
+    residual_flat = residual.view(edge_count, 2, 10 * 32)
+    out = torch.empty_like(residual)
+    out_flat = out.view(edge_count, 2, 10 * 32)
+    for focus in range(2):
+        torch.addmm(
+            residual_flat[:, focus, : 4 * 32],
+            grad_flat[:, focus, : 4 * 32],
+            w0_t[focus],
+            out=out_flat[:, focus, : 4 * 32],
+        )
+        torch.addmm(
+            residual_flat[:, focus, 4 * 32 :],
+            grad_flat[:, focus, 4 * 32 :],
+            wpair_t[focus],
+            out=out_flat[:, focus, 4 * 32 :],
+        )
+    return out
+
+
+class _DispatchRecorder(TorchDispatchMode):
+    def __init__(self):
+        super().__init__()
+        self.calls = []
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        kwargs = kwargs or {}
+        self.calls.append((func, args, kwargs))
+        return func(*args, **kwargs)
+
+
+class TestSO2InplaceResidual(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(20260630)
+        self.edge_count = 7
+        self.grad_out = torch.randn(
+            self.edge_count,
+            2,
+            10,
+            32,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        self.residual = torch.randn_like(self.grad_out)
+        self.w0_t = torch.randn(
+            2,
+            4 * 32,
+            4 * 32,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        self.wpair_t = torch.randn(
+            2,
+            6 * 32,
+            6 * 32,
+            dtype=torch.float32,
+            device="cpu",
+        )
+
+    def test_inplace_result_matches_out_of_place_equations(self):
+        candidate = _load_candidate()
+        expected = _out_of_place_reference(
+            self.grad_out,
+            self.residual,
+            self.w0_t,
+            self.wpair_t,
+        )
+        residual = self.residual.clone()
+        storage_ptr = residual.untyped_storage().data_ptr()
+
+        actual = candidate.neo_so2_linear_backward_residual_inplace(
+            residual,
+            self.grad_out,
+            self.w0_t,
+            self.wpair_t,
+        )
+
+        self.assertIs(actual, residual)
+        self.assertEqual(actual.untyped_storage().data_ptr(), storage_ptr)
+        self.assertEqual(actual.stride(), (2 * 10 * 32, 10 * 32, 32, 1))
+        torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
+
+    def test_only_four_inplace_addmm_ops_cross_the_dispatch_boundary(self):
+        candidate = _load_candidate()
+        residual = self.residual.clone()
+        recorder = _DispatchRecorder()
+
+        with recorder:
+            candidate.neo_so2_linear_backward_residual_inplace(
+                residual,
+                self.grad_out,
+                self.w0_t,
+                self.wpair_t,
+            )
+
+        addmm_calls = [
+            args
+            for func, args, _kwargs in recorder.calls
+            if func is torch.ops.aten.addmm_.default
+        ]
+        self.assertEqual(len(addmm_calls), 4)
+        forbidden = {
+            torch.ops.aten.clone.default,
+            torch.ops.aten.contiguous.default,
+            torch.ops.aten.copy_.default,
+        }
+        self.assertFalse(
+            any(func in forbidden for func, _args, _kwargs in recorder.calls)
+        )
+
+        for residual_block, grad_block, weight in addmm_calls:
+            width = residual_block.shape[1]
+            self.assertIn(width, (4 * 32, 6 * 32))
+            self.assertEqual(residual_block.stride(), (2 * 10 * 32, 1))
+            self.assertEqual(grad_block.stride(), (2 * 10 * 32, 1))
+            self.assertEqual(weight.stride(), (width, 1))
+            self.assertTrue(candidate._has_direct_cublas_layout(residual_block))
+            self.assertTrue(candidate._has_direct_cublas_layout(grad_block))
+            self.assertTrue(candidate._has_direct_cublas_layout(weight))
+
+    def test_layout_predicate_and_fp32_scope_are_explicit(self):
+        candidate = _load_candidate()
+
+        self.assertFalse(hasattr(candidate, "is_cublas_compatible_matrix"))
+        self.assertIn("PyTorch 2.10", candidate._has_direct_cublas_layout.__doc__)
+        self.assertIn("dtype/device", candidate._has_direct_cublas_layout.__doc__)
+        self.assertIn("highest", candidate.__doc__)
+        self.assertIn("TF32", candidate.__doc__)
+
+    def test_meta_execution_preserves_outer_shape_and_stride_contract(self):
+        candidate = _load_candidate()
+        residual = torch.empty(11, 2, 10, 32, dtype=torch.float32, device="meta")
+        grad_out = torch.empty_like(residual)
+        w0_t = torch.empty(2, 4 * 32, 4 * 32, dtype=torch.float32, device="meta")
+        wpair_t = torch.empty(
+            2,
+            6 * 32,
+            6 * 32,
+            dtype=torch.float32,
+            device="meta",
+        )
+
+        result = candidate.neo_so2_linear_backward_residual_inplace(
+            residual,
+            grad_out,
+            w0_t,
+            wpair_t,
+        )
+
+        self.assertIs(result, residual)
+        self.assertEqual(result.shape, (11, 2, 10, 32))
+        self.assertEqual(result.stride(), (2 * 10 * 32, 10 * 32, 32, 1))
+
+    def test_aliasing_grad_out_is_rejected_for_final_layer_safety(self):
+        candidate = _load_candidate()
+        shared = self.residual.clone()
+
+        with self.assertRaisesRegex(ValueError, "must not alias"):
+            candidate.neo_so2_linear_backward_residual_inplace(
+                shared,
+                shared,
+                self.w0_t,
+                self.wpair_t,
+            )
+
+    def test_exact_cross_storage_grad_out_alias_is_rejected(self):
+        candidate = _load_candidate()
+        byte_count = self.residual.numel() * self.residual.element_size()
+        backing = bytearray(byte_count)
+        residual = torch.frombuffer(
+            backing,
+            dtype=torch.float32,
+            count=self.residual.numel(),
+        ).view_as(self.residual)
+        grad_out = torch.frombuffer(
+            backing,
+            dtype=torch.float32,
+            count=self.grad_out.numel(),
+        ).view_as(self.grad_out)
+        self.assertNotEqual(
+            residual.untyped_storage()._cdata,
+            grad_out.untyped_storage()._cdata,
+        )
+        self.assertEqual(residual.data_ptr(), grad_out.data_ptr())
+        self.assertFalse(torch._C._overlaps(residual, grad_out))
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "residual and grad_out must not alias",
+        ):
+            candidate.neo_so2_linear_backward_residual_inplace(
+                residual,
+                grad_out,
+                self.w0_t,
+                self.wpair_t,
+            )
+
+    def test_residual_overlapping_w0_is_rejected(self):
+        candidate = _load_candidate()
+        storage = torch.randn(
+            2 * 4 * 32 * 4 * 32,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        residual = storage[: self.residual.numel()].view_as(self.residual)
+        w0_t = storage.view(2, 4 * 32, 4 * 32)
+        self.assertTrue(residual.is_contiguous())
+        self.assertTrue(w0_t.is_contiguous())
+        self.assertTrue(torch._C._overlaps(residual, w0_t))
+
+        with self.assertRaisesRegex(ValueError, "residual and w0_t must not alias"):
+            candidate.neo_so2_linear_backward_residual_inplace(
+                residual,
+                self.grad_out,
+                w0_t,
+                self.wpair_t,
+            )
+
+    def test_partial_cross_storage_w0_alias_is_rejected(self):
+        candidate = _load_candidate()
+        overlap_elements = 17
+        weight_elements = self.w0_t.numel()
+        weight_offset = self.residual.numel() - overlap_elements
+        backing = bytearray(
+            (weight_offset + weight_elements) * self.residual.element_size()
+        )
+        residual = torch.frombuffer(
+            backing,
+            dtype=torch.float32,
+            count=self.residual.numel(),
+        ).view_as(self.residual)
+        w0_t = torch.frombuffer(
+            backing,
+            dtype=torch.float32,
+            count=weight_elements,
+            offset=weight_offset * self.residual.element_size(),
+        ).view_as(self.w0_t)
+        self.assertNotEqual(
+            residual.untyped_storage()._cdata,
+            w0_t.untyped_storage()._cdata,
+        )
+        self.assertNotEqual(residual.data_ptr(), w0_t.data_ptr())
+        self.assertFalse(torch._C._overlaps(residual, w0_t))
+
+        with self.assertRaisesRegex(ValueError, "residual and w0_t must not alias"):
+            candidate.neo_so2_linear_backward_residual_inplace(
+                residual,
+                self.grad_out,
+                w0_t,
+                self.wpair_t,
+            )
+
+    def test_residual_overlapping_wpair_is_rejected(self):
+        candidate = _load_candidate()
+        storage = torch.randn(
+            2 * 6 * 32 * 6 * 32,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        residual = storage[: self.residual.numel()].view_as(self.residual)
+        wpair_t = storage.view(2, 6 * 32, 6 * 32)
+        self.assertTrue(residual.is_contiguous())
+        self.assertTrue(wpair_t.is_contiguous())
+        self.assertTrue(torch._C._overlaps(residual, wpair_t))
+
+        with self.assertRaisesRegex(ValueError, "residual and wpair_t must not alias"):
+            candidate.neo_so2_linear_backward_residual_inplace(
+                residual,
+                self.grad_out,
+                self.w0_t,
+                wpair_t,
+            )
+
+    def test_partial_cross_storage_wpair_alias_is_rejected(self):
+        candidate = _load_candidate()
+        overlap_elements = 17
+        weight_elements = self.wpair_t.numel()
+        weight_offset = self.residual.numel() - overlap_elements
+        backing = bytearray(
+            (weight_offset + weight_elements) * self.residual.element_size()
+        )
+        residual = torch.frombuffer(
+            backing,
+            dtype=torch.float32,
+            count=self.residual.numel(),
+        ).view_as(self.residual)
+        wpair_t = torch.frombuffer(
+            backing,
+            dtype=torch.float32,
+            count=weight_elements,
+            offset=weight_offset * self.residual.element_size(),
+        ).view_as(self.wpair_t)
+        self.assertNotEqual(
+            residual.untyped_storage()._cdata,
+            wpair_t.untyped_storage()._cdata,
+        )
+        self.assertNotEqual(residual.data_ptr(), wpair_t.data_ptr())
+        self.assertFalse(torch._C._overlaps(residual, wpair_t))
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "residual and wpair_t must not alias",
+        ):
+            candidate.neo_so2_linear_backward_residual_inplace(
+                residual,
+                self.grad_out,
+                self.w0_t,
+                wpair_t,
+            )
+
+    def test_non_fp32_and_noncanonical_layouts_are_rejected(self):
+        candidate = _load_candidate()
+        with self.assertRaisesRegex(TypeError, "float32"):
+            candidate.neo_so2_linear_backward_residual_inplace(
+                self.residual.double(),
+                self.grad_out.double(),
+                self.w0_t.double(),
+                self.wpair_t.double(),
+            )
+
+        residual = torch.empty(
+            2,
+            self.edge_count,
+            10,
+            32,
+            dtype=torch.float32,
+            device="cpu",
+        ).transpose(0, 1)
+        self.assertEqual(residual.shape, self.residual.shape)
+        with self.assertRaisesRegex(ValueError, "contiguous"):
+            candidate.neo_so2_linear_backward_residual_inplace(
+                residual,
+                self.grad_out,
+                self.w0_t,
+                self.wpair_t,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/model/sezm_cute/test_so2_manual_adjoints.py b/source/tests/pt/model/sezm_cute/test_so2_manual_adjoints.py
new file mode 100644
index 0000000000..f1692be871
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_manual_adjoints.py
@@ -0,0 +1,365 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Parity tests for the real-module adjoints used by Neo SO2 backward."""
+
+from __future__ import (
+    annotations,
+)
+
+import importlib
+from types import (
+    SimpleNamespace,
+)
+
+import pytest
+import torch
+
+from deepmd.pt.model.descriptor.sezm_nn.activation import (
+    SwiGLU,
+)
+from deepmd.pt.model.descriptor.sezm_nn.grid_net import (
+    SO3GridNet,
+)
+from deepmd.pt.model.descriptor.sezm_nn.norm import (
+    EquivariantRMSNorm,
+)
+from deepmd.pt.model.descriptor.sezm_nn.so3 import (
+    FocusLinear,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2
+
+
+@pytest.fixture(autouse=True)
+def _construct_modules_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(env, "DEVICE", torch.device("cpu"))
+
+
+def _randn(*shape: int, requires_grad: bool = False) -> torch.Tensor:
+    return torch.randn(
+        *shape,
+        dtype=torch.float64,
+        device="cpu",
+        requires_grad=requires_grad,
+    )
+
+
+def test_equivariant_rmsnorm_manual_adjoint_matches_real_module() -> None:
+    norm = EquivariantRMSNorm(
+        lmax=2,
+        channels=4,
+        n_focus=2,
+        eps=2.0e-6,
+        dtype=torch.float64,
+        trainable=False,
+    )
+    with torch.no_grad():
+        norm.adam_scale.copy_(
+            torch.linspace(
+                0.5,
+                1.5,
+                norm.adam_scale.numel(),
+                dtype=torch.float64,
+                device="cpu",
+            ).reshape_as(norm.adam_scale)
+        )
+        norm.bias.copy_(
+            torch.linspace(
+                -0.2,
+                0.2,
+                norm.bias.numel(),
+                dtype=torch.float64,
+                device="cpu",
+            ).reshape_as(norm.bias)
+        )
+    x = _randn(3, 9, 2, 4, requires_grad=True)
+    grad_out = _randn(*x.shape)
+
+    reference_grad = torch.autograd.grad(norm(x), x, grad_out)[0]
+    actual_grad = so2._equivariant_rmsnorm_backward(
+        norm,
+        x.detach(),
+        grad_out,
+    )
+
+    torch.testing.assert_close(actual_grad, reference_grad, atol=1.0e-12, rtol=1.0e-12)
+
+
+def test_focus_linear_manual_adjoint_matches_real_module() -> None:
+    linear = FocusLinear(
+        in_channels=6,
+        out_channels=5,
+        n_focus=2,
+        dtype=torch.float64,
+        bias=True,
+        trainable=False,
+        seed=None,
+    )
+    x = _randn(7, 2, 6, requires_grad=True)
+    grad_out = _randn(7, 2, 5)
+
+    reference = linear(x)
+    reference_grad = torch.autograd.grad(reference, x, grad_out)[0]
+
+    torch.testing.assert_close(so2._focus_linear_forward(linear, x), reference)
+    torch.testing.assert_close(
+        so2._focus_linear_backward_input(linear, grad_out),
+        reference_grad,
+        atol=1.0e-12,
+        rtol=1.0e-12,
+    )
+
+
+def test_swiglu_manual_adjoint_matches_real_module() -> None:
+    activation = SwiGLU()
+    x = _randn(6, 3, 10, requires_grad=True)
+    grad_out = _randn(6, 3, 5)
+
+    reference = activation(x)
+    reference_grad = torch.autograd.grad(reference, x, grad_out)[0]
+
+    torch.testing.assert_close(so2._swiglu_forward(x), reference)
+    torch.testing.assert_close(
+        so2._swiglu_backward_input(x.detach(), grad_out),
+        reference_grad,
+        atol=1.0e-12,
+        rtol=1.0e-12,
+    )
+
+
+def test_grid_cross_glu_flat_manual_adjoint_matches_real_module() -> None:
+    net = SO3GridNet(
+        lmax=3,
+        kmax=1,
+        channels=4,
+        n_focus=2,
+        mode="cross",
+        op_type="glu",
+        dtype=torch.float64,
+        layout="flat",
+        coefficient_layout="packed",
+        residual_scale_init=0.25,
+        trainable=False,
+        seed=None,
+    )
+    coeff_dim = net.projector.coeff_dim // net.n_frames
+    query = _randn(3, coeff_dim, net.n_focus * net.channels, requires_grad=True)
+    context = _randn(3, coeff_dim, net.n_focus * net.channels, requires_grad=True)
+    grad_out = _randn(3, coeff_dim, net.n_focus * net.channels)
+
+    reference_query, reference_context = torch.autograd.grad(
+        net(query, context),
+        (query, context),
+        grad_out,
+    )
+    actual_query, actual_context = so2._so3_grid_cross_glu_flat_backward(
+        net,
+        query.detach(),
+        context.detach(),
+        grad_out,
+    )
+
+    torch.testing.assert_close(
+        actual_query,
+        reference_query,
+        atol=1.0e-11,
+        rtol=1.0e-11,
+    )
+    torch.testing.assert_close(
+        actual_context,
+        reference_context,
+        atol=1.0e-11,
+        rtol=1.0e-11,
+    )
+
+
+def _cute_cuda_runtime_available() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        importlib.import_module("cutlass.cute")
+        importlib.import_module("cuda.bindings.driver")
+    except Exception:  # pragma: no cover - runtime dependent
+        return False
+    return True
+
+
+@pytest.mark.skipif(
+    not _cute_cuda_runtime_available(),
+    reason="SM90 message-grid differential requires CUDA and CuTe DSL",
+)
+def test_sm90_message_grid_forward_and_adjoint_match_real_module() -> None:
+    if tuple(torch.cuda.get_device_capability()) != (9, 0):
+        pytest.skip("SM90 message-grid differential requires an SM90 GPU")
+
+    from deepmd.pt_expt.kernels.cute.sezm.so2.message_grid import (
+        run_packed_message_grid_forward,
+    )
+    from deepmd.pt_expt.kernels.cute.sezm.so2.sm90.message_grid_readout import (
+        prepare_sm90_message_grid_state,
+        run_sm90_message_grid_backward,
+    )
+
+    prior_precision = torch.get_float32_matmul_precision()
+    prior_tf32 = torch.backends.cuda.matmul.allow_tf32
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.set_float32_matmul_precision("highest")
+    try:
+        net = SO3GridNet(
+            lmax=3,
+            kmax=1,
+            channels=32,
+            n_focus=2,
+            mode="cross",
+            op_type="glu",
+            dtype=torch.float32,
+            layout="flat",
+            coefficient_layout="packed",
+            residual_scale_init=0.25,
+            trainable=False,
+            seed=None,
+        ).to("cuda")
+        generator = torch.Generator(device="cuda").manual_seed(20260816)
+        query = (
+            0.1
+            * torch.randn(
+                5,
+                16,
+                64,
+                generator=generator,
+                device="cuda",
+                dtype=torch.float32,
+            )
+        ).requires_grad_(True)
+        context = (
+            0.1
+            * torch.randn(
+                16,
+                5,
+                64,
+                generator=generator,
+                device="cuda",
+                dtype=torch.float32,
+            ).permute(1, 0, 2)
+        ).requires_grad_(True)
+        assert context.stride() == (64, 5 * 64, 1)
+        grad_out = torch.randn(
+            query.shape,
+            generator=generator,
+            device="cuda",
+            dtype=torch.float32,
+        )
+
+        reference = net(query, context)
+        reference_query, reference_context = torch.autograd.grad(
+            reference,
+            (query, context),
+            grad_out,
+        )
+        state = prepare_sm90_message_grid_state(net)
+        actual, product = run_packed_message_grid_forward(
+            net,
+            query.detach(),
+            context.detach(),
+            return_product=True,
+            sm90_state=state,
+        )
+        actual_query, actual_context = run_sm90_message_grid_backward(
+            net,
+            query.detach(),
+            context.detach(),
+            grad_out,
+            product,
+            state,
+        )
+        torch.cuda.synchronize()
+
+        torch.testing.assert_close(actual, reference, atol=5.0e-5, rtol=5.0e-5)
+        torch.testing.assert_close(
+            actual_query,
+            reference_query,
+            atol=5.0e-5,
+            rtol=5.0e-5,
+        )
+        torch.testing.assert_close(
+            actual_context,
+            reference_context,
+            atol=5.0e-5,
+            rtol=5.0e-5,
+        )
+    finally:
+        torch.backends.cuda.matmul.allow_tf32 = prior_tf32
+        torch.set_float32_matmul_precision(prior_precision)
+
+
+@pytest.mark.skipif(
+    not _cute_cuda_runtime_available(),
+    reason="Q/K CuTe adjoint regression requires CUDA and CuTe DSL",
+)
+def test_qk_manual_adjoint_matches_autograd_with_sparse_edges() -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.qk_edge import (
+        compile_neo_qk_edge_backward,
+        compile_neo_qk_node_input_adjoint,
+    )
+
+    torch.manual_seed(20260813)
+    device = torch.device("cuda", torch.cuda.current_device())
+    node_count = 7
+    src = torch.tensor([0, 1, 1], dtype=torch.int32, device=device)
+    dst = torch.tensor([2, 2, 3], dtype=torch.int32, device=device)
+    edge_count = src.numel()
+    eps = 1.0e-5
+    scale = 32.0**-0.5
+    x_wide = torch.randn(
+        node_count,
+        16,
+        64,
+        dtype=torch.float32,
+        device=device,
+        requires_grad=True,
+    )
+    q_weight = torch.randn(32, 2, 32, dtype=torch.float32, device=device)
+    k_weight = torch.randn_like(q_weight)
+    norm_scale = torch.randn(2, 32, dtype=torch.float32, device=device)
+    grad_logits = torch.randn(edge_count, 2, dtype=torch.float32, device=device)
+
+    x_l0 = x_wide[:, 0, :].reshape(node_count, 2, 32)
+    x_norm = (
+        x_l0
+        * torch.rsqrt(x_l0.square().mean(dim=-1, keepdim=True) + eps)
+        * norm_scale.unsqueeze(0)
+    )
+    q_node = torch.einsum("nfi,ifo->nfo", x_norm, q_weight)
+    k_node = torch.einsum("nfi,ifo->nfo", x_norm, k_weight)
+    logits = (q_node[dst.long()] * k_node[src.long()]).sum(dim=-1) * scale
+    reference = torch.autograd.grad(logits, x_wide, grad_logits)[0]
+    compile_identity = (
+        torch.cuda.current_device(),
+        *torch.cuda.get_device_capability(device),
+    )
+
+    runner = SimpleNamespace(
+        so2=SimpleNamespace(
+            attn_q_proj=SimpleNamespace(weight=q_weight.reshape(32, 64)),
+            attn_k_proj=SimpleNamespace(weight=k_weight.reshape(32, 64)),
+            attn_qk_norm=SimpleNamespace(adam_scale=norm_scale, eps=eps),
+        ),
+        node_count=node_count,
+        edge_count=edge_count,
+        x_wide=x_wide.detach(),
+        q_node=q_node.detach().contiguous(),
+        k_node=k_node.detach().contiguous(),
+        src_i32=src,
+        dst_i32=dst,
+        qk_edge_backward=compile_neo_qk_edge_backward(scale, compile_identity),
+        qk_node_input_adjoint=compile_neo_qk_node_input_adjoint(
+            eps,
+            compile_identity,
+        ),
+    )
+    actual = so2._qk_manual_backward(runner, grad_logits)
+    torch.cuda.synchronize(device)
+
+    torch.testing.assert_close(actual, reference, atol=5.0e-5, rtol=5.0e-5)
diff --git a/source/tests/pt/model/sezm_cute/test_so2_packed_impl.py b/source/tests/pt/model/sezm_cute/test_so2_packed_impl.py
new file mode 100644
index 0000000000..95099e323c
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_packed_impl.py
@@ -0,0 +1,106 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Implementation differentials for the live packed Neo Wigner path."""
+
+from __future__ import (
+    annotations,
+)
+
+import importlib
+
+import pytest
+import torch
+
+TOL = 5.0e-5
+
+
+def _cute_runtime_skip_reason() -> str | None:
+    if not torch.cuda.is_available():
+        return "packed implementation differentials require CUDA"
+    try:
+        importlib.import_module("cutlass.cute")
+        importlib.import_module("cuda.bindings.driver")
+    except Exception as exc:  # pragma: no cover - runtime dependent
+        return f"packed implementation differentials require CuTe DSL: {exc}"
+    return None
+
+
+def _randn(
+    *shape: int,
+    generator: torch.Generator,
+    scale: float = 0.1,
+) -> torch.Tensor:
+    return scale * torch.randn(
+        *shape,
+        generator=generator,
+        device="cuda",
+        dtype=torch.float32,
+    )
+
+
+def _dense_wigner(edge_count: int, generator: torch.Generator) -> torch.Tensor:
+    from deepmd.pt_expt.kernels.cute.sezm.so2 import wigner_layout as layout
+
+    dense = torch.zeros(
+        edge_count,
+        16,
+        16,
+        device="cuda",
+        dtype=torch.float32,
+    )
+    for start, stop in zip(
+        layout.FULL_BLOCK_OFFSETS[:-1],
+        layout.FULL_BLOCK_OFFSETS[1:],
+        strict=True,
+    ):
+        dense[:, start:stop, start:stop] = _randn(
+            edge_count,
+            stop - start,
+            stop - start,
+            generator=generator,
+        )
+    return dense
+
+
+_CUTE_SKIP_REASON = _cute_runtime_skip_reason()
+
+
+@pytest.mark.skipif(
+    _CUTE_SKIP_REASON is not None,
+    reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable",
+)
+def test_panel_native_quaternion_backward_matches_dense_torch():
+    from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+        WignerDCalculator,
+    )
+    from deepmd.pt_expt.kernels.cute.sezm import (
+        wignerd,
+    )
+    from deepmd.pt_expt.kernels.cute.sezm.so2 import wigner_layout as layout
+
+    generator = torch.Generator(device="cuda").manual_seed(20260702)
+    q_value = _randn(3, 4, generator=generator, scale=1.0)
+    q_value = q_value / q_value.norm(dim=-1, keepdim=True)
+    grad_panel = _randn(3, 46, generator=generator)
+
+    q_panel = q_value.detach().requires_grad_(True)
+    panel = wignerd._wignerd_panel_op(q_panel)
+    grad_q_panel = torch.autograd.grad((panel * grad_panel).sum(), q_panel)[0]
+
+    q_dense = q_value.detach().requires_grad_(True)
+    dense, dense_t = WignerDCalculator(lmax=3, dtype=q_dense.dtype).to("cuda")(q_dense)
+    entries = tuple(layout.iter_packed_entries())
+    rows = [entry.full_row for entry in entries]
+    cols = [entry.full_col for entry in entries]
+    dense_weight = torch.zeros_like(dense)
+    dense_weight[:, rows, cols] = grad_panel
+    dense_loss = (dense * dense_weight).sum() + 0.0 * dense_t.sum()
+    grad_q_dense = torch.autograd.grad(dense_loss, q_dense)[0]
+
+    torch.testing.assert_close(
+        panel,
+        dense[..., rows, cols],
+        rtol=TOL,
+        atol=TOL,
+    )
+    torch.testing.assert_close(grad_q_panel, grad_q_dense, rtol=TOL, atol=TOL)
diff --git a/source/tests/pt/model/sezm_cute/test_so2_phase_c_inplace_adjoint.py b/source/tests/pt/model/sezm_cute/test_so2_phase_c_inplace_adjoint.py
new file mode 100644
index 0000000000..657aad1ac9
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_phase_c_inplace_adjoint.py
@@ -0,0 +1,320 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Storage-contract tests for the exact in-place Phase-C adjoint."""
+
+from __future__ import (
+    annotations,
+)
+
+import itertools
+from dataclasses import (
+    fields,
+    replace,
+)
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability() not in {(8, 0), (9, 0)},
+    reason="in-place Phase-C tests require sm_80 or sm_90",
+)
+
+OUTPUT_FIELDS = (
+    "grad_stack",
+    "grad_wigner_dt",
+    "grad_logits",
+    "grad_edge",
+    "grad_z_partial",
+    "grad_z",
+    "grad_focus_src",
+)
+NONALIASED_OUTPUT_FIELDS = OUTPUT_FIELDS[1:]
+INPUT_FIELDS = (
+    "grad_out",
+    "stack",
+    "wigner_dt",
+    "alpha",
+    "focus_alpha",
+    "dst_ptr",
+    "rotate_inv_rescale",
+    "edge_gate",
+    "z_bias_raw",
+    "group_max",
+    "denom",
+    "focus_src",
+    "focus_weight",
+    "focus_scale",
+)
+REDZONE_ELEMENTS = 4
+SENTINEL = 937.25
+
+
+def _phase_c_api():
+    pytest.importorskip("cutlass")
+    from deepmd.pt_expt.kernels.cute.sezm.so2.phase_c import (
+        CuteNeoPhaseCBackwardLayout,
+        NeoPhaseCBackwardLayoutOutputs,
+    )
+
+    return CuteNeoPhaseCBackwardLayout, NeoPhaseCBackwardLayoutOutputs
+
+
+def _runner():
+    Runner, _ = _phase_c_api()
+    return Runner(
+        focus_eps=1.0e-8,
+        focus_tau=1.0,
+        focus_label_smoothing=0.0,
+    )
+
+
+def _inputs(dst_ptr_values: list[int]):
+    device = torch.device("cuda")
+    generator = torch.Generator(device=device).manual_seed(
+        20260722 + dst_ptr_values[-1]
+    )
+    node_count = len(dst_ptr_values) - 1
+    edge_count = dst_ptr_values[-1]
+    rand = lambda *shape: torch.randn(  # noqa: E731
+        *shape,
+        device=device,
+        dtype=torch.float32,
+        generator=generator,
+    )
+    return {
+        "grad_out": rand(node_count, 16, 64),
+        "stack": rand(edge_count, 2, 10, 32),
+        "wigner_dt": rand(edge_count, 46),
+        "alpha": torch.softmax(rand(edge_count, 2), dim=0),
+        "focus_alpha": torch.softmax(rand(edge_count, 2), dim=1),
+        "dst_ptr": torch.tensor(dst_ptr_values, device=device, dtype=torch.int32),
+        "rotate_inv_rescale": rand(16),
+        "edge_gate": rand(edge_count).abs(),
+        "z_bias_raw": rand(2),
+        "group_max": rand(node_count, 2),
+        "denom": rand(node_count, 2).abs().add_(0.5),
+        "focus_src": rand(edge_count, 2, 32),
+        "focus_weight": rand(32, 2),
+        "focus_scale": rand(2, 32),
+    }
+
+
+def _clone_inputs(inputs):
+    return {name: value.clone() for name, value in inputs.items()}
+
+
+def _allocate_outputs(inputs):
+    _, Outputs = _phase_c_api()
+    edge_count = inputs["stack"].shape[0]
+    node_count = inputs["grad_out"].shape[0]
+    options = {"device": inputs["stack"].device, "dtype": torch.float32}
+    return Outputs(
+        grad_stack=inputs["stack"],
+        grad_wigner_dt=torch.empty(edge_count, 46, **options),
+        grad_logits=torch.empty(edge_count, 2, **options),
+        grad_edge=torch.empty(edge_count, **options),
+        grad_z_partial=torch.empty(node_count, 2, **options),
+        grad_z=torch.empty(2, **options),
+        grad_focus_src=torch.empty(2, edge_count, 32, **options),
+    )
+
+
+def _call(runner, inputs, outputs):
+    return runner(
+        inputs["grad_out"],
+        inputs["stack"],
+        inputs["wigner_dt"],
+        inputs["alpha"],
+        inputs["focus_alpha"],
+        inputs["dst_ptr"],
+        inputs["rotate_inv_rescale"],
+        inputs["edge_gate"],
+        inputs["z_bias_raw"],
+        inputs["group_max"],
+        inputs["denom"],
+        inputs["focus_src"],
+        inputs["focus_weight"],
+        inputs["focus_scale"],
+        outputs,
+    )
+
+
+def _redzoned_like(tensor: torch.Tensor):
+    flat = torch.full(
+        (tensor.numel() + 2 * REDZONE_ELEMENTS,),
+        SENTINEL,
+        device=tensor.device,
+        dtype=tensor.dtype,
+    )
+    value = flat[REDZONE_ELEMENTS : REDZONE_ELEMENTS + tensor.numel()].view_as(tensor)
+    return value, (flat[:REDZONE_ELEMENTS], flat[-REDZONE_ELEMENTS:])
+
+
+def _misaligned_like(tensor: torch.Tensor):
+    flat = torch.empty(tensor.numel() + 1, device=tensor.device, dtype=tensor.dtype)
+    value = flat[1:].view_as(tensor)
+    value.copy_(tensor)
+    assert value.data_ptr() % 16 != 0
+    return value
+
+
+def _assert_outputs_close(actual, expected):
+    for field in fields(type(expected)):
+        torch.testing.assert_close(
+            getattr(actual, field.name),
+            getattr(expected, field.name),
+            atol=1.0e-6,
+            rtol=1.0e-6,
+        )
+
+
+@pytest.mark.parametrize(
+    "dst_ptr_values",
+    (
+        [0, 0, 2, 2, 7],
+        [0, 1, 1, 4, 12, 12],
+        [0, 0, 37, 37],
+    ),
+)
+def test_exact_inplace_matches_independent_run_and_preserves_redzones(
+    dst_ptr_values,
+):
+    runner = _runner()
+    reference_inputs = _inputs(dst_ptr_values)
+    reference_outputs = _allocate_outputs(reference_inputs)
+    _call(runner, reference_inputs, reference_outputs)
+
+    inputs = _inputs(dst_ptr_values)
+    stack, stack_redzones = _redzoned_like(inputs["stack"])
+    stack.copy_(inputs["stack"])
+    inputs["stack"] = stack
+    outputs = _allocate_outputs(inputs)
+    redzones = {"grad_stack": stack_redzones}
+    for field_name in NONALIASED_OUTPUT_FIELDS:
+        value, field_redzones = _redzoned_like(getattr(outputs, field_name))
+        outputs = replace(outputs, **{field_name: value})
+        redzones[field_name] = field_redzones
+
+    _call(runner, inputs, outputs)
+    torch.cuda.synchronize()
+
+    assert outputs.grad_stack.data_ptr() == inputs["stack"].data_ptr()
+    _assert_outputs_close(outputs, reference_outputs)
+    for field_name, (prefix, suffix) in redzones.items():
+        assert torch.equal(prefix, torch.full_like(prefix, SENTINEL)), field_name
+        assert torch.equal(suffix, torch.full_like(suffix, SENTINEL)), field_name
+
+
+def test_grad_stack_must_be_the_exact_input_view():
+    inputs = _inputs([0, 3, 7])
+    outputs = replace(_allocate_outputs(inputs), grad_stack=inputs["stack"].clone())
+
+    with pytest.raises(ValueError, match="exact in-place stack view"):
+        _call(_runner(), inputs, outputs)
+
+
+def test_partial_or_shifted_stack_alias_is_rejected():
+    inputs = _inputs([0, 3, 7])
+    edge_count = inputs["stack"].shape[0]
+    base = torch.empty(
+        edge_count + 1,
+        2,
+        10,
+        32,
+        device="cuda",
+        dtype=torch.float32,
+    )
+    base[1:].copy_(inputs["stack"])
+    inputs["stack"] = base[1:]
+    outputs = replace(_allocate_outputs(inputs), grad_stack=base[:-1])
+
+    with pytest.raises(ValueError, match="exact in-place stack view"):
+        _call(_runner(), inputs, outputs)
+
+
+@pytest.mark.parametrize(
+    "tensor_name",
+    INPUT_FIELDS + tuple(f"outputs.{name}" for name in NONALIASED_OUTPUT_FIELDS),
+)
+def test_compiled_tensors_require_16_byte_alignment(tensor_name):
+    inputs = _inputs([0, 3, 7])
+    outputs = _allocate_outputs(inputs)
+
+    if tensor_name.startswith("outputs."):
+        field_name = tensor_name.removeprefix("outputs.")
+        outputs = replace(
+            outputs,
+            **{field_name: _misaligned_like(getattr(outputs, field_name))},
+        )
+    else:
+        inputs[tensor_name] = _misaligned_like(inputs[tensor_name])
+        if tensor_name == "stack":
+            outputs = replace(outputs, grad_stack=inputs["stack"])
+
+    with pytest.raises(ValueError, match="must be 16-byte aligned"):
+        _call(_runner(), inputs, outputs)
+
+
+@pytest.mark.parametrize("field_name", NONALIASED_OUTPUT_FIELDS)
+def test_nonaliased_outputs_reject_input_overlap(field_name):
+    inputs = _inputs([0, 3, 7])
+    outputs = _allocate_outputs(inputs)
+    target = getattr(outputs, field_name)
+    overlapping = inputs["stack"].view(-1)[: target.numel()].view_as(target)
+    outputs = replace(outputs, **{field_name: overlapping})
+
+    with pytest.raises(ValueError, match="must not overlap"):
+        _call(_runner(), inputs, outputs)
+
+
+@pytest.mark.parametrize(
+    ("first_name", "second_name"),
+    tuple(itertools.combinations(NONALIASED_OUTPUT_FIELDS, 2)),
+)
+def test_nonaliased_output_pairs_must_not_overlap(first_name, second_name):
+    inputs = _inputs([0, 3, 7])
+    outputs = _allocate_outputs(inputs)
+    first = getattr(outputs, first_name)
+    second = getattr(outputs, second_name)
+    slab = torch.empty(
+        max(first.numel(), second.numel()),
+        device="cuda",
+        dtype=torch.float32,
+    )
+    outputs = replace(
+        outputs,
+        **{
+            first_name: slab[: first.numel()].view_as(first),
+            second_name: slab[: second.numel()].view_as(second),
+        },
+    )
+
+    with pytest.raises(ValueError, match="must not overlap output"):
+        _call(_runner(), inputs, outputs)
+
+
+def test_retained_runner_supports_dynamic_edge_counts():
+    runner = _runner()
+    for dst_ptr_values in ([0, 2, 7], [0, 0, 1, 9, 9]):
+        inputs = _inputs(dst_ptr_values)
+        original = _clone_inputs(inputs)
+        outputs = _allocate_outputs(inputs)
+        _call(runner, inputs, outputs)
+        first = {
+            field.name: getattr(outputs, field.name).clone()
+            for field in fields(outputs)
+        }
+
+        inputs = original
+        outputs = _allocate_outputs(inputs)
+        _call(runner, inputs, outputs)
+        torch.cuda.synchronize()
+        for field_name, expected in first.items():
+            torch.testing.assert_close(
+                getattr(outputs, field_name),
+                expected,
+                atol=1.0e-6,
+                rtol=1.0e-6,
+            )
diff --git a/source/tests/pt/model/sezm_cute/test_so2_radial_focus_fusion.py b/source/tests/pt/model/sezm_cute/test_so2_radial_focus_fusion.py
new file mode 100644
index 0000000000..d85adaae4d
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_radial_focus_fusion.py
@@ -0,0 +1,135 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""SM80/SM90 differential coverage for Phase-C focus-source load fusion."""
+
+from __future__ import (
+    annotations,
+)
+
+import unittest
+
+import pytest
+import torch
+
+
+def _has_supported_gpu() -> bool:
+    return torch.cuda.is_available() and torch.cuda.get_device_capability() in {
+        (8, 0),
+        (9, 0),
+    }
+
+
+def test_source_csr_preserves_equal_source_order() -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.so2.radial_phase_a import (
+        build_source_csr,
+    )
+
+    src = torch.tensor([2, 0, 2, 1, 0, 2], dtype=torch.int64, device="cpu")
+    source_csr = build_source_csr(src, node_count=3)
+
+    torch.testing.assert_close(
+        source_csr.source_order,
+        torch.tensor([1, 4, 3, 0, 2, 5], dtype=torch.int32, device="cpu"),
+    )
+
+
+@pytest.mark.parametrize("src_values", ([-1, 0], [0, 3]))
+def test_source_csr_rejects_out_of_range_sources(src_values: list[int]) -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.so2.radial_phase_a import (
+        build_source_csr,
+    )
+
+    src = torch.tensor(src_values, dtype=torch.int64, device="cpu")
+    with pytest.raises(RuntimeError, match="source indices"):
+        build_source_csr(src, node_count=3)
+
+
+def test_source_csr_eager_validation_reports_value_error() -> None:
+    from deepmd.pt_expt.kernels.cute.sezm.so2.radial_phase_a import (
+        build_source_csr,
+    )
+
+    src = torch.tensor([0, 3], dtype=torch.int64, device="cpu")
+    with pytest.raises(ValueError, match="source indices"):
+        build_source_csr(src, node_count=3, validate_sources=True)
+
+
+@unittest.skipUnless(_has_supported_gpu(), "requires an SM80 or SM90 CUDA device")
+class TestRadialFocusSourceFusion(unittest.TestCase):
+    def test_fused_load_matches_standalone_scalar_lane_add(self) -> None:
+        from deepmd.pt_expt.kernels.cute.sezm.so2.radial_phase_a import (
+            build_source_csr,
+            prepare_batched_radial_projection_weight,
+            run_neo_radial_phase_a_backward_node_tiled,
+        )
+
+        torch.manual_seed(18072026)
+        device = torch.device("cuda")
+        edge_count = 7
+        node_count = 3
+
+        src = torch.tensor([0, 2, 1, 0, 2, 1, 2], device=device, dtype=torch.int64)
+        source_csr = build_source_csr(src, node_count)
+        grad_stack = torch.randn(edge_count, 2, 10, 32, device=device)
+        grad_focus_src = torch.randn(2, edge_count, 32, device=device)
+        grad_logits = torch.randn(edge_count, 2, device=device)
+        radial_compact = torch.randn(edge_count, 25, device=device)
+        combined_weight = torch.randn(128, 25, device=device)
+        attention_weight = torch.randn(32, 2, device=device)
+        channel_basis = torch.randn(64, device=device)
+        x_wide = torch.randn(node_count, 16 * 64, device=device)
+        d_full = torch.randn(edge_count, 46, device=device)
+        projection_weight = prepare_batched_radial_projection_weight(
+            combined_weight,
+            attention_weight,
+        )
+
+        standalone_grad_stack = grad_stack.clone()
+        standalone_grad_stack[:, :, 0, :].add_(grad_focus_src.permute(1, 0, 2))
+        standalone = run_neo_radial_phase_a_backward_node_tiled(
+            standalone_grad_stack.view(edge_count, 2 * 10 * 32),
+            grad_logits,
+            radial_compact,
+            channel_basis,
+            x_wide,
+            source_csr.source_order,
+            source_csr.source_ptr,
+            d_full,
+            grad_focus_src_focus=torch.zeros_like(grad_focus_src),
+            batched_radial_projection_weight=projection_weight,
+        )
+        fused = run_neo_radial_phase_a_backward_node_tiled(
+            grad_stack.view(edge_count, 2 * 10 * 32),
+            grad_logits,
+            radial_compact,
+            channel_basis,
+            x_wide,
+            source_csr.source_order,
+            source_csr.source_ptr,
+            d_full,
+            grad_focus_src_focus=grad_focus_src,
+            batched_radial_projection_weight=projection_weight,
+        )
+
+        torch.testing.assert_close(
+            fused.grad_x_wide,
+            standalone.grad_x_wide,
+            rtol=0.0,
+            atol=0.0,
+        )
+        torch.testing.assert_close(
+            fused.grad_d_full,
+            standalone.grad_d_full,
+            rtol=0.0,
+            atol=0.0,
+        )
+        torch.testing.assert_close(
+            fused.grad_radial_m0,
+            standalone.grad_radial_m0,
+            rtol=0.0,
+            atol=0.0,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/model/sezm_cute/test_so2_radial_pitch68_safe.py b/source/tests/pt/model/sezm_cute/test_so2_radial_pitch68_safe.py
new file mode 100644
index 0000000000..d3b296d964
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_radial_pitch68_safe.py
@@ -0,0 +1,163 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Safety coverage for the fixed pitch-68 radial backward path."""
+
+from __future__ import (
+    annotations,
+)
+
+import ast
+import unittest
+
+from ._paths import (
+    CUTE_ROOT,
+)
+
+WRAPPER_PATH = CUTE_ROOT / "so2" / "radial_phase_a.py"
+KERNEL_PATH = CUTE_ROOT / "so2" / "kernels" / "radial_phase_a_backward.py"
+
+WARP_HELPERS = {"_warp_owned_grad_compact", "_warp_owned_grad_d"}
+
+
+def _call_name(node: ast.Call) -> str | None:
+    if isinstance(node.func, ast.Name):
+        return node.func.id
+    if isinstance(node.func, ast.Attribute):
+        return node.func.attr
+    return None
+
+
+def _is_constexpr_if(node: ast.If) -> bool:
+    test = node.test
+    return isinstance(test, ast.Call) and _call_name(test) == "const_expr"
+
+
+def _function(tree: ast.AST, name: str) -> ast.FunctionDef:
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name == name:
+            return node
+    raise AssertionError(f"missing function {name}")
+
+
+def _dynamic_if_ancestors(
+    node: ast.AST,
+    parents: dict[ast.AST, ast.AST],
+    boundary: ast.FunctionDef,
+) -> list[ast.If]:
+    result: list[ast.If] = []
+    current = parents.get(node)
+    while current is not None and current is not boundary:
+        if isinstance(current, ast.If) and not _is_constexpr_if(current):
+            result.append(current)
+        current = parents.get(current)
+    return result
+
+
+class TestRadialPitch68SafeStatic(unittest.TestCase):
+    def test_pitch68_specialization_contract(self) -> None:
+        wrapper_source = WRAPPER_PATH.read_text(encoding="utf-8")
+        kernel_source = KERNEL_PATH.read_text(encoding="utf-8")
+        kernel_tree = ast.parse(kernel_source)
+        pitch_assignment = next(
+            node
+            for node in kernel_tree.body
+            if isinstance(node, ast.Assign)
+            and any(
+                isinstance(target, ast.Name) and target.id == "SHARED_ROW_PITCH"
+                for target in node.targets
+            )
+        )
+        pitch_expression = pitch_assignment.value
+
+        self.assertIsInstance(pitch_expression, ast.BinOp)
+        self.assertIsInstance(pitch_expression.op, ast.Add)
+        self.assertIsInstance(pitch_expression.left, ast.Name)
+        self.assertEqual(pitch_expression.left.id, "HIDDEN")
+        self.assertIsInstance(pitch_expression.right, ast.Constant)
+        self.assertEqual(pitch_expression.right.value, 4)
+        self.assertNotIn("pitch68_safe", wrapper_source)
+        self.assertNotIn("pitch68_safe", kernel_source)
+
+    def test_warp_collectives_are_not_runtime_predicated(self) -> None:
+        source = KERNEL_PATH.read_text(encoding="utf-8")
+        tree = ast.parse(source)
+        parents = {
+            child: parent
+            for parent in ast.walk(tree)
+            for child in ast.iter_child_nodes(parent)
+        }
+        kernel = _function(tree, "neo_radial_phase_a_backward_node_kernel")
+
+        helper_calls = [
+            node
+            for node in ast.walk(kernel)
+            if isinstance(node, ast.Call) and _call_name(node) in WARP_HELPERS
+        ]
+        self.assertEqual(
+            sorted(_call_name(node) for node in helper_calls),
+            sorted(WARP_HELPERS),
+        )
+        for call in helper_calls:
+            self.assertEqual(
+                _dynamic_if_ancestors(call, parents, kernel),
+                [],
+                f"{_call_name(call)} must be reached by every warp lane",
+            )
+
+        compact_call = next(
+            node
+            for node in helper_calls
+            if _call_name(node) == "_warp_owned_grad_compact"
+        )
+        panel_call = next(
+            node for node in helper_calls if _call_name(node) == "_warp_owned_grad_d"
+        )
+        self.assertEqual(ast.unparse(compact_call.args[3]), "safe_compact_idx")
+        self.assertEqual(ast.unparse(panel_call.args[2]), "safe_panel_idx")
+
+        for helper_name in WARP_HELPERS:
+            helper = _function(tree, helper_name)
+            reductions = [
+                node
+                for node in ast.walk(helper)
+                if isinstance(node, ast.Call)
+                and _call_name(node) == "warp_reduction_sum"
+            ]
+            self.assertEqual(len(reductions), 1)
+            self.assertEqual(
+                _dynamic_if_ancestors(reductions[0], parents, helper),
+                [],
+                f"warp collective in {helper_name} must be unconditional",
+            )
+
+        assignments = [
+            node for node in ast.walk(kernel) if isinstance(node, ast.Assign)
+        ]
+        compact_write = next(
+            node
+            for node in assignments
+            if ast.unparse(node.targets[0]) == "grad_compact[compact_idx]"
+        )
+        panel_write = next(
+            node
+            for node in assignments
+            if ast.unparse(node.targets[0]) == "params.grad_d_full[edge, panel_idx]"
+        )
+        self.assertEqual(
+            {
+                ast.unparse(guard.test)
+                for guard in _dynamic_if_ancestors(compact_write, parents, kernel)
+            },
+            {"compact_idx < COMPACT_WIDTH", "subgroup_lane == 0"},
+        )
+        self.assertEqual(
+            {
+                ast.unparse(guard.test)
+                for guard in _dynamic_if_ancestors(panel_write, parents, kernel)
+            },
+            {"panel_idx < PACKED_WIGNER_VALUES", "subgroup_lane == 0"},
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/model/sezm_cute/test_so2_radial_projection.py b/source/tests/pt/model/sezm_cute/test_so2_radial_projection.py
new file mode 100644
index 0000000000..13f84f8af9
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_radial_projection.py
@@ -0,0 +1,204 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""CPU contracts for the batched Neo radial adjoint projection."""
+
+from __future__ import (
+    annotations,
+)
+
+import unittest
+from types import (
+    SimpleNamespace,
+)
+from unittest import (
+    mock,
+)
+
+import torch
+
+from deepmd.pt_expt.kernels.cute.sezm.so2 import radial_phase_a as _RADIAL_PROJECTION
+from deepmd.pt_expt.kernels.cute.sezm.so2 import runner as _SO2_RUNNER
+
+COMPACT_WIDTH = _RADIAL_PROJECTION.COMPACT_WIDTH
+FOCUS_COUNT = _RADIAL_PROJECTION.FOCUS_COUNT
+FOCUS_HIDDEN = _RADIAL_PROJECTION.FOCUS_HIDDEN
+PROJECTION_INPUT_WIDTH = _RADIAL_PROJECTION.PROJECTION_INPUT_WIDTH
+RADIAL_WIDTH = _RADIAL_PROJECTION.RADIAL_WIDTH
+_project_batched_radial_adjoint = _RADIAL_PROJECTION._project_batched_radial_adjoint
+prepare_batched_radial_projection_weight = (
+    _RADIAL_PROJECTION.prepare_batched_radial_projection_weight
+)
+_batched_radial_projection_weight = _SO2_RUNNER._batched_radial_projection_weight
+
+
+class TestBatchedRadialProjection(unittest.TestCase):
+    def test_precombined_single_gemm_matches_two_projection_algebra(self):
+        torch.manual_seed(20260718)
+        edge_count = 7
+        grad_compact = torch.randn(
+            edge_count, COMPACT_WIDTH, dtype=torch.float32, device="cpu"
+        )
+        grad_logits = torch.randn(
+            edge_count, FOCUS_COUNT, dtype=torch.float32, device="cpu"
+        )
+        combined_weight = torch.randn(
+            RADIAL_WIDTH,
+            COMPACT_WIDTH,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        attention_weight = torch.randn(
+            FOCUS_HIDDEN,
+            FOCUS_COUNT,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        projection_weight = prepare_batched_radial_projection_weight(
+            combined_weight,
+            attention_weight,
+        )
+
+        expected = grad_compact @ combined_weight.transpose(0, 1)
+        expected[:, :FOCUS_HIDDEN].add_(grad_logits @ attention_weight.transpose(0, 1))
+        workspace = torch.full(
+            (edge_count, FOCUS_COUNT * 10 * FOCUS_HIDDEN),
+            torch.nan,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        actual = torch.empty(
+            edge_count,
+            RADIAL_WIDTH,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        workspace[:, :COMPACT_WIDTH].copy_(grad_compact)
+        workspace[:, COMPACT_WIDTH:PROJECTION_INPUT_WIDTH].copy_(grad_logits)
+
+        with (
+            mock.patch.object(torch, "cat", wraps=torch.cat) as cat,
+            mock.patch.object(torch, "mm", wraps=torch.mm) as mm,
+        ):
+            _project_batched_radial_adjoint(
+                projection_weight,
+                actual,
+                workspace,
+            )
+
+        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=2e-5)
+        self.assertEqual(mm.call_count, 1)
+        self.assertEqual(cat.call_count, 0)
+        gemm_input, gemm_weight = mm.call_args.args
+        self.assertEqual(tuple(gemm_input.shape), (edge_count, PROJECTION_INPUT_WIDTH))
+        self.assertEqual(
+            tuple(gemm_weight.shape), (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH)
+        )
+        self.assertEqual(gemm_input.stride(), (workspace.shape[1], 1))
+        self.assertEqual(
+            gemm_input.untyped_storage().data_ptr(),
+            workspace.untyped_storage().data_ptr(),
+        )
+        self.assertIs(mm.call_args.kwargs["out"], actual)
+        torch.testing.assert_close(
+            workspace[:, :COMPACT_WIDTH],
+            grad_compact,
+            rtol=0.0,
+            atol=0.0,
+        )
+        torch.testing.assert_close(
+            workspace[:, COMPACT_WIDTH:PROJECTION_INPUT_WIDTH],
+            grad_logits,
+            rtol=0.0,
+            atol=0.0,
+        )
+        self.assertTrue(torch.isnan(workspace[:, PROJECTION_INPUT_WIDTH:]).all())
+
+    def test_precombined_weight_layout_and_cache_contract(self):
+        combined_weight = torch.randn(
+            RADIAL_WIDTH,
+            COMPACT_WIDTH,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        attention_weight = torch.randn(
+            FOCUS_HIDDEN,
+            FOCUS_COUNT,
+            dtype=torch.float32,
+            device="cpu",
+        )
+        owner = SimpleNamespace()
+
+        first = _batched_radial_projection_weight(
+            owner,
+            combined_weight,
+            attention_weight,
+        )
+        second = _batched_radial_projection_weight(
+            owner,
+            combined_weight,
+            attention_weight,
+        )
+
+        self.assertIs(first, second)
+        self.assertEqual(tuple(first.shape), (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH))
+        self.assertTrue(first.is_contiguous())
+        torch.testing.assert_close(
+            first[:COMPACT_WIDTH],
+            combined_weight.transpose(0, 1),
+            rtol=0.0,
+            atol=0.0,
+        )
+        attention_panel = torch.zeros(
+            FOCUS_COUNT,
+            2 * FOCUS_HIDDEN,
+            device="cpu",
+        )
+        attention_panel[:, :FOCUS_HIDDEN].copy_(attention_weight.transpose(0, 1))
+        torch.testing.assert_close(
+            first[COMPACT_WIDTH:, : 2 * FOCUS_HIDDEN],
+            attention_panel,
+            rtol=0.0,
+            atol=0.0,
+        )
+        self.assertTrue(torch.count_nonzero(first[COMPACT_WIDTH:, 64:]) == 0)
+
+        attention_weight.add_(1.0)
+        updated = _batched_radial_projection_weight(
+            owner,
+            combined_weight,
+            attention_weight,
+        )
+        self.assertIsNot(first, updated)
+
+    def test_projection_requires_strict_fp32_and_full_consumed_workspace(self):
+        combined_weight = torch.randn(
+            RADIAL_WIDTH,
+            COMPACT_WIDTH,
+            device="cpu",
+        )
+        attention_weight = torch.randn(
+            FOCUS_HIDDEN,
+            FOCUS_COUNT,
+            device="cpu",
+        )
+        with self.assertRaisesRegex(TypeError, "torch.float32"):
+            prepare_batched_radial_projection_weight(
+                combined_weight.double(),
+                attention_weight.double(),
+            )
+
+        projection_weight = prepare_batched_radial_projection_weight(
+            combined_weight,
+            attention_weight,
+        )
+        edge_count = 2
+        with self.assertRaisesRegex(ValueError, "consumed_workspace must have shape"):
+            _project_batched_radial_adjoint(
+                projection_weight,
+                torch.empty(edge_count, RADIAL_WIDTH, device="cpu"),
+                torch.empty(edge_count, PROJECTION_INPUT_WIDTH, device="cpu"),
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/model/sezm_cute/test_so2_structural_gate_vec4_sm80.py b/source/tests/pt/model/sezm_cute/test_so2_structural_gate_vec4_sm80.py
new file mode 100644
index 0000000000..e3dd389da1
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_structural_gate_vec4_sm80.py
@@ -0,0 +1,550 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Contracts for the shared structural-gate vec4 path."""
+
+from __future__ import (
+    annotations,
+)
+
+import ast
+import importlib
+import math
+import random
+import struct
+import unittest
+
+from ._paths import (
+    CUTE_ROOT,
+)
+
+KERNEL_PATH = CUTE_ROOT / "so2" / "kernels" / "structural_gate_sm80.py"
+DYNAMIC_EDGE_COUNTS = (0, 1, 7, 8, 9, 31, 32, 33)
+
+
+def _f32(value: float) -> float:
+    return struct.unpack("f", struct.pack("f", value))[0]
+
+
+def _sigmoid_f32(value: float) -> float:
+    value = _f32(value)
+    return _f32(_f32(1.0) / _f32(_f32(1.0) + _f32(math.exp(_f32(-value)))))
+
+
+def _mul_add_f32(lhs: float, rhs: float, addend: float) -> float:
+    return _f32(_f32(lhs * rhs) + addend)
+
+
+def _reference_forward(
+    residual: list[float],
+    y: list[float],
+    logits: list[float],
+    edge_count: int,
+) -> list[float]:
+    out = residual.copy()
+    for edge in range(edge_count):
+        for focus in range(2):
+            row = edge * 2 + focus
+            row_base = row * 10 * 32
+            logits_base = (focus * edge_count + edge) * 3 * 32
+            gates = [
+                [
+                    _sigmoid_f32(logits[logits_base + gate * 32 + channel])
+                    for channel in range(32)
+                ]
+                for gate in range(3)
+            ]
+            for channel in range(32):
+                index = row_base + channel
+                y0 = y[index]
+                out[index] = _mul_add_f32(y0, _sigmoid_f32(y0), out[index])
+            for degree in range(1, 10):
+                gate = gates[(degree - 1) % 3]
+                for channel in range(32):
+                    index = row_base + degree * 32 + channel
+                    out[index] = _mul_add_f32(
+                        y[index],
+                        gate[channel],
+                        out[index],
+                    )
+    return out
+
+
+def _vec4_forward_inplace_model(
+    out: list[float],
+    y: list[float],
+    logits: list[float],
+    edge_count: int,
+) -> None:
+    rows = edge_count * 2
+    for block_row in range((rows + 15) // 16):
+        for row_slot in range(16):
+            row = block_row * 16 + row_slot
+            if row >= rows:
+                continue
+            edge, focus = divmod(row, 2)
+            row_base = row * 10 * 32
+            logits_base = (focus * edge_count + edge) * 3 * 32
+            for channel_group in range(8):
+                channel_base = channel_group * 4
+                for lane in range(4):
+                    index = row_base + channel_base + lane
+                    y0 = y[index]
+                    out[index] = _mul_add_f32(y0, _sigmoid_f32(y0), out[index])
+                for gate_index in range(3):
+                    gates = [
+                        _sigmoid_f32(
+                            logits[logits_base + gate_index * 32 + channel_base + lane]
+                        )
+                        for lane in range(4)
+                    ]
+                    for repeat in range(3):
+                        degree = 1 + gate_index + repeat * 3
+                        for lane in range(4):
+                            index = row_base + degree * 32 + channel_base + lane
+                            out[index] = _mul_add_f32(
+                                y[index],
+                                gates[lane],
+                                out[index],
+                            )
+
+
+def _backward_reference(
+    grad_out: list[float],
+    y: list[float],
+    logits: list[float],
+    edge_count: int,
+) -> tuple[list[float], list[float]]:
+    grad_y = [0.0] * len(y)
+    grad_logits = [0.0] * len(logits)
+    for edge in range(edge_count):
+        for focus in range(2):
+            row = edge * 2 + focus
+            row_base = row * 10 * 32
+            logits_base = (focus * edge_count + edge) * 3 * 32
+            for channel in range(32):
+                y0 = _f32(y[row_base + channel])
+                sig0 = _sigmoid_f32(y0)
+                grad0 = _f32(grad_out[row_base + channel])
+                inner = _f32(_f32(1.0) + _f32(y0 * _f32(_f32(1.0) - sig0)))
+                grad_y[row_base + channel] = _f32(_f32(grad0 * sig0) * inner)
+                for gate_index in range(3):
+                    gate_offset = logits_base + gate_index * 32 + channel
+                    gate = _sigmoid_f32(logits[gate_offset])
+                    grad_logit = _f32(0.0)
+                    for repeat in range(3):
+                        degree = 1 + gate_index + repeat * 3
+                        index = row_base + degree * 32 + channel
+                        gout = _f32(grad_out[index])
+                        grad_y[index] = _f32(gout * gate)
+                        term = _f32(gout * _f32(y[index]))
+                        term = _f32(term * gate)
+                        term = _f32(term * _f32(_f32(1.0) - gate))
+                        grad_logit = _f32(grad_logit + term)
+                    grad_logits[gate_offset] = grad_logit
+    return grad_y, grad_logits
+
+
+def _vec4_backward_model(
+    grad_out: list[float],
+    y: list[float],
+    logits: list[float],
+    edge_count: int,
+) -> tuple[list[float], list[float]]:
+    grad_y = [0.0] * len(y)
+    grad_logits = [0.0] * len(logits)
+    rows = edge_count * 2
+    for block_row in range((rows + 15) // 16):
+        for row_slot in range(16):
+            row = block_row * 16 + row_slot
+            if row >= rows:
+                continue
+            edge, focus = divmod(row, 2)
+            row_base = row * 10 * 32
+            logits_base = (focus * edge_count + edge) * 3 * 32
+            for channel_group in range(8):
+                channel_base = channel_group * 4
+                for lane in range(4):
+                    channel = channel_base + lane
+                    y0 = _f32(y[row_base + channel])
+                    sig0 = _sigmoid_f32(y0)
+                    grad0 = _f32(grad_out[row_base + channel])
+                    inner = _f32(_f32(1.0) + _f32(y0 * _f32(_f32(1.0) - sig0)))
+                    grad_y[row_base + channel] = _f32(_f32(grad0 * sig0) * inner)
+                for gate_index in range(3):
+                    for lane in range(4):
+                        channel = channel_base + lane
+                        gate_offset = logits_base + gate_index * 32 + channel
+                        gate = _sigmoid_f32(logits[gate_offset])
+                        grad_logit = _f32(0.0)
+                        for repeat in range(3):
+                            degree = 1 + gate_index + repeat * 3
+                            index = row_base + degree * 32 + channel
+                            gout = _f32(grad_out[index])
+                            grad_y[index] = _f32(gout * gate)
+                            term = _f32(gout * _f32(y[index]))
+                            term = _f32(term * gate)
+                            term = _f32(term * _f32(_f32(1.0) - gate))
+                            grad_logit = _f32(grad_logit + term)
+                        grad_logits[gate_offset] = grad_logit
+    return grad_y, grad_logits
+
+
+def _module_constants(tree: ast.Module) -> dict[str, object]:
+    def constant_value(node: ast.expr, namespace: dict[str, object]) -> object:
+        if isinstance(node, ast.Constant):
+            return node.value
+        if isinstance(node, ast.Name):
+            return namespace[node.id]
+        if isinstance(node, ast.BinOp):
+            lhs = constant_value(node.left, namespace)
+            rhs = constant_value(node.right, namespace)
+            if isinstance(node.op, ast.FloorDiv):
+                return lhs // rhs
+            if isinstance(node.op, ast.Mult):
+                return lhs * rhs
+        raise ValueError("not a supported module constant")
+
+    constants: dict[str, object] = {}
+    namespace: dict[str, object] = {}
+    for node in tree.body:
+        if not isinstance(node, ast.Assign) or len(node.targets) != 1:
+            continue
+        target = node.targets[0]
+        if not isinstance(target, ast.Name):
+            continue
+        try:
+            value = constant_value(node.value, namespace)
+        except (KeyError, TypeError, ValueError, ZeroDivisionError):
+            continue
+        constants[target.id] = value
+        namespace[target.id] = value
+    return constants
+
+
+def _load_wrapper_module():
+    return importlib.import_module(
+        "deepmd.pt_expt.kernels.cute.sezm.so2.structural_gate"
+    )
+
+
+class TestSM80StructuralGateVec4Contract(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.kernel_source = KERNEL_PATH.read_text()
+        cls.constants = _module_constants(ast.parse(cls.kernel_source))
+
+    def test_vector_layout_and_launch_are_fixed(self):
+        self.assertEqual(self.constants["FOCUS_COUNT"], 2)
+        self.assertEqual(self.constants["REDUCED_COUNT"], 10)
+        self.assertEqual(self.constants["CHANNELS"], 32)
+        self.assertEqual(self.constants["VECTOR_WIDTH"], 4)
+        self.assertEqual(self.constants["CHANNEL_GROUPS"], 8)
+        self.assertEqual(self.constants["ROWS_PER_BLOCK"], 16)
+        self.assertEqual(self.constants["THREADS"], 128)
+        self.assertIn("residual.element_type.width * VECTOR_WIDTH", self.kernel_source)
+        self.assertIn("cute.make_tiled_copy_tv", self.kernel_source)
+        self.assertIn("(1, CHANNEL_GROUPS)", self.kernel_source)
+        self.assertIn("(1, VECTOR_WIDTH)", self.kernel_source)
+
+    def test_kernel_is_strict_fp32_edge_major_only(self):
+        for forbidden in (
+            "Float16",
+            "BFloat16",
+            "TensorFloat32",
+        ):
+            self.assertNotIn(forbidden, self.kernel_source)
+        self.assertIn("cutlass.Float32", self.kernel_source)
+        self.assertIn('"assumed_align": 16', self.kernel_source)
+        self.assertIn("_guard_vec4_dispatch", self.kernel_source)
+        self.assertIn("runtime_policy.SUPPORTED_SO2_CAPABILITIES", self.kernel_source)
+
+
+class TestSM80StructuralGateVec4Arithmetic(unittest.TestCase):
+    def test_forward_dynamic_edges_and_block_tails(self):
+        for edge_count in DYNAMIC_EDGE_COUNTS:
+            with self.subTest(edge_count=edge_count):
+                rng = random.Random(20260721 + edge_count)
+                values = edge_count * 2 * 10 * 32
+                residual = [_f32(rng.uniform(-2.0, 2.0)) for _ in range(values)]
+                y = [_f32(rng.uniform(-2.0, 2.0)) for _ in range(values)]
+                logits = [
+                    _f32(rng.uniform(-4.0, 4.0)) for _ in range(2 * edge_count * 3 * 32)
+                ]
+                expected = _reference_forward(residual, y, logits, edge_count)
+                actual = residual.copy()
+                _vec4_forward_inplace_model(actual, y, logits, edge_count)
+                self.assertEqual(actual, expected)
+
+    def test_backward_dynamic_edges_and_block_tails(self):
+        for edge_count in DYNAMIC_EDGE_COUNTS:
+            with self.subTest(edge_count=edge_count):
+                rng = random.Random(20260722 + edge_count)
+                values = edge_count * 2 * 10 * 32
+                grad_out = [_f32(rng.uniform(-2.0, 2.0)) for _ in range(values)]
+                y = [_f32(rng.uniform(-2.0, 2.0)) for _ in range(values)]
+                logits = [
+                    _f32(rng.uniform(-4.0, 4.0)) for _ in range(2 * edge_count * 3 * 32)
+                ]
+                expected_grad_y, expected_grad_logits = _backward_reference(
+                    grad_out,
+                    y,
+                    logits,
+                    edge_count,
+                )
+                actual_grad_y, actual_grad_logits = _vec4_backward_model(
+                    grad_out,
+                    y,
+                    logits,
+                    edge_count,
+                )
+                self.assertEqual(actual_grad_y, expected_grad_y)
+                self.assertEqual(actual_grad_logits, expected_grad_logits)
+
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+@unittest.skipUnless(torch is not None, "PyTorch is required")
+class TestSM80StructuralGateVec4Dispatch(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.dispatch = staticmethod(
+            _load_wrapper_module()._dispatch_aligned_vec4_kernel
+        )
+
+    def test_aligned_fp32_tensor_dispatches(self):
+        calls = []
+        tensor = torch.empty(16, dtype=torch.float32, device="cpu")
+
+        def kernel(value):
+            calls.append(value)
+            return "dispatched"
+
+        result = self.dispatch(kernel, ("value",), tensor)
+        self.assertEqual(result, "dispatched")
+        self.assertEqual(calls, [tensor])
+
+    def test_misaligned_storage_offset_is_rejected_before_dispatch(self):
+        calls = []
+        tensor = torch.empty(17, dtype=torch.float32, device="cpu")[1:]
+        self.assertTrue(tensor.is_contiguous())
+        self.assertEqual(tensor.storage_offset(), 1)
+
+        with self.assertRaisesRegex(ValueError, "storage offsets divisible by 4"):
+            self.dispatch(lambda value: calls.append(value), ("value",), tensor)
+        self.assertEqual(calls, [])
+
+    def test_misaligned_pointer_is_rejected_before_dispatch(self):
+        class MisalignedTensor:
+            dtype = torch.float32
+            shape = (4,)
+
+            @staticmethod
+            def numel():
+                return 4
+
+            @staticmethod
+            def is_contiguous():
+                return True
+
+            @staticmethod
+            def stride():
+                return (1,)
+
+            @staticmethod
+            def storage_offset():
+                return 0
+
+            @staticmethod
+            def data_ptr():
+                return 4
+
+        calls = []
+        with self.assertRaisesRegex(ValueError, "16-byte-aligned"):
+            self.dispatch(
+                lambda value: calls.append(value),
+                ("value",),
+                MisalignedTensor(),
+            )
+        self.assertEqual(calls, [])
+
+    def test_noncompact_and_non_fp32_tensors_are_rejected(self):
+        noncompact = torch.empty(4, 8, dtype=torch.float32, device="cpu").T
+        with self.assertRaisesRegex(ValueError, "compact tensors"):
+            self.dispatch(lambda value: value, ("value",), noncompact)
+
+        with self.assertRaisesRegex(TypeError, "requires float32"):
+            self.dispatch(
+                lambda value: value,
+                ("value",),
+                torch.empty(4, dtype=torch.float64, device="cpu"),
+            )
+
+    def test_empty_tensor_does_not_launch(self):
+        calls = []
+        result = self.dispatch(
+            lambda value: calls.append(value),
+            ("value",),
+            torch.empty(0, dtype=torch.float32, device="cpu"),
+        )
+        self.assertIsNone(result)
+        self.assertEqual(calls, [])
+
+
+@unittest.skipUnless(
+    torch is not None
+    and torch.cuda.is_available()
+    and tuple(torch.cuda.get_device_capability())
+    in {(8, 0), (8, 6), (8, 9), (9, 0), (10, 0), (12, 0)},
+    "A supported Neo SO2 CUDA runtime is required",
+)
+class TestSM80StructuralGateVec4CudaDifferential(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.structural_gate_sm80 import (
+            compile_neo_gate_split_structural_vec4_sm80_backward,
+            compile_neo_gate_split_structural_vec4_sm80_forward,
+        )
+
+        capability = tuple(torch.cuda.get_device_capability())
+        compile_identity = (torch.cuda.current_device(), *capability)
+        cls.vec4_forward = staticmethod(
+            compile_neo_gate_split_structural_vec4_sm80_forward(compile_identity)
+        )
+        cls.vec4_backward = staticmethod(
+            compile_neo_gate_split_structural_vec4_sm80_backward(compile_identity)
+        )
+
+    def test_forward_dynamic_edges_tails_and_empty(self):
+        for edge_count in DYNAMIC_EDGE_COUNTS:
+            with self.subTest(edge_count=edge_count):
+                torch.manual_seed(20260721 + edge_count)
+                residual = torch.randn(
+                    edge_count,
+                    2,
+                    10,
+                    32,
+                    device="cuda",
+                    dtype=torch.float32,
+                )
+                y = torch.randn_like(residual)
+                logits = torch.randn(
+                    2,
+                    edge_count,
+                    3 * 32,
+                    device="cuda",
+                    dtype=torch.float32,
+                )
+                gates = torch.sigmoid(
+                    logits.permute(1, 0, 2).reshape(edge_count, 2, 3, 32)
+                )
+                expected = residual.clone()
+                expected[:, :, 0, :].add_(torch.nn.functional.silu(y[:, :, 0, :]))
+                for degree in range(1, 10):
+                    expected[:, :, degree, :].add_(
+                        y[:, :, degree, :] * gates[:, :, (degree - 1) % 3, :]
+                    )
+
+                vec4 = residual.clone()
+                self.vec4_forward(
+                    vec4.view(edge_count * 2, 10 * 32),
+                    y.view(edge_count * 2, 10 * 32),
+                    logits,
+                    vec4.view(edge_count * 2, 10 * 32),
+                )
+                torch.testing.assert_close(vec4, expected, atol=5e-5, rtol=5e-5)
+
+    def test_backward_dynamic_edges_tails_and_empty(self):
+        gate_indices = torch.tensor(
+            [0, 1, 2, 0, 1, 2, 0, 1, 2],
+            device="cuda",
+        )
+        for edge_count in DYNAMIC_EDGE_COUNTS:
+            with self.subTest(edge_count=edge_count):
+                torch.manual_seed(20260722 + edge_count)
+                grad_out = torch.randn(
+                    edge_count,
+                    2,
+                    10,
+                    32,
+                    device="cuda",
+                    dtype=torch.float32,
+                )
+                y = torch.randn_like(grad_out)
+                logits = torch.randn(
+                    2,
+                    edge_count,
+                    3 * 32,
+                    device="cuda",
+                    dtype=torch.float32,
+                )
+                vec4_grad_y = torch.empty_like(y)
+                vec4_grad_logits = torch.empty_like(logits)
+                self.vec4_backward(
+                    grad_out.view(edge_count * 2, 10 * 32),
+                    y.view(edge_count * 2, 10 * 32),
+                    logits,
+                    vec4_grad_y.view(edge_count * 2, 10 * 32),
+                    vec4_grad_logits,
+                )
+
+                y_reference = y.detach().clone().requires_grad_(True)
+                logits_reference = logits.detach().clone().requires_grad_(True)
+                gates = torch.sigmoid(
+                    logits_reference.permute(1, 0, 2).reshape(
+                        edge_count,
+                        2,
+                        3,
+                        32,
+                    )
+                )
+                output = torch.cat(
+                    (
+                        torch.nn.functional.silu(y_reference[:, :, :1, :]),
+                        y_reference[:, :, 1:, :] * gates.index_select(2, gate_indices),
+                    ),
+                    dim=2,
+                )
+                expected_grad_y, expected_grad_logits = torch.autograd.grad(
+                    output,
+                    (y_reference, logits_reference),
+                    grad_out,
+                )
+                torch.testing.assert_close(
+                    vec4_grad_y,
+                    expected_grad_y,
+                    atol=5e-5,
+                    rtol=5e-5,
+                )
+                torch.testing.assert_close(
+                    vec4_grad_logits,
+                    expected_grad_logits,
+                    atol=5e-5,
+                    rtol=5e-5,
+                )
+
+    def test_misaligned_cuda_view_is_rejected(self):
+        edge_count = 1
+        storage = torch.empty(
+            edge_count * 2 * 10 * 32 + 1,
+            device="cuda",
+            dtype=torch.float32,
+        )
+        residual = storage[1:].view(edge_count * 2, 10 * 32)
+        y = torch.empty_like(residual)
+        logits = torch.empty(
+            2,
+            edge_count,
+            3 * 32,
+            device="cuda",
+            dtype=torch.float32,
+        )
+        with self.assertRaisesRegex(ValueError, "storage offsets divisible by 4"):
+            self.vec4_forward(residual, y, logits, residual)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/source/tests/pt/model/sezm_cute/test_so2_structural_memory_reuse.py b/source/tests/pt/model/sezm_cute/test_so2_structural_memory_reuse.py
new file mode 100644
index 0000000000..92ec6ded3e
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_so2_structural_memory_reuse.py
@@ -0,0 +1,214 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Behavioral contracts for SO2 structural storage reuse."""
+
+from __future__ import (
+    annotations,
+)
+
+from types import (
+    SimpleNamespace,
+)
+
+from deepmd.pt_expt.kernels.cute.sezm.so2 import linear as so2_linear
+from deepmd.pt_expt.kernels.cute.sezm.so2 import structural_gate as so2_gate_structural
+
+from .test_so2 import (
+    _SO2,
+    NeoFullCuteBackward,
+    NeoSO2BackwardWorkspace,
+    NeoSO2RuntimeConfig,
+    StackCache,
+    _validate_runtime_config,
+)
+
+
+def _valid_structural_config() -> NeoSO2RuntimeConfig:
+    return NeoSO2RuntimeConfig(per_focus_so2_fwd_pair=True)
+
+
+def _storage_ptr(tensor) -> int:
+    return tensor.untyped_storage().data_ptr()
+
+
+def _workspace_inputs(edge_count: int):
+    torch = _SO2.torch
+    return {
+        "edge_count": edge_count,
+        "node_count": 3,
+        "d_full": torch.empty(edge_count, 46, dtype=torch.float32, device="cpu"),
+        "dt_full": torch.empty(edge_count, 46, dtype=torch.float32, device="cpu"),
+        "radial": torch.empty(edge_count, 4, 32, dtype=torch.float32, device="cpu"),
+    }
+
+
+def _so2_linear(torch):
+    return SimpleNamespace(
+        lmax=3,
+        mmax=1,
+        in_channels=32,
+        out_channels=32,
+        n_focus=2,
+        mlp_bias=False,
+        weight_m0=torch.randn(4 * 32, 2 * 4 * 32, device="cpu"),
+        weight_m=(torch.randn(3 * 32, 2 * 2 * 3 * 32, device="cpu"),),
+    )
+
+
+def test_structural_memory_reuse_profile_is_valid() -> None:
+    config = _valid_structural_config()
+    assert _validate_runtime_config(config, compute_capability=(8, 0)) is None
+
+
+def test_final_so2_single_input_fold_matches_residual_plus_linear():
+    torch = _SO2.torch
+    torch.manual_seed(20260705)
+    x_local = torch.randn(3, 2, 10, 32, dtype=torch.float32, device="cpu")
+    linear = _so2_linear(torch)
+
+    linear_only = so2_linear.run_neo_so2_linear_manual(linear, x_local)
+    folded = so2_linear.run_neo_so2_linear_manual(
+        linear,
+        x_local,
+        add_residual=True,
+    )
+
+    torch.testing.assert_close(folded, linear_only + x_local, atol=5e-5, rtol=5e-5)
+
+
+def test_per_focus_pair_forward_matches_batched_forward():
+    torch = _SO2.torch
+    torch.manual_seed(20260718)
+    x_local = torch.randn(17, 2, 10, 32, dtype=torch.float32, device="cpu")
+    linear = _so2_linear(torch)
+
+    batched = so2_linear.run_neo_so2_linear_manual(linear, x_local)
+    per_focus = so2_linear.run_neo_so2_linear_manual(
+        linear,
+        x_local,
+        per_focus_pair=True,
+    )
+
+    torch.testing.assert_close(per_focus, batched, atol=5e-5, rtol=5e-5)
+
+
+def test_structural_workspace_reuses_saved_slabs_and_gate_panel():
+    torch = _SO2.torch
+    edge_count = 5
+    phase_c_stack = torch.empty(
+        edge_count, 2, 10, 32, dtype=torch.float32, device="cpu"
+    )
+    phase_c_y = torch.empty_like(phase_c_stack)
+    radial_scratch = torch.empty(
+        2, edge_count, 3 * 32, dtype=torch.float32, device="cpu"
+    )
+
+    workspace = NeoSO2BackwardWorkspace(
+        torch,
+        **_workspace_inputs(edge_count),
+        phase_c_stack=phase_c_stack,
+        phase_c_y=phase_c_y,
+        radial_scratch=radial_scratch,
+        structural_memory_reuse=True,
+    )
+
+    stack_storage = _storage_ptr(phase_c_stack)
+    assert _storage_ptr(workspace.grad_stack_focus) == stack_storage
+    assert _storage_ptr(workspace.grad_y) == stack_storage
+    assert _storage_ptr(workspace.grad_x_rot) == stack_storage
+    assert _storage_ptr(workspace.grad_d) == stack_storage
+    assert _storage_ptr(workspace.grad_radial_flat) == _storage_ptr(radial_scratch)
+    assert _storage_ptr(workspace.grad_mixed_slab) == _storage_ptr(phase_c_y)
+    assert workspace.grad_gate_logits is None
+
+
+def test_single_input_workspace_omits_phase_c_y_and_keeps_mixed_output_distinct():
+    torch = _SO2.torch
+    edge_count = 5
+    phase_c_stack = torch.empty(
+        edge_count, 2, 10, 32, dtype=torch.float32, device="cpu"
+    )
+    radial_scratch = torch.empty(edge_count, 4, 32, dtype=torch.float32, device="cpu")
+
+    workspace = NeoSO2BackwardWorkspace(
+        torch,
+        **_workspace_inputs(edge_count),
+        phase_c_stack=phase_c_stack,
+        phase_c_y=None,
+        radial_scratch=radial_scratch,
+        structural_memory_reuse=True,
+        phase_c_single_input_reuse=True,
+    )
+
+    assert _storage_ptr(workspace.grad_stack_focus) == _storage_ptr(phase_c_stack)
+    assert _storage_ptr(workspace.grad_mixed_slab) != _storage_ptr(phase_c_stack)
+    assert _storage_ptr(workspace.grad_radial_flat) == _storage_ptr(radial_scratch)
+    assert workspace.grad_mixed_slab.shape == phase_c_stack.shape
+
+
+def test_runner_routes_saved_structural_buffers_into_lazy_workspace():
+    torch = _SO2.torch
+    edge_count = 5
+    runner = object.__new__(NeoFullCuteBackward)
+    runner.torch = torch
+    runner.config = _valid_structural_config()
+    runner.edge_count = edge_count
+    runner.node_count = 3
+    runner.d = torch.empty(edge_count, 46, dtype=torch.float32, device="cpu")
+    runner.dt = torch.empty_like(runner.d)
+    runner.radial = torch.empty(edge_count, 4, 32, dtype=torch.float32, device="cpu")
+    runner.phase_c_stack = torch.empty(
+        edge_count, 2, 10, 32, dtype=torch.float32, device="cpu"
+    )
+    runner.phase_c_y = None
+    first_logits = torch.empty(2, edge_count, 3 * 32, dtype=torch.float32, device="cpu")
+    runner.stack_caches = [
+        StackCache(
+            torch.empty_like(runner.phase_c_stack), first_logits, object(), False
+        ),
+        StackCache(
+            torch.empty_like(runner.phase_c_stack),
+            torch.empty_like(first_logits),
+            object(),
+            False,
+        ),
+        StackCache(torch.empty_like(runner.phase_c_stack), None, object(), True),
+    ]
+    runner._backward_workspace = None
+
+    workspace = runner.ensure_backward_workspace()
+
+    assert runner.ensure_backward_workspace() is workspace
+    assert _storage_ptr(workspace.grad_stack_focus) == _storage_ptr(
+        runner.phase_c_stack
+    )
+    assert _storage_ptr(workspace.grad_radial_flat) == _storage_ptr(first_logits)
+    assert _storage_ptr(workspace.grad_mixed_slab) != _storage_ptr(runner.phase_c_stack)
+
+
+def test_structural_backward_can_overwrite_consumed_logits():
+    torch = _SO2.torch
+    edge_count = 3
+    grad_out = torch.randn(edge_count, 2, 10, 32, dtype=torch.float32, device="cpu")
+    y = torch.randn_like(grad_out)
+    logits = torch.randn(2, edge_count, 3 * 32, dtype=torch.float32, device="cpu")
+    grad_y = torch.empty_like(y)
+
+    def fake_kernel(grad_out_flat, y_flat, logits_in, grad_y_flat, grad_logits):
+        assert logits_in is logits
+        assert grad_logits is logits
+        grad_y_flat.copy_(grad_out_flat + y_flat)
+        grad_logits.fill_(7.0)
+
+    result = so2_gate_structural.run_structural_gate_backward(
+        fake_kernel,
+        grad_out,
+        y,
+        logits,
+        grad_y,
+        grad_logits=None,
+        overwrite_logits=True,
+    )
+
+    assert result is logits
+    torch.testing.assert_close(grad_y, grad_out + y)
diff --git a/source/tests/pt/model/sezm_cute/test_tiled_grid_product.py b/source/tests/pt/model/sezm_cute/test_tiled_grid_product.py
new file mode 100644
index 0000000000..3453dc1a85
--- /dev/null
+++ b/source/tests/pt/model/sezm_cute/test_tiled_grid_product.py
@@ -0,0 +1,379 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Strict-FP32 differentials for the tiled Neo output-grid product."""
+
+from __future__ import (
+    annotations,
+)
+
+import ast
+import importlib
+
+import pytest
+import torch
+
+from ._paths import (
+    CUTE_ROOT,
+)
+
+TOL = 5.0e-5
+PACKED_COEFF_DIM = 48
+GRID_SIZE = 152
+SUPPORTED_HIDDEN_CHANNELS = (96, 192)
+TILED_KERNEL_PATH = CUTE_ROOT / "output_grid" / "kernels" / "tiled_product.py"
+MESSAGE_GRID_PATH = CUTE_ROOT / "so2" / "kernels" / "message_grid_product.py"
+
+
+def _cute_runtime_skip_reason() -> str | None:
+    if not torch.cuda.is_available():
+        return "tiled output-grid differentials require CUDA"
+    if torch.cuda.get_device_capability()[0] < 8:
+        return "tiled output-grid differentials require compute capability 8.0+"
+    try:
+        importlib.import_module("cutlass.cute")
+        importlib.import_module("cuda.bindings.driver")
+    except Exception as exc:  # pragma: no cover - runtime dependent
+        return f"tiled output-grid differentials require CuTe DSL: {exc}"
+    return None
+
+
+_CUTE_SKIP_REASON = _cute_runtime_skip_reason()
+
+
+def _sm80_skip_reason() -> str | None:
+    if _CUTE_SKIP_REASON is not None:
+        return _CUTE_SKIP_REASON
+    if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6)}:
+        return "specialized output-grid differentials require sm80 or sm86"
+    return None
+
+
+_SM80_SKIP_REASON = _sm80_skip_reason()
+
+
+def _method_node(tree: ast.AST, name: str) -> ast.FunctionDef:
+    matches = [
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.FunctionDef) and node.name == name
+    ]
+    assert len(matches) == 1
+    return matches[0]
+
+
+def test_sm80_c96_n48_panel_tiles_shared_input_before_partition_b() -> None:
+    """Guard the CuTe B-fragment layout contract without requiring CUDA."""
+    tree = ast.parse(TILED_KERNEL_PATH.read_text(encoding="utf-8"))
+    method = _method_node(tree, "_backproject_panel_accumulate")
+
+    assignments = {
+        node.targets[0].id: node.value
+        for node in ast.walk(method)
+        if isinstance(node, ast.Assign)
+        and len(node.targets) == 1
+        and isinstance(node.targets[0], ast.Name)
+    }
+    shared_b = assignments["sB"]
+    assert isinstance(shared_b, ast.Call)
+    assert ast.unparse(shared_b.func) == "cute.local_tile"
+    assert ast.unparse(shared_b.args[0]) == "panel_b"
+    keywords = {
+        keyword.arg: ast.unparse(keyword.value) for keyword in shared_b.keywords
+    }
+    assert keywords == {
+        "tiler": "self.cta_tiler",
+        "coord": "(0, 0, None)",
+        "proj": "(None, 1, 1)",
+    }
+
+    partition_b = assignments["tSsB"]
+    assert isinstance(partition_b, ast.Call)
+    assert ast.unparse(partition_b.func) == "thr_mma.partition_B"
+    assert ast.unparse(partition_b.args[0]) == "sB"
+    assert "thr_mma.partition_B(panel_b)" not in ast.unparse(method)
+
+
+def test_sm80_c96_n48_panel_retains_one_ordered_accumulator() -> None:
+    """Guard the strict-FP32 K-tile order of the panel decomposition."""
+    tree = ast.parse(TILED_KERNEL_PATH.read_text(encoding="utf-8"))
+    parent = _method_node(tree, "_panel_adjoint_backward")
+    helper = _method_node(tree, "_backproject_panel_accumulate")
+    parent_source = ast.unparse(parent)
+    helper_source = ast.unparse(helper)
+
+    assert "for grid_tile in cutlass.range_constexpr(GRID_TILES)" in parent_source
+    assert parent_source.count("self._backproject_panel_accumulate(") == 2
+    assert "tCrOut_left.fill(0.0)" in parent_source
+    assert "tCrOut_right.fill(0.0)" in parent_source
+    assert ".fill(0.0)" not in helper_source
+    assert "panel_k_start = grid_tile * (TILE_M // self.tile_k)" in helper_source
+    assert "logical_k_tile = cutlass.Int32(0)" in helper_source
+    assert "logical_k_tile = logical_k_tile + 1" in helper_source
+
+
+def test_packed_message_grid_initializes_the_panel_adjoint_selector() -> None:
+    """Keep the manually constructed backward operation compile-complete."""
+    tree = ast.parse(MESSAGE_GRID_PATH.read_text(encoding="utf-8"))
+    factory = _method_node(tree, "_make_grid_operation")
+    assigned_attributes = {
+        target.attr
+        for node in ast.walk(factory)
+        if isinstance(node, ast.Assign)
+        for target in node.targets
+        if isinstance(target, ast.Attribute)
+        and isinstance(target.value, ast.Name)
+        and target.value.id == "operation"
+    }
+    assert "sm80_c96_n48_panel" in assigned_attributes
+    assert "channel_tile_start" in assigned_attributes
+    assert "panel_adjoint" not in assigned_attributes
+
+
+def _reference(left, right, to_grid, from_grid):
+    left_grid = torch.einsum("gj,njc->ngc", to_grid, left)
+    right_grid = torch.einsum("gj,njc->ngc", to_grid, right)
+    return torch.einsum("jg,ngc->njc", from_grid, left_grid * right_grid)
+
+
+def _reference_backward(grad_out, left, right, to_grid, from_grid):
+    grad_product = torch.einsum("jg,njc->ngc", from_grid, grad_out)
+    left_grid = torch.einsum("gj,njc->ngc", to_grid, left)
+    right_grid = torch.einsum("gj,njc->ngc", to_grid, right)
+    grad_left = torch.einsum("gj,ngc->njc", to_grid, grad_product * right_grid)
+    grad_right = torch.einsum("gj,ngc->njc", to_grid, grad_product * left_grid)
+    return grad_left, grad_right
+
+
+def _inputs(nodes: int, hidden_channels: int):
+    generator = torch.Generator(device="cuda").manual_seed(
+        20260703 + nodes + hidden_channels
+    )
+    left = 0.1 * torch.randn(
+        nodes,
+        PACKED_COEFF_DIM,
+        hidden_channels,
+        device="cuda",
+        generator=generator,
+    )
+    right = 0.1 * torch.randn(
+        left.shape,
+        device="cuda",
+        generator=generator,
+    )
+    to_grid = 0.1 * torch.randn(
+        GRID_SIZE,
+        PACKED_COEFF_DIM,
+        device="cuda",
+        generator=generator,
+    )
+    from_grid = 0.1 * torch.randn(
+        PACKED_COEFF_DIM,
+        GRID_SIZE,
+        device="cuda",
+        generator=generator,
+    )
+    return left, right, to_grid, from_grid
+
+
+def test_architecture_policy_shares_sm80_backend_with_sm86():
+    from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import (
+        PORTABLE_TILED_BACKEND,
+        PYTORCH_BACKEND,
+        output_grid_arch_key,
+        select_output_grid_backend,
+    )
+
+    for hidden_channels in SUPPORTED_HIDDEN_CHANNELS:
+        assert select_output_grid_backend((8, 0), hidden_channels) == (
+            PORTABLE_TILED_BACKEND
+        )
+        assert select_output_grid_backend((8, 6), hidden_channels) == (
+            PORTABLE_TILED_BACKEND
+        )
+        assert select_output_grid_backend((9, 0), hidden_channels) == (
+            PORTABLE_TILED_BACKEND
+        )
+        for capability in ((7, 5), (8, 9), (10, 0), (12, 0)):
+            assert (
+                select_output_grid_backend(capability, hidden_channels)
+                == PYTORCH_BACKEND
+            )
+    assert select_output_grid_backend((8, 0), 128) == PYTORCH_BACKEND
+    assert select_output_grid_backend((8, 6), 128) == PYTORCH_BACKEND
+    assert select_output_grid_backend((9, 0), 128) == PYTORCH_BACKEND
+    assert output_grid_arch_key((8, 0)) == "sm80"
+    assert output_grid_arch_key((8, 6)) == "sm80"
+    assert output_grid_arch_key((9, 0)) == "sm90"
+
+
+@pytest.mark.skipif(
+    _CUTE_SKIP_REASON is not None,
+    reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable",
+)
+class TestTiledOutputGridProductForwardCuda:
+    @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+    @pytest.mark.parametrize("nodes", [1, 7, 65])
+    def test_matches_strict_fp32_reference(
+        self,
+        nodes: int,
+        hidden_channels: int,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import (
+            run_tiled_output_grid_product,
+        )
+
+        left, right, to_grid, from_grid = _inputs(nodes, hidden_channels)
+        expected = _reference(left, right, to_grid, from_grid)
+        actual = run_tiled_output_grid_product(left, right, to_grid, from_grid)
+        torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+
+    @pytest.mark.skipif(
+        _SM80_SKIP_REASON is not None,
+        reason=_SM80_SKIP_REASON or "sm80-family GPU is unavailable",
+    )
+    @pytest.mark.parametrize("nodes", [1, 7, 65])
+    def test_sm80_c96_n48_matches_strict_fp32_reference(self, nodes: int):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import (
+            run_tiled_output_grid_product,
+        )
+
+        left, right, to_grid, from_grid = _inputs(nodes, 96)
+        expected = _reference(left, right, to_grid, from_grid)
+        actual = run_tiled_output_grid_product(
+            left,
+            right,
+            to_grid,
+            from_grid,
+            use_sm80_c96_n48=True,
+        )
+        torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+
+    @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+    def test_one_compile_accepts_symbolic_node_counts(self, hidden_channels: int):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import (
+            _compiled_tiled_forward,
+            run_tiled_output_grid_product,
+        )
+
+        before = _compiled_tiled_forward.cache_info()
+        for nodes in (3, 19):
+            left, right, to_grid, from_grid = _inputs(nodes, hidden_channels)
+            actual = run_tiled_output_grid_product(
+                left,
+                right,
+                to_grid,
+                from_grid,
+            )
+            torch.testing.assert_close(
+                actual,
+                _reference(left, right, to_grid, from_grid),
+                atol=TOL,
+                rtol=TOL,
+            )
+        after = _compiled_tiled_forward.cache_info()
+        assert after.misses - before.misses <= 1
+        assert after.hits > before.hits
+
+    @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+    def test_launches_on_current_non_default_stream(self, hidden_channels: int):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import (
+            run_tiled_output_grid_product,
+        )
+
+        left, right, to_grid, from_grid = _inputs(11, hidden_channels)
+        stream = torch.cuda.Stream()
+        with torch.cuda.stream(stream):
+            actual = run_tiled_output_grid_product(
+                left,
+                right,
+                to_grid,
+                from_grid,
+            )
+            expected = _reference(left, right, to_grid, from_grid)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL)
+
+
+@pytest.mark.skipif(
+    _CUTE_SKIP_REASON is not None,
+    reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable",
+)
+class TestTiledOutputGridProductBackwardCuda:
+    @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+    @pytest.mark.parametrize("nodes", [1, 7, 65])
+    def test_first_backward_matches_strict_fp32_reference(
+        self,
+        nodes: int,
+        hidden_channels: int,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import (
+            run_tiled_output_grid_product_backward,
+        )
+
+        left, right, to_grid, from_grid = _inputs(nodes, hidden_channels)
+        grad_out = torch.randn_like(left)
+        expected = _reference_backward(grad_out, left, right, to_grid, from_grid)
+        actual = run_tiled_output_grid_product_backward(
+            grad_out,
+            left,
+            right,
+            to_grid,
+            from_grid,
+        )
+        torch.testing.assert_close(actual[0], expected[0], atol=TOL, rtol=TOL)
+        torch.testing.assert_close(actual[1], expected[1], atol=TOL, rtol=TOL)
+
+    @pytest.mark.skipif(
+        _SM80_SKIP_REASON is not None,
+        reason=_SM80_SKIP_REASON or "sm80-family GPU is unavailable",
+    )
+    @pytest.mark.parametrize("nodes", [1, 7, 65])
+    def test_sm80_c96_n48_panel_matches_strict_fp32_reference(
+        self,
+        nodes: int,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import (
+            run_tiled_output_grid_product_backward,
+        )
+
+        left, right, to_grid, from_grid = _inputs(nodes, 96)
+        grad_out = torch.randn_like(left)
+        expected = _reference_backward(grad_out, left, right, to_grid, from_grid)
+        actual = run_tiled_output_grid_product_backward(
+            grad_out,
+            left,
+            right,
+            to_grid,
+            from_grid,
+            use_sm80_c96_n48_panel=True,
+        )
+        torch.testing.assert_close(actual[0], expected[0], atol=TOL, rtol=TOL)
+        torch.testing.assert_close(actual[1], expected[1], atol=TOL, rtol=TOL)
+
+    @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS)
+    def test_one_backward_compile_accepts_symbolic_node_counts(
+        self,
+        hidden_channels: int,
+    ):
+        from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import (
+            _compiled_tiled_backward,
+            run_tiled_output_grid_product_backward,
+        )
+
+        before = _compiled_tiled_backward.cache_info()
+        for nodes in (3, 19):
+            left, right, to_grid, from_grid = _inputs(nodes, hidden_channels)
+            grad_out = torch.randn_like(left)
+            expected = _reference_backward(grad_out, left, right, to_grid, from_grid)
+            actual = run_tiled_output_grid_product_backward(
+                grad_out,
+                left,
+                right,
+                to_grid,
+                from_grid,
+            )
+            torch.testing.assert_close(actual[0], expected[0], atol=TOL, rtol=TOL)
+            torch.testing.assert_close(actual[1], expected[1], atol=TOL, rtol=TOL)
+        after = _compiled_tiled_backward.cache_info()
+        assert after.misses - before.misses <= 1
+        assert after.hits > before.hits
diff --git a/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py b/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py
new file mode 100644
index 0000000000..edc9d32114
--- /dev/null
+++ b/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py
@@ -0,0 +1,363 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Focused contracts for Neo's C=96 block-FFN grid fusion."""
+
+from __future__ import (
+    annotations,
+)
+
+import pytest
+import torch
+
+from deepmd.pt.model.descriptor.sezm import (
+    DescrptSeZM,
+)
+from deepmd.pt.model.descriptor.sezm_nn.grid_net import (
+    GridBranch,
+    GridMLP,
+    S2GridNet,
+)
+
+
+def test_single_branch_bypasses_router_and_accepts_fused_middle(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    branch = (
+        GridBranch(
+            channels=2,
+            n_branches=1,
+            n_frames=3,
+            dtype=torch.float32,
+            trainable=False,
+            seed=7,
+        )
+        .to("cpu")
+        .eval()
+    )
+    monkeypatch.setattr(
+        branch.router,
+        "forward",
+        lambda value: pytest.fail("single-branch router must not run"),
+    )
+    calls = 0
+
+    def fused_middle(
+        left: torch.Tensor,
+        right: torch.Tensor,
+    ) -> torch.Tensor:
+        nonlocal calls
+        calls += 1
+        return left * right
+
+    left = torch.randn(2, 4, 1, 6, device="cpu")
+    out = branch(
+        left,
+        torch.randn_like(left),
+        torch.randn(2, 1, 4, device="cpu"),
+        to_grid=lambda value: value,
+        from_grid=lambda value: value,
+        pair_grid=fused_middle,
+    )
+
+    assert calls == 1
+    assert out.shape == left.shape
+
+
+def test_single_branch_keeps_pytorch_middle_fallback(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    branch = (
+        GridBranch(
+            channels=2,
+            n_branches=1,
+            n_frames=1,
+            dtype=torch.float32,
+            trainable=False,
+            seed=11,
+        )
+        .to("cpu")
+        .eval()
+    )
+    original_forward = branch.router.forward
+    router_calls = 0
+
+    def tracked_router(value: torch.Tensor) -> torch.Tensor:
+        nonlocal router_calls
+        router_calls += 1
+        return original_forward(value)
+
+    monkeypatch.setattr(branch.router, "forward", tracked_router)
+    calls = {"to_grid": 0, "from_grid": 0}
+
+    def to_grid(value: torch.Tensor) -> torch.Tensor:
+        calls["to_grid"] += 1
+        return value
+
+    def from_grid(value: torch.Tensor) -> torch.Tensor:
+        calls["from_grid"] += 1
+        return value
+
+    left = torch.randn(2, 4, 1, 2, device="cpu")
+    out = branch(
+        left,
+        torch.randn_like(left),
+        torch.randn(2, 1, 4, device="cpu"),
+        to_grid=to_grid,
+        from_grid=from_grid,
+    )
+
+    assert calls == {"to_grid": 2, "from_grid": 1}
+    assert router_calls == 1
+    assert out.shape == left.shape
+
+
+def test_single_branch_fused_pair_matches_routed_fallback(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    branch = (
+        GridBranch(
+            channels=2,
+            n_branches=1,
+            n_frames=1,
+            dtype=torch.float32,
+            trainable=False,
+            seed=17,
+        )
+        .to("cpu")
+        .eval()
+    )
+    left = torch.randn(2, 4, 1, 2, device="cpu")
+    right = torch.randn_like(left)
+    scalar_pair = torch.randn(2, 1, 4, device="cpu")
+    original_forward = branch.router.forward
+    router_calls = 0
+
+    def tracked_router(value: torch.Tensor) -> torch.Tensor:
+        nonlocal router_calls
+        router_calls += 1
+        return original_forward(value)
+
+    monkeypatch.setattr(branch.router, "forward", tracked_router)
+    routed = branch(
+        left,
+        right,
+        scalar_pair,
+        to_grid=lambda value: value,
+        from_grid=lambda value: value,
+    )
+    shortcut = branch(
+        left,
+        right,
+        scalar_pair,
+        to_grid=lambda value: value,
+        from_grid=lambda value: value,
+        pair_grid=lambda lhs, rhs: lhs * rhs,
+    )
+
+    assert router_calls == 1
+    torch.testing.assert_close(shortcut, routed)
+
+
+def test_multi_branch_preserves_softmax_router(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    branch = GridBranch(
+        channels=2,
+        n_branches=2,
+        n_frames=1,
+        dtype=torch.float32,
+        trainable=False,
+        seed=13,
+    ).to("cpu")
+    original_softmax = torch.softmax
+    softmax_calls = 0
+
+    def tracked_softmax(
+        value: torch.Tensor,
+        dim: int,
+    ) -> torch.Tensor:
+        nonlocal softmax_calls
+        softmax_calls += 1
+        return original_softmax(value, dim=dim)
+
+    monkeypatch.setattr(torch, "softmax", tracked_softmax)
+    left = torch.randn(2, 4, 1, 2, device="cpu")
+    out = branch(
+        left,
+        torch.randn_like(left),
+        torch.randn(2, 1, 4, device="cpu"),
+        to_grid=lambda value: value,
+        from_grid=lambda value: value,
+        pair_grid=lambda left, right: pytest.fail(
+            "multi-branch routing must keep the generic path"
+        ),
+    )
+
+    assert softmax_calls == 1
+    assert out.shape == left.shape
+    assert torch.isfinite(out).all()
+
+
+def test_single_branch_training_uses_frozen_router_fallback(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    branch = (
+        GridBranch(
+            channels=2,
+            n_branches=1,
+            n_frames=1,
+            dtype=torch.float32,
+            trainable=True,
+            seed=17,
+        )
+        .to("cpu")
+        .train()
+    )
+    original_softmax = torch.softmax
+    softmax_calls = 0
+
+    def tracked_softmax(
+        value: torch.Tensor,
+        dim: int,
+    ) -> torch.Tensor:
+        nonlocal softmax_calls
+        softmax_calls += 1
+        return original_softmax(value, dim=dim)
+
+    monkeypatch.setattr(torch, "softmax", tracked_softmax)
+    left = torch.randn(2, 4, 1, 2, device="cpu", requires_grad=True)
+    right = torch.randn_like(left)
+    out = branch(
+        left,
+        right,
+        torch.randn(2, 1, 4, device="cpu"),
+        to_grid=lambda value: value,
+        from_grid=lambda value: value,
+    )
+    out.sum().backward()
+
+    assert softmax_calls == 1
+    assert branch.router.weight.grad is None
+    assert left.grad is not None
+
+
+def test_single_branch_eval_with_trainable_parameters_uses_router_fallback() -> None:
+    branch = (
+        GridBranch(
+            channels=2,
+            n_branches=1,
+            n_frames=1,
+            dtype=torch.float32,
+            trainable=True,
+            seed=19,
+        )
+        .to("cpu")
+        .eval()
+    )
+    left = torch.randn(2, 4, 1, 2, device="cpu", requires_grad=True)
+    out = branch(
+        left,
+        torch.randn_like(left),
+        torch.randn(2, 1, 4, device="cpu"),
+        to_grid=lambda value: value,
+        from_grid=lambda value: value,
+    )
+    out.sum().backward()
+
+    assert branch.router.weight.grad is None
+    assert left.grad is not None
+
+
+@pytest.mark.parametrize("training", (False, True))
+def test_grid_net_with_trainable_parameters_keeps_differentiable_path(
+    training: bool,
+) -> None:
+    net = S2GridNet(
+        lmax=1,
+        mmax=1,
+        channels=2,
+        n_focus=1,
+        mode="self",
+        op_type="mlp",
+        dtype=torch.float32,
+        layout="ndfc",
+        coefficient_layout="packed",
+        grid_method="e3nn",
+        trainable=True,
+        seed=23,
+    ).to("cpu")
+    net.train(training)
+
+    query = torch.randn(2, 4, 1, 4, device="cpu", requires_grad=True)
+    net(query).sum().backward()
+    assert query.grad is not None
+
+
+def test_frozen_eval_grid_net_offers_fused_product(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    net = (
+        S2GridNet(
+            lmax=1,
+            mmax=1,
+            channels=2,
+            n_focus=1,
+            mode="self",
+            op_type="mlp",
+            dtype=torch.float32,
+            layout="ndfc",
+            coefficient_layout="packed",
+            grid_method="e3nn",
+            trainable=False,
+            seed=29,
+        )
+        .to("cpu")
+        .eval()
+    )
+    calls = 0
+
+    def tracked_product(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
+        nonlocal calls
+        calls += 1
+        return net._from_grid(net._to_grid(left) * net._to_grid(right))
+
+    monkeypatch.setattr(net, "_pair_grid", tracked_product)
+    output = net(torch.randn(2, 4, 1, 4, device="cpu"))
+
+    assert calls == 1
+    assert output.shape == (2, 4, 1, 2)
+
+
+def test_exact_neo_has_two_c96_block_products_and_c192_readout() -> None:
+    descriptor = DescrptSeZM(
+        ntypes=2,
+        sel=4,
+        channels=32,
+        lmax=3,
+        mmax=1,
+        n_blocks=2,
+        so2_layers=3,
+        n_focus=2,
+        message_node_so3=True,
+        ffn_neurons=0,
+        ffn_so3_grid=True,
+        grid_branch=[0, 0, 1],
+        ffn_blocks=1,
+        so3_readout="mlp",
+        use_amp=False,
+        precision="float32",
+        trainable=False,
+        seed=42,
+    )
+
+    assert len(descriptor.blocks) == 2
+    for block in descriptor.blocks:
+        assert len(block.ffns) == 1
+        grid_op = block.ffns[0].act.grid_op
+        assert isinstance(grid_op, GridBranch)
+        assert grid_op.n_branches == 1
+        assert grid_op.channels == 96
+
+    readout_grid_op = descriptor.output_ffn.act.grid_op
+    assert isinstance(readout_grid_op, GridMLP)
+    assert readout_grid_op.hidden_channels == 192
diff --git a/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py b/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py
new file mode 100644
index 0000000000..02327ab143
--- /dev/null
+++ b/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py
@@ -0,0 +1,246 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""CPU algebra checks for the Neo degree-zero readout boundary."""
+
+from __future__ import (
+    annotations,
+)
+
+import torch
+
+COEFF_DIM = 16
+N_FRAMES = 3
+GRID_SIZE = 152
+HIDDEN_CHANNELS = 192
+
+
+def _readout_reference(left, right, to_grid, from_grid):
+    left_grid = torch.einsum("gj,njh->ngh", to_grid, left)
+    right_grid = torch.einsum("gj,njh->ngh", to_grid, right)
+    return torch.einsum("g,ngh->nh", from_grid[0], left_grid * right_grid)
+
+
+def _gram_reference(left, right, gram):
+    transformed_right = torch.einsum("ij,njh->nih", gram, right)
+    return torch.sum(left * transformed_right, dim=1)
+
+
+def _output_ffn():
+    from deepmd.pt.model.descriptor.sezm_nn.ffn import (
+        EquivariantFFN,
+    )
+
+    return (
+        EquivariantFFN(
+            lmax=3,
+            channels=32,
+            hidden_channels=96,
+            kmax=1,
+            grid_mlp=True,
+            grid_branch=0,
+            dtype=torch.float32,
+            s2_activation=False,
+            ffn_so3_grid=True,
+            activation_function="silu",
+            glu_activation=True,
+            mlp_bias=False,
+            trainable=False,
+            seed=17,
+        )
+        .to("cpu")
+        .eval()
+    )
+
+
+def test_final_slice_has_exactly_zero_l_greater_than_zero_output_cotangent():
+    module = _output_ffn()
+    value = torch.randn(
+        2,
+        COEFF_DIM,
+        1,
+        32,
+        device="cpu",
+        requires_grad=True,
+    )
+    ffn_out = module(value)
+    seed = torch.randn(2, 32, device="cpu")
+
+    cotangent = torch.autograd.grad(
+        (value + ffn_out)[:, 0, 0, :],
+        ffn_out,
+        seed,
+    )[0]
+
+    assert torch.equal(cotangent[:, 0, 0, :], seed)
+    assert torch.count_nonzero(cotangent[:, 1:, :, :]) == 0
+
+
+def test_degree_zero_nonzero_frame_projector_rows_are_structural_zero():
+    projector = _output_ffn().act.projector
+    from_grid = projector.from_grid_mat.reshape(
+        COEFF_DIM,
+        N_FRAMES,
+        GRID_SIZE,
+    )
+
+    assert projector.frame_set == [0, -1, 1]
+    assert torch.count_nonzero(from_grid[0, 0]) > 0
+    assert torch.count_nonzero(from_grid[0, 1]) == 0
+    assert torch.count_nonzero(from_grid[0, 2]) == 0
+
+
+def test_scalar_readout_matches_row_zero_of_generic_grid_product():
+    generator = torch.Generator().manual_seed(20260704)
+    left = 0.1 * torch.randn(
+        3,
+        COEFF_DIM * N_FRAMES,
+        HIDDEN_CHANNELS,
+        device="cpu",
+        generator=generator,
+    )
+    right = 0.1 * torch.randn(
+        left.shape,
+        device="cpu",
+        generator=generator,
+    )
+    projector = _output_ffn().act.projector
+    to_grid = projector.to_grid_mat
+    from_grid = projector.from_grid_mat
+
+    left_grid = torch.einsum("gj,njh->ngh", to_grid, left)
+    right_grid = torch.einsum("gj,njh->ngh", to_grid, right)
+    generic = torch.einsum(
+        "jg,ngh->njh",
+        from_grid,
+        left_grid * right_grid,
+    )
+
+    actual = _readout_reference(left, right, to_grid, from_grid)
+
+    assert torch.allclose(actual, generic[:, 0, :], atol=5.0e-5, rtol=5.0e-5)
+
+
+def test_dense_gram_matches_both_projected_input_adjoints():
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import (
+        build_readout_l0_gram,
+    )
+
+    generator = torch.Generator().manual_seed(20260718)
+    left = 0.1 * torch.randn(
+        2,
+        COEFF_DIM * N_FRAMES,
+        HIDDEN_CHANNELS,
+        device="cpu",
+        generator=generator,
+    )
+    right = 0.1 * torch.randn(left.shape, device="cpu", generator=generator)
+    dq0 = 0.1 * torch.randn(
+        left.shape[0],
+        HIDDEN_CHANNELS,
+        device="cpu",
+        generator=generator,
+    )
+    projector = _output_ffn().act.projector
+    gram = build_readout_l0_gram(
+        projector.to_grid_mat,
+        projector.from_grid_mat,
+    )
+
+    left_ref = left.detach().clone().requires_grad_(True)
+    right_ref = right.detach().clone().requires_grad_(True)
+    q0 = _readout_reference(
+        left_ref,
+        right_ref,
+        projector.to_grid_mat,
+        projector.from_grid_mat,
+    )
+    expected_left, expected_right = torch.autograd.grad(
+        q0,
+        (left_ref, right_ref),
+        dq0,
+    )
+    actual_left = dq0[:, None, :] * torch.einsum(
+        "ij,njh->nih",
+        gram,
+        right,
+    )
+    actual_right = dq0[:, None, :] * torch.einsum(
+        "ji,njh->nih",
+        gram,
+        left,
+    )
+
+    assert gram.shape == (COEFF_DIM * N_FRAMES, COEFF_DIM * N_FRAMES)
+    assert gram.dtype == torch.float32
+    assert gram.is_contiguous()
+    assert not gram.requires_grad
+    torch.testing.assert_close(actual_left, expected_left, atol=5.0e-5, rtol=5.0e-5)
+    torch.testing.assert_close(
+        actual_right,
+        expected_right,
+        atol=5.0e-5,
+        rtol=5.0e-5,
+    )
+
+
+def test_dense_gram_forward_matches_row_zero_grid_projection():
+    from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import (
+        build_readout_l0_gram,
+    )
+
+    generator = torch.Generator().manual_seed(20260720)
+    left = 0.1 * torch.randn(
+        3,
+        COEFF_DIM * N_FRAMES,
+        HIDDEN_CHANNELS,
+        device="cpu",
+        generator=generator,
+    )
+    right = 0.1 * torch.randn(left.shape, device="cpu", generator=generator)
+    projector = _output_ffn().act.projector
+    gram = build_readout_l0_gram(
+        projector.to_grid_mat,
+        projector.from_grid_mat,
+    )
+
+    expected = _readout_reference(
+        left,
+        right,
+        projector.to_grid_mat,
+        projector.from_grid_mat,
+    )
+    actual = _gram_reference(left, right, gram)
+
+    torch.testing.assert_close(actual, expected, atol=5.0e-5, rtol=5.0e-5)
+
+
+def test_dense_gram_forward_is_channelwise_without_cross_channel_mixing():
+    generator = torch.Generator().manual_seed(20260721)
+    left = torch.randn(
+        2,
+        COEFF_DIM * N_FRAMES,
+        HIDDEN_CHANNELS,
+        device="cpu",
+        generator=generator,
+    )
+    right = torch.randn(left.shape, device="cpu", generator=generator)
+    gram = torch.randn(
+        COEFF_DIM * N_FRAMES,
+        COEFF_DIM * N_FRAMES,
+        device="cpu",
+        generator=generator,
+    )
+    changed_channel = 37
+
+    baseline = _gram_reference(left, right, gram)
+    changed_right = right.clone()
+    changed_right[:, :, changed_channel].mul_(1.5)
+    changed = _gram_reference(left, changed_right, gram)
+
+    unaffected = torch.ones(HIDDEN_CHANNELS, dtype=torch.bool, device="cpu")
+    unaffected[changed_channel] = False
+    torch.testing.assert_close(changed[:, unaffected], baseline[:, unaffected])
+    torch.testing.assert_close(
+        changed[:, changed_channel],
+        1.5 * baseline[:, changed_channel],
+    )
diff --git a/source/tests/pt/model/test_descriptor_sezm_train_paths.py b/source/tests/pt/model/test_descriptor_sezm_train_paths.py
index cdf53b1d5e..7f62809d43 100644
--- a/source/tests/pt/model/test_descriptor_sezm_train_paths.py
+++ b/source/tests/pt/model/test_descriptor_sezm_train_paths.py
@@ -143,34 +143,34 @@ def test_triton_mode_gate_binds_each_stage(
         assert conv.triton_infer_level == infer_level
         # Rotation and flash wrappers retain their eager implementations when
         # Triton is unavailable, so their binding follows the mode gates alone.
-        assert (conv._rotate_to_local_fn is not None) is requested
-        assert (conv._rotate_back_fn is not None) is requested
-        assert (conv._flash_atten_fn is not None) is requested
-        assert conv._flash_atten_trains is (requested and training)
+        assert (conv.triton_l_1_rotate_to_local is not None) is requested
+        assert (conv.triton_l_1_rotate_back is not None) is requested
+        assert (conv.flash_attention is not None) is requested
+        assert conv.flash_attention_supports_training is (requested and training)
         # Segment softmax uses fp32 accumulation and preserves fp64 compute
         # through the descriptor's reference path.
-        assert (conv._segment_softmax_fn is not None) is (
+        assert (conv.triton_l_1_segment_softmax is not None) is (
             requested and SEGMENT_SOFTMAX_TRITON_AVAILABLE and precision == "float32"
         )
         # The rotate-mix front end is bound by a profitability bound on the
         # hidden width, which this narrow block sits below.
         assert conv.hidden_channels < 128
-        assert conv._triton_rotate_mix is None
+        assert conv.triton_l_1_rotate_mix is None
         # The CUDA gate is off, so the value stream stays on the stages.
-        assert conv._cuda_value_train is None
+        assert conv.cuda_train_value is None
 
     for module in descriptor.modules():
         if isinstance(module, SO2Linear):
             # The fused GEMM additionally needs every |m| block width to align
             # to its BN=64 tile, which a narrow block does not satisfy.
             aligned = slices_supported(module._block_diag_slices)
-            assert (module._block_diag_gemm is not None) is (
+            assert (module.triton_l_1_block_diag_gemm is not None) is (
                 requested and SO2_BLOCK_GEMM_TRITON_AVAILABLE and aligned
             )
         if isinstance(module, DynamicRadialDegreeMixer):
             # The callable contains its eager fallback, so construction binds it
             # whenever either mode requests the stage.
-            assert (module._radial_mix_block is not None) is requested
+            assert (module.triton_l_1_radial_mix is not None) is requested
         if isinstance(module, GatedActivation):
             assert module.triton_train_level == train_level
             assert module.triton_infer_level == infer_level
@@ -212,13 +212,13 @@ def test_cuda_train_gate_binds_the_value_stream(monkeypatch) -> None:
     ]
     assert convolutions
     for conv in convolutions:
-        assert conv._cuda_value_train is not None
+        assert conv.cuda_train_value is not None
         # The two layers are independent: the CUDA value stream does not
         # switch on any Triton stage, and the attention span stays dense
         # until the Triton gate asks for it.
         assert conv.triton_train_level == 0
-        assert conv._segment_softmax_fn is None
-        assert conv._flash_atten_trains is False
+        assert conv.triton_l_1_segment_softmax is None
+        assert conv.flash_attention_supports_training is False
 
 
 def test_cuda_triton_train_reuses_packed_wigner_runs(monkeypatch) -> None:
@@ -232,14 +232,14 @@ def test_cuda_triton_train_reuses_packed_wigner_runs(monkeypatch) -> None:
     monkeypatch.setenv("DP_TRITON_TRAIN", "1")
 
     descriptor = _make_descriptor(2, [20], 4.0).train()
-    assert descriptor._packed_wigner_train
+    assert descriptor.cuda_train_covers_all_blocks
     assert not descriptor._build_full_wigner()
 
     for block in descriptor.blocks:
         conv = block.so2_conv
-        assert conv._cuda_value_train is not None
-        assert conv._flash_atten_fn is not None
-        assert conv._flash_atten_trains
+        assert conv.cuda_train_value is not None
+        assert conv.flash_attention is not None
+        assert conv.flash_attention_supports_training
 
 
 @pytest.mark.parametrize("gate_name", ["DP_TRITON_TRAIN", "DP_TRITON_INFER"])
@@ -264,7 +264,7 @@ def test_grid_pair_train_follows_its_gate_and_slot_bound(
             and GRID_PAIR_TRITON_AVAILABLE
             and slots >= 75
         )
-        assert (net._grid_pair_train_fn is not None) is expected
+        assert (net.triton_train_l_1_grid_pair is not None) is expected
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
@@ -328,11 +328,11 @@ def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None:
             for module in accelerated.modules()
             if isinstance(module, SO2Convolution)
         )
-        if conv._cuda_conv_fn is None:
+        if conv.cuda_infer_l_2_conv is None:
             pytest.skip("the descriptor layout has no fused CUDA convolution")
-        assert conv._flash_atten_fn is not None
-        assert conv._cuda_value_train is None
-        assert not accelerated._wigner_free_conv
+        assert conv.flash_attention is not None
+        assert conv.cuda_train_value is None
+        assert not accelerated.cuda_infer_l_2_covers_all_blocks
         assert accelerated._build_full_wigner()
 
         output, gradient = self._inference_step(accelerated)
@@ -364,9 +364,9 @@ def test_training_step_matches_the_dense_path(self, monkeypatch, path) -> None:
             module for module in fused.modules() if isinstance(module, SO2Convolution)
         )
         if path in ("cuda", "cuda-triton"):
-            assert conv._cuda_value_train is not None
+            assert conv.cuda_train_value is not None
         if path in ("triton", "cuda-triton"):
-            assert conv._segment_softmax_fn is not None
+            assert conv.triton_l_1_segment_softmax is not None
         fused_objective, fused_gradient = self._step(fused)
 
         np.testing.assert_allclose(
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index f359f073d6..ffc5d75931 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -569,7 +569,7 @@ def test_reference_matches_module_eager_path(self):
             with self.subTest(lmax=lmax):
                 mixer = self._mixer(lmax)
                 # Force the dense scatter path regardless of the ambient flag.
-                mixer._radial_mix_block = None
+                mixer.triton_l_1_radial_mix = None
                 x_local, radial_feat, compact = self._inputs(mixer, seed=lmax)
                 with torch.no_grad():
                     module_out = mixer(x_local, radial_feat)
@@ -935,7 +935,7 @@ def test_prepared_weights_follow_device_without_entering_state_dict(self) -> Non
         conv = self._build_conv(*self.CASES[1]).cpu()
         value_path = make_triton_value_path(conv)
         self.assertIsNotNone(value_path)
-        conv._triton_value_path = value_path
+        conv.triton_infer_l_2_value = value_path
 
         buffer_names = (
             "_triton_w0_all",
@@ -1093,7 +1093,7 @@ def test_wigner_calculator_matches_reference_chain(self):
                         .to("cuda")
                         .eval()
                     )
-                self.assertTrue(fused_calc._use_triton_monomials)
+                self.assertTrue(fused_calc.triton_infer_l_1_monomials)
                 got, got_t = fused_calc(q)
                 want, want_t = ref_calc(q)
                 torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-5)
diff --git a/source/tests/pt/model/test_mlp.py b/source/tests/pt/model/test_mlp.py
index 5f067cb0e6..def412d0be 100644
--- a/source/tests/pt/model/test_mlp.py
+++ b/source/tests/pt/model/test_mlp.py
@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import itertools
+import os
 import unittest
+from unittest import (
+    mock,
+)
 
 import numpy as np
 import torch
@@ -11,6 +15,7 @@
     NativeLayer,
     NativeNet,
 )
+from deepmd.pt.model.network import mlp as mlp_module
 from deepmd.pt.model.network.mlp import (
     MLP,
     EmbeddingNet,
@@ -84,6 +89,114 @@ def test_jit(self) -> None:
             ml1 = MLPLayer.deserialize(ml.serialize())
             model = torch.jit.script(ml1)
 
+    def test_thin_so2_eval_linear_traces_without_addmm(self) -> None:
+        from torch.fx.experimental.proxy_tensor import (
+            make_fx,
+        )
+
+        layer = MLPLayer(
+            5,
+            8,
+            bias=True,
+            activation_function="none",
+            precision="float32",
+            trainable=False,
+        ).to(env.DEVICE)
+        layer.eval()
+        mlp_module.enable_neo_cute_compile_visible_linears(layer)
+        value = torch.randn(7, 5, device=env.DEVICE, dtype=torch.float32)
+        environment = {
+            "DP_CUTE_INFER": "1",
+            "DP_CUTE_SO2_THIN_WRAPPER": "1",
+        }
+
+        with mock.patch.dict(os.environ, environment, clear=False):
+            graph = make_fx(layer)(value)
+            with mock.patch.object(
+                mlp_module.F,
+                "linear",
+                side_effect=AssertionError("thin eval path reached aten::linear"),
+            ):
+                actual = layer(value)
+
+        expected = mlp_module.F.linear(value, layer.matrix.t(), layer.bias)
+        torch.testing.assert_close(actual, expected)
+        targets = {
+            node.target for node in graph.graph.nodes if node.op == "call_function"
+        }
+        self.assertIn(torch.ops.aten.mm.default, targets)
+        self.assertNotIn(torch.ops.aten.addmm.default, targets)
+
+    def test_thin_so2_eval_linear_is_scoped_to_marked_model(self) -> None:
+        layer = MLPLayer(
+            5,
+            8,
+            bias=True,
+            activation_function="none",
+            precision="float32",
+            trainable=False,
+        ).to(env.DEVICE)
+        layer.eval()
+        value = torch.randn(7, 5, device=env.DEVICE, dtype=torch.float32)
+        environment = {
+            "DP_CUTE_INFER": "1",
+            "DP_CUTE_SO2_THIN_WRAPPER": "1",
+        }
+
+        with (
+            mock.patch.dict(os.environ, environment, clear=False),
+            mock.patch.object(
+                mlp_module,
+                "_matmul_bias",
+                side_effect=AssertionError("unmarked MLP reached Neo-only topology"),
+            ),
+        ):
+            actual = layer(value)
+
+        expected = mlp_module.F.linear(value, layer.matrix.t(), layer.bias)
+        torch.testing.assert_close(actual, expected)
+
+    def test_thin_so2_linear_matches_shared_device_policy(self) -> None:
+        from deepmd.pt_expt.kernels.cute.sezm import (
+            runtime_policy,
+        )
+
+        with (
+            mock.patch.dict(
+                os.environ,
+                {"DP_CUTE_INFER": "1"},
+                clear=True,
+            ),
+            mock.patch.object(torch.cuda, "is_available", return_value=True),
+        ):
+            for capability in ((8, 0), (8, 6), (8, 9), (9, 0), (12, 0)):
+                with (
+                    self.subTest(capability=capability),
+                    mock.patch.object(
+                        torch.cuda, "get_device_capability", return_value=capability
+                    ) as get_capability,
+                ):
+                    device = torch.device("cuda", 1)
+                    self.assertEqual(
+                        mlp_module._use_so2_compile_visible_linear(device),
+                        runtime_policy.is_so2_thin_wrapper_enabled(capability),
+                    )
+                    get_capability.assert_called_once_with(device)
+
+        with (
+            mock.patch.dict(
+                os.environ,
+                {
+                    "DP_CUTE_INFER": "1",
+                    "DP_CUTE_SO2_THIN_WRAPPER": "0",
+                },
+                clear=True,
+            ),
+            mock.patch.object(torch.cuda, "is_available", return_value=True),
+            mock.patch.object(torch.cuda, "get_device_capability", return_value=(8, 0)),
+        ):
+            self.assertFalse(mlp_module._use_so2_compile_visible_linear())
+
 
 class TestMLP(unittest.TestCase):
     def setUp(self) -> None:
diff --git a/source/tests/pt/model/test_nlist_backend.py b/source/tests/pt/model/test_nlist_backend.py
index d38262dacd..6d243a6b4e 100644
--- a/source/tests/pt/model/test_nlist_backend.py
+++ b/source/tests/pt/model/test_nlist_backend.py
@@ -10,6 +10,9 @@
 """
 
 import copy
+from types import (
+    SimpleNamespace,
+)
 
 import numpy as np
 import pytest
@@ -172,6 +175,55 @@ def test_self_built_model_forces_native(pt_files, monkeypatch) -> None:
         assert deep_eval._nlist_builder is None
 
 
+def test_edge_strategy_freezes_parameters_only_during_inference_call() -> None:
+    from deepmd.pt.infer.deep_eval import DeepEval as PTDeepEval
+
+    class ProbeModel(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.ones((), device="cpu"))
+            self.saw_frozen = False
+
+        def get_sel(self) -> list[int]:
+            return [4]
+
+        def forward_common_lower(self, *args, **kwargs) -> dict[str, torch.Tensor]:
+            self.saw_frozen = not self.weight.requires_grad
+            return {}
+
+    class ProbeBuilder:
+        def build(self, coord, atype, *args, **kwargs) -> SimpleNamespace:
+            return SimpleNamespace(
+                coord=coord,
+                atype=atype,
+                edge_index=torch.empty((2, 0), dtype=torch.long, device="cpu"),
+                edge_vec=torch.empty((0, 3), dtype=coord.dtype, device="cpu"),
+                edge_scatter_index=torch.empty((2, 0), dtype=torch.long, device="cpu"),
+                edge_mask=torch.empty((0,), dtype=torch.bool, device="cpu"),
+            )
+
+    inner = ProbeModel()
+    deep_eval = object.__new__(PTDeepEval)
+    deep_eval.dp = ModelWrapper(inner)
+    deep_eval._uses_edge_schema = True
+    deep_eval._nlist_builder = ProbeBuilder()
+    deep_eval.rcut = 4.0
+
+    PTDeepEval._eval_lower_strategy(
+        deep_eval,
+        torch.zeros((1, 1, 3), device="cpu"),
+        torch.zeros((1, 1), dtype=torch.long, device="cpu"),
+        None,
+        None,
+        None,
+        None,
+        False,
+    )
+
+    assert inner.saw_frozen
+    assert inner.weight.requires_grad
+
+
 # --- equivalence with the native dense builder ------------------------------
 
 
diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py
index ef0beb3c68..b44df52f28 100644
--- a/source/tests/pt/model/test_sezm_model.py
+++ b/source/tests/pt/model/test_sezm_model.py
@@ -40,6 +40,7 @@
 from deepmd.pt.model.model.sezm_model import (
     InnerPotential,
     SeZMModel,
+    _sort_edge_tensors_by_destination,
 )
 from deepmd.pt.model.model.sezm_native_spin_model import (
     SeZMNativeSpinModel,
@@ -230,6 +231,39 @@ def setUp(self) -> None:
         self.device = env.DEVICE
         torch.manual_seed(2024)
 
+    def test_destination_edge_sort_is_empty_safe_and_keeps_alignment(self) -> None:
+        empty_index = torch.empty((2, 0), dtype=torch.long, device="cpu")
+        empty_vec = torch.empty((0, 3), device="cpu")
+        empty_mask = torch.empty((0,), dtype=torch.bool, device="cpu")
+        empty_scatter = torch.empty((2, 0), dtype=torch.long, device="cpu")
+
+        empty_result = _sort_edge_tensors_by_destination(
+            empty_index,
+            empty_vec,
+            empty_mask,
+            empty_scatter,
+        )
+        self.assertEqual(tuple(empty_result[0].shape), (2, 0))
+
+        edge_index = torch.tensor([[2, 0, 1, 0], [1, 0, 1, 0]], device="cpu")
+        edge_vec = torch.arange(12, dtype=torch.float32, device="cpu").view(4, 3)
+        edge_mask = torch.tensor([True, False, True, True], device="cpu")
+        edge_scatter = edge_index + 10
+        sorted_index, sorted_vec, sorted_mask, sorted_scatter = (
+            _sort_edge_tensors_by_destination(
+                edge_index,
+                edge_vec,
+                edge_mask,
+                edge_scatter,
+            )
+        )
+        permutation = torch.tensor([1, 3, 2, 0], device="cpu")
+
+        self.assertTrue(torch.equal(sorted_index, edge_index[:, permutation]))
+        self.assertTrue(torch.equal(sorted_vec, edge_vec[permutation]))
+        self.assertTrue(torch.equal(sorted_mask, edge_mask[permutation]))
+        self.assertTrue(torch.equal(sorted_scatter, edge_scatter[:, permutation]))
+
     @staticmethod
     def _randomize_params(model: torch.nn.Module, seed: int = 1234) -> None:
         """Fill all parameters with small random values.
@@ -429,6 +463,111 @@ def _make_frame_with_natoms(
         )
         return coord, atype, box
 
+    def test_neo_cute_flag_skips_ineligible_kernel(self) -> None:
+        """Reject an unsupported model before entering the per-block CuTe path."""
+        from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as cute_so2
+
+        cpu = torch.device("cpu")
+        with (
+            mock.patch.object(env, "DEVICE", cpu),
+            mock.patch.object(self, "device", cpu),
+            torch.device(cpu),
+        ):
+            coord, atype, _, _, _, _ = self._make_tiny_frame()
+            params = self._build_model_params(use_compile=False)
+            params["descriptor"]["seed"] = None
+            params["descriptor"]["use_env_seed"] = False
+            params["fitting_net"]["seed"] = None
+            model = get_sezm_model(params).to(device=cpu)
+            self._randomize_params(model)
+            model.eval()
+            for parameter in model.parameters():
+                parameter.requires_grad_(False)
+
+            extended_coord, extended_atype, mapping, nlist = (
+                extend_input_and_build_neighbor_list(
+                    coord,
+                    atype,
+                    model.get_rcut(),
+                    model.get_sel(),
+                    mixed_types=model.mixed_types(),
+                    box=None,
+                )
+            )
+            edge = edge_schema_from_extended(
+                extended_coord,
+                extended_atype,
+                nlist,
+                mapping,
+            )
+
+            def run_lower() -> dict[str, torch.Tensor]:
+                return model.forward_common_lower(
+                    edge.coord,
+                    edge.atype,
+                    edge.edge_index,
+                    edge.edge_vec,
+                    edge.edge_scatter_index,
+                    edge.edge_mask,
+                )
+
+            with mock.patch.dict(
+                os.environ,
+                {"DP_CUTE_INFER": "0"},
+                clear=False,
+            ):
+                reference = run_lower()
+            with (
+                mock.patch.dict(
+                    os.environ,
+                    {"DP_CUTE_INFER": "1"},
+                    clear=False,
+                ),
+                mock.patch.object(
+                    cute_so2,
+                    "maybe_run_cute_so2",
+                    wraps=cute_so2.maybe_run_cute_so2,
+                ) as maybe_run,
+            ):
+                actual = run_lower()
+
+            self.assertEqual(maybe_run.call_count, 0)
+            self.assertEqual(actual.keys(), reference.keys())
+            for name, expected in reference.items():
+                if not isinstance(expected, torch.Tensor):
+                    continue
+                self.assertTrue(
+                    torch.equal(actual[name], expected),
+                    msg=f"DP_CUTE_INFER changed ineligible-model {name}",
+                )
+
+    def test_neo_cute_sort_guard_uses_post_cast_geometry_dtype(self) -> None:
+        """Raw FP64 coordinates do not hide an otherwise eligible FP32 SO2."""
+        from deepmd.pt.model.model import (
+            sezm_model,
+        )
+
+        descriptor = mock.Mock()
+        descriptor.compute_dtype = torch.float32
+        descriptor.is_cute_infer_packed_wigner_candidate.return_value = True
+        device = torch.device("cuda")
+
+        with mock.patch.object(
+            sezm_model, "_neo_cute_infer_enabled", return_value=True
+        ):
+            actual = sezm_model._neo_cute_so2_requires_sorted_edges(
+                descriptor,
+                training=False,
+                device=device,
+            )
+
+        self.assertTrue(actual)
+        descriptor.is_cute_infer_packed_wigner_candidate.assert_called_once_with(
+            device,
+            torch.float32,
+            torch.float32,
+        )
+
     def test_trace_pad_dim_trim_returns_contiguous(self) -> None:
         """Trimmed trace inputs stay contiguous so strides mirror runtime layout.
 
@@ -590,6 +729,105 @@ def test_compile_cache_slots_and_eval_shape_change(self) -> None:
             model_cmp.compiled_core_compute_cache[eval_key], callable_eval_first
         )
 
+    def test_load_state_dict_invalidates_local_and_shared_compile_caches(self) -> None:
+        """Checkpoint loads must discard every callable that captured old state."""
+        from deepmd.pt.model.model import sezm_model as sezm_model_module
+
+        model = get_sezm_model(self._build_model_params(use_compile=True))
+        self.assertTrue(
+            all(
+                getattr(hook, "__self__", None) is None
+                for hook in model._load_state_dict_post_hooks.values()
+            )
+        )
+        local_callable = object()
+        shared_callable = object()
+        model.compiled_core_compute_cache[(False, False)] = local_callable
+        model._task_buf_order_cache[(False, False)] = ("task",)
+        object.__setattr__(model, "compiled_embedding", local_callable)
+        object.__setattr__(model, "_embedding_task_buf_order", ("task",))
+        object.__setattr__(model, "compiled_dens_compute", local_callable)
+        model._dens_compiled = True
+        model._deepmd_cute_so2_state = False
+
+        with (
+            mock.patch.dict(
+                sezm_model_module._SEZM_COMPILE_CACHE,
+                {("shared",): shared_callable},
+                clear=True,
+            ),
+            mock.patch.dict(
+                sezm_model_module._SEZM_TASK_BUF_ORDER,
+                {("shared",): ("task",)},
+                clear=True,
+            ),
+        ):
+            model.load_state_dict(model.state_dict())
+            self.assertEqual(model.compiled_core_compute_cache, {})
+            self.assertEqual(model._task_buf_order_cache, {})
+            self.assertIsNone(model.compiled_embedding)
+            self.assertIsNone(model._embedding_task_buf_order)
+            self.assertIsNone(model.compiled_dens_compute)
+            self.assertFalse(model._dens_compiled)
+            self.assertFalse(hasattr(model, "_deepmd_cute_so2_state"))
+            self.assertEqual(sezm_model_module._SEZM_COMPILE_CACHE, {})
+            self.assertEqual(sezm_model_module._SEZM_TASK_BUF_ORDER, {})
+
+    @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
+    def test_eval_compile_retraces_after_load_state_dict(self) -> None:
+        """The public eval path must not reuse a make_fx graph after reloading."""
+        coord, atype, box, _, _, _ = self._make_tiny_frame()
+        initial = get_sezm_model(self._build_model_params(use_compile=False))
+        replacement = get_sezm_model(self._build_model_params(use_compile=False))
+        self._randomize_params(initial, seed=1234)
+        self._randomize_params(replacement, seed=5678)
+        with mock.patch.dict(os.environ, {"DP_COMPILE_INFER": "1"}, clear=False):
+            compiled = get_sezm_model(self._build_model_params(use_compile=True))
+        compiled.load_state_dict(initial.state_dict())
+        initial.eval()
+        replacement.eval()
+        compiled.eval()
+
+        # Exercise the public make_fx/AOT cache lifecycle without depending on
+        # Inductor's dynamic-View lowering for this intentionally tiny frame.
+        with mock.patch(
+            "torch._inductor.compile_fx.compile_fx_inner",
+            side_effect=lambda graph, inputs: graph.forward,
+        ):
+            before = compiled(coord, atype, box=box)
+            old_callable = compiled.compiled_core_compute_cache[(False, False)]
+
+            compiled.load_state_dict(replacement.state_dict())
+            self.assertEqual(compiled.compiled_core_compute_cache, {})
+            expected = replacement(coord, atype, box=box)
+            actual = compiled(coord, atype, box=box)
+        self.assertIsNot(
+            compiled.compiled_core_compute_cache[(False, False)],
+            old_callable,
+        )
+        self.assertFalse(torch.equal(actual["energy"], before["energy"]))
+        _assert_close_with_strict_warning(
+            actual["energy"],
+            expected["energy"],
+            atol=1.0e-6,
+            rtol=1.0e-6,
+            msg="eval energy mismatch after checkpoint reload",
+        )
+        _assert_close_with_strict_warning(
+            actual["force"],
+            expected["force"],
+            atol=1.0e-6,
+            rtol=1.0e-6,
+            msg="eval force mismatch after checkpoint reload",
+        )
+        _assert_close_with_strict_warning(
+            actual["virial"],
+            expected["virial"],
+            atol=1.0e-5,
+            rtol=1.0e-5,
+            msg="eval virial mismatch after checkpoint reload",
+        )
+
     @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
     def test_charge_spin_condition_matches_compile(self) -> None:
         """Charge/spin conditions should work through the compiled energy path."""
@@ -731,6 +969,49 @@ def test_fixed_edge_geometry_matches_standard_cache(self) -> None:
         torch.testing.assert_close(cache_std.D_full, cache_sparse.D_full[:n_real])
         torch.testing.assert_close(cache_std.Dt_full, cache_sparse.Dt_full[:n_real])
 
+        from deepmd.pt.model.descriptor.sezm_nn import edge_cache as edge_cache_module
+
+        with (
+            mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=False),
+            mock.patch.object(
+                edge_cache_module,
+                "_build_edge_wigner",
+                wraps=edge_cache_module._build_edge_wigner,
+            ) as build_edge_wigner,
+        ):
+            sorted_edge_index, sorted_edge_vec, sorted_edge_mask, _ = (
+                _sort_edge_tensors_by_destination(
+                    edge_index,
+                    edge_vec,
+                    edge_mask,
+                    edge_index,
+                )
+            )
+            cache_sfpg = build_edge_cache_from_edges(
+                type_ebed=type_ebed,
+                atype_flat=atype_loc.reshape(-1),
+                edge_index=sorted_edge_index,
+                edge_vec=sorted_edge_vec,
+                edge_mask=sorted_edge_mask,
+                compute_dtype=descriptor.compute_dtype,
+                eps=descriptor.eps,
+                deg_norm_floor=descriptor.deg_norm_floor,
+                inner_clamp=descriptor.inner_clamp,
+                bridging_switch=torch.ones_like,
+                edge_envelope=descriptor.edge_envelope,
+                radial_basis=descriptor.radial_basis,
+                has_exclude_types=False,
+                edge_type_keep_mask=descriptor._edge_type_keep_mask,
+                random_gamma=False,
+                wigner_calc=descriptor.wigner_calc,
+                packed_wigner_candidate=True,
+                destinations_sorted=True,
+            )
+
+        self.assertFalse(build_edge_wigner.call_args.kwargs["packed_wigner"])
+        self.assertEqual(cache_sfpg.D_full.dim(), 3)
+        self.assertIsNotNone(cache_sfpg.edge_src_gate)
+
     @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
     def test_eval_compile_policy(self) -> None:
         """Eval should stay eager by default and compile only with env override."""
diff --git a/source/tests/pt/model/test_sezm_parallel.py b/source/tests/pt/model/test_sezm_parallel.py
index bceec317d8..5cbbdd7336 100644
--- a/source/tests/pt/model/test_sezm_parallel.py
+++ b/source/tests/pt/model/test_sezm_parallel.py
@@ -17,6 +17,9 @@
 
 import ctypes
 import unittest
+from unittest import (
+    mock,
+)
 
 import numpy as np
 import torch
@@ -44,6 +47,29 @@
 _SENDLIST_KEEPALIVE: list[np.ndarray] = []
 
 
+class TestSeZMDenseCommIsolation(unittest.TestCase):
+    def test_local_nlist_features_are_not_passed_to_ghost_exchange(self) -> None:
+        model = _build_model(torch.device("cpu"))
+        descriptor = model.atomic_model.descriptor
+        coord = torch.tensor(
+            [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]],
+            dtype=torch.float64,
+            device="cpu",
+        )
+        atype = torch.tensor([[0, 1, 0]], dtype=torch.int64, device="cpu")
+        # Two local atoms and a mapped ghost atom.
+        nlist = torch.tensor([[[1, 2], [0, 2]]], dtype=torch.int64, device="cpu")
+        mapping = torch.tensor([[0, 1, 0]], dtype=torch.int64, device="cpu")
+        expected = descriptor(coord, atype, nlist, mapping=mapping)[0]
+        with mock.patch.object(
+            sezm_block,
+            "exchange_ghost_features",
+            side_effect=AssertionError("local-only features reached ghost exchange"),
+        ):
+            actual = descriptor(coord, atype, nlist, mapping=mapping, comm_dict={})[0]
+        torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
 def _tiny_parallel_model_params(**overrides) -> dict:
     """Minimal fp64 SeZM config exercising message passing and FiLM/GIE seeds."""
     descriptor = {
diff --git a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py
index a72698fbd4..96cfc2b330 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py
@@ -1,6 +1,13 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """End-to-end parity of the pt_expt DPA4 accelerated inference paths."""
 
+from types import (
+    SimpleNamespace,
+)
+from typing import (
+    Any,
+)
+
 import numpy as np
 import pytest
 import torch
@@ -10,6 +17,9 @@
 except ImportError:
     pass
 
+from deepmd.dpmodel.utils.neighbor_graph import (
+    NeighborGraph,
+)
 from deepmd.pt_expt.descriptor.dpa4 import (
     DescrptDPA4,
 )
@@ -32,6 +42,13 @@
     so2_conv,
     zonal_scatter,
 )
+from deepmd.pt_expt.kernels.cute.sezm import (
+    runtime_policy,
+)
+from deepmd.pt_expt.kernels.cute.sezm.output_grid import (
+    readout_l0,
+)
+from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2
 from deepmd.pt_expt.kernels.cutile import (
     CUTILE_AVAILABLE,
 )
@@ -86,6 +103,234 @@ def _make_descriptor(
     )
 
 
+def make_neo_descriptor(
+    sel: int | list[int] = 4,
+    *,
+    edge_norm: bool = True,
+) -> DescrptDPA4:
+    """Build the exact Neo layout served by the current CuTe kernels."""
+    return DescrptDPA4(
+        ntypes=2,
+        sel=sel,
+        channels=32,
+        lmax=3,
+        mmax=1,
+        n_blocks=2,
+        so2_layers=3,
+        n_focus=2,
+        message_node_so3=True,
+        ffn_neurons=0,
+        ffn_so3_grid=True,
+        grid_branch=[0, 0, 1],
+        ffn_blocks=1,
+        so3_readout="mlp",
+        use_amp=False,
+        random_gamma=False,
+        edge_norm=edge_norm,
+        precision="float32",
+        trainable=False,
+        seed=42,
+    ).eval()
+
+
+def test_cute_contract_recognizes_pt_expt_neo_modules(monkeypatch) -> None:
+    """Keep CuTe eligibility independent of the concrete PT module classes."""
+    for name in (
+        "DP_TRITON_INFER",
+        "DP_CUDA_INFER",
+        "DP_CUTILE_INFER",
+        "DP_CUTE_INFER",
+    ):
+        monkeypatch.setenv(name, "0")
+
+    descriptor = make_neo_descriptor()
+
+    assert descriptor.blocks
+    assert all(so2.is_supported_neo_so2_block(block) for block in descriptor.blocks)
+    assert readout_l0.has_neo_readout_contract(descriptor.output_ffn)
+
+
+def test_packed_wigner_graph_compacts_and_reuses_csr(monkeypatch) -> None:
+    """Reuse canonical graph metadata without sorting the edge stream again."""
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    descriptor = _make_descriptor(2, [4], 4.0).eval()
+    monkeypatch.setattr(
+        descriptor,
+        "is_cute_infer_packed_wigner_candidate",
+        lambda *args: True,
+    )
+    graph = NeighborGraph(
+        n_node=torch.tensor([4], dtype=torch.int64),
+        edge_index=torch.tensor(
+            [[3, 0, 2, 0, 0], [0, 1, 2, 0, 0]],
+            dtype=torch.int64,
+        ),
+        edge_vec=torch.arange(15, dtype=torch.float32).reshape(5, 3),
+        edge_mask=torch.tensor([True, True, True, False, False]),
+        destination_order=torch.arange(5, dtype=torch.int64),
+        destination_row_ptr=torch.tensor([0, 1, 2, 3, 3], dtype=torch.int64),
+        source_order=torch.tensor([1, 2, 0, 3, 4], dtype=torch.int64),
+        source_row_ptr=torch.tensor([0, 1, 1, 2, 3], dtype=torch.int64),
+        destination_sorted=True,
+    )
+    monkeypatch.setattr(
+        torch,
+        "argsort",
+        lambda *args, **kwargs: pytest.fail("canonical CSR must not be rebuilt"),
+    )
+
+    packed_graph = descriptor.prepare_packed_wigner_graph(graph, n_nodes=4)
+
+    assert packed_graph is not None
+    torch.testing.assert_close(
+        packed_graph.edge_index,
+        torch.tensor([[3, 0, 2], [0, 1, 2]], dtype=torch.int64),
+    )
+    torch.testing.assert_close(
+        packed_graph.edge_vec,
+        torch.arange(9, dtype=torch.float32).reshape(3, 3),
+    )
+    assert torch.all(packed_graph.edge_mask)
+    assert packed_graph.destination_sorted
+    torch.testing.assert_close(
+        packed_graph.destination_row_ptr,
+        torch.tensor([0, 1, 2, 3, 3], dtype=torch.int64),
+    )
+    torch.testing.assert_close(
+        packed_graph.source_order,
+        torch.tensor([1, 2, 0], dtype=torch.int64),
+    )
+    torch.testing.assert_close(
+        packed_graph.source_row_ptr,
+        torch.tensor([0, 1, 1, 2, 3], dtype=torch.int64),
+    )
+
+    edge_cache = SimpleNamespace(
+        src=packed_graph.edge_index[0],
+        dst=packed_graph.edge_index[1],
+        csr_cache={
+            "dst": (
+                packed_graph.destination_order,
+                packed_graph.destination_row_ptr,
+            ),
+            "src": (packed_graph.source_order, packed_graph.source_row_ptr),
+        },
+        destinations_sorted=True,
+        D_packed=torch.empty(3, 46),
+        edge_src_gate=None,
+    )
+    metadata = descriptor.prepare_cute_infer_so2_metadata(
+        edge_cache,
+        n_nodes=4,
+    )
+    assert metadata is not None
+    destination_row_ptr, source_order, source_row_ptr = metadata
+    torch.testing.assert_close(
+        destination_row_ptr,
+        packed_graph.destination_row_ptr.to(torch.int32),
+    )
+    torch.testing.assert_close(
+        source_order,
+        packed_graph.source_order.to(torch.int32),
+    )
+    torch.testing.assert_close(
+        source_row_ptr,
+        packed_graph.source_row_ptr.to(torch.int32),
+    )
+
+
+def test_packed_wigner_graph_without_csr_builds_metadata_once(monkeypatch) -> None:
+    """Build each CSR ordering once when the graph does not provide one."""
+    monkeypatch.setenv("DP_CUTE_INFER", "1")
+    descriptor = _make_descriptor(2, [4], 4.0).eval()
+    monkeypatch.setattr(
+        descriptor,
+        "is_cute_infer_packed_wigner_candidate",
+        lambda *args: True,
+    )
+    graph = NeighborGraph(
+        n_node=torch.tensor([2], dtype=torch.int64),
+        edge_index=torch.tensor([[0, 1], [1, 0]], dtype=torch.int64),
+        edge_vec=torch.ones(2, 3),
+        edge_mask=torch.ones(2, dtype=torch.bool),
+    )
+    argsort = torch.argsort
+    sort_count = 0
+
+    def count_argsort(*args: Any, **kwargs: Any) -> torch.Tensor:
+        nonlocal sort_count
+        sort_count += 1
+        return argsort(*args, **kwargs)
+
+    monkeypatch.setattr(torch, "argsort", count_argsort)
+
+    packed_graph = descriptor.prepare_packed_wigner_graph(graph, n_nodes=2)
+
+    assert packed_graph is not None
+    edge_cache = SimpleNamespace(
+        src=packed_graph.edge_index[0],
+        dst=packed_graph.edge_index[1],
+        csr_cache={
+            "dst": (
+                packed_graph.destination_order,
+                packed_graph.destination_row_ptr,
+            ),
+            "src": (packed_graph.source_order, packed_graph.source_row_ptr),
+        },
+        destinations_sorted=True,
+        D_packed=torch.empty(2, 46),
+        edge_src_gate=None,
+    )
+
+    assert descriptor.prepare_cute_infer_so2_metadata(edge_cache, n_nodes=2) is not None
+    assert sort_count == 2
+
+
+def test_packed_wigner_cache_routes_pt_expt_block_to_cute(monkeypatch) -> None:
+    """Use packed Wigner storage as the prevalidated per-block dispatch token."""
+    block = make_neo_descriptor().blocks[0]
+    destination_row_ptr = torch.tensor([0, 1], dtype=torch.int32)
+    source_order = torch.tensor([0], dtype=torch.int32)
+    source_row_ptr = torch.tensor([0, 1], dtype=torch.int32)
+    edge_cache = SimpleNamespace(
+        D_packed=torch.empty(1, 46),
+        cute_infer_so2_metadata=(
+            destination_row_ptr,
+            source_order,
+            source_row_ptr,
+        ),
+    )
+    expected = torch.empty(1)
+
+    def run_cute(
+        candidate_block: Any,
+        x: torch.Tensor,
+        candidate_edge_cache: Any,
+        radial_feat: torch.Tensor,
+        *,
+        dst_ptr: torch.Tensor,
+        source_order: torch.Tensor,
+        source_ptr: torch.Tensor,
+    ) -> torch.Tensor:
+        del x, radial_feat
+        assert candidate_block is block
+        assert candidate_edge_cache is edge_cache
+        assert dst_ptr is destination_row_ptr
+        assert source_order is edge_cache.cute_infer_so2_metadata[1]
+        assert source_ptr is source_row_ptr
+        return expected
+
+    monkeypatch.setattr(so2, "maybe_run_cute_so2", run_cute)
+
+    actual = block._run_so2_unit_impl(
+        torch.empty(1),
+        edge_cache,
+        torch.empty(1),
+    )
+
+    assert actual is expected
+
+
 @pytest.mark.parametrize(
     ("precision", "expected_bound"),  # descriptor dtype and CUDA binding eligibility
     [("float32", True), ("float64", False)],
@@ -128,17 +373,23 @@ def test_fp32_only_cuda_bindings(
 
     assert len(initial_embeddings) == 1
     assert grid_nets
-    assert (descriptor._cuda_radial_fn is not None) is expected_bound
-    assert all(module._cuda_scatter is expected_bound for module in initial_embeddings)
+    assert (descriptor.cuda_infer_l_1_radial is not None) is expected_bound
     assert all(
-        (module._grid_pair_fn is not None) is expected_bound for module in grid_nets
+        module.cuda_infer_l_1_scatter is expected_bound for module in initial_embeddings
+    )
+    assert all(
+        (module.cuda_infer_l_1_grid_pair is not None) is expected_bound
+        for module in grid_nets
     )
     cpu_zonal = torch.empty(1)
-    assert all(not module._can_fuse_scatter(cpu_zonal) for module in initial_embeddings)
+    assert all(
+        not module.can_run_cuda_infer_l_1_scatter(cpu_zonal)
+        for module in initial_embeddings
+    )
     for module in initial_embeddings:
-        module._force_fused_scatter = True
+        module.force_cuda_infer_l_1_scatter = True
     assert all(
-        module._can_fuse_scatter(cpu_zonal) is expected_bound
+        module.can_run_cuda_infer_l_1_scatter(cpu_zonal) is expected_bound
         for module in initial_embeddings
     )
 
@@ -230,17 +481,79 @@ def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None:
             for module in accelerated.modules()
             if isinstance(module, SO2Convolution)
         )
-        if conv._cuda_conv_fn is None:
+        if conv.cuda_infer_l_2_conv is None:
             pytest.skip("the descriptor layout has no fused CUDA convolution")
-        assert conv._flash_atten_fn is not None
-        assert conv._cuda_value_train is None
-        assert not accelerated._wigner_free_conv
+        assert conv.flash_attention is not None
+        assert conv.cuda_train_value is None
+        assert not accelerated.cuda_infer_l_2_covers_all_blocks
         assert accelerated._build_full_wigner()
 
         output, gradient = self._step(accelerated)
         np.testing.assert_allclose(output, dense_output, rtol=2e-4, atol=2e-5)
         np.testing.assert_allclose(gradient, dense_gradient, rtol=2e-4, atol=2e-5)
 
+    @pytest.mark.parametrize("edge_norm", [True, False])
+    def test_cute_neo_forward_and_coordinate_gradient(
+        self,
+        monkeypatch,
+        edge_norm: bool,
+    ) -> None:
+        """Exercise packed Wigner and SO2 through the PT-expt descriptor."""
+        try:
+            import cuda.bindings.driver  # noqa: F401
+            import cutlass.cute  # noqa: F401
+            import tvm_ffi  # noqa: F401
+
+            from deepmd.pt_expt.kernels.cute.sezm import (
+                wignerd,
+            )
+        except ImportError:
+            pytest.skip("the CuTe DSL runtime is unavailable")
+        capability = tuple(torch.cuda.get_device_capability(self.device))
+        if not runtime_policy.is_supported_so2_capability(capability):
+            pytest.skip(f"the CuTe SO2 path does not support {capability}")
+
+        for name in (
+            "DP_TRITON_INFER",
+            "DP_CUDA_INFER",
+            "DP_CUTILE_INFER",
+            "DP_CUTE_INFER",
+        ):
+            monkeypatch.setenv(name, "0")
+        data = make_neo_descriptor(self.sel_mix, edge_norm=edge_norm).serialize()
+        reference = DescrptDPA4.deserialize(data).to(self.device).eval()
+        reference_output, reference_gradient = self._step(reference)
+
+        monkeypatch.setenv("DP_CUTE_INFER", "1")
+        monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False)
+        accelerated = DescrptDPA4.deserialize(data).to(self.device).eval()
+        dispatch_count = {"wigner": 0, "so2": 0}
+        run_wignerd = wignerd.run_cute_wignerd
+        run_so2 = so2.maybe_run_cute_so2
+
+        def count_wignerd(*args: Any, **kwargs: Any) -> Any:
+            result = run_wignerd(*args, **kwargs)
+            dispatch_count["wigner"] += result is not None
+            return result
+
+        def count_so2(*args: Any, **kwargs: Any) -> Any:
+            result = run_so2(*args, **kwargs)
+            dispatch_count["so2"] += result is not None
+            return result
+
+        monkeypatch.setattr(wignerd, "run_cute_wignerd", count_wignerd)
+        monkeypatch.setattr(so2, "maybe_run_cute_so2", count_so2)
+        output, gradient = self._step(accelerated)
+
+        assert dispatch_count == {"wigner": 1, "so2": len(accelerated.blocks)}
+        np.testing.assert_allclose(output, reference_output, rtol=5e-5, atol=5e-5)
+        np.testing.assert_allclose(
+            gradient,
+            reference_gradient,
+            rtol=5e-5,
+            atol=5e-5,
+        )
+
     @pytest.mark.parametrize(
         "backend", ["triton", "cuda", "cutile"]
     )  # inference kernels
@@ -287,18 +600,18 @@ def test_forward_and_coordinate_gradient(
             if isinstance(module, WignerDCalculator)
         )
         if backend == "triton":
-            assert so2._flash_atten_fn is not None
-            assert so2._triton_value_path is not None
-            assert wigner._use_triton_monomials
+            assert so2.flash_attention is not None
+            assert so2.triton_infer_l_2_value is not None
+            assert wigner.triton_infer_l_1_monomials
         elif backend == "cuda":
-            if so2._cuda_conv_fn is None:
+            if so2.cuda_infer_l_2_conv is None:
                 pytest.skip("DPA4 CUDA operators are unavailable")
-            assert accelerated._cuda_radial_fn is not None
-            assert accelerated._cuda_wigner_fn is not None
+            assert accelerated.cuda_infer_l_1_radial is not None
+            assert accelerated.cuda_infer_l_1_wigner is not None
         else:
-            assert so2._flash_atten_fn is not None
-            assert so2._cutile_value_path is not None
-            assert wigner._use_cutile_monomials
+            assert so2.flash_attention is not None
+            assert so2.cutile_infer_value is not None
+            assert wigner.cutile_infer_monomials
 
         coord_ref, atype, nlist = self._inputs()
         output_ref = reference(coord_ref, atype, nlist)[0]
diff --git a/source/tests/pt_expt/descriptor/test_dpa4_ckpt_triton.py b/source/tests/pt_expt/descriptor/test_dpa4_ckpt_triton.py
index 200a947aa6..930d0a957f 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4_ckpt_triton.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4_ckpt_triton.py
@@ -120,9 +120,9 @@ def test_triton_eager_fallback_parity(self, monkeypatch) -> None:
         monkeypatch.setenv("DP_TRITON_INFER", "1")
         m = DescrptDPA4.deserialize(self.data).to(self.device).eval()
         so2 = next(x for x in m.modules() if isinstance(x, SO2Convolution))
-        assert so2.use_triton_infer
-        assert so2._rotate_to_local_fn is not None
-        assert so2._rotate_back_fn is not None
+        assert so2.triton_infer_level >= 1
+        assert so2.triton_l_1_rotate_to_local is not None
+        assert so2.triton_l_1_rotate_back is not None
         mixers = [x for x in m.modules() if isinstance(x, DynamicRadialDegreeMixer)]
         assert mixers and all(x.triton_infer_level >= 1 for x in mixers)
 
diff --git a/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py b/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py
index 397ccdc40a..8ae2abb50b 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py
@@ -44,7 +44,7 @@
     DynamicRadialDegreeMixer,
     SO2Convolution,
     SO2Linear,
-    _active_triton_level,
+    active_triton_level,
 )
 from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import (
     op_available as cuda_value_available,
@@ -143,34 +143,34 @@ def test_triton_mode_gate_binds_each_stage(
         assert conv.triton_infer_level == infer_level
         # Rotation and flash wrappers retain their eager implementations when
         # Triton is unavailable, so their binding follows the mode gates alone.
-        assert (conv._rotate_to_local_fn is not None) is requested
-        assert (conv._rotate_back_fn is not None) is requested
-        assert (conv._flash_atten_fn is not None) is requested
-        assert conv._flash_atten_trains is (requested and training)
+        assert (conv.triton_l_1_rotate_to_local is not None) is requested
+        assert (conv.triton_l_1_rotate_back is not None) is requested
+        assert (conv.flash_attention is not None) is requested
+        assert conv.flash_attention_supports_training is (requested and training)
         # Segment softmax uses fp32 accumulation and preserves fp64 compute
         # through the descriptor's reference path.
-        assert (conv._segment_softmax_fn is not None) is (
+        assert (conv.triton_l_1_segment_softmax is not None) is (
             requested and SEGMENT_SOFTMAX_TRITON_AVAILABLE and precision == "float32"
         )
         # The rotate-mix front end is bound by a profitability bound on the
         # hidden width, which this narrow block sits below.
         assert conv.hidden_channels < 128
-        assert conv._triton_rotate_mix is None
+        assert conv.triton_l_1_rotate_mix is None
         # The CUDA gate is off, so the value stream stays on the stages.
-        assert conv._cuda_value_train is None
+        assert conv.cuda_train_value is None
 
     for module in descriptor.modules():
         if isinstance(module, SO2Linear):
             # The fused GEMM additionally needs every |m| block width to align
             # to its BN=64 tile, which a narrow block does not satisfy.
             aligned = slices_supported(module._block_diag_slices)
-            assert (module._block_diag_gemm is not None) is (
+            assert (module.triton_l_1_block_diag_gemm is not None) is (
                 requested and SO2_BLOCK_GEMM_TRITON_AVAILABLE and aligned
             )
         if isinstance(module, DynamicRadialDegreeMixer):
             # The callable contains its eager fallback, so construction binds it
             # whenever either mode requests the stage.
-            assert (module._radial_mix_block is not None) is requested
+            assert (module.triton_l_1_radial_mix is not None) is requested
         if isinstance(module, GatedActivation):
             assert module.triton_train_level == train_level
             assert module.triton_infer_level == infer_level
@@ -191,9 +191,9 @@ def test_triton_mode_gate_binds_each_stage(
             if isinstance(
                 module, (SO2Convolution, SO2Linear, DynamicRadialDegreeMixer)
             ):
-                assert _active_triton_level(module) == active_level
+                assert active_triton_level(module) == active_level
             if isinstance(module, SO2Convolution):
-                assert module._rotation_active() is bool(active_level)
+                assert module.rotation_kernel_active() is bool(active_level)
             if isinstance(module, GatedActivation):
                 level = (
                     module.triton_train_level
@@ -216,13 +216,13 @@ def test_cuda_train_gate_binds_the_value_stream(monkeypatch) -> None:
     ]
     assert convolutions
     for conv in convolutions:
-        assert conv._cuda_value_train is not None
+        assert conv.cuda_train_value is not None
         # The two layers are independent: the CUDA value stream does not
         # switch on any Triton stage, and the attention span stays dense
         # until the Triton gate asks for it.
         assert conv.triton_train_level == 0
-        assert conv._segment_softmax_fn is None
-        assert conv._flash_atten_trains is False
+        assert conv.triton_l_1_segment_softmax is None
+        assert conv.flash_attention_supports_training is False
 
 
 def test_cuda_triton_train_reuses_packed_wigner_runs(monkeypatch) -> None:
@@ -236,14 +236,14 @@ def test_cuda_triton_train_reuses_packed_wigner_runs(monkeypatch) -> None:
     monkeypatch.setenv("DP_TRITON_TRAIN", "1")
 
     descriptor = _make_descriptor(2, [20], 4.0).train()
-    assert descriptor._packed_wigner_train
+    assert descriptor.cuda_train_covers_all_blocks
     assert not descriptor._build_full_wigner()
 
     for block in descriptor.blocks:
         conv = block.so2_conv
-        assert conv._cuda_value_train is not None
-        assert conv._flash_atten_fn is not None
-        assert conv._flash_atten_trains
+        assert conv.cuda_train_value is not None
+        assert conv.flash_attention is not None
+        assert conv.flash_attention_supports_training
 
 
 @pytest.mark.parametrize("gate_name", ["DP_TRITON_TRAIN", "DP_TRITON_INFER"])
@@ -270,7 +270,7 @@ def test_grid_pair_train_follows_its_gate_and_slot_bound(
             and GRID_PAIR_TRITON_AVAILABLE
             and slots >= 75
         )
-        assert (net._grid_pair_train_fn is not None) is expected
+        assert (net.triton_train_l_1_grid_pair is not None) is expected
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
@@ -330,9 +330,9 @@ def test_training_step_matches_the_dense_path(self, monkeypatch, path) -> None:
             module for module in fused.modules() if isinstance(module, SO2Convolution)
         )
         if path in ("cuda", "cuda-triton"):
-            assert conv._cuda_value_train is not None
+            assert conv.cuda_train_value is not None
         if path in ("triton", "cuda-triton"):
-            assert conv._segment_softmax_fn is not None
+            assert conv.triton_l_1_segment_softmax is not None
         fused_objective, fused_gradient = self._step(fused)
 
         np.testing.assert_allclose(
diff --git a/source/tests/pt_expt/utils/test_serialization_kernel_levels.py b/source/tests/pt_expt/utils/test_serialization_kernel_levels.py
index 53651103ca..ad929b5a4b 100644
--- a/source/tests/pt_expt/utils/test_serialization_kernel_levels.py
+++ b/source/tests/pt_expt/utils/test_serialization_kernel_levels.py
@@ -225,7 +225,7 @@ def test_packed_weights_only_prepare_for_bound_dpa4_value_path(
 ) -> None:
     model = torch.nn.Module()
     conv = torch.nn.Module()
-    conv._triton_value_path = object() if has_value_path else None
+    conv.triton_infer_l_2_value = object() if has_value_path else None
     model.add_module("conv", conv)
 
     calls = []