Compiling LLMs into a MegaKernel: A path to low-latency inference 1 year ago
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