NVIDIA 用多智能体 AI 将 CUDA Tile 内核从 Python 翻译到 Rust

cuTile Rust（cutile-rs

）是一个基于 tile 的系统，用于以 Rust 编程语言安全、符合习惯地编写 GPU 内核。它将 Rust 的所有权模型扩展到基于 tile 的 GPU 内核，把可变输出拆分为互不相交的部分，并在内核启动之间保持主机端的所有权契约。它还允许程序员在需要更低层控制时局部选择退出，从而能够直接执行 Tile IR 操作。

TileGym CUDA tile 内核库已经积累了大量用 CUDA Tile Python（cuTile Python）和 Triton-TileIR（nvtriton

）编写的生产内核。为了让所有这些内核也能在 Rust 中使用，我们的团队构建了一个 AI agent 技能，用于将 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 仓库发布，因此你可以将其应用于你自己的内核。

Tile IR 前端之间的内核翻译

cuTile Python、Triton-TileIR 和 cuTile Rust 是同一个 IR 之上的三个前端：CUDA Tile IR，即 cuda_tile

方言。三者都馈入同一个 tileiras

编译器，该编译器执行 tile 级优化并生成 GPU 二进制。这一共享基础使得在 CUDA Tile 家族之间进行翻译既实用，也同样重要的是，可验证。

cuTile Python ─┐
Triton-TileIR ─┼─► CUDA Tile IR (cuda_tile dialect) ─► tileiras ─► cubin
cuTile Rust ─┘

TileGym 的生产级 tile 内核是针对前两个前端编写的。由于三者最终都汇聚到同一个 IR，将内核移植到 cuTile Rust 并不是一个重新优化的问题。它只是用更安全的主机语言重新表达同一个 tile 程序，底层使用相同的编译器和相同的性能模型。共享的 IR 使得翻译过程可以被校验。

一次忠实的移植应当复现参考内核的 IR 结构：相同的内存操作族、相同的 tile 形状、相同的归约。由于三个前端都发出相同的方言，这可以直接验证——在运行任何测试之前，先 dump 参考内核的 Tile IR 和翻译后内核的 Tile IR，然后对它们进行“diff”。

这使得我们能够从结构上而不仅仅是功能上检查 agent 的输出。一个错误但看似合理的翻译（例如带有错误 cost hint 的 TMA 加载，或丢失了整除性属性）可能通过测试，但在测试覆盖范围之外仍然是错误的，并且可能带来性能回退。这些问题可以通过与参考 IR 对比来轻松检查并修复。IR diff 阶段是本文所述流水线的核心。

关于 Rust 前端，还有两个额外的方面对本次讨论很重要。首先，Rust 源代码是提前编译的。Tile 形状和元素类型由 rustc 检查

。crate 内嵌了内核 AST，在首次启动时，运行时会用具体的 const-generic 值对其进行特化，并编译出一个 cubin（之后会被缓存）。GPU 二进制本身仍然是 JIT 编译的，但隐式性消失了：除非内核签名声明了特化，否则不会有任何特化。其次，在 TileGym 中，cuTile Rust 只是另一个后端。tilegym.set_backend("cutile-rs")

将相同的算子 API 路由到 Rust 内核。

让特化显式化

两个前端在特化发生的位置上有所不同。cuTile Python JIT 在调用时对它所看到的内容进行特化。cuTile Rust 只对内核签名所声明的内容进行特化。大部分翻译工作来自于把 Python 源码中隐含的内容显式地写出来。主要情况总结在下表中。

表 1. cuTile Python-Rust 翻译差异示例

下一节使用一个真实的内核示例来说明这些差异。

Softmax 翻译示例

这个示例内核被有意设计得很简单，以便你可以逐行比较两个版本。首先，是 cuTile Python：

@ct.kernel
def softmax_kernel(output, input, TILE_SIZE: Constant[int]):
row_idx = ct.bid(0) # one CTA per row
row = ct.load(input, index=(row_idx, 0), shape=(1, TILE_SIZE),
padding_mode=ct.PaddingMode.NEG_INF)
row = ct.astype(row, ct.float32)
row_max = ct.max(row, axis=1, keepdims=True)
numerator = ct.exp(row - row_max)
denominator = ct.sum(numerator, axis=1, keepdims=True)
out = numerator / denominator
out = ct.astype(out, input.dtype)
ct.store(output, index=(row_idx, 0), tile=out)

以及 cuTile Rust 中的同一个内核：

#[cutile::module]
pub mod softmax_module {
use cutile::core::*;
#[cutile::entry()]
pub fn softmax_kernel<E: ElementType, const TILE_SIZE: i32>(
output: &mut Tensor<E, { [1, TILE_SIZE] }>, // one row per CTA
input: &Tensor<E, { [-1, -1] }>,
) {
let row_idx = get_tile_block_id().0; // ct.bid(0)
// ct.load(..., padding_mode=NEG_INF): a safe partition view whose ragged
// columns pad with -inf, then a load of this CTA's row.
let token: Token = get_tensor_token(input);
let row_view: Partition<E, { [1, TILE_SIZE] }> = make_partition_view(
input, const_shape![1, TILE_SIZE], padding::NegInf, dim_map::Identity, token);
let row: Tile<E, { [1, TILE_SIZE] }> = row_view.load([row_idx, 0i32]);
let row: Tile<f32, { [1, TILE_SIZE] }> = convert_tile(row); // ct.astype(f32)
let row_max: Tile<f32, { [1] }> = reduce_max(row, 1i32);
let shifted = row - row_max.reshape(const_shape![1, 1])
.broadcast(const_shape![1, TILE_SIZE]);
let numerator: Tile<f32, { [1, TILE_SIZE] }> = exp(shifted);
let denominator: Tile<f32, { [1] }> = reduce_sum(numerator, 1i32);
let out = numerator / denominator.reshape(const_shape![1, 1])
.broadcast(const_shape![1, TILE_SIZE]);
let out: Tile<E, { [1, TILE_SIZE] }> = convert_tile(out); // ct.astype(dtype)
output.store(out); // ct.store
}
}

你可以直接阅读这些对应关系。它们之所以如此清晰，是因为两个前端都只是同一组 Tile IR 操作之上的薄层：

Constant[int]

参数会变成 const 泛型（const TILE_SIZE: i32

），由宿主通过同一个 Tile IR JIT 按每次启动的形状进行实例化。ct.load(..., padding_mode=NEG_INF)

会变成两个显式步骤。首先构建一个make_partition_view(..., padding::NegInf, ...)

，然后是一个Partition::load

——即参考 IR 中包含的同一个由 TMA 支持的视图加载，其中不规则的尾部被填充为-inf

。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 中并非所有已发布的 kernel 都使用这种完全安全的风格。每个移植都必须精确复现参考 kernel 的 Tile IR，因此当只有未检查的 API 能复现它时，移植就会使用该 API。我们仍在将这些 kernel 迁移到本文所示的安全接口上。

跨越 C ABI

示例 kernel.rs

已经是一个完整、一等公民的 cuTile Rust kernel。Rust 应用程序可以依赖 cutile

crate，包含该 kernel 模块，并通过该 crate 的类型化 API 直接启动其入口（所有权检查、tile 类型等等），无需涉及 FFI。

C-ABI 层服务于一个更窄的目的：将这些 kernel 接入 TileGym Python 调度和测试框架（并且通过同样的机制，接入任何非 Rust 宿主）。

每个算子从聚合的 cdylib

中导出一个 C 符号（整个库对应一个 libcutile_kernels.so

）。张量以普通描述符结构体（ptr, ndim, shape[], strides[]

）的形式跨越边界，在 Rust 和 Python 之间镜像对应：

#[unsafe(no_mangle)]
pub unsafe extern "C" fn cutile_softmax(
out: *const TensorDesc, inp: *const TensorDesc,
n_rows: i32, tile_size: i32, device_id: i32, raw_stream: u64,
) -> i32 {
let out_d = unsafe { &*out };
let inp_d = unsafe { &*inp };
let device = Device::new(device_id as usize).expect("device");
let stream = unsafe { Stream::borrow_raw(raw_stream as *mut c_void, &device) };
let mut y = unsafe { borrow_f32(out_d, device_id as usize) };
let x = unsafe { borrow_f32(inp_d, device_id as usize) };
let y_part = (&mut *y).partition([1, tile_size as usize]);
match softmax_kernel(y_part, &*x).sync_on(&stream) {
Ok(_) => 0,
Err(_) => -1,
}
}

在 Python 侧，cffi

从 cdef

字符串中绑定该符号，该字符串是签名的唯一真实来源。包装器是一个带有验证检查的薄层：

_FFI_CDEF = """
int32_t cutile_softmax(
const TensorDesc* out, const TensorDesc* inp,
int32_t n_rows, int32_t tile_size,
int32_t device_id, uint64_t raw_stream);
"""
def softmax(x):
x = x.contiguous(); m, n = x.shape
y = torch.empty_like(x)
rc = lib.cutile_softmax(_desc(y), _desc(x), m, next_pow2(n),
x.device.index or 0,
torch.cuda.current_stream().cuda_stream)
assert rc == 0
return y

