CUDA Python: Streams, Overlap and Multi-GPU Patterns

A stream is an ordered queue of GPU work. Work within one stream runs in order, but work in different streams can run at the same time, so a copy can overlap a kernel. All examples ran on a Tesla T4 in Google Colab with CUDA 13.0, numba 0.61.2 and numba-cuda 0.28.2. Colab has a single GPU, so the multi-GPU code at the end ran on one device: it shows the pattern, not a multi-GPU speedup.

Work on a stream runs in the background

An event is a marker in a stream. query() asks whether the GPU has reached it yet, without waiting.

# async_launch.py - work on a stream runs in the background while Python carries on
import numpy as np
from numba import cuda

@cuda.jit
def busy(x, rounds):
    i = cuda.grid(1)
    if i < x.size:
        v = x[i]
        for _ in range(rounds):           # deliberately slow arithmetic
            v = v * 0.999 + 0.001         # float64 on purpose: slow work
        x[i] = v

x = cuda.to_device(np.ones(1 << 20, dtype=np.float32))
busy[4096, 256](x, 1)                     # compile first
cuda.synchronize()

s = cuda.stream()
done = cuda.event()
busy[4096, 256, s](x, 20000)              # queued on stream s
done.record(s)                            # a marker after the kernel in the same stream
print("right after the launch, finished:", done.query())
s.synchronize()                           # block until everything on s has run
print("after s.synchronize(), finished:", done.query())

$ python async_launch.py
right after the launch, finished: False
after s.synchronize(), finished: True

The launch returned while the kernel was still running, so the event reported False. Python could do other work until s.synchronize(), which blocks until everything queued on s is finished. All three runs printed the same two lines.

Reading results too early

An asynchronous copy to pinned memory returns immediately, and so does everything queued before it.

# stale.py - read an async copy's destination too early, then after synchronising
import numpy as np
from numba import cuda

@cuda.jit
def busy(x, rounds):
    i = cuda.grid(1)
    if i < x.size:
        v = x[i]
        for _ in range(rounds):
            v = v * 0.999 + 0.001         # float64 on purpose: slow work
        x[i] = v

@cuda.jit
def set_value(x, value):
    i = cuda.grid(1)
    if i < x.size:
        x[i] = value

scratch = cuda.to_device(np.ones(1 << 20, dtype=np.float32))
d = cuda.device_array(1 << 20, dtype=np.float32)
out = cuda.pinned_array(1 << 20, dtype=np.float32)
out[:] = 0.0
busy[4096, 256](scratch, 1); set_value[4096, 256](d, 0.0); cuda.synchronize()

s = cuda.stream()
busy[4096, 256, s](scratch, 20000)        # slow work queued first
set_value[4096, 256, s](d, 3.0)
d.copy_to_host(out, stream=s)             # async copy: returns at once
print("read too early:", out[0], out[-1])
s.synchronize()
print("after s.synchronize():", out[0], out[-1])

$ python stale.py
read too early: 0.0 0.0
after s.synchronize(): 3.0 3.0

The first read returned the old zeros, with no error and no warning, because the slow kernel was still running and the copy hadn’t happened yet. After s.synchronize() the array holds the 3.0 the kernel wrote. Any code that reads a destination array before synchronising has a race that can pass a quick test and fail under load.

What the hardware can overlap

Whether copies and kernels can run at once depends on the GPU’s copy engines.

# engines.py - can this GPU copy and compute at the same time?
from numba import cuda

dev = cuda.get_current_device()
print("device:", dev.name)
print("copy engines (ASYNC_ENGINE_COUNT):", dev.ASYNC_ENGINE_COUNT)
print("concurrent kernels:", bool(dev.CONCURRENT_KERNELS))
print("GPUs visible:", len(cuda.gpus))

$ python engines.py
device: Tesla T4
copy engines (ASYNC_ENGINE_COUNT): 3
concurrent kernels: True
GPUs visible: 1

