deepmodeling / deepmodeling/deepmd-kit

Move torch-only kernels under pt_expt

Open
#5,782 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

## Scope

Move the torch-only `deepmd/kernels/` package to `deepmd/pt_expt/kernels/`.

The package is implemented through `torch.library.triton_op`, `wrap_triton`, and `torch.ops.deepmd.*`; it is not backend-agnostic and cannot be consumed by the dpmodel, TensorFlow, JAX, or Paddle backends. `pt_expt` is the long-term owner and the design target for these exportable kernels.

## Dependency direction

During the legacy `pt` backend's deprecation window, `pt` should import kernels from `pt_expt`. The dependency must not point from `pt_expt` back to `pt`, so removing the legacy backend also removes the temporary imports cleanly.

## Invariant

`deepmd/dpmodel/**` must not import the torch-specific kernels package.

This issue tracks the relocation separately from DPA1 CUDA performance PR #5758 and should include the existing DPA4 kernels in the same package-level move.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.