The TileGym team implemented a multi-agent AI skill that automatically translates tile-based GPU kernels written in cuTile Python and Triton-TileIR (nvtriton) into cuTile Rust (cutile-rs). Using the skill, they ported all 24 public TileGym operators (about 40 GPU kernels, with some operators having multiple kernel variants) and achieved device-time parity with the cuTile Python baseline: a geometric mean of 0.995 over 347 paired configurations measured on an NVIDIA DGX B200.
What is cuTile Rust and why port kernels?
cuTile Rust (cutile-rs) is a tile-based system for writing safe, idiomatic GPU kernels in Rust. It extends Rust’s ownership model to tile-based kernels by partitioning mutable outputs into disjoint pieces and preserving host-side ownership contracts across kernel launches. It also exposes unsafe APIs for direct Tile IR operations when low-level control is required.
TileGym’s production tile kernels have been written primarily against cuTile Python and Triton-TileIR. Making those kernels available in Rust lets Rust-hosted workflows use the same kernels while retaining the same compiler and performance model underneath, because all three front ends target the same CUDA Tile IR (cuda_tile dialect) and the same tileiras compiler.
Shared Tile IR makes translation practical and verifiable
cuTile Python, Triton-TileIR and cuTile Rust are three front ends that emit the same cuda_tile dialect. All feed the tileiras compiler, which performs tile-level optimizations and emits the GPU binary. Because the IR is shared, translation is not a re-optimization problem but re-expressing the same tile program in a safer host language. A faithful port must reproduce the reference kernel’s IR structure (memory-op families, tile shapes, reductions), and the pipeline verifies that by dumping and diffing the reference and translated Tile IR before tests run.
Structural IR checks catch plausible-but-wrong translations (for example, a TMA load with an incorrect cost hint or a lost divisibility attribute) that could pass functional tests yet break in production or cause performance regressions.
The main translation gap: implicit JIT vs explicit AOT
The largest practical difference between cuTile Python and cuTile Rust is where specialization happens:
- cuTile Python (implicit JIT) specializes kernels dynamically at call time. Untaken branches and baked-in constants in the lowered cuda_tile IR disappear, and dtype combinations may be compiled on demand.
- cuTile Rust (AOT Rust source) requires every specialization to be declared in the kernel signature (const generics, type parameters). The runtime still JIT-compiles the GPU binary at first launch, but the kernel’s source-level implicitness is gone.
Consequences for translation include splitting one Python kernel with implicit branches into multiple structural Rust entries, making dtype support an explicit ABI extension, and replacing the JIT’s input validation with two defensive layers: semantic checks in the Python wrapper and ABI checks behind the FFI.
Softmax example — a direct mapping
A short softmax example shows a line-by-line correspondence:
- Constant[int] parameters in cuTile Python become const generics (const TILE_SIZE: i32) in Rust.
- ct.load(..., padding_mode=NEG_INF) maps to building an explicit partition view with padding::NegInf and calling Partition::load (the same TMA-backed view load the reference IR uses).
- ct.astype becomes convert_tile, bid(0) becomes get_tile_block_id(), and keepdims reductions become explicit reduce_* plus reshape and broadcast.
Because both front ends are thin surfaces over the same Tile IR ops, the generated Rust code compiles to the same op inventory; the IR diff confirms that (view load, reduce_max/reduce_sum on the same axis, view store, TMA at both ends).
Crossing the C ABI and integrating with TileGym
Rust kernels can be invoked directly through the cutile crate’s typed API. To plug Rust kernels into TileGym’s Python dispatch and tests, each operator exports one C symbol from an aggregated cdylib (one libcutile_kernels.so for the whole library). Tensors cross the boundary as a simple descriptor struct (ptr, ndim, shape[], strides[]) mirrored in Rust and Python.
On the Python side, cffi binds the C symbol via a cdef string that serves as the single source of truth for the signature. The Python wrapper performs basic validation and calls the C-ABI launcher without copying or taking ownership. On the Rust side, borrow_f32 wraps the PyTorch device pointer in a ManuallyDrop<Tensor> so Rust doesn’t free memory it does not own; the kernel launches asynchronously on the caller’s CUDA stream, so from PyTorch’s perspective the extension acts like any other op.
TileGym’s backend tracks source freshness and compiles cutile Rust lazily, so editing a kernel.rs file causes the shared library to be rebuilt automatically on the next call—iteration in Rust is comparable to iteration in Python.
How the agent skill is organized
The tilegym-converting-python-to-rust skill (in NVIDIA/TileGym) is organized so the top-level agent is a pure orchestrator and does no engineering work itself. It spawns specialized subagents; each subagent loads only the instruction files and reference documents it needs. Key subagents and responsibilities:
- Analyzer: resolves the JIT-hidden spec by dumping reference Tile IR for each variant, benchmarking cuTile Python and Triton-TileIR references, and writing analysis.json (variants, constants, dtypes, tolerances, launch grids, autotune space, chosen baseline). Downstream stages route from this file.
- Kernel writer: produces kernel.rs only (no host code). It follows 49 coding rules distilled from prior failures, runs an in-Rust pipeline test (no FFI/Python), and clears an IR self-check against the analyzer’s reference dump.
- Host/FFI builder: produces the C-ABI launcher and Python wrapper, runs the operator’s full TileGym test suite across dtypes and shapes, and only an ALL_PASS unlocks benchmarking.
- Performance validator: runs CUPTI device-time measurements per-config paired against the reference on the same GPU; the geometric mean must be within 5% of the reference.
Two diagnostic specialists run only on failure:
- IR-diff analyst: diffs reference Tile IR against generated IR variant-by-variant and classifies divergences; this separates mistranslation from upstream compiler bugs.
- Residual-performance investigator: analyzes correct but slow kernels, looking both at device-side codegen/memory-op families and host-side launch/autotune logic.
Design constraints that make the pipeline reliable:
- Subagents communicate only through artifacts with fixed schemas, never through freeform conversation.
- Every stage ends with a machine-checkable validator block and VERDICT line; the orchestrator routes purely on these verdicts.
- The shared cuda_tile dialect makes IR diff the backbone of verification, used both by the kernel writer’s self-check and the IR-diff analyst.
- Bounded retries per stage prevent thrashing: a conversion either converges within the allotted attempts or stops with a diagnosis on disk.
Orchestrator loop and artifacts
A conversion run includes preflight (scripts/preflight.sh) that validates env vars and toolchain paths; spawn with minimal pointers so each subagent only reads what it needs; mechanical validators (each subagent must return a <VALIDATOR_OUTPUT> block and a VERDICT: line); routing by a decision table; and hard spawn caps for limited retries. The final validate_kernel.sh rechecks the 17-file output contract across stages before accepting results.
The skill directory packages stage instructions, shared knowledge, and validator scripts so each subagent loads only its required files. Notable contents: SKILL.md, agents/*.md, references (coding-rules.md, op-mapping.md, ir-diff-checklist.md, pipeline.md, performance-checklist.md), concepts/, scripts/ and examples/{softmax,bmm}/.
The coding rules codify past failures (e.g., assume_div_by applies only to pointers, broadcasts must be preceded by reshape, reduction-axis bookkeeping must be exact for every tile rank).
Benchmarking results
Conversions done by the skill reduce token-costs roughly by half on average, validate numerical correctness, and meet the performance criterion of geomean ≥ 0.95 relative to cuTile Python. The CI benchmark pipeline used CUPTI device-time on an NVIDIA DGX B200 with one exclusive GPU per backend and 347 paired configurations across 24 operators; for each configuration the best measurement across four CI runs is reported.
The overall geometric mean is 0.995—parity with cuTile Python. All 24 operators cleared the 0.95 check; about one-third were faster in Rust, with the largest wins on element-wise and normalization kernels. Each conversion is packaged as a six-file changeset to keep reviews mechanical.
Note: CUPTI device time isolates the kernel execution itself; wall-clock time includes launch and scheduling overheads and may differ, particularly for very short kernels. The CI focuses on device time to compare kernels themselves.
cuTile Rust can also emit Tile IR directly via its unsafe API surface, so in principle a kernel could be written to match another front end’s emitted Tile IR exactly; however, such kernels are less interpretable, so the skill generates idiomatic Rust that reproduces the Tile IR.
Getting started
The skill and converted operators live in the TileGym repo at skills/tilegym-converting-python-to-rust/. Converted kernels are under src/tilegym/ops/cutile_rs/, 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 begin, point an agent at the repo and request “add a cutile-rs backend for <op>”. The pipeline handles analysis, kernel authoring, FFI/host construction, correctness validation, and benchmarking. See the TileGym README for more details.
Closing remarks
By leveraging a shared Tile IR and a verdict-driven multi-agent workflow, the team converted 24 TileGym operators unattended and produced cuTile Rust kernels that are verifiably correct and perform at parity with existing cuTile Python implementations. The approach emphasizes machine-checkable artifacts, structured blame assignment, and bounded retries to make automated kernel translation repeatable rather than accidental.
