GPU Kernels in Python with Numba and CuPy

Python can drive the GPU in two styles. Numba compiles a Python function into a CUDA kernel that you launch over thousands of threads. CuPy gives you NumPy-like arrays that live in GPU memory. Everything below ran on a Tesla T4 in Google Colab, with CUDA 13.0, Python 3.13, numba 0.61.2 with numba-cuda 0.28.2, and cupy-cuda13x 14.0.1, all preinstalled on Colab. numba-cuda is the package that now provides numba.cuda, replacing the CUDA target that used to ship inside numba. Elsewhere, install the same versions with pip install "numba-cuda[cu13]==0.28.2" cupy-cuda13x==14.0.1. According to the package metadata, the cu13 extra pulls in the CUDA 13 runtime and compiler libraries. You still need an NVIDIA driver, and this install line was not tested outside Colab.

Check what Numba can see

Before writing a kernel, ask Numba which GPU it found.

# device.py - what Numba sees on this machine
from importlib.metadata import version
from numba import cuda

print("CUDA available:", cuda.is_available())
dev = cuda.get_current_device()
print("device:", dev.name)
print("compute capability:", dev.compute_capability)
print("SMs:", dev.MULTIPROCESSOR_COUNT)
print("max threads per block:", dev.MAX_THREADS_PER_BLOCK)
print("numba", version("numba"), "| numba-cuda", version("numba-cuda"))

$ python device.py
CUDA available: True
device: Tesla T4
compute capability: ComputeCapability(major=7, minor=5)
SMs: 40
max threads per block: 1024
numba 0.61.2 | numba-cuda 0.28.2

The T4 has 40 SMs (streaming multiprocessors), the GPU’s equivalent of CPU cores, and a block can hold at most 1,024 threads. Compute capability 7.5 is the Turing generation. If CUDA available prints False, no kernel below will run, so fix the driver or runtime first.

Your first kernel: one thread per element

A kernel is a function that every thread runs at once. Each thread uses cuda.grid(1) to find the one element it owns.

# vec_add.py - the first Numba kernel: one thread per element
import numpy as np
from numba import cuda

@cuda.jit
def vec_add(a, b, c):
    i = cuda.grid(1)          # this thread's global index
    if i < c.size:            # the grid can be bigger than the data
        c[i] = a[i] + b[i]

n = 1_000_000
a = np.arange(n, dtype=np.float32)
b = np.full(n, 2.0, dtype=np.float32)

d_a, d_b = cuda.to_device(a), cuda.to_device(b)   # host -> device
d_c = cuda.device_array_like(a)                   # uninitialised GPU memory

threads = 256
blocks = (n + threads - 1) // threads
vec_add[blocks, threads](d_a, d_b, d_c)
c = d_c.copy_to_host()                            # device -> host

print("launch:", blocks, "blocks x", threads, "threads =", blocks * threads)
print("c[:3]:", c[:3], "c[-1]:", c[-1])
print("matches NumPy:", np.array_equal(c, a + b))

$ python vec_add.py
launch: 3907 blocks x 256 threads = 1000192
c[:3]: [2. 3. 4.] c[-1]: 1000001.0
matches NumPy: True

The program follows the three steps of nearly every GPU job: copy the inputs to the device, launch the kernel, and copy the result back. The launch syntax vec_add[blocks, threads] sets the grid shape. One million elements need 3,907 blocks of 256, which is 1,000,192 threads, so 192 threads have no element. c[-1] is 999999 + 2, and the whole result agrees with NumPy.

The spare threads are real

The grid is rounded up to whole blocks, so it is almost always a little bigger than the data. This kernel counts the threads that fail the guard.

# spare.py - count the launched threads that have no element to work on
import numpy as np
from numba import cuda

@cuda.jit
def scale(x, out, spare):
    i = cuda.grid(1)
    if i >= out.size:
        cuda.atomic.add(spare, 0, 1)   # this thread has nothing to do
        return
    out[i] = 2.0 * x[i]

n = 100_000
x = cuda.to_device(np.ones(n, dtype=np.float32))
out = cuda.device_array(n, dtype=np.float32)
spare = cuda.to_device(np.zeros(1, dtype=np.int32))

threads = 256
blocks = (n + threads - 1) // threads
scale[blocks, threads](x, out, spare)
print(blocks * threads, "threads for", n, "elements:", spare.copy_to_host()[0], "spare")
print("sum of out:", out.copy_to_host().sum())

$ python spare.py
100096 threads for 100000 elements: 96 spare
sum of out: 200000.0

391 blocks of 256 make 100,096 threads for 100,000 elements, and exactly 96 of them took the early return. cuda.atomic.add is needed because those 96 threads all update the same counter at the same moment, and a plain spare[0] += 1 would lose updates.

What happens without the guard

Drop the if and the spare threads write anyway. Here the array is a 1,000-element view into a 1,024-element buffer, so you can see where those writes land.

# noguard.py - drop the bounds guard and watch the neighbouring data change
import numpy as np
from numba import cuda

@cuda.jit
def fill_twos(out):
    i = cuda.grid(1)
    out[i] = 2.0                  # no "if i < out.size" guard

buf = cuda.to_device(np.full(1024, -1.0, dtype=np.float32))
out = buf[:1000]                  # a 1000-element view; the last 24 floats are not ours

fill_twos[4, 256](out)            # 1024 threads for 1000 elements
host = buf.copy_to_host()
print("last two of out:", host[998:1000])
print("first four after it:", host[1000:1004])
print("floats overwritten past the end:", int((host[1000:] == 2.0).sum()), "of 24")

$ python noguard.py
/usr/local/lib/python3.13/dist-packages/numba_cuda/numba/cuda/dispatcher.py:716: NumbaPerformanceWarning: Grid size 4 will likely result in GPU under-utilization due to low occupancy.
  warn(errors.NumbaPerformanceWarning(msg))
last two of out: [2. 2.]
first four after it: [2. 2. 2. 2.]
floats overwritten past the end: 24 of 24

The kernel didn’t crash and Numba didn’t complain about the out-of-bounds writes. The 24 spare threads simply wrote 2.0 into the 24 floats that follow the view, which here is someone else’s data. In a real program that’s silent memory corruption. The only warning printed is about performance. numba-cuda 0.28.2 prints it for any grid of fewer than 128 blocks, and it matters here because 4 blocks can’t keep 40 SMs busy. The path in that warning depends on where Python is installed, so yours will differ.

The same work with CuPy

CuPy arrays live on the GPU and support NumPy expressions. For anything an expression can’t say, ElementwiseKernel takes a one-line C body and handles the indexing and the guard for you.

# cupy_add.py - the same vector add with CuPy: arrays first, kernels optional
import cupy as cp
import numpy as np

n = 1_000_000
a = cp.arange(n, dtype=cp.float32)        # allocated on the GPU
b = cp.full(n, 2.0, dtype=cp.float32)

c = a + b                                 # CuPy generates and launches the kernel
add = cp.ElementwiseKernel("float32 x, float32 y", "float32 z", "z = x + y", "add")
d = add(a, b)                             # a kernel body you write, CuPy launches

print(type(c), c.device, c[:3])
print("same result:", bool(cp.array_equal(c, d)))
host = cp.asnumpy(c)                      # device -> host
print(type(host), host[-1], np.array_equal(host, np.arange(n) + 2))

$ python cupy_add.py
<class 'cupy.ndarray'> <CUDA Device 0> [2. 3. 4.]
same result: True
<class 'numpy.ndarray'> 1000001.0 True

a + b ran on the GPU without any kernel code, and the custom kernel gave the same answer. The result is a cupy.ndarray on device 0 until cp.asnumpy copies it into an ordinary numpy.ndarray. Printing c[:3] also copied those three values to the host so Python could show them.

Passing NumPy arrays straight in

Numba will accept host arrays in a launch, which looks convenient.

# hostwarn.py - pass NumPy arrays straight to a kernel and read the warning
import numpy as np
from numba import cuda

@cuda.jit
def vec_add(a, b, c):
    i = cuda.grid(1)
    if i < c.size:
        c[i] = a[i] + b[i]

