cuTile Rust: 安全、无数据竞态的 GPU 内核

cuTile Rust (cutile-rs) 通过将 Rust 的所有权和借用规则扩展到 GPU 启动边界，实现了内存安全且无数据竞态的 GPU 内核开发。通过在主机端将可变张量划分为不相交的分片，该系统确保每个 GPU tile 程序都能获得其特定数据段的独占可变引用，从而在编译时消除数据竞态。

内存安全与基于 Tile 的模型

cuTile Rust 用一种安全的、基于 tile 的编程模型取代了传统的 SIMT (Single Instruction, Multiple Threads) 编程——后者通常需要手动管理共享内存和使用 unsafe 代码。

跨启动边界的所有权

该系统通过执行严格的访问规则来维持 Rust 的安全保证：

  • 可变张量 (Mutable Tensors): 这些张量在内核启动前被划分为不相交的分片。每个 tile 被授予其分配的分片的独占 &mut 视图。
  • 不可变张量 (Immutable Tensors): 这些张量作为只读的 & 引用在所有 tile 之间共享。
  • 启动器 (Launchers): 生成的启动器在 GPU 工作进行时保留所有权，支持同步启动、异步流水线和 CUDA graph replay。

通过 CUDA Tile IR 进行 JIT 编译

内核使用 #[cutile::module] 宏进行定义，该宏会捕获 Rust AST 并将其嵌入到主机二进制文件中。在运行时，cuTile Rust 通过 CUDA Tile IR 将此 AST JIT 编译为 GPU cubin。虽然安全的 surface API 处理了大部分操作，但对于需要直接访问 Tile IR 操作层面的用户，也可以通过 unsafe intrinsics 进行底层控制。

性能基准测试

在 NVIDIA B200 上进行的评估表明，cuTile Rust 在不产生可感知的运行时开销的情况下提供了安全性，其性能水平足以与经过手工优化的底层实现相媲美。

  • 逐元素操作 (Element-wise Operations): 达到了 7 TB/s，约为峰值内存带宽的 91%。
  • GEMM (General Matrix Multiply): 达到了 2 PFlop/s（达到 dense f16 峰值的 92%），其表现与底层 Tile IR 变体仅差 0.3%，并与 cuBLAS 保持竞争关系。
  • 推理性能: 通过与 Hugging Face 合作，使用 cuTile Rust 构建的 Grout 推理引擎在 RTX 5090 上为 Qwen3-4B 实现了 171 tokens/s，并在 B200 上为 Qwen3-32B 在 batch-1 decode 期间实现了 82 tokens/s。

技术要求与兼容性

cuTile Rust 专为计算能力为 sm_80 (Ampere) 或更高版本的 NVIDIA GPU 设计。

硬件与软件依赖

  • GPU 架构: 最低要求 sm_80。支持 sm_100+ 需要 CUDA 13.1+；sm_8x 需要 CUDA 13.2；sm_90 需要 CUDA 13.3。
  • CUDA Toolkit: 建议使用版本 13.3，以利用 FP4 packing 和 block-scaled MMA 等特性。
  • Rust 版本: 1.89+。
  • 操作系统: Linux (已在 Ubuntu 24.04 上测试)。

与其他 Rust GPU 项目的对比

虽然像 cuda-oxide 这样的其他项目提供了针对 SIMT 风格内核的 Rust-to-CUDA 编译，但 cuTile Rust 专注于基于 tile 的方法。这种权衡放弃了显式的 warp 原语和手动共享内存控制，以换取一种可以在编译时进行安全性检查的语义模型。

项目状态与生态系统

cuTile Rust 目前是一个处于活跃开发早期阶段的研究项目。用户应预见潜在的 API 破坏性变更和功能不完整的情况。该项目已在 crates.io 上以 cutile 发布，并提供 Nix flake 以实现简化的环境搭建。

工作区架构

  • cutile: 用于编写和执行内核的主要面向用户 crate。
  • cutile-compiler: 通过 cutile-ir 处理 Rust 内核到可执行文件的编译。
  • cuda-async / cuda-core: 分别提供异步 CUDA 执行和符合惯用法且安全的 CUDA API。

Sources