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

@@ -1334,6 +1334,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,

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

@@ -1169,13 +1169,9 @@ def _block_layout_lmax(coeff_index: Tensor, dim_full: int) -> int:
     """Return ``lmax`` if ``(coeff_index, dim_full)`` is the m-major ``mmax=1``
     layout that the block-diagonal kernels assume, else ``-1``.
 
-    Detection uses ONLY shapes / python ints -- never tensor *values* -- so it is
-    safe under ``make_fx`` / fake-tensor tracing (the production compiled
-    inference path). The test is: ``dim_full`` is a perfect square ``(lmax+1)^2``
-    and ``Dm == 3*lmax+1``. For a fixed ``lmax`` the reduced size ``Dm`` is
-    strictly increasing in ``mmax`` (``lmax+1``, ``3*lmax+1``, ``5*lmax-1``, ...),
-    so ``Dm == 3*lmax+1`` uniquely pins ``mmax == 1``; combined with the model's
-    canonical ``build_m_major_index`` ordering this fully determines the layout.
+    Shape detection handles fake-tensor tracing; real tensors are additionally
+    checked against ``build_m_major_index(lmax, 1)`` because the block kernels
+    ignore ``coeff_index`` values.
     """
     dim_full = int(dim_full)
     root = math.isqrt(dim_full)
@@ -1188,9 +1184,24 @@ def _block_layout_lmax(coeff_index: Tensor, dim_full: int) -> int:
         return -1
     if lmax < 1 or numel != 3 * lmax + 1:
         return -1
+    if not _canonical_block_index_values(coeff_index, lmax):
+        return -1
     return lmax
 
 
+def _canonical_block_index_values(coeff_index: Tensor, lmax: int) -> bool:
+    """Return whether real ``coeff_index`` values match the m-major mmax=1 layout."""
+    if getattr(coeff_index, "fake_mode", None) is not None:
+        return True
+    if coeff_index.device.type == "meta":
+        return True
+    try:
+        expected = build_m_major_index(int(lmax), 1, device=coeff_index.device)
+        return torch.equal(coeff_index, expected)
+    except Exception:  # pragma: no cover - exotic tensor subclasses
+        return False
+
+
 def _launch_bd_to_local_fwd(
     x: Tensor, src: Tensor, wigner: Tensor, lmax: int
 ) -> Tensor: