GPU Acceleration of Awkward Arrays: Using Python cuda.compute

arXiv:2610.08238v1 [cs.DC] 6 Oct 2026

GPU Acceleration of Awkward Arrays: cuda.compute Maksym Naumchyk1 1

,∗

and Ianna Osborne1

Using Python

,∗∗

Princeton University Abstract. Awkward Array is a widely used library in high-energy physics

(HEP) for representing and manipulating nested, variable-length data in Python. Previous CHEP contributions have explored GPU acceleration for Awkward Array, demonstrating the feasibility and performance benefits of CUDA-based backend while also identifying limitations related to irregular data access, finegrained kernel launches, and composability of operations. In this contribution, we present recent developments that build directly on these earlier efforts by introducing a CUDA execution model for Awkward Array based on the Python CUDA Core Compute Libraries (CCCL). Using CCCL, we eliminate the need for custom CUDA kernels and can instead use a high-level Python interface. The CCCL-based approach also enables fusion of multiple Awkward operations into a reduced number of CUDA kernels, addressing kernel launch overhead observed in earlier GPU implementations. Lazy execution allows expression graphs to be constructed and optimized prior to kernel generation, improving performance for analysis workflows involving jagged arrays, combinatorial operations, and reductions. In contrast to earlier approaches, this design also emphasizes extensibility, allowing user-defined Python code to be incorporated into GPU execution paths with minimal boilerplate and without breaking existing analysis semantics. We present performance studies that demonstrate improvements over previously reported eager GPU execution strategies for representative HEP analysis patterns. These developments extend the GPU capabilities of Awkward Array toward a more composable and sustainable backend, aligned with the needs of Python-based analysis at the HL-LHC and beyond. We thank the NVIDIA team for support and collaboration in integrating CCCL into Awkward, developing CUDA kernels and providing guidance on GPU optimization strategies. Their contributions are gratefully acknowledged.

1 Prologue Each of the different Awkward functions (sum, mean, etc.) consists of one or multiple kernels. Historically, these kernels were hand-written directly in CUDA C++. However, this implementation had its own drawbacks: raw CUDA is hard to write, hard to optimise, and hard to maintain. So the idea was to replace the existing kernels with cuda.compute — a Python interface to CUDA’s parallel algorithms. This approach would greatly reduce the codebase ∗ e-mail: maxymnaumchyk@gmail.com ∗∗ e-mail: Ianna.Osborne@cern.ch

Figure 1. Using non-rectangular arrays with Awkward

and make it much easier to maintain and develop new kernels. It would even bring performance improvements coming from using optimised cuda.compute algorithms and reducing kernel launches. The idea took shape at the ACAT workshop [1], where the limitations of the existing approach became the starting point for this work. The next two sections introduce Awkward itself and CUDA kernels in more detail. If you are already familiar with these concepts, jump straight to Section 4.

2 Introduction or is your data feeling awkward? Let us start from afar. Numbers. You have data you want to analyse. One of the most common tools to do that now is the NumPy library[2]. It lets you structure your data into arrays. For example, you can have an array like that: n u m p y _ a r r a y = np . a r r a y ( [ 0 . 0 , 1 . 1 , 2 . 2 , 3 . 3 , 4 . 4 , 5 . 5 , 6 . 6 , 7.7 , 8.8 , 9.9]) Real-world data is not always so simple. Take this sentence, for example: five words, each with three to eight letters. If we want to convert this information into arrays of numbers, we can get something like this: [ [ ’T ’ , [’t ’, [ ’s ’ , [ ’f ’ , [ ’e ’ ,

’a ’ , ’h ’ , ’e ’ , ’o ’ , ’x ’ ,

’k ’ , ’e ’ ] , ’i ’ , ’s ’] , ’n ’ , ’ t ’ , ’e ’ , ’n ’ , ’c ’ , ’e ’ ] , ’r ’] , ’ a ’ , ’m’ , ’ p ’ , ’ l ’ , ’ e ’ ]

] Now, if we try using NumPy to store this sentence, we will run into an error: V a l u e E r r o r : The r e q u e s t e d a r r a y h a s an i n h o m o g e n e o u s s h a p e What does it mean? Inhomogeneous shape. It means that each of our subarrays (words) has a different length. Such arrays are sometimes called non-rectangular. Some libraries call them jagged or ragged arrays [3]. We call them Awkward Arrays [4], and we have a library with the same name where you can use them! (See Figure 1.)

Figure 2. Inner structure of Awkward arrays

Figure 3. Finding a sum on axis 1 (for each subarray)

Let us take a more detailed look at how the Awkward arrays work from the inside. Internally, they are built from the same NumPy arrays. We flatten all of our data into a single array. This is marked as "values" in Figure 2. We also create one new NumPy array that stores information about the "raggedness" of the data — basically, the beginning and end index of each of our subarrays. We call them "offsets". So, from Figure 2: the first subarray has a start offset of "0" and an end offset of "3". The second array has both a start and end offset of "3", meaning that it is an empty array, and so on. As you can see, each subarray has its own pair of offsets. Now, an Awkward user does not need to know about such specifics. Why are we mentioning this here then? It will come in handy later on in this article. We loaded our data. Now we finally want to do something with it. Awkward has a huge variety of NumPy-like operations. You can calculate combinations, mean, argmax of arrays, and many more. Let us take a simple summation as an example. Using it is as easy as shown in Figure 3.

