Luminal 是一款高性能的通用推理编译器。

使用方法

use luminal::prelude::*;
// 创建计算图
let mut cx = Graph::new();
let a = cx.tensor((3, 1));
let b = cx.tensor((1, 4));

let c = a.matmul(b).output();

// 编译
cx.build_search_space::<NativeRuntime>();
let mut rt = cx.search(NativeRuntime::default(), 1);

// 设置输入张量
rt.set_data(a, vec![1.0, 2.0, 3.0]);
rt.set_data(b, vec![1.0, 2.0, 3.0, 3.0]);

// 运行
rt.execute(&cx.dyn_map);

// 获取输出张量
println!("结果: {:?}", rt.get_f32(c));

快速入门

Llama 3 8B

以下是一个使用 Luminal 在 CUDA 上本地运行 Llama 3 8B 的快速示例：

cd ./examples/llama
cargo run --release

特性

速度

Luminal 可以在 H100 上以接近理论最大性能的 80% 运行 Q8 版本的 Llama 3 8B。我们的目标是成为在任何设备上运行任何模型时最快的机器学习框架。

简洁性

Luminal 的核心始终是极简的。你可以在一个下午内理解整个核心库。

RISC 风格架构

Luminal 中的所有操作都可以归结为 14 种基本算子：

• 一元算子：Log2, Exp2, Sin, Sqrt, Recip
• 二元算子：Add, Mul, Mod, LessThan
• 其他算子：SumReduce, MaxReduce, Iota, Gather, Cast

这些算子足以支持 Transformer、卷积神经网络以及几乎所有流行的模型。

搜索

最好的启发式就是没有启发式。我们尝试搜索所有可能的决策，以赋予编译器最大的灵活性来发现复杂的优化。这使我们能够自动推导出 Flash Attention 等复杂重写规则。同时，这也让我们能够在未来的很长一段时间内保持代码体积非常小，并且在性能上超越那些拥有大量手写内核的大型框架。

原生实现

当前的机器学习生态系统过于碎片化，而解决方案并不是再增加一层抽象。Luminal 使用 Rust 编写，直接与 CUDA / Metal API 交互。没有任何间接层或抽象、Docker 容器或虚拟环境，只是一份静态链接的 Rust crate。

与 PyTorch 对比验证

正确性至关重要。我们编写了尽可能多的测试用例，以覆盖所有算子，并验证它们的行为与 PyTorch 的等效实现一致。（仍需改进！）

理念

为什么它看起来与其他深度学习库如此不同？

大多数深度学习库都是“急切执行”模式，即每次调用算子都会直接对数据进行操作。在 PyTorch 中，当你看到 x + y 时，加法运算实际上就在那里立即发生。这种方式非常适合调试，因为它完全符合大多数开发者的预期。

然而，这种模式并不利于性能。对开发者来说合理的方式并不一定适合机器，就像没有人会手写汇编代码一样。大多数库试图通过添加算子融合或 JIT 编译来改变编译流程，使其更适合机器。但事实证明，即使对于 PyTorch 来说，这也是 极其 困难 的事情 甚至 难以实现！

编译一切

Luminal 的核心原则之一是提前编译。只要可能，就将所有内容推到编译时，不留给运行时处理。Luminal 的方法更类似于 XLA 和 tinygrad。这里一切都是静态的。当你写下表达式 x + y 时，实际上并不会进行任何计算。该操作会被记录到一个有向无环计算图中，供后续执行。只有当调用 graph.execute() 时，计算才会真正发生。但这不就是惰性执行吗？ 是的！但在 Luminal 中，一切都以这种方式进行。所有的神经网络都被构建为一个或几个静态计算图，经过编译后稍后再执行。

那么为什么呢？

这样做的一个后果是，实际运行的计算可能会与编写的代码截然不同。由于我们有一个完整的神经网络被表示在计算图中，编译器拥有全局视角。这意味着我们可以将大部分机器学习的复杂性交给编译器来处理。例如，设备、数据类型和执行调度都由编译器负责。甚至自动微分也是由编译器完成的！

现在我们可以做到：

• 大规模的内核融合
• 根据形状在运行时编译特定的内核
• 设备和数据类型的处理由编译器完成（只需运行 CUDA 编译器将图转换为使用 CUDA 内核，再运行 fp16 编译器将其转换为半精度内核）
• 网络可以用通用代码编写，但却能在高度定制化的架构上快速编译和运行（试着写一个既能支持 TF32 数据类型又能运行在 TPU 上的 PyTorch 网络吧；准备好迎接 if 语句地狱……）

