feat(vlasov1d): add native CUDA/MPI backend with kinetic ions

.github/workflows/cpu-tests.yaml

@@ -153,9 +153,39 @@ jobs:
     - uses: ./.github/actions/setup_repo
       with:
         python-version: '3.11'
+    - name: Install native validation dependencies
+      run: |
+        sudo apt-get update
+        packages=(pkg-config libfftw3-dev libopenmpi-dev openmpi-bin)
+        if ! command -v cmake >/dev/null; then packages+=(cmake); fi
+        sudo apt-get install -y --no-install-recommends "${packages[@]}"
+        uv sync --extra dev --extra native --frozen
     - name: Test with pytest
+      env:
+        OMPI_MCA_rmaps_base_oversubscribe: '1'
       run: uv run --locked pytest tests/test_vlasov1d --ignore=tests/test_vlasov1d/test_config_regression.py
 
+  test-native-vlasov1d:
+    needs: changes
+    if: needs.changes.outputs.run_all == 'true' || needs.changes.outputs.vlasov1d == 'true'
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v7
+    - uses: ./.github/actions/setup_repo
+      with:
+        python-version: '3.11'
+    - name: Install native validation dependencies
+      run: |
+        sudo apt-get update
+        packages=(pkg-config libfftw3-dev libopenmpi-dev openmpi-bin)
+        if ! command -v cmake >/dev/null; then packages+=(cmake); fi
+        sudo apt-get install -y --no-install-recommends "${packages[@]}"
+        uv sync --extra dev --extra native --frozen
+    - name: Test native backend
+      env:
+        OMPI_MCA_rmaps_base_oversubscribe: '1'
+      run: uv run --locked pytest tests/test_vlasov1d/test_native.py tests/test_vlasov1d/test_observations.py
+
   test-farsight1d:
     needs: changes
     if: needs.changes.outputs.run_all == 'true' || needs.changes.outputs.farsight1d == 'true'

adept/_base_.py

