fix

Author: OutisLi

deepmd/kernels/cuda/dpa1/graph_compress.py

@@ -35,8 +35,6 @@
     "op_available",
 ]
 
-_TILE_EDGES = 128  # matches kTileEdges in dpa1_graph_compress.cu
-
 # Instantiated table widths in dpa1_graph_compress.cu; each divides the 256
 # threads per block, so the moment walk splits channels evenly. An arbitrary
 # ``ng`` is served by the smallest width that is at least ``ng``, zero-padding

deepmd/kernels/cuda/dpa1/graph_energy_force.py

@@ -245,9 +245,13 @@ def _cpu(
         protection,
         nnei,
     )
+    # Assemble force / virial in the descriptor weight precision, matching
+    # ``_fake`` and the CUDA operator; the sub-operator would otherwise return
+    # fp64 whenever the edge inputs are fp64.
+    fprec = w1.dtype
     force, atom_virial, virial = torch.ops.deepmd.edge_force_virial(
-        g_e,
-        edge_vec,
+        g_e.to(fprec),
+        edge_vec.to(fprec),
         edge_index,
         edge_mask,
         n_node,

deepmd/kernels/cuda/graph_fitting.py

@@ -16,10 +16,10 @@
 
 Usage and pitfalls
 ------------------
-* The forward's second output packs the layer activations and their
-  derivatives as one flat buffer ``[h chunks | adot chunks]``; it is an
-  autograd save, never a user-facing value, and receives no gradient
-  (``set_materialize_grads(False)``).
+* The forward's second output packs the layer activation derivatives as one
+  flat buffer of ``adot chunks`` (the activations themselves are a forward-only
+  transient); it is an autograd save, never a user-facing value, and receives
+  no gradient (``set_materialize_grads(False)``).
 * The head bias is passed as a device tensor, not a Python float: reading a
   value host-side (``.item()``) inside the dispatch path would fail under
   symbolic tracing (``GuardOnDataDependentSymNode``) and force a GPU sync per
@@ -115,7 +115,7 @@ def _forward_fake(
     total_width = sum(int(w.shape[1]) for w in ws)
     return (
         x.new_empty(n_node, 1, dtype=torch.float64),
-        x.new_empty(2 * n_node * total_width, dtype=torch.float32),
+        x.new_empty(n_node * total_width, dtype=torch.float32),
     )
 
 
@@ -163,7 +163,7 @@ def _cpu_forward(
     bias_atom_e: torch.Tensor,
     act: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    hs, adots = [], []
+    adots = []
     cur = x.to(torch.float32)
     for w, b, idt, res in zip(ws, bs, idts, resnets):
         pre = cur @ w
@@ -178,17 +178,15 @@ def _cpu_forward(
         if idt.numel():
             a = a * idt
             adot = adot * idt
-        y = a + cur if (res and w.shape[0] == w.shape[1]) else a
-        hs.append(y)
         adots.append(adot)
-        cur = y
+        cur = a + cur if (res and w.shape[0] == w.shape[1]) else a
     e = (cur @ w_head[:, None]).to(torch.float64)
     if b_head.numel():
         e = e + b_head.to(torch.float64)
     e = e + bias_atom_e[atype][:, None]
-    # Chunk layout mirrors the CUDA op: [h chunks | adot chunks], each a
-    # contiguous row-major (N, w_l) block.
-    saved = torch.cat([t.reshape(-1) for t in hs + adots])
+    # Chunk layout mirrors the CUDA op: adot chunks only, each a contiguous
+    # row-major (N, w_l) block; the backward needs only these derivatives.
+    saved = torch.cat([t.reshape(-1) for t in adots])
     return e, saved
 
 
@@ -201,13 +199,11 @@ def _cpu_backward(
     w_head: torch.Tensor,
 ) -> torch.Tensor:
     n_node = x.shape[0]
-    total_width = sum(int(w.shape[1]) for w in ws)
     offset = [0]
     for w in ws:
         offset.append(offset[-1] + int(w.shape[1]))
-    base = n_node * total_width
     adots = [
-        saved[base + offset[l] * n_node : base + offset[l + 1] * n_node].reshape(
+        saved[offset[l] * n_node : offset[l + 1] * n_node].reshape(
             n_node, int(ws[l].shape[1])
         )
         for l in range(len(ws))

deepmd/kernels/triton/dpa1/sweep_tile_configs.py

@@ -73,9 +73,6 @@
 # the per-program register tile small, large values amortize launch overhead.
 _BLOCK_E_CANDIDATES = (1, 2, 4, 8, 16, 32)
 _WARP_CANDIDATES = (2, 4, 8)
-# Reject a configuration whose forward+backward is slower than this multiple of
-# the best observed time; a large backward regression signals a register spill.
-_REJECT_RATIO = 3.0
 _REL_TOL = 1e-5
 # An edge_conv candidate is only recorded when it beats the universal default by
 # more than ``1 - _EDGE_WIN_RATIO`` (5%); otherwise the default is kept, so a
@@ -186,16 +183,20 @@ def fwd_bwd(bn: int, nw: int) -> None:
 
 def _make_edge_inputs(
     n_edge: int, n_node: int, ng: int, h1: int, device: torch.device
-) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+    p = 4096
     z2 = torch.randn(n_edge, ng, dtype=torch.float32, device=device)
     h1t = torch.randn(n_edge, h1, dtype=torch.float32, device=device)
     idt = torch.ones(ng, dtype=torch.float32, device=device)
+    tt = torch.randn(p, ng, dtype=torch.float32, device=device) * 0.3
+    idx = torch.randint(0, p, (n_edge,), dtype=torch.int64, device=device)
+    sw = torch.rand(n_edge, dtype=torch.float32, device=device)
     rr = torch.randn(n_edge, 4, dtype=torch.float32, device=device)
     dst = torch.randint(0, n_node, (n_edge,), dtype=torch.int64, device=device)
     edge_mask = (torch.rand(n_edge, dtype=torch.float32, device=device) > 0.05).to(
         dtype=torch.float32
     )
-    return z2, h1t, idt, rr, dst, edge_mask
+    return z2, h1t, idt, tt, idx, sw, rr, dst, edge_mask
 
 
 def sweep_edge(
@@ -238,22 +239,31 @@ def sweep_edge(
     device = device or torch.device("cuda")
     torch.backends.cuda.matmul.allow_tf32 = False
     act = 0
-    z2, h1t, idt, rr, dst, em = _make_edge_inputs(n_edge, n_node, ng, h1, device)
-    ref = _edge_conv_reference(z2, h1t, idt, rr, dst, em, n_node, resnet_mult, act)
+    # The launch configuration is keyed by (ng, H1) and is independent of the
+    # tebd-input mode; the strip gate (gated = 1) is the register-heaviest case,
+    # so its optimum is a safe upper bound for concat (gated = 0).
+    gated = 1
+    # Shared leading arguments of the three edge_conv entry points, ordered to
+    # match their signatures so each call splats this tuple and appends only its
+    # launch configuration.
+    edge_args = (
+        *_make_edge_inputs(n_edge, n_node, ng, h1, device),
+        n_node,
+        resnet_mult,
+        act,
+        gated,
+    )
+    ref = _edge_conv_reference(*edge_args)
     gout = torch.randn_like(ref)
 
     def fwd_bwd(be: int, nw: int) -> None:
-        _edge_conv_fwd_impl(z2, h1t, idt, rr, dst, em, n_node, resnet_mult, act, be, nw)
-        _edge_conv_bwd_impl(
-            gout, z2, h1t, idt, rr, dst, em, n_node, resnet_mult, act, be, nw
-        )
+        _edge_conv_fwd_impl(*edge_args, be, nw)
+        _edge_conv_bwd_impl(gout, *edge_args, be, nw)
 
     results: dict[tuple[int, int], float] = {}
     for be, nw in itertools.product(_BLOCK_E_CANDIDATES, _WARP_CANDIDATES):
         try:
-            out = _edge_conv_fwd_impl(
-                z2, h1t, idt, rr, dst, em, n_node, resnet_mult, act, be, nw
-            )
+            out = _edge_conv_fwd_impl(*edge_args, be, nw)
             rel = (out - ref).abs().max().item() / ref.abs().max().item()
             if rel > _REL_TOL:
                 continue

deepmd/kernels/triton/env_mat.py

@@ -475,7 +475,6 @@ def _env_mat_reference(
     use_exp_switch: bool,
 ) -> tuple[Tensor, Tensor, Tensor]:
     """Eager environment matrix, numerically identical to ``prod_env_mat``."""
-    nf, nloc, nnei = nlist.shape
     mask, _, d, length_safe, _ = _geometry(coord, nlist)
     q = length_safe + protection
     sw_raw, _ = _switch_reference(length_safe, rcut_smth, rcut, use_exp_switch)

deepmd/pt/model/descriptor/se_atten.py

@@ -667,15 +667,10 @@ def forward(
                     rr,
                     gated=0,
                 )
-                # ``gg`` (the pair representation g2) is not formed on this path.
-                gg = torch.empty(
-                    nframes,
-                    nloc,
-                    self.nnei,
-                    self.filter_neuron[-1],
-                    dtype=xyz_scatter.dtype,
-                    device=xyz_scatter.device,
-                )
+                # ``gg`` (the pair representation g2) is not formed on this path;
+                # it is returned as ``None``, so a zero-element placeholder keeps
+                # ``gg`` a tensor without the full (nf, nloc, nnei, ng) allocation.
+                gg = xyz_scatter.new_empty(0)
                 g2_is_none = True
             else:
                 # nfnl x nnei x ng
@@ -755,15 +750,10 @@ def forward(
                     rr,
                     gated=1,
                 )
-                # ``gg`` (the pair representation g2) is not formed on this path.
-                gg = torch.empty(
-                    nframes,
-                    nloc,
-                    self.nnei,
-                    self.filter_neuron[-1],
-                    dtype=xyz_scatter.dtype,
-                    device=xyz_scatter.device,
-                )
+                # ``gg`` (the pair representation g2) is not formed on this path;
+                # it is returned as ``None``, so a zero-element placeholder keeps
+                # ``gg`` a tensor without the full (nf, nloc, nnei, ng) allocation.
+                gg = xyz_scatter.new_empty(0)
                 g2_is_none = True
             else:
                 # (nf x nl) x nnei x ng

deepmd/pt_expt/descriptor/se_atten_v2.py

@@ -71,6 +71,44 @@ def fused_energy_force_graph(self, *args: Any, **kwargs: Any) -> Any:
 
         return DescrptDPA1.fused_energy_force_graph(self, *args, **kwargs)
 
+    # Graph-lower routing and the fused-kernel helpers it dispatches to,
+    # forwarded to DescrptDPA1 so the accelerated CUDA / Triton / tabulated
+    # graph paths serve se_atten_v2 rather than only the dpmodel reference.
+    def call_graph(self, *args: Any, **kwargs: Any) -> Any:
+        from deepmd.pt_expt.descriptor.dpa1 import (
+            DescrptDPA1,
+        )
+
+        return DescrptDPA1.call_graph(self, *args, **kwargs)
+
+    def _call_graph_cuda(self, *args: Any, **kwargs: Any) -> Any:
+        from deepmd.pt_expt.descriptor.dpa1 import (
+            DescrptDPA1,
+        )
+
+        return DescrptDPA1._call_graph_cuda(self, *args, **kwargs)
+
+    def _call_graph_cuda_compress(self, *args: Any, **kwargs: Any) -> Any:
+        from deepmd.pt_expt.descriptor.dpa1 import (
+            DescrptDPA1,
+        )
+
+        return DescrptDPA1._call_graph_cuda_compress(self, *args, **kwargs)
+
+    def _call_graph_triton(self, *args: Any, **kwargs: Any) -> Any:
+        from deepmd.pt_expt.descriptor.dpa1 import (
+            DescrptDPA1,
+        )
+
+        return DescrptDPA1._call_graph_triton(self, *args, **kwargs)
+
+    def _call_graph_compress_reference(self, *args: Any, **kwargs: Any) -> Any:
+        from deepmd.pt_expt.descriptor.dpa1 import (
+            DescrptDPA1,
+        )
+
+        return DescrptDPA1._call_graph_compress_reference(self, *args, **kwargs)
+
     def _call_triton(self, *args: Any, **kwargs: Any) -> Any:
         from deepmd.pt_expt.descriptor.dpa1 import (
             DescrptDPA1,

deepmd/pt_expt/model/edge_transform_output.py

@@ -105,7 +105,7 @@ def edge_energy_deriv(
     ):
         g_e = g_e.to(force_precision)
         edge_vec = edge_vec.to(force_precision)
-    if cuda_infer_level() >= 1 and fused_scatter_available():
+    if cuda_infer_level() >= 1 and not create_graph and fused_scatter_available():
         n_cap = node_capacity if node_capacity is not None else int(n_node.sum())
         force, atom_virial, virial = fused_edge_force_virial(
             g_e, edge_vec, edge_index, edge_mask, n_node, n_cap, do_atomic_virial

pyproject.toml

@@ -305,7 +305,7 @@ environment-pass = [
 ]
 before-all = [
     """if [ ! -z "${DP_PKG_NAME}" ]; then sed -i "s/name = \\"deepmd-kit\\"/name = \\"${DP_PKG_NAME}\\"/g" pyproject.toml; fi""",
-    """{ if [ -n "${CUDA_VERSION}" ] && [ "$(uname -m)" = "x86_64" ] ; then yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-nvcc-${CUDA_VERSION/./-} cuda-cudart-devel-${CUDA_VERSION/./-}; fi }""",
+    """{ if [ -n "${CUDA_VERSION}" ] && [ "$(uname -m)" = "x86_64" ] ; then yum config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-nvcc-${CUDA_VERSION/./-} cuda-cudart-devel-${CUDA_VERSION/./-}; fi }""",
 ]
 before-build = [
 ]

source/CMakeLists.txt

@@ -356,6 +356,16 @@ if(ENABLE_PYTORCH AND NOT DEEPMD_C_ROOT)
     endif()
   endif()
   find_package(Torch REQUIRED)
+  # TorchConfig populates TORCH_CUDA_LIBRARIES only when PyTorch itself was
+  # built with CUDA. The fused graph-lower operators in op/pt include ATen CUDA
+  # headers and link libtorch_cuda, so they can be compiled only against a
+  # CUDA-enabled PyTorch: a CPU-only wheel ships the ATen/c10 CUDA headers but
+  # neither the generated c10/cuda/impl/cuda_cmake_macros.h nor libtorch_cuda.
+  if(TORCH_CUDA_LIBRARIES)
+    set(DEEPMD_TORCH_HAS_CUDA TRUE)
+  else()
+    set(DEEPMD_TORCH_HAS_CUDA FALSE)
+  endif()
   if(NOT Torch_VERSION VERSION_LESS "2.1.0")
     set_if_higher(CMAKE_CXX_STANDARD 17)
   elseif(NOT Torch_VERSION VERSION_LESS "1.5.0")