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

deepmd/pt/model/descriptor/sezm.py

@@ -132,6 +132,52 @@
 
 
 @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.
+
+    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.
+    """
+    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.  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):
+        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 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
+    )
+
+
 @BaseDescriptor.register("sezm")
 @BaseDescriptor.register("DPA4")
 @BaseDescriptor.register("dpa4")
@@ -1126,6 +1172,13 @@ 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; 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
         # distance and are cheap enough that the compiler inlines them into
@@ -1427,6 +1480,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_ok(),
             )
 
         ebed_dim_0 = self.node_init_dim  # (node_init_lmax+1)^2
@@ -1743,6 +1797,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_ok(),
                 node_partial_exchange=node_partial_exchange,
             )
 
@@ -2062,6 +2117,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:
@@ -2134,6 +2198,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,
@@ -93,6 +97,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 +130,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(
@@ -328,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.
@@ -435,13 +447,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)
@@ -457,6 +470,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,
     )
 
@@ -482,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,
@@ -584,13 +599,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 ===
@@ -634,6 +650,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,
     )
@@ -647,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.
@@ -669,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
     -------
@@ -695,9 +719,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(
@@ -714,6 +753,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 +821,7 @@ def _finalize_edge_cache(
         csr_cache={},
         edge_quat=edge_quat,
         node_gate=node_gate,
+        D_packed=D_packed,
     )
 
 
@@ -938,15 +979,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 +1011,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,
     )

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

@@ -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(