Rust + Burn = C.U.M.

Compute-Optimized Unified Modeling

Photo by Lorenzo Herrera on Unsplash

Rust-Powered AI: Building a ChatGPT with CubeCL and Burn

Rust is quickly gaining traction as a language for machine learning. Its safety and performance make it an attractive alternative to Python-based frameworks. Recently, Tracel AI released CubeCL and Burn — an end-to-end Rust stack for high-performance deep learning.

Burn is a next-generation deep learning framework in pure Rust, and CubeCL is a GPU compute language that integrates with Burn to run kernels on any hardware. Together, they promise to let developers write ChatGPT-style transformer models entirely in Rust, leveraging GPU acceleration without sacrificing safety or portability.

In this article, we’ll explore what these projects are, how they work under the hood, and walk through how one might implement a transformer-based language model with them. We’ll also compare this Rust ML stack to mainstream tools like PyTorch, looking at performance, ergonomics, and ecosystem maturity.

What are CubeCL and Burn?

CubeCL is a “multi-platform high-performance compute language extension for Rust”. Think of it as a Rust-native way to write GPU kernels: you annotate Rust functions with a #[cube] macro and the CubeCL compiler turns them into optimized GPU code. It supports CUDA (NVIDIA), ROCm (AMD), Vulkan (via wgpu), Metal (Apple), and even WebGPU backends. The goal is to let you program GPUs with zero-cost Rust abstractions – automatic vectorization, compile-time optimizations, and JIT autotuning – instead of hand-writing CUDA or other shading languages. CubeCL’s design uses Rust’s proc-macros: it parses your kernel code in Rust, then emits a new Rust function that, when run, generates optimized GPU IR at runtime. This gives features like comptime evaluation, automatic vectorization of loops, and custom kernel fusion. In short, CubeCL solves the problem of writing portable, high-performance GPU code in Rust: you write one version in Rust and CubeCL will JIT-compile it to run efficiently on AMD GPUs, NVIDIA GPUs, even in a browser via WebGPU.

Burn is built on top of CubeCL and other backends. It’s a fully Rust deep learning framework (dynamic graphs with static performance). Burn’s tagline is “flexibility, efficiency and portability”. Burn provides standard neural-network building blocks (tensors, layers, loss functions, optimizers, etc.), but abstracts over a choice of compute backend. You might use CPU backends (like an ndarray-based backend or a PyTorch/libtorch backend via tch-rs), or GPU backends (WGPU/Vulkan or CUDA). The key design is that Burn keeps your model code generic over a Backend type – this lets you write one model and run it on different hardware just by swapping a type parameter. Burn’s novel twist is its tensor-stream architecture and Just-In-Time compiler: it tracks tensor usage precisely (thanks to Rust’s ownership) and fuses operations at runtime to squeeze out performance. The result is dynamic graphs and shapes (you can’t easily do that in pure static-graph frameworks) but without the usual overhead: Burn compiles combined kernels on-the-fly and auto-tunes them for your hardware.

Both projects aim to solve common problems in ML development. CubeCL tackles hardware portability and manual GPU-kernel writing: instead of maintaining separate CUDA/Metal/Vulkan code, you get one high-level Rust kernel. Burn addresses the “dependency hell” and performance tuning headaches of existing frameworks: it promises “blazingly fast” execution with minimal user tweaking, cross-platform inference (cloud to mobile), and Rust’s compile-time safety. Together, CubeCL and Burn form a Rust-native ML stack that’s designed for speed and safety.

CubeCL: Writing GPU Kernels in Rust

Under the hood, CubeCL takes Rust code and lowers it into GPU compute kernels. You decorate functions with #[cube], and CubeCL’s proc macro rewrites them. For example, a simple element-wise GELU implementation might look like this in CubeCL:

use cubecl::prelude::*;

#[cube(launch_unchecked)]
fn gelu_array<F: Float>(input: &Array<Line<F>>, output: &mut Array<Line<F>>)
{
    if ABSOLUTE_POS < input.len() {
        output[ABSOLUTE_POS] = gelu_scalar(input[ABSOLUTE_POS]);
    }
}

#[cube]
fn gelu_scalar<F: Float>(x: Line<F>) -> Line<F> {
    // compute sqrt(2) at compile time
    let sqrt2 = F::new(comptime!(2.0f32.sqrt()));
    let tmp = x / Line::new(sqrt2);
    x * (Line::erf(tmp) + 1.0) / 2.0
}

