告别 CUDA，无需 Triton！Mirage：人人都能成为 GPU 编程大师？

AIxiv 专栏 | 机器之心报道

近年来，随着 GPU 加速器的不断发展以及以大语言模型为代表的生成式 AI 应用的不断推广，开发高性能 GPU 内核来优化 PyTorch程序的计算效率变得越来越重要。然而，编写高性能 GPU 内核需要高水平的 GPU 专业知识和大量的工程开发经验，这使得这项任务成为了一项门槛极高的工作，只有专门的 GPU 专家才能胜任。

为了解决这一问题，来自卡内基梅隆大学的 Catalyst Group 团队发布了 Mirage 项目，这是一款基于 SuperOptimization技术的 PyTorch 算子编译器，它可以自动生成高效的 GPU 内核，而无需用户编写任何 CUDA 或 Triton 代码。

Mirage 的核心优势在于它能够自动搜索可能的 GPU 内核实现，并从中选择最优的方案。 它的搜索空间不仅包括现有的手动设计的注意力内核（如 FlashAttention 和 FlashDecoding），还包括在某些场景中比目前的手写版本快多达 3.5 倍的其他实现。

Mirage 的使用方式非常简单： 用户只需编写几行Python 代码来描述所需的计算过程，Mirage 就会自动生成相应的 GPU 内核，并将其集成到 PyTorch 程序中。

以下是一个使用 Mirage 生成 FlashAttention 内核的示例：

“`python
import mirage as mi

graph = mi.newkernelgraph()
Q = graph.newinput(dims=(64, 1, 128), dtype=mi.float16)
K = graph.newinput(dims=(64, 128, 4096), dtype=mi.float16)
V = graph.new_input(dims=(64, 4096, 128), dtype=mi.float16)

A = graph.matmul(Q, K)
S = graph.softmax(A)
O = graph.matmul(S, V)

optimized_graph = graph.superoptimize()
“`

Mirage 生成的 GPU 内核可以直接在 PyTorch 张量上操作，并可以在 PyTorch 程序中直接调用：

“`python
import torch

input_tensors = [
torch.randn(64, 1, 128, dtype=torch.float16, device=’cuda:0′),
torch.randn(64, 128, 4096, dtype=torch.float16, device=’cuda:0′),
torch.randn(64, 4096, 128, dtype=torch.float16, device=’cuda:0′)
]

Launch the Mirage-generated kernel to perform attention

output = optimizedgraph(inputtensors)
“`

Mirage 的出现为机器学习系统工程师带来了以下三个主要优势：

更高的生产力： 随着 GPU架构的不断发展，现代 GPU 编程需要持续学习大量的专业知识。Mirage 的目标是提高机器学习系统工程师的生产力，他们只需在 PyTorch 层面描述所需的计算，Mirage 便会自动生成适用于各种 GPU 架构的高性能实现。
更佳的性能： Mirage的 SuperOptimization 技术能够自动搜索最优的 GPU 内核实现，并生成比手动编写代码更快的内核。
更低的开发成本： Mirage 可以帮助机器学习开发者节省大量的时间和精力，从而将更多的时间投入到模型设计和算法优化等更重要的工作中。

Mirage 的出现标志着 GPU 编程领域的一次重大变革，它将使更多的人能够轻松地开发高性能 GPU 内核，并加速人工智能应用的落地。 随着 Mirage 的不断发展，我们有理由相信，未来每个人都能够成为 GPU 编程大师，并为人工智能的未来贡献自己的力量。

参考文献：