Worms Are Hip: Hijacking the GPU

For a long time, the kernel boundary has been a convenient place to divide responsibility. The host prepares work, the runtime submits a kernel, the GPU executes it, and once that kernel finishes, the system moves on to whatever comes next. That model is simple, useful, and easy to reason about.

But the kernel boundary is also a control boundary.

Hexagon-MLIR helped make that obvious in a different context. Qualcomm’s work targets an NPU, not a GPU, but the compiler lesson carries over. When operations remain separate, intermediate values often have to be materialized, moved, and read back before the next operation can begin. Fusion removes some of those boundaries and gives the compiler a larger region over which it can reason about memory placement, tiling, vectorization, reuse, and scheduling.

Megakernels push that idea much further.

Hazy Research has been exploring kernels large enough to cover substantial portions of model execution. Their work goes beyond fusing a few operations together. The megakernel can contain an instruction representation and an interpreter that coordinates work across the GPU, allowing different parts of the machine to load, compute, store, and communicate without returning to the host between every former kernel boundary.

That does not mean the GPU suddenly owns the entire schedule. In Hazy’s work, a lot of that schedule is still prepared ahead of time on the CPU. The interesting part is that the GPU can now execute a much larger plan internally instead of requiring a new native dispatch for every piece of the graph.

Traditionally, execution might look like this:

host
  ↓
kernel A
  ↓
host/runtime coordination
  ↓
kernel B
  ↓
host/runtime coordination
  ↓
kernel C

A megakernel can move much of that coordination into one larger resident program:

host
  ↓
megakernel
  ↓
interpreter
  ↓
many pieces of work
  ↓
return

Pero, why? Because removing those boundaries also means inheriting the responsibilities that used to live around them. Synchronization still has to happen (unfortunately)…Dependencies still have to be respected. Shared memory, registers, and SM time are still limited. Workers can race, stall, or deadlock.

More control means more responsibility.

Hazy’s own writing about megakernels makes that clear. Larger scheduling regions brought more complicated synchronization, data structures, races, and deadlocks. The kernel can absorb more of the runtime, but somebody still has to build the runtime it absorbed.

There have been several experiments that keep a program resident on an AMD RX 7900 XTX instead of repeatedly launching separate kernels from the host. The compute units are reserved up front, and software running on the GPU can accept things such as tinygrad UOps, assign work to wave groups, manage local memory, track dependencies, and load more compiled work without requiring a new native dispatch every time.

A megakernel takes a larger chunk of computation and coordinates it internally. Flapping Airplanes’ Worm approach begins by treating the GPU more like a machine that stays alive and keeps accepting work.

Once a resident program can consume tasks, track dependencies, assign workers, and continue scheduling instead of exiting after one fixed computation, the line between kernel and runtime starts to get lost in complexity, as more of the scheduling domain that normally belongs to the host runtime, driver, firmware, or hardware scheduler moves inside the resident program.

The interesting part is not simply moving everything onto the GPU. Different layers know different things. The compiler may know the whole graph. The host runtime may know which requests actually arrived and which shapes or routes became concrete. Device-side software may know which workers are free right now. Hardware sees the machine at a much finer timescale than software ever will.

A basic problem is deciding which scheduling decisions belong at which layer.

Moving a decision closer to the hardware only helps when that layer has information the host cannot use as effectively. Otherwise, all we have done is move complexity.

This matters because modern GPUs are already full of resources that can overlap. Tensor cores can compute while memory engines move data. Different warps can take on producer and consumer roles. Different SMs can work on separate parts of a larger pipeline. Communication can overlap with local computation.

That connects back to something I have written about before: capacity is not coordination. A machine can have enormous parallel capacity and still waste most of it if the schedule is poor.

ThunderKittens, HipKittens, and TileLang are some of the many doing exciting stuff in this area. Some give programmers more direct control over tiles, workers, and pipelines. Some push more responsibility into the compiler. Worm goes further by experimenting with a program that remains resident and keeps scheduling new work from inside the GPU. Kernel boundaries exist because they make a hard problem much easier to manage. Removing them increases complexity fast.

Once a kernel can stay resident, interpret work, coordinate resources, and keep accepting new tasks, it starts looking less like a function the runtime launches and more like part of the runtime itself.

That gets us one step closer to hijacking the whole GPU.

添加评论
点赞收藏
点踩分享查看原文
评论
?
参与讨论