From cde9690d2321156e5e436fb5a11b1f07e94334f5 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sat, 10 Oct 2026 20:51:50 +0800
Subject: [PATCH 01/15] perf(pt): banded Wigner-D storage for the fused
 inference consumers (DP_PACKED_D_INFER)

---
 .../pt/model/descriptor/sezm_nn/edge_cache.py |  41 ++++-
 .../pt/model/descriptor/sezm_nn/embedding.py  |  21 ++-
 deepmd/pt/model/descriptor/sezm_nn/so2.py     |   7 +-
 deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 165 +++++++++++++-----
 .../kernels/triton/sezm/so2_value_path.py     | 149 +++++++++++-----
 deepmd/pt_expt/kernels/utils.py               |  20 +++
 6 files changed, 305 insertions(+), 98 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
index 4a67e2d1f9..dcb09f4c82 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
@@ -25,6 +25,9 @@
     rearrange,
 )
 
+from deepmd.pt_expt.kernels.utils import (
+    use_packed_d_infer,
+)
 from .utils import (
     get_promoted_dtype,
     nvtx_range,
@@ -93,6 +96,12 @@ class EdgeFeatureCache(NamedTuple):
     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.
+    D_packed
+        Structural Wigner band with shape ``(E, 3*D-2)`` (the ``m = 0/-1/+1``
+        rows per degree block, flash packed convention). Present only when the
+        banded inference storage is enabled and supported; then ``D_full`` and
+        ``Dt_full`` are ``None`` and the fused reduced-layout consumers address
+        this band instead. None otherwise.
     D_to_m_cache
         Lazy cache for projected D matrices keyed by a normalized
         ``"lmax:mmax"`` identifier.
@@ -120,6 +129,7 @@ class EdgeFeatureCache(NamedTuple):
     csr_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None
     edge_quat: torch.Tensor | None = None
     node_gate: torch.Tensor | None = None
+    D_packed: torch.Tensor | None = None
 
 
 def cached_edge_csr(
@@ -435,7 +445,7 @@ def build_edge_cache(
 
     # === Step 6. Edge quaternion -> Wigner-D blocks ===
     with nvtx_range("wigner_d"):
-        D_full, Dt_full, edge_quat = _build_edge_wigner(
+        D_full, Dt_full, edge_quat, D_packed = _build_edge_wigner(
             edge_vec=edge_vec,
             edge_len=edge_len,
             eps=eps,
@@ -457,6 +467,7 @@ def build_edge_cache(
         D_full=D_full,
         Dt_full=Dt_full,
         edge_quat=edge_quat,
+        D_packed=D_packed,
         deg_norm_floor=deg_norm_floor,
     )
 
@@ -584,7 +595,7 @@ def build_edge_cache_from_edges(
 
     # === Step 4. Edge quaternion -> Wigner-D blocks ===
     with nvtx_range("wigner_d"):
-        D_full, Dt_full, edge_quat = _build_edge_wigner(
+        D_full, Dt_full, edge_quat, D_packed = _build_edge_wigner(
             edge_vec=edge_vec,
             edge_len=edge_len,
             eps=eps,
@@ -634,6 +645,7 @@ def build_edge_cache_from_edges(
         D_full=D_full,
         Dt_full=Dt_full,
         edge_quat=edge_quat,
+        D_packed=D_packed,
         deg_norm_floor=deg_norm_floor,
         node_gate=node_gate,
     )
@@ -695,9 +707,23 @@ 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, edge_quat, None
+    # ``wigner_calc`` may arrive as the calculator instance (callable through
+    # ``nn.Module.__call__``) or as a bound forward; accept both.
+    calc = getattr(wigner_calc, "__self__", wigner_calc)
+    if (
+        use_packed_d_infer()
+        and callable(getattr(calc, "forward_band", None))
+        and calc.band_supported()
+    ):
+        # Banded inference storage: the fused reduced-layout consumers (the
+        # rotate-mix operator and the flash aggregation) read only the three
+        # structural rows per degree block, so the dense surfaces are never
+        # materialized. Reference paths that need the dense layout must not
+        # enable the flag.
+        return None, None, edge_quat, calc.forward_band(edge_quat)
     D_full, Dt_full = wigner_calc(edge_quat)
-    return D_full, Dt_full, edge_quat
+    return D_full, Dt_full, edge_quat, None
 
 
 def _finalize_edge_cache(
@@ -714,6 +740,7 @@ def _finalize_edge_cache(
     edge_quat: torch.Tensor,
     deg_norm_floor: float,
     node_gate: torch.Tensor | None = None,
+    D_packed: torch.Tensor | None = None,
 ) -> EdgeFeatureCache:
     """
     Assemble the shared `EdgeFeatureCache` layout.
@@ -781,6 +808,7 @@ def _finalize_edge_cache(
         csr_cache={},
         edge_quat=edge_quat,
         node_gate=node_gate,
+        D_packed=D_packed,
     )
 
 
@@ -938,15 +966,19 @@ def edge_cache_to_dtype(
     _D_full = cache.D_full
     _Dt_full = cache.Dt_full
     _edge_quat = cache.edge_quat
+    _D_packed = cache.D_packed
     D_full: torch.Tensor | None = None
     Dt_full: torch.Tensor | None = None
     edge_quat: torch.Tensor | None = None
+    D_packed: 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 _edge_quat is not None:
         edge_quat = _edge_quat.to(dtype=dtype)
+    if _D_packed is not None:
+        D_packed = _D_packed.to(dtype=dtype)
 
     # CSR views contain only integer topology. Preserve them across the dtype
     # conversion so every accelerated consumer shares the per-step sort.
@@ -966,4 +998,5 @@ def edge_cache_to_dtype(
         csr_cache=None if cache.csr_cache is None else dict(cache.csr_cache),
         edge_quat=edge_quat,
         node_gate=cache.node_gate,
+        D_packed=D_packed,
     )
diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
index 296205002a..cc14ba70bb 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
@@ -272,12 +272,21 @@ def forward(
         # 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)
+            if getattr(edge_cache, "D_packed", None) is not None:
+                # The m = 0 band segment holds D[r0(l), d] per column d; the
+                # zonal entry Dt[row(l, m), col(l, 0)] equals D[r0(l), row(l, m)],
+                # so indexing the segment by the non-scalar rows reproduces
+                # the gather below without the dense surface.
+                zonal_coupling = getattr(edge_cache, "D_packed", None)[
+                    :, self.non_scalar_row_index
+                ]
+            else:
+                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 ===
         # Each non-scalar packed row reuses the radial feature of its degree l.
diff --git a/deepmd/pt/model/descriptor/sezm_nn/so2.py b/deepmd/pt/model/descriptor/sezm_nn/so2.py
index 4803b5d7c5..453ccd4725 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/so2.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/so2.py
@@ -2016,7 +2016,12 @@ def forward_attention_flash(
         dst = edge_cache.dst
         n_node = x.shape[0]
         order, row_ptr = cached_edge_csr(edge_cache, "dst", n_node)
-        rotation = edge_cache.Dt_full
+        # The banded inference storage serves both rotation directions: the
+        # flash packed convention stores the structural band of D, whose
+        # entries are exactly the Dt columns the aggregation reads.
+        rotation = getattr(edge_cache, "D_packed", None)
+        if rotation is None:
+            rotation = edge_cache.Dt_full
         if rotation is None:
             rotation = self._cuda_value_train.edge_runs(edge_cache)
         pre_gate = self._flash_atten_fn(
diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index b283e3a7e5..bd31b5cf39 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -524,84 +524,95 @@ def __init__(
                 persistent=True,
             )
 
-    def forward(
+    def _compute_all_blocks(
         self, edge_quaternion: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Build packed block-diagonal Wigner-D matrices from edge quaternions.
+    ) -> list[tuple[int, torch.Tensor]]:
+        """Compute every per-degree block below the polynomial tail.
 
-        Parameters
-        ----------
-        edge_quaternion
-            Unit quaternions with shape ``(E, 4)`` representing the global->local
-            edge rotation.
-
-        Returns
-        -------
-        tuple[torch.Tensor, torch.Tensor]
-            ``(D_full, Dt_full)`` with shape ``(E, (lmax+1)^2, (lmax+1)^2)``.
+        Returns ``[(l, block)]`` with ``block`` of shape
+        ``(E, 2l+1, 2l+1)``, degrees ascending; degree 0 (the constant 1)
+        is not included.
         """
-        edge_quaternion = quaternion_normalize(
-            edge_quaternion.to(dtype=self.dtype),
-            eps=self.eps,
-        )
-        n_edge = edge_quaternion.shape[0]
-        D_full = torch.zeros(
-            n_edge,
-            self.dim_full,
-            self.dim_full,
-            dtype=edge_quaternion.dtype,
-            device=edge_quaternion.device,
-        )
-        D_full[:, 0, 0] = 1.0
-
+        blocks: list[tuple[int, torch.Tensor]] = []
         if self.lmax >= 1:
             with nvtx_range("WignerD/l1"):
-                D_full[:, 1:4, 1:4] = self._compute_l1_block(edge_quaternion)
-
+                blocks.append((1, self._compute_l1_block(edge_quaternion)))
         if self.lmax >= 2:
             with nvtx_range("WignerD/l2"):
-                D_full[:, 4:9, 4:9] = self._compute_l2_block(edge_quaternion)
-
+                blocks.append((2, self._compute_l2_block(edge_quaternion)))
         if self.lmax >= 3:
             if self.lmax >= 4:
                 with nvtx_range("WignerD/l3l4"):
                     D_l3, D_l4 = self._compute_l3l4_blocks(edge_quaternion)
-                    D_full[:, 9:16, 9:16] = D_l3
-                    D_full[:, 16:25, 16:25] = D_l4
+                blocks.append((3, D_l3))
+                blocks.append((4, D_l4))
             else:
                 with nvtx_range("WignerD/l3"):
-                    D_full[:, 9:16, 9:16] = self._compute_l3_block(edge_quaternion)
-
+                    blocks.append((3, self._compute_l3_block(edge_quaternion)))
         if self.lmax >= 5:
             if self.lmax >= 6:
                 with nvtx_range("WignerD/l5l6"):
                     D_l5, D_l6 = self._compute_l5l6_blocks(edge_quaternion)
-                    D_full[:, 25:36, 25:36] = D_l5
-                    D_full[:, 36:49, 36:49] = D_l6
+                blocks.append((5, D_l5))
+                blocks.append((6, D_l6))
             else:
                 with nvtx_range("WignerD/l5"):
-                    D_full[:, 25:36, 25:36] = self._compute_l5_block(edge_quaternion)
-
+                    blocks.append((5, self._compute_l5_block(edge_quaternion)))
         if self.lmax >= 7:
             if self.lmax >= 8:
                 with nvtx_range("WignerD/l7l8"):
                     D_l7, D_l8 = self._compute_l7l8_blocks(edge_quaternion)
-                    D_full[:, 49:64, 49:64] = D_l7
-                    D_full[:, 64:81, 64:81] = D_l8
+                blocks.append((7, D_l7))
+                blocks.append((8, D_l8))
             else:
                 with nvtx_range("WignerD/l7"):
-                    D_full[:, 49:64, 49:64] = self._compute_l7_block(edge_quaternion)
-
+                    blocks.append((7, self._compute_l7_block(edge_quaternion)))
         if self.lmax >= 9:
             if self.lmax >= 10:
                 with nvtx_range("WignerD/l9l10"):
                     D_l9, D_l10 = self._compute_l9l10_blocks(edge_quaternion)
-                    D_full[:, 81:100, 81:100] = D_l9
-                    D_full[:, 100:121, 100:121] = D_l10
+                blocks.append((9, D_l9))
+                blocks.append((10, D_l10))
             else:
                 with nvtx_range("WignerD/l9"):
-                    D_full[:, 81:100, 81:100] = self._compute_l9_block(edge_quaternion)
+                    blocks.append((9, self._compute_l9_block(edge_quaternion)))
+        return blocks
+
+    def forward(
+        self, edge_quaternion: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Build packed block-diagonal Wigner-D matrices from edge quaternions.
+
+        Parameters
+        ----------
+        edge_quaternion
+            Unit quaternions with shape ``(E, 4)`` representing the global->local
+            edge rotation.
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor]
+            ``(D_full, Dt_full)`` with shape ``(E, (lmax+1)^2, (lmax+1)^2)``.
+        """
+        edge_quaternion = quaternion_normalize(
+            edge_quaternion.to(dtype=self.dtype),
+            eps=self.eps,
+        )
+        n_edge = edge_quaternion.shape[0]
+        D_full = torch.zeros(
+            n_edge,
+            self.dim_full,
+            self.dim_full,
+            dtype=edge_quaternion.dtype,
+            device=edge_quaternion.device,
+        )
+        D_full[:, 0, 0] = 1.0
+
+        for l_degree, block in self._compute_all_blocks(edge_quaternion):
+            base = l_degree * l_degree
+            size = 2 * l_degree + 1
+            D_full[:, base : base + size, base : base + size] = block
 
         if self.lmax >= self.poly_lmin:
             with nvtx_range("WignerD/polynomial"):
@@ -641,6 +652,66 @@ def forward(
             Dt_full = Dt_full.contiguous()
         return D_full, Dt_full
 
+    def band_supported(self) -> bool:
+        """Return whether :meth:`forward_band` covers this configuration.
+
+        The band layout stores three structural rows per degree block, which
+        only the per-degree ladder below the polynomial tail can produce.
+        """
+        return self.lmax < self.poly_lmin
+
+    def forward_band(self, edge_quaternion: torch.Tensor) -> torch.Tensor:
+        """
+        Build the structural Wigner band in the flash packed convention.
+
+        The returned tensor has shape ``(E, 3 * D - 2)`` with
+        ``D = (lmax + 1)^2`` and holds, per output column ``d`` of degree
+        ``l = degree(d)``, the three Wigner entries the reduced ``mmax = 1``
+        consumers read:
+
+        - slots ``[0, D)``: ``D[l^2 + l, d]`` (the ``m = 0`` row),
+        - slots ``[D, 2D - 2)``: ``D[l^2 + l - 1, d]`` (``m = -1`` row,
+          degrees ``l >= 1`` only, stored at slot ``D + d - 1``),
+        - slots ``[2D - 1, 3D - 2)``: ``D[l^2 + l + 1, d]`` (``m = +1`` row,
+          stored at slot ``2D + d - 2``).
+
+        The stored values are the same floats the dense layout keeps at those
+        ``(row, column)`` pairs; only the storage order differs, so consumers
+        addressing the band produce bit-identical results. The band is a
+        projection for the reduced (``mmax = 1``) consumers, not the full
+        matrix: the remaining block rows are never read by them, and the
+        dense layout cannot be reconstructed from the band alone.
+
+        Parameters
+        ----------
+        edge_quaternion
+            Unit quaternions with shape ``(E, 4)``.
+
+        Returns
+        -------
+        torch.Tensor
+            Band tensor with shape ``(E, 3 * dim_full - 2)``.
+        """
+        if not self.band_supported():
+            raise RuntimeError(
+                f"banded Wigner storage requires lmax < {self.poly_lmin}, "
+                f"got lmax={self.lmax}"
+            )
+        edge_quaternion = quaternion_normalize(
+            edge_quaternion.to(dtype=self.dtype),
+            eps=self.eps,
+        )
+        n_edge = edge_quaternion.shape[0]
+        blocks = self._compute_all_blocks(edge_quaternion)
+        m0_rows = [edge_quaternion.new_ones((n_edge, 1))]
+        mm_rows: list[torch.Tensor] = []
+        mp_rows: list[torch.Tensor] = []
+        for l_degree, block in blocks:
+            m0_rows.append(block[:, l_degree, :])
+            mm_rows.append(block[:, l_degree - 1, :])
+            mp_rows.append(block[:, l_degree + 1, :])
+        return torch.cat(m0_rows + mm_rows + mp_rows, dim=1).contiguous()
+
     def forward_zonal(
         self,
         edge_quaternion: torch.Tensor,
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..e98a627c93 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -516,6 +516,7 @@ def _rotate_mix_fwd_kernel(
         CW: tl.constexpr,  # true C_wide; BC = next_power_of_2(CW) lanes with mask
         BC: tl.constexpr,
         RANK: tl.constexpr,
+        PACKED: tl.constexpr,
     ):
         """One program per edge, channels vectorized.
 
@@ -538,7 +539,12 @@ def _rotate_mix_fwd_kernel(
         cmask = chan < CW
         src = tl.load(src_ptr + edge).to(tl.int64)
         x_base = x_ptr + src * x_sn
-        d_base = w_ptr + edge * DIM * DIM
+        # PACKED: wigner holds the (E, 3*D-2) structural band (flash packed
+        # convention); the band slots replace the dense (row, column) offsets.
+        if PACKED:
+            d_base = w_ptr + edge * (3 * DIM - 2)
+        else:
+            d_base = w_ptr + edge * DIM * DIM
 
         # === Phase 1. Rotate to the local frame (registers) ===
         xrows = ()
@@ -555,12 +561,20 @@ def _rotate_mix_fwd_kernel(
             acc0 = tl.zeros((BC,), dtype=tl.float32)
             accm = tl.zeros((BC,), dtype=tl.float32)
             accp = tl.zeros((BC,), dtype=tl.float32)
+            if PACKED:
+                off0, offm, offp = base, DIM + base - 1, 2 * DIM + base - 2
+            else:
+                off0, offm, offp = (
+                    r0 * DIM + base,
+                    (r0 - 1) * DIM + base,
+                    (r0 + 1) * DIM + base,
+                )
             for j in tl.static_range(2 * l + 1):
                 xv = xrows[l * l + j]
-                acc0 += tl.load(d_base + r0 * DIM + base + j) * xv
+                acc0 += tl.load(d_base + off0 + j) * xv
                 if l >= 1:
-                    accm += tl.load(d_base + (r0 - 1) * DIM + base + j) * xv
-                    accp += tl.load(d_base + (r0 + 1) * DIM + base + j) * xv
+                    accm += tl.load(d_base + offm + j) * xv
+                    accp += tl.load(d_base + offp + j) * xv
             rows0 = rows0 + (acc0,)
             if l >= 1:
                 rows_m = rows_m + (accm,)
@@ -648,6 +662,7 @@ def _rotate_mix_bwd_kernel(
         CW: tl.constexpr,
         BC: tl.constexpr,
         RANK: tl.constexpr,
+        PACKED: tl.constexpr,
     ):
         """Backward of the fused front end (one program per edge).
 
@@ -674,7 +689,10 @@ def _rotate_mix_bwd_kernel(
                 tl.load(cb_ptr + r * CW + chan, mask=cmask, other=0.0).to(tl.float32),
             )
         x_base = x_ptr + src * x_sn
-        d_base = w_ptr + edge * DIM * DIM
+        if PACKED:
+            d_base = w_ptr + edge * (3 * DIM - 2)
+        else:
+            d_base = w_ptr + edge * DIM * DIM
         if RANK == 0:
             kc_base = kc_ptr + edge * NS0 * CW
             gkc_base = gkc_ptr + edge * NS0 * CW
@@ -698,12 +716,20 @@ def _rotate_mix_bwd_kernel(
             acc0 = tl.zeros((BC,), dtype=tl.float32)
             accm = tl.zeros((BC,), dtype=tl.float32)
             accp = tl.zeros((BC,), dtype=tl.float32)
+            if PACKED:
+                off0, offm, offp = base, DIM + base - 1, 2 * DIM + base - 2
+            else:
+                off0, offm, offp = (
+                    r0 * DIM + base,
+                    (r0 - 1) * DIM + base,
+                    (r0 + 1) * DIM + base,
+                )
             for j in tl.static_range(2 * l + 1):
                 xv = xrows[l * l + j]
-                acc0 += tl.load(d_base + r0 * DIM + base + j) * xv
+                acc0 += tl.load(d_base + off0 + j) * xv
                 if l >= 1:
-                    accm += tl.load(d_base + (r0 - 1) * DIM + base + j) * xv
-                    accp += tl.load(d_base + (r0 + 1) * DIM + base + j) * xv
+                    accm += tl.load(d_base + offm + j) * xv
+                    accp += tl.load(d_base + offp + j) * xv
             rows0 = rows0 + (acc0,)
             if l >= 1:
                 rows_m = rows_m + (accm,)
@@ -755,7 +781,10 @@ def _rotate_mix_bwd_kernel(
                         )
 
         # === Phase 2. Rotation backward with g_local formed on the fly ===
-        gd_base = gw_ptr + edge * DIM * DIM
+        if PACKED:
+            gd_base = gw_ptr + edge * (3 * DIM - 2)
+        else:
+            gd_base = gw_ptr + edge * DIM * DIM
         for l in tl.static_range(L + 1):
             base = l * l
             r0 = base + l
@@ -797,16 +826,24 @@ def _rotate_mix_bwd_kernel(
                         gp += keff * gy[NS0 + L + o]
             for j in tl.static_range(2 * l + 1):
                 col = base + j
+                if PACKED:
+                    a0, am, ap = col, DIM + col - 1, 2 * DIM + col - 2
+                else:
+                    a0, am, ap = (
+                        r0 * DIM + col,
+                        (r0 - 1) * DIM + col,
+                        (r0 + 1) * DIM + col,
+                    )
                 xv = xrows[l * l + j]
-                w0 = tl.load(d_base + r0 * DIM + col)
+                w0 = tl.load(d_base + a0)
                 gx_row = w0 * g0
-                tl.store(gd_base + r0 * DIM + col, tl.sum(g0 * xv))
+                tl.store(gd_base + a0, tl.sum(g0 * xv))
                 if l >= 1:
-                    wmv = tl.load(d_base + (r0 - 1) * DIM + col)
-                    wpv = tl.load(d_base + (r0 + 1) * DIM + col)
+                    wmv = tl.load(d_base + am)
+                    wpv = tl.load(d_base + ap)
                     gx_row += wmv * gm + wpv * gp
-                    tl.store(gd_base + (r0 - 1) * DIM + col, tl.sum(gm * xv))
-                    tl.store(gd_base + (r0 + 1) * DIM + col, tl.sum(gp * xv))
+                    tl.store(gd_base + am, tl.sum(gm * xv))
+                    tl.store(gd_base + ap, tl.sum(gp * xv))
                 tl.store(
                     gxe_ptr + edge * DIM * CW + col * CW + chan, gx_row, mask=cmask
                 )
@@ -831,6 +868,7 @@ def _rotate_mix_bwd_block_kernel(
         CP: tl.constexpr,  # next power of two >= CW (vector lane count)
         RANK: tl.constexpr,
         BLOCK_E: tl.constexpr,
+        PACKED: tl.constexpr,
     ):
         """Edge-block variant of the rotate+mix backward.
 
@@ -874,8 +912,12 @@ def _rotate_mix_bwd_block_kernel(
 
         src = tl.load(src_ptr + eq, mask=e_mask, other=0).to(tl.int64)
         x_base = x_ptr + (src * x_sn)[:, None]
-        d_base = w_ptr + eq * DIM * DIM
-        gd_base = gw_ptr + eq * DIM * DIM
+        if PACKED:
+            d_base = w_ptr + eq * (3 * DIM - 2)
+            gd_base = gw_ptr + eq * (3 * DIM - 2)
+        else:
+            d_base = w_ptr + eq * DIM * DIM
+            gd_base = gw_ptr + eq * DIM * DIM
         gxe_base = gxe_ptr + (eq * DIM * CW)[:, None]
         if RANK == 0:
             kc_base = kc_ptr + (eq * NS0 * CW)[:, None]
@@ -918,17 +960,21 @@ def _rotate_mix_bwd_block_kernel(
             xl0 = tl.zeros((BLOCK_E, CP), dtype=tl.float32)
             xlm = tl.zeros((BLOCK_E, CP), dtype=tl.float32)
             xlp = tl.zeros((BLOCK_E, CP), dtype=tl.float32)
+            if PACKED:
+                off0, offm, offp = base, DIM + base - 1, 2 * DIM + base - 2
+            else:
+                off0, offm, offp = (
+                    r0 * DIM + base,
+                    (r0 - 1) * DIM + base,
+                    (r0 + 1) * DIM + base,
+                )
             for j in tl.static_range(2 * l + 1):
                 xv = xrows[j]
-                w0 = tl.load(d_base + r0 * DIM + base + j, mask=e_mask, other=0.0)
+                w0 = tl.load(d_base + off0 + j, mask=e_mask, other=0.0)
                 xl0 += w0[:, None] * xv
                 if l >= 1:
-                    wm = tl.load(
-                        d_base + (r0 - 1) * DIM + base + j, mask=e_mask, other=0.0
-                    )
-                    wp = tl.load(
-                        d_base + (r0 + 1) * DIM + base + j, mask=e_mask, other=0.0
-                    )
+                    wm = tl.load(d_base + offm + j, mask=e_mask, other=0.0)
+                    wp = tl.load(d_base + offp + j, mask=e_mask, other=0.0)
                     xlm += wm[:, None] * xv
                     xlp += wp[:, None] * xv
 
@@ -981,24 +1027,24 @@ def _rotate_mix_bwd_block_kernel(
             # Rotation backward: node gradient rows and Wigner gradients.
             for j in tl.static_range(2 * l + 1):
                 col = base + j
+                if PACKED:
+                    a0, am, ap = col, DIM + col - 1, 2 * DIM + col - 2
+                else:
+                    a0, am, ap = (
+                        r0 * DIM + col,
+                        (r0 - 1) * DIM + col,
+                        (r0 + 1) * DIM + col,
+                    )
                 xv = xrows[j]
-                w0 = tl.load(d_base + r0 * DIM + col, mask=e_mask, other=0.0)
+                w0 = tl.load(d_base + a0, mask=e_mask, other=0.0)
                 gx_row = w0[:, None] * g0
-                tl.store(gd_base + r0 * DIM + col, tl.sum(g0 * xv, axis=1), mask=e_mask)
+                tl.store(gd_base + a0, tl.sum(g0 * xv, axis=1), mask=e_mask)
                 if l >= 1:
-                    wm = tl.load(d_base + (r0 - 1) * DIM + col, mask=e_mask, other=0.0)
-                    wp = tl.load(d_base + (r0 + 1) * DIM + col, mask=e_mask, other=0.0)
+                    wm = tl.load(d_base + am, mask=e_mask, other=0.0)
+                    wp = tl.load(d_base + ap, mask=e_mask, other=0.0)
                     gx_row += wm[:, None] * gm + wp[:, None] * gp
-                    tl.store(
-                        gd_base + (r0 - 1) * DIM + col,
-                        tl.sum(gm * xv, axis=1),
-                        mask=e_mask,
-                    )
-                    tl.store(
-                        gd_base + (r0 + 1) * DIM + col,
-                        tl.sum(gp * xv, axis=1),
-                        mask=e_mask,
-                    )
+                    tl.store(gd_base + am, tl.sum(gm * xv, axis=1), mask=e_mask)
+                    tl.store(gd_base + ap, tl.sum(gp * xv, axis=1), mask=e_mask)
                 tl.store(gxe_base + col * CW + chan[None, :], gx_row, mask=em)
 
     @triton.jit
@@ -2481,7 +2527,13 @@ def _rotate_mix_impl(
     # The source CSR view rides through the forward untouched so the autograd
     # context can hand it to the backward's segment reduction.
     del src_order, src_rowptr
+    packed = wigner.dim() == 2
     if not _use_triton(x):
+        if packed:
+            raise RuntimeError(
+                "the banded Wigner layout requires the Triton rotate-mix kernels; "
+                "unset DP_PACKED_D_INFER or enable the Triton inference path"
+            )
         return _rotate_mix_reference(x, src, wigner, kc, cb, lmax, n_focus, rank)
     n_edge = src.shape[0]
     c_wide = int(x.shape[2])
@@ -2506,6 +2558,7 @@ def _rotate_mix_impl(
         CW=c_wide,
         BC=triton.next_power_of_2(c_wide),
         RANK=int(rank),
+        PACKED=packed,
         num_warps=warps,
         num_stages=stages,
     )
@@ -2523,7 +2576,13 @@ def _rotate_mix_bwd_impl(
     n_focus: int,
     rank: int,
 ) -> tuple[Tensor, Tensor, Tensor]:
+    packed = wigner.dim() == 2
     if not _use_triton(x):
+        if packed:
+            raise RuntimeError(
+                "the banded Wigner layout requires the Triton rotate-mix kernels; "
+                "unset DP_PACKED_D_INFER or enable the Triton inference path"
+            )
         return _rotate_mix_backward_reference(
             grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank
         )
@@ -2563,6 +2622,7 @@ def _rotate_mix_bwd_impl(
             CP=triton.next_power_of_2(c_wide),
             RANK=int(rank),
             BLOCK_E=block_e,
+            PACKED=packed,
             num_warps=warps,
             num_stages=stages,
         )
@@ -2586,6 +2646,7 @@ def _rotate_mix_bwd_impl(
         CW=c_wide,
         BC=triton.next_power_of_2(c_wide),
         RANK=int(rank),
+        PACKED=packed,
         num_warps=warps,
         num_stages=stages,
     )
@@ -5032,6 +5093,14 @@ def _is_supported(conv: SO2Convolution) -> bool:
 _mixing_stack_op.register_autocast("cuda", torch.bfloat16)
 
 
+def _packed_d(edge_cache) -> Tensor:
+    """Return the banded Wigner storage when present, else the dense surface."""
+    band = getattr(edge_cache, "D_packed", None)
+    if band is not None:
+        return band
+    return edge_cache.D_full
+
+
 class _TritonRotateMix:
     """Per-convolution entry running rotate-to-local + degree mixing fused.
 
@@ -5100,7 +5169,7 @@ def __call__(
             src,
             src_order,
             src_rowptr,
-            edge_cache.D_full,
+            _packed_d(edge_cache),
             kc.contiguous(),
             cb.contiguous(),
             conv.lmax,
@@ -5332,7 +5401,7 @@ def __call__(
             src,
             src_order,
             src_rowptr,
-            edge_cache.D_full,
+            _packed_d(edge_cache),
             kc.contiguous(),
             cb.contiguous(),
             conv.lmax,
diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py
index c6e54bd17f..ca0b796bad 100644
--- a/deepmd/pt_expt/kernels/utils.py
+++ b/deepmd/pt_expt/kernels/utils.py
@@ -380,3 +380,23 @@ def use_amp_infer() -> bool:
         ``True`` when ``DP_AMP_INFER`` is set to a truthy value.
     """
     return os.environ.get("DP_AMP_INFER", "0").strip().lower() in _INFER_TRUE
+
+
+def use_packed_d_infer() -> bool:
+    """Return whether inference stores the Wigner rotations in banded form.
+
+    The flag is controlled by the ``DP_PACKED_D_INFER`` environment variable
+    and is read when the edge cache is built. When truthy (and the model
+    configuration supports it), the edge cache stores only the ``3 * D - 2``
+    structural band entries of the block-diagonal Wigner matrices per edge
+    instead of the dense ``(E, D, D)`` surfaces; the fused Triton consumers
+    (rotate-mix, flash aggregation) address the band directly, and paths that
+    still need the dense layout unpack it on demand. The stored values are the
+    same floats, so consumer outputs are bit-identical to the dense mode.
+
+    Returns
+    -------
+    bool
+        ``True`` when ``DP_PACKED_D_INFER`` is set to a truthy value.
+    """
+    return os.environ.get("DP_PACKED_D_INFER", "0").strip().lower() in _INFER_TRUE

From 3c65fbde48272baef45f59fa344f843f05adc908 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sat, 10 Oct 2026 22:04:37 +0800
Subject: [PATCH 02/15] test(pt): bitwise coverage for the banded Wigner
 storage

---
 .../pt/model/test_descriptor_sezm_triton.py   | 157 ++++++++++++++++++
 1 file changed, 157 insertions(+)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index f359f073d6..4191348f0a 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2053,3 +2053,160 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+class TestSeZMTritonWignerBand(unittest.TestCase):
+    """Cross-check the banded Wigner storage against the dense layout.
+
+    The band keeps the three structural rows per degree block in the flash
+    packed convention; every consumer addressing it must stay bit-identical
+    to the dense ``(E, D, D)`` layout it replaces.
+    """
+
+    N_EDGE = 2048
+    N_NODE = 128
+
+    def _quaternion(self, n, seed):
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+        quaternion = torch.randn(n, 4, device="cuda", generator=generator)
+        return quaternion / quaternion.norm(dim=-1, keepdim=True)
+
+    def _band_segments(self, band, dim):
+        return band[:, :dim], band[:, dim : 2 * dim - 1], band[:, 2 * dim - 1 :]
+
+    @_GPU_KERNELS
+    def test_band_matches_dense_blocks(self):
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            WignerDCalculator,
+        )
+
+        for lmax in (1, 2, 3, 4, 5):
+            with self.subTest(lmax=lmax):
+                calc = WignerDCalculator(lmax, dtype=torch.float32).to("cuda")
+                quaternion = self._quaternion(64, 17 + lmax)
+                dense, _ = calc(quaternion)
+                band = calc.forward_band(quaternion)
+                dim = (lmax + 1) ** 2
+                self.assertEqual(tuple(band.shape), (64, 3 * dim - 2))
+                m0, mm, mp = self._band_segments(band, dim)
+                self.assertTrue(torch.equal(m0[:, 0], dense[:, 0, 0]))
+                for degree in range(1, lmax + 1):
+                    base, size = degree * degree, 2 * degree + 1
+                    r0 = base + degree
+                    cols = slice(base, base + size)
+                    self.assertTrue(torch.equal(m0[:, cols], dense[:, r0, cols]))
+                    self.assertTrue(
+                        torch.equal(
+                            mm[:, base - 1 : base - 1 + size],
+                            dense[:, r0 - 1, cols],
+                        )
+                    )
+                    self.assertTrue(
+                        torch.equal(
+                            mp[:, base - 1 : base - 1 + size],
+                            dense[:, r0 + 1, cols],
+                        )
+                    )
+
+    @_GPU_KERNELS
+    def test_rotate_mix_packed_bitwise(self):
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            WignerDCalculator,
+        )
+        from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as svp
+
+        lmax, n_focus, focus_dim = 3, 2, 32
+        c_wide = n_focus * focus_dim
+        dim = (lmax + 1) ** 2
+        calc = WignerDCalculator(lmax, dtype=torch.float32).to("cuda")
+        quaternion = self._quaternion(self.N_EDGE, 23)
+        dense = calc(quaternion)[0]
+        band = calc.forward_band(quaternion)
+
+        generator = torch.Generator(device="cuda").manual_seed(29)
+        x = (
+            torch.randn(self.N_NODE, dim, c_wide, device="cuda", generator=generator)
+            * 0.5
+        ).contiguous()
+        src = torch.randint(
+            0, self.N_NODE, (self.N_EDGE,), device="cuda", generator=generator
+        )
+        kc = torch.randn(self.N_EDGE, 25, device="cuda", generator=generator) * 0.2
+        cb = (torch.randn(1, c_wide, device="cuda", generator=generator) * 0.3).reshape(
+            -1
+        )
+
+        u0 = svp._rotate_mix_impl(x, src, None, None, dense, kc, cb, lmax, n_focus, 1)
+        u1 = svp._rotate_mix_impl(x, src, None, None, band, kc, cb, lmax, n_focus, 1)
+        self.assertTrue(torch.equal(u0, u1))
+        grad_u = torch.randn(u0.shape, device="cuda", generator=generator)
+        gx0, gw0, gkc0 = svp._rotate_mix_bwd_impl(
+            grad_u, x, src, dense, kc, cb, lmax, n_focus, 1
+        )
+        gx1, gw1, gkc1 = svp._rotate_mix_bwd_impl(
+            grad_u, x, src, band, kc, cb, lmax, n_focus, 1
+        )
+        self.assertTrue(torch.equal(gx0, gx1))
+        self.assertTrue(torch.equal(gkc0, gkc1))
+        # The dense Wigner gradient is zero off the structural band; compare
+        # the three band segments slot by slot.
+        m0 = [gw0[:, 0, 0:1]]
+        mm: list[torch.Tensor] = []
+        mp: list[torch.Tensor] = []
+        for degree in range(1, lmax + 1):
+            base, size = degree * degree, 2 * degree + 1
+            r0 = base + degree
+            m0.append(gw0[:, r0, base : base + size])
+            mm.append(gw0[:, r0 - 1, base : base + size])
+            mp.append(gw0[:, r0 + 1, base : base + size])
+        self.assertTrue(torch.equal(torch.cat(m0, dim=1), gw1[:, :dim]))
+        self.assertTrue(torch.equal(torch.cat(mm, dim=1), gw1[:, dim : 2 * dim - 1]))
+        self.assertTrue(torch.equal(torch.cat(mp, dim=1), gw1[:, 2 * dim - 1 :]))
+
+    @_GPU_KERNELS
+    def test_value_path_packed_bitwise(self):
+        from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import (
+            make_triton_value_path,
+        )
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            WignerDCalculator,
+        )
+
+        case = (3, 64, 2, 96, 4, "degree_channel", 1)
+        value_path_cls = TestSeZMTritonValuePath
+        conv = value_path_cls._build_conv(self, *case)
+        fused = make_triton_value_path(conv)
+        self.assertIsNotNone(fused)
+        x, cache, radial = value_path_cls._edge_inputs(self, conv, case[0], case[1])
+
+        calc = WignerDCalculator(case[0], dtype=torch.float32).to("cuda")
+        # Deterministic rotations; both layouts carry the same entries.
+        quaternion = self._quaternion(cache.D_full.shape[0], 31)
+        dense, _ = calc(quaternion)
+        band = calc.forward_band(quaternion)
+
+        import copy
+
+        # One shared cotangent so both layouts differentiate the same output.
+        cotangent = None
+
+        def run(rotation_layout):
+            nonlocal cotangent
+            source = copy.copy(cache)
+            source.D_full = dense if rotation_layout == "dense" else None
+            source.D_packed = None if rotation_layout == "dense" else band
+            rad = radial.clone().requires_grad_(True)
+            out, _ = fused(x, source, rad)
+            if cotangent is None:
+                cotangent = torch.randn(
+                    out.shape,
+                    device=out.device,
+                    generator=torch.Generator(device=out.device).manual_seed(37),
+                )
+            (grad_radial,) = torch.autograd.grad(out, rad, cotangent)
+            return out.detach(), grad_radial
+
+        out0, grad0 = run("dense")
+        out1, grad1 = run("band")
+        self.assertTrue(torch.equal(out0, out1))
+        self.assertTrue(torch.equal(grad0, grad1))

From 788964354b6fcb55abd13c6039e65c85458d2954 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 10 Oct 2026 14:06:41 +0000
Subject: [PATCH 03/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/descriptor/sezm_nn/edge_cache.py     | 1 +
 source/tests/pt/model/test_descriptor_sezm_triton.py | 6 +++---
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
index dcb09f4c82..e13833b8ef 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
@@ -28,6 +28,7 @@
 from deepmd.pt_expt.kernels.utils import (
     use_packed_d_infer,
 )
+
 from .utils import (
     get_promoted_dtype,
     nvtx_range,
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 4191348f0a..be248c7f01 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2165,12 +2165,12 @@ def test_rotate_mix_packed_bitwise(self):
 
     @_GPU_KERNELS
     def test_value_path_packed_bitwise(self):
-        from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import (
-            make_triton_value_path,
-        )
         from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
             WignerDCalculator,
         )
+        from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import (
+            make_triton_value_path,
+        )
 
         case = (3, 64, 2, 96, 4, "degree_channel", 1)
         value_path_cls = TestSeZMTritonValuePath

From 93472fa410e48bcfd28c8111ad66872b2950764d Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 11 Oct 2026 22:28:02 +0800
Subject: [PATCH 04/15] perf(pt): op-bound banded Wigner assembly with analytic
 backward

---
 .../pt/model/descriptor/sezm_nn/edge_cache.py |   2 +-
 deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 232 ++++++++++++++++++
 2 files changed, 233 insertions(+), 1 deletion(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
index e13833b8ef..29e56d316d 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
@@ -722,7 +722,7 @@ def _build_edge_wigner(
         # structural rows per degree block, so the dense surfaces are never
         # materialized. Reference paths that need the dense layout must not
         # enable the flag.
-        return None, None, edge_quat, calc.forward_band(edge_quat)
+        return None, None, edge_quat, calc.forward_band_fused(edge_quat)
     D_full, Dt_full = wigner_calc(edge_quat)
     return D_full, Dt_full, edge_quat, None
 
diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index bd31b5cf39..515035f38e 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -712,6 +712,51 @@ def forward_band(self, edge_quaternion: torch.Tensor) -> torch.Tensor:
             mp_rows.append(block[:, l_degree + 1, :])
         return torch.cat(m0_rows + mm_rows + mp_rows, dim=1).contiguous()
 
+    def band_op_supported(self) -> bool:
+        """Return whether the banded assembly can run behind the fused op.
+
+        The op boundary makes the band a save point for the compiled
+        backward: without it the whole quaternion -> monomial -> product
+        chain is replayed (once per backward section) instead of the tiny
+        band being kept. The op covers the CUDA fp32 inference ladder the
+        production model uses (degrees 1..3 with the Triton monomial
+        kernels); anything else keeps the plain composition.
+        """
+        return (
+            self.lmax in (1, 2, 3)
+            and not self.training
+            and bool(getattr(self, "_use_triton_monomials", False))
+            and "exp_l2" in getattr(self, "_monomial_exponents_flat", {})
+            and "exp_l3" in getattr(self, "_monomial_exponents_flat", {})
+            and hasattr(self, "_l2_monomial_coeff")
+        )
+
+    def forward_band_fused(self, edge_quaternion: torch.Tensor) -> torch.Tensor:
+        """
+        Build the structural band through the custom-op boundary.
+
+        Same values as :meth:`forward_band`; the normalization stays in the
+        traced graph so its own cotangent keeps flowing, and the band
+        production becomes opaque to the compiler, which then saves the
+        ``(E, 3D-2)`` band for the backward instead of replaying the whole
+        quaternion chain per backward section.
+        """
+        if not self.band_op_supported():
+            return self.forward_band(edge_quaternion)
+        q_unit = quaternion_normalize(
+            edge_quaternion.to(dtype=self.dtype), eps=self.eps
+        )
+        return wigner_band_op(
+            q_unit,
+            self.l1_perm,
+            self.l1_sign_outer,
+            self._l2_monomial_coeff,
+            self.small_order_kernels.C_l3,
+            self._monomial_exponents_flat["exp_l2"],
+            self._monomial_exponents_flat["exp_l3"],
+            int(self.lmax),
+        )
+
     def forward_zonal(
         self,
         edge_quaternion: torch.Tensor,
@@ -2029,3 +2074,190 @@ def deserialize(cls, data: dict[str, Any]) -> WignerDCalculator:
         raise NotImplementedError(
             "WignerDCalculator.deserialize should be called by parent with lmax/dtype"
         )
+
+
+def _band_group_offsets(lmax: int) -> tuple[list[int], int]:
+    """Per-degree offsets of the three band row groups and the full dim.
+
+    The band packs, for degrees ``l = 1..lmax``, the ``m = 0`` row of each
+    block (preceded by the degree-0 constant), then all ``m = -1`` rows,
+    then all ``m = +1`` rows; degrees ascend inside every group.
+    """
+    dim = (lmax + 1) ** 2
+    widths = [2 * ell + 1 for ell in range(1, lmax + 1)]
+    bases = [sum(widths[:i]) for i in range(len(widths))]
+    return bases, dim
+
+
+def _wigner_band_impl(
+    quat: torch.Tensor,
+    l1_perm: torch.Tensor,
+    l1_sign: torch.Tensor,
+    coeff_l2: torch.Tensor,
+    c_l3: torch.Tensor,
+    exp_l2: list[int],
+    exp_l3: list[int],
+    lmax: int,
+) -> torch.Tensor:
+    """Band assembly behind an opaque boundary (input must be unit norm).
+
+    Mirrors :meth:`WignerDCalculator.forward_band` for the ``lmax <= 3``
+    ladder: the l=1 block is the permuted signed rotation matrix, the
+    l=2/l=3 blocks are the monomial evaluations contracted with the frozen
+    coefficient tables.
+    """
+    from deepmd.pt_expt.kernels.triton.sezm.wigner_monomials import (
+        wigner_monomials,
+    )
+
+    n_edge = quat.shape[0]
+    m0_rows = [quat.new_ones((n_edge, 1))]
+    mm_rows: list[torch.Tensor] = []
+    mp_rows: list[torch.Tensor] = []
+    if lmax >= 1:
+        rot = quaternion_to_rotation_matrix(quat)
+        block = rot.index_select(-2, l1_perm).index_select(-1, l1_perm)
+        block = block * l1_sign
+        m0_rows.append(block[:, 1, :])
+        mm_rows.append(block[:, 0, :])
+        mp_rows.append(block[:, 2, :])
+    if lmax >= 2:
+        mono = wigner_monomials(quat, exp_l2, 4)
+        block = torch.matmul(mono, coeff_l2).view(n_edge, 5, 5)
+        m0_rows.append(block[:, 2, :])
+        mm_rows.append(block[:, 1, :])
+        mp_rows.append(block[:, 3, :])
+    if lmax >= 3:
+        mono = wigner_monomials(quat, exp_l3, 6)
+        block = torch.matmul(mono, c_l3.transpose(0, 1)).view(n_edge, 7, 7)
+        m0_rows.append(block[:, 3, :])
+        mm_rows.append(block[:, 2, :])
+        mp_rows.append(block[:, 4, :])
+    return torch.cat(m0_rows + mm_rows + mp_rows, dim=1).contiguous()
+
+
+_wigner_band_op = torch.library.triton_op("sezm_triton::wigner_band", mutates_args=())(
+    _wigner_band_impl
+)
+
+
+@_wigner_band_op.register_fake
+def _(
+    quat: torch.Tensor,
+    l1_perm: torch.Tensor,
+    l1_sign: torch.Tensor,
+    coeff_l2: torch.Tensor,
+    c_l3: torch.Tensor,
+    exp_l2: list[int],
+    exp_l3: list[int],
+    lmax: int,
+) -> torch.Tensor:
+    dim = (lmax + 1) ** 2
+    return quat.new_empty((quat.shape[0], 3 * dim - 2))
+
+
+def _wigner_band_setup_context(
+    ctx: Any, inputs: tuple[Any, ...], output: torch.Tensor
+) -> None:
+    quat, l1_perm, l1_sign, coeff_l2, c_l3, exp_l2, exp_l3, lmax = inputs
+    ctx.save_for_backward(quat, l1_perm, l1_sign, coeff_l2, c_l3)
+    ctx.exp_l2 = exp_l2
+    ctx.exp_l3 = exp_l3
+    ctx.lmax = lmax
+
+
+def _wigner_band_backward(ctx: Any, grad_band: torch.Tensor) -> tuple:
+    """Analytic band cotangent: rows -> monomial space -> quaternion.
+
+    The three stored rows of each degree block are contiguous block rows,
+    so each degree contracts its ``3 x (2l+1)`` gradient slice against the
+    matching coefficient rows in one small product, hands the result to the
+    monomial operator's own backward, and the l=1 rows differentiate the
+    rotation-matrix formula directly.
+    """
+    from deepmd.pt_expt.kernels.triton.sezm.wigner_monomials import (
+        _monomials_bwd_op,
+    )
+
+    quat, l1_perm, l1_sign, coeff_l2, c_l3 = ctx.saved_tensors
+    lmax = ctx.lmax
+    bases, dim = _band_group_offsets(lmax)
+    grad_quat = torch.zeros_like(quat)
+    if lmax >= 2:
+        width = 5
+        g2 = torch.cat(
+            [
+                grad_band[
+                    :,
+                    dim + bases[1] : dim + bases[1] + width,
+                ],
+                grad_band[:, 1 + bases[1] : 1 + bases[1] + width],
+                grad_band[
+                    :,
+                    2 * dim - 1 + bases[1] : 2 * dim - 1 + bases[1] + width,
+                ],
+            ],
+            dim=1,
+        )
+        # block rows (1, 2, 3) are flat columns [5, 20) of the product
+        g_mono = torch.matmul(g2, coeff_l2[:, 5 : 5 + 3 * width].t())
+        # the monomial backward runs as its own operator: routing through
+        # autograd would re-execute the monomial forward inside the traced
+        # backward and hand back the replay we are removing
+        grad_quat = grad_quat + _monomials_bwd_op(g_mono, quat, ctx.exp_l2, 4)
+    if lmax >= 3:
+        width = 7
+        g3 = torch.cat(
+            [
+                grad_band[
+                    :,
+                    dim + bases[2] : dim + bases[2] + width,
+                ],
+                grad_band[:, 1 + bases[2] : 1 + bases[2] + width],
+                grad_band[
+                    :,
+                    2 * dim - 1 + bases[2] : 2 * dim - 1 + bases[2] + width,
+                ],
+            ],
+            dim=1,
+        )
+        # block rows (2, 3, 4) are flat columns [14, 35) of the product
+        g_mono = torch.matmul(g3, c_l3[2 * width : 5 * width, :])
+        grad_quat = grad_quat + _monomials_bwd_op(g_mono, quat, ctx.exp_l3, 6)
+    with torch.enable_grad():
+        q = quat.detach().requires_grad_(True)
+        if lmax >= 1:
+            # rows (0, 1, 2) of the l=1 block come from the three groups
+            g1 = torch.cat(
+                [
+                    grad_band[:, dim + bases[0] : dim + bases[0] + 3],
+                    grad_band[:, 1 + bases[0] : 1 + bases[0] + 3],
+                    grad_band[:, 2 * dim - 1 + bases[0] : 2 * dim - 1 + bases[0] + 3],
+                ],
+                dim=1,
+            ).view(-1, 3, 3)
+            rot = quaternion_to_rotation_matrix(q)
+            block1 = rot.index_select(-2, l1_perm).index_select(-1, l1_perm) * l1_sign
+            grad_quat = grad_quat + torch.autograd.grad(block1, q, g1)[0]
+    return (grad_quat, None, None, None, None, None, None, None)
+
+
+torch.library.register_autograd(
+    "sezm_triton::wigner_band",
+    _wigner_band_backward,
+    setup_context=_wigner_band_setup_context,
+)
+
+
+def wigner_band_op(
+    quat: torch.Tensor,
+    l1_perm: torch.Tensor,
+    l1_sign: torch.Tensor,
+    coeff_l2: torch.Tensor,
+    c_l3: torch.Tensor,
+    exp_l2: list[int],
+    exp_l3: list[int],
+    lmax: int,
+) -> torch.Tensor:
+    """Public entry for the op-bound banded Wigner assembly."""
+    return _wigner_band_op(quat, l1_perm, l1_sign, coeff_l2, c_l3, exp_l2, exp_l3, lmax)

From 8ccbbf1cba7ba8fa0f03b37c8e27abea3d0cfd81 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 11 Oct 2026 22:28:02 +0800
Subject: [PATCH 05/15] perf(pt): reuse the packed band m=0 group for the GIE
 zonal coupling

---
 deepmd/pt/model/descriptor/sezm.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index 7a3373da5e..9fbbaad30f 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -2134,6 +2134,13 @@ def _build_gie_zonal_coupling(
             shared = self._shared_wigner_runs(edge_cache, calc.lmax)
             if shared is not None:
                 return shared
+            packed = getattr(edge_cache, "D_packed", None)
+            if packed is not None and calc.lmax <= self.lmax:
+                # The packed band's m=0 group already carries these entries
+                # (band columns 1..(lmax+1)^2-1 equal the zonal coupling's
+                # packed-row order), so no second quaternion -> monomial ->
+                # block run is scheduled next to the banded assembly.
+                return packed[:, 1 : (calc.lmax + 1) ** 2]
             return calc.forward_zonal(self._edge_quaternion(edge_cache), lmin=1)
         if self.gie_zonal_wigner_calc is None:
             return None

From 85d18f42f8ad044db99fef4b4f4d22adcbee4291 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 11 Oct 2026 22:28:02 +0800
Subject: [PATCH 06/15] test(pt): cover the op-bound band assembly and zonal
 reuse

---
 .../pt/model/test_descriptor_sezm_triton.py   | 97 +++++++++++++++++++
 1 file changed, 97 insertions(+)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index be248c7f01..d0a59db3ee 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2210,3 +2210,100 @@ def run(rotation_layout):
         out1, grad1 = run("band")
         self.assertTrue(torch.equal(out0, out1))
         self.assertTrue(torch.equal(grad0, grad1))
+
+
+class TestSeZMWignerBandOp(unittest.TestCase):
+    """The op-bound banded assembly and its analytic backward."""
+
+    def _make_calc(self, lmax):
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            WignerDCalculator,
+        )
+
+        calc = WignerDCalculator(lmax=lmax, dtype=torch.float32).cuda()
+        calc.eval()
+        return calc
+
+    def _op_band(self, calc, quat):
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            quaternion_normalize,
+            wigner_band_op,
+        )
+
+        return wigner_band_op(
+            quaternion_normalize(quat, eps=calc.eps),
+            calc.l1_perm,
+            calc.l1_sign_outer,
+            calc._l2_monomial_coeff,
+            calc.small_order_kernels.C_l3,
+            calc._monomial_exponents_flat["exp_l2"],
+            calc._monomial_exponents_flat["exp_l3"],
+            int(calc.lmax),
+        )
+
+    @_GPU_KERNELS
+    def test_matches_composition_with_grads(self):
+        import os
+        from unittest.mock import patch
+
+        generator = torch.Generator(device="cuda").manual_seed(23)
+        # the gate is conservative: the op promises the deployed ladder
+        # (all three degree tables present), lower lmax falls back
+        for lmax in (3,):
+            with self.subTest(lmax=lmax):
+                with patch.dict(os.environ, {"DP_TRITON_INFER": "3"}):
+                    calc = self._make_calc(lmax)
+                self.assertTrue(calc.band_op_supported())
+                quat = torch.randn(4096, 4, device="cuda", generator=generator)
+                quat = quat / quat.norm(dim=-1, keepdim=True)
+                quat = quat + 0.01 * torch.randn(
+                    4096, 4, device="cuda", generator=generator
+                )
+
+                q1 = quat.clone().requires_grad_(True)
+                band = self._op_band(calc, q1)
+                q2 = quat.clone().requires_grad_(True)
+                ref = calc.forward_band(q2)
+                torch.testing.assert_close(band, ref, atol=1e-6, rtol=1e-6)
+
+                cotangent = torch.randn(
+                    ref.shape,
+                    device="cuda",
+                    generator=torch.Generator(device="cuda").manual_seed(37),
+                )
+                g1 = torch.autograd.grad(band, q1, cotangent)[0]
+                g2 = torch.autograd.grad(ref, q2, cotangent)[0]
+                torch.testing.assert_close(g1, g2, atol=1e-5, rtol=1e-5)
+
+    @_GPU_KERNELS
+    def test_compile_matches_eager(self):
+        generator = torch.Generator(device="cuda").manual_seed(29)
+        calc = self._make_calc(3)
+        quat = torch.randn(2048, 4, device="cuda", generator=generator)
+        quat = quat / quat.norm(dim=-1, keepdim=True)
+
+        def fn(q):
+            return self._op_band(calc, q)
+
+        q_e = quat.clone().requires_grad_(True)
+        out_e = fn(q_e)
+        cot = torch.randn_like(out_e)
+        g_e = torch.autograd.grad(out_e, q_e, cot)[0]
+
+        q_c = quat.clone().requires_grad_(True)
+        out_c = torch.compile(fn)(q_c)
+        g_c = torch.autograd.grad(out_c, q_c, cot)[0]
+        torch.testing.assert_close(out_c, out_e, atol=1e-6, rtol=1e-6)
+        torch.testing.assert_close(g_c, g_e, atol=1e-5, rtol=1e-5)
+
+    @_GPU_KERNELS
+    def test_zonal_coupling_equals_band_m0_group(self):
+        """The GIE zonal entries are exactly the band's m=0 columns."""
+        generator = torch.Generator(device="cuda").manual_seed(31)
+        calc = self._make_calc(3)
+        quat = torch.randn(2048, 4, device="cuda", generator=generator)
+        quat = quat / quat.norm(dim=-1, keepdim=True)
+        zonal = calc.forward_zonal(quat, lmin=1)
+        band = calc.forward_band(quat)
+        dim = (calc.lmax + 1) ** 2
+        torch.testing.assert_close(zonal, band[:, 1:dim], atol=1e-6, rtol=1e-6)

From e29fbbf182fc080bce7cf666403f58500d1bc880 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 11 Oct 2026 14:34:02 +0000
Subject: [PATCH 07/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/pt/model/test_descriptor_sezm_triton.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index d0a59db3ee..2d343ad4c0 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2244,7 +2244,9 @@ def _op_band(self, calc, quat):
     @_GPU_KERNELS
     def test_matches_composition_with_grads(self):
         import os
-        from unittest.mock import patch
+        from unittest.mock import (
+            patch,
+        )
 
         generator = torch.Generator(device="cuda").manual_seed(23)
         # the gate is conservative: the op promises the deployed ladder

From 3ac4f666b1412da2607f4da7bfe31d597f4455be Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Wed, 14 Oct 2026 21:27:06 +0800
Subject: [PATCH 08/15] fix(pt): pass the banded-layout flag in the rotate-mix
 tuning sweep launches

---
 deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py
index 44c0a91aac..d0e8c7446e 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py
@@ -902,6 +902,8 @@ def launch(warps: int, stages: int) -> None:
             CW=c_wide,
             BC=triton.next_power_of_2(c_wide),
             RANK=1,
+            # sweep the dense layout; the banded form shares this entry
+            PACKED=False,
             num_warps=warps,
             num_stages=stages,
         )
@@ -1002,6 +1004,8 @@ def launch(block_e: int, warps: int, stages: int) -> tuple[torch.Tensor, ...]:
             CP=triton.next_power_of_2(c_wide),
             RANK=1,
             BLOCK_E=block_e,
+            # sweep the dense layout; the banded form shares this entry
+            PACKED=False,
             num_warps=warps,
             num_stages=stages,
         )

From 29d4be56ca5c397dd93ab4b5ec0e04bbd6baf64d Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Thu, 15 Oct 2026 20:30:26 +0800
Subject: [PATCH 09/15] fix(pt): match the fp16x3 operator's two-output
 contract in the tuning sweep

---
 deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py
index d0e8c7446e..2c69572625 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py
@@ -2008,7 +2008,9 @@ def wins_against_fp32() -> bool:
         x_fp32_speed, z_fp32_speed, u_fp32_speed = _mixing_stack_op(
             u0, alpha, w0_all, w1_all, gw_all, lmax, cf, True
         )
-        x_fp16_speed, z_fp16_speed, u_fp16_speed = _mixing_stack_fp16x3_op(
+        # the fp16x3 operator returns (x_local, z_all); its backward
+        # recovers the final identity input internally, unlike the fp32 one
+        x_fp16_speed, z_fp16_speed = _mixing_stack_fp16x3_op(
             u0, alpha, w0_all, w1_all, gw_all, lmax, cf, True
         )
 

From 25bb484ad181043533c54ce27dd57d7cf65352c7 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Fri, 16 Oct 2026 20:03:44 +0800
Subject: [PATCH 10/15] fix(pt): gate banded Wigner storage on every consumer
 addressing the band

---
 deepmd/pt/model/descriptor/sezm.py            | 28 +++++++++++++++++++
 .../pt/model/descriptor/sezm_nn/edge_cache.py | 12 ++++++++
 2 files changed, 40 insertions(+)

diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index 9fbbaad30f..676a8c4630 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -132,6 +132,27 @@
 
 
 @BaseDescriptor.register("SeZM")
+def block_supports_packed_wigner(so2_conv: Any) -> bool:
+    """Whether one block's selected rotation consumers address the band.
+
+    The banded layout serves the Triton value path and the Triton
+    rotate-mix path (whose kernels address the band directly), together
+    with the attention aggregation and the GIE coupling that read the same
+    tensor. Any dense-layout consumer -- the cuTile value path, the fused
+    CUDA convolution, the CUDA training value path, or the plain SO2
+    reference path (no Triton binding at all) -- keeps the dense surfaces.
+    """
+    return (
+        getattr(so2_conv, "_cutile_value_path", None) is None
+        and getattr(so2_conv, "_cuda_conv_fn", None) is None
+        and getattr(so2_conv, "_cuda_value_train", None) is None
+        and (
+            getattr(so2_conv, "_triton_value_path", None) is not None
+            or getattr(so2_conv, "_triton_rotate_mix", None) is not None
+        )
+    )
+
+
 @BaseDescriptor.register("sezm")
 @BaseDescriptor.register("DPA4")
 @BaseDescriptor.register("dpa4")
@@ -1126,6 +1147,11 @@ def __init__(
             and block.so2_conv._flash_atten_trains
             for block in self.blocks
         )
+        # Banded storage additionally requires every selected rotation
+        # consumer to address the band (see ``block_supports_packed_wigner``).
+        self._packed_wigner_consumers = bool(self.blocks) and all(
+            block_supports_packed_wigner(block.so2_conv) for block in self.blocks
+        )
 
         # The envelope and the radial basis are both functions of the pair
         # distance and are cheap enough that the compiler inlines them into
@@ -1427,6 +1453,7 @@ def forward(
                 random_gamma=self.random_gamma and self.training,
                 wigner_calc=self.wigner_calc,
                 build_wigner=self._build_full_wigner(),
+                packed_wigner=self._packed_wigner_consumers,
             )
 
         ebed_dim_0 = self.node_init_dim  # (node_init_lmax+1)^2
@@ -1743,6 +1770,7 @@ def forward_with_edges(
                 random_gamma=self.random_gamma and self.training,
                 wigner_calc=self.wigner_calc,
                 build_wigner=self._build_full_wigner(),
+                packed_wigner=self._packed_wigner_consumers,
                 node_partial_exchange=node_partial_exchange,
             )
 
diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
index 29e56d316d..d14ab03cf8 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
@@ -339,6 +339,7 @@ def build_edge_cache(
     random_gamma: bool,
     wigner_calc: WignerCalculatorFn,
     build_wigner: bool = True,
+    packed_wigner: bool = False,
 ) -> EdgeFeatureCache:
     """
     Build the global edge cache from DeePMD padded neighbor list.
@@ -453,6 +454,7 @@ def build_edge_cache(
             random_gamma=random_gamma,
             wigner_calc=wigner_calc,
             build_full=build_wigner,
+            packed_wigner=packed_wigner,
         )  # (E, D, D), (E, D, D), (E, 4)
 
     edge_type_feat = build_edge_type_feat(type_ebed, src, dst)  # (E, C)
@@ -494,6 +496,7 @@ def build_edge_cache_from_edges(
     random_gamma: bool,
     wigner_calc: WignerCalculatorFn,
     build_wigner: bool = True,
+    packed_wigner: bool = False,
     node_partial_exchange: Callable[[torch.Tensor], torch.Tensor] | None = None,
     fused_radial: FusedRadialFn | None = None,
     fused_wigner: WignerCalculatorFn | None = None,
@@ -603,6 +606,7 @@ def build_edge_cache_from_edges(
             random_gamma=random_gamma,
             wigner_calc=fused_wigner if fused_wigner is not None else wigner_calc,
             build_full=build_wigner,
+            packed_wigner=packed_wigner,
         )  # (E, D, D), (E, D, D), (E, 4)
 
     # === Step 5. Edge type features ===
@@ -660,6 +664,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.
@@ -682,6 +687,12 @@ 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 every selected rotation consumer addresses the banded
+        layout. The banded storage is only chosen when this holds; any
+        dense-layout consumer (the cuTile value path, the fused CUDA
+        convolution, or the plain SO2 reference path) keeps the dense
+        surfaces even with ``DP_PACKED_D_INFER`` set.
 
     Returns
     -------
@@ -714,6 +725,7 @@ def _build_edge_wigner(
     calc = getattr(wigner_calc, "__self__", wigner_calc)
     if (
         use_packed_d_infer()
+        and packed_wigner
         and callable(getattr(calc, "forward_band", None))
         and calc.band_supported()
     ):

From b615b426d3b3be3c7a84956af641aa0b9e5de381 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Fri, 16 Oct 2026 20:03:44 +0800
Subject: [PATCH 11/15] test(pt): cover the packed-Wigner consumer gate matrix

---
 .../pt/model/test_descriptor_sezm_triton.py   | 49 +++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 2d343ad4c0..05d4d2ab9e 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2051,6 +2051,55 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self):
         self.assertTrue(tc.has_tile_config("flash_bwd_block", key))
 
 
+class TestSeZMPackedWignerConsumerGate(unittest.TestCase):
+    """The banded layout requires every rotation consumer to address it."""
+
+    def _conv(self, **kw):
+        from types import SimpleNamespace
+
+        base = {
+            "_cutile_value_path": None,
+            "_cuda_conv_fn": None,
+            "_cuda_value_train": None,
+            "_triton_value_path": None,
+            "_triton_rotate_mix": None,
+        }
+        base.update(kw)
+        return SimpleNamespace(**base)
+
+    def test_gate_matrix(self):
+        from deepmd.pt.model.descriptor.sezm import (
+            block_supports_packed_wigner,
+        )
+
+        self.assertTrue(
+            block_supports_packed_wigner(self._conv(_triton_value_path=object()))
+        )
+        self.assertTrue(
+            block_supports_packed_wigner(self._conv(_triton_rotate_mix=object()))
+        )
+        # cuTile, the fused CUDA convolution and the CUDA training path all
+        # address the dense layout and must block the banded storage
+        self.assertFalse(
+            block_supports_packed_wigner(
+                self._conv(_triton_value_path=object(), _cutile_value_path=object())
+            )
+        )
+        self.assertFalse(
+            block_supports_packed_wigner(
+                self._conv(_triton_value_path=object(), _cuda_conv_fn=object())
+            )
+        )
+        self.assertFalse(
+            block_supports_packed_wigner(
+                self._conv(_triton_value_path=object(), _cuda_value_train=object())
+            )
+        )
+        # no accelerated binding: the plain SO2 reference path reads the
+        # dense surface
+        self.assertFalse(block_supports_packed_wigner(self._conv()))
+
+
 if __name__ == "__main__":
     unittest.main()
 

From 8605f6c9df9a3514e6f670903cf80b1feee92d6e Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 16 Oct 2026 12:04:56 +0000
Subject: [PATCH 12/15] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/pt/model/test_descriptor_sezm_triton.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 05d4d2ab9e..79fdc1d931 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2055,7 +2055,9 @@ class TestSeZMPackedWignerConsumerGate(unittest.TestCase):
     """The banded layout requires every rotation consumer to address it."""
 
     def _conv(self, **kw):
-        from types import SimpleNamespace
+        from types import (
+            SimpleNamespace,
+        )
 
         base = {
             "_cutile_value_path": None,

From 3291de47ab9119557584eb5cc2c1c39528bafc44 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 18 Oct 2026 17:48:33 +0800
Subject: [PATCH 13/15] fix(pt): evaluate the packed-Wigner gate per cache
 build against current consumers

The gate is now a call-time check over the currently selected forward
and inverse consumers: without a bound flash aggregation the message is
rotated back through the dense inverse matrix, so no-attention and
projected-attention configurations keep the dense surfaces; training
mode reads the training level and its rotate-mix binding instead of a
construction-time boolean, so level combinations like inference 2 /
training 0 followed by .train() no longer pass a stale verdict.  The
gate matrix tests now cover both mode cases and the aggregation
routes.
---
 deepmd/pt/model/descriptor/sezm.py            | 70 ++++++++++-----
 .../pt/model/test_descriptor_sezm_triton.py   | 87 +++++++++++++++----
 2 files changed, 121 insertions(+), 36 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index 676a8c4630..b0ea412c71 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -67,6 +67,8 @@
 )
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
+    triton_infer_level,
+    triton_train_level,
     use_amp_infer,
 )
 from deepmd.utils.bridging import (
@@ -132,24 +134,41 @@
 
 
 @BaseDescriptor.register("SeZM")
-def block_supports_packed_wigner(so2_conv: Any) -> bool:
-    """Whether one block's selected rotation consumers address the band.
+def block_supports_packed_wigner(so2_conv: Any, training: bool) -> bool:
+    """Whether one block's *currently selected* rotation consumers address
+    the band.
 
     The banded layout serves the Triton value path and the Triton
-    rotate-mix path (whose kernels address the band directly), together
-    with the attention aggregation and the GIE coupling that read the same
-    tensor. Any dense-layout consumer -- the cuTile value path, the fused
-    CUDA convolution, the CUDA training value path, or the plain SO2
-    reference path (no Triton binding at all) -- keeps the dense surfaces.
+    rotate-mix path (whose kernels address the band directly) together
+    with the Triton flash aggregation and the GIE coupling that read the
+    same tensor.  Every dense-layout consumer keeps the dense surfaces:
+    the cuTile value path, the fused CUDA convolution, the CUDA training
+    value path, the plain SO2 reference path (no Triton binding selected
+    at the *current* level and mode), and the aggregation routes that
+    rotate messages back through ``Dt_full`` (no attention heads, or an
+    attention layout with value/output projections, where the flash
+    operator does not bind).  The check is evaluated per cache build, so
+    mode and level changes between builds are honored.
     """
-    return (
-        getattr(so2_conv, "_cutile_value_path", None) is None
-        and getattr(so2_conv, "_cuda_conv_fn", None) is None
-        and getattr(so2_conv, "_cuda_value_train", None) is None
-        and (
-            getattr(so2_conv, "_triton_value_path", None) is not None
-            or getattr(so2_conv, "_triton_rotate_mix", None) is not None
-        )
+    if getattr(so2_conv, "_cutile_value_path", None) is not None:
+        return False
+    if getattr(so2_conv, "_cuda_conv_fn", None) is not None:
+        return False
+    if getattr(so2_conv, "_cuda_value_train", None) is not None:
+        return False
+    # Inverse-rotation consumer: without a bound flash aggregation the
+    # message is rotated back through the dense inverse matrix.
+    if getattr(so2_conv, "_flash_atten_fn", None) is None:
+        return False
+    level = triton_train_level() if training else triton_infer_level()
+    if level < 1:
+        return False
+    # Forward-rotation consumer selected at the current level.
+    if training:
+        return getattr(so2_conv, "_triton_rotate_mix", None) is not None
+    return getattr(so2_conv, "_triton_value_path", None) is not None or (
+        getattr(so2_conv, "_triton_value_path", None) is None
+        and getattr(so2_conv, "_triton_rotate_mix", None) is not None
     )
 
 
@@ -1148,9 +1167,11 @@ def __init__(
             for block in self.blocks
         )
         # Banded storage additionally requires every selected rotation
-        # consumer to address the band (see ``block_supports_packed_wigner``).
-        self._packed_wigner_consumers = bool(self.blocks) and all(
-            block_supports_packed_wigner(block.so2_conv) for block in self.blocks
+        # consumer to address the band; the decision is re-evaluated at
+        # each cache build so train/eval and level changes are honored
+        # (see ``block_supports_packed_wigner``).
+        self._blocks_have_so2_conv = bool(self.blocks) and all(
+            hasattr(block, "so2_conv") for block in self.blocks
         )
 
         # The envelope and the radial basis are both functions of the pair
@@ -1453,7 +1474,7 @@ def forward(
                 random_gamma=self.random_gamma and self.training,
                 wigner_calc=self.wigner_calc,
                 build_wigner=self._build_full_wigner(),
-                packed_wigner=self._packed_wigner_consumers,
+                packed_wigner=self._packed_wigner_ok(),
             )
 
         ebed_dim_0 = self.node_init_dim  # (node_init_lmax+1)^2
@@ -1770,7 +1791,7 @@ def forward_with_edges(
                 random_gamma=self.random_gamma and self.training,
                 wigner_calc=self.wigner_calc,
                 build_wigner=self._build_full_wigner(),
-                packed_wigner=self._packed_wigner_consumers,
+                packed_wigner=self._packed_wigner_ok(),
                 node_partial_exchange=node_partial_exchange,
             )
 
@@ -2090,6 +2111,15 @@ def _edge_quaternion(self, edge_cache: EdgeFeatureCache) -> torch.Tensor:
             )
         return edge_quat
 
+    def _packed_wigner_ok(self) -> bool:
+        """Every block's currently selected consumers address the band."""
+        if not self._blocks_have_so2_conv:
+            return False
+        return all(
+            block_supports_packed_wigner(block.so2_conv, self.training)
+            for block in self.blocks
+        )
+
     def _build_full_wigner(self) -> bool:
         """Return whether the active execution path needs dense Wigner blocks."""
         if not self._need_full_wigner:
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 79fdc1d931..ac2d270971 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -23,6 +23,7 @@
 """
 
 import math
+import os
 import typing
 import unittest
 from unittest import (
@@ -2052,7 +2053,7 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self):
 
 
 class TestSeZMPackedWignerConsumerGate(unittest.TestCase):
-    """The banded layout requires every rotation consumer to address it."""
+    """The banded layout requires every *currently selected* consumer."""
 
     def _conv(self, **kw):
         from types import (
@@ -2065,41 +2066,95 @@ def _conv(self, **kw):
             "_cuda_value_train": None,
             "_triton_value_path": None,
             "_triton_rotate_mix": None,
+            "_flash_atten_fn": None,
         }
         base.update(kw)
         return SimpleNamespace(**base)
 
+    def _with_levels(self, infer, train):
+        old_i = os.environ.get("DP_TRITON_INFER")
+        old_t = os.environ.get("DP_TRITON_TRAIN")
+        os.environ["DP_TRITON_INFER"] = infer
+        os.environ["DP_TRITON_TRAIN"] = train
+
+        def restore():
+            for key, old in (("DP_TRITON_INFER", old_i), ("DP_TRITON_TRAIN", old_t)):
+                if old is None:
+                    os.environ.pop(key, None)
+                else:
+                    os.environ[key] = old
+
+        self.addCleanup(restore)
+
     def test_gate_matrix(self):
         from deepmd.pt.model.descriptor.sezm import (
             block_supports_packed_wigner,
         )
 
-        self.assertTrue(
-            block_supports_packed_wigner(self._conv(_triton_value_path=object()))
-        )
-        self.assertTrue(
-            block_supports_packed_wigner(self._conv(_triton_rotate_mix=object()))
+        self._with_levels("2", "0")
+        # Dense-layout value/aggregation consumers block the band.
+        self.assertFalse(
+            block_supports_packed_wigner(self._conv(_triton_value_path=object()), False)
         )
-        # cuTile, the fused CUDA convolution and the CUDA training path all
-        # address the dense layout and must block the banded storage
+        # The envelope aggregation (no attention heads) rotates back
+        # through Dt_full even with a packed-capable forward path.
         self.assertFalse(
             block_supports_packed_wigner(
-                self._conv(_triton_value_path=object(), _cutile_value_path=object())
+                self._conv(
+                    _triton_value_path=object(),
+                    _flash_atten_fn=None,
+                ),
+                False,
             )
         )
-        self.assertFalse(
+        # cuTile, the fused CUDA convolution and the CUDA training path all
+        # address the dense layout and must block the banded storage.
+        for dense in (
+            "_cutile_value_path",
+            "_cuda_conv_fn",
+            "_cuda_value_train",
+        ):
+            with self.subTest(dense=dense):
+                kw = {
+                    "_triton_value_path": object(),
+                    "_flash_atten_fn": object(),
+                    dense: object(),
+                }
+                self.assertFalse(block_supports_packed_wigner(self._conv(**kw), False))
+        # The plain SO2 reference path (no Triton binding at all) reads the
+        # dense surface.
+        self.assertFalse(block_supports_packed_wigner(self._conv(), False))
+        # A fully packed-capable selection passes in eval.
+        self.assertTrue(
             block_supports_packed_wigner(
-                self._conv(_triton_value_path=object(), _cuda_conv_fn=object())
+                self._conv(_triton_value_path=object(), _flash_atten_fn=object()),
+                False,
             )
         )
-        self.assertFalse(
+        self.assertTrue(
             block_supports_packed_wigner(
-                self._conv(_triton_value_path=object(), _cuda_value_train=object())
+                self._conv(_triton_rotate_mix=object(), _flash_atten_fn=object()),
+                False,
             )
         )
-        # no accelerated binding: the plain SO2 reference path reads the
-        # dense surface
-        self.assertFalse(block_supports_packed_wigner(self._conv()))
+
+    def test_gate_follows_mode_and_levels(self):
+        """The gate is evaluated against the current mode and levels."""
+        from deepmd.pt.model.descriptor.sezm import (
+            block_supports_packed_wigner,
+        )
+
+        packed_conv = self._conv(_triton_value_path=object(), _flash_atten_fn=object())
+        # Inference level 2, training level 0: eval passes, train rejects.
+        self._with_levels("2", "0")
+        self.assertTrue(block_supports_packed_wigner(packed_conv, False))
+        self.assertFalse(block_supports_packed_wigner(packed_conv, True))
+        # Inference level 0, training level 1: eval rejects; a
+        # training-bound rotate-mix block passes in train mode.
+        self._with_levels("0", "1")
+        self.assertFalse(block_supports_packed_wigner(packed_conv, False))
+        train_conv = self._conv(_triton_rotate_mix=object(), _flash_atten_fn=object())
+        self.assertTrue(block_supports_packed_wigner(train_conv, True))
 
 
 if __name__ == "__main__":

From 813b37c418a851c11427a8c4809ad56a258d8ae6 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 18 Oct 2026 18:16:26 +0800
Subject: [PATCH 14/15] fix(pt): derive the packed-Wigner gate from the
 per-block captured levels

Dispatch resolves Triton levels through the values captured on each
block at construction (active_triton_level), so the gate now reads
those same numbers instead of the ambient environment: a block built
inside an inference-2/training-0 scope can no longer pass the gate in
training mode after the ambient training level changes, and the
converse eval case is covered as well.  Training-mode aggregation
additionally requires the flash operator to have been bound for
training (_flash_atten_trains), matching what the training
dispatch actually uses.
---
 deepmd/pt/model/descriptor/sezm.py            | 30 ++++++-----
 .../pt/model/test_descriptor_sezm_triton.py   | 51 +++++++++++++++----
 2 files changed, 59 insertions(+), 22 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index b0ea412c71..65c3a1aaf9 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -67,8 +67,6 @@
 )
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
-    triton_infer_level,
-    triton_train_level,
     use_amp_infer,
 )
 from deepmd.utils.bridging import (
@@ -143,12 +141,12 @@ def block_supports_packed_wigner(so2_conv: Any, training: bool) -> bool:
     with the Triton flash aggregation and the GIE coupling that read the
     same tensor.  Every dense-layout consumer keeps the dense surfaces:
     the cuTile value path, the fused CUDA convolution, the CUDA training
-    value path, the plain SO2 reference path (no Triton binding selected
-    at the *current* level and mode), and the aggregation routes that
-    rotate messages back through ``Dt_full`` (no attention heads, or an
-    attention layout with value/output projections, where the flash
-    operator does not bind).  The check is evaluated per cache build, so
-    mode and level changes between builds are honored.
+    value path, the plain SO2 reference path, and the aggregation routes
+    that rotate messages back through ``Dt_full`` (no attention heads,
+    a projected-attention layout, or a training flash that did not bind
+    for training).  Levels derive from the per-block values captured at
+    construction -- the same numbers ``active_triton_level`` resolves
+    dispatch with -- and the train/eval mode is read per cache build.
     """
     if getattr(so2_conv, "_cutile_value_path", None) is not None:
         return False
@@ -157,13 +155,21 @@ def block_supports_packed_wigner(so2_conv: Any, training: bool) -> bool:
     if getattr(so2_conv, "_cuda_value_train", None) is not None:
         return False
     # Inverse-rotation consumer: without a bound flash aggregation the
-    # message is rotated back through the dense inverse matrix.
-    if getattr(so2_conv, "_flash_atten_fn", None) is None:
+    # message is rotated back through the dense inverse matrix.  In
+    # training the aggregation is packed-capable only when the flash
+    # operator was bound for training too.
+    flash = getattr(so2_conv, "_flash_atten_fn", None)
+    if flash is None:
         return False
-    level = triton_train_level() if training else triton_infer_level()
+    if training and not getattr(so2_conv, "_flash_atten_trains", False):
+        return False
+    # Dispatch resolves against the per-block levels captured at
+    # construction (``active_triton_level``), not the ambient environment;
+    # eligibility must derive from the same captured numbers.
+    level = so2_conv.triton_train_level if training else so2_conv.triton_infer_level
     if level < 1:
         return False
-    # Forward-rotation consumer selected at the current level.
+    # Forward-rotation consumer selected at the captured level.
     if training:
         return getattr(so2_conv, "_triton_rotate_mix", None) is not None
     return getattr(so2_conv, "_triton_value_path", None) is not None or (
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index ac2d270971..17ee7242ff 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2067,6 +2067,9 @@ def _conv(self, **kw):
             "_triton_value_path": None,
             "_triton_rotate_mix": None,
             "_flash_atten_fn": None,
+            "_flash_atten_trains": True,
+            "triton_infer_level": 2,
+            "triton_train_level": 0,
         }
         base.update(kw)
         return SimpleNamespace(**base)
@@ -2138,23 +2141,51 @@ def test_gate_matrix(self):
             )
         )
 
-    def test_gate_follows_mode_and_levels(self):
-        """The gate is evaluated against the current mode and levels."""
+    def test_gate_follows_mode_and_captured_levels(self):
+        """The gate derives from the per-block captured levels, like dispatch."""
         from deepmd.pt.model.descriptor.sezm import (
             block_supports_packed_wigner,
         )
 
-        packed_conv = self._conv(_triton_value_path=object(), _flash_atten_fn=object())
-        # Inference level 2, training level 0: eval passes, train rejects.
-        self._with_levels("2", "0")
+        # Captured inference 2 / training 0: eval passes, train rejects
+        # (the captured training level selects the dense path regardless
+        # of what the ambient environment now says).
+        packed_conv = self._conv(
+            _triton_value_path=object(),
+            _flash_atten_fn=object(),
+            _flash_atten_trains=True,
+            triton_infer_level=2,
+            triton_train_level=0,
+        )
+        self._with_levels("1", "1")  # ambient disagrees on purpose
         self.assertTrue(block_supports_packed_wigner(packed_conv, False))
         self.assertFalse(block_supports_packed_wigner(packed_conv, True))
-        # Inference level 0, training level 1: eval rejects; a
-        # training-bound rotate-mix block passes in train mode.
-        self._with_levels("0", "1")
-        self.assertFalse(block_supports_packed_wigner(packed_conv, False))
-        train_conv = self._conv(_triton_rotate_mix=object(), _flash_atten_fn=object())
+        # Captured inference 0 / training 1: eval rejects; a training-
+        # bound rotate-mix block whose flash binds for training passes.
+        train_conv = self._conv(
+            _triton_rotate_mix=object(),
+            _flash_atten_fn=object(),
+            _flash_atten_trains=True,
+            triton_infer_level=0,
+            triton_train_level=1,
+        )
+        self.assertFalse(block_supports_packed_wigner(train_conv, False))
         self.assertTrue(block_supports_packed_wigner(train_conv, True))
+        # Training flash bound for inference only: the training-mode
+        # aggregation rotates back through the dense inverse, so the band
+        # is rejected even with a training-bound rotate-mix forward.
+        self.assertFalse(
+            block_supports_packed_wigner(
+                self._conv(
+                    _triton_rotate_mix=object(),
+                    _flash_atten_fn=object(),
+                    _flash_atten_trains=False,
+                    triton_infer_level=2,
+                    triton_train_level=1,
+                ),
+                True,
+            )
+        )
 
 
 if __name__ == "__main__":

From d2616eb2ef0d3484aea8d08897a0e4955c5cdc80 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Wed, 21 Oct 2026 20:17:41 +0800
Subject: [PATCH 15/15] fix(pt): restore SeZM descriptor registration

The module-level packed-Wigner gate helper accidentally carried the
register("SeZM") decorator, so the descriptor registry mapped "SeZM"
to a plain function and factory construction raised TypeError. Move the
decorator back onto DescrptSeZM and cover construction through the
registry for all four type keys.
---
 deepmd/pt/model/descriptor/sezm.py            |  2 +-
 source/tests/pt/model/test_descriptor_sezm.py | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index 65c3a1aaf9..b197444dff 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -131,7 +131,6 @@
     )
 
 
-@BaseDescriptor.register("SeZM")
 def block_supports_packed_wigner(so2_conv: Any, training: bool) -> bool:
     """Whether one block's *currently selected* rotation consumers address
     the band.
@@ -178,6 +177,7 @@ def block_supports_packed_wigner(so2_conv: Any, training: bool) -> bool:
     )
 
 
+@BaseDescriptor.register("SeZM")
 @BaseDescriptor.register("sezm")
 @BaseDescriptor.register("DPA4")
 @BaseDescriptor.register("dpa4")
diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py
index 5bf465da0f..62a85ea8d2 100644
--- a/source/tests/pt/model/test_descriptor_sezm.py
+++ b/source/tests/pt/model/test_descriptor_sezm.py
@@ -699,6 +699,24 @@ def test_forward_with_descriptor_variants(self) -> None:
             with self.subTest(mode=name):
                 self._assert_forward_backward_smoke(**model_kwargs)
 
+    def test_registry_construction(self) -> None:
+        """Every type key must resolve to the descriptor class in the registry.
+
+        A module-level helper once ended up decorated with
+        ``register("SeZM")``, which silently replaced the class with a
+        plain function in the registry and broke factory construction.
+        """
+        from deepmd.pt.model.descriptor.base_descriptor import (
+            BaseDescriptor,
+        )
+
+        for type_name in ("SeZM", "sezm", "DPA4", "dpa4"):
+            with self.subTest(type=type_name):
+                cls = BaseDescriptor.get_class_by_type(type_name)
+                self.assertIs(cls, DescrptSeZM)
+                model = cls(**_descriptor_kwargs())
+                self.assertIsInstance(model, DescrptSeZM)
+
     def test_fixed_basis_types_freeze_the_basis(self) -> None:
         """``/fix`` basis types keep the basis parameters out of training only."""
         for family in ("bessel", "gaussian"):