PyKokkos: Python Interface for Kokkos

• PyKokkos is a Python interface for Kokkos that provides performance portability by translating Python kernels into backend-specific C++ code via pybind11.
• It supports rapid prototyping and diverse computational kernels—from structured loops to irregular particle and grid operations—achieving notable speedups.
• PyKokkos integrates with Kokkos abstractions to enable single-source code execution on heterogeneous architectures while managing synchronization and memory optimizations.

PyKokkos is a Python interface for the Kokkos C++ parallel programming framework that has been used in recent high-performance computing work as a performance-portable kernel layer for heterogeneous systems. In the published applications that discuss it directly, PyKokkos translates Python kernel descriptions into Kokkos/C++ through pybind11, memoizes compiled kernels after the first invocation, and targets OpenMP on CPUs, CUDA on NVIDIA GPUs, and HIP on AMD GPUs. Its documented role is therefore dual: it supports rapid algorithmic exploration in Python while also acting as the portability mechanism that allows a single kernel code base to execute on multiple architectures without vendor-specific rewrites (Almgren-Bell et al., 9 Aug 2025, Kosmacher et al., 17 Jun 2026).

1. Software model and execution pathway

In the available literature, PyKokkos appears as a Python-front-end realization of Kokkos abstractions rather than as an independent execution model. In the collisional-plasma solver, particle state is stored in CuPy arrays, with positions and velocities in an $N \times 6$ array and particle weights, source terms, and energies in an $N \times 4$ array, while PyKokkos supplies the GPU kernel abstraction over those arrays. In the fast Ewald summation code, the implementation is likewise in Python and combines CuPy for arrays with PyKokkos to invoke Kokkos kernels. Both accounts describe a workflow in which Python manages data and experimentation while compiled kernels are lowered to backend-specific execution spaces (Almgren-Bell et al., 9 Aug 2025, Kosmacher et al., 17 Jun 2026).

This organization makes PyKokkos significant for performance portability rather than merely for language convenience. The plasma study presents it as the reason that the exact same code base can run on both CUDA and HIP backends, and the Ewald study describes it as enabling Kokkos in Python “with minimal overheads using just-in-time compiling.” A plausible implication is that PyKokkos is most naturally understood as a single-source kernel interface whose operational semantics are inherited from Kokkos, while Python remains the orchestration environment (Almgren-Bell et al., 9 Aug 2025, Kosmacher et al., 17 Jun 2026).

2. Kokkos abstractions exposed through Python

The clearest published account of PyKokkos’ abstraction layer comes from the Ewald-summation work. There, the programming model is described by the mapping GPU block $\rightarrow$ Kokkos team, shared memory $\rightarrow$ Kokkos scratchpad memory, 2D thread rows $\rightarrow$ Kokkos vector ranges, row reductions $\rightarrow$ Kokkos vector reductions, and CPU teams $\rightarrow$ OpenMP threads. On CPUs, the authors set block size to one and do not use explicit SIMD vectorization, instead relying on the compiler (Kosmacher et al., 17 Jun 2026).

The plasma application complements this by showing PyKokkos in kernels that are intentionally irregular and synchronization-heavy. These include particle creation through a single atomic write index shared by all threads, per-cell candidate-list construction with atomics, cell-local binary pairing for Coulombic collisions, and atomic segmented reductions for particle-to-cell accumulation. Taken together, these reports indicate that PyKokkos is not confined to regular dense loop nests; it is used for hierarchical parallelism, scratchpad-memory staging, atomics, reductions, and per-cell irregularity, which are precisely the patterns for which Kokkos-style abstractions are typically invoked (Almgren-Bell et al., 9 Aug 2025, Kosmacher et al., 17 Jun 2026).

3. Particle kernels in collisional plasma simulation

