CUDA Python: Memory Transfers and Benchmarking

A GPU benchmark that forgets how CUDA schedules work measures the wrong thing. This lesson times the same kinds of work several ways and shows which numbers to trust. Everything ran on a Tesla T4 in Google Colab with CUDA 13.0, numba 0.61.2, numba-cuda 0.28.2 and cupy-cuda13x 14.0.1. Timings change from run to run, so read the ratios rather than the exact milliseconds. Each example was run three times, and the prose gives the range.

A launch returns before the work is done

Time a kernel launch on its own, then time it again up to the moment the GPU finishes.

# nosync.py - time a kernel without and with synchronising
import time
import numpy as np
from numba import cuda

@cuda.jit
def scale(x, out):
    i = cuda.grid(1)
    if i < out.size:
        out[i] = 2.0 * x[i]

n = 1 << 26                                    # 64 Mi floats, 256 MiB per array
x = cuda.to_device(np.ones(n, dtype=np.float32))
out = cuda.device_array_like(x)
blocks = (n + 255) // 256
scale[blocks, 256](x, out)                     # warm-up: compile once
cuda.synchronize()

t0 = time.perf_counter()
scale[blocks, 256](x, out)                     # returns as soon as it is queued
t1 = time.perf_counter()
cuda.synchronize()                             # wait for the GPU to finish
t2 = time.perf_counter()
print(f"launch returned after {(t1 - t0) * 1e3:.3f} ms")
print(f"kernel finished after {(t2 - t0) * 1e3:.3f} ms")

$ python nosync.py
launch returned after 0.104 ms
kernel finished after 2.230 ms

The launch came back in about 0.1 ms, but the kernel needed about 2.2 ms (2.23 to 2.47 ms across three runs). A launch only queues the work, and Python carries on at once. Time without cuda.synchronize() and you’d report a kernel twenty times faster than it is.

Time on the GPU’s clock with events

CUDA events are timestamps the GPU records in its own work queue, so they measure only the GPU work, with no Python overhead in between.

# events.py - time a kernel on the GPU's own clock and turn it into GB/s
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 << 24
a = cuda.to_device(np.ones(n, dtype=np.float32))
b = cuda.to_device(np.ones(n, dtype=np.float32))
c = cuda.device_array_like(a)
blocks = (n + 255) // 256
vec_add[blocks, 256](a, b, c)                  # warm-up (JIT compile)

start, stop = cuda.event(), cuda.event()
iters = 20
start.record()
for _ in range(iters):
    vec_add[blocks, 256](a, b, c)
stop.record()
stop.synchronize()                             # wait until the stop event has happened
ms = cuda.event_elapsed_time(start, stop) / iters
gbs = 3 * n * 4 / (ms / 1e3) / 1e9             # two reads + one write per element
print(f"vec_add on {n:,} floats: {ms:.3f} ms, {gbs:.0f} GB/s")

$ python events.py
vec_add on 16,777,216 floats: 0.768 ms, 262 GB/s

Each call reads two arrays and writes one, 192 MiB in total, in 0.768 ms. That’s 262 GB/s, about 82% of the T4’s 320 GB/s peak memory bandwidth, and it stayed within 1 GB/s across three runs. Averaging over 20 launches smooths out the noise of a single one. stop.synchronize() waits only until the GPU reaches the stop event.

The first call pays for compilation

@cuda.jit compiles lazily, the first time the kernel runs.

# warmup.py - the first launch pays for compilation
import time
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]

x = cuda.to_device(np.ones(1 << 20, dtype=np.float32))
out = cuda.device_array_like(x)
for call in range(1, 4):
    t0 = time.perf_counter()
    vec_add[4096, 256](x, x, out)
    cuda.synchronize()
    print(f"call {call}: {(time.perf_counter() - t0) * 1e3:.1f} ms")

$ python warmup.py
call 1: 477.0 ms
call 2: 0.2 ms
call 3: 0.1 ms

The first call took about 480 ms (477 to 487 ms across runs), and the next ones took a fraction of a millisecond. That first call is JIT compilation plus one-time compiler setup: loading NVVM and libdevice, then loading the compiled module onto the GPU. It isn’t the first contact with the GPU, because cuda.to_device() had already created the CUDA context before the timer started. Always make one untimed call before you start the timer, as the earlier examples do.

Where the time goes when data lives on the host

Most real work starts in a NumPy array and ends in one. Events placed around each step show the cost of each part, and the result is copied back twice: into an existing array, then into a new one.

# breakdown.py - where the time goes when the data starts and ends on the host
import numpy as np
from numba import cuda

@cuda.jit
def scale(x, y):
    i = cuda.grid(1)
    if i < y.size:
        y[i] = 1.5 * x[i]

n = 1 << 24                                    # 16 Mi floats = 64 MiB
blocks = (n + 255) // 256
x, y = np.ones(n, dtype=np.float32), np.ones(n, dtype=np.float32)
d_x, d_y = cuda.to_device(x), cuda.device_array(n, dtype=np.float32)
scale[blocks, 256](d_x, d_y)                   # warm-up
ev = [cuda.event() for _ in range(5)]

ev[0].record()
d_x.copy_to_device(x)                          # host -> device
ev[1].record()
scale[blocks, 256](d_x, d_y)
ev[2].record()
d_y.copy_to_host(y)                            # device -> host, into the existing y
ev[3].record()
d_y.copy_to_host()                             # device -> host, into a brand-new array
ev[4].record()
ev[4].synchronize()
for name, a in [("H2D copy", 0), ("kernel", 1), ("D2H into y", 2), ("D2H into new array", 3)]:
    print(f"{name}: {cuda.event_elapsed_time(ev[a], ev[a + 1]):.2f} ms")

$ python breakdown.py
H2D copy: 14.74 ms
kernel: 0.63 ms
D2H into y: 14.52 ms
D2H into new array: 24.85 ms

Moving 64 MiB to the GPU took about 15 ms (14.6 to 14.8 ms across three runs), while the kernel took about 0.6 ms, so each copy cost over twenty times the computation. The two copies back move the same bytes but differ in where they land. Copying into y, an array that already exists, took 14.3 to 14.6 ms, the same as the copy in. copy_to_host() with no argument took 24.9 to 25.1 ms, because it allocates a new NumPy array for the result first. In a loop, pass an existing array. The pinned version in Practice 2 reuses one buffer, and both directions drop to about 5 ms.

Pinned memory copies faster

An ordinary NumPy array is pageable: the operating system can move it, so the driver first copies it into a locked staging buffer. A pinned array is locked in RAM already.

# pinned.py - pageable vs pinned host memory for the same 256 MiB copy
import time
import numpy as np
from numba import cuda

n = 1 << 26                                    # 64 Mi floats = 256 MiB
d = cuda.device_array(n, dtype=np.float32)
pageable = np.ones(n, dtype=np.float32)
pinned = cuda.pinned_array(n, dtype=np.float32)
pinned[:] = 1.0

for name, host in [("pageable", pageable), ("pinned", pinned)]:
    d.copy_to_device(host)                     # warm-up
    cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(5):
        d.copy_to_device(host)
    cuda.synchronize()
    seconds = (time.perf_counter() - t0) / 5
    print(f"{name}: {n * 4 / seconds / 1e9:.1f} GB/s host to device")

$ python pinned.py
pageable: 4.6 GB/s host to device
pinned: 12.2 GB/s host to device

Pinned memory ran at 12.2 GB/s against 4.6 GB/s, about 2.7 times faster, and the numbers barely moved across runs. Both are far below the 262 GB/s the kernel reached inside GPU memory, because the copies cross the PCIe bus. Pinned memory is a limited resource, so pin the buffers you reuse, not everything.

Move the data once

Twenty steps of a cheap kernel, first with a round trip every step, then with the data kept on the GPU.

# resident.py - copy the data every step, or keep it on the GPU
import time
import numpy as np
from numba import cuda

@cuda.jit
def step(x):
    i = cuda.grid(1)
    if i < x.size:
        x[i] = x[i] * 0.5 + 1.0

n, steps = 1 << 24, 20
blocks = (n + 255) // 256
x = np.zeros(n, dtype=np.float32)
step[blocks, 256](cuda.to_device(x))           # compile before timing

t0 = time.perf_counter()
h = x.copy()
for _ in range(steps):                         # a round trip every step
    d = cuda.to_device(h)
    step[blocks, 256](d)
    h = d.copy_to_host()
t1 = time.perf_counter()
d = cuda.to_device(x)
for _ in range(steps):                         # the data stays on the device
    step[blocks, 256](d)
h2 = d.copy_to_host()
t2 = time.perf_counter()
print(f"copy every step: {(t1 - t0) * 1e3:.0f} ms")
print(f"copy once: {(t2 - t1) * 1e3:.0f} ms")
print("same answer:", np.array_equal(h, h2))

$ python resident.py
copy every step: 856 ms
copy once: 48 ms
same answer: True

Keeping the data on the device was about 18 times faster (856 to 898 ms against 48 to 58 ms), with the same answer. The kernel work is identical in both loops. The first loop adds 38 copies (40 against 2), and every one of its 20 steps also allocates a fresh device array in to_device and a fresh NumPy array in copy_to_host, which is the “allocating inside the loop” gotcha below. When a pipeline has several GPU steps, chain them on device arrays and copy once at each end.

