Translating CUDA Tile Operations from Python to Rust Using Agentic AI

NVIDIA Labs built an AI agent skill that translated all 24 public TileGym operators from cuTile Python and Triton-TileIR into cuTile Rust, reaching 99.5% of cuTile Python performance on average across roughly 40 GPU kernels. The multi-agent pipeline covers analysis, device kernel, host and FFI code, and benchmarking, with validator scripts and Tile IR diffs deciding whether each conversion proceeds. The skill ships in the TileGym repository so developers can apply it to their own kernels.

cuTile Rust https://github.com/nvlabs/cutile-rs 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 https://github.com/NVIDIA/cutile-python and Triton-TileIR https://github.com/triton-lang/Triton-to-tile-IR nvtriton . To make all of these kernels available in Rust as well, our team built an AI agent https://www.nvidia.com/en-us/ai/ 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 https://www.nvidia.com/en-us/glossary/mixture-of-experts/ 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. Kernel translation between Tile IR front ends cuTile Python, Triton-TileIR, and cuTile Rust are three front ends over the same IR: CUDA Tile IR https://github.com/NVIDIA/cuda-tile , 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. cuTile Python ─┐ Triton-TileIR ─┼─► CUDA Tile IR cuda tile dialect ─► tileiras ─► cubin cuTile Rust ─┘ 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 https://crates.io/crates/cutile/ 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. Making specialization explicit 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. | cuTile Python implicit JIT | cuTile Rust AOT Rust source | Consequence for translation | |---|---|---| | Untaken if ct.Constant branches are dropped before compilation | Both branches must type-check | One Python kernel becomes multiple structural Rust entries for example, layer norm splits into 2-D nchw and 1-D w1 entries because the branch changes tile rank | | Any dtype combination compiles on demand | The FFI dispatches over a fixed symbol/dtype table | Supporting a dtype is an explicit ABI extension; the shared table spans f32/f16/bf16/i32/i64/f8e5m2/f8e4m3fn | | The JIT type system is the input validation | Past the C ABI there is no safety net, so a wrong stride is a silent corruption, not an exception | Two defensive layers: semantic checks in the Python wrapper, ABI checks null/dtype/device behind the FFI with named return codes | Table 1. Examples of cuTile Python-Rust translation gaps The following section illustrates these differences using a real kernel example. Softmax translation example This example kernel is intentionally simple so you can compare the two versions line by line. First, in cuTile Python: 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 And the same kernel in 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 } } You can read the correspondences directly. They’re this clean because both front ends are thin surfaces over the same Tile IR ops: - Constant int parameters become const generics const TILE SIZE: i32 , instantiated per launch shape by the host through the same Tile IR JIT. - ct.load ..., padding mode=NEG INF becomes two explicit steps. First build a make partition view ..., padding::NegInf, ... , then a Partition::load —the same TMA-backed view load that the reference IR contains, with the ragged tail padded to -inf . - ct.bid 0 maps to get tile block id . - What cuTile Python keeps implicit becomes an explicit type. Every intermediate is a Tile<f32, { 1, TILE SIZE } , and a keepdims=True reduction becomes a reduce followed by an explicit reshape and broadcast . The IR diff then confirms that Rust compiles to the same op inventory as the Python original: one view load, reduce max/reduce sum on the right axis, one view store, and TMA on both ends. Note that not all of the shipped kernels in TileGym use this fully safe style yet. Each port has to reproduce the reference kernel Tile IR exactly, so where only an unchecked API reproduces it, the port uses that API. We’re still migrating those kernels onto the safe surface shown in this post. Crossing the C ABI The example kernel.rs is already a complete, first-class cuTile Rust kernel. A Rust application can depend on the cutile crate, include the kernel module, and launch its entry directly through the crate typed API ownership checks, tile types, and all with no FFI involved. The C-ABI layer serves a narrower purpose: plugging those kernels into the TileGym Python dispatch and test framework and, by the same mechanism, any non-Rust host . Each operator exports one C symbol from the aggregated cdylib one libcutile kernels.so for the whole library . Tensors cross as a plain descriptor struct ptr, ndim, shape , strides mirrored between Rust and Python: js 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, } } On the Python side, cffi binds that symbol from a cdef string that is the single source of truth for the signature. The wrapper is a thin layer with validation checks: js 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 Note that the launcher never copies, never allocates, and never takes ownership. borrow f32 wraps the PyTorch device pointer in a ManuallyDrop<Tensor , so Rust can hand the kernel its tensors without ever freeing memory it does not own, and the kernel launches asynchronously on the caller CUDA stream. From a PyTorch perspective, this looks like any other extension op. This is also friction-free within TileGym, because cuTile Rust compiles lazily. The backend tracks source freshness, so editing any kernel.rs or the crate manifest makes the next call automatically rebuild the shared library before dispatch, with no explicit cargo build in the develop-test cycle. Iterating on a tile kernel in Rust is as easy as it is in Python: change the kernel, run the test, and the new binary is already in place. How does the agent skill work? The tilegym-converting-python-to-rust agent skill https://github.com/NVIDIA/TileGym/tree/main/skills/tilegym-converting-cutile-triton-to-cutile-rs , shipped in the NVIDIA/TileGym https://github.com/NVIDIA/TileGym/tree/main/skills/tilegym-converting-cutile-triton-to-cutile-rs GitHub repo, is built around one design decision: the agent that loads it does no engineering work at all . Reading SKILL.md turns the top-level agent into a pure orchestrator whose only authority is routing; the work happens in specialized subagents it spawns, each loading only the reference documents its stage needs. We’ll walk through each subagent type and its role in the conversion. The analyzer solves the “JIT hides the spec” problem. In a reference kernel, constants are baked in once the DSL lowers to the cuda tile dialect, untaken branches vanish, and launch parameters live in host code. The analyzer also selects the baseline: an operator often has both a cuTile Python and a Triton-TileIR implementation, so the analyzer benchmarks each, compares them, and selects the faster one per structural variant as the reference the port must match. Before any Rust exists, it dumps the Tile IR of that reference for each variant the ground truth for the kernel writer and writes analysis.json , a machine-readable spec of variants, constants, dtypes, tolerances, launch grids, autotune space, and the chosen baseline. Everything downstream routes from this file. The kernel writer produces kernel.rs and nothing else. Barred from host code, its failures stay attributable. Its hard problem is the translation gap itself, distilled into the skill’s 49 coding rules. It proves its work twice. First functionally, with an in-Rust pipeline test that runs the kernel with no FFI and no Python, so a numerics bug can’t hide behind host plumbing. Second structurally, by clearing an IR self-check against the analyzer’s reference dump. The host/FFI builder makes validated kernels callable from TileGym the C-ABI launcher plus the Python wrapper and owns the correctness checks, runs the operator’s real TileGym test suite across all dtypes and shapes, and only its ALL PASS verdict unlocks benchmarking. This is the first point where the full stack kernel, launcher, and wrapper runs end-to-end. The performance validator runs the CUPTI benchmark protocol device-time measurement, per-config pairing against the reference on the same GPU and requires the geometric mean to land within 5% of the reference. Its job is not to optimize but to measure honestly. Two specialists join only on failure. Neither edits code; both diagnose by reading IR. The IR-diff analyst is spawned when a correctness test fails or a benchmark looks off. It diffs the reference Tile IR against the generated IR variant by variant and classifies each divergence. Crucially, this separates a mistranslation route back to the kernel writer with a specific fix from an upstream compiler bug that no kernel change can fix. The residual-performance investigator takes a correct kernel that is slow on some input shapes and root-causes the gap on both sides of the boundary: the device side memory-op family, codegen and the host side launch configuration, autotune, and wrapper logic . It emits a report that the kernel writer acts on. Two key reasons motivate this design. First, a full conversion runs on the order of millions of tokens. Second, the split isolates blame. Because the kernel is proven in isolation before any host code exists, a later failure has a tractable owner. Three choices make this split work. Subagents communicate only through artifacts with fixed schemas, never through conversation. Each stage ends with a machine-checkable verdict the orchestrator routes on without reading prose. And the shared cuda tile dialect makes IR diff the backbone of verification—used both as the kernel writer’s self-check before tests run and as the IR-diff analyst’s deep comparison when something fails—rejecting structurally wrong translations a reduction on the wrong axis, a lost mask that would otherwise pass as plausible. The orchestrator loop A conversion run is a small state machine, and the orchestrator’s own instructions fit in a lean SKILL.md . The steps are detailed in Figure 1 and following. 1. Preflight: scripts/preflight.sh verifies env vars and toolchain paths. A non-zero exit stops the run: the environment is unusable, and no amount of agent effort fixes a missing compiler. 2. Spawn with minimal pointers: Every subagent is spawned from one template whose prompt contains only two elements: the stage Step-0 file list its own instruction file plus the reference docs that stage needs and the concrete paths of prior-stage artifacts. The orchestrator never pastes instructions into prompts. Each subagent reads its own files, so the context of each stage holds only what that stage needs. 3. Mechanical validator: Every subagent return must end with a literal <VALIDATOR OUTPUT block and one VERDICT: line. The orchestrator checks the exit codes inside the block, then routes purely on the verdict; it never infers a fix from prose. A malformed return earns exactly one same-agent repair respawn, never an escalation. 4. Route by table: Figure 1 is the whole decision function. Verdicts advance the green path, failure routes carry a machine-readable owner tag host → the builder respawns itself; kernel → the IR-diff analyst assigns the owner; env → stop , and a failed perf benchmark routes once through the residual-perf investigator. A missing owner tag is itself a failure. The orchestrator stops rather than guessing because a host fault misrouted to the kernel stage wastes an entire retry. 5. Hard spawn caps: The xN in each box of Figure 1 caps the attempts one analysis, two kernel-writer and two host-builder attempts, one diagnosis, two benchmark runs, and one optional perf pass . A run either converges within the budget or stops with a diagnosis on disk; it cannot thrash. 6. Final aggregate: Only after the route reaches completion does validate kernel.sh recheck the full 17-file output contract across all stages: reports, IR dumps, correctness, and performance logs. On disk, the skill packages each agent’s role, the shared knowledge, and the validators separately, so every subagent loads only what it needs: skills/tilegym-converting-python-to-rust/ ├── SKILL.md entry point + orchestration contract ├── agents/ .md one instruction file per stage ├── references/ │ ├── coding-rules.md numbered rules each from a real failure │ ├── op-mapping.md ct. - cutile-rs API table │ ├── ir-diff-checklist.md what counts as a critical IR divergence │ ├── pipeline.md in-Rust pipeline test harness │ └── performance-checklist.md benchmark protocol ├── concepts/ tensor-vs-pointer, FFI bridge, transpose ├── scripts/ diff ir.sh + validate .sh per agent └── examples/{softmax,bmm}/ two fully worked conversions The coding rules are the distilled failure history. Each one exists because an early conversion produced a kernel that compiled but was wrong without it. They range from the narrow to the structural: assume div by applies only to pointers and never to Tensor entries, a broadcast must be preceded by a reshape, and reduction-axis bookkeeping must be exact for every tile rank. How the harness takes effect Three layers turn the markdown into a running system: activation, contracts, and the outer driver. Activation: The runtime activates the skill by matching a task against its description “convert, port, or translate Triton-TileIR or cuTile Python GPU kernels to cuTile Rust” . On match, the top-level agent loads SKILL.md only , a lean file that turns it into the orchestrator. It never reads the subagent files; those load inside the subagents themselves, alongside just the reference documents their stage needs. Contracts: Between layers, everything is a file or a fixed-format string. Spawn prompts are minimal pointers, stage outputs are artifacts with schemas, and returns are a validator block plus a verdict line. The entire authority of the orchestrator is routing, and the entire authority of the validator scripts is exit codes. Nothing in the loop depends on one LLM interpreting the prose of another LLM. This is what makes 24 unattended conversions repeatable rather than lucky. The outer driver: In production, a one-shot driver wraps the skill to make each conversion a hands-off batch job. It creates a fresh checkout on a per-operator branch, hides any pre-existing implementation of the target operator so the agent must translate , launches the agent in a container detached from the operator’s terminal, and polls progress from outside. When the run ends, the driver applies the captured repo diff and runs the acceptance check: TileGym correctness is green, proof that the cuTile Rust backend actually executed, and CUPTI geomean speedup ≥ 0.95 against the cuTile Python baseline. Only a green outcome autocommits. A thin batch driver runs the operator list with at most two attempts each and pushes the branches that pass; failed conversions land as a diagnosis trail. Benchmarking results With the tilegym-converting-python-to-rust skill, kernel conversion becomes far more efficient. Token cost drops to about half on average, every operator is validated for numerical correctness, and each one hits a geomean speedup of ≥0.95 versus cuTile Python. The final performance numbers come from the CI benchmark pipeline itself: CUPTI device time on NVIDIA DGX B200 https://www.nvidia.com/en-us/data-center/dgx-b200/ one exclusive GPU per backend, 347 paired configurations across the 24 operators . For each configuration, the best measurement is taken across four CI runs. Overall geomean is 0.995, parity with cuTile Python. The shared-IR architecture primarily explains these results. Both front ends feed the same tile program into a shared optimizer, and a faithful translation inherits the reference’s performance by construction. All 24 operators clear the 0.95 check, and about a third come out ahead of the reference, with the largest wins on element-wise and normalization kernels. Each conversion lands as a standard six-file changeset, so review stays mechanical. Figure 2 reports CUPTI device time, which isolates the kernel itself. Wall-clock time and device time answer different questions on submicrosecond kernels. Wall-clock time includes launch and scheduling cost and reflects what the user experiences, while CUPTI device time compares the kernels in isolation. We measure wall clock and report device time to keep the operator-to-operator comparison about the kernels. cuTile Rust can also emit Tile IR directly. The DSL exposes the Tile IR instruction set as part of its unsafe API surface. In principle, you could write a kernel to match the Tile IR emitted by other front ends exactly. However, such kernels become uninterpretable, so the skill is biased to generate idiomatic code. Because the experiments capture only device time, we expect that matching the emitted Tile IR exactly would match performance exactly across front ends. Get started with cuTile Rust agent skills The agent skill that translates cuTile Python and Triton-TileIR kernels into cuTile Rust and all converted operators ships with TileGym. Access the skill through skills/tilegym-converting-python-to-rust/ https://github.com/NVIDIA/TileGym/tree/main/skills/tilegym-converting-cutile-triton-to-cutile-rs . It includes per-stage agent instructions, a coding rulebook, concept guides, validator scripts, and worked softmax and bmm examples. Access the kernels through src/tilegym/ops/cutile rs/ https://github.com/NVIDIA/TileGym/tree/main/src/tilegym/ops/cutile rs , including one <op kernel/ per operator plus the aggregated cutile kernels crate. Requirements: CUDA 13.1+, a Blackwell GPU for the perf check, Rust 1.89+, and the tileiras compiler. To get started, point any agent at the repo and ask it to “add a cutile-rs backend for <op “. The pipeline handles analysis, kernel, FFI, correctness, and benchmarking. For more details refer to the TileGym README https://github.com/NVIDIA/TileGym 5-enable-the-cutile-rs-rust-backend-optional on GitHub.