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/16] 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/16] 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/16] [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/16] 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/16] 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/16] 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/16] [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/16] 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/16] 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/16] 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/16] 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/16] [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/16] 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/16] 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/16] 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"):

From 0eb7874e367b80a17ee71b917eb752b0e1d20b63 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Thu, 22 Oct 2026 20:41:43 +0800
Subject: [PATCH 16/16] fix(pt): keep the packed Wigner layout inference-only

Training builds now always keep the dense Wigner surfaces: the
registered parameter-gradient consumers (the channel-basis backward of
the rotate-mix operator among them) address the dense surface, so a
training cache must never carry the band even when the bound kernels
could read it. The gate returns False for training before any consumer
check, and the gate matrix test asserts training rejection for
packed-capable blocks as well.

End-to-end coverage with DP_PACKED_D_INFER actually set drives the whole
chain (edge-cache build, per-block gate, consumers): the flag-on eval
cache carries the (E, 3D-2) band and never materializes the dense
surfaces, the same model in training builds the dense layout, and the
flag-off outputs are bitwise-equal. The unsupported branches are covered
as well: forward_band_fused falls back to the plain composition when the
op gate rejects (training), and forward_band raises for lmax >=
poly_lmin.

Docstrings: wigner_band_op gains full parameter documentation, the
packed_wigner contract lists the full exclusion set, _build_edge_wigner
documents the fourth return value (D_packed), the m = -1 slot range is
corrected to [D, 2D - 1), and the flag description no longer claims an
on-demand unpack that does not exist.
---
 deepmd/pt/model/descriptor/sezm.py            |  35 +++---
 .../pt/model/descriptor/sezm_nn/edge_cache.py |  25 ++--
 deepmd/pt/model/descriptor/sezm_nn/wignerd.py |  34 +++++-
 deepmd/pt_expt/kernels/utils.py               |   9 +-
 source/tests/pt/model/test_descriptor_sezm.py | 107 ++++++++++++++++++
 .../pt/model/test_descriptor_sezm_triton.py   |  42 +++++--
 6 files changed, 213 insertions(+), 39 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py
index b197444dff..345a8f4fdc 100644
--- a/deepmd/pt/model/descriptor/sezm.py
+++ b/deepmd/pt/model/descriptor/sezm.py
@@ -133,20 +133,25 @@
 
 def block_supports_packed_wigner(so2_conv: Any, training: bool) -> bool:
     """Whether one block's *currently selected* rotation consumers address
-    the band.
+    the band (inference only).
 
     The banded layout serves the Triton value path and the Triton
     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, 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.
+    the cuTile value path, the fused CUDA convolution, the plain SO2
+    reference path, and the aggregation routes that rotate messages back
+    through ``Dt_full`` (no attention heads or a projected-attention
+    layout).  Training always keeps the dense layout: the registered
+    parameter-gradient consumers (the channel-basis backward of the
+    rotate-mix operator among them) address the dense surface, so a
+    training cache must never carry the band even when the bound kernels
+    could read it.  The level derives from the per-block value captured
+    at construction -- the same number ``active_triton_level`` resolves
+    dispatch with -- and the gate is evaluated per cache build.
     """
+    if training:
+        return False
     if getattr(so2_conv, "_cutile_value_path", None) is not None:
         return False
     if getattr(so2_conv, "_cuda_conv_fn", None) is not None:
@@ -154,23 +159,15 @@ 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.  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
-    if training and not getattr(so2_conv, "_flash_atten_trains", False):
+    # message is rotated back through the dense inverse matrix.
+    if getattr(so2_conv, "_flash_atten_fn", None) is None:
         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:
+    if so2_conv.triton_infer_level < 1:
         return False
     # 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 (
         getattr(so2_conv, "_triton_value_path", None) is None
         and getattr(so2_conv, "_triton_rotate_mix", None) is not None
diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
index d14ab03cf8..93e131ab08 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
@@ -665,7 +665,12 @@ def _build_edge_wigner(
     wigner_calc: WignerCalculatorFn,
     build_full: bool = True,
     packed_wigner: bool = False,
-) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]:
+) -> tuple[
+    torch.Tensor | None,
+    torch.Tensor | None,
+    torch.Tensor,
+    torch.Tensor | None,
+]:
     """
     Build packed Wigner-D blocks from edge vectors.
 
@@ -690,16 +695,22 @@ def _build_edge_wigner(
     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.
+        dense-layout consumer keeps the dense surfaces even with
+        ``DP_PACKED_D_INFER`` set.  See
+        :func:`deepmd.pt.model.descriptor.sezm.block_supports_packed_wigner`
+        for the full exclusion set (the cuTile value path, the fused CUDA
+        convolution, the CUDA training value path, the plain SO2 reference
+        path, and the aggregation routes that rotate messages back through
+        ``Dt_full``; training is dense by construction).
 
     Returns
     -------
-    tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]
+    tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor, torch.Tensor | None]
         Packed Wigner-D matrices ``(D_full, Dt_full)`` with shape ``(E, D, D)``
-        (or ``None`` when ``build_full`` is False) and the quaternion used to
-        build them with shape ``(E, 4)``.
+        (or ``None`` when ``build_full`` is False), the quaternion used to
+        build them with shape ``(E, 4)``, and the banded storage ``D_packed``
+        with shape ``(E, 3 * D - 2)`` when the banded layout was selected
+        (``None`` otherwise, including every training build).
     """
     # === Step 1. Build edge-aligned quaternions ===
     edge_quat = build_edge_quaternion(
diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
index 515035f38e..142f5c18c9 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py
@@ -670,7 +670,7 @@ def forward_band(self, edge_quaternion: torch.Tensor) -> torch.Tensor:
         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,
+        - slots ``[D, 2D - 1)``: ``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``).
@@ -2259,5 +2259,35 @@ def wigner_band_op(
     exp_l3: list[int],
     lmax: int,
 ) -> torch.Tensor:
-    """Public entry for the op-bound banded Wigner assembly."""
+    """Public entry for the op-bound banded Wigner band assembly.
+
+    Same values as :meth:`WignerDCalculator.forward_band` behind the
+    ``sezm_triton::wigner_band`` operator boundary, so the compiled graph
+    saves the ``(E, 3 * dim_full - 2)`` band for the backward instead of
+    replaying the quaternion -> monomial -> product chain.
+
+    Parameters
+    ----------
+    quat
+        Unit quaternions with shape ``(E, 4)``.
+    l1_perm
+        Degree-1 row permutation index with shape ``(3,)``.
+    l1_sign_outer
+        Sign outer product for the degree-1 block with shape ``(3, 3)``.
+    coeff_l2
+        Degree-2 monomial coefficient table.
+    c_l3
+        Degree-3 case-coefficient table (``CaseCoefficients.C_l3``).
+    exp_l2
+        Flattened degree-2 monomial exponents.
+    exp_l3
+        Flattened degree-3 monomial exponents.
+    lmax
+        Maximum spherical harmonic degree; 1..3.
+
+    Returns
+    -------
+    torch.Tensor
+        Band tensor with shape ``(E, 3 * (lmax + 1)**2 - 2)``.
+    """
     return _wigner_band_op(quat, l1_perm, l1_sign, coeff_l2, c_l3, exp_l2, exp_l3, lmax)
diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py
index ca0b796bad..156deddbfc 100644
--- a/deepmd/pt_expt/kernels/utils.py
+++ b/deepmd/pt_expt/kernels/utils.py
@@ -390,9 +390,12 @@ def use_packed_d_infer() -> bool:
     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.
+    (rotate-mix, flash aggregation) address the band directly. There is no
+    unpack: the dense surfaces simply never exist in that mode, the Triton
+    launchers reject a banded input on the dense-addressing paths, and
+    reference (non-Triton) execution must not enable the flag. The stored
+    values are the same floats the dense layout keeps at those positions, so
+    consumer outputs are bit-identical to the dense mode.
 
     Returns
     -------