Common gotchas

• Timing without synchronising: you measure the launch, not the kernel.
• Timing the first call: you measure the compiler. Call once before the timer starts, and again after changing argument dtypes.
• Measuring only the kernel: if production copies data every call, benchmark the copies too, because they’re usually the larger part.
• Allocating inside the loop: a fresh host or device array per iteration adds cost that a reused buffer avoids.

Practice

1. Run the events.py measurement at 65,536, 1,048,576 and 16,777,216 floats. At which size does the launch overhead dominate?

# q1_small.py - events.py at three sizes: when does launch overhead win?
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]

for n in (1 << 16, 1 << 20, 1 << 24):
    a = cuda.to_device(np.ones(n, dtype=np.float32))
    b = cuda.to_device(np.ones(n, dtype=np.float32))
    c = cuda.device_array_like(a)
    blocks = (n + 255) // 256
    vec_add[blocks, 256](a, b, c)
    start, stop = cuda.event(), cuda.event()
    start.record()
    for _ in range(20):
        vec_add[blocks, 256](a, b, c)
    stop.record()
    stop.synchronize()
    ms = cuda.event_elapsed_time(start, stop) / 20
    print(f"{n:,} floats: {ms:.4f} ms, {3 * n * 4 / (ms / 1e3) / 1e9:.0f} GB/s")

$ python q1_small.py
65,536 floats: 0.0466 ms, 17 GB/s
1,048,576 floats: 0.0548 ms, 230 GB/s
16,777,216 floats: 0.7688 ms, 262 GB/s

2. Rewrite breakdown.py with pinned host buffers for both the input and the output, and reuse the output buffer.

# q2_pinned_breakdown.py - breakdown.py with pinned host buffers
import numpy as np
from numba import cuda

@cuda.jit
def scale(x, y):
    i = cuda.grid(1)
    if i < y.size:
        y[i] = 1.5 * x[i]

n = 1 << 24
blocks = (n + 255) // 256
x = cuda.pinned_array(n, dtype=np.float32)     # page-locked input
x[:] = 1.0
y = cuda.pinned_array(n, dtype=np.float32)     # page-locked output, allocated once
d_x, d_y = cuda.to_device(x), cuda.device_array(n, dtype=np.float32)
scale[blocks, 256](d_x, d_y)
ev = [cuda.event() for _ in range(4)]

ev[0].record()
d_x.copy_to_device(x)
ev[1].record()
scale[blocks, 256](d_x, d_y)
ev[2].record()
d_y.copy_to_host(y)
ev[3].record()
ev[3].synchronize()
for name, a, b in [("H2D copy", 0, 1), ("kernel", 1, 2), ("D2H copy", 2, 3)]:
    print(f"{name}: {cuda.event_elapsed_time(ev[a], ev[b]):.2f} ms")

$ python q2_pinned_breakdown.py
H2D copy: 5.56 ms
kernel: 0.69 ms
D2H copy: 5.21 ms

3. Call the warmup.py kernel twice with float32 arrays, then twice with float64 arrays. What happens on the first float64 call?

# q3_types.py - a new argument type means a new compilation
import time
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]

for dtype in (np.float32, np.float32, np.float64, np.float64):
    x = cuda.to_device(np.ones(1 << 20, dtype=dtype))
    out = cuda.device_array_like(x)
    t0 = time.perf_counter()
    vec_add[4096, 256](x, x, out)
    cuda.synchronize()
    print(f"{np.dtype(dtype).name}: {(time.perf_counter() - t0) * 1e3:.1f} ms")
print("compiled signatures:", len(vec_add.signatures))

$ python q3_types.py
float32: 489.3 ms
float32: 0.2 ms
float64: 61.5 ms
float64: 0.3 ms
compiled signatures: 2

4. Time a + b on two CuPy arrays with cupyx.profiler.benchmark and convert the result to GB/s.

# q4_cupy_bench.py - let CuPy's benchmark helper do the warm-up and the GPU timing
import cupy as cp
from cupyx.profiler import benchmark

n = 1 << 24
a = cp.ones(n, dtype=cp.float32)
b = cp.ones(n, dtype=cp.float32)

result = benchmark(lambda: a + b, n_repeat=20, n_warmup=3)
ms = result.gpu_times.mean() * 1e3
print(f"a + b on {n:,} floats: {ms:.3f} ms on the GPU clock")
print(f"bandwidth: {3 * n * 4 / (ms / 1e3) / 1e9:.0f} GB/s")

$ python q4_cupy_bench.py
a + b on 16,777,216 floats: 0.790 ms on the GPU clock
bandwidth: 255 GB/s