In "A Portable Multi-GPU Solver for Collisional Plasmas with Coulombic Interactions" (Almgren-Bell et al., 9 Aug 2025), PyKokkos is the main Python-based performance-portable implementation layer for the particle kernels of a multi-GPU particle-in-cell solver for low-temperature plasmas. The broader solver advances a kinetic electron distribution with DSMC/PIC methods, using a kinetic description for electrons and a fluid approximation for heavy species. The paper states that “for every Numba or CUDA GPU kernel, we also have a PyKokkos implementation.” The covered kernels include the DSMC collision kernel for electron-heavy collisions, the recombination kernel for the three-body reaction $C5$, the particle-to-cell reduction kernels, and the Coulombic collision kernel for electron-electron interactions.

For the DSMC collision kernel, each particle evaluates collision probabilities against species-specific cross sections using

$$\pi_{lck}=1-\exp(-\Delta t\, \|v_l\|_2\, \sigma_{ck}(\|v_l\|_2)\, n_k),$$

then decides whether a collision occurs, updates velocities by conservation laws, possibly creates a new particle for ionization, advects the particle, and applies boundary conditions. Recombination $C5$ is implemented as two kernels: first, build per-cell catalyte lists with atomics; second, match primaries to catalytes and update the particle data. The Coulombic kernel models electron-electron collisions with a Takizuka–Abe style binary-collision method by grouping particles by cell, pairing them within each cell, and applying collision updates. The particle-to-cell kernel uses atomic segmented reductions to compute cell-wise quantities such as electron density and electron temperature (Almgren-Bell et al., 9 Aug 2025).

A central implementation result in this work is that atomics are not treated as a last resort but as a deliberate algorithmic choice. For the DSMC collision kernel, particle creation uses a single atomic write index into preallocated buffers, which is reported to yield up to a $N \times 4$0 speedup. For the particle-to-cell kernel, the atomic-reduction strategy makes both the CUDA and PyKokkos versions $N \times 4$1–$N \times 4$2 faster than the non-atomic alternative for problem sizes from $N \times 4$3 to $N \times 4$4 particles (Almgren-Bell et al., 9 Aug 2025).

4. PyKokkos in fast Ewald summation for Stokes flow

In "A performance portable fast Ewald summation for Stokes flow" (Kosmacher et al., 17 Jun 2026), PyKokkos serves as the implementation substrate for custom near-field and far-field kernels in a periodic Stokes solver. The paper uses the standard Ewald decomposition

$N \times 4$5

with near-field particle-to-particle interactions and a far-field pipeline

$N \times 4$6

All interactions and their variants are implemented using Kokkos abstractions via PyKokkos, specifically P2P, P2G, G2P, and the diagonal Fourier convolution. FFT and IFFT are not implemented with PyKokkos; instead, the code uses vendor libraries for single- and multi-device transforms (Kosmacher et al., 17 Jun 2026).

The P2P kernel uses cell lists and several variants organized around assigning a target cell to a GPU block or Kokkos team, with shared-memory versions using Kokkos scratchpad to stage source data. The most prominent methodological contribution concerns P2G. A baseline P2G loops over source particles, computes contributions to all $N \times 4$7 neighboring grid points, and uses atomicAdd for each grid write, so atomic contention becomes the bottleneck. The paper studies four P2G variants—P2G-base, P2G-source, P2G-grid, and P2G-hybrid—with the hybrid algorithm combining source ordering, shared-memory reuse, and reduced atomic frequency. In P2G-hybrid, the source cell is assigned to a GPU block, source data and per-source window coefficients are loaded into shared memory, threads are assigned grid points in the local colleague region, each thread accumulates locally over sources, and final writes use atomics only once per grid-point aggregate. At $N \times 4$8, tolerance $N \times 4$9, and optimal $\rightarrow$0, the paper reports that P2G-hybrid achieves a $\rightarrow$1 speedup over P2G-base (Kosmacher et al., 17 Jun 2026).

5. Portability, performance, and scaling