diff --git a/source/tests/pt/model/test_descriptor_sezm.py b/source/tests/pt/model/test_descriptor_sezm.py
index 62a85ea8d2..e59861aaf5 100644
--- a/source/tests/pt/model/test_descriptor_sezm.py
+++ b/source/tests/pt/model/test_descriptor_sezm.py
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import importlib.util
 import itertools
 import math
 import os
@@ -3134,5 +3135,111 @@ def test_sfpg_leaks_reopen_when_gate_is_disabled(self) -> None:
         )
 
 
+@unittest.skipUnless(
+    torch.cuda.is_available() and importlib.util.find_spec("triton") is not None,
+    "the packed Wigner layout engages on CUDA with Triton",
+)
+class TestSeZMPackedWignerEndToEnd(_SeZMTestCase):
+    """End-to-end coverage with ``DP_PACKED_D_INFER`` actually set.
+
+    Drives the whole chain with the flag on: the edge-cache build, the
+    per-block gate in ``sezm.py``, and the consumers in ``so2.py`` /
+    ``embedding.py`` all run against the banded storage, and the outputs
+    must match the same model built through the dense layout.
+    """
+
+    def _kwargs(self):
+        return {
+            "rcut": 3.0,
+            "sel": [1, 1],
+            "ntypes": 2,
+            "l_schedule": [1, 1],
+            "mmax": 1,
+            "channels": 32,
+            "focus_dim": 32,
+            "n_focus": 2,
+            "n_radial": 3,
+            "radial_mlp": [6],
+            "ffn_neurons": 8,
+            "ffn_blocks": 1,
+            "random_gamma": False,
+            "n_atten_head": 1,
+            "mlp_bias": False,
+            "precision": "float32",
+            "trainable": True,
+            "radial_so2_mode": "degree_channel",
+            "radial_so2_rank": 2,
+        }
+
+    def test_flag_on_builds_band_and_matches_dense(self):
+        import deepmd.pt.model.descriptor.sezm_nn.edge_cache as edge_cache_mod
+
+        saved_flag = os.environ.get("DP_PACKED_D_INFER")
+        saved_level = os.environ.get("DP_TRITON_INFER")
+        os.environ["DP_PACKED_D_INFER"] = "1"
+        os.environ["DP_TRITON_INFER"] = "2"
+
+        def restore():
+            for key, old in (
+                ("DP_PACKED_D_INFER", saved_flag),
+                ("DP_TRITON_INFER", saved_level),
+            ):
+                if old is None:
+                    os.environ.pop(key, None)
+                else:
+                    os.environ[key] = old
+
+        self.addCleanup(restore)
+
+        coord = torch.tensor(
+            [[[0.0, 0.0, 0.0], [0.0, 0.0, 1.5]]], dtype=torch.float32, device="cuda"
+        )
+        atype = torch.tensor([[0, 1]], dtype=torch.long, device="cuda")
+        nlist = torch.tensor([[[0, 1], [1, -1]]], dtype=torch.long, device="cuda")
+
+        records = []
+        real_build = edge_cache_mod._build_edge_wigner
+
+        def spy(*args, **kwargs):
+            built = real_build(*args, **kwargs)
+            records.append(tuple(None if t is None else tuple(t.shape) for t in built))
+            return built
+
+        model = DescrptSeZM(**self._kwargs()).cuda()
+        model.eval()
+
+        with mock.patch.object(edge_cache_mod, "_build_edge_wigner", spy):
+            with torch.no_grad():
+                packed_out = model(coord, atype, nlist, mapping=None, comm_dict=None)[0]
+            # The flag-on cache carries the band and never materializes the
+            # dense surfaces: D = (E, D_dim, D_dim) with D_dim = (lmax+1)^2,
+            # band = (E, 3 * D_dim - 2).
+            d_dense, dt_dense, _quat, d_packed = records[-1]
+            self.assertIsNone(d_dense)
+            self.assertIsNone(dt_dense)
+            self.assertIsNotNone(d_packed)
+            self.assertEqual(d_packed, (3, 10))
+
+            # Training keeps the dense layout even with the flag on: the
+            # parameter-gradient consumers address the dense surface.
+            model.train()
+            model(coord, atype, nlist, mapping=None, comm_dict=None)
+            d_dense_t, dt_dense_t, _quat_t, d_packed_t = records[-1]
+            self.assertIsNone(d_packed_t)
+            self.assertIsNotNone(d_dense_t)
+            self.assertIsNotNone(dt_dense_t)
+            model.eval()
+
+            # Same model, flag off: the dense layout through the same
+            # consumers, bitwise-equal outputs.
+            os.environ.pop("DP_PACKED_D_INFER", None)
+            with torch.no_grad():
+                dense_out = model(coord, atype, nlist, mapping=None, comm_dict=None)[0]
+            d_dense_o, _dt, _quat_o, d_packed_o = records[-1]
+            self.assertIsNone(d_packed_o)
+            self.assertIsNotNone(d_dense_o)
+        self.assertTrue(torch.equal(packed_out, dense_out))
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 17ee7242ff..6a3c10a81d 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2160,8 +2160,11 @@ def test_gate_follows_mode_and_captured_levels(self):
         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))
