cuTile Rust(cutile-rs)是一个基于分块(tile)的系统,用于在Rust编程语言中以安全、符合语言习惯的方式编写GPU内核。它将Rust的所有权模型扩展到基于分块的GPU内核,把可变输出拆分为互不重叠的部分,并在内核启动过程中保持主机端的所有权契约。它还允许程序员在需要更底层控制时局部退出该模型,直接执行Tile IR操作。

TileGym CUDA分块内核库已经积累了大量用CUDA Tile Python(cuTile Python)和Triton-TileIR(nvtriton)编写的生产级内核。为了让所有这些内核也能在Rust中使用,我们的团队构建了一个AI智能体技能,用于将cuTile Python和Triton-TileIR内核翻译为cuTile Rust。

借助这项技能,我们将全部24个公开的TileGym运算符移植到了cuTile Rust,平均达到了cuTile Python性能的99.5%。这些运算符总共包含大约40个GPU内核,涵盖从逐元素运算到闪存注意力(flash-attention)解码、多头潜在注意力(MLA)以及专家混合(MoE)模型。需要说明的是,部分运算符需要多种内核变体。

每次转换都从该运算符现有的参考实现(cuTile Python或Triton-TileIR)开始,经过一个有边界的多智能体流水线,涵盖分析、设备内核、主机与FFI代码以及基准测试。每个阶段都以可机器检查的结论收尾,由验证脚本和Tile IR差异比对决定转换是否可以继续推进。主要挑战在于:cuTile Python的JIT编译在调用时隐式地对每个内核进行特化,而Rust则要求在内核签名中显式声明每一种特化。

本文将介绍我们如何开发一套多智能体工作流,将cuTile Python和Triton-TileIR内核翻译为cuTile Rust,并在每个阶段检查正确性与性能。文章将说明在真实内核中差距具体表现在哪里,该技能是如何构建的,从而使任何阶段都无需被盲目信任,以及最终生成的内核相对于参考实现的性能表现。该技能已随TileGym代码库发布,你可以将其应用于自己的内核。

分块IR前端之间的内核翻译

cuTile Python、Triton-TileIR和cuTile Rust是同一中间表示(IR)之上的三种前端:CUDA Tile IR,即cuda_tile方言。这三者都输入同一个tileiras编译器,由其执行分块级别的优化并生成GPU二进制文件。这一共享基础使得在CUDA Tile系列之间进行翻译既可行,同样重要的是,也可验证。

cuTile Python ─┐ Triton-TileIR ─┼─? CUDA Tile IR(cuda_tile方言)─? tileiras ─? cubin cuTile Rust ─┘

TileGym的生产级分块内核是基于前两种前端编写的。由于这三者最终都汇聚到同一IR,将内核移植到cuTile Rust并非重新优化的问题,而是用一种更安全的主机语言重新表达同一个分块程序,底层使用相同的编译器和相同的性能模型。共享的IR使翻译可以被检验。

一次忠实的移植应当复现参考内核的IR结构:相同的内存操作族、相同的分块形状、相同的规约操作。由于三种前端都生成相同的方言,这可以通过导出参考内核的Tile IR和翻译后内核的Tile IR并进行“差异比对”来直接验证,甚至无需运行任何测试。

这使得可以对智能体的输出进行结构层面的检查,而不仅仅是功能层面的检查。一个看似合理但实际错误的翻译(例如TMA加载的成本提示错误,或丢失了整除性属性)可能通过测试,但仍然是不正确的,且可能带来性能回归。这类问题可以通过与参考IR比对轻松发现并修复。IR差异比对阶段是本文所述流水线的核心。

关于Rust前端还有两个值得说明的要点。首先,Rust源代码是提前编译的。分块形状和元素类型由rustc检查。crate中嵌入了内核的抽象语法树(AST),在首次启动时,运行时会用具体的常量泛型值对其进行特化,并编译出一个cubin(此后缓存复用)。GPU二进制本身仍然是JIT编译的,但隐式特化不复存在:除非内核签名中声明,否则不会进行任何特化。其次,在TileGym中,cuTile Rust只是另一个后端。tilegym.set_backend("cutile-rs")会将同一套运算符API路由到Rust内核。

使特化显式化

两种前端的区别在于特化发生的位置。cuTile Python的JIT会根据调用时看到的内容进行特化;cuTile Rust则只根据内核签名中声明的内容进行特化。翻译工作大部分来自于把Python源代码中隐式留下的内容明确写出来。主要情况已在下表中总结。

以下部分将通过一个真实内核示例说明这些差异。

Softmax翻译示例

该示例内核刻意保持简单,以便你逐行比较两个版本。首先是cuTile Python版本:

(代码略,与原文一致)

以下是同一内核的cuTile Rust版本:

(代码略,与原文一致)

你可以直接对照阅读这些对应关系。它们如此清晰,是因为两种前端都是同一批Tile IR操作之上的薄层封装:

Constant[int]参数变为常量泛型(const TILE_SIZE: i32),由主机端根据每次启动的形状通过同一个Tile IR JIT进行实例化。

ct.load(..., padding_mode=NEG_INF)变为两个显式步骤:先构建一个make_partition_view(..., padding::NegInf, ...),再进行Partition::load——与参考IR中包含的、由TMA支持的视图加载相同,边缘不足部分填充为负无穷。

ct.bid(0)对应get_tile_block_id()。

cuTile Python中隐式的内容会变成显式的类型。每个中间量都是Tile<f32, {[1, TILE_SIZE]}>,而带有keepdims=True的规约操作会变成一次reduce_*操作,再加上显式的reshape和broadcast。

随后的IR差异比对确认Rust编译出的操作清单与Python原版一致:一次视图加载、在正确轴上的reduce_max/reduce_sum、一次视图存储,两端都使用了TMA。需要说明的是,并非所有已发布的TileGym内核都已采用这种完全安全的风格。每次移植都必须精确复现参考内核的Tile IR,因此在只有非安全API才能复现该结构的情况下,移植工作会使用该API。我们仍在将这些内核迁移到本文展示的安全接口上。

跨越C ABI

示例中的kernel.rs本身已经是一个完整的、一等公民的cuTile Rust内核。一个Rust应用程序可以依赖cutile这个crate,引入该内核模块,并通过crate的类型化API(所有权检查、分块类型一应俱全)直接启动其入口,完全不涉及FFI。

C-ABI层的作用范围更窄:将这些内核接入TileGym的Python调度与测试框架(同理,也可用于任何非Rust的主机)。

每个运算符从聚合的cdylib(整个库对应一个libcutile_kernels.so)中导出一个C符号。张量以一个普通的描述符结构体(ptr、ndim、shape[]、strides[])跨越语言边界,在Rust和Python之间镜像对应。

(代码略,与原文一致)

在Python一侧,cffi通过一个cdef字符串绑定该符号,该字符串是签名的唯一可信来源。封装层是一层带有校验检查的薄封装。

(代码略,与原文一致)

需要说明的是,启动器从不进行拷贝,从不进行分配,也从不获取所有权。borrow_f32将PyTorch的设备指针包装在一个ManuallyDrop<Tensor>中,因此Rust可以将张量交给内核,而不会释放自己并不拥有的内存,内核在调用方的CUDA流上异步启动。从PyTorch的角度看,这与其他任何扩展算子并无二致。

这在TileGym内部也几乎没有摩擦,因为cuTile Rust是惰性编译的。后端会跟踪源代码的新鲜度,因此编辑任何kernel.rs(或crate的清单文件)都会使下一次调用自动在调度前重新构建共享库,开发测试循环中无需显式执行cargo build。在Rust中迭代一个分块内核,和在Python中一样简单:修改内核、运行测试,新的二进制文件已经就位。

智能体技能是如何运作的