Figure 4. The whole ak.sum() operation using raw CUDA[5] kernels

As a result, we get the number "3" for the sum of all elements of the first subarray (0+1+2), zero for an empty array, and so on. Now, once again, let us look deeper into how the summation operation works inside Awkward.

3 CUDA kernels Figure 4 shows how Awkward handled a summation operation just a while back. You usually have a couple of C++ functions that parallelise their specific tasks. Here we have two such functions: awkward_reduce_sum_a and awkward_reduce_sum_b. Such functions are also called kernels. Almost all of the code in these kernels exists just to set things up. Computing thread positions, tracking subarray boundaries, and keeping threads in sync. Only one line at the end does the real accumulating work: atomicAdd (& t o p t r [ p a r e n t ] , temp [ t h r e a d _ i d ] ) ; This line adds each thread’s partial sum into the correct output slot. Using raw CUDA is cumbersome. You need to be an expert to write CUDA code. You need to constantly think about the optimisations that could take place. Maintaining and updating such code is also time-consuming. A new CUDA package can help us with that.

4 The cuda.compute package Recently, NVIDIA has been developing a new package called cuda.compute. It provides a simple Python interface to write CUDA kernels. Let us take our previous summation operation and rewrite it using cuda.compute (Figure 5). The code itself is short: we pass in the values array together with the start and end offsets (Section 2), an initial value h_init for the reduction (summation in our case), and call segmented_reduce. This is one of the algorithms cuda.compute provides. This maps so well onto Awkward arrays because of their structure. Recall from Section 2 that an Awkward array is really just a flat "values" array and a set of offsets, marking where each subarray begins and ends. A segmented reduction is built to do exactly this. Given a flat array and a set of segment boundaries, it reduces each segment independently and in parallel.

Figure 5. Implementation of a summation operation using cuda.compute

The "segments" that cuda.compute expects are precisely the subarrays that Awkward’s offsets already describe, so no restructuring or custom indexing logic is needed. This gives us two practical advantages. First, the whole reduction runs as a single kernel launch, rather than the several kernels that our old raw CUDA implementation needed. Second, the same segmented_reduce call can be reused for other reduction-like operations (mean, min, max, etc.). Just by changing the operator, instead of hand-writing a new kernel for each one. With that, we successfully rewrote the Awkward summation operation. Outside, we only use Python and cuda.compute algorithms. Inside, these algorithms use highly optimised C++ libraries: CUB and Thrust. This way we leave the development of new algorithms and all the optimisations to NVIDIA developers and concentrate only on the correct kernel behaviours.

5 Speed comparison In this section, let us explore how the cuda.compute-based implementation performs compared to the old raw CUDA kernels. Figure 6 shows the use of another Awkward operation (argmax) for arrays of different sizes. In terms of performance, the cuda.compute-based implementation is around 12 times faster than the raw CUDA kernels it replaces, and around 400 times faster than the equivalent implementation on a CPU. The bigger the arrays, the better the performance gain we get. This benchmark [6] was run on a personal computer with an Intel Core i7-12700 CPU and an RTX 3080 GPU. One of the biggest reasons for 12 times improvement is "kernel fusion". Figure 7 shows that the old implementation of the argmax operation needs three kernels, and the new cuda.compute implementation uses just one kernel (the same segmented_reduce we saw before). Every kernel launch has a fixed cost, but the bigger issue is memory. Each kernel has to write its results back to the GPU’s global memory before the next kernel can read them

Figure 6. Benchmarking the speed of ak.argmax() using arrays of different sizes

Figure 7. Benchmarking the ak.argmax() operation using NVIDIA Nsight Systems

in. Memory access is typically the slowest part of any GPU operation. With three separate kernels, this happens twice. Once between each pair of kernels. Fusing all three steps into a single segmented_reduce call removes both of these round trips, keeping intermediate values on-chip instead. Generally, when working with CUDA, the fewer kernels, the better. On top of that, segmented_reduce itself benefits from years of tuning inside CUB and Thrust. Careful use of shared memory, coalesced memory access and load balancing across thread blocks are built into cuda.compute. Let us perform some more benchmarks. We will take around 40 kernels and compare their speeds for two different implementations. For the tests, we will be using big sample arrays (several GBs). Figure 8 shows an improvement in speed of up to 15 times. The cause of this is the same as described above. Fewer kernel launches and better-optimised memory access. For more benchmarks using actual physics analysis, see the work of Kelci Mohrman et al. [7] or the SciPy blog from Ianna Osborne and Ashwin Srinath [8].

Figure 8. Comparing the speed of 40 different kernels.

