Compiling LLMs into a MegaKernel: A path to low-latency inference
zhihaojia.medium.com
Compiling LLMs into a MegaKernel: A path to low-latency inference
1–10 of 79 posts
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#2Ollama integration?
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#3Next step - compile straight to verilog so I can buy some LLMs on aliexpress
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#4> 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.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#5This 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:
FlashDMoE: Fast Distributed MoE in a Single Kernel - https://arxiv.org/pdf/2506.04667
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#6[dead]
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#7The Qwen 8B number, if verified, is very impressive. Much more practical than the previous megakernel one.
That's being said, these one-persisted kernel on each SM reminds me Larrabee, and now wondering what the world will be if we just do traditional process-thread-simd path rather than CUDA path.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#8> 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.
Systems might want to anticipate changes in LLM architectures (even small changes can make a big difference kernel wise), so it's good to not "bake" too much in ahead of time.
That said, at some point it just depends where the costs lie and it might make sense hiring some GPU engineers to do what they did here for whatever architecture you're optimising for.
Not as low-hanging as you might imagine.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#9> 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 attention, and CUDA or Triton for custom computation. This fragmentation makes it difficult to consolidate the entire inference pipeline into a single, unified kernel.
So my naive assumption is that yes it is obvious, but nontrivial.
Re: Compiling LLMs into a MegaKernel: A path to low-latency inference
#10> 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.
It really is not obvious. These launches are asynchronous, and data movement / computation is overlapped properly through CUDA APIs. Even per-kernel launch cost is reduced with the cudagraph introduction.
CUDA programming model relies on each kernel to be computationally expensive to make sense, and these are not true for token generation of LLM. And we are talking about network evaluation at higher than 1000 per second, whereas previously besides recommendation systems, network evaluation we are look at is ~100 per second at most.
Also, nobody remember Alex's "One Weird Trick" paper, which slices matmul into pieces to overlap device-to-device transfer v.s. computation. That is 10 years ago.