n = 1_000_000
a = np.ones(n, dtype=np.float32)
b = np.ones(n, dtype=np.float32)
c = np.zeros(n, dtype=np.float32)
vec_add[(n + 255) // 256, 256](a, b, c)   # host arrays: Numba copies them for you
print("c[0]:", c[0])                       # and copies c back afterwards

$ python hostwarn.py
/usr/local/lib/python3.13/dist-packages/numba_cuda/numba/cuda/cudadrv/devicearray.py:1089: NumbaPerformanceWarning: Host array used in CUDA kernel will incur copy overhead to/from device.
  warn(NumbaPerformanceWarning(msg))
c[0]: 2.0

The answer is right, because Numba copied all three arrays to the GPU before the launch and copied them all back afterwards. The warning is there because that happens on every call: in a loop you pay for those transfers again and again. Copy once with cuda.to_device and keep the device arrays, as vec_add.py does.

Common gotchas

• No bounds checking: a missing if i < x.size corrupts memory silently, as noguard.py shows.
• Host arrays in a launch: they work but copy in both directions on every call.
• Uninitialised output: device_array_like doesn’t zero its memory. Write every element or start from cuda.to_device(np.zeros(...)).
• A tiny grid: fewer blocks than the GPU has SMs leaves most of it idle.

Practice

1. Change vec_add.py to use 128 threads per block. Predict the number of blocks and spare threads before you run it.

# p1_blocks.py - the same kernel with 128 threads per block
import numpy as np
from numba import cuda

@cuda.jit
def vec_add(a, b, c):
    i = cuda.grid(1)
    if i < c.size:
        c[i] = a[i] + b[i]

n = 1_000_000
a = np.arange(n, dtype=np.float32)
b = np.full(n, 2.0, dtype=np.float32)
d_a, d_b = cuda.to_device(a), cuda.to_device(b)
d_c = cuda.device_array_like(a)

threads = 128
blocks = (n + threads - 1) // threads
vec_add[blocks, threads](d_a, d_b, d_c)
print("launch:", blocks, "blocks x", threads, "threads =", blocks * threads)
print("spare threads:", blocks * threads - n)
print("matches NumPy:", np.array_equal(d_c.copy_to_host(), a + b))

$ python p1_blocks.py
launch: 7813 blocks x 128 threads = 1000064
spare threads: 64
matches NumPy: True

2. Run spare.py with n = 1000. How many threads are spare, and what else gets printed?

# p2_spare.py - predict the spare threads for 1,000 elements in 256-thread blocks
import numpy as np
from numba import cuda

@cuda.jit
def scale(x, out, spare):
    i = cuda.grid(1)
    if i >= out.size:
        cuda.atomic.add(spare, 0, 1)
        return
    out[i] = 2.0 * x[i]

n = 1000
x = cuda.to_device(np.ones(n, dtype=np.float32))
out = cuda.device_array(n, dtype=np.float32)
spare = cuda.to_device(np.zeros(1, dtype=np.int32))

threads = 256
blocks = (n + threads - 1) // threads
scale[blocks, threads](x, out, spare)
print(blocks * threads, "threads for", n, "elements:", spare.copy_to_host()[0], "spare")

$ python p2_spare.py
/usr/local/lib/python3.13/dist-packages/numba_cuda/numba/cuda/dispatcher.py:716: NumbaPerformanceWarning: Grid size 4 will likely result in GPU under-utilization due to low occupancy.
  warn(errors.NumbaPerformanceWarning(msg))
1024 threads for 1000 elements: 24 spare

3. Put the bounds guard back into noguard.py and confirm that the neighbouring floats keep their value of -1.0.

# p3_guard.py - noguard.py with the bounds guard put back
import numpy as np
from numba import cuda

@cuda.jit
def fill_twos(out):
    i = cuda.grid(1)
    if i < out.size:
        out[i] = 2.0

buf = cuda.to_device(np.full(1024, -1.0, dtype=np.float32))
out = buf[:1000]

fill_twos[4, 256](out)
host = buf.copy_to_host()
print("last two of out:", host[998:1000])
print("first four after it:", host[1000:1004])
print("floats overwritten past the end:", int((host[1000:] == 2.0).sum()), "of 24")

$ python p3_guard.py
/usr/local/lib/python3.13/dist-packages/numba_cuda/numba/cuda/dispatcher.py:716: NumbaPerformanceWarning: Grid size 4 will likely result in GPU under-utilization due to low occupancy.
  warn(errors.NumbaPerformanceWarning(msg))
last two of out: [2. 2.]
first four after it: [-1. -1. -1. -1.]
floats overwritten past the end: 0 of 24

4. Write SAXPY (z = a * x + y, with a scalar a) as a CuPy ElementwiseKernel and check it against the same CuPy expression.

# p4_saxpy.py - SAXPY (z = a*x + y) as a CuPy ElementwiseKernel
import cupy as cp

saxpy = cp.ElementwiseKernel(
    "float32 a, float32 x, float32 y", "float32 z", "z = a * x + y", "saxpy")

x = cp.arange(5, dtype=cp.float32)
y = cp.ones(5, dtype=cp.float32)
z = saxpy(cp.float32(3.0), x, y)
print(z)
print("matches cupy expression:", bool(cp.array_equal(z, 3.0 * x + y)))

$ python p4_saxpy.py
[ 1.  4.  7. 10. 13.]
matches cupy expression: True