Skip to content

fix(lammps): expose Kokkos pair style through plugin - #5980

Draft
njzjz-bot wants to merge 5 commits into
deepmodeling:masterfrom
njzjz-bot:agent/lammps-kokkos-plugin
Draft

njzjz-bot wants to merge 5 commits into
deepmodeling:masterfrom
njzjz-bot:agent/lammps-kokkos-plugin

Conversation

@njzjz-bot

@njzjz-bot njzjz-bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Summary

  • register deepmd/kk and deepmd/kk/device from the dynamic LAMMPS plugin when it is built with DEEPMD_LAMMPS_KOKKOS=ON
  • expose DP_ENABLE_LAMMPS_KOKKOS=1 through the Python build backend so wheel builds can request the existing CMake option
  • compile the plugin target through Kokkos's downstream compiler launcher when Kokkos support is enabled
  • reject the build-backend toggle when no LAMMPS plugin version is selected
  • add focused tests for the Python build-backend toggle

Why

The Kokkos pair style is already implemented and the plugin CMake target already has a DEEPMD_LAMMPS_KOKKOS option, but deepmdplugin.cpp currently does not register the Kokkos pair style. As a result, compiling those sources into the module is not sufficient for pair_style deepmd/kk to become available through the dynamic plugin.

This PR provides the prerequisite wiring needed before enabling Kokkos in PyPI wheel builds.

PyPI / CUDA architecture investigation

The requested distribution target was a small representative CUDA set (sm_70, sm_89, sm_120 plus a compute_80 PTX fallback). I did not turn that set on in the release wheel in this PR because the current packaging stack needs additional coordination with the PyPI LAMMPS host:

  1. Kokkos rejects multiple configured GPU architectures in one build, so a Kokkos/LAMMPS build cannot be configured as VOLTA70 + ADA89 + BLACKWELL120 in the same way as an ordinary CUDA fatbin.
  2. DeePMD-kit currently fetches stable_22Jul2025_update2 for the wheel plugin headers, and the current PyPI lammps wheel uses stable_22Jul2025_update4; both LAMMPS releases bundle Kokkos 4.6.02. Native BLACKWELL120 support was added in Kokkos 4.7. A newer external Kokkos could be evaluated, but the LAMMPS host and plugin must use a compatible configuration together.
  3. The current PyPI lammps wheel is not built with PKG_KOKKOS. Enabling only the DeePMD plugin side would therefore produce a wheel that installs but cannot safely run deepmd/kk with the default PyPI LAMMPS host.

A safe packaging follow-up is to make the PyPI LAMMPS host Kokkos-enabled with the same Kokkos toolchain as the DeePMD plugin, then determine whether a supported fat-binary strategy is viable or whether matched architecture variants are required. The desired native targets remain Volta70, Ada89 (including RTX 4090), and Blackwell120, with compute_80 PTX considered as a forward-compatibility fallback rather than a substitute for Kokkos's architecture configuration.

Validation

  • added focused tests for DP_ENABLE_LAMMPS_KOKKOS
  • inspected the dynamic plugin registration path and Kokkos pair-style instantiations
  • inspected the PyPI wheel build backend and LAMMPS wheel packaging path
  • verified that Kokkos 4.6.02 exports kokkos_compilation(TARGET ...) for downstream CUDA compilation
  • full Kokkos wheel compilation is not exercised by the existing release-wheel CI yet because no matching Kokkos-enabled PyPI LAMMPS host is currently supplied

Agent: ChatGPT
Model: GPT-5.6 Sol

Agent: ChatGPT
Model: GPT-5.6 Sol
Agent: ChatGPT
Model: GPT-5.6 Sol
Agent: ChatGPT
Model: GPT-5.6 Sol

@njzjz-bot njzjz-bot left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed this draft against current master. The original Kokkos plugin-registration portion has since landed in a more complete form via #5988, and the Kokkos runtime-linking assumptions in this old branch were superseded by #5989. Current master registers all six Kokkos aliases (deepmd/kk, /device, /host, and the DPA4Spin counterparts) and deliberately avoids linking a second Kokkos::kokkos runtime into the plugin because that can create a duplicate CUDA/Kokkos runtime and corrupt plugin shutdown.

That means this branch should not be merged or rebased wholesale. The only concept here that is not currently present on master is the Python build-backend DP_ENABLE_LAMMPS_KOKKOS environment toggle/tests. If that toggle is still wanted, it should be split/reimplemented narrowly on current master and validated against the now-shared-runtime CMake path rather than carrying forward this branch's old CMake/plugin code. The branch is also currently non-mergeable and its historical Test Python workflow is red.

I left one inline note at the obsolete Kokkos compilation/linkage area. Because this is still a draft (and the executing account is also the PR author), no APPROVE/REQUEST_CHANGES disposition is submitted.


Agent: ChatGPT
Model: GPT-5.6 Sol
GitHub account: njzjz-bot
Reviewed head: 0f8e3bb
Trigger: scheduled all-PR monitoring

target_compile_definitions(${libname} PRIVATE LMP_KOKKOS)
target_include_directories(${libname} PRIVATE ${LAMMPS_HEADER_DIR}/KOKKOS)
target_link_libraries(${libname} PUBLIC Kokkos::kokkos)
kokkos_compilation(TARGET ${libname})

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This build path is obsolete relative to current master. The surrounding branch still links Kokkos::kokkos directly, but merged #5989 replaced that with imported compile/interface properties specifically because a runtime LAMMPS plugin must share liblammps' Kokkos runtime; linking another Kokkos runtime can create a second CUDA/Kokkos instance and corrupt shutdown. Do not carry this kokkos_compilation() hunk forward by rebasing the old CMake block. If the Python DP_ENABLE_LAMMPS_KOKKOS toggle is still desired, apply only that narrow frontend change on top of the current shared-runtime CMake implementation and validate a real plugin build there.


Agent: ChatGPT
Model: GPT-5.6 Sol
GitHub account: njzjz-bot
Reviewed head: 0f8e3bb
Trigger: scheduled all-PR monitoring

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant