Fearless Concurrency on the GPU

• 作者: Melih Elibol, Jared Roesch, Isaac Gelado, Eric Buehler, Michael Garland(NVIDIA / HuggingFace)
• arXiv: 2606.15991v1
• PDF: 2606.15991v1.pdf

摘要

Rust 让带静态安全保证的系统编程在 CPU 上变得实用,但在 Rust 里写自定义 GPU kernel 时,程序员仍然被抛在语言的所有权(ownership)保证之外。本文提出 cuTileRust——一个基于 tile(分块)的系统,让人能在 Rust 里安全、地道地编写 GPU kernel。

cuTileRust 把 Rust 的所有权规则延伸到基于 tile 的 GPU kernel:可变输出被切分成 互不相交的片段,kernel 启动时保持主机侧(host)的所有权契约,并且当需要更底层的 控制时程序员可以在局部“退出“这套约束。系统还提供了一个可组合的主机端执行模型, 涵盖同步启动、异步流水线以及 CUDA graph 重放。

关键结果

• 在 NVIDIA B200 GPU 上,cuTileRust 的逐元素(element-wise)操作达到 7 TB/s,GEMM 达到 2 PFlop/s(为 cuBLAS 的 96%),与 cuTile Python 在测量噪声范围内持平。
• Grout 是一个基于 cuTileRust 的推理引擎,在端到端的 Qwen3 推理路径上完整 检验了这套系统。
• 在 batch-1 decode 下,Grout 对 Qwen3-4B 在 RTX 5090 上达到 171 tokens/s, 对 Qwen3-32B 在 B200 上达到 82 tokens/s——与 vLLM、SGLang 相当,并与 HBM roofline 的合理性校验一致。

笔记

核心思想是:GPU 加速的应用会把代码分成主机端(CPU)和设备端(GPU)两部分; 把 Rust 的所有权模型映射到基于 tile 的 kernel 上,就能让编译器强制保证可变访问 互不相交(disjoint),同时又不损失手写 CUDA 的性能。