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/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

source/op/pt/CMakeLists.txt

@@ -1,6 +1,20 @@
 file(GLOB OP_SRC print_summary.cc comm.cc tabulate_multi_device.cc)
-# Fused graph-lower inference operators (CUDA / cuBLAS); GPU builds only.
+# Fused graph-lower inference operators (CUDA / cuBLAS); GPU builds only. The
+# CUDA language is scoped per directory, so it must be enabled here: the sibling
+# GPU library turns it on only within its own subtree, which does not cover this
+# target.
 if(USE_CUDA_TOOLKIT)
+  find_package(CUDAToolkit REQUIRED)
+  if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+    # CUDA 12.9 and 13.x CCCL fail to compile CUB/Thrust with -arch=all.
+    if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "12.9" AND CUDAToolkit_VERSION
+                                                            VERSION_LESS "14.0")
+      set(CMAKE_CUDA_ARCHITECTURES all-major)
+    else()
+      set(CMAKE_CUDA_ARCHITECTURES all)
+    endif()
+  endif()
+  enable_language(CUDA)
   list(
     APPEND
     OP_SRC

source/op/pt/graph_fitting.cu

@@ -266,6 +266,11 @@ torch::Tensor graph_fitting_backward(torch::Tensor d_e,
                                      std::vector<int64_t> resnets,
                                      torch::Tensor w_head) {
   const long n_node = x.size(0);
+  auto d_x = torch::empty_like(x);
+  // Guard the empty system before the division by ``n_node`` below.
+  if (n_node == 0) {
+    return d_x;
+  }
   auto stream = at::cuda::getCurrentCUDAStream();
   const int n_layer = (int)ws.size();
   const long total_width = saved.numel() / (2 * n_node);
@@ -274,10 +279,6 @@ torch::Tensor graph_fitting_backward(torch::Tensor d_e,
     offset[l + 1] = offset[l] + ws[l].size(1);
   }
   auto f32 = x.options().dtype(torch::kFloat32);
-  auto d_x = torch::empty_like(x);
-  if (n_node == 0) {
-    return d_x;
-  }
   auto d_e_c = d_e.contiguous();
 
   long width_max = 0;