Hacker News Reader

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

1. In MPK, each task is mapped to an individual SM. The amount of work handled by a task is similar to that of a thread block in the traditional kernel-per-operator approach.

2. TL;DR: MPK automatically analyzes inter-task dependencies by tracking the input and output tensors associated with each task. A longer version: Longer version: MPK uses imap, omap, and fmap (see Section 2 of the Mirage paper) to determine each task’s input and output tensors. A dependency is introduced between task A and task B if A produces any tensor elements that B consumes—that is, if A's outputs overlap with B's inputs.

> Again taking matmul as an example: a given output output tile requires the correspond M_BLOCK rows of the A matrix. If the A matrix was itself an output of a prior matmul (+ nonlinearity), the dependees would be all of output tile tasks corresponding to those M_BLOCK rows of the operator that produced A?

Exactly. In this case, all output tile tasks that consume those M_BLOCK rows of A will depend on all tasks responsible for producing the corresponding parts of A in the previous operator.

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

Thanks for reproducing our results!

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

You are right that CUDA graph can help reduce launch overhead but does not support overlapping computation/communication across layers, since data dependencies are described at the kernel level.

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

Ooops, missed one sentence in my previous response. Stanford's MegaKernel project tackles a similar challenge but focuses on manual CUDA implementation. While MPK takes a compiler-driven approach—users express their LLMs at the PyTorch level, and MPK automatically compiles them into optimized megakernels. Our goal is to make programming megakernels much more accessible.

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

Thanks for the great feedback! Stanford's MegaKernel project tackles a similar challenge but focuses on manual CUDA implementation. While MPK takes a compiler-driven approach—users express their LLMs at the PyTorch level, and MPK automatically compiles them into optimized megakernels. Our goal is to make programming megakernels much more accessible.

I completely agree that CUDA can be a limiting factor, especially for latency-sensitive workloads. As GPUs are becoming larger and faster, it's increasingly difficult to write standalone kernels that fully utilize hardware resources—particularly when optimizing for low latency with small batch sizes.

> What are the chances we see your work land in PyTorch as an experimental backend?

We're definitely excited about that direction. We believe MPK can help PyTorch support megakernel generation, and we’re actively exploring how to make that happen. Stay tuned!

> P.S. minor typo, your first two paragraphs under part 1 are nearly identical.

Thanks for pointing it out--I meant to remove the duplicate paragraph when finalizing the post.

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

Yes, it would be a lot of fun if MPK can enable torch.compile to generate megakernels. Torch-generated kernels are currently too slow for latency-sensitive workloads.

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

Thanks for the feedback! Yes, we believe the approach is general and applicable to other ML workloads.

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

The task implementations used by MPK are currently optimized for A100. While the Mirage compiler can generate task implementations for other architectures such as Hopper and Blackwell, but we haven't integrated things together yet. This is on the very top of our todo list. Stay tuned!

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

JAX's operator fusion (https://apxml.com/courses/advanced-jax/chapter-2-optimizing-...) can fuse a few local operators (e.g., matmul and elementwise computation) into a single kernel. But JAX's approach cannot fuse an entire LLM with hundreds of operators into a single kernel because many operators involve loop transformations.

MPK takes a different approach where instead of incrementally fusing local operators, it decomposes operators into a task graph and builds a runtime system within a single kernel to execute all tasks specified in the task graph.

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

The github repo includes a tutorial for using MPK: https://github.com/mirage-project/mirage/tree/mpk

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

Thanks a lot for your positive feedback! We believe that MPK can enhance existing LLM serving systems, especially for low-latency LLM serving. We are very excited about the opportunity to collaborate with others on direction.

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

Thanks for reading the post and github README. Supporting training is definitely feasible but the benefit may not be as significant as low-latency inference since training generally involves much larger kernels, making kernel launch overhead less significant.

Thanks for sharing the FlashDMoE work. Our next step is to support MoE models. Stay tuned!

zhihaojiaon Compiling LLMs into a MegaKernel: A path to low-latency inference

This is the writer of the blog post. You are right that Stanford's work is a parallel effort. The main difference is that our focus is on compilation: making it easier to generate megakernels automatically.

zhihaojiaon Accelerating LLM Serving with Speculative Inference and Token Tree Verification

SpecInfer is a system that accelerates generative LLM serving with speculative inference and token tree verification. The key idea is to use an LLM as a token tree verifier instead of an incremental decoder. We show that this reduces LLM inference latency by 2.8x.