cuTile Rust（cutile-rs）是一套在 Rust 中编写 GPU 内核的 tile 级系统。它把 Rust 的所有权模型扩展到基于 tile 的内核上：把可变输出切分成互不重叠的片段，并让 host 侧的所有权契约跨内核启动保持不变；当需要更底层的控制时，程序员可以局部退出这套约束，直接执行 Tile IR 操作。
这套前端要解决的问题是生态衔接。TileGym CUDA tile 内核库已经积累了相当多生产内核，分别用 cuTile Python 与 Triton-TileIR 编写。为了让这些内核同样能在 Rust 中使用，团队构建了一个 AI 智能体技能（skill），把 cuTile Python 与 Triton-TileIR 内核翻译成 cuTile Rust。
结果是：借助该技能，团队把 TileGym 全部 24 个公开算子移植到 cuTile Rust，合计约 40 个 GPU 内核，覆盖从逐元素运算到 flash-attention decode、多头潜在注意力（MLA）与 MoE 模型，平均达到 cuTile Python 性能的 99.5%。部分算子需要多个内核变体。
之所以能自动校验，是因为三种前端本就共享同一层中间表示。cuTile Python、Triton-TileIR 与 cuTile Rust 都是 CUDA Tile IR 上 cuda_tile dialect 的前端，最终都交给同一个 tileiras 编译器完成 tile 级优化并生成 GPU 二进制。这让翻译结果可以做结构性检查，而不只是功能性测试：把参考内核与翻译内核的 Tile IR 分别 dump 出来做 diff，在跑任何测试之前就能发现问题。一个看起来合理但实际错误的翻译——例如 TMA load 的 cost hint 写错，或丢了整除性属性——可能通过测试，却在测试覆盖之外出错，甚至带来性能回退。
每次转换都从算子已有的参考实现出发，走一条有界（bounded）的多智能体流水线：分析、device kernel、host 与 FFI 代码、基准测试，每个阶段都以机器可判定的结论收尾，由校验脚本与 Tile IR diff 决定这次转换能否继续。主要难点在于两种前端的特化时机不同：cuTile Python 的 JIT 在调用时隐式特化每个内核，而 Rust 要求把每一种特化都写进内核签名。例如 Python 里未被走到的 ct.Constant 分支会在编译前被丢弃，Rust 里两个分支都必须通过类型检查，于是一个 Python 内核可能拆成多个结构性 Rust 入口——layer_norm 就因为分支改变了 tile rank，拆成 2D 的 nchw 与 1D 的 w1。
本文用一个 softmax 内核逐行对照展示了这种差异：Constant[int] 参数变成 const 泛型，ct.load 的 NEG_INF 填充变成显式的 make_partition_view 加 Partition::load，ct.bid(0) 对应 get_tile_block_id()，带 keepdims 的归约则展开成 reduce_* 加 reshape 与 broadcast。工程上，Rust 源码提前编译，tile 形状与元素类型由 rustc 检查；crate 内嵌内核 AST，首次启动时由运行时用具体 const 泛型值特化并编译 cubin 后缓存。dtype 支持变成显式的 ABI 扩展，共享符号/类型表覆盖 f32、f16、bf16、i32、i64、f8e5m2 与 f8e4m3fn；跨 C ABI 时张量以描述符结构体（指针、维数、形状、步长）在 Rust 与 Python 之间传递，Python 包装层做语义检查，FFI 之后再做空指针、dtype 与设备检查。在 TileGym 中，cuTile Rust 只是又一个后端，调用 tilegym.set_backend("cutile-rs") 即可把同一套算子 API 路由到 Rust 内核。该技能已随 TileGym 仓库发布，可用于自己的内核。
WeChat
Profile