团队用多智能体技能把24个TileGym算子从cuTile Python/Triton-TileIR移植到cuTile Rust,平均达到参考性能的99.5%
核心概要
该工作构建了一个多智能体AI技能,将cuTile Python与Triton-TileIR编写的GPU tile内核翻译为cuTile Rust,并据此把全部24个公开TileGym算子(约40个内核,覆盖逐元素运算到flash-attention decode、MLA与MoE)移植完成,在NVIDIA DGX B200上以CUPTI设备时间测得几何平均0.995,即平均达到cuTile Python性能的99.5%,且每个阶段都以可机器校验的判定和Tile IR diff把关。
深度剖析
提出并落地了一个多智能体翻译流水线,把cuTile Python与Triton-TileIR内核转换为cuTile Rust,覆盖分析、设备内核、host与FFI代码、基准测试四个环节。 此前跨CUDA Tile前端移植需要人工重写;这里把流程拆成专职子智能体,顶层智能体只做路由,每个阶段只加载自己需要的参考文档。 文中给出技能目录结构(SKILL.md、agents/*.md、references/、scripts/、examples/),并说明每个子智能体返回必须以字面块加一行VERDICT结束,编排器只依据判定路由。
以共享的CUDA Tile IR(cuda_tile方言)为验证骨架,用Tile IR diff在跑测试之前就做结构性校验。 功能测试可能放过“看似合理但错误”的翻译(如TMA load代价提示错误或丢失整除属性);IR diff能直接比对内存操作族、tile形状与归约,把这类问题提前暴露。 文中以softmax为例说明Rust与Python编译出相同算子清单:一次view load、正确轴上的reduce_max/reduce_sum、一次view store、两端TMA;并指出IR diff同时用作内核作者自检与IR-diff分析师的深度比对。
用该技能完成全部24个公开TileGym算子的移植,约40个GPU内核,平均达到cuTile Python性能的99.5%。 结果来自CI基准流水线本身:NVIDIA DGX B200上每个后端独占一块GPU,24个算子共347组配对配置,每组取四次CI运行中的最佳测量。 整体几何平均0.995,全部24个算子通过0.95门槛,约三分之一超过参考实现,逐元素与归一化内核收益最大;每次转换产出标准六文件变更集。
把翻译难点定位为“隐式特化”差异,并给出可操作的对应关系。 cuTile Python的JIT在调用时隐式特化,未走到的ct.Constant分支会被丢弃;cuTile Rust要求签名中声明所有特化,因此一个Python内核可能拆成多个Rust入口(如layer_norm按分支拆成2-Dnchw与1-Dw1)。 文中以softmax逐行对照展示映射:Constant[int]变const泛型、ct.load加NEG_INF填充变为make_partition_view加Partition::load、ct.bid(0)变get_tile_block_id()、keepdims归约变为reduce_*加显式reshape与broadcast。
启示与展望
该结果面向使用CUDA Tile家族(cuTile Python、Triton-TileIR、cuTile Rust)并希望把内核库迁移到Rust的团队,前提是环境满足CUDA 13.1+、用于性能检查的Blackwell GPU、Rust 1.89+以及tileiras编译器。技能随TileGym仓库发布,路径为skills/tilegym-converting-python-to-rust/,内核位于src/tilegym/ops/cutile_rs/;使用者可让智能体为算子添加cutile-rs后端,流水线会处理分析、内核、FFI、正确性与基准测试。性能结论的适用范围是CUPTI设备时间下的算子间比较,作者说明其测量墙钟时间但报告设备时间,以让比较聚焦于内核本身。
作者指出并非TileGym中所有已发布内核都已采用文中展示的完全安全风格:由于每个移植必须精确复现参考内核的Tile IR,当只有非受检API能复现时移植就使用该API,团队仍在把这些内核迁移到安全接口。此外,cuTile Rust可直接发射Tile IR,理论上可写出与其它前端完全一致的内核,但这类内核会变得难以解读,因此技能偏向生成惯用代码;作者预期精确匹配发射的Tile IR即可精确匹配性能,但实验只捕获设备时间。读者还需留意:性能数字来自CI基准流水线,依赖特定硬件与工具链版本;文中提到一次完整转换消耗的token量在百万量级,且编排器设有硬性生成次数上限,运行要么在预算内收敛,要么带着诊断停止。