注意，启动器从不复制、从不分配、也从不获取所有权。borrow_f32

将 PyTorch 设备指针包装在 ManuallyDrop<Tensor> 中

，因此 Rust 可以将张量交给内核，而无需释放它并不拥有的内存，并且内核在调用者 CUDA 流上异步启动。从 PyTorch 的角度来看，这看起来就像任何其他扩展操作一样。

这在 TileGym 中也是无摩擦的，因为 cuTile Rust 是惰性编译的。后端跟踪源文件的新鲜度，因此编辑任何 kernel.rs

（或 crate 清单）都会使下一次调用在分派前自动重新构建共享库，在开发-测试周期中无需显式执行 cargo build

。在 Rust 中迭代 tile 内核就像在 Python 中一样简单：修改内核，运行测试，新的二进制文件就已经就位了。

智能体技能如何工作？

tilegym-converting-python-to-rust 智能体技能，随 NVIDIA/TileGym GitHub 仓库一起发布，围绕一个设计决策构建：加载它的智能体完全不进行任何工程工作。 阅读 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-diff 分析师。它逐变体地将参考 Tile IR 与生成的 IR 进行差异比对，并对每一处分歧进行分类。关键在于，这能将误译（带着具体修复方案退回给内核编写者）与上游编译器缺陷（任何内核改动都无法修复）区分开来。

残余性能调查员接手一个在某些输入形状上运行缓慢但结果正确的内核，并从边界两侧对差距进行根因分析：设备侧（内存操作族、代码生成）和主机侧（启动配置、自动调优和包装逻辑）。它生成一份报告，供内核编写者据此行动。

有两个关键原因促成了这一设计。第一，一次完整转换运行大约需要数百万个 token。第二，这种拆分隔离了责任归属。由于内核在任何主机代码存在之前就已单独得到验证，后续的失败就有了可追溯的责任方。

三项选择使这种拆分得以奏效。子代理仅通过具有固定 schema 的产物进行通信，绝不通过对话。每个阶段都以一个机器可校验的判定结束，编排器据此进行路由，而无需阅读散文式描述。而共享的 cuda_tile

方言使 IR diff 成为验证的骨干——既用作内核编写者在测试运行前的自检，也用作 IR-diff 分析师在出现失败时的深度比对——从而拒绝那些结构上错误的翻译（在错误的轴上做归约、丢失掩码），否则这些翻译会被当作看似合理而通过。

编排器循环

一次转换运行是一个小型状态机，编排器自身的指令可以容纳在一个精简的 SKILL.md

中。各步骤详见图 1 及下文。

图 1. 转换流水线涉及带有限重试的判定驱动路由预检：scripts/preflight.sh

验证env

变量和工具链路径。非零退出会停止运行：环境不可用，再多的代理努力也无法修复缺失的编译器。以最小指针启动：每个子代理都从一个模板启动，其提示仅包含两个元素：该阶段的 Step-0 文件列表（它自己的指令文件加上该阶段所需的参考文档）以及先前阶段产物的具体路径。编排器从不将指令粘贴到提示中。每个子代理读取自己的文件，因此每个阶段的上下文只包含该阶段所需的内容。机械校验器：每个子代理的返回必须以一个字面量<VALIDATOR_OUTPUT>

块和一行VERDICT:

结尾。编排器检查块内的退出码，然后纯粹根据判定进行路由；它从不从散文式描述中推断修复方案。格式错误的返回只会获得一次同一代理的修复重启，绝不升级。按表路由：图 1 就是整个决策函数。判定推进绿色路径，失败路由带有机器可读的责任方标签（host

→ 构建器重启自身；kernel

→ IR-diff 分析师分配责任方；env

→ 停止），而失败的性能基准测试会通过残余性能调查器路由一次。缺失所有者标签本身即为失败。编排器会停止而非猜测，因为将主机故障错误路由到内核阶段会浪费整个重试。硬性生成上限：图1中每个框内的xN限制了尝试次数（一次分析、两次内核编写器和两次主机构建器尝试、一次诊断、两次基准测试运行，以及一次可选的性能通过）。一次运行要么在预算内收敛，要么在磁盘上留下诊断后停止；它无法反复折腾。最终聚合：只有当路由到达完成状态后，validate_kernel.sh

才会重新检查所有阶段的完整17文件输出契约：报告、IR转储、正确性以及性能日志。

