Show HN: cuTile Rust：在 Rust 中实现安全、无数据竞争的 GPU 内核

背景

GPU 编程长期以来一直是高性能计算和 AI 领域的痛点。传统的 CUDA 或 ROCm 编程模型虽然强大，但缺乏内存安全保证，极易引发数据竞争（Data Races）和悬空指针等严重错误。尽管 Rust 以其“零成本抽象”和内存安全特性在系统级编程中崭露头角，但将其所有权（Ownership）和借用检查（Borrow Checking）机制延伸至 GPU 执行边界一直是一个巨大的挑战。

cuTile Rust（cutile-rs）正是为了解决这一矛盾而诞生的研究项目。它旨在将 Rust 的所有权纪律扩展到 GPU 启动边界，允许开发者使用惯用的 Rust 代码编写内存安全、无数据竞争的 GPU 内核。该项目由 Melih Elibol、Jared Roesch、Isaac Gelado、Eric Buehler 和 Michael Garland 等人开发，并发表了题为《Fearless Concurrency on the GPU》（GPU 上的无畏并发）的论文。

核心内容

cuTile Rust 是一个基于 Tile（瓦片/分块）的系统，它通过一种创新的方式将 Rust 的类型系统映射到 GPU 硬件抽象层。其核心机制包括以下几个方面：

1. 所有权在 GPU 边界的延伸

cuTile Rust 的核心创新在于它将 Rust 的所有权规则应用于 GPU 内核的启动过程：

• 可变张量分区：在启动前，可变的（mutable）张量被划分为不相交（disjoint）的块。这确保了不同线程块（Thread Block）或线程之间不会访问同一块内存，从而在编译期消除了数据竞争的可能性。
• 不可变张量共享：只读的（immutable）张量可以被安全地共享。
• 启动器保留所有权：生成的主机端启动器（Launcher）在 GPU 工作处于飞行状态（in-flight）时，依然保留所有权语义，支持同步启动、异步流水线以及 CUDA Graph 重放。

2. JIT 编译与 AST 嵌入

cuTile Rust 使用 #[cutile::module] 宏将捕获的 Rust AST（抽象语法树）嵌入到主机二进制文件中。当需要执行内核时，cuTile Rust 会通过 CUDA Tile IR 对该 AST 进行 JIT 编译，生成 GPU 的 cubin 二进制文件。这种设计使得内核代码可以直接作为 Rust 代码编写，同时享受 JIT 编译的灵活性。如果开发者需要更低级别的细粒度控制，仍然可以通过局部退出（local opt-outs）机制访问底层接口。

3. 代码示例解析

以下代码展示了如何使用 cuTile Rust 编写一个简单的向量加法内核：

use cutile::prelude::*;

#[cutile::module]
mod kernel {
    use cutile::core::*;

    #[cutile::entry()]
    fn add<const B: i32>(
        z: &mut Tensor<f32, { [B] }>, // 独占的可变输出张量，大小为 B
        x: &Tensor<f32, { [-1] }>,   // 共享的只读输入张量
        y: &Tensor<f32, { [-1] }>,   // 共享的只读输入张量
    ) {
        let tx = load_tile_like(x, z); // 加载与输出分块匹配的输入瓦片
        let ty = load_tile_like(y, z);
        z.store(tx + ty);              // 存储结果
    }
}

fn main() -> Result<(), Error> {
    let x = api::ones::<f32>(&[1024]);
    let y = api::ones::<f32>(&[1024]);
    // 将可变输出张量 z 分区为 128 元素的块
    let z = api::zeros::<f32>(&[1024]).partition([128]);
    
    // 调用内核，.sync() 会触发 JIT 编译并执行
    let (_z, _x, _y) = kernel::add(z, x, y).sync()?;
    Ok(())
}

• 内核签名：z 是独占的可变输出，x 和 y 是共享的只读输入。这种签名直接将访问纪律带入设备代码。
• 自动网格推断：启动网格（Launch Grid）(8, 1, 1) 是根据分区自动推断的。由于总大小为 1024，分区大小为 128，因此需要 $1024 \div 128 = 8$ 个瓦片（Tiles）。
• 执行流程：主机代码构建惰性张量操作，对可变输出进行分区，调用 .sync() 后，cuTile Rust 负责 JIT 编译并执行内核。

4. 性能表现

在 NVIDIA B200 GPU 上，cuTile Rust 展示了极具竞争力的性能：

• 内存带宽利用率：对于逐元素操作（Element-wise operations），达到了 7 TB/s，约为峰值内存带宽的 91%。
• 计算性能：对于 GEMM（通用矩阵乘法），达到了 2 PFlop/s，约为密集 f16 峰值性能的 92%。这一结果与 cuBLAS 相当。
• 安全性开销：微基准测试显示，cuTile Rust 在提供安全保证的同时，没有引入可测量的运行时开销。在 M=N=K=8192 的持久化 GEMM 测试中，安全 Rust 版本达到了 2.07 PFlop/s，仅比对应的低级 Tile IR 变体低 0.3%。

5. 实际应用案例：Grout 推理引擎

论文还评估了由 Hugging Face 与 cuTile Rust 合作构建的 Qwen3 推理引擎 Grout。

• 在 NVIDIA GeForce RTX 5090 上，Qwen3-4B 的 batch-1 解码速度达到 171 tokens/s。
• 在 B200 上，Qwen3-32B 的解码速度达到 82 tokens/s。
• 根据 HBM 屋顶线（Roofline）分析，Grout 在内存受限的推理任务中达到了具有竞争力的最先进（SOTA）性能。

关键要点

• 内存安全与零开销：cuTile Rust 通过编译期的所有权检查消除了 GPU 编程中的数据竞争风险，且在 B200 上的基准测试表明其运行时开销几乎为零（<0.3%）。
• 基于 Tile 的抽象：系统围绕张量分区（Tensor Partitions）和张量核心导向的操作构建，自动处理网格推断和内存分区，简化了内核开发。
• JIT 编译架构：利用 #[cutile::module] 宏嵌入 Rust AST，并通过 CUDA Tile IR 进行 JIT 编译，既保留了 Rust 的开发体验，又实现了 GPU 代码的高效生成。
• 广泛的硬件支持：
  • 支持 NVIDIA GPU 计算能力 sm_80 及以上（最低支持 sm_80，推荐 sm_100+ 即 Blackwell 架构）。
  • 需要 CUDA 13.3+（支持 sm_80+ 及 CUDA Tile IR 13.3 特性如 FP4 打包和块缩放 MMA）。
  • 需要 Rust 1.89+。
  • 主要支持 Linux（已在 Ubuntu 24.04 上测试）。
• 早期阶段但活跃开发：该项目目前处于早期阶段，API 可能会发生破坏性变更，存在已知 Bug 和功能不完整的情况，但作者鼓励社区通过反馈和贡献（参考 CONTRIBUTING.md）共同塑造其发展方向。
• 生态系统集成：提供了 Nix flake 以简化开发和环境设置，并包含完整的测试套件、示例（如 saxpy, async_gemm）和基准测试工具。

意义与影响