This Rust code, once compiled and launched, ends up running as an optimized GPU shader. CubeCL handles the details: it will generate a WGSL or CUDA kernel, handle memory transfers, and even fuse operations. For example, if you write your own gelu like above, at runtime CubeCL can auto-generate a single fused kernel to compute the whole thing, rivaling a hand-written CUDA implementation. It also auto-vectorizes the kernel (using SIMD lanes in “Line<F>” vectors) and auto-tunes the launch dimensions. All the while, the CubeCL runtime reuses memory buffers and does lazy evaluation to minimize copies.

CubeCL is designed for cube-like GPU topologies. A GPU executes work in 3D blocks (“cubes”), each made of smaller units (“threads”) and groups (“planes” or warps). CubeCL’s naming reflects this: you launch kernels in “CubeCount” and “CubeDim” to match GPU threadblocks. Its Topology concept (see figure below) maps these cubes and sub-cubes so that kernels can efficiently synchronize or share memory on an SM.

Figure: CubeCL’s cube-and-hypercube topology (adapted) shows how computation is divided into 3D blocks of work. This 3×3×3 cube has 27 units, and hypercubes stack these cubes in 3D, matching GPU thread block hierarchy.

Behind the scenes, CubeCL’s two-step macro system first parses your GPU kernel code into a syntax tree, then generates a new Rust function that creates the kernel IR when executed. This means all Rust optimizations (like comptime! and generics) happen naturally. The proc macro approach also lets CubeCL automatically pick vector lengths: it inspects your code and figures out how many lanes of SIMD to use in each dimension. If you use 32-bit floats, for example, it might pack 4 floats into a single Line<f32> SIMD vector.

Finally, CubeCL supports multiple platforms out of the box. You get the same code base running on:

• CUDA (NVIDIA) via a tiny C++ JIT and PTX generation.
• ROCm/HIP (AMD) with HIP C++.
• WebGPU (cross-vendor) using WGSL shaders (runs on Desktop and Web).
• Metal (Apple) through wgpu.
• Vulkan (Linux/Windows) via SPIR-V (wgpu).

This means a single #[cube] kernel can run on an NVIDIA GPU, an AMD GPU, even in your browser. If a backend doesn’t support a given instruction (like tensor cores on WebGPU), CubeCL will throw a runtime error or fall back. But usually it “just works” – the launch function in your code will pick the right backend for the detected device.

In summary, CubeCL abstracts away the gritty details of GPU programming. You don’t hand-code CUDA kernels or manage streams; you write regular Rust functions and trust CubeCL to JIT-compile them. The payoff is a maintainable GPU layer for Burn and other Rust numeric code, bringing portability and performance together.

Burn: A Dynamic Deep Learning Framework

Burn is the high-level side of the stack: it provides tensors, layers, and training infrastructure. At its core Burn has a Tensor<B, D> type parameterized by a backend B: Backend and a dimension D. For example, Tensor<Wgpu, 2> might be a 2-D float tensor on a WGPU (Vulkan) device. You write your model generically over any B: Backend. This backend-agnostic design allows you to swap out an ndarray (CPU) backend for a Wgpu (GPU) backend with a type change. Burn even provides a Router backend decorator to combine multiple devices: you could run some layers on CPU and others on GPU by using Router<(Cuda, NdArray)> and specifying devices at runtime.

Burn’s architecture is heavily inspired by PyTorch and other ML frameworks, but with Rust idioms. For example, models are modules that implement a forward method. The burn::nn library includes layers like Linear, Conv2d, ReLU, GELU, LayerNorm, and even a full MultiHeadAttention module. The MultiHeadAttention struct holds Linear layers for query, key, value, and output, plus dropout and activation – essentially the same building blocks from “Attention is All You Need”. To use it you’d create a MultiHeadAttentionConfig and call MultiHeadAttention::new(config). Then in your model’s forward pass you’d pass tensors of shape [batch, seq_len, d_model] to it. The output is also [batch, seq_len, d_model]. Importantly, Burn also supports a forward_cache method on attention to cache past keys/values for fast autoregressive generation (like GPT caching).

Under the hood, Burn uses automatic differentiation (reverse-mode) to compute gradients. As of recent releases, Burn completely rewrote its autograd engine to be more “Rusty” and memory-efficient. It now allows any backend to support backprop by tracking operations and doing in-place updates when safe. In practice, this means you write normal Rust code for the forward pass, then call .backward() on a loss tensor to propagate gradients to all parameters. Burn even avoids unnecessary copies: if a tensor has a single owner, Burn will do the operation in-place rather than allocate a new tensor, similar to PyTorch’s patterns but built into the backend API. This leads to lower memory usage (Burn’s authors note especially on CPU Burn uses memory far more efficiently than PyTorch).

Burn provides all the usual training tools. You have Optimizers like Adam, and Loss functions like CrossEntropyLoss in the burn::loss module. A typical training loop looks like:

use burn::tensor::Tensor;
use burn::optim::{Adam, Optimizer};
use burn::loss::CrossEntropyLoss;
// Assume `model` implements Module<Tensor<f32>> and is mutable
let mut optimizer = Adam::new(&model, Default::default());
let loss_fn = CrossEntropyLoss::new();

// For each epoch...
for (batch_x, batch_y) in data_loader {
    let preds = model.forward(batch_x.clone());
    let loss = loss_fn.forward(preds.clone(), batch_y.clone());
    optimizer.zero_grad();
    loss.backward();      // auto-diff computes gradients
    optimizer.step();     // update model weights
}

This is nearly identical to PyTorch code. The .forward and .backward magic is handled by Burn; the user never deals with pointers or backprop manually. Burn’s training loop is ergonomic: a recent blog example showed exactly this pattern. You can even visualize training with Burn’s built-in terminal dashboard (powered by the ratatui crate) to track metrics in real time.

Because Burn is fully written in Rust, it also integrates deeply with systems. It supports exporting and importing models via ONNX and PyTorch formats. That means you can take a pretrained PyTorch transformer, export it to ONNX or safetensors, and load its weights into a Burn-defined model. For deployment, Burn supports running in WebAssembly: the WGPU backend can run in browsers or any WebGPU runtime, and an ndarray CPU backend works in no_std environments. For example, Burn demo apps exist for MNIST in the browser using WGPU.

Performance is a core goal. Burn’s JIT and fusion backend (called burn_fusion) attempt to combine adjacent operations into single kernels to minimize data movement. All first-party GPU backends use this fusion by default. The team is also working on automatic gradient checkpointing to trade compute for memory when needed. In practice, Burn has shown impressive raw speed. The developers highlight (and benchmarks confirm) that Burn’s matrix multiply on NVIDIA GPUs can match or even beat NVIDIA’s own cuBLAS library. We’ll look at that next.

Building a Transformer (ChatGPT) in Rust

With CubeCL and Burn, building a transformer model looks conceptually similar to PyTorch, but in Rust. You’d start by defining your model struct, using Burn’s built-in layers. For example, a tiny GPT-like model might be:

use burn::tensor::Tensor;
use burn::module::{Module, Param};
use burn::nn::{Embedding, MultiHeadAttention, LayerNorm, Linear};

struct GPTConfig {
    vocab_size: usize,
    d_model: usize,
    n_heads: usize,
    n_layers: usize,
    // ... other hyperparameters ...
}

struct GPT2<B: Backend> {
    wte: Embedding<B>,            // token embedding
    wpe: Embedding<B>,            // positional embedding
    blocks: Vec<TransformerBlock<B>>,
    ln_f: LayerNorm<B>,
    lm_head: Linear<B>,           // final projection
}

impl<B: Backend> GPT2<B> {
    pub fn new(config: &GPTConfig) -> Self {
        let wte = Embedding::new(config.vocab_size, config.d_model);
        let wpe = Embedding::new(config.max_len, config.d_model);
        let mut blocks = Vec::new();
        for _ in 0..config.n_layers {
            blocks.push(TransformerBlock::new(config.d_model, config.n_heads));
        }
        let ln_f = LayerNorm::new(config.d_model);
        let lm_head = Linear::new(config.d_model, config.vocab_size);
        Self { wte, wpe, blocks, ln_f, lm_head }
    }
}

impl<B: Backend> Module<Tensor<B>> for GPT2<B> {
    fn forward(&self, input_ids: Tensor<B, 2>) -> Tensor<B, 3> {
        // input_ids: [batch, seq_len]
        let seq_len = input_ids.dims()[1];
        let positions = Tensor::arange(0, seq_len, &input_ids.device());
        let tok_emb = self.wte.forward(input_ids.clone());     // [batch, seq, d_model]
        let pos_emb = self.wpe.forward(positions.expand(input_ids.dims())); 
        let mut x = tok_emb + pos_emb;                          // add token+pos embeddings

        for block in &self.blocks {
            x = block.forward(x);
        }
        let x = self.ln_f.forward(x);                            // final layer norm
        self.lm_head.forward(x)                                 // project to vocab logits [batch, seq, vocab]
    }
}

Here, TransformerBlock would be a small module we define that internally does one self-attention layer followed by a feed-forward network, with residuals. For example:

struct TransformerBlock<B: Backend> {
    attn: MultiHeadAttention<B>,
    attn_ln: LayerNorm<B>,
    ff: (Linear<B>, GeLU, Linear<B>),  // feed-forward MLP: Linear->GeLU->Linear
    ff_ln: LayerNorm<B>,
}

impl<B: Backend> TransformerBlock<B> {
    pub fn new(d_model: usize, n_heads: usize) -> Self {
        let attn = MultiHeadAttentionConfig::new(d_model, n_heads).init();
        let attn_ln = LayerNormConfig::new(d_model).init();
        let ff = (
            Linear::new(d_model, 4*d_model),   // expand hidden size 4x
            GeLU::new(),
            Linear::new(4*d_model, d_model),
        );
        let ff_ln = LayerNormConfig::new(d_model).init();
        Self { attn, attn_ln, ff, ff_ln }
    }
}

impl<B: Backend> Module<Tensor<B>> for TransformerBlock<B> {
    fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
        // Self-attention sublayer
        let y = self.attn.forward(MhaInput::new(x.clone(), x.clone(), x.clone()));
        let x = (x + y).apply(self.attn_ln.clone());
        // Feed-forward sublayer
        let y = (self.ff.1.clone())(self.ff.0.forward(x.clone()));
        let y = self.ff.2.forward(y);
        (x + y).apply(self.ff_ln.clone())
    }
}

(Note: The actual API names might differ; this is illustrative.) Each forward returns a new Tensor, and Burn tracks the computation to do backprop later.

Because Burn’s tensors carry the device and backend in their type, you would compile or run your code with a chosen backend, e.g. GPT2::<burn::backend::Wgpu> for GPU training. The code above is pure Rust – no Python or Torch dependencies. For training, you’d pick an optimizer and loss (e.g. CrossEntropyLoss on the final logits). Thanks to Burn’s AutodiffModule trait, you can then do:

let mut model = GPT2::<burn::backend::Wgpu>::new(&config);
let mut opt = Adam::new(&model, Default::default());

for batch in dataloader {
    let inputs = batch.inputs.clone();
    let targets = batch.targets.clone();
    let logits = model.forward(inputs);
    let loss = CrossEntropyLoss::new().forward(logits.clone(), targets.clone());

    opt.zero_grad();
    loss.backward();    // compute gradients through the whole model
    opt.step();         // update weights
}

Despite the change of language, this is almost exactly what you’d write in Python with PyTorch. Burn handles the details: the .backward() triggers its autograd engine, which uses the chosen backend (via CubeCL for GPU) to compute gradients of every Linear and Embedding weight.

During inference (text generation), you can feed tokens one by one. Burn’s MultiHeadAttention supports a forward_cache API which reuses previous key/value computations so that each new token isn’t O(n²) on the entire prefix. This mimics how GPT generation is usually implemented. CubeCL comes into play here for acceleration: whenever Burn does a tensor operation (e.g. a matrix multiply or a GELU activation), on a GPU backend it will use CubeCL to run a fused kernel. For example, multiplying two tensors is handled by CubeCL’s BLAS-like kernels. Hence the whole transformer runs on the GPU via CubeCL (or the CPU if you chose a CPU backend).

Performance and Maturity Compared to PyTorch

The Rust ML ecosystem is still young compared to PyTorch’s decade of development. PyTorch has a huge collection of ready models and a large user community. Burn is just emerging, with around 12.5k GitHub stars and a growing contributor base. For now, many pretrained models still live in Python, and the tooling (like tokenizers, datasets, HuggingFace integration) is more limited for Rust. However, Burn is actively working on import/export features — you can load ONNX or PyTorch weights directly and run them with Burn, bridging the gap.

On developer experience, Rust offers some advantages. There is no Global Interpreter Lock or Python dependency hell — you compile your model into a native binary or WebAssembly. Rust’s strong typing and borrow checker catch many errors at compile time (for example, mixing up tensor shapes or forgetting to clone). Many developers find this safer and easier to refactor. Burn itself strives to be ergonomic: its API is modular (like layers and Sequential containers shown in [41]) and it even has nice utilities like a terminal training dashboard. But Rust code can be more verbose than Python, and compilation times can be longer. The Rust ecosystem around ML is also smaller, so you might have to write more custom code yourself than rely on a plug-and-play library.

On performance, early indicators are very promising for Burn/CubeCL. A recent benchmark from Tracel showed that Burn’s own matrix-multiply kernels (entirely open-source Rust) outperform NVIDIA’s highly tuned cuBLAS on some GPUs. In tests on an NVIDIA RTX GPU, Burn’s “Simple” GEMM algorithm was “remarkably fast and stable, nearly always outperforming the cuBLAS/CUTLASS reference”. (A “MultiRow” variant even topped the charts across the board.) We show one example benchmark below:

Burn’s Rust-based MATMUL vs NVIDIA’s CUDA/cuBLAS on a modern GPU. Burn’s custom JIT-compiled kernels (green) match or beat NVIDIA’s libraries (blue) across various matrix sizes. The graph is from Tracel’s July 2025 report.

That result is striking: it means that a full Rust ML stack (Burn/CubeCL) can match the performance of NVIDIA’s proprietary libraries, while also running on non-NVIDIA hardware (AMD, Intel, etc.) thanks to Vulkan support. In practice, Burn can dispatch operations to whichever backend is fastest. So if you have an AMD GPU, you could use the ROCm or Vulkan path and still get accelerated training.

By contrast, PyTorch is extremely mature: it has years of optimizations on NVIDIA and Intel hardware, and countless model implementations. In benchmarks on standard tasks, PyTorch (LibTorch + cuBLAS) still often wins by a small margin on NVIDIA GPUs. But PyTorch requires the CUDA toolchain, GPUs with specific drivers, and sometimes fiddling with library versions. Burn/CubeCL aims for broader compatibility: for example, a Burn model can run in-browser on Apple Silicon or Windows with equal ease, whereas PyTorch would not.

In terms of ergonomics, PyTorch and Burn share similar high-level design: dynamic graphs, autograd, familiar layers. But Burn’s code is in Rust. This means all your model code is compiled ahead of time, which catches many bugs early. On the other hand, dynamic Python allows rapid iteration in notebooks, which Rust currently lacks. Burn is improving here (it has a “Burn Book” and REPL workflows planned), but Python’s REPL and dynamic scripting remain easier for quick experiments.

Finally, ecosystem and tooling: PyTorch has integrations with ecosystem tools (TensorBoard, Optuna, etc.), whereas Burn is building that out. But Burn has some unique features: for example, because it’s Rust, it can compile down to WebAssembly and run fully on-device or embedded (no host Python). It also supports advanced backend composition (the “Router” and “Remote” backends) to distribute work, which is harder to set up in PyTorch without additional frameworks.

Conclusion

CubeCL and Burn represent an exciting new frontier: a fully Rust native stack for building and training neural networks. We’ve seen that CubeCL abstracts over GPUs in a portable way, and that Burn provides the familiar deep-learning building blocks with Rust’s safety and performance. A ChatGPT-style transformer can be implemented in Rust by composing Burn’s layers (embeddings, attention, MLPs, etc.), training with Burn’s autograd, and running efficiently on GPUs via CubeCL. Early benchmarks indicate that this Rust combo can rival or even beat traditional ML toolchains on speed, while offering advantages in deployment and safety.

For developers interested in machine learning and Rust, Burn and CubeCL open up new possibilities. You can write GPU code in Rust without leaving the language, and you can train models without Python. The projects are still evolving — not everything in PyTorch exists yet — but their design is modern and ambitious. As one user commented, it’s “insane how it took this long to get a fully cross-platform ML runtime that didn’t rely on proprietary vendor lock-in”. Indeed, combining Rust’s memory safety and ownership with automatic differentiation and JIT acceleration is a compelling idea.

The road to a full-fledged Rust AI ecosystem will take time, but CubeCL and Burn have built strong foundations. They invite contributions (e.g. more kernel algorithms, layers, and tutorials) to reach feature parity with established frameworks. Until then, they already allow you to explore Rust-powered deep learning: building transformer models, training them on GPUs, or even running them in the browser — all with a Rust toolchain. For the future of high-performance, safe, and portable AI, CubeCL and Burn are charting a promising course.

Sources:

Official docs and GitHub for CubeCL and Burn, Tracel AI blog posts, and community articles provide details on design goals and performance. The Burn documentation and tutorials (get-started guide, examples) illustrate usage. All quoted text and benchmarks are from those sources as cited.