Python GPU programming: CuPy vs Numba vs PyCUDA

Prerequisites

This article assumes you have read Your first CUDA program, including the host-device execution model, the <<<blocks, threads>>> launch syntax, and the bounds-check pattern for mapping threads to array elements. That article already showed a minimal Numba @cuda.jit vector addition; this one goes much further, covering CuPy and PyCUDA as well, and the operational details (compilation strategy, streams, error handling, profiling, interop) that a single toy example glosses over.

You do not need to have written C++ before. Every example here is Python.

Why Python at all

CUDA C++ gives you full control, but most day-to-day GPU work in data science and machine learning never touches nvcc directly. Three Python libraries expose the CUDA programming model at different levels of abstraction, and picking the right one for a given task matters as much as picking CUDA over plain Python in the first place.

graph TD
  A["CUDA C++
(nvcc, full control)"] --> B["PyCUDA
raw CUDA C strings
compiled via SourceModule"]
  B --> C["Numba @cuda.jit
Python subset
compiled to PTX via NVVM"]
  C --> D["CuPy
NumPy-compatible arrays
library calls + RawKernel"]
  style A fill:#e8744f,color:#fff
  style B fill:#e8744f,color:#fff
  style C fill:#4a90d9,color:#fff
  style D fill:#00CC96,color:#fff

Moving down this ladder trades control for productivity. CuPy at the top of the ladder (in terms of abstraction, drawn at the bottom of the diagram) lets you swap numpy for cupy and get GPU-accelerated array operations with almost no code changes. Numba lets you write a kernel body in a restricted subset of Python and get it JIT-compiled to run on the device. PyCUDA hands you the closest thing to raw CUDA C++ available from Python: you write actual CUDA C source as a string and compile it at runtime.

None of these libraries replace the mental model from the rest of this series. Every one of them still allocates device memory, transfers data across PCIe, launches a grid of threads, and requires you to reason about thread indices and bounds checks. They differ in how much of that machinery they hide from you and how the kernel source code itself is written.

The same program, three ways: vector addition

To compare the three libraries directly, here is c[i] = a[i] + b[i] for one million elements, implemented identically in each. Pay attention to what each version makes explicit versus implicit: allocation, host-to-device transfer, kernel launch configuration, synchronization, and result validation.

CuPy: NumPy-compatible arrays

import cupy as cp
import numpy as np

n = 1_000_000

# Allocation + host-to-device transfer happen together: creating a cupy
# array from a numpy array copies it to the device immediately.
h_a = np.arange(n, dtype=np.float32)
h_b = (np.arange(n, dtype=np.float32) * 2)

d_a = cp.asarray(h_a)   # host -> device copy
d_b = cp.asarray(h_b)

# No explicit kernel launch: cp.add dispatches to a prebuilt elementwise
# CUDA kernel that CuPy ships with the library.
d_c = cp.add(d_a, d_b)

cp.cuda.Stream.null.synchronize()   # wait for the default stream to finish

h_c = cp.asnumpy(d_c)                # device -> host copy
expected = h_a + h_b
assert np.allclose(h_c, expected)
print(f"All {n} elements correct via CuPy.")

There is no threadsPerBlock, no blockIdx.x * blockDim.x + threadIdx.x, and no bounds check in your code. CuPy’s built-in elementwise kernels already handle launch configuration and bounds internally. This is the entire point of the library: for operations CuPy already implements (arithmetic, reductions, linear algebra, FFTs, sorting, and more), you write array expressions, not kernels.

Numba: JIT-compiled Python kernels

import numpy as np
from numba import cuda

@cuda.jit                            # compiled to PTX the first time it is called
def vec_add_kernel(a, b, c):
    i = cuda.grid(1)                 # blockIdx.x * blockDim.x + threadIdx.x
    if i < a.size:                   # explicit bounds check, same as CUDA C
        c[i] = a[i] + b[i]

n = 1_000_000
h_a = np.arange(n, dtype=np.float32)
h_b = (np.arange(n, dtype=np.float32) * 2)

d_a = cuda.to_device(h_a)            # explicit host -> device transfer
d_b = cuda.to_device(h_b)
d_c = cuda.device_array(n, dtype=np.float32)

threads_per_block = 256
blocks = (n + threads_per_block - 1) // threads_per_block  # explicit launch config
vec_add_kernel[blocks, threads_per_block](d_a, d_b, d_c)

cuda.synchronize()                    # explicit device-wide sync

h_c = d_c.copy_to_host()             # explicit device -> host transfer
expected = h_a + h_b
assert np.allclose(h_c, expected)
print(f"All {n} elements correct via Numba.")

Numba sits in the middle: you still write the kernel body, compute the global index yourself, and guard against out-of-bounds threads, exactly like CUDA C. What Numba removes is the C compiler toolchain: there is no separate .cu file and no nvcc invocation. The function is compiled the first time it runs.

PyCUDA: raw CUDA C, orchestrated from Python

import numpy as np
import pycuda.autoinit          # initializes a CUDA context on import
import pycuda.driver as cuda_drv
from pycuda.compiler import SourceModule

kernel_source = """
__global__ void vec_add(const float *a, const float *b, float *c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        c[i] = a[i] + b[i];
    }
}
"""

mod = SourceModule(kernel_source)     # compiles the CUDA C string via NVRTC/nvcc
vec_add = mod.get_function("vec_add")

n = 1_000_000
h_a = np.arange(n, dtype=np.float32)
h_b = (np.arange(n, dtype=np.float32) * 2)
h_c = np.empty_like(h_a)

d_a = cuda_drv.mem_alloc(h_a.nbytes)  # explicit device allocation
d_b = cuda_drv.mem_alloc(h_b.nbytes)
d_c = cuda_drv.mem_alloc(h_c.nbytes)

cuda_drv.memcpy_htod(d_a, h_a)        # explicit host -> device copy
cuda_drv.memcpy_htod(d_b, h_b)

threads_per_block = 256
blocks = (n + threads_per_block - 1) // threads_per_block
vec_add(d_a, d_b, d_c, np.int32(n),
        block=(threads_per_block, 1, 1), grid=(blocks, 1))

cuda_drv.Context.synchronize()        # explicit sync

cuda_drv.memcpy_dtoh(h_c, d_c)        # explicit device -> host copy
expected = h_a + h_b
assert np.allclose(h_c, expected)
print(f"All {n} elements correct via PyCUDA.")

PyCUDA is the least abstracted of the three: you write actual CUDA C, manage raw device pointers, and pass kernel arguments positionally with explicit NumPy-typed scalars (np.int32(n)) so PyCUDA marshals them correctly. This is close to writing CUDA C++ directly, with Python replacing the host-side main() and build system.

JIT vs. AOT, and why warm-up matters

CuPy’s elementwise operations, Numba’s @cuda.jit functions, and PyCUDA’s SourceModule all compile at runtime, not ahead of time. The first call to any of them pays a compilation tax:

• CuPy’s built-in kernels are precompiled into the library itself for common operations, but custom kernels written with cp.ElementwiseKernel, cp.RawKernel, or cp.RawModule compile via NVRTC (NVIDIA’s runtime compiler) the first time they run for a given argument type and shape.
• Numba’s @cuda.jit compiles the decorated Python function to PTX through Numba’s NVVM-based backend, specialized to the argument types seen on the first call. A second call with different dtypes triggers a new compilation.
• PyCUDA’s SourceModule compiles the CUDA C string you provide via NVRTC or nvcc when you construct it, not on every call, but that construction step itself is a JIT compile that happens when your Python process runs.

None of these are ahead-of-time (AOT) compiled by default: an AOT approach would compile once at build time (using nvcc to produce a .cubin or .ptx file, or Numba’s separate AOT compilation utilities) and load the precompiled binary at runtime with no compilation step at all. AOT trades flexibility (you cannot change the kernel source without recompiling and redeploying) for zero first-call latency, which matters for short-lived processes or strict latency SLAs.

Because JIT compilation happens on first use, always warm up before timing anything:

import time
import cupy as cp

def timed_run(fn, *args, warmup=3, iters=20):
    for _ in range(warmup):
        fn(*args)
    cp.cuda.Stream.null.synchronize()

    start = time.perf_counter()
    for _ in range(iters):
        fn(*args)
    cp.cuda.Stream.null.synchronize()
    elapsed = time.perf_counter() - start
    return elapsed / iters

Skipping the warm-up loop means your “average” time is dominated by a one-time compilation cost that will not recur in a long-running process. Numba caches compiled functions in memory for the lifetime of the process by argument-type signature, and can persist that cache to disk across process restarts with @cuda.jit(cache=True). CuPy similarly caches compiled kernels both in memory and, by default, in an on-disk kernel cache directory, so repeated runs of the same script after the first are faster to reach steady state.

Streams, errors, and profiling

All three libraries expose CUDA streams, but with different APIs:

# CuPy: streams are objects; use them as context managers
stream = cp.cuda.Stream()
with stream:
    d_c = cp.add(d_a, d_b)   # enqueued on `stream`, not the default stream
stream.synchronize()

# Numba: create a stream and pass it to array transfers and kernel launches
stream = cuda.stream()
d_a = cuda.to_device(h_a, stream=stream)
vec_add_kernel[blocks, threads_per_block, stream](d_a, d_b, d_c)
stream.synchronize()

# PyCUDA: streams come from the driver API
stream = cuda_drv.Stream()
cuda_drv.memcpy_htod_async(d_a, h_a, stream)
vec_add(d_a, d_b, d_c, np.int32(n),
        block=(256, 1, 1), grid=(blocks, 1), stream=stream)
stream.synchronize()

The streams and asynchronous execution article covers what streams actually buy you (overlapping transfers with compute); the point here is only that the concept and the underlying CUDA stream object are identical across all three libraries. They are wrapping the same driver and runtime API underneath.

Error handling also differs in style but not in substance. CuPy raises Python exceptions (cupy.cuda.runtime.CUDARuntimeError) from failed CUDA API calls. Numba raises numba.cuda.cudadrv.driver.CudaAPIError for driver-level failures, and a kernel that itself indexes out of bounds may simply produce garbage or, with NUMBA_CUDA_DEBUGINFO/compute-sanitizer, be caught during debugging. PyCUDA raises pycuda._driver.LogicError or similar for API misuse. In every case, wrap early development iterations with compute-sanitizer (covered in debugging and profiling) because none of these Python layers catch out-of-bounds device memory access on their own; that is a hardware-level fault, not a Python-level one.

Profiling is unified at the tool level: nsys profile python my_script.py and ncu python my_script.py work identically across CuPy, Numba, and PyCUDA, because Nsight instruments the CUDA driver and runtime API calls, not any particular Python library. NVTX ranges (cupy.cuda.nvtx, torch.cuda.nvtx, or the standalone nvtx Python package) let you annotate regions of Python code so they show up as labeled blocks in the Nsight Systems timeline, which is invaluable once a script mixes many small GPU operations.

Interop between libraries

A single project often uses more than one of these libraries, plus a deep learning framework, on the same GPU arrays. This works because of two interoperability protocols:

• __cuda_array_interface__: a lightweight protocol (device pointer, shape, dtype, strides) that CuPy, Numba, and several other libraries implement. Any two libraries that both support it can share a device array without copying it.
• DLPack: a more general tensor-exchange format used by PyTorch, CuPy, and others, again enabling zero-copy handoff of GPU memory between frameworks.

import cupy as cp
from numba import cuda

d_arr = cp.arange(1000, dtype=cp.float32)

# Numba can wrap a CuPy array directly via the array interface, no copy:
numba_view = cuda.as_cuda_array(d_arr)

@cuda.jit
def scale_in_place(arr, factor):
    i = cuda.grid(1)
    if i < arr.size:
        arr[i] *= factor

scale_in_place[4, 256](numba_view, 2.0)
cuda.synchronize()

# d_arr itself was modified in place; no data left the GPU.
print(cp.asnumpy(d_arr)[:5])

This is the practical payoff of learning more than one of these libraries: you can use CuPy for the 95% of an array-processing pipeline that maps cleanly onto NumPy semantics, and drop into a Numba @cuda.jit kernel only for the custom stencil, scatter, or fusion operation CuPy does not already provide, without ever leaving the GPU or writing a separate .cu file.

Raw kernels and modules: when you need actual CUDA C

Sometimes you need a custom CUDA C kernel but still want it to interoperate cleanly with CuPy or PyCUDA arrays, rather than writing a full Numba kernel in restricted Python. Two mechanisms exist for this:

• CuPy RawKernel / RawModule: write real CUDA C++ as a Python string, compiled via NVRTC, callable with CuPy arrays as arguments directly.
• PyCUDA SourceModule: the same idea, shown above for vector addition, generalizes to any CUDA C source, including multiple kernels and __device__ helper functions in one string.

import cupy as cp

add_kernel = cp.RawKernel(r'''
extern "C" __global__
void vec_add(const float* a, const float* b, float* c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        c[i] = a[i] + b[i];
    }
}
''', 'vec_add')

n = 1_000_000
d_a = cp.arange(n, dtype=cp.float32)
d_b = cp.arange(n, dtype=cp.float32) * 2
d_c = cp.empty(n, dtype=cp.float32)

threads_per_block = 256
blocks = (n + threads_per_block - 1) // threads_per_block
add_kernel((blocks,), (threads_per_block,), (d_a, d_b, d_c, n))

cp.cuda.Stream.null.synchronize()
assert cp.allclose(d_c, d_a + d_b)

extern "C" is required so NVRTC does not name-mangle vec_add, which would break the string-based lookup RawKernel uses to find the entry point. This pattern gives you CuPy’s ergonomic array handling and memory pool alongside hand-written CUDA C for the one operation that needs it.

A note on Numba’s CUDA packaging

The CUDA target bundled with the main Numba package is deprecated. Active CUDA development moved to NVIDIA’s separate numba-cuda package, which deliberately preserves the familiar from numba import cuda import path. That compatibility means source code can look unchanged even though the package providing the implementation has changed.

Install and migration details remain version-sensitive: supported Python, CUDA toolkit, and driver ranges can change independently. Before setting up an environment, check the current Numba-CUDA documentation and migration guide, then pin the Python packages and document the driver/toolkit versions together. Do not assume that installation instructions from an older numba.cuda tutorial still describe the maintained package.

Decision matrix

In practice

Start with CuPy for anything that looks like array math; it gets you GPU acceleration with the least code change and the lowest maintenance burden, and its prebuilt kernels are written and tuned by people who specialize in exactly that operation. Reach for Numba when CuPy does not already provide the operation you need and you want to stay in Python. Reach for PyCUDA, or CuPy’s RawKernel, when you already have CUDA C code, need behavior CUDA C exposes that a Python-level abstraction does not, or are integrating with an existing CUDA C++ codebase. None of these choices are permanent or mutually exclusive within one project; it is common to mix all three, and to eventually promote a hot, stable kernel into a compiled CUDA C++ library once its interface stops changing.

What comes next

You have now seen the same computation expressed at three different levels of the Python-to-CUDA abstraction ladder, and you know how to warm up, stream, and profile each one. The next article, CUDA thread hierarchy in depth, returns to CUDA C and goes deeper into multi-dimensional grids, blocks, and warps: the exact same concepts cuda.grid(1) and blockIdx.x * blockDim.x + threadIdx.x were hiding behind a one-line call in the examples above.