← 文章 / 编程开发
NVIDIA 开发者博客 4小时前 · 2026-09-09 02:11:20 · 3 阅读

CUDA Rust 来了:GPU Kernel 编程的两条路线

2026 年 9 月,NVIDIA 宣布将投入原生 Rust GPU 编程。CUDA C++ 和 CUDA Python 已经是成熟的企业级工具链,而 NVIDIA 将在 2027 年及以后持续发展和完善 CUDA Rust。

AI 的系统层涵盖推理引擎、服务基础设施、驱动程序和 Agent 运行时,而且随着模型和技术的变化不断迭代。其中越来越多部分是用 Rust 编写的——它能在编译期捕获整类 bug,同时不牺牲性能。

NVIDIA 也在做同样的转变。Nova Linux 驱动是 Rust 写的,NVIDIA Dynamo 构建在 Rust 核心之上,NVTX 也提供了 Rust 绑定。

唯一的例外是 GPU kernel。你可以用 Rust 启动 kernel,但 kernel 本身往往得用别的语言来写。

NVIDIA CUDA Rust 填补了这一空白。GPU kernel 可以直接用 Rust 编写,原生编译到 PTX,而不是对其他语言代码的封装。

与 CUDA 本身的两条路径相对应,用 Rust 写 kernel 也有两条路径。SIMT 就是你在 CUDA C++ 或 numba-cuda 中已经在用的模型:描述单个线程做什么,然后启动成千上万个线程。Tile 是较新的编程模型,同样提供 C++Python 接口。所有这些前端都让你只描述一个数据 tile 的计算逻辑,剩下的工作交给 Tile IR 编译器完成。

选择从哪个入手时,优先考虑 Tile。编译器会决定 tile 如何映射到不同架构,源代码中不必包含架构相关的选择;当你需要那种级别的控制权、或者想自己管理内存和线程时,再退回到 SIMT。

选语言和选模型是两个独立的问题。选择最适合你现有技术栈的 CUDA 接口即可。下面两个项目面向的就是以 Rust 为技术栈的场景。我们计划支持跨语言互操作,所以无论选哪个都不会把你挡在其他选择之外。

下面展示同一份 kernel 在两条路线上的写法,功能是对 1024 个 float 做逐元素加法。两者都是完整可运行的程序,输出结果也相同,可以对照着看,体会差异在哪里。

SIMT 路线:cuda-oxide

cuda-oxide 是一个自定义的 rustc codegen 后端。它拦截编译流程,将 #[kernel] 函数经 Rust MIR 送入社区 Pliron IR 框架,再经 LLVM IR 一路降到 PTX;其余代码交给标准后端处理。Pliron 之上的 GPU 方言由我们定义,方言及其所有转换均保留在 Rust 侧,直到标准 LLVM 后端接管。

环境要求:Linux、计算能力 8.0 及以上的 GPU、CUDA 工具包(12.x 或更高)、clang 及其 libclang 头文件、固定版本的 nightly 工具链。cargo oxide doctor 会检查以上全部依赖,包括可选的系统级 LLVM。安装 cargo-oxide(驱动构建流程的 Cargo 子命令):

cargo +nightly-2026-04-03 install --git https://github.com/NVlabs/cuda-oxide.git cargo-oxide

然后初始化项目并运行。模板就是一个完整的向量加法程序:

cargo oxide new vecadd_demo
cd vecadd_demo
cargo oxide doctor
cargo oxide run

第一次 cargo oxide run 会构建 codegen 后端,耗时较长属正常现象;后续运行可复用缓存。

输出 PASSED: all 1024 elements correct。以上就是完成该任务的完整程序,即 cargo oxide new 生成的原始代码,此处补上了注释:

use cuda_device::{kernel, launch_bounds, launch_contract, thread, DisjointSlice};
use cuda_host::cuda_module;
use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig1D};
 
