feat(pt): Add Optimized DPA4 Neo inference path using CuTe

.github/workflows/test_cuda.yml

@@ -58,14 +58,15 @@ jobs:
       - 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)')
-          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
+          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,cute,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
+      - run: python -c "import cutlass.cute"
       - run: python -m pytest source/tests --ignore=source/tests/pd
         env:
           NUM_WORKERS: 0

deepmd/kernels/cute/neo/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Neo-specialized CuTe kernels and PyTorch integration."""

deepmd/kernels/cute/neo/compile_cache.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Device-aware caching for architecture-specific CuTe compilation."""
+
+from __future__ import (
+    annotations,
+)
+
+from collections.abc import (
+    Callable,
+)
+from contextlib import (
+    nullcontext,
+)
+from functools import (
+    lru_cache,
+    wraps,
+)
+from typing import (
+    Any,
+    TypeVar,
+    cast,
+)
+
+_T = TypeVar("_T", bound=Callable[..., Any])
+
+
+def current_cuda_compile_identity() -> tuple[int, int, int]:
+    """Return the current CUDA device and its compute capability."""
+    import torch
+
+    device_index = torch.cuda.current_device()
+    major, minor = torch.cuda.get_device_capability(device_index)
+    return device_index, major, minor
+
+
+def device_aware_lru_cache(
+    *,
+    maxsize: int,
+    identity_getter: Callable[[], tuple[int, int, int]] = current_cuda_compile_identity,
+) -> Callable[[_T], _T]:
+    """Cache a compile factory separately for each CUDA device architecture."""
+
+    def decorate(function: _T) -> _T:
+        @lru_cache(maxsize=maxsize)
+        def cached(
+            identity: tuple[int, int, int],
+            args: tuple[Any, ...],
+            kwargs: tuple[tuple[str, Any], ...],
+        ) -> Any:
+            import torch
+
+            device_index = identity[0]
+            device_count = getattr(torch.cuda, "device_count", None)
+            device_is_visible = device_count is None or device_index < device_count()
+            compile_device = (
+                torch.cuda.device(device_index)
+                if torch.cuda.is_available() and device_is_visible
+                else nullcontext()
+            )
+            with compile_device:
+                return function(*args, **dict(kwargs))
+
+        @wraps(function)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return cached(
+                identity_getter(),
+                args,
+                tuple(sorted(kwargs.items())),
+            )
+
+        wrapper.cache_clear = cached.cache_clear
+        wrapper.cache_info = cached.cache_info
+        wrapper.cache_parameters = cached.cache_parameters
+        wrapper._deepmd_cute_cached = True
+        return cast("_T", wrapper)
+
+    return decorate