The plasma study frames PyKokkos as central to portability across NVIDIA and AMD GPUs. The same particle-kernel code runs on NVIDIA Volta V100 and AMD MI250X systems, enabling direct comparison of portability and performance. The authors report that the MI250X is slightly faster for most kernels but shows more sensitivity to register pressure. For the DSMC collision kernel and the Coulombic kernel, the analysis emphasizes memory traffic, register pressure, and atomic overheads; the cross-platform difference is attributed primarily to register pressure rather than atomics, because removing atomics does not materially change the V100/MI250X performance ratio whereas reducing register usage does. On V100 GPUs, distributed-memory scaling remains reasonable when the particle-to-cell ratio is large and the solver uses subcycling: with $\rightarrow$2, efficiency remains about $\rightarrow$3 at 8 GPUs and $\rightarrow$4 at 16 GPUs, whereas for $\rightarrow$5 it drops more sharply because the replicated mesh and MPI communication become more dominant (Almgren-Bell et al., 9 Aug 2025).

The Ewald study extends the performance picture beyond GPUs to CPUs. Its P2P interaction achieves around $\rightarrow$6 compute efficiency on NVIDIA H200, $\rightarrow$7 on NVIDIA A100, $\rightarrow$8 on AMD MI300, $\rightarrow$9 on Grace CPU, and $\rightarrow$0 on AMD Epyc CPU. For the full code path, the overall Ewald solver processes approximately 8 million particles per second on a H200 GPU and about a half-million particles per second on a Grace CPU, for nine digits of accuracy. The distributed implementation performs a weak-scaling test on up to 256 million particles over 64 GPUs and reports bounded communication cost for all stages except the all-to-all particle sorting step, which the authors note can be reduced to neighbor communication in the relevant time-stepping regime (Kosmacher et al., 17 Jun 2026).

Across both applications, PyKokkos is therefore presented not as merely “good enough” but as competitive for custom kernels, including irregular particle algorithms, while preserving a single code path across CUDA, HIP, and OpenMP targets (Almgren-Bell et al., 9 Aug 2025, Kosmacher et al., 17 Jun 2026).

6. Position within the Kokkos ecosystem and recurring tradeoffs

A recurrent misconception is to treat PyKokkos as a thin convenience wrapper whose primary value is programmer productivity. The plasma paper rejects that framing explicitly: PyKokkos is “not just a convenience wrapper” but the portability mechanism of the portable multi-GPU solver, and its kernels are reported as competitive and in some cases comparable to CUDA versions. The Ewald paper shows a complementary division of labor: PyKokkos implements the custom interaction kernels, while FFT and IFFT remain delegated to vendor libraries. This suggests a characteristic deployment pattern in which PyKokkos covers application-specific kernels, and highly specialized transforms or communication-heavy library routines remain external (Almgren-Bell et al., 9 Aug 2025, Kosmacher et al., 17 Jun 2026).

That pattern aligns with the broader Kokkos literature. In HEP track-reconstruction mini-applications, Kokkos is described as a template-library portability solution whose backend switching is “relatively seamless,” with close-to-native GPU performance on NVIDIA and reasonable performance on AMD GPUs (Kwok et al., 2024). In structured-grid and AMR codes such as K-Athena and AthenaK, Kokkos supports a single-source code base across CPUs and GPUs, but those papers are explicit that portable performance is not automatic: execution-policy choice, memory layout, compiler vectorization behavior, memory traffic, and code structure remain decisive determinants of realized performance (Grete et al., 2019, Stone et al., 2024). PETSc extends the same idea at the solver-library level by separating the application programming model from the library’s internal implementation and by exposing zero-copy Kokkos views for callbacks, which suggests a plausible integration path for PyKokkos-style kernels inside larger solver ecosystems (Mills et al., 2020).

Within that landscape, PyKokkos occupies a specific niche. It brings Kokkos’ execution-space and memory-space abstractions into Python for kernels that remain performance-critical, irregular, or architecture-sensitive, while preserving rapid prototyping and a single kernel description across heterogeneous backends. The published evidence to date is concentrated in particle, particle-grid, and reduction-heavy workloads, where synchronization, register pressure, and backend-specific hardware behavior remain central technical concerns (Almgren-Bell et al., 9 Aug 2025, Kosmacher et al., 17 Jun 2026).