随NVIDIA/TileGym GitHub仓库一同发布的tilegym-converting-python-to-rust智能体技能,围绕一个设计决策构建:加载该技能的智能体本身不做任何工程工作。读取SKILL.md会将顶层智能体变成一个纯粹的编排者,其唯一权限是路由;实际工作发生在它派生出的专门子智能体中,每个子智能体只加载自己所在阶段所需的参考文档。下面我们将逐一介绍每类子智能体及其在转换过程中的角色。

分析器解决的是“JIT隐藏了规格说明”的问题。在参考内核中,一旦DSL被降低到cuda_tile方言,常量就已被固化,未走到的分支会消失,启动参数则存在于主机代码中。分析器还会选择基准实现:一个运算符往往同时存在cuTile Python和Triton-TileIR两种实现,因此分析器会对二者分别进行基准测试、比较,并为每种结构变体选出速度更快的一方,作为移植必须匹配的参考。

在任何Rust代码存在之前,它会为每种变体导出该参考实现的Tile IR(作为内核编写者的基准真值),并写出analysis.json,这是一份机器可读的规格说明,包含变体、常量、数据类型、容差、启动网格、自动调优空间以及所选定的基准实现。下游的一切都以此文件为路由依据。

内核编写者只生产kernel.rs,别无其他。由于被禁止接触主机代码,它的失败原因始终可归因。它面对的核心难题正是翻译差距本身,被提炼为该技能中的49条编码规则。它需要两次证明自己的工作:首先是功能层面,通过一个纯Rust内的流水线测试来运行内核,不涉及FFI也不涉及Python,因此数值错误无法躲藏在主机层的管道代码背后;其次是结构层面,需要通过针对分析器参考导出结果的IR自检。

主机/FFI构建者使经过验证的内核可以被TileGym调用(C-ABI启动器加上Python封装),并负责正确性检查,在所有数据类型和形状上运行该运算符真实的TileGym测试套件,只有全部通过(ALL_PASS)的结论才能解锁基准测试。这是整个技术栈(内核、启动器、封装层)第一次端到端运行。

性能验证器运行CUPTI基准测试协议(设备时间测量,在同一块GPU上按配置与参考实现配对比较),要求几何平均值落在参考实现的5%以内。它的职责不是优化,而是诚实地测量。

只有在失败时才会加入两个专家角色。二者都不修改代码,都是通过阅读IR进行诊断。当正确性测试失败或基准测试结果异常时,会派生出IR差异分析师。它会逐个变体比对参考Tile IR与生成的IR,并对每处分歧进行分类。关键在于,这一步能区分是翻译错误(返回给内核编写者,给出具体修复方案)还是内核修改无法解决的上游编译器缺陷。

残余性能调查员负责处理某些输入形状下速度较慢、但结果正确的内核,从边界两侧追溯差距根源:设备侧(内存操作类型、代码生成)和主机侧(启动配置、自动调优、封装层逻辑)。它会输出一份报告,供内核编写者据此采取行动。

这一设计出于两个关键原因。首先,一次完整转换的token消耗量级达到数百万。其次,这种拆分实现了责任隔离。由于内核在任何主机代码存在之前就已被独立验证,后续出现的失败会有一个可追溯的明确责任方。

有三个选择使这种拆分得以运作。子智能体只通过具有固定模式(schema)的产物进行沟通,而非通过对话。每个阶段都以一个机器可检查的结论收尾,编排者据此路由,无需阅读散文式描述。共享的cuda_tile方言使IR差异比对成为验证的支柱——既作为内核编写者在测试运行前的自检手段,也作为IR差异分析师在出现问题时进行深入比对的手段——从而拒绝那些结构上错误、但表面看似合理的翻译(比如在错误的轴上进行规约、丢失了掩码)。

编排循环

一次转换运行是一个精简的状态机,编排者自身的指令集也能容纳在一个精简的SKILL.md中。具体步骤如图1及其后所述。

