Skip to content

Move torch-only kernels under pt_expt #5780

Description

@OutisLi

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.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions