Translating CUDA Tile Operations from Python to Rust Using Agentic AI

來源原標題

Translating CUDA Tile Operations from Python to Rust Using Agentic AI

作者

Tanya Lenz

原始標籤

Agentic AI / Generative AI · Developer Tools & Techniques · AI Agent · CUDA Tile · cuTile · Mixture of Experts (MoE) · Python

原始連結

https://developer.nvidia.com/blog/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. 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: @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(...