使用 AI Agent 将 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 智能体 技能，将 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 方言) ─► tileiras ─► cubin cuTile Rust ─┘
TileGym 生产级的 tile 内核是针对前两种前端编写的。由于三者最终汇聚于同一套 IR，将内核移植到 cuTile Rust 并非重新优化问题。它是在相同编译器支撑下，用更安全的宿主语言重新表达同一个 tile 程序，且底层性能模型保持一致。共享的 IR 使得翻译过程可检查。
忠实移植应复现参考内核的 IR 结构：相同的内存操作族、相同的 tile 形状以及相同的归约。由于三种前端输出的方言一致，可以直接导出参考内核和移植内核的 Tile IR 并进行“差分”对比，甚至在运行任何测试之前即可完成验证。
这使得我们不仅能从功能上，还能从结构上检查智能体（Agent）的输出。例如，一个看似合理实则错误的翻译（如 TMA 加载使用了错误的 cost hint，或遗漏了可整除属性）可能通过测试，但在测试覆盖范围之外却是不正确的，并可能导致性能回退。通过与参考 IR 对比，这些问题可以轻松发现并修复。IR 差分阶段是本文描述流水线中的核心环节。
对于本次讨论，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 只根据 kernel 签名中的声明来特化。因此，大部分翻译工作都花在了把 Python 源码中隐含的信息明确写出来上。主要情况总结如下表。
| cuTile Python（隐式 JIT） | cuTile Rust（AOT Rust 源码） | 对翻译的影响 |
|---|---|---|
未执行的 ct.Constant 分支会在编译前被剔除 | 两个分支都必须通过类型检查 | 一个 Python kernel 可能对应多个结构不同的 Rust 入口（例如 layer_norm 被拆成 2-D nchw 和 1-D w1 两个入口，因为分支会改变 tile 的秩） |
任意 dtype 组合都可按需编译 | FFI 只按固定的 symbol/dtype 表分发 | 支持某种 dtype 是一次显式的 ABI 扩展；共享表涵盖 f32/f16/bf16/i32/i64/f8e5m2/f8e4m3fn |
| JIT 类型系统本身就承担了输入校验 | 越过 C ABI 之后没有任何安全网，错误的 stride 会导致静默的数据损坏而非异常 | 设置两层防御：Python 包装层做语义检查，FFI 之后再做 ABI 检查（null/dtype/device）并返回具名错误码 |
下一节用一个真实的 kernel 示例来展示这些差异。
Softmax 翻译示例
这个示例 kernel 特意写得很简单，方便你逐行对比两个版本。先是 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] }>, // 每个 CTA 一行
input: &Tensor<E, { [-1, -1] }>,
) {
let row_idx = get_tile_block_id().0; // 等价于 ct.bid(0)
// 对应 ct.load(..., padding_mode=NEG_INF)：构建一个安全的分区视图，
// 将不规则列以 -inf 填充，然后加载当前 CTA 的行数据。
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 中隐含的语义在 Rust 中变成了显式类型。每个中间值都是
Tile<f32, {[1, TILE_SIZE]}>，而keepdims=True的归约操作则转化为reduce_*后紧跟显式的reshape和broadcast。
IR diff 结果证实，Rust 编译出的操作清单与 Python 原实现完全一致：一次 view 加载，在右轴上的 reduce_max/reduce_sum，一次 view 存储，以及两端的 TMA。需要注意的是，TileGym 中现有的内核尚未全部采用这种完全安全的风格。由于每个移植项目都必须精确复现参考内核的 Tile IR，因此当只有非安全 API 才能复现该 IR 时，移植代码会使用非安全 API。目前我们仍在将这些内核迁移到本博客展示的安全接口上。
跨越 C ABI
示例中的 kernel.rs 已是一个完整的、一等公民级的 cuTile Rust 内核。Rust 应用程序可以依赖 cutile crate，引入内核模块，并通过该 crate 的强类型 API 直接启动其入口点（包含所有权检查、tile 类型等），全程无需涉及 FFI。
C-ABI 层服务于更窄的场景：将这些内核接入 TileGym 的 Python 调度与测试框架（以及通过相同机制接入任何非 Rust 宿主）。
每个算子都会从一个聚合的 cdylib（整个库共一个 libcutile_kernels.so）中导出一个 C 符号。张量通过一个普通的描述符结构体（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 就能把张量交给 kernel，而无需释放自己并不拥有的内存；kernel 则在调用方的 CUDA stream 上异步启动。从 PyTorch 的角度看，它就是一个普通的扩展算子。
在 TileGym 中使用同样毫无摩擦，因为 cuTile Rust 采用惰性编译。后端会追踪源码的新旧状态，所以只要修改了任何 kernel.rs（或 crate 的 manifest），下一次调用就会在分发前自动重新构建共享库，开发-测试流程中完全不需要手动执行 cargo build。在 Rust 里迭代 tile kernel 和在 Python 里一样轻松：改完 kernel，跑一下测试，新的二进制就已经就位了。
agent skill 是如何工作的？
随 NVIDIA/TileGym GitHub 仓库一起发布的 tilegym-converting-python-to-rust agent skill，其核心只有一个设计决策：加载它的 agent 本身不做任何工程工作。读取 SKILL.md 后，顶层 agent 就变成了纯粹的任务编排器，只负责路由；真正的活儿由它派生的专职 subagent 完成，每个 subagent 只加载自己阶段所需的参考文档。下面我们逐一介绍各 subagent 的类型及其在转换流程中的角色。
分析器解决了“JIT 遮蔽规格”的难题。在参考内核中，当 DSL 降低为 cuda_tile 方言时，常量已固化，未执行的分支消失，而启动参数则位于宿主代码中。分析器还负责选取基线：由于算子通常同时拥有 cuTile Python 和 Triton-TileIR 两种实现，分析器会对两者进行基准测试，比较结果，并为每个结构变体选择速度更快的一方作为移植必须匹配的参考基准。
在任何 Rust 代码生成之前，它会为该参考基准的每个变体导出 Tile IR（这是内核编写者的真值依据），并生成 analysis.json，这是一个机器可读的规格文件，涵盖变体、常量、数据类型、容差、启动网格、自动调优空间以及选定的基线。后续所有流程均由此文件驱动。
内核编写者仅生成 kernel.rs，不涉及其他内容。由于被禁止接触宿主代码，其错误边界清晰可溯。其核心难点在于翻译鸿沟本身，这被提炼为技能集中的 49 条编码规则。它通过两次验证来证明工作质量：首先进行功能测试，利用纯 Rust 管道测试运行内核，全程不依赖 FFI 和 Python，确保数值错误无法隐藏在宿主接口背后；其次进行结构验证，通过对照分析器的参考 dump 完成 IR 自检。
宿主/FFI 构建器负责将经过验证的内核暴露给 TileGym（包括 C-ABI 启动器和 Python 封装器），并独占正确性检查职责。它会跨所有数据类型和形状运行算子的真实 TileGym 测试套件，只有当其判定为 ALL_PASS 时，基准测试才被解锁。这是整个技术栈（内核、启动器和封装器）首次实现端到端运行的环节。
性能验证器执行 CUPTI 基准测试协议（包括设备耗时测量和在同一 GPU 上与参考基准按配置配对标比对），要求几何均值落在参考基准的 5% 误差范围内。其职责并非优化，而是诚实地测量。
只有在校验失败时，两名专家介入。二者均不修改代码，仅通过阅读 IR 进行诊断。IR 差异分析器在正确性测试失败或基准测试结果异常时启动。它逐一比对参考 Tile IR 与生成的 IR 差异，并对每种分歧进行分类。关键在于，它能区分翻译错误（需反馈给内核编写者进行特定修复）与上游编译器 Bug（无法通过修改内核解决），从而明确责任归属。
负责性能剩余量调查的子智能体针对那些在特定输入形状下表现缓慢的正确内核，深入定位边界两侧的性能差距根源：设备端（内存操作族、代码生成）与宿主端（启动配置、自动调优及包装逻辑）。它生成的报告由内核编写者据此采取行动。 这一设计有两大核心动因。首先，完整的转换流程消耗的令牌量达数百万级别。其次，这种拆分实现了责任隔离。由于内核在宿主端代码存在之前已单独通过验证，后续出现的任何故障都有明确的、可处理的归属对象。 三个关键选择支撑了这一拆分架构。子智能体仅通过具有固定架构的工件通信，而非通过对话。每个阶段都以机器可校验的判定结束，编排者据此进行路由，而无需阅读散文。共享的cuda_tile 方言使 IR diff 成为验证的支柱——既作为内核编写者在测试运行前的自检手段，也作为 IR-diff 分析者在出现失败时的深度对比工具——从而拒绝那些结构上错误的翻译（例如在错误轴上执行归约、丢失掩码等），这些错误若不加甄别，往往会被误认为合理的输出。
编排者循环
转换运行是一个小型状态机，编排者自身的指令封装在一个精简的 SKILL.md 中。具体步骤详见图 1 及后续内容。
- 预检：
scripts/preflight.sh验证env变量和工具链路径。非零退出码将停止运行：环境不可用，无论智能体如何努力，都无法解决编译器缺失的问题。 - 以最少的指引启动：每个子代理都由同一个模板生成，提示词只包含两项内容：该阶段的 Step-0 文件清单（即它自己的指令文件加上该阶段所需的参考文档），以及上一阶段产物的具体路径。编排器从不把指令直接塞进提示词。每个子代理自行读取自己的文件，因此每个阶段的上下文只保留该阶段所需的内容。
- 机械式校验器：每个子代理的返回必须以字面量
<VALIDATOR_OUTPUT>块和一行VERDICT:结尾。编排器检查块内的退出码，然后完全依据裁决结果来路由，绝不从散文式的文字中推断修复方案。格式错误的返回只会触发一次同代理内的修复重试，绝不升级。 - 按表路由：图 1 就是全部的决策逻辑。裁决沿绿色路径推进，失败路由带有机器可读的责任标签（
host→ 构建器自行重启；kernel→ IR-diff 分析师指派责任方；env→ 停止），性能基准测试失败则路由给残差性能调查员处理一次。缺少责任标签本身就是一种失败。编排器宁可停止也不瞎猜，因为把宿主机故障错误地路由到 kernel 阶段，会浪费一整轮重试。 - 硬性启动上限：图 1 每个方框中的 xN 限制了尝试次数（一次分析、两次 kernel 编写、两次宿主机构建、一次诊断、两次基准测试运行，以及一次可选的性能调优）。运行要么在预算内收敛，要么带着落盘的诊断结果停止，不会陷入无效循环。
- 最终聚合：只有当路由走到完成节点后，
validate_kernel.sh才会对所有阶段的完整 17 文件输出契约做最终复查：报告、IR 转储、正确性和性能日志。
在磁盘上，这个 skill 将每个代理的角色、共享知识和校验器分开打包，让每个子代理只加载自己需要的内容：
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/ # 张量与指针、FFI 桥接、转置
├── scripts/ # 每个 agent 配有 diff_ir.sh + validate_*.sh
└── examples/{softmax,bmm}/ # 两个完整的转换示例
这些编码规则是从失败历史中提炼出来的。每一条规则的存在，都因为早期某次转换在缺少它时生成了一个能编译但结果错误的 kernel。它们涵盖从局部到结构的各个层面：assume_div_by 只能用于指针，绝不能用于 Tensor 条目；broadcast 前必须先做 reshape；并且对于每个 tile rank，reduction-axis 的记录必须精确无误。
框架如何生效
三层结构将 markdown 变成了一个可运行的系统：激活、契约，以及外层驱动器。
激活：运行时通过任务与描述匹配来激活该 skill（"将 Triton-TileIR 或 cuTile Python GPU kernel 转换、移植或翻译到 cuTile Rust"）。一旦匹配成功，顶层 agent 仅加载 SKILL.md，这是一个精简文件，使其转变为编排者。它从不读取子 agent 文件；那些文件在子 agent 内部加载，并伴随该阶段所需的参考文档。
契约：各层之间，一切都以文件或固定格式字符串的形式存在。Spawn 提示是指向最小化指针，阶段输出是带有 schema 的工件，返回则是一个校验器块加一行判定。编排者的全部权威在于路由，校验脚本的全部权威在于退出码。循环中的任何环节都不依赖一个 LLM 去解读另一个 LLM 的散文。正是这一点让 24 次无人值守的转换具备可重复性，而非碰运气。
外层驱动：生产环境中，一次性驱动程序封装该技能，将每次转换设置为无人值守的批处理任务。它在每个算子分支上创建全新检出，隐藏目标算子的已有实现（迫使代理执行翻译），在脱离算子终端的容器中启动代理，并在外部轮询进度。运行结束后，驱动程序应用捕获的仓库差异并执行验收检查：TileGym 正确性通过、证明 cuTile Rust 后端确实执行，且相对于 cuTile Python 基线，CUPTI 几何平均加速比 ≥ 0.95。仅当结果通过时才自动提交。轻量批处理驱动运行算子列表，每个算子最多尝试两次，推送通过的分支；失败的转换生成诊断轨迹。
基准测试结果
使用 tilegym-converting-python-to-rust 技能后，内核转换效率大幅提升。Token 成本平均降低约一半，每个算子均通过数值正确性验证，且相较于 cuTile Python，每个算子的几何平均加速比均 ≥0.95。最终性能数据直接来自 CI 基准测试流水线：在 NVIDIA DGX B200 上的 CUPTI 设备时间（每个后端独占一块 GPU，覆盖 24 个算子共 347 组配对配置）。每组配置取四次 CI 运行中的最佳测量值。
