any detailed tutorial about how to use it?
Compiling LLMs into a MegaKernel: A path to low-latency inference
21–30 of 79 posts
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#22After working pretty closely with vLLM and SGLang over the past few months, this is EXACTLY what I had envisioned what a successor project would look like - analyzing an operation dependency graph and then fusing (or, at a minimum, scheduling tasks smarter). Congrats to the team.
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.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#23any detailed tutorial about how to use it?
The github repo includes a tutorial for using MPK: https://github.com/mirage-project/mirage/tree/mpk
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#24Does anyone have an intuition on why this offers significant gains over CUDA Graphs?. The CPU launch cost of a graph is tiny which implies most of the work has been offloaded to the GPU's own scheduler. I'd expect that some I/O marshalling at kernel boundaries could be avoided with megakernels. Maybe some loop fusion? Are there any more interesting optimizations they enable?
> The CPU launch cost of a graph is tiny
Absolutely not; it’s comparable to the launch overhead of a kernel.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#25> Traditional LLM systems often rely on sequences of GPU kernel launches and external communication calls, resulting in underutilized hardware. What? Why? This seems like an obvious optimization if it's possible.
From the article > Despite these advantages, compiling an LLM into a megakernel is highly challenging. Existing high-level ML frameworks — such as PyTorch, Triton, and TVM — do not natively support end-to-end megakernel generation. Additionally, modern LLM systems are built from a diverse collection of specialized kernel libraries: NCCL or NVSHMEM for communication, FlashInfer or FlashAttention for efficient attentio…
Your naive assumption is the right one. It’s quite hard to do this. Even doing it automatically like it’s done here runs into problems with trying to figure out data dependencies and synchronization across nontrivial computation.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#26Isn’t fusing ops at a fine-grained level also the core benefit of JAX over TensorFlow? How does this work compare to JAX?
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#27Certainly an important discovery for utilizing these models on scaled hardware. This approach could certainly be applied beyond LLMs to other types of neural networks. That would be an interesting space to explore.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#28Next step - compile straight to verilog so I can buy some LLMs on aliexpress
https://riscv.org/blog/2021/02/hardware-description-language... That was one of the promising ideas before AI & GPUs come to the scene. As CPUs are stagnant, and naturally people want further optimize the middle layers software and hardware. But I suspect parallel computing in GPU style is going to dominate acclerated computing. General purpose CPUs are going to stay to become the little brain that orchestrates GPUs.…
> General purpose CPUs are going to stay to become the little brain that orchestrates GPUs
Brings the deterministic compute to the indeterministic.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#29This is very cool. I enjoyed going through the writeup and GitHub README. I was wondering if these same optimizations can be brought to bear on training as well, rather than only inference. I guess the challenge here is fusing backward computations with gradient communication. I also saw that this currently does not handle dynamic workloads such as MoE. I recently came across this paper that does exactly this: FlashD…
Personally I think its a bit of a waste to invest time into gradient training optimizations. A lot of training tasks IRL have discrete values in nature, which can't be trained with gradients.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#30This project is from CMU. Hazy Research at Stanford talked about the megakernel too: https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles Good to see the competition in this area. (Edited): Related paper covering the larger "mirage" project, but this doesn't cover the "megakernel" approach: https://arxiv.org/abs/2405.05751
Hazy Research also has ThunderKittens, pretty cool library. There's a lot of effort to really formalize, pipeline, divide and conquer in the current NVIDIA GPU model for maximize GPU efficiency, and to write compilers/DSL's for things, it seems.