// === 设备端代码 - 这里所有内容都会被编译为 PTX ===
// 该宏还会生成下面用到的宿主端 API:
// `load`、`prepare_vecadd`,以及安全的 `vecadd` 启动方法。
#[cuda_module]
mod kernels {
    use super::*;
 
    #[kernel] // GPU 入口点
    #[launch_bounds(256)] // 每个 block 的最大线程数;让编译器据此分配寄存器
    #[launch_contract(domain = 1, block = (256, 1, 1))] // 一维索引,256 线程的 block
    pub fn vecadd(a: &[f32], b: &[f32], mut c: DisjointSlice<f32>) {
        let idx = thread::index_1d();
        let idx_raw = idx.get(); // 原始 usize 值,用于读取输入数组
        if let Some(c_elem) = c.get_mut(idx) {
            *c_elem = a[idx_raw] + b[idx_raw];
        }
    }
}
 
fn main() -> Result<(), Box<dyn std::error::Error>> {
    // === 宿主端初始化:设备、流和缓冲区 ===
    let ctx = CudaContext::new(0)?;
    let stream = ctx.default_stream();
 
    const N: usize = 1024;
    let a_host: Vec<f32> = (0..N).map(|i| i as f32).collect();
    let b_host: Vec<f32> = (0..N).map(|i| (i * 2) as f32).collect();
 
    let a_dev = DeviceBuffer::from_host(&stream, &a_host)?;
    let b_dev = DeviceBuffer::from_host(&stream, &b_host)?;
    let mut c_dev = DeviceBuffer::<f32>::zeroed(&stream, N)?;
 
    // === 加载、准备、启动 ===
    // SAFETY:本包拥有上面 kernels 模块生成的嵌入式设备代码包。
    let module = unsafe { kernels::load(&ctx)? };
 
    // 4 个 block,每个 256 线程,0 字节动态共享内存。`prepare_vecadd` 会
    // 将其与上面的契约以及设备的实际限制进行校验。
    // 下面安全的 `vecadd` 调用接收该 token,替代原始 config 的位置。
    let prepared = module.prepare_vecadd(LaunchConfig1D::new((N as u32).div_ceil(256), 256, 0))?;
    module.vecadd(&stream, &prepared, &a_dev, &b_dev, &mut c_dev)?;
 
    // === 读回并校验 ===
    // 将结果拷贝回宿主端并同步,确保读到 `c_host` 时启动已完成。
    let c_host = c_dev.to_host_vec(&stream)?;
    let errors = (0..N)
        .filter(|&i| (c_host[i] - (a_host[i] + b_host[i])).abs() > 1e-5)
        .count();
 
    if errors == 0 {
        println!("PASSED: all {} elements correct", N);
    } else {
        eprintln!("FAILED: {} errors", errors);
        std::process::exit(1);
    }
    Ok(())
}

主机代码和设备代码放在同一个文件里,一条命令即可构建,不需要单独的 kernel crate。

先看 kernel 的函数签名,安全性论证全在里面。ab 是普通的共享切片,所有线程都可读。cDisjointSlice<f32>,这种类型让每个线程只能独占访问属于自己的那一个元素。之所以需要它,是因为 &mut [f32] 不适合这个场景:所有线程都要拿到同一个 &mut,Rust 会正确地拒绝。DisjointSlice 把这一个可变借用拆成了每线程各自一份。

thread::index_1d() 返回的是一个索引类型,而不是裸整数,并且 c.get_mut(idx) 只接受这种类型。返回值是 Option,越界情况就变成了一个需要你处理的分支,而不是事后才发现的内存错误。

启动是经过校验的,而不是凭信任。#[launch_contract] 声明这个 kernel 按一维索引、以 256 线程的 block 运行。prepare_vecadd 会依据这份声明以及设备的实际限制来校验你传入的 LaunchConfig1D,并返回安全版 vecadd 方法所需的证明。没有契约的 kernel 只暴露原始的 unsafe 启动方法,因为一个裸的 LaunchConfig 说明不了它要启动的 kernel 的任何信息。

Tile 路线:cutile-rs

