An Async Runtime for Mojo GPU Coroutines

I am continuing my journey of learning and writing a native TLS library for Mojo. And while working on it, specifically the GPU part, I ran into a problem - how to efficiently “synchronize” GPU workload while running multiple, different GPU-based functions concurrently.

`Mojo` GPU basics intro

Before we dive into what I’ve done, it’s essential to first understand what GPU programming in Mojo looks like and how it works. There is a good set of tutorials with puzzles on the official website, which I really encourage anyone to take a look at. It’s an amazing starting point for everyone who wants to start learning GPU programming.

As a quick refresher, here is how a GPU program generally works, from the Modular GPU fundamentals docs:

… GPU programming follows a distinct pattern where work is divided between the CPU and GPU:

The CPU (host) manages program flow and coordinates GPU operations.

The GPU (device) executes parallel computations across many threads.

You must explicitly manage data exchange between host and device memory.

A GPU program generally follows these steps:

Initialize data in host (CPU) memory.

Allocate device (GPU) memory and transfer data from host to device memory.

Execute a kernel function on the GPU to process the data.

Transfer results back from device to host memory.

This process typically runs asynchronously, allowing the CPU to perform other tasks while the GPU processes data. Any time that the CPU needs to ensure that the GPU has completed an operation, such as before it copies kernel results from device memory, it must first explicitly synchronize with the GPU …

And here is a simple example

from max.gpu.host import DeviceContext
from std.gpu import global_idx

def square_kernel(buf: Pointer[Float32, MutAnyOrigin], size: Int32):
    var idx = global_idx.x
    if idx < Int(size):
        var value = buf[unsafe_offset=idx]
        buf[unsafe_offset=idx] = value * value


def square[
    size: Int
](ctx: DeviceContext, input: Array[Float32, size]) raises -> Array[Float32, size]:
    var device_buffer = ctx.enqueue_create_buffer[DType.float32](size)
    ctx.enqueue_copy(
        dst_buf=device_buffer, src_ptr=input.unsafe_ptr()
    )

    ctx.enqueue_function[square_kernel](
        device_buffer, Int32(size), grid_dim=1, block_dim=size
    )

    var result = Array[Float32, size](uninitialized=True)
    ctx.enqueue_copy(
        dst_ptr=result.unsafe_ptr(), src_buf=device_buffer
    )
    ctx.synchronize()
    return result^

As explicitly stated, working with GPU primitives requires you to always synchronize them, since the whole model of execution between CPU and GPU is async. Synchronization is done by calling ctx.synchronize(), which blocks execution of the current CPU thread until all queued operations on the associated DeviceContext stream have completed.

Problem statement

From what we’ve just learned, in a scenario where your application has to perform aes encryption and calculate sha256, you would probably end up with some code like this:

from max.gpu.host import DeviceContext

def aes_kernel(...):
    ...

def aes(ctx: DeviceContext, ...) raises:
    ...
    ctx.enqueue_function[aes_kernel](...)
    ...
    ctx.synchronize()
    ...


def sha256_kernel(...):
    ...

def sha256(ctx: DeviceContext, ...) raises:
    ...
    ctx.enqueue_function[sha256_kernel](...)
    ...
    ctx.synchronize()
    ...


def main(): raises:
    with DeviceContext() as ctx:
        aes(ctx)
        sha256(ctx)

As you can notice, with respect to executing aes and sha256, the functions are fully sequential and blocking, and there is no concurrent execution of these two functions. Here is what that looks like step by step:

CPU queue

GPU queue

main() calls aes(ctx) — sha256 has not been called yet, everything is strictly sequential. — CPU queue: aes — GPU queue: empty

aes() spawns aes_kernel onto the GPU, then calls ctx.synchronize() directly — this blocks the CPU thread. — CPU queue: ctx.synchronize(), aes₁ — GPU queue: aes_kernel

ctx.synchronize() returns — aes_kernel has completed. aes() finishes the rest of its body. — CPU queue: aes₁ — GPU queue: empty

aes(ctx) has fully returned. Only now does main() call sha256(ctx) — nothing ever overlapped. — CPU queue: sha256 — GPU queue: empty

sha256() spawns sha256_kernel onto the GPU, then calls ctx.synchronize() directly — blocks again. — CPU queue: ctx.synchronize(), sha256₁ — GPU queue: sha256_kernel

ctx.synchronize() returns — sha256_kernel has completed. sha256() finishes the rest of its body. — CPU queue: sha256₁ — GPU queue: empty

