Accelerating Python: A Guide to GPU Optimization with CuPy and Numba CUDA

Author: Marco Celorio (CINECA)

Abstract

Python has become a leading language for scientific computing thanks to its read ability, flexibility, and rich ecosystem of numerical libraries. However, exploiting the full performance potential of modern graphics processing unit (GPU) architectures requires more than conventional Python execution. Effective acceleration typically relies on a hybrid programming model, where Python manages high-level workflow, data movement, and application logic, while computationally intensive kernels are offloaded to GPU via
frameworks such as NVIDIA’s Compute Unified Device Architecture (CUDA) for massively parallel execution.

This work presents a comprehensive, engineering-focused workflow for accelerating Python applications on NVIDIA GPU systems using CuPy, Numba CUDA, CuPy RawKernel/RawModule, CUDA-aware MPI, and NCCL. The approach is demonstrated through two systematic case studies. The first, dense matrix multiplication, scales from educational custom kernels and shared-memory tiling to production-grade CuPy interfaces. The second, a 2D Laplace/Jacobi solver, progresses from a baseline NumPy reference through vectorized CuPy, fused CuPy, Numba CUDA, and custom CuPy RawKernel implementations, culminating in distributed multi-GPU implementations based on CUDA-aware MPI and NCCL.

Together, the case studies contrast a compute-intensive library-suited operation with a memory-bandwidth-limited iterative stencil, illustrating how the appropriate optimization strategy depends on the workload.

Rather than offering simple application programming interface (API) coverage, this guide emphasizes practical software engineering and high-performance computing (HPC) principles, including GPU data residency, asynchronous execution timing, distributed halo exchanges, and NVTX-driven profiling