显著提升小LLM的性能，所有kernel集成一个巨大kernel！
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
以下文章来源于微信公众号：oldpan博客
作者：小潘潘
链接：https ://mp. weixin.qq. com/s/kyXsoBC_l6aP0qYp6Uuyjg
本文仅用于学术分享，如有侵权，请联系后台作删文处理
导读
在传统LLM推理依赖多次GPUkernel启动和通信，导致延迟高、效率低。本项目提出一种自动编译器，可将推理流程融合为单个megakernel，在一次GPU启动中完成计算与通信，将延迟降低1. 2–6. 7倍。使用简便，仅需少量Python代码即可实现端到端加速。
一个LLM编译器，它可以自动将LLM推理转换为单个megakernel——一个融合的GPUkernel，能在一次启动中执行所有必要的计算和通信。这种端到端的GPU融合方法将LLM推理延迟降低了1. 2-6. 7倍。我们的编译器使用简单——只需几十行Python代码就可以将LLM编译成高性能megakernel。
核心理念是什么？传统的LLM系统通常依赖于一系列GPUkernel启动和外部通信调用，导致硬件利用率不足。我们的编译器可以自动将这些操作（跨越多层、多次迭代和多个GPU）融合为一个megakernel。这种设计消除了启动开销，实现了细粒度的软件流水线，并使计算与跨GPU通信重叠。
项目地址:
https ://github. com/mirage-project/mirage/tree/mpk
降低LLM推理延迟最有效的方法之一是将所有计算和通信融合到单个megakernel——也称为persistentkernel中。在这种设计中，系统只需启动一个GPUkernel就能执行整个模型——从逐层计算到跨GPU通信——无需中断。这种方法带来几个关键性能优势：
消除kernel启动开销，即使在多GPU设置中，也可以避免重复的kernel调用；
实现跨层软件流水线，使kernel能在计算当前层的同时开始加载下一层的数据；
重叠计算和通信，因为megakernel可以同时执行计算操作和跨GPU通信以隐藏延迟。
我们能通过编译来自动化这个过程吗？基于这个问题，我们来自CMU、UW、Berkeley、NVIDIA和清华的团队开发了**MiragePersistentKernel[ 24](MPK)**——一个编译器和运行时系统，可以自动将多GPULLM推理转换为高性能megakernel。MPK释放了端到端GPU融合的优势，同时只需要开发者付出最少的手动努力。
MPK的一个关键优势是通过消除kernel启动开销并最大程度地重叠计算、数据加载和跨GPU通信，实现了LLM推理的极低延迟。
图1. 比较MPK与现有系统的LLM解码延迟。我们使用了39个token的提示词并生成了512个token，未使用推测解码。
图1展示了MPK与现有LLM推理系统在单GPU和多GPU配置下的性能对比。在单个NVIDIAA10040GBGPU上，MPK将每个token的解码延迟从14. 5毫秒（这是像vLLM和SGLang这样的优化系统所能达到的水平）降低到了12. 5毫秒，接近10毫秒的理论下限（基于以1. 6TB/s的内存带宽加载16GB的权重）。
除了单GPU优化之外，MPK将计算和跨GPU通信融合到单个megakernel中。这种设计使MPK能够最大程度地重叠计算和通信。因此，MPK相比当前系统的性能提升会随着GPU数量的增加而增加，这使其在多GPU部署中特别有效。
本博客的其余部分将深入探讨MPK的工作原理：
第1部分介绍MPK编译器，它将LLM的计算图转换为优化的任务图；
第2部分介绍MPK运行时，它在megakernel中执行这个任务图以实现高吞吐量和低延迟。
图2. MPK编译器将LLM的computationgraph（在PyTorch中定义）转换为优化的细粒度taskgraph，以暴露最大并行性。右侧展示了一个替代方案——但这是次优的taskgraph，它引入了不必要的数据依赖和全局同步障碍，限制了跨层的流水线机会。
在MPKtaskgraph中：
每个task（如图2中的矩形所示）代表分配给单个GPUstreamingmultiprocessor(SM)的计算或通信单元。
每个event（显示为圆圈）代表任务之间的同步点。
每个task都有一个指向triggeringevent的出边，当所有相关task完成时该event被激活。
每个task还有一个来自dependentevent的入边，表示task可以在event激活后立即开始执行。
MPK包含一个on-GPU运行时系统，它在单个GPUmegakernel中完全执行taskgraph，实现了对任务执行和调度的细粒度控制，在推理过程中无需任何kernel启动。
每个worker在一个SM上运行并维护一个专用任务队列。它遵循一个简单但高效的执行循环：
从队列中获取下一个任务。
执行任务（例如矩阵乘法、attention或跨GPU数据传输）。
任务完成时通知触发事件。
重复以上步骤。
这种设计确保workers保持充分利用，同时使任务执行能够在各层和操作之间异步进行。
出队已满足依赖关系的已激活事件（即所有前置任务已完成）。
启动依赖于已激活事件的任务集。
这种分散式调度机制最小化了协调开销，同时实现了跨SM的可扩展执行。
图3. MPK运行时在megakernel中执行taskgraph。
图3展示了MPK的执行时间线。每个矩形代表在worker上运行的任务；每个圆圈代表一个事件。当任务完成时，它会增加其对应触发事件的计数器。当事件计数器达到预定阈值时，该事件被视为已激活并进入scheduler的事件队列。然后scheduler启动依赖于该事件的所有下游任务。
这种设计实现了细粒度软件流水线和计算与通信的重叠。例如：
不同层的Matmul任务可以与attention任务并行执行。
一旦获得部分matmul结果，就可以开始Allreduce通信。
由于所有调度和任务转换都发生在单个kernel上下文中，任务之间的开销极低——通常仅为1-2微秒——从而实现了多层、多GPULLM工作负载的高效执行。
我们对MPK的愿景是使megakernel编译既易于使用又具有高性能。目前，您只需几十行Python代码就可以将LLM编译成megakernel——主要是用于指定megakernel的输入和输出。我们对这个方向感到兴奋，还有更多值得探索的地方。以下是我们正在积极开发的几个关键领域：
处理工作负载动态性。MPK目前构建静态taskgraph，这限制了它处理动态工作负载（如**mixture-of-experts(MoE)**模型）的能力。我们正在开发新的编译策略，使MPK能够支持megakernels内的动态控制流和条件执行。
高级调度和任务分配：MPK在task级别实现了新的细粒度调度。虽然我们当前的实现使用简单的轮询调度来在SM之间分配任务，但我们看到了高级调度策略的令人兴奋的机会——例如优先级感知或吞吐量优化策略——用于延迟SLO驱动服务或混合批处理等用例。
我们相信MPK代表了GPU上LLM推理工作负载编译和执行方式的根本性转变，我们渴望与社区合作推进这一愿景。
让LLM始终运行在同一个kernel里，极大减少了kernellaunch和数据搬运开销，实现LLM推理的低延迟与高吞吐。个人觉得比较适合1B、3B、7B这种可以单卡容纳的小模型，这种级别的LLM有overhead可以被巨大kernel优化，但是比较大的模型（32B、70B）相应的overhead不明显，一个kernel和多个kernel对性能的影响不大。
而且这个方式不够灵活，对资源利用会比较大（在编译kernel的时候为了性能会强行设置然后assert一些资源，和之前trt的做法一些，有些牺牲显存换取性能的意思，这个会更明显些），多gpu适配难度比较大。
比较适合小的LLM。
PyTorch:https ://pytorch. org/
Triton:https ://github. com/triton-lang/triton
TVM:https ://tvm. apache.org/
NCCL:https ://github. com/NVIDIA/nccl
NVSHMEM:https ://developer. nvidia.com/nvshmem
FlashInfer:https ://github. com/flashinfer-ai/flashinfer
FlashAttention:https ://github. com/Dao-AILab/flash-attention
Triton:https ://github. com/triton-lang/triton
**MiragePersistentKernel:https ://github. com/mirage-project/mirage
Miragekernelsuperoptimizer:https ://github. com/mirage-project/mirage
欢迎加入《AI未来星球》，一起成长
扫描下方二维码即可加入~
真诚分享AI落地过程（AI商机->项目签约->算法开发->产品开发->实施运维）中的各方面经验和踩过的坑。
你可以获得什么？
1、大白之前花费10W+购买，AI行业各场景私有数据集下载，星球内倾情分享；2、AI行业研发、产品、商业落地问题咨询（目前AI公司创业中），都可获高质量解答，有效期一年，无限次提问，有问必答。3、定期邀请AI行业各类嘉宾分享，创业/商业等方面的经验！
帮助你解决遇到的实际问题，升职加薪！
大家一起加油！
