ci: split CUDA tests into separate jobs

.github/workflows/test_cuda.yml

@@ -14,11 +14,10 @@ concurrency:
   cancel-in-progress: true
 name: Test CUDA
 jobs:
-  test_cuda:
-    name: Test Python and C++ on CUDA
+  test_python:
+    name: Test Python on CUDA
     runs-on: gpu
-    # The full CUDA suite (serial pytest ~3h43m + C++ + LAMMPS) exceeds the
-    # default 360-min job limit; the self-hosted GPU runner has no hard cap.
+    # The serial Python test suite can take several hours on the GPU runner.
     timeout-minutes: 480
     # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845
     # container:
@@ -71,6 +70,51 @@ jobs:
           XLA_PYTHON_CLIENT_PREALLOCATE: false
           XLA_PYTHON_CLIENT_ALLOCATOR: platform
           FLAGS_use_stride_compute_kernel: 0
+
+  test_cc:
+    name: Test C++ on CUDA
+    runs-on: gpu
+    if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group')
+    steps:
+      # Jobs run on separate runners, so the C++ job needs its own complete
+      # CUDA toolchain and Python dependency installation.
+      - uses: actions/checkout@v7
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.11"
+          # cache: 'pip'
+      - name: Install wget and unzip
+        run: sudo apt-get update && sudo apt-get install -y wget unzip
+      - uses: lukka/get-cmake@latest
+        with:
+          useLocalCache: true
+          useCloudCache: false
+      - run: |
+          UBUNTU_VERSION=$(lsb_release -rs | tr -d '.')
+          OS_NAME="ubuntu${UBUNTU_VERSION}"
+          echo "Current OS: ${OS_NAME}"
+
+          wget https://developer.download.nvidia.com/compute/cuda/repos/${OS_NAME}/x86_64/cuda-keyring_1.1-1_all.deb \
+          && sudo dpkg -i cuda-keyring_1.1-1_all.deb \
+          && sudo apt-get update \
+          && sudo apt-get -y install cuda-toolkit-12-9 cudnn9-cuda-12
+          echo "CUDA_PATH=/usr/local/cuda-12.9" >> $GITHUB_ENV
+          echo "/usr/local/cuda-12.9/bin" >> $GITHUB_PATH
+        # if: false # skip as we use nvidia image
+      - run: python -m pip install -U uv
+      - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax_gpu
+      - run: |
+          export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
+          export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
+          pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn "paddlepaddle-gpu==3.4.0.dev20260310"
+          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
+          # See https://github.com/jax-ml/jax/issues/29042
+          source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
+        env:
+          DP_VARIANT: cuda
+          DP_ENABLE_NATIVE_OPTIMIZATION: 1
+          DP_ENABLE_PYTORCH: 1
+      - run: dp --version
       - name: Convert models
         run: source/tests/infer/convert-models.sh
       - run: |
@@ -99,12 +143,12 @@ jobs:
           CUDA_VISIBLE_DEVICES: 0
   pass:
     name: Pass testing on CUDA
-    needs: [test_cuda]
+    needs: [test_python, test_cc]
     runs-on: ubuntu-slim
     if: always()
     steps:
       - name: Decide whether the needed jobs succeeded or failed
         uses: re-actors/alls-green@release/v1
         with:
           jobs: ${{ toJSON(needs) }}
-          allowed-skips: test_cuda
+          allowed-skips: test_python, test_cc