cutile-rs 的抽象层次更高。你操作的是 tile 而不是标量。每个 tile block 只以单一逻辑线程的方式执行一次 kernel 主体,作用于数据的一个子张量,由编译器决定背后需要多少真实 GPU 线程。#[cutile::module] 宏会把 kernel 的 AST 嵌入宿主二进制文件,并在首次需要该 kernel 时通过 CUDA Tile IR(NVIDIA 的 tile 级编译器 IR)进行 JIT 编译。

它的要求比 SIMT 路线宽松:需要一块计算能力 8.0 以上的 GPU、CUDA 13.3、stable Rust 1.89 或更新版本,以及 Linux,但不依赖 nightly 工具链,也不用自己编 LLVM。

cutile 已发布到 crates.io,无需克隆仓库:

cargo new vecadd_demo
cd vecadd_demo
cargo add cutile

下面是同样的逐元素加法,改用 tile 写法。把它粘贴到 src/main.rs,然后 cargo run

use cutile::prelude::*;
 
// 该宏将此模块的 AST 捕获进宿主二进制。内核在首次实际启动时通过 CUDA Tile IR 进行 JIT 编译。
#[cutile::module]
mod kernel {
    use cutile::core::*;
 
    #[cutile::entry()]
    fn add<const B: i32>(
        // B 是 tile 宽度,属于静态维度。不同的 B 会产生不同的特化版本。
        z: &mut Tensor<f32, { [B] }>, // 独占输出,B 个元素的子张量
        x: &Tensor<f32, { [-1] }>,    // 共享输入;-1 为动态维度,启动时解析
        y: &Tensor<f32, { [-1] }>,
    ) {
        // 该函数体对每个可变子张量执行一次,作为单个逻辑线程运行。
        // Tile 内核从 x、y 加载的是整个 tile,而非标量。
        let tx = load_tile_like(x, z); // x 中与 z 的此子张量对齐的切片
        let ty = load_tile_like(y, z);
        z.store(tx + ty); // 对整个 tile 逐元素计算
    }
}
 
fn main() -> Result<(), Error> {
    let device = Device::new(0)?;
    let stream = device.new_stream()?;
 
    // 以下均为惰性求值,此时尚未有任何操作触及 GPU。
    let x = api::ones::<f32>(&[1024]);
    let y = api::ones::<f32>(&[1024]);
 
    // 分区操作一次完成三件事:让每个 tile 独占其 128 元素分片、将网格固定为 1024/128 = 8 个 tile、并确定 B 的值。
    let z = api::zeros::<f32>(&[1024]).partition([128]);
 
    let c: Vec<f32> = kernel::add(z, x, y) // 接管三个张量的所有权
        .first()                           // ……并返回它们;从中取出输出
        .unpartition()                     // 移除宿主侧分区包装器,数据不移动
        .to_host_vec()                     // 记录回拷操作
        .sync_on(&stream)?;                // 此刻才真正执行上述所有操作
 
    let errors = c.iter().filter(|&&v| (v - 2.0).abs() > 1e-5).count();
    if errors == 0 {
        println!("PASSED: all {} elements correct", c.len());
    } else {
        eprintln!("FAILED: {errors} errors");
    }
    Ok(())
}

PASSED: all 1024 elements correct

Tile 轨道在 stable Rust 上给出了同样的结果,签名也做了同样的安全性论证。这次不需要 DisjointSlice。宿主端的分区操作只针对可变 tensor,它为每个 tile block 分出一块可写的子 tensor,其他 tile block 无法与之重叠。这种排他性正是 &mut 本身所保证的。

输入 shape 中的 -1 是一个哨兵值而非实际尺寸。该维度在启动时从 tensor 中读取,因此 shape 可以变化而无需重新编译。

宿主端真正值得注意的是 .partition([128]),它同时承担了三件事。让排他性落到实处——每个 tile 独占自己的 128 元素块,其他 tile 无法触及。固定启动几何——1024 除以 128 得到 8 个 tile 的 grid。

