Thousands of lines of hand-written CUDA C++, now Python. Less code, and it runs faster.
A single collision event in a particle detector holds a variable number of particles, each with a variable number of measurements:
[[1.1, 2.2, 3.3], [], [4.4, 5.5]]
Awkward Array is a Python library for manipulating nested, variable-length (“ragged”) data like this with NumPy-like idioms. It stores that data flat, as one content buffer holding every value contiguously plus an offsets array marking where each sublist begins and ends:
content: [1.1, 2.2, 3.3, 4.4, 5.5]
offsets: [0, 3, 3, 5] # the empty middle list spans no elements
Nothing is wasted on padding, but every operation must then be written in terms of those two buffers rather than a simple shape. That is where dense GPU tensor frameworks stop helping: they want rectangles. For several years, Awkward’s answer was a dictionary of hand-written CUDA C++ kernels compiled at runtime with CuPy.
The Awkward Array and NVIDIA teams have now rebuilt that layer on cuda.compute, which brings the CUDA C++ parallel-algorithm libraries CUB and Thrust (the reductions, scans, and sorts that power production GPU software) into Python as ordinary callables. Instead of writing a kernel, the backend now composes ones that already exist. Four results stand out.
ak.min, the minimum over each ragged sublist, took three kernel launches in CUDA C++: initialize a scratch buffer, reduce within each block using shared memory and explicit thread synchronization, and copy the result out. Each communicated with the next through global memory.
It is now a single call to a library primitive:
import cupy as cp, numpy as np
from cuda.compute import OpKind, segmented_reduce
def awkward_reduce_min(
toptr: cp.ndarray,
fromptr: cp.ndarray,
offsets: cp.ndarray,
outlength: int,
identity: float,
) -> None:
toptr[:outlength] = identity
segmented_reduce(
fromptr,
toptr,
offsets[:-1],
offsets[1:],
OpKind.MINIMUM,
np.asarray(identity, dtype=fromptr.dtype),
outlength,
)segmented_reduce performs one reduction per segment, and offsets[:-1] and offsets[1:] hand it the start and end of each sublist directly. No CUDA C++, no synchronization, no scratch buffers: the library handles all of it.
Awkward’s user-facing API is unchanged. Underneath, 2.10.0 routes 106 of its 133 GPU kernels (80%) through cuda.compute, up from none in 2.8.11. Every reduction and the sort now run through it with no hand-written implementation remaining. The 27 still in CUDA C++ handle structural work like jagged indexing, padding, and validity checks. Their control flow depends on the ragged layout itself, which fits the segmented primitives less naturally. That migration is ongoing. Counting only code that must actually be maintained, the switch to cuda.compute has so far cut hand-written CUDA C++ from 8,288 lines to 2,170 — a 74% reduction.
We might expect to pay something for the abstraction.
Over 5,000,000 ragged sublists, ak.argmin takes 0.96 ms per call as a hand-written CUDA kernel and 0.39 ms through cuda.compute: 2.5x faster, with identical output.
That speedup is inherited rather than hand-tuned, which is exactly the point of building on CUB: segmented reductions are difficult to write well by hand, and CUB’s have been tuned per architecture for years. Awkward gets that tuning now, and the next architecture’s when it ships, without changing its own code.
The first two results came from the migration. For this one, a user drops down to cuda.compute and composes the primitives by hand.
Awkward evaluates eagerly: every operation returns a real array, so a chain of them writes an intermediate to global memory at each step and reads it back at the next. cuda.compute algorithms instead accept iterators that are evaluated lazily as the algorithm runs, letting many logical steps ride along inside a single pass. That is kernel fusion, and it saves both the memory traffic and the launch overhead.
The opposite-sign di-muon invariant mass, a standard reconstruction in particle physics, combines a few measured quantities for every pair of particles in an event. In Awkward, it is one line:
mu1, mu2 = ak.unzip(ak.combinations(muons, 2))
mass = np.sqrt(
2 * mu1.pt * mu2.pt * (np.cosh(mu1.eta - mu2.eta) - np.cos(mu1.phi - mu2.phi))
)Evaluated step by step, that chain of arithmetic and trigonometric operations becomes one or more kernels per step, with every intermediate written out as a full-length array and read back.
Written by hand as one cuda.compute call, the whole formula becomes a single operator: a gpu_struct keeps each particle’s fields together, a ZipIterator combines them, and a PermutationIterator produces each pair on demand. The operator sees one complete pair at a time, and no intermediate is ever built.
Over all such pairs in a CMS open-data sample:
| kernel launches | memory operations | GPU time | |
|---|---|---|---|
| step by step | 88 | 212 | 25.7 ms |
| fused into one call | 1 | 45 | 10.2 ms |
2.5x faster, with nothing allocated in between.
We ran the ADL benchmark queries (a standard set of physics-analysis tasks) on CMS 2012 open data, against released Awkward 2.8.11 on its hand-written CuPy backend, the last version before cuda.compute. For the two combinatoric queries below, the Awkward expression is unchanged; only the backend underneath differs, so the comparison isolates the migration itself.
Measured on the GPU compute stage at 100k, 1M, and 10M events, both speed up by margins that grow with the data:
cuda.compute’s time stays approximately constant across those sizes while the hand-written implementation grows super-linearly, so the gap widens as data grows rather than closing.
A stated goal of the Awkward Array project is to let physicists and data analysts write high-performance code in Python without GPU expertise. The old backend required contributors to understand CUDA thread hierarchies, atomics, and shared-memory behavior before they could add or fix a kernel. The new one asks for an ordinary Python function and a call to the right primitive. Domain scientists can read it, review it, and unit-test its logic without a GPU.
Awkward knows the problem. cuda.compute knows the hardware. The result is Python that’s simpler and faster than the CUDA C++ it replaced.
Full methodology, per-query results, and the code-counting rules are in “GPU-Accelerated Awkward Arrays with CUDA Python” by Ashwin Srinath (NVIDIA) and Ianna Osborne (Princeton University), Proceedings of the 24th Python in Science Conference (SciPy 2026). All measurements were taken on an NVIDIA RTX PRO 6000 Blackwell Server Edition with CUDA 13.2, cuda.compute 1.1.0, and CuPy 14.1.1; every benchmark was run twice on independent machines, with structural counts identical and timings agreeing to within a few percent. Migration progress is tracked in scikit-hep/awkward#3793.
Much of the kernel migration was implemented by Maxym Naumchyk. Thanks also to the cuda.compute and CUB/Thrust developers at NVIDIA and to the Scikit-HEP community. This work was supported in part by NSF grants OAC-1450377, OAC-1836650, OAC-2103945, PHY-2121686, and PHY-2323298.
TL;DR: If you have GPU code in your project, setup a GitHub hosted GPU runner today. It is fairly quick to do and will free you from having to run tests manually.
Writing automated tests for your code base and certainly for the more complex parts of it has become as normal as brushing your teeth in the morning. Having a system that automatically runs a project’s tests for every Pull Request is completely normal. However, until recently it was very complex and expensive to setup a system that can run tests on a system with a GPU. This means that, when dealing with GPU related code, we were thrown back into the dark ages where you had to rely on manual testing.
In this blog post I will describe how we set up a GitHub Action based GPU runner for the scikit-learn project and the things we learnt along the way. The goal is to give you some additional information and details about the setup we now use.
All workflows for your GitHub project are executed on a runner. Normally all your workflows run on the default runner, but you can have additional runners too. If you wanted to you could host a runner yourself on your own infrastructure. Until now this was the only way to get access to a runner with a GPU. However, hosting your own runner is complicated and comes with pitfalls regarding security.
Since about April 2024 GitHub has made larger runners with a GPU generally available.
To use these you will have to setup a credit card for your organisation. Configure a spending limit so that you do not end up getting surprised with a very large bill. For scikit-learn we currently use a limit of $50.
When adding a new GitHub hosted runner make sure to select the “Partner” tab when choosing the VM’s image. You need to select the “NVIDIA GPU-Optimized Image for AI and HPC” image in order to be able to choose the GPU runner later on.
The group the runner is assigned to can be configured to only allow particular repositories and workflows to use the runner group. It makes sense to only enable the runner group for the repository in which you plan to use it. Limiting which workflows your runner will pick up requires an additional level of indirection in your workflow setup, so I will not cover it in this blog post.
Name your runner group cuda-gpu-runner-group to match the name used in the examples
below.
The GPU runner uses a disk image provided by NVIDIA. This means that there are some differences to the image that the default runner uses.
The gh command-line utility is not installed by default. Keep this in mind
if you want to do things like removing a label from the Pull Request or
other such tasks.
The biggest difference to the standard image is that the GPU image contains a conda installation, but the file permissions do not allow the workflow user to modify the existing environment or create new environments. As a result for scikit-learn we install conda a second time via miniforge. The conda environment is created from a lockfile, so we do not need to run the dependency solver.
A key difference between the GPU runner and the default runner is that a project has to pay for the time of the GPU runner. This means that you might want to execute your GPU workflow only for some Pull Requests instead of all of them.
The GPU available in the runner is not very powerful, this means it is not that attractive of a target for people who are looking to abuse free GPU resources. Nevertheless, once in a while someone might try. Another reason to not run the GPU workflow by default.
A nice way to deal with running the workflow only after some form of human review
is to use a label. To mark a Pull Request (PR) for execution on the GPU runner a
reviewer applies a particular label. Applying a label does not cause a notification
to be sent to all PR participants, unlike using a special comment to trigger the
workflow.
In the following example the CUDA CI label is used to mark a PR for execution and
the runs-on directive is used to select the GPU runner. This is a snippet from
the full GPU workflow used in the scikit-learn repository.
name: CUDA GPU
on:
pull_request:
types:
- labeled
jobs:
tests:
if: contains(github.event.pull_request.labels.*.name, 'CUDA CI')
runs-on:
group: cuda-gpu-runner-group
steps:
- uses: actions/setup-python@v5
with:
python-version: '3.12.3'
- name: Checkout main repository
uses: actions/checkout@v4
...
In order to remove the label again we need a workflow with elevated permissions. It needs to be able to edit a Pull Request. This privilege is not available for workflows triggered from Pull Requests from forks. Instead the workflow has to run in the context of the main repository and should only do the minimum amount of work.
on:
# Using `pull_request_target` gives us the possibility to get a API token
# with write permissions
pull_request_target:
types:
- labeled
# In order to remove the "CUDA CI" label we need to have write permissions for PRs
permissions:
pull-requests: write
jobs:
label-remover:
if: contains(github.event.pull_request.labels.*.name, 'CUDA CI')
runs-on: ubuntu-20.04
steps:
- uses: actions-ecosystem/action-remove-labels@v1
with:
labels: CUDA CI
This snippet is from the label remover workflow we use in scikit-learn.
For scikit-learn we have been using the GPU runner for about six weeks. So far we have stayed below the $50 monthly spending limit we set. This includes some runs to debug the setup at the start.
One of the scikit-learn contributors created a Colab notebook that people can use to setup and run the scikit-learn test suite on Colab. This is useful for contributors who do not have easy access to a GPU. They can test their changes or debug failures without having to wait for a maintainer to label the Pull Request. We plan to add a workflow that comments on PRs with information on how to use this notebook to increase its discoverability.
Overall it was not too difficult to setup the GPU runner. It took a little bit of fiddling to deal with the differences in VM image content as well as a few iterations for how to setup the workflow, given we wanted to manually trigger them.
The GPU runner has been reliably working and picking up work when requested. It saves us (the maintainers) a lot of time, as we do not have to checkout a PR locally and run the tests by hand.
The costs so far have been manageable and it has been worth spending the money as it removes a repetitive and tedious manual task from the reviewing workflow. However, it does require having the funds and a credit card.
]]>