Mirage 项目使用教程

1. 项目介绍

Mirage 是一个多层次的张量代数超级优化器，旨在自动生成 PyTorch 程序的高效 GPU 内核。通过超级优化技术，Mirage 能够搜索与输入程序功能等效的潜在 GPU 内核空间，并发现高度优化的内核候选。这种方法使得 Mirage 能够找到新的自定义内核，这些内核在性能上优于现有的专家设计的内核。

2. 项目快速启动

安装 Mirage

最快捷的安装方式是从 pip 安装最新稳定版本：

pip install mirage-project

或者，你也可以从源代码安装：

git clone --recursive https://github.com/mirage-project/mirage.git
cd mirage
pip install -e . -v

快速启动示例

Mirage 可以自动生成任意 PyTorch 程序的高效 GPU 内核。以下是一个使用 Mirage 生成融合 RMSNorm 和 Linear 内核的示例，以加速 Transformer 计算。

import torch
import mirage

# 定义一个函数来获取 Mirage 生成的优化内核
def get_mirage_kernel(batch_size, output_dim):
    graph = mirage.new_kernel_graph()
    X = graph.new_input(dims=(batch_size, 4096), dtype=mirage.float16)
    Wqkv = graph.new_input(dims=(4096, output_dim), dtype=mirage.float16)
    Y = graph.rms_norm(X, normalized_shape=(4096,))
    Z = graph.matmul(Y, Wqkv)
    graph.mark_output(Y)
    return graph.superoptimize()

# 使用 Mirage 生成内核
kernel_1 = get_mirage_kernel(batch_size, output_dim=Wqkv.shape[-1])
kernel_2 = get_mirage_kernel(batch_size, output_dim=W13.shape[-1])

# 在 PyTorch 程序中使用生成的内核
Z = kernel_1(inputs=[X, Wqkv])
O = attention(Z)
V = kernel_2(inputs=[Z, W13])
V1, V3 = V.chunk(2, -1)
output = torch.matmul(torch.silu(V1) * V3, W2)

3. 应用案例和最佳实践

加速 Transformer 计算

Mirage 特别适用于加速 Transformer 模型中的计算。通过生成融合 RMSNorm 和 Linear 的内核，Mirage 可以将 Transformer 层的计算速度提升 1.5 到 1.7 倍。

自定义内核生成

Mirage 不仅限于预定义的操作，用户可以通过编写几行 Python 代码来描述所需的计算，Mirage 将自动生成优化的 GPU 内核。

4. 典型生态项目

PyTorch

Mirage 与 PyTorch 深度集成，生成的内核可以直接在 PyTorch 程序中使用，无需额外修改。

CUDA

Mirage 生成的内核基于 CUDA，充分利用 NVIDIA GPU 的计算能力，提供高性能的计算加速。

TensorFlow

虽然 Mirage 主要针对 PyTorch，但其优化技术也可以应用于 TensorFlow，通过适当的接口转换，Mirage 生成的内核可以在 TensorFlow 中使用。

通过以上步骤，你可以快速上手并使用 Mirage 项目，享受其带来的高性能计算加速。