Translating CUDA Tile Operations from Python to Rust Using Agentic AI

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.

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, 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 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.

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:

@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:

#[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:

_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, shipped in the NVIDIA/TileGym 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 (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/. 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/, 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 on GitHub.