Mirage Persistent Kernel：把 LLM 推理压进单个 GPU 内核

它解决什么：多内核启动的延迟税

LLM 推理通常由几十个独立内核组成，每个内核负责一个算子，比如 RMSNorm、线性层、注意力。每个内核启动都有固定开销，GPU 在等待启动时处于空闲。Mirage Persistent Kernel（MPK）的目标是把这些内核融合成一个 megakernel，即一个在单次启动中完成所有计算和通信的 GPU 内核。README 声称这能将推理延迟降低 1.2 到 6.7 倍。这个数字来自项目方，我没有独立验证。MPK 面向的是多 GPU 场景，因为通信（如 all-reduce）通常需要多次内核启动和同步，融合后可以减少这些开销。它适合已经用 PyTorch 和 Hugging Face 模型、但延迟不达标的工程师。如果你只是跑单 GPU 小模型，内核启动开销可能不是瓶颈，MPK 的复杂度就不划算。

工作机制：持久内核与任务图

MPK 的核心是持久内核（persistent kernel），它不是一个一次性的内核，而是一个常驻 GPU 的运行时。你定义计算图，MPK 把它编译成一组任务，这些任务在 GPU 上由调度器分配。关键参数是 num_workers、num_local_schedulers 和 num_remote_schedulers，它们必须匹配物理 SM 数量，公式是 num_workers + (num_local_schedulers + num_remote_schedulers) / 4。这意味着你需要知道 GPU 的 SM 数，比如 A100 有 108 个 SM，你得据此设置参数。计算图由融合算子组成，例如 rmsnorm_linear_layer 把 RMSNorm 和线性层合并成一个任务。每个算子调用需要指定 grid_dim 和 block_dim，即线程块和线程数。README 建议线程块总数是 worker 数的倍数，以避免某些 worker 空闲。编译后调用 mpk.compile() 生成 CUDA 内核，然后 mpk() 执行。整个流程中，MPK 维护两个元张量：step 记录解码步数，tokens 存储提示和生成的 token。

快速上手：从源码安装到运行 demo

目前没有预编译的 wheel 包，README 明确说正在开发中。安装必须从源码编译，命令如下：git clone --recursive --branch mpk https://www.github.com/mirage-project/mirage，然后 cd mirage 并 pip install -e . -v。安装后需要设置环境变量 export MIRAGE_HOME=$(pwd)。仓库附带一个 Qwen3-8B 的 demo 脚本。先用原生 Triton 和 FlashInfer 内核运行：python demo/qwen3/demo.py。然后用 MPK 编译并执行：python demo/qwen3/demo.py --use-mirage。如果要可视化执行时间线，加 --profiling 参数。这个 demo 展示了如何用几十行 Python 定义内核的输入和输出。你需要定义一个 PersistentKernel 对象，指定 world_size 和 mpi_rank，然后通过 attach_input 绑定已有的 PyTorch 张量，或用 new_tensor 分配新张量。io_category 参数必须设为 cuda_tensor 或 nvshmem_tensor，后者用于跨 GPU 通信。整个过程对熟悉 PyTorch 的人不算陌生，但 SM 数量的要求增加了配置负担。

关键 API 与张量管理

MPK 的 API 设计围绕张量生命周期。attach_input 把外部 PyTorch 张量引入 megakernel，你需要给每个张量一个 name，MPK 在生成的 CUDA 代码中用它引用。new_tensor 分配新张量，必须指定 dims、dtype 和 io_category。io_category 只有两种选择：cuda_tensor 用于本 GPU 内存，nvshmem_tensor 用于远程 GPU 访问。这意味着跨 GPU 通信（如 all-reduce）必须使用 nvshmem，这要求你的环境支持 NVSHMEM，一个基于 CUDA 的对称内存库。编译时，MPK 会把计算图转换成 CUDA 内核，但 README 没有说明编译时间有多长，也没有说明是否支持动态形状。从代码看，tokens 张量形状是 [num_requests, seq_length]，seq_length 可能随解码变化，但 MPK 如何调整任务图没有细节。这是一个潜在的局限：如果输入序列长度变化大，megakernel 可能无法高效适应。