6 Kernel development process At the time of writing this article, around 90% of all the Awkward kernels are already migrated. For the remaining kernels, using raw CUDA is still beneficial. cupy . RawKernel ( ) The most prominent examples of such cases are the kernels that raise errors. cuda.compute’s algorithms have no built-in way for a thread to report a failure back to the caller. Awkward’s raw CUDA kernels handle this with a shared error flag instead. It is also worth noting that, in some cases, using CuPy alone can be beneficial. def awkward_ListArray_rpad_and_clip_length_axis1 ( tomin , f r o m s t a r t s , f r o m s t o p s , t a r g e t , l e n s t a r t s ): counts = fromstops [ : l e n s t a r t s ] − fromstarts [ : l e n s t a r t s ] t o m i n [ : 1 ] = cp . sum ( cp . maximum ( c o u n t s , t a r g e t ) ) . a s t y p e ( t o m i n . d t y p e , copy= F a l s e ) This is just an example of one of the kernels that Awkward uses. The exact implementation details are not important here. Just note that different cupy operations are used here, such as cp.sum(), cp.maximum(), and slicing. The cuda.compute package is still new and is currently being actively developed. That is why some workarounds like these can take place now. We (the Awkward library team) are working very closely with the NVIDIA team to address any bugs or performance issues we encounter. Speaking of which, CUDA Python (a parent package of cuda.compute) just had its first release [9] (at the time of writing this article).

7 Next steps – improving performance In this section, we will explore how to continue improving the GPU performance of Awkward Arrays. • First, we will extensively benchmark the implementations of new kernels. Validation by Awkward users is also needed to make sure no calculations are affected by the changes. Based on these results, we can identify the existing bottlenecks and work on them. • Right now, all the cuda.compute kernels are JIT (just-in-time) compiled. We will be integrating ahead-of-time compilation when this gets officially supported.

• In the example that was used in this article (an implementation of summation), we used only one segmented_reduce operation. More are required for some of the more complex Awkward functions. Going one step further, we can fuse even the cuda.compute operations themselves — for example, using Parrot [10], another NVIDIA library. Since it is still in early development, other solutions are possible, such as lazy execution built in directly inside Awkward.

8 Acknowledgments This work was supported by the National Science Foundation under Cooperative Agreement PHY-2323298 (IRIS-HEP). On behalf of the Awkward team, we thank the NVIDIA team (Jake Hemstad, Ashwin Srinath, et al.) for their help and support with any obstacles we encountered during this project. Thanks to my mentor, Ianna Osborne, for helping and advising me on this project. Thanks to the IRIS-HEP grant for sponsoring me on this journey, giving me an opportunity to work with amazing people of the HEP community and to present the results of my work at the 28th Conference on Computing in High Energy and Nuclear Physics (CHEP 2026) in Bangkok.

References [1] Ianna Osborne. "Advancing Awkward Arrays for High-Performance CPU and GPU Processing." Advanced Computing and Analysis Techniques in Physics Research (ACAT 2025), Hamburg, 2025, https://indico.cern.ch/event/1488410/contributions/6562844/. [2] Harris, Charles R., et al. "Array Programming with NumPy." Nature, vol. 585, no. 7825, 2020, pp. 357–362, https://arxiv.org/abs/2006.10256. [3] "Ragged Tensors." TensorFlow, https://www.tensorflow.org/guide/ragged_tensor. [4] "Awkward Array." Awkward Array Documentation, https://awkward-array.org/doc/ main/. [5] Awkward Array Team. "Raw CUDA Kernel Implementation of ak.sum() Operation." GitHub, https://github.com/scikit-hep/awkward/blob/v2.9.0/src/awkward/_connect/cuda/ cuda_kernels/awkward_reduce_sum.cu. [6] Awkward Array Team. "Benchmarking for an argmax Operation." GitHub, https://github. com/scikit-hep/awkward/pull/3800. [7] Mohrman, Kelci Ann, et al. "GPU Acceleration of End-User Analyses at the LHC." 28th Conference on Computing in High Energy and Nuclear Physics (CHEP 2026), Bangkok, 2026, https://indico.cern.ch/event/1471803/contributions/6968223/. [8] Osborne, Ianna, and Ashwin Srinath. "Rewriting Awkward Array’s GPU Kernels in Python with NVIDIA’s cuda.compute." Scientific Python Blog, 26 Aug. 2026, https: //blog.scientific-python.org/awkward/rewriting-gpu-kernels-cuda-compute/. [9] Bentz, Jonathan. "NVIDIA CUDA 13.3 Enhances GPU Development with Tile Programming in C++, Compiler Autotuning, and Python Updates." NVIDIA Technical Blog, 26 May 2026, https://developer.nvidia.com/blog/nvidia-cuda-13-3-enhances-gpudevelopment-with-tile-programming-in-c-compiler-autotuning-and-python-updates. [10] NVIDIA. "Parrot Documentation." https://nvlabs.github.io/parrot/.

Related documents

• GPU Acceleration of Awkward Arrays: Using Python cuda.compute#1217973
• Taming Bitwise Behavior in GPU Kernels with Tensor Core: Black-Box Reconstruction, Compiler Enforcement, and Static Verification#673469
• FastPair: GPU-Optimized String Decoding#1084928
• FastPair: GPU-Optimized String Decoding#919348
• PAGE404: Breaking GPU Isolation Through Hardware-assisted Page Migration#1356298