预检:scripts/preflight.sh负责验证环境变量和工具链路径。非零退出码将终止本次运行:如果环境不可用,再多的智能体努力也无法修复缺失的编译器这类问题。

以最少提示信息进行派生:每个子智能体都从同一个模板派生,其提示词只包含两项内容:该阶段的Step-0文件清单(自身的指令文件加上该阶段所需的参考文档),以及前一阶段产物的具体路径。编排者从不将指令直接粘贴进提示词中。每个子智能体读取自己的文件,因此每个阶段的上下文中只包含该阶段所需的信息。

机械化验证器:每个子智能体的返回结果都必须以一个字面意义上的<VALIDATOR_OUTPUT>代码块结尾,并附带一行VERDICT结论。编排者检查该代码块内的退出码,然后纯粹依据结论进行路由;它从不从散文描述中推断修复方案。格式不正确的返回结果只会获得一次同一智能体的修复重试机会,绝不会被上报升级。

按表路由:图1即是整个决策函数。结论沿绿色路径推进,失败路由携带机器可读的责任方标签(主机问题→构建者自我重试;内核问题→由IR差异分析师指派责任方;环境问题→停止),一次失败的性能基准测试会经过一次残余性能调查员的处理。缺失责任方标签本身就被视为一种失败。编排者会选择停止而非猜测,因为若将主机端故障误路由到内核阶段,会浪费一整次重试机会。

严格的派生上限:图1每个方框中的xN限定了尝试次数(一次分析、两次内核编写者尝试、两次主机构建者尝试、一次诊断、两次基准测试运行,以及一次可选的性能优化)。一次运行要么在预算内收敛,要么带着落盘的诊断结果停止;它不能无限空转。

最终汇总:只有当路由到达完成状态后,validate_kernel.sh才会重新检查所有阶段共计17个文件的完整输出契约:报告、IR导出结果、正确性和性能日志。

在磁盘上,该技能将每个智能体的角色、共享知识以及验证器分别打包,因此每个子智能体只加载自己所需的内容:

(目录结构略,与原文一致)

这些编码规则是失败历史的提炼。每一条都源于某次早期转换生成的内核虽能编译,但若无该规则便是错误的。这些规则从细节到结构性问题都有涉及:assume_div_by只适用于指针,绝不适用于Tensor条目;broadcast之前必须先进行reshape;每种分块阶数下,规约轴的记录都必须精确无误。

该框架如何生效

有三个层次将这份Markdown文档转变为一个真正运行的系统:激活、契约以及外层驱动程序。

激活:运行时通过将任务与技能描述(“将Triton-TileIR或cuTile Python的GPU内核转换、移植或翻译为cuTile Rust”)进行匹配来激活该技能。一旦匹配成功,顶层智能体只会加载SKILL.md,这份精简文件将其转变为编排者。它从不读取子智能体的文件;那些文件只在子智能体内部加载,同时加载的还有该阶段所需的参考文档。

契约:各层之间传递的一切要么是文件,要么是固定格式的字符串。派生提示词只是最少量的指引,各阶段的输出是带有模式(schema)的产物,返回结果则是一个验证器代码块加一行结论。整个循环中没有任何环节依赖某个大语言模型去解读另一个大语言模型的自然语言描述。这正是24次无人值守转换能够可重复、而非纯靠运气成功的原因。

外层驱动程序:在生产环境中,一个一次性驱动程序会包装该技能,使每次转换成为一个无需人工干预的批处理任务。它会为每个运算符创建一个全新的独立分支检出,隐藏该目标运算符任何已有的实现(迫使智能体必须进行翻译),在与该运算符终端相分离的容器中启动智能体,并从外部轮询进度。运行结束后,驱动程序会应用捕获到的仓库差异,并执行验收检查:TileGym正确性为绿色(证明cuTile Rust后端确实被执行),且CUPTI几何平均加速比相对于cuTile Python基准≥0.95。只有绿色结果才会自动提交。一个精简的批处理驱动程序会运行整份运算符清单,每个运算符最多尝试两次,并推送通过验收的分支;转换失败的会留下一份诊断记录。