-        # Captured inference 0 / training 1: eval rejects; a training-
-        # bound rotate-mix block whose flash binds for training passes.
+        # Captured inference 0 / training 1: eval rejects; training never
+        # carries the band -- the parameter-gradient consumers (the
+        # channel-basis backward) address the dense surface -- so even a
+        # training-bound rotate-mix block whose flash binds for training
+        # keeps the dense layout in training mode.
         train_conv = self._conv(
             _triton_rotate_mix=object(),
             _flash_atten_fn=object(),
@@ -2170,16 +2173,13 @@ def test_gate_follows_mode_and_captured_levels(self):
             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(train_conv, True))
         self.assertFalse(
             block_supports_packed_wigner(
                 self._conv(
-                    _triton_rotate_mix=object(),
+                    _triton_value_path=object(),
                     _flash_atten_fn=object(),
-                    _flash_atten_trains=False,
+                    _flash_atten_trains=True,
                     triton_infer_level=2,
                     triton_train_level=1,
                 ),
@@ -2378,6 +2378,32 @@ def _op_band(self, calc, quat):
             int(calc.lmax),
         )
 
+    @_GPU_KERNELS
+    def test_fused_falls_back_to_composition_when_op_unsupported(self):
+        """Training keeps the plain composition (the op is eval-only)."""
+        import torch as torch_mod
+
+        calc = self._make_calc(3)
+        calc.train()
+        self.assertFalse(calc.band_op_supported())
+        quat = torch_mod.randn(16, 4, device="cuda")
+        quat = quat / quat.norm(dim=-1, keepdim=True)
+        fused = calc.forward_band_fused(quat)
+        plain = calc.forward_band(quat)
+        self.assertTrue(torch.equal(fused, plain))
+
+    def test_band_rejects_lmax_at_poly_lmin(self):
+        """`lmax >= poly_lmin` has no per-degree ladder below the tail."""
+        from deepmd.pt.model.descriptor.sezm_nn.wignerd import (
+            WignerDCalculator,
+        )
+
+        calc = WignerDCalculator(lmax=11, dtype=torch.float32).cpu()
+        self.assertFalse(calc.band_supported())
+        quat = torch.randn(4, 4, dtype=torch.float32, device="cpu")
+        with self.assertRaisesRegex(RuntimeError, "lmax < 11"):
+            calc.forward_band(quat)
+
     @_GPU_KERNELS
     def test_matches_composition_with_grads(self):
         import os