Hacker News Board Hacker News, browsed like a classic message board
Live data from Hacker News

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

zhihaojia.medium.com

Post reply on HN
Print view View original

51–60 of 79 posts

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

#51

Post by zhihaojia »

skavi wrote: post #45
saagarjha wrote: post #24

Earlier quoted context omitted.

> The CPU launch cost of a graph is tiny Absolutely not; it’s comparable to the launch overhead of a kernel.

Fair enough. I should have clarified that “approximately the cost of a single kernel launch” is pretty much what I meant by “tiny”. What I was getting at was that a “megakernel” and a captured graph should have similar launch costs.

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.

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

#52

Post by fc417fc802 »

stirfish wrote: post #42
fc417fc802 wrote: post #41

Earlier quoted context omitted.

Isn't that easily accomplished by setting up a local deployment and then yanking the network cable? Anything that can quickly run a capable LLM is going to be a pretty beefy box though. More like LLM in an expensive space heater.

I was thinking more like those Bitcoin mining usb Asics that used to be a thing, but instead of becoming ewaste, you can still use them to talk with chatgpt 2 or whatever. I'm picturing an llm appliance.

There is no magic ASIC that can get around needing to do hundreds of watts worth of computations and having on the order of hundreds of gigabytes of very fast memory. Otherwise the major players would be doing that instead of (quite literally) investing in nuclear reactors to power their future data center expansions.

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

#53

Post by stirfish »

fc417fc802 wrote: post #52
stirfish wrote: post #42

Earlier quoted context omitted.

I was thinking more like those Bitcoin mining usb Asics that used to be a thing, but instead of becoming ewaste, you can still use them to talk with chatgpt 2 or whatever. I'm picturing an llm appliance.

There is no magic ASIC that can get around needing to do hundreds of watts worth of computations and having on the order of hundreds of gigabytes of very fast memory. Otherwise the major players would be doing that instead of (quite literally) investing in nuclear reactors to power their future data center expansions.

That's a really good point. I wasn't thinking further than ollama on my MacBook, but I'm not deploying my laptop into production.

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

#54

Post by gongy »

The improvement is real!

And unlike a lot of research, the code actually runs well. I can reproduce the results using Modal GPUs, leaving the code here: https://github.com/mirage-project/mirage/pull/327/files

Triton + FlashInfer: Prompt length 39, generate length 264, per-token latency 19.189573345762312 ms

MPK: Prompt length 39, generate length 334, per-token latency 7.71875 ms

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

#55

Post by bronxbomber92 »

A question for the author(s) since they seem to be very responsive to this thread :).

1. How fine grain is each task? In a traditional matrix multiplication kernel, for example, each thread block is responsible for a small output tile of the resulting matrix. In Mirage's mega kernel, would there correspondingly be a task for each small output tile?

2. How does the Mirage compiler form the task graph? Does it have domain knowledge of every operator's data flow at the granularity of individual elements? 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?

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

#57

Post by zhihaojia »

gongy wrote: post #54

The improvement is real! And unlike a lot of research, the code actually runs well. I can reproduce the results using Modal GPUs, leaving the code here: https://github.com/mirage-project/mirage/pull/327/files Triton + FlashInfer: Prompt length 39, generate length 264, per-token latency 19.189573345762312 ms MPK: Prompt length 39, generate length 334, per-token latency 7.71875 ms

Thanks for reproducing our results!

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

#58

Post by zhihaojia »

bronxbomber92 wrote: post #55

A question for the author(s) since they seem to be very responsive to this thread :). 1. How fine grain is each task? In a traditional matrix multiplication kernel, for example, each thread block is responsible for a small output tile of the resulting matrix. In Mirage's mega kernel, would there correspondingly be a task for each small output tile? 2. How does the Mirage compiler form the task graph? Does it have dom…

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.

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

#59

Post by touisteur »

skavi wrote: post #45
saagarjha wrote: post #24

Earlier quoted context omitted.

> The CPU launch cost of a graph is tiny Absolutely not; it’s comparable to the launch overhead of a kernel.

Fair enough. I should have clarified that “approximately the cost of a single kernel launch” is pretty much what I meant by “tiny”. What I was getting at was that a “megakernel” and a captured graph should have similar launch costs.

It's not so much kernel overhead than memory traffic between global memory and L2 cache / shared memory that you (or at least I) target with fused kernel approach. Kernel launch overhead can be drastically reduced with cuda-graph indeed.

I'm not sure it applies so well in LLMs though (should read the paper...).

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

#60

Post by otabdeveloper4 »

fc417fc802 wrote: post #52
stirfish wrote: post #42

Earlier quoted context omitted.

I was thinking more like those Bitcoin mining usb Asics that used to be a thing, but instead of becoming ewaste, you can still use them to talk with chatgpt 2 or whatever. I'm picturing an llm appliance.

There is no magic ASIC that can get around needing to do hundreds of watts worth of computations and having on the order of hundreds of gigabytes of very fast memory. Otherwise the major players would be doing that instead of (quite literally) investing in nuclear reactors to power their future data center expansions.

24 gigabytes is more than enough to run a local LLM for a small household or business.

This is "gaming PC" territory, not "space heater". I mean people already have PS5's and whatnot in their homes.

The hundreds of gigabytes thing exists because the big cloud LLM providers went down the increasing parameter count path. That way is a dead end and we've reached negative returns already.

Prompt engineering + finetunes is the future, but you need developer brains for that, not TFLOPs.

Post reply on HN