Translating CUDA Tile Operations from Python to Rust Using Agentic AI

PUBLISHER ARTICLE PREVIEW

From the original article

cuTile Rust (cutile-rs) is a tile-based system for safe, idiomatic GPU kernel authoring in the Rust programming language. Extending the Rust ownership model to tile-based GPU kernels, it splits mutable outputs into disjoint pieces and preserves the host-side ownership contract across kernel launches. It also allows programmers to opt out locally when they need lower-level control, enabling direct execution of Tile IR operations.

The TileGym CUDA tile kernel library has accumulated a large library of production kernels written in CUDA Tile Python (cuTile Python) and Triton-TileIR (nvtriton). To make all of these kernels available in Rust as well, our team built an AI agent skill that translates cuTile Python and Triton-TileIR kernels into cuTile Rust.

Using this skill, we ported all 24 public TileGym operators to cuTile Rust and reached 99.5% of cuTile Python performance on average. They contain roughly 40 GPU kernels in total, ranging from element-wise operations to flash-attention decode, Multi-head Latent Attention (MLA), and mixture-of-experts (MoE) models. Note that some operators need multiple kernel variants.

Each conversion starts from whichever reference implementation the operator has (cuTile Python or Triton-TileIR) and runs through a bounded multi-agent pipeline covering analysis, the device kernel, host and FFI code, and benchmarking. Every stage ends in a machine-checkable verdict, with validator scripts and Tile IR diffs deciding whether a conversion moves forward. The main challenge is that cuTile Python JIT compilation specializes each kernel implicitly at call time, whereas Rust requires that you declare every specialization in the kernel’s signature.

This post explains how we developed a multi-agent workflow to translate cuTile Python and Triton-TileIR kernels into cuTile Rust, with checks for correctness and performance at each stage. It covers what the gap looks like in a real kernel, how the skill is structured so that no stage has to be taken on trust, and how the resulting kernels perform against their references.

The skill ships in the TileGym repo, so you can apply it to your own kernels.

cuTile Python, Triton-TileIR, and cuTile Rust are three front ends over the same IR: CUDA Tile IR, the cuda_tile dialect. All three feed the same tileiras compiler, which performs the tile-level optimizations and emits the GPU binary. This shared foundation makes translating across the CUDA Tile family practical and, just as important, verifiable.

The TileGym production tile kernels are written against the first two front ends. Because all three meet at the same IR, porting a kernel to cuTile Rust is not a re-optimization problem. It is re-expressing the same tile program in a safer host language, with the same compiler and the same performance model underneath. The shared IR makes translation checkable.

A faithful port should reproduce the reference kernel’s IR structure: the same memory-op families, same tile shapes, and same reductions. Because all three front ends emit the same dialect, this can be directly verified by dumping the reference kernel Tile IR and the translated kernel Tile IR and “diffing” them before a single test is executed.

This enables checking the agent’s output structurally, not just functionally. A wrong-but-plausible translation (a TMA load with wrong cost hint or a dropped divisibility attribute, for example) can pass tests yet still be incorrect outside of test coverage and may bring performance regressions. These issues can be easily checked and fixed by comparing with the reference IR.

The IR diff stage is central to the pipeline described in this post.

Two additional aspects of the Rust front end are important to note for this discussion. First, the Rust source is compiled ahead of time. Tile shapes and element types are checked by rustc. The crate embeds the kernel AST, and at first launch the runtime specializes it with the concrete const-generic values and compiles a cubin (cached thereafter).

The GPU binary itself is still JIT-compiled, but the implicitness is gone: nothing is specialized unless the kernel signature declares it. Second, in TileGym, cuTile Rust is simply another backend. tilegym.set_backend("cutile-rs") routes the same operator API to the Rust kernels.

The two front ends differ in where specialization happens. cuTile Python JIT specializes on whatever it sees at call time. cuTile Rust specializes only on what the kernel signature declares. Most of the translation work comes from spelling out what the Python source leaves implicit. The main cases are summarized in the following table.

This example kernel is intentionally simple so you can compare the two versions line by line. First, in cuTile Python:

You can read the correspondences directly. They’re this clean because both front ends are thin surfaces over the same Tile IR ops:

The IR diff then