真正的限制：配置复杂且依赖特定环境

MPK 的配置参数不是自动推导的。你必须手动设置 num_workers、num_local_schedulers 和 num_remote_schedulers，且它们必须精确匹配物理 SM 数。公式中的除法意味着某些组合可能不合法，比如 num_local_schedulers + num_remote_schedulers 必须是 4 的倍数。这增加了部署难度，尤其是在异构 GPU 集群上，每个 GPU 的 SM 数可能不同。另一个限制是 MPK 目前只支持两个元张量：step 和 tokens。如果你的模型需要额外的状态（比如 KV cache 的偏移），你需要在计算图中手动管理。README 没有提到 KV cache 的处理方式，这让我怀疑 MPK 是否适用于长序列生成，因为 KV cache 通常需要动态分配。此外，MPK 依赖 nvshmem，这不是所有集群都支持。如果集群没有配置 NVSHMEM，跨 GPU 通信就无法工作。最后，1.2 到 6.7 倍的加速是项目方声称的，我没有看到基准测试的细节，比如 GPU 型号、批量大小、序列长度。这些变量对结果影响很大，你必须在自己的硬件上验证。

替代方案：Triton 与 FlashInfer 的对比

MPK 的替代方案是使用 Triton 和 FlashInfer，这也是 demo 脚本默认的运行方式。Triton 是一个 GPU 编程语言，允许你编写自定义算子，但每个算子仍然是一个独立内核。FlashInfer 是一个注意力内核库，优化了注意力计算，但同样没有融合整个模型。区别在于融合粒度：Triton 和 FlashInfer 融合单个算子或小算子组，MPK 融合整个模型。这意味着 MPK 减少了内核启动次数，但增加了编译复杂度和运行时调度开销。如果你只是需要优化注意力，FlashInfer 可能就足够了，它不需要你配置 SM 数量。如果你需要自定义算子，Triton 更灵活，但你需要手动管理多 GPU 通信。MPK 的卖点是把通信也融合进去，但代价是你必须接受它的抽象。另一个替代是手写 CUDA，但 README 没有提供这样的示例，而且手写 megakernel 的难度极高。

维护与升级成本，以及许可证

MPK 的代码库以 Cuda 为主要语言，这意味着你或你的团队需要 CUDA 知识来调试问题。目前没有预编译 wheel，每次安装都要从源码编译，这增加了 CI/CD 的复杂度。仓库默认分支是 mpk，不是 master 或 main，这表明项目仍处于积极开发阶段。最近一次发布是 osdi26-ae-v1，这是 OSDI 2026 的 artifact 版本，说明项目与学术论文绑定。这意味着 API 可能会随论文修订而变动，生产环境需要锁定版本。许可证是 Apache-2.0，这是一个宽松的许可证，允许商业使用和修改，但你需要保留版权声明。没有看到关于贡献或代码审查流程的细节，只有 Slack 频道和 issue 链接。升级成本方面，由于依赖 nvshmem 和特定 CUDA 版本，升级到新 GPU 或 CUDA 版本可能需要重新编译和调整参数。

编辑结论

Mirage Persistent Kernel 适合那些已经用 PyTorch 和 Hugging Face 模型、且推理延迟受多内核启动开销主导的团队。它不适合刚接触 GPU 编程的人，因为你需要理解 SM 数量、调度器和 nvshmem 等概念。在采用前，先确认你的 GPU 驱动支持 nvshmem，并在你的模型和批量大小上复现 demo 中的延迟收益。MPK 目前没有预编译 wheel，安装必须从源码编译，这要求你的 CUDA 环境与仓库版本匹配。如果你只需要单 GPU 上的算子融合，直接使用 Triton 或 FlashInfer 更简单；只有当跨 GPU 通信成为瓶颈时，MPK 的 megakernel 思路才值得投入。最终判断：这是一个研究驱动的项目，OSDI 论文背书，但生产部署前必须自行验证其 1.2 到 6.7 倍的加速是否在你的工作负载上成立。

官方来源

1. Official documentation
2. Official README
3. Project repository
4. Release notes