fixup

Author: OutisLi

deepmd/pt/infer/deep_eval.py

@@ -325,6 +325,12 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None:
                         "vesin[torch]`) or use nlist_backend='native' (or 'auto')."
                     )
                 builder = VesinNeighborList()
+            elif DEVICE.type != "cuda":
+                raise ValueError(
+                    "nlist_backend='nv' requires CUDA inference tensors; "
+                    f"current DEVICE is {DEVICE!s}. Use nlist_backend='native' "
+                    "(or 'auto') for CPU inference."
+                )
             elif not is_nv_available():
                 raise ImportError(
                     "nlist_backend='nv' was requested but 'nvalchemi-toolkit-ops'"
@@ -338,7 +344,7 @@ def _setup_nlist_backend(self, nlist_backend: str) -> None:
             # Pick the first available O(N) builder; nv is GPU-only.
             if is_vesin_torch_available():
                 builder = VesinNeighborList()
-            elif is_nv_available() and torch.cuda.is_available():
+            elif is_nv_available() and DEVICE.type == "cuda":
                 builder = NvNeighborList()
         self._nlist_builder = builder
 

deepmd/pt/model/descriptor/sezm.py

@@ -619,7 +619,7 @@ def __init__(
         self.kmax = int(kmax)
         if self.kmax < 0:
             raise ValueError("`kmax` must be non-negative")
-        if self.kmax > int(lmax):
+        if self.kmax > self.lmax:
             raise ValueError("`kmax` must be <= `lmax`")
         self.ebed_dims = [get_so3_dim_of_lmax(l) for l in self.l_schedule]
         self._init_node_l_schedules(extra_node_l)
@@ -844,8 +844,14 @@ def __init__(
 
         blocks: list[SeZMInteractionBlock] = []
         for block_idx, (l_b, node_l_b, m_b) in enumerate(
-            zip(self.l_schedule, self.node_l_schedule, self.m_schedule)
+            zip(
+                self.l_schedule,
+                self.node_l_schedule,
+                self.m_schedule,
+                strict=True,
+            )
         ):
+            k_b = min(self.kmax, l_b)
             blocks.append(
                 SeZMInteractionBlock(
                     lmax=l_b,
@@ -877,7 +883,7 @@ def __init__(
                     message_node_so3=self.message_node_so3,
                     ffn_s2_activation=self.ffn_s2_activation,
                     ffn_so3_grid=self.ffn_so3_grid,
-                    kmax=self.kmax,
+                    kmax=k_b,
                     so2_lebedev_quadrature=self.so2_lebedev_quadrature,
                     ffn_lebedev_quadrature=self.ffn_lebedev_quadrature,
                     n_atten_head=self.n_atten_head,
@@ -1511,12 +1517,14 @@ def _build_gie_zonal_coupling(
             mp_row_index,
             mp_m0_col_index,
         ]
-        edge_len = safe_norm(edge_cache.edge_vec, self.eps)
-        edge_quat = build_edge_quaternion(
-            edge_cache.edge_vec,
-            edge_len=edge_len,
-            eps=self.eps,
-        )
+        edge_quat = edge_cache.edge_quat
+        if edge_quat is None:
+            edge_len = safe_norm(edge_cache.edge_vec, self.eps)
+            edge_quat = build_edge_quaternion(
+                edge_cache.edge_vec,
+                edge_len=edge_len,
+                eps=self.eps,
+            )
         extra_coupling = self.gie_zonal_wigner_calc.forward_zonal(
             edge_quat,
             lmin=self.lmax + 1,
@@ -1676,7 +1684,7 @@ def _init_lm_schedules(
             raise ValueError("`m_schedule` must have the same length as `l_schedule`")
         if any(x < 0 for x in self.m_schedule):
             raise ValueError("`m_schedule` entries must be non-negative")
-        if any(m > l for m, l in zip(self.m_schedule, self.l_schedule)):
+        if any(m > l for m, l in zip(self.m_schedule, self.l_schedule, strict=True)):
             raise ValueError(
                 "`m_schedule` entries must satisfy `m_schedule[i] <= l_schedule[i]`"
             )

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

@@ -638,7 +638,8 @@ def _use_infer_activation_checkpoint(self, *tensors: torch.Tensor) -> bool:
         """
         return (
             not self.training
-            and os.environ.get("DP_ACT_INFER") == "1"
+            and os.environ.get("DP_ACT_INFER", "").strip().lower()
+            in {"1", "true", "yes", "on"}
             and os.environ.get("DP_COMPILE_INFER", "").strip().lower()
             not in {"1", "true", "yes", "on"}
             and torch.is_grad_enabled()

deepmd/pt/model/descriptor/sezm_nn/cute/so2_rotation.py

@@ -830,10 +830,12 @@ def _rb_backward(ctx: Any, grad_out: Tensor) -> tuple:
 
 
 # === Public API ==============================================================
-def _cute_usable(*tensors: Tensor) -> bool:
+def _cute_usable(channels: int, *tensors: Tensor) -> bool:
     """Return True when the CuTe fast path is available for these tensors."""
     if not SEZM_CUTE_AVAILABLE:
         return False
+    if int(channels) < _TN or int(channels) % _TN != 0:
+        return False
     return all(
         t.is_cuda and t.dtype == torch.float32 for t in tensors if t.is_floating_point()
     )
@@ -872,7 +874,7 @@ def rotate_to_local_cute(
     Experimental path that is not used in production. See the module docstring
     for the benchmark conclusion and why the Triton kernels were chosen instead.
     """
-    if _cute_usable(x, wigner) and src.numel() > 0:
+    if _cute_usable(x.shape[2], x, wigner) and src.numel() > 0:
         return torch.ops.sezm_cute.rotate_to_local(
             x, src, wigner, coeff_index, int(dim_full)
         )
@@ -909,7 +911,7 @@ def rotate_back_cute(
     Experimental path that is not used in production. See the module docstring
     for the benchmark conclusion and why the Triton kernels were chosen instead.
     """
-    if _cute_usable(x_local, wigner) and x_local.shape[0] > 0:
+    if _cute_usable(x_local.shape[2], x_local, wigner) and x_local.shape[0] > 0:
         return torch.ops.sezm_cute.rotate_back(
             x_local, wigner, coeff_index, int(dim_full)
         )

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

@@ -71,6 +71,9 @@ class EdgeFeatureCache(NamedTuple):
         Used for efficient batched rotation. None if not available.
     Dt_full
         Transpose of D_full with shape (E, D, D). None if not available.
+    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_to_m_cache
         Lazy cache for projected D matrices keyed by a normalized
         ``"lmax:mmax"`` identifier.
@@ -103,6 +106,7 @@ class EdgeFeatureCache(NamedTuple):
     D_to_m_cache: dict[str, torch.Tensor] | None = None
     Dt_from_m_cache: dict[str, torch.Tensor] | None = None
     edge_src_gate: torch.Tensor | None = None
+    edge_quat: torch.Tensor | None = None
 
 
 def compute_edge_src_gate(
@@ -337,13 +341,13 @@ def build_edge_cache(
 
     # === Step 6. Edge quaternion -> Wigner-D blocks ===
     with nvtx_range("wigner_d"):
-        D_full, Dt_full = _build_edge_wigner(
+        D_full, Dt_full, edge_quat = _build_edge_wigner(
             edge_vec=edge_vec,
             edge_len=edge_len,
             eps=eps,
             random_gamma=random_gamma,
             wigner_calc=wigner_calc,
-        )  # (E, D, D), (E, D, D)
+        )  # (E, D, D), (E, D, D), (E, 4)
 
     edge_type_feat = build_edge_type_feat(type_ebed, src, dst)  # (E, C)
 
@@ -357,6 +361,7 @@ def build_edge_cache(
         edge_env=edge_env,
         D_full=D_full,
         Dt_full=Dt_full,
+        edge_quat=edge_quat,
         deg_norm_floor=deg_norm_floor,
     )
 
@@ -460,13 +465,13 @@ def build_edge_cache_from_edges(
 
     # === Step 4. Edge quaternion -> Wigner-D blocks ===
     with nvtx_range("wigner_d"):
-        D_full, Dt_full = _build_edge_wigner(
+        D_full, Dt_full, edge_quat = _build_edge_wigner(
             edge_vec=edge_vec,
             edge_len=edge_len,
             eps=eps,
             random_gamma=random_gamma,
             wigner_calc=wigner_calc,
-        )  # (E, D, D), (E, D, D)
+        )  # (E, D, D), (E, D, D), (E, 4)
 
     # === Step 5. Edge type features ===
     edge_type_feat = build_edge_type_feat(type_ebed, src, dst)
@@ -498,6 +503,7 @@ def build_edge_cache_from_edges(
         edge_env=edge_env,
         D_full=D_full,
         Dt_full=Dt_full,
+        edge_quat=edge_quat,
         deg_norm_floor=deg_norm_floor,
         edge_src_gate=edge_src_gate,
     )
@@ -510,7 +516,7 @@ def _build_edge_wigner(
     eps: float,
     random_gamma: bool,
     wigner_calc: WignerCalculatorFn,
-) -> tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     Build packed Wigner-D blocks from edge vectors.
 
@@ -530,8 +536,9 @@ def _build_edge_wigner(
 
     Returns
     -------
-    tuple[torch.Tensor, torch.Tensor]
-        Packed Wigner-D matrices ``(D_full, Dt_full)`` with shape ``(E, D, D)``.
+    tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        Packed Wigner-D matrices ``(D_full, Dt_full)`` with shape ``(E, D, D)``
+        and the quaternion used to build them with shape ``(E, 4)``.
     """
     # === Step 1. Build edge-aligned quaternions ===
     edge_quat = build_edge_quaternion(
@@ -550,7 +557,8 @@ def _build_edge_wigner(
         edge_quat = quaternion_multiply(quaternion_z_rotation(gamma), edge_quat)
 
     # === Step 3. Convert quaternions to packed Wigner-D blocks ===
-    return wigner_calc(edge_quat)
+    D_full, Dt_full = wigner_calc(edge_quat)
+    return D_full, Dt_full, edge_quat
 
 
 def _finalize_edge_cache(
@@ -564,6 +572,7 @@ def _finalize_edge_cache(
     edge_env: torch.Tensor,
     D_full: torch.Tensor,
     Dt_full: torch.Tensor,
+    edge_quat: torch.Tensor,
     deg_norm_floor: float,
     edge_src_gate: torch.Tensor | None = None,
 ) -> EdgeFeatureCache:
@@ -590,6 +599,9 @@ def _finalize_edge_cache(
         Packed Wigner-D matrices with shape (E, D, D).
     Dt_full
         Transposed packed Wigner-D matrices with shape (E, D, D).
+    edge_quat
+        Global-to-local quaternions used to build the Wigner-D matrices with
+        shape (E, 4).
     deg_norm_floor
         Floor added to the envelope-squared degree before the inverse-sqrt
         normalization. A tiny ``eps`` reproduces the legacy behavior; an
@@ -627,6 +639,7 @@ def _finalize_edge_cache(
         D_to_m_cache={},
         Dt_from_m_cache={},
         edge_src_gate=edge_src_gate,
+        edge_quat=edge_quat,
     )
 
 
@@ -661,6 +674,7 @@ def _get_empty_edge_cache(
     """
     empty_long = torch.empty(0, dtype=torch.long, device=device)
     empty_vec = torch.empty(0, 3, dtype=dtype, device=device)
+    empty_quat = torch.empty(0, 4, dtype=dtype, device=device)
     empty_rbf = torch.empty(0, n_radial, dtype=dtype, device=device)
     empty_type_feat = torch.empty(0, n_channel, dtype=dtype, device=device)
     deg = torch.zeros(n_nodes, dtype=dtype, device=device)
@@ -679,6 +693,7 @@ def _get_empty_edge_cache(
         D_to_m_cache={},
         Dt_from_m_cache={},
         edge_src_gate=None,
+        edge_quat=empty_quat,
     )
 
 
@@ -835,15 +850,19 @@ def edge_cache_to_dtype(
     _D_full = cache.D_full
     _Dt_full = cache.Dt_full
     _edge_src_gate = cache.edge_src_gate
+    _edge_quat = cache.edge_quat
     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
     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_src_gate is not None:
         edge_src_gate = _edge_src_gate.to(dtype=dtype)
+    if _edge_quat is not None:
+        edge_quat = _edge_quat.to(dtype=dtype)
 
     return EdgeFeatureCache(
         src=cache.src,
@@ -859,4 +878,5 @@ def edge_cache_to_dtype(
         D_to_m_cache=None if cache.D_to_m_cache is None else {},
         Dt_from_m_cache=None if cache.Dt_from_m_cache is None else {},
         edge_src_gate=edge_src_gate,
+        edge_quat=edge_quat,
     )

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

@@ -72,8 +72,10 @@
     SO3Linear,
 )
 from .triton import (
-    rotate_back,
-    rotate_to_local,
+    rotate_back_block,
+    rotate_back_dense,
+    rotate_to_local_block,
+    rotate_to_local_dense,
 )
 from .utils import (
     ATTN_RES_MODES,
@@ -1334,6 +1336,7 @@ def __init__(
                     dtype=self.compute_dtype,
                     layout="flat",
                     grid_resolution_list=self.s2_grid_resolution,
+                    coefficient_layout="m_major",
                     grid_method=self.s2_grid_method,
                     grid_branches=node_wise_branches,
                     mlp_bias=self.mlp_bias,
@@ -1443,26 +1446,18 @@ def forward(
             D_full = edge_cache.D_full
             x_dst_local: torch.Tensor | None = None
             if self.use_triton_infer and not self.training:
-                # ``rotate_to_local`` / ``rotate_back`` pick the kernel from the
-                # coefficient layout, not from this flag: the block-diagonal
-                # kernel for the canonical m-major ``mmax == 1`` layout used here
-                # (with ``lmax`` inferred from ``ebed_dim_full``), and a dense
-                # kernel for any other ``(lmax, mmax)``. Both compose with the
-                # traced force path through their functional custom-op autograd.
-                x_local = rotate_to_local(
-                    x,
-                    src,
-                    D_full,
-                    self.coeff_index_m,
-                    self.ebed_dim_full,
+                # This module owns ``coeff_index_m``.  Use the block kernels only
+                # for the canonical m-major ``mmax == 1`` layout; the general
+                # kernels honor ``coeff_index_m`` for every other layout.
+                rotate_to_local_fn = (
+                    rotate_to_local_block if self.mmax == 1 else rotate_to_local_dense
+                )
+                x_local = rotate_to_local_fn(
+                    x, src, D_full, self.coeff_index_m, self.ebed_dim_full
                 )  # (E, D_m, C_wide)
                 if self.node_wise_grid_product is not None:
-                    x_dst_local = rotate_to_local(
-                        x,
-                        dst,
-                        D_full,
-                        self.coeff_index_m,
-                        self.ebed_dim_full,
+                    x_dst_local = rotate_to_local_fn(
+                        x, dst, D_full, self.coeff_index_m, self.ebed_dim_full
                     )  # (E, D_m, C_wide)
             else:
                 D_m_prime = project_D_to_m(
@@ -1620,7 +1615,10 @@ def apply_bias_correction(
         with nvtx_range("SO2Conv/rotate_back"):
             Dt_full = edge_cache.Dt_full
             if self.use_triton_infer and not self.training:
-                x_message = rotate_back(
+                rotate_back_fn = (
+                    rotate_back_block if self.mmax == 1 else rotate_back_dense
+                )
+                x_message = rotate_back_fn(
                     x_local,
                     Dt_full,
                     self.coeff_index_m,

deepmd/pt/model/descriptor/sezm_nn/triton/__init__.py

@@ -1,17 +1,21 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Hardware-accelerated SeZM/DPA4 operators.
 
-This package hosts clean, ``torch.compile``-composable Triton implementations of
-SeZM hot paths. The first member is the fused SO(2)/Wigner rotation pair used by
-the SO(2) convolution (``rotate_to_local`` / ``rotate_back``).
+This package hosts ``make_fx``-composable Triton implementations of SeZM hot
+paths.  The SO(2) rotation API exposes a general dense path that honors arbitrary
+coefficient indices and a block path for the canonical m-major ``mmax=1`` layout.
 """
 
 from .so2_rotation import (
-    rotate_back,
-    rotate_to_local,
+    rotate_back_block,
+    rotate_back_dense,
+    rotate_to_local_block,
+    rotate_to_local_dense,
 )
 
 __all__ = [
-    "rotate_back",
-    "rotate_to_local",
+    "rotate_back_block",
+    "rotate_back_dense",
+    "rotate_to_local_block",
+    "rotate_to_local_dense",
 ]