Skip to main content

The Megakernel

A megakernel combines work from several kernels into one persistent device launch. Extending it to an entire token loop also requires on-device sampling and stop handling. Related persistent-kernel work includes Mirage 1 and MPK 2.

poot has no megakernel. The kernel IR has scoped fences and compare-and-swap, and the PTX and ROCm runtimes expose cooperative-launch mechanisms. Those pieces do not constitute a megakernel: a recorder alone cannot make arbitrary host-loaded kernels callable from a persistent kernel, and a multi-token loop needs sampling, cancellation and streaming contracts. The SPIR-V backends, wgpu and raw Vulkan, have no grid-wide synchronization primitive, which this strategy would require. Execution is native replay or cached dispatch encoding.

Footnotes​

  1. Wu, M. et al., "Mirage: A Multi-Level Superoptimizer for Tensor Programs," OSDI, 2024. https://arxiv.org/abs/2405.05751 ↩

  2. Wu, M. et al., "MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs," arXiv:2512.22219, 2025. https://arxiv.org/abs/2512.22219 ↩