在磁盘上，该技能将每个代理的角色、共享知识以及验证器分别打包，因此每个子代理只加载其所需内容：

skills/tilegym-converting-python-to-rust/
├── SKILL.md # 入口点 + 编排契约
├── agents/*.md # 每个阶段一个指令文件
├── references/
│ ├── coding-rules.md # 编号规则（每条均来自真实失败案例）
│ ├── op-mapping.md # ct.* -> cutile-rs API 对照表
│ ├── ir-diff-checklist.md # 何种情况算作关键 IR 分歧
│ ├── pipeline.md # Rust 内流水线测试框架
│ └── performance-checklist.md # 基准测试协议
├── concepts/ # 张量 vs 指针、FFI 桥接、转置
├── scripts/ # diff_ir.sh + 每个 agent 的 validate_*.sh
└── examples/{softmax,bmm}/ # 两个完整转换示例

编码规则是失败历史的提炼。每一条规则的存在，都是因为早期转换产生了一个在没有它的情况下能编译但结果错误的内核。这些规则从狭窄到结构性不等：assume_div_by

仅适用于指针，绝不适用于 Tensor

条目，广播之前必须进行重塑，并且对于每个瓦片秩，归约轴的簿记必须精确。

测试框架如何生效

三层将 markdown 转化为运行系统：激活、契约和外部驱动程序。

激活： 运行时通过将任务与其描述（“将 Triton-TileIR 或 cuTile Python GPU 内核转换、移植或翻译为 cuTile Rust”）进行匹配来激活技能。匹配后，顶层代理仅加载 SKILL.md，这是一个精简文件，将其转变为编排器。它从不读取子代理文件；这些文件在子代理内部加载，同时仅加载其阶段所需的参考文档。

契约： 层与层之间，一切都是文件或固定格式字符串。生成提示是最小指针，阶段输出是带有模式的工件，返回是验证器块加判定行。编排器的全部权限是路由，验证器脚本的全部权限是退出代码。循环中没有任何东西依赖于一个 LLM 解释另一个 LLM 的散文。这就是使 24 次无人值守转换可重复而非侥幸的原因。

外部驱动程序： 在生产中，一次性驱动程序包装该技能，使每次转换成为无人值守的批处理作业。它在每个操作员分支上创建一个新的检出，隐藏目标操作员的任何预先存在的实现（因此代理必须进行翻译），在脱离操作员终端的容器中启动代理，并从外部轮询进度。运行结束时，驱动程序应用捕获的仓库差异并运行验收检查：TileGym 正确性为绿色，证明 cuTile Rust 后端实际执行，以及 CUPTI 几何平均加速比 ≥ 0.95 相对于 cuTile Python 基线。只有绿色结果才会自动提交。一个轻量级批处理驱动程序运行操作员列表，每个最多尝试两次，并推送通过的分支；失败的转换作为诊断轨迹落地。

基准测试结果

使用 tilegym-converting-python-to-rust 技能，内核转换变得高效得多。令牌成本平均降至约一半，每个操作员都经过数值正确性验证，并且每个操作员相对于 cuTile Python 的几何平均加速比达到 ≥0.95。最终性能数字来自 CI 基准测试管道本身：NVIDIA DGX B200 上的 CUPTI 设备时间（每个后端一个独占 GPU，24 个操作员共 347 个配对配置）。对于每个配置，最佳测量值取自四次 CI 运行。

图 2. 在 NVIDIA DGX B200 上，合成 cuTile Rust 与 cuTile Python 在 TileGym 各算子上的对比 总体几何平均值为 0.995，与 cuTile Python 持平。共享 IR 架构是这些结果的主要原因。两个前端将相同的 tile 程序送入同一个优化器，而忠实的翻译在构造上就继承了参考实现的性能。全部 24 个算子都通过了 0.95 的检查，约三分之一的表现优于参考实现，其中提升最大的是逐元素和归一化内核。每次转换都作为一个标准的六文件变更集落地，因此审查保持机械化。

图 2 报告的是 CUPTI 设备时间，它隔离了内核本身。在亚微秒级内核上，挂钟时间和设备时间回答的是不同的问题。挂钟时间包含启动和调度开销，反映用户实际体验，而 CUPTI 设备时间则在隔离状态下比较内核。我们测量挂钟时间并报告设备时间，是为了让算子之间的比较聚焦于内核本身。

cuTile Rust 也可以直接生成 Tile IR。该 DSL 将 Tile IR 指令集作为其 unsafe 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。