grid 由分区直接推导而来,不需要单独计算再与 kernel 的索引对照。它还提供了 B——调用处无需显式写出,因为 launcher 直接从分区中读取 tile 宽度。正因如此,&mut 输出必须先经过分区才能传入。

再看启动的返回值。宿主端调用的 add 是宏生成的 launcher,而非上面那个设备函数。它接管全部三个 tensor 的所有权,GPU 完成后以 tuple 形式归还。.first() 的作用就是从中取回输出。

.sync_on(&stream) 之前,什么都不会执行。前面所有操作都只是惰性描述——被记录而非提交,包括 oneszeros、kernel 调用乃至拷回宿主端。整个程序是一条链,只有一个同步点。

编译器能捕获什么

两个 kernel 对内存做了同样的声明:输入为共享,输出归唯一写者独占。区别仅在于声明的粒度,以及是否需要一个专用类型来实现。

这之所以重要,是因为成千上万的线程以不确定顺序访问同一组 buffer。当两个线程命中同一地址且其中一个在写时,顺序就决定了结果。这类 bug 极少能按需复现,往往测试通过后才在生产环境中暴露。

把 SIMT kernel 自己的输出 buffer 作为它的一个输入传进去,编译不通过,不管这个 kernel 实际上会不会产生竞争:

module.vecadd(&stream, &prepared, &c_dev, &b_dev, &mut c_dev)?;

error[E0502]: cannot borrow `c_dev` as mutable because it is also borrowed as immutable

Tile 这边同样的别名操作也过不了编译:

let z = api::zeros::<f32>(&[1024]);
kernel::add(z.partition([128]), z, y)

error[E0382]: use of moved value: `z`

两个例子都在编译期就抓住了经典的别名错误,只是划线位置不同。cuda-oxide 检查的是每次 launch 调用;cutile-rs 的所有权机制则贯穿整个 launch 边界,跟随 tensor 本身,这是更强的保证。

Tile 不给你碰共享内存或线程索引的机会,因为这两者都由编译器掌管。一个 tile block 就是一个逻辑线程,所以根本没有线程可以让你产生竞争。这正是它“天然安全”的原因,也是你付出的代价。SIMT 保留了这些控制权,但目前用共享内存得写 unsafe。共享内存是高性能 SIMT kernel 的基石,让这条路变得安全是正在进行的工作。

两个项目的现状

两个项目都处于早期阶段,都还不能用于生产环境。cuda-oxide 还在早期 alpha。cutile-rs 走得更远一些,已发布到 crates.io,NVIDIA 之外已有实际使用,包括 HuggingFace 的 Grout 推理引擎和 mistral.rs。功能覆盖还不完整,API 也会变动。遇到不好用的地方,欢迎反馈。

Cargo 和 crates 给人的预期是上手就能用。GPU 编程历史上恰恰相反,缩小这段距离也是我们工作的一部分。SIMT 这条路目前还需要锁定 nightly 工具链,这正是我们希望将来不再要求你做的事。

Rust 跑在 GPU 上并不是新鲜事。这个领域在我们之前就有出色的工作,而且还在继续发展。cuda-oxide 手册里的生态附录梳理了我们与 Rust-GPU、rust-cuda、CubeCL 等项目的相对位置,我们也一直在和 rust-cuda 的维护者合作,共同推进两个项目的成熟。

真正的新意,在于我们背后投入的工程,以及对未来方向的清晰判断。

现在就能做什么

动手试试现有的东西,也欢迎加入我们一起来做。项目还很早期,完全开放,你现在构建的东西将塑造接下来的方向。

Rust 社区

NVIDIA 很高兴能深度融入 Rust 社区,共同推动原生 Rust GPU 编程。rust-cuda、rust-gpu、cudarc 等项目率先打通了 GPU 与 Rust 的结合,背后的团队(包括 VectorWare)至今仍在影响着我们对自身工作的思考——我们在与 Rust 社区共建的过程中,始终受他们启发。

原始来源: NVIDIA 开发者博客

评论 (0)