基准测试结果

借助tilegym-converting-python-to-rust这项技能,内核转换的效率大幅提升。平均token消耗降至约一半,每个运算符都通过了数值正确性验证,并且每个运算符相对于cuTile Python都实现了≥0.95的几何平均加速比。最终的性能数据来自CI基准测试流水线本身:在NVIDIA DGX B200上测得的CUPTI设备时间(每个后端独占一块GPU,在24个运算符上共347组配对配置)。每种配置都取四次CI运行中的最佳测量结果。

总体几何平均值为0.995,与cuTile Python基本持平。共享IR架构是这一结果的主要原因。两种前端都将同一个分块程序输入同一个共享优化器,一次忠实的翻译从构造上就继承了参考实现的性能。全部24个运算符都通过了0.95的检查门槛,其中约三分之一的表现优于参考实现,逐元素运算和归一化内核上的提升幅度最大。每次转换都会形成一个标准的六文件变更集,因此代码审查工作可以保持机械化、程式化。

图2报告的是CUPTI设备时间,它单独隔离出了内核本身的表现。对于亚微秒级的内核而言,挂钟时间和设备时间回答的是不同的问题:挂钟时间包含了启动和调度开销,反映的是用户的实际体验;而CUPTI设备时间则是对内核本身的独立比较。我们同时测量了挂钟时间,但报告设备时间,以确保运算符之间的比较聚焦于内核本身。

cuTile Rust也可以直接生成Tile IR。该DSL将Tile IR指令集作为其非安全API的一部分暴露出来。原则上,你可以编写一个内核,使其生成的Tile IR与其他前端生成的结果完全一致。但这样的内核会变得难以理解,因此该技能在设计上倾向于生成符合语言习惯的代码。由于本次实验只测量了设备时间,我们预期,如果精确匹配所生成的Tile IR,各前端之间的性能也将精确匹配。

开始使用cuTile Rust智能体技能

将cuTile Python和Triton-TileIR内核翻译为cuTile Rust的智能体技能,以及所有已转换的运算符,均已随TileGym一同发布。可通过skills/tilegym-converting-python-to-rust/访问该技能,其中包含各阶段的智能体指令、编码规则手册、概念指南、验证脚本,以及softmax和bmm两个完整示例。可通过src/tilegym/ops/cutile_rs/访问这些内核,其中包含每个运算符对应的一个<op>_kernel/目录,以及聚合而成的cutile_kernels这个crate。环境要求:CUDA 13.1及以上版本、用于性能检查的Blackwell架构GPU、Rust 1.89及以上版本,以及tileiras编译器。

要开始使用,只需让任意智能体指向该代码仓库,并要求它“为<op>添加一个cutile-rs后端”。整个流水线会自动处理分析、内核编写、FFI、正确性验证和基准测试等环节。更多细节请参阅GitHub上的TileGym README文档。

Q&A

Q1:cuTile Rust是什么?它解决了什么问题?

A:cuTile Rust(cutile-rs)是一个基于分块的系统,用于在Rust中安全地编写GPU内核。它将Rust所有权模型扩展到GPU内核开发中,并通过AI智能体技能,将已有的cuTile Python和Triton-TileIR内核自动翻译为Rust版本,同时保证正确性和性能。

Q2:这套多智能体翻译系统性能表现如何?

A:团队将全部24个TileGym公开运算符移植到cuTile Rust,平均达到cuTile Python性能的99.5%。总体几何平均加速比为0.995,全部运算符都通过了0.95的性能门槛,约三分之一表现优于参考实现。

Q3:cuTile Python和cuTile Rust翻译的主要难点是什么?

A:主要难点在于特化方式不同。cuTile Python的JIT在调用时隐式特化内核,而cuTile Rust要求在内核签名中显式声明所有特化内容,因此翻译工作大部分是把Python中隐式的部分显式地写出来。

NVIDIA