当前进展

• 搜索功能已部分合并。我们正处于 1.0 和 2.0（搜索）之间，预计将在未来一个月左右完成。
• 支持 Metal 和 CUDA，分别用于在 Mac 和 Nvidia GPU 上以全精度和半精度运行模型。
• 提供基于计算图的完整训练支持，包含自动微分。
• Llama 3、Phi 3、Whisper 和 Yolo v8 已在 examples/ 目录中实现。请参阅上方说明以了解如何运行。
• 我们在 luminal_nn 中提供了一个小型的神经网络模块库，其中包括 Transformer 模块。
• 在 hl_ops 中实现了大量的高级算子。我们正努力匹配 PyTorch API 中最常用的约 80% 的功能。

路线图上的其他事项：

• 扩展搜索空间，以更灵活地利用 Tensor Core
• 使 CUDA 功能与 Metal 功能对齐
• 添加 Blackwell 的内置函数，如 TMEM 和 TMA
• 构建 ROCm 后端
• 构建基准测试套件，用于与其他库进行比较
• 实现分布式数据、流水线和张量并行
• 在 LLM 推理和训练方面超越 PyTorch 2.0 的性能
• 编写量子光子回退编码器的编译器
• 构建戴森球群

许可证

根据您的选择，本文件依据 Apache 证书版本 2.0（http://www.apache.org/licenses/LICENSE-2.0）或 MIT 证书（http://opensource.org/licenses/MIT）授权。除遵循这些条款外，不得复制、修改或分发本文件。

Luminal 快速上手指南

Luminal 是一个高性能的通用推理编译器，采用 Rust 编写，直接调用 CUDA/Metal API，旨在通过静态计算图和编译时搜索优化，在任何设备上实现极致的机器学习推理速度。

环境准备

• 操作系统:
  • Linux (推荐用于 NVIDIA GPU)
  • macOS (支持 Apple Silicon Metal)
• 硬件要求:
  • NVIDIA GPU (需安装 CUDA Toolkit) 或 Apple M 系列芯片。
• 前置依赖:
  • Rust: 必须安装最新稳定版 Rust 工具链。
  • CUDA (仅限 NVIDIA): 确保 nvcc 在路径中且版本与显卡驱动兼容。
  • Xcode Command Line Tools (仅限 macOS): 用于 Metal 开发。

安装 Rust (如未安装):

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
source $HOME/.cargo/env

国内加速建议: 如果下载 Rust 组件缓慢，可配置国内镜像源：

export RUSTUP_DIST_SERVER=https://mirrors.tuna.tsinghua.edu.cn/rustup
export RUSTUP_UPDATE_ROOT=https://mirrors.tuna.tsinghua.edu.cn/rustup/rustup

安装步骤

Luminal 作为一个 Rust crate 发布，你可以直接在项目中添加依赖，或克隆仓库运行示例。

方式一：作为依赖库安装

在你的 Rust 项目 Cargo.toml 中添加：

[dependencies]
luminal = "0.2" # 请根据 crates.io 最新版本调整

方式二：克隆源码运行示例

获取源代码并验证环境：

git clone https://github.com/jafioti/luminal.git
cd luminal
cargo build --release

基本使用

Luminal 采用“先定义计算图，后编译执行”的模式。以下是最小化的使用示例，展示如何创建张量、构建矩阵乘法图、编译并运行。

use luminal::prelude::*;

fn main() {
    // 1. 创建计算图
    let mut cx = Graph::new();
    
    // 定义输入张量形状
    let a = cx.tensor((3, 1));
    let b = cx.tensor((1, 4));

    // 定义操作：矩阵乘法并标记为输出
    let c = a.matmul(b).output();

    // 2. 编译阶段
    // 构建搜索空间并使用原生运行时进行搜索优化
    cx.build_search_space::<NativeRuntime>();
    let mut rt = cx.search(NativeRuntime::default(), 1);

    // 3. 设置输入数据
    rt.set_data(a, vec![1.0, 2.0, 3.0]);
    rt.set_data(b, vec![1.0, 2.0, 3.0, 3.0]);

    // 4. 执行计算
    rt.execute(&cx.dyn_map);

    // 5. 获取并打印结果
    println!("Result: {:?}", rt.get_f32(c));
}

运行大模型示例 (Llama 3 8B)

如果你已配置好 CUDA 环境，可以直接运行官方提供的 Llama 3 示例：

cd ./examples/llama
cargo run --release

注意: 首次运行时编译器会进行搜索和优化，可能需要一定时间，后续运行将直接使用编译好的内核。