sha256(ctx) returns. Two synchronize() calls, two real GPU syncs, and aes/sha256 never ran concurrently. — CPU queue: empty — GPU queue: empty

sha256 doesn’t even get queued until aes has fully returned — each ctx.synchronize() blocks the CPU thread in place, so the GPU sits idle between the two kernels while the host waits.

Using well-known `async`/`await` syntax

So I’ve decided, why not apply the well-known asynchronous programming approach and Mojo’s already existing (but not yet fully matured) async/await and Coroutine capabilities.

That’s exactly what became warp — a small, single-threaded runtime dedicated to scheduling GPU-bound coroutines, which coalesces DeviceContext.synchronize() calls across them instead of leaving every function to manage synchronization on its own.

The rewrite of the aes/sha256 example from before is pretty mechanical: every function that used to call ctx.synchronize() directly becomes async def, and every ctx.synchronize() becomes await ctx.synchronize() — where ctx is no longer a raw DeviceContext, but the gpu_async.Context handed to you by the Executor.

from max.gpu.host import DeviceContext
from std.gpu import global_idx
from gpu_async import Context, Executor

def aes_kernel(...):
    ...

async def aes(ctx: Context, ...) raises:
    ...
    ctx.gpu_ctx().enqueue_function[aes_kernel](...)
    ...
    await ctx.synchronize()
    ...


def sha256_kernel(...):
    ...

async def sha256(ctx: Context, ...) raises:
    ...
    ctx.gpu_ctx().enqueue_function[sha256_kernel](...)
    ...
    await ctx.synchronize()
    ...


def main() raises:
    with DeviceContext() as ctx:
        var executor = Executor(ctx)
        var context = executor.context()

        var t1 = executor.add(aes(context, ...))
        var t2 = executor.add(sha256(context, ...))

        executor.wait()
        t1^.wait()
        t2^.wait()

Notice what aes and sha256 no longer do: neither of them calls DeviceContext.synchronize() directly, and neither of them knows or cares that the other one exists. They just declare, via await ctx.synchronize(), the points at which they’re willing to sync, and the runtime decides when to call DeviceContext.synchronize().

Here is the actual execution flow, step by step:

CPU queue

GPU queue

Both aes and sha256 are queued and ready to run. — CPU queue: aes, sha256 — GPU queue: empty

resume(aes) — spawns aes_kernel onto the GPU queue, then awaits ctx.synchronize(). — CPU queue: sha256, ctx.synchronize(), aes₁ — GPU queue: aes_kernel

resume(sha256) — spawns sha256_kernel onto the GPU queue, then awaits ctx.synchronize(); a sync is already pending, so none is added. — CPU queue: ctx.synchronize(), aes₁, sha256₁ — GPU queue: aes_kernel, sha256_kernel

The pending ctx.synchronize() is reached — real GPU sync fires; aes_kernel and sha256_kernel complete. — CPU queue: aes₁, sha256₁ — GPU queue: empty

resume(aes₁) — returns, task complete. — CPU queue: sha256₁ — GPU queue: empty

resume(sha256₁) — returns. All done: two synchronize() call sites, only one real GPU sync. — CPU queue: empty — GPU queue: empty

Two await ctx.synchronize() call sites between the two functions, but only one real DeviceContext.synchronize() call ever happens.

Compare that to the sync case: aes_kernel and sha256_kernel are both enqueued before that single sync fires, so — unlike the sync version, where the GPU sat idle between two separate blocking syncs — they now run concurrently on the GPU, and most likely in parallel.

And that’s the main idea: pass the decision of when to call DeviceContext.synchronize() off to the runtime, across multiple concurrent GPU tasks.

Final words

warp is still at the proof-of-concept stage. On top of that, Mojo’s async/await and Coroutine support, which the whole executor is built on, is itself still unstable and actively evolving. The language is young, the async story in particular is one of the least mature parts of it, and I’d fully expect breaking changes to land under it before it settles down.

With all that said, I’d genuinely encourage you to go try it out, poke at it, break it. If you run into something - a rough edge, a use case it doesn’t handle, an idea for where this could go next - open an issue, I’d love to hear about it.

An Async Runtime for Mojo GPU Coroutines

Mojo GPU basics intro

Problem statement

Using well-known async/await syntax

Final words

`Mojo` GPU basics intro

Using well-known `async`/`await` syntax