@@ -53,6 +53,8 @@ class ADEPTModule:
 
     """
 
+    jit_compile = True
+
     def __init__(self, cfg) -> None:
         self.cfg = cfg
         self.state = None
@@ -352,6 +354,10 @@ def _get_adept_module_(self, cfg: dict) -> ADEPTModule:
         else:
             raise NotImplementedError("This solver approach has not been implemented yet")
 
+        if cfg.get("backend", "jax") == "native":
+            from adept.vlasov1d import native_backend
+
+            this_module = native_backend(this_module)
         return this_module(cfg)
 
     def _setup_(self, cfg: dict, td: str, adept_module: ADEPTModule = None, log: bool = True) -> dict[str, Module]:
@@ -370,6 +376,10 @@ def _setup_(self, cfg: dict, td: str, adept_module: ADEPTModule = None, log: boo
         if adept_module is None:
             self.adept_module = self._get_adept_module_(cfg)
         else:
+            if cfg.get("backend", "jax") == "native":
+                from adept.vlasov1d import native_backend
+
+                adept_module = native_backend(adept_module)
             self.adept_module = adept_module(cfg)
 
         # dump raw config
@@ -401,11 +411,14 @@ def _setup_(self, cfg: dict, td: str, adept_module: ADEPTModule = None, log: boo
         self.adept_module.init_diffeqsolve()
         modules = self.adept_module.init_modules()
 
-        self._compat_execution, self.compatibility_fallback_reason = LegacyPreparedExecution.try_create(
-            original_cfg,
-            legacy_module=self.adept_module,
-            custom_module=adept_module is not None,
-        )
+        if not self.adept_module.jit_compile:
+            self.compatibility_fallback_reason = "native execution owns its worker lifecycle"
+        else:
+            self._compat_execution, self.compatibility_fallback_reason = LegacyPreparedExecution.try_create(
+                original_cfg,
+                legacy_module=self.adept_module,
+                custom_module=adept_module is not None,
+            )
         if self._compat_execution is not None:
             self.execution_backend = "prepared"
 
@@ -414,6 +427,9 @@ def _setup_(self, cfg: dict, td: str, adept_module: ADEPTModule = None, log: boo
         return modules
 
     def _execute_simulation(self, modules: dict | None, args: dict | None) -> dict:
+        if not self.adept_module.jit_compile:
+            self.execution_backend = "native"
+            return self.adept_module(modules, args)
         if self._compat_execution is not None:
             reason = self._compat_execution.fallback_reason(
                 state=self.adept_module.state,

adept/_vlasov1d/builder.py

@@ -83,6 +83,11 @@ def prepare(self, spec: SimulationSpec, *, key: int | jax.Array) -> PreparedSimu
             raise ValueError(f"Vlasov1DBuilder cannot prepare solver {spec.solver!r}")
         if spec.schema_version != "1":
             raise ValueError(f"Unsupported Vlasov1D specification schema version {spec.schema_version!r}")
+        if spec.config_dict().get("backend", "jax") == "native":
+            raise NotImplementedError(
+                "backend: native runs through ergoExo (including custom BaseVlasov1D subclasses); "
+                "Vlasov1DBuilder prepares a JAX program"
+            )
         if not jax.config.jax_enable_x64:
             raise RuntimeError("vlasov-1d requires float64; enable jax_enable_x64 before preparation")
         normalized_key, seed, provenance = normalize_key(key)

adept/_vlasov1d/datamodel.py

@@ -340,8 +340,29 @@ class Vlasov1DConfig(BaseModel):
     grid: GridConfig
     save: SaveConfig
     solver: str
+    backend: Literal["jax", "native"] = "jax"
+    native: "NativeSolverConfig | None" = None
     mlflow: MLFlowConfig
     drivers: EMDriverSetConfig
     terms: TermsConfig
     diagnostics: DiagnosticsConfig = DiagnosticsConfig()
     iaw_diagnostics: IAWDiagnosticsConfig | None = None
+
+
+class NativeSolverConfig(BaseModel):
+    """Execution options for the optional CUDA/MPI Vlasov backend."""
+
+    model_config = ConfigDict(extra="forbid")
+    library: str | None = None
+    device: Literal["cuda", "cpu"] = "cuda"
+    ranks: int = Field(default=1, ge=1)
+    launcher: list[str] = Field(default_factory=list)
+    workdir: str | None = None
+    halo_capacity: int = Field(default=64, ge=2)
+    scalar_stride: int = Field(default=1, ge=1)
+    custom_save_times: Literal["preserve", "steps"] = "preserve"
+    host_mpi: bool = False
+    timeout: float | None = Field(default=None, gt=0)
+
+
+Vlasov1DConfig.model_rebuild()

adept/_vlasov1d/native/CMakeLists.txt

@@ -0,0 +1,21 @@
+cmake_minimum_required(VERSION 3.24)
+project(adept_vlasov_native LANGUAGES CXX)
+option(VLASOV_CUDA "Use CUDA/cuFFT (OFF builds the FFTW validation backend)" ON)
+find_package(MPI REQUIRED COMPONENTS CXX)
+add_library(adept_vlasov SHARED api.cpp core/backend.cpp)
+target_compile_features(adept_vlasov PRIVATE cxx_std_17)
+target_link_libraries(adept_vlasov PRIVATE MPI::MPI_CXX)
+set_target_properties(adept_vlasov PROPERTIES POSITION_INDEPENDENT_CODE ON)
+if(VLASOV_CUDA)
+  enable_language(CUDA)
+  find_package(CUDAToolkit REQUIRED)
+  set_source_files_properties(core/backend.cpp PROPERTIES LANGUAGE CUDA)
+  set_target_properties(adept_vlasov PROPERTIES CUDA_STANDARD 17 CUDA_ARCHITECTURES 80)
+  target_compile_definitions(adept_vlasov PRIVATE VLASOV_CUDA)
+  target_compile_options(adept_vlasov PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:--extended-lambda>)
+  target_link_libraries(adept_vlasov PRIVATE CUDA::cudart CUDA::cufft)
+else()
+  find_package(PkgConfig REQUIRED)
+  pkg_check_modules(FFTW REQUIRED IMPORTED_TARGET fftw3)
+  target_link_libraries(adept_vlasov PRIVATE PkgConfig::FFTW)
+endif()

adept/_vlasov1d/native/README.md

@@ -0,0 +1,49 @@
+# Native Vlasov backend
+
+Select `backend: native` in the usual `solver: vlasov-1d` configuration and run
+through `ergoExo`. See the [configuration and build instructions](../../../docs/source/solvers/vlasov1d/config.md#execution-backend).
+An existing `BaseVlasov1D` subclass can still supply initialization, save
+callbacks and postprocessing. C++/CUDA/MPI advances the distribution; JAX runs
+the original save callbacks, without automatic differentiation of the solve.
+
+`core/` is the canonical numerical implementation. The standalone program in
+`benchmarks/vlasov_cuda` includes these same sources. Build the shared library
+on the target system with CMake; `VLASOV_CUDA=OFF` builds the FFTW validation
+implementation. The library is loaded only by MPI workers, not the outer
+Python driver.
+
+The outer driver performs initialization on the CPU and stages the initial
+state in bounded blocks. One persistent worker owns the solver on each GPU.
+At a save point, DLPack exposes synchronized native buffers as JAX arrays on
+a velocity-slab mesh. CUDA views require zero-copy import. The CPU validation
+build permits a copy when FFTW's host alignment is insufficient for JAX.
+Callbacks complete and saved outputs reach the host before the next native
+mutation. No borrowed state view may survive `av_close`; the library checks
+this lifetime rule explicitly.
+
+Workers stream saved leaves into host `.npy` memory maps. The driver assembles
+distributed saved outputs and returns the usual Diffrax `Solution` containing
+NumPy arrays. Existing xarray/NetCDF writers consume it. The module retains
+the backing workspace while its returned saves are in use. MLflow artifacts
+include actual save schedules, worker statistics/logs and the library hash.
+
+Standard saves round to actual timestep boundaries and merge duplicates.
+Scalars default to every step; `native.scalar_stride` can reduce the cadence.
+Custom callbacks preserve their requested times by default because some
+project postprocessors reconstruct those axes. Off-step custom saves require
+a previous-state copy and interpolation. `custom_save_times: steps` avoids
+that work when the postprocessor consumes actual result times. Postprocessors
+that reconstruct uniform scalar axes also need actual result times if the
+last scalar interval is shorter than the requested stride.
+
+The validated physics is spectral spatial streaming, ADEPT's local cubic C1
+velocity interpolation, sixth-order splitting, Poisson fields, and optional
+constant-rate Chang-Cooper Dougherty collisions. Kinetic ions use the existing
+`terms.species` entries, with independent velocity grids, charges and masses.
+Each stage sums every species' charge density into one MPI field reduction,
+then advances every species under that shared field. Collisions act on the
+reference species, as in ADEPT's JAX path; inter-species collisional exchange
+is not included. Rebuild the library when updating ADEPT (current ABI: 3).
+Unsupported physics and
+trainable modules fail explicitly. Performance and reproduction evidence is
+under [the benchmark results](../../../benchmarks/vlasov_cuda/results/2026-10-24-adept).

adept/_vlasov1d/native/__init__.py

@@ -0,0 +1 @@
+"""Optional native Vlasov timestepping with ADEPT preprocessing and saves."""