fix

Author: OutisLi

.github/workflows/build_cc.yml

@@ -38,7 +38,7 @@ jobs:
           wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
           && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
           && sudo apt-get update \
-          && sudo apt-get -y install cuda-cudart-dev-12-2 cuda-nvcc-12-2 libcublas-dev-12-2 libcusparse-dev-12-2
+          && sudo apt-get -y install cuda-cudart-dev-12-2 cuda-nvcc-12-2 libcublas-dev-12-2 libcusparse-dev-12-2 libcusolver-dev-12-2
         if: matrix.variant == 'cuda120'
         env:
           DEBIAN_FRONTEND: noninteractive

.github/workflows/codeql.yml

@@ -42,7 +42,7 @@ jobs:
           wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
            && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
            && sudo apt-get update \
-           && sudo apt-get -y install cuda-cudart-dev-12-2 cuda-nvcc-12-2 libcublas-dev-12-2 libcusparse-dev-12-2
+           && sudo apt-get -y install cuda-cudart-dev-12-2 cuda-nvcc-12-2 libcublas-dev-12-2 libcusparse-dev-12-2 libcusolver-dev-12-2
           python -m pip install uv
           uv pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
         env:

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,

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/./-} libcublas-devel-${CUDA_VERSION/./-} libcusparse-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/./-} libcublas-devel-${CUDA_VERSION/./-} libcusparse-devel-${CUDA_VERSION/./-} libcusolver-devel-${CUDA_VERSION/./-}; fi }""",
 ]
 before-build = [
 ]

source/install/docker_package_c.sh

@@ -7,9 +7,10 @@ PYTORCH_DEPENDENCY_GROUP=${PYTORCH_DEPENDENCY_GROUP:-pin_pytorch_cpu}
 PYTORCH_TORCH_BACKEND=${PYTORCH_TORCH_BACKEND:-cpu}
 TENSORFLOW_DEPENDENCY_GROUP=${TENSORFLOW_DEPENDENCY_GROUP:-pin_tensorflow_cpu}
 # CUDA version tag (e.g. "12-9") parsed from the manylinux CUDA image name. It
-# selects the cuBLAS/cuSPARSE development packages installed inside the
-# container: the fused PyTorch CUDA operators pull in ATen CUDA headers that
-# require cublas and cusparse headers absent from the base image.
+# selects the cuBLAS/cuSPARSE/cuSOLVER development packages installed inside the
+# container: the fused PyTorch CUDA operators pull in ATen CUDA headers
+# (CUDAContextLight.h) that include cublas_v2.h, cusparse.h and cusolverDn.h,
+# none of which ship in the base image.
 CUDA_VERSION_DASH=$(printf '%s' "${MANYLINUX_CUDA_IMAGE}" | sed -n 's/.*cuda\([0-9][0-9]*\)_\([0-9][0-9]*\).*/\1-\2/p')
 
 docker run --rm -v "${SCRIPT_PATH}/../..":/root/deepmd-kit -w /root/deepmd-kit \
@@ -27,8 +28,8 @@ docker run --rm -v "${SCRIPT_PATH}/../..":/root/deepmd-kit -w /root/deepmd-kit \
 	"${MANYLINUX_CUDA_IMAGE}" \
 	/bin/bash -lc 'set -euo pipefail
             if [ -n "${CUDA_VERSION_DASH:-}" ]; then
-              yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo
-              yum install -y "libcublas-devel-${CUDA_VERSION_DASH}" "libcusparse-devel-${CUDA_VERSION_DASH}"
+              yum config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo
+              yum install -y "libcublas-devel-${CUDA_VERSION_DASH}" "libcusparse-devel-${CUDA_VERSION_DASH}" "libcusolver-devel-${CUDA_VERSION_DASH}"
             fi
             export PATH="$(dirname "${PYTHON_BIN}"):${PATH}"
             "${PYTHON_BIN}" -m pip install uv