The T4 reports 3 asynchronous engines and supports running kernels concurrently, so a copy in each direction can run alongside a kernel. Only one GPU is visible here, which is why the multi-GPU section below runs on one device.

A pipeline on one, two and four streams

Split 256 MiB into 16 chunks. Each chunk is copied in, processed and copied back, all on one stream. Then the chunks are dealt round-robin to several streams, each with its own device buffers.

# overlap.py - the same chunked pipeline on one, two and four streams
import time
import numpy as np
from numba import cuda

@cuda.jit
def work(a, out):
    i = cuda.grid(1)
    if i < out.size:
        v = a[i]
        for _ in range(500):              # enough work to be worth overlapping
            v = v * np.float32(0.99) + np.float32(0.01)   # float32, not float64
        out[i] = v

N, CHUNK = 1 << 26, 1 << 22                # 256 MiB in 16 chunks of 16 MiB
src = cuda.pinned_array(N, dtype=np.float32)
src[:] = 1.0
dst = cuda.pinned_array(N, dtype=np.float32)

def run(k):
    streams = [cuda.stream() for _ in range(k)]
    bufs = [(cuda.device_array(CHUNK, np.float32), cuda.device_array(CHUNK, np.float32))
            for _ in range(k)]
    cuda.synchronize()
    t0 = time.perf_counter()
    for c in range(N // CHUNK):
        s, (d_in, d_out), lo = streams[c % k], bufs[c % k], c * CHUNK
        d_in.copy_to_device(src[lo:lo + CHUNK], stream=s)
        work[CHUNK // 256, 256, s](d_in, d_out)
        d_out.copy_to_host(dst[lo:lo + CHUNK], stream=s)
    cuda.synchronize()
    return (time.perf_counter() - t0) * 1e3

run(1)                                     # compile and warm up
for k in (1, 2, 4):
    print(f"{k} stream(s): {run(k):.0f} ms")

$ python overlap.py
1 stream(s): 103 ms
2 stream(s): 63 ms
4 stream(s): 50 ms

One stream took 103 ms (101 to 103 ms across three runs). With two streams the same work took 53 to 63 ms, and with four it took 45 to 50 ms, roughly half. On one stream every chunk waits for the previous chunk’s copy back. With several, one chunk’s copies run while another chunk’s kernel computes. The next program times the two halves on their own, with events, to show why the saving stops at about half.

# kernel_only.py - time overlap.py's kernels and copies separately, with events
import numpy as np
from numba import cuda

@cuda.jit
def work(a, out):
    i = cuda.grid(1)
    if i < out.size:
        v = a[i]
        for _ in range(500):
            v = v * np.float32(0.99) + np.float32(0.01)
        out[i] = v

N, CHUNK = 1 << 26, 1 << 22
src = cuda.pinned_array(N, dtype=np.float32)
src[:] = 1.0
dst = cuda.pinned_array(N, dtype=np.float32)
d_in, d_out = cuda.to_device(src[:CHUNK]), cuda.device_array(CHUNK, np.float32)
work[CHUNK // 256, 256](d_in, d_out)           # compile
ev = [cuda.event() for _ in range(3)]
ev[0].record()
for c in range(N // CHUNK):                    # the 16 kernels, no copies
    work[CHUNK // 256, 256](d_in, d_out)
ev[1].record()
for c in range(N // CHUNK):                    # the 32 copies, no kernels
    d_in.copy_to_device(src[c * CHUNK:(c + 1) * CHUNK])
    d_out.copy_to_host(dst[c * CHUNK:(c + 1) * CHUNK])
ev[2].record()
ev[2].synchronize()
print(f"kernels alone: {cuda.event_elapsed_time(ev[0], ev[1]):.0f} ms")
print(f"copies alone: {cuda.event_elapsed_time(ev[1], ev[2]):.0f} ms")

$ python kernel_only.py
kernels alone: 53 ms
copies alone: 45 ms

The 32 copies alone took 45 ms (45 to 46 ms across three runs). The 16 kernels alone took 53 ms in this run, but they moved more between runs, from 53 to 90 ms, so read that number as approximate. With the two halves costing about the same, one stream pays for both, about 100 ms. Overlapping can at best hide the shorter half behind the longer one, which leaves about 50 ms, and the two- and four-stream times land close to that. Four streams beat two in all three runs, by 5 to 13 ms.

Two details make this work. The host arrays are pinned, and each stream has its own d_in and d_out. If two streams shared one device buffer, a chunk could overwrite data another stream was still reading.

Spreading work over several GPUs

A multi-GPU job splits the data, gives each device its slice, and collects the results only after every device has started.

# multigpu.py - split one array across every visible GPU (on Colab: one T4)
import numpy as np
from numba import cuda

@cuda.jit
def square(x):
    i = cuda.grid(1)
    if i < x.size:
        x[i] = x[i] * x[i]

data = np.arange(1_000_000, dtype=np.float32)
parts = np.array_split(data, len(cuda.gpus))     # one slice per device

pending = []
for dev_id, part in enumerate(parts):            # 1) give every GPU its work
    with cuda.gpus[dev_id]:                      # make this GPU current
        s = cuda.stream()
        d = cuda.to_device(part, stream=s)
        square[(part.size + 255) // 256, 256, s](d)
        pending.append((dev_id, s, d))

results = []
for dev_id, s, d in pending:                     # 2) only then collect the results
    with cuda.gpus[dev_id]:
        results.append(d.copy_to_host(stream=s))
        s.synchronize()
print("GPUs used:", len(parts), "| slice sizes:", [p.size for p in parts])
print("matches NumPy:", np.array_equal(np.concatenate(results), data * data))

$ python multigpu.py
GPUs used: 1 | slice sizes: [1000000]
matches NumPy: True

With one T4, len(cuda.gpus) is 1, so the whole array became a single slice on device 0. That is the honest limit of this run: it shows the code is correct, not that it scales. On a machine with four GPUs the same program makes four slices, and the devices should compute in parallel: each kernel launch returns at once, so the first loop moves on to the next GPU while the previous one is still working. Two caveats, neither tested here because only one GPU was available. First, data is an ordinary pageable array, so each to_device call holds the host while the driver stages the data; pinned host buffers, as in the pipeline above, would let the copies overlap too. Second, collecting inside the first loop would make each GPU wait for the one before it, which is why the results are gathered in a second loop.

Common gotchas

• Reading before synchronising: the destination still holds old data, and nothing warns you.
• Pageable host arrays: asynchronous copies stop overlapping. Use cuda.pinned_array.
• Shared buffers across streams: each in-flight chunk needs its own device memory.
• Float64 by accident: a Python float such as 0.99 in a kernel makes the arithmetic float64, which runs far slower on a T4. Write np.float32(0.99) to stay in float32.
• Expecting more streams to keep helping: once copies and compute overlap, extra streams add little.

Practice

1. Run overlap.py with ordinary NumPy arrays instead of pinned ones. What happens to the benefit of extra streams?

# r1_pageable.py - overlap.py with ordinary (pageable) NumPy host arrays
import time
import numpy as np
from numba import cuda

@cuda.jit
def work(a, out):
    i = cuda.grid(1)
    if i < out.size:
        v = a[i]
        for _ in range(500):
            v = v * np.float32(0.99) + np.float32(0.01)
        out[i] = v

N, CHUNK = 1 << 26, 1 << 22
src = np.ones(N, dtype=np.float32)          # pageable, not pinned
dst = np.empty(N, dtype=np.float32)

def run(k):
    streams = [cuda.stream() for _ in range(k)]
    bufs = [(cuda.device_array(CHUNK, np.float32), cuda.device_array(CHUNK, np.float32))
            for _ in range(k)]
    cuda.synchronize()
    t0 = time.perf_counter()
    for c in range(N // CHUNK):
        s, (d_in, d_out), lo = streams[c % k], bufs[c % k], c * CHUNK
        d_in.copy_to_device(src[lo:lo + CHUNK], stream=s)
        work[CHUNK // 256, 256, s](d_in, d_out)
        d_out.copy_to_host(dst[lo:lo + CHUNK], stream=s)
    cuda.synchronize()
    return (time.perf_counter() - t0) * 1e3

run(1)
for k in (1, 2, 4):
    print(f"{k} stream(s): {run(k):.0f} ms")

$ python r1_pageable.py
1 stream(s): 180 ms
2 stream(s): 166 ms
4 stream(s): 159 ms

2. In stale.py, wait for the copy with an event recorded right after it instead of synchronising the whole stream.

# r2_event.py - wait for one copy with an event instead of the whole stream
import numpy as np
from numba import cuda

@cuda.jit
def set_value(x, value):
    i = cuda.grid(1)
    if i < x.size:
        x[i] = value

d = cuda.device_array(1 << 20, dtype=np.float32)
out = cuda.pinned_array(1 << 20, dtype=np.float32)
out[:] = 0.0
set_value[4096, 256](d, 0.0)
cuda.synchronize()

s = cuda.stream()
set_value[4096, 256, s](d, 3.0)
d.copy_to_host(out, stream=s)
copied = cuda.event()
copied.record(s)                          # marks the point just after the copy
copied.synchronize()                      # wait for that point only
print("after copied.synchronize():", out[0], out[-1])

$ python r2_event.py
after copied.synchronize(): 3.0 3.0

3. Measure how long the launch in async_launch.py takes compared with the wait in s.synchronize().

# r3_wait.py - how long does the host actually wait in synchronize()?
import time
import numpy as np
from numba import cuda

@cuda.jit
def busy(x, rounds):
    i = cuda.grid(1)
    if i < x.size:
        v = x[i]
        for _ in range(rounds):
            v = v * 0.999 + 0.001         # float64 on purpose: slow work
        x[i] = v

x = cuda.to_device(np.ones(1 << 20, dtype=np.float32))
busy[4096, 256](x, 1)
cuda.synchronize()

s = cuda.stream()
t0 = time.perf_counter()
busy[4096, 256, s](x, 20000)
t1 = time.perf_counter()
s.synchronize()
t2 = time.perf_counter()
print(f"launch took {(t1 - t0) * 1e3:.2f} ms, waiting took {(t2 - t1) * 1e3:.0f} ms")

$ python r3_wait.py
launch took 0.22 ms, waiting took 195 ms

4. Split the data in multigpu.py into three slices and deal them round-robin over however many GPUs exist.

# r4_roundrobin.py - deal three slices out to however many GPUs there are
import numpy as np
from numba import cuda

@cuda.jit
def square(x):
    i = cuda.grid(1)
    if i < x.size:
        x[i] = x[i] * x[i]

data = np.arange(1_000_000, dtype=np.float32)
parts = np.array_split(data, 3)                  # more slices than GPUs
pending = []
for n, part in enumerate(parts):
    dev_id = n % len(cuda.gpus)                  # round-robin over the devices
    with cuda.gpus[dev_id]:
        s = cuda.stream()
        d = cuda.to_device(part, stream=s)
        square[(part.size + 255) // 256, 256, s](d)
        pending.append((dev_id, s, d))

results = []
for dev_id, s, d in pending:
    with cuda.gpus[dev_id]:
        results.append(d.copy_to_host(stream=s))
        s.synchronize()
print("slice -> GPU:", [dev_id for dev_id, _, _ in pending])
print("matches NumPy:", np.array_equal(np.concatenate(results), data * data))

$ python r4_roundrobin.py
slice -> GPU: [0, 0, 0]
matches NumPy: True