perf(pt): banded Wigner-D storage for the fused inference consumers

deepmd/pt/model/descriptor/sezm.py

@@ -123,6 +123,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")
@@ -1086,6 +1107,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
@@ -1351,6 +1377,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
@@ -1652,6 +1679,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,
             )
 
@@ -2036,6 +2064,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

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

@@ -25,6 +25,10 @@
     rearrange,
 )
 
+from deepmd.pt_expt.kernels.utils import (
+    use_packed_d_infer,
+)
+
 from .utils import (
     get_promoted_dtype,
     nvtx_range,
@@ -80,6 +84,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.
@@ -118,6 +128,7 @@ class EdgeFeatureCache(NamedTuple):
     csr_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None
     edge_src_gate: torch.Tensor | None = None
     edge_quat: torch.Tensor | None = None
+    D_packed: torch.Tensor | None = None
 
 
 def cached_edge_csr(
@@ -299,6 +310,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.
@@ -406,13 +418,14 @@ 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,
             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)
@@ -428,6 +441,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,
     )
 
@@ -452,6 +466,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,
@@ -547,13 +562,14 @@ 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,
             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 ===
@@ -588,6 +604,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,
         edge_src_gate=edge_src_gate,
     )
@@ -601,6 +618,7 @@ def _build_edge_wigner(
     random_gamma: bool,
     wigner_calc: WignerCalculatorFn,
     build_full: bool = True,
+    packed_wigner: bool = False,
 ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]:
     """
     Build packed Wigner-D blocks from edge vectors.
@@ -623,6 +641,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
     -------
@@ -649,9 +673,24 @@ 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 packed_wigner
+        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_fused(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(
@@ -668,6 +707,7 @@ def _finalize_edge_cache(
     edge_quat: torch.Tensor,
     deg_norm_floor: float,
     edge_src_gate: torch.Tensor | None = None,
+    D_packed: torch.Tensor | None = None,
 ) -> EdgeFeatureCache:
     """
     Assemble the shared `EdgeFeatureCache` layout.
@@ -735,6 +775,7 @@ def _finalize_edge_cache(
         csr_cache={},
         edge_src_gate=edge_src_gate,
         edge_quat=edge_quat,
+        D_packed=D_packed,
     )
 
 
@@ -892,10 +933,12 @@ def edge_cache_to_dtype(
     _Dt_full = cache.Dt_full
     _edge_src_gate = cache.edge_src_gate
     _edge_quat = cache.edge_quat
+    _D_packed = cache.D_packed
     D_full: torch.Tensor | None = None
     Dt_full: torch.Tensor | None = None
     edge_src_gate: 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:
@@ -904,6 +947,8 @@ def edge_cache_to_dtype(
         edge_src_gate = _edge_src_gate.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.
@@ -923,4 +968,5 @@ def edge_cache_to_dtype(
         csr_cache=None if cache.csr_cache is None else dict(cache.csr_cache),
         edge_src_gate=edge_src_gate,
         edge_quat=edge_quat,
+        D_packed=D_packed,
     )

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.

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

@@ -2027,7 +2027,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(