Burn 0.22.0: Faster Builds, Easier Extensions, and Smarter Autotuning

Burn 0.22 makes it easier to build applications, integrate custom kernels, and put models to work. In the previous release, we introduced burn-dispatch and started working toward removing backend generics from the user API with two goals in mind: simplify application code and improve compilation times. With this release, your code no longer carries backend type parameters, execution is instead selected through the device. In our tests, these changes made rebuilds up to 15× faster.

Beyond the API changes, we've improved execution across CPUs and GPUs with adaptive memory management, smarter autotuning, graph replay, and changes to the compiler and kernel foundations. The new CubeCL Environment lets applications reuse warmed compilation and autotuning caches on compatible targets. This release also adds support for adapting existing models with LoRA and QLoRA, ONNX export, and expanded support for remote compute. This post covers the major changes. The full list of improvements and fixes is available in the release notes.

Simpler Models, Faster Builds

Until now, choosing a backend affected the types throughout a Burn application. A model carried a B: Backend type, from its layers all the way down to their respective parameters, and tensor functions had to propagate it. This gave us a flexible architecture, but it also tied application code to a large dependency chain during compilation.

In 0.22, that backend parameter disappears from the user-facing API. The change is visible even in a small model:

 #[derive(Module, Debug)]
-pub struct Model<B: Backend> {
-    linear: Linear<B>,
+pub struct Model {
+    linear: Linear,
 }

-impl<B: Backend> Model<B> {
-    pub fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
+impl Model {
+    pub fn forward(&self, input: Tensor<2>) -> Tensor<2> {
         self.linear.forward(input)
     }
 }

The device now selects how and where execution happens. Enable the corresponding Cargo features, then choose a device when initializing the model and its inputs:

use burn::tensor::Device;

let device = Device::cuda(0);
let device = Device::wgpu(Default::default());
let device = Device::flex();

Training follows the same pattern, which means the B: AutodiffBackend no longer exists in application code. Instead, autodiff context is configured on the device and inherited by the tensors.

Underneath, operations follow Tensor → Bridge → Dispatch → Backend. The bridge hides the concrete backend representation from the high-level tensor API, breaking the dependency chain that previously reached into application code. This type erasure reduces recompilation when you make changes to your project. The dispatch, on the other hand, routes each operation to the corresponding backend, allowing a single application to easily enable mutiple backends. The Backend trait remains at the core of Burn to support backend implementations and custom operations.

To give you an example of how this affects compilation times, we tested model changes in two separate projects. One uses a small CNN with Burn's built-in learner, the other uses a transformer with a custom training loop. Adding or removing a hidden layer in the CNN benchmark reduced the median release rebuild from 28.42 seconds to 4.57 seconds. For the transformer, alternating equivalent feedforward expressions reduced its rebuild from 14.73 seconds to 1.00 second. With these two setups, we can see how compilation improves across different models and dependency footprints.

Easier Custom Operations

Custom operations are still implemented using backend primitives, even though the application code no longer carries backend generics. If you already had a backend extension trait, its definition and implementations remain pretty much the same. The #[backend_extension] macro connects it to the dispatch system, so you can expose your operation through a tensor function without backend generics. Here is an excerpt from our custom CubeCL kernel example, which combines matrix multiplication, bias addition, and ReLU:

use burn::{
    backend::{Dispatch, backend_extension, tensor::FloatTensor},
    tensor::Tensor,
};

#[backend_extension(Autodiff, Cube, Fusion)]
pub trait Backend: burn::backend::Backend {
    #[fusion(dtype = lhs, shape = output_shape(lhs, rhs, bias))]
    fn fused_matmul_add_relu(
        lhs: FloatTensor<Self>,
        rhs: FloatTensor<Self>,
        bias: FloatTensor<Self>,
    ) -> FloatTensor<Self>;
}

pub fn matmul_add_relu_custom(
    lhs: Tensor<3>,
    rhs: Tensor<3>,
    bias: Tensor<3>,
) -> Tensor<3> {
    Tensor::from_dispatch(Dispatch::fused_matmul_add_relu(
        lhs.into_dispatch(),
        rhs.into_dispatch(),
        bias.into_dispatch(),
    ))
}

The output_shape helper describes the output before the kernel runs. With that metadata, the macro generates registration for lazy execution through Fusion, alongside dispatch routing. The full example includes the kernel and its backward implementation, so callers can use the same function for training and inference.

This gives extension authors control over the implementation, including a custom backward pass. By default, custom operations act as boundaries in the fusion graph. Operations on either side can still be fused, but combining them with the custom operation requires a dedicated fusion optimization. The extension guide explains the complete workflow. Burn's linear algebra and signal processing functionality now use this approach through the burn-linalg and burn-signal extension crates.

Faster Execution, Better Memory Management

While building the new AI engine on top of Burn (more on that in What's Next), we developed a way to tune memory pools using allocation statistics collected during a dry run. That work became the basis for CubeCL's new memory management. We also improved performance across CubeCL and CubeK, including several convolution optimizations.

To measure the impact, we compared Burn 0.21 and 0.22 on two training workloads, using the same two projects, running in FP32 with CUDA on an NVIDIA GeForce RTX 4050 Laptop GPU. Each workload ran three times per version, with 50 warmup steps followed by 500 measured steps. The graphs show medians across those runs, with error bars for the observed min-max ranges.

The convolution optimizations helped increase training throughput by 80% in our CNN benchmark. This work includes faster convolution gradient computation, an improved direct kernel for transposed convolution, and a dedicated BatchNorm training operation with an optimized backward pass.

We also improved how fusion handles tensor layouts to avoid unnecessary copies and allow more operations to be fused after an unfusable operation.

Memory That Adapts to the Workload

CubeCL already reused device memory through pools, avoiding repeated allocations. But when the pool sizes don't match the workload, they can reserve more memory than the tensors need. The new adaptive memory pools adjust their allocation sizes around the workload, without requiring users to configure pool sizes for each model. backing allocations around the workload, without requiring users to configure pool sizes for each model.

As larger allocations arrive, the pools adjust their page sizes and release outdated pages once they become empty. The runtime can also move live allocations to free them sooner, while keeping frequently changing small allocations in a separate pool. Memory reports now show both used and reserved memory, so you can see how much of the reserved space your workload actually needs. This is available through Burn's Device API, which exposes more CubeCL utilities in this release.

As the benchmarks show, peak VRAM usage after warmup fell from 956 to 486 MiB for the CNN, which is nearly half. For the transformer, median peak VRAM fell from 3,486 to 2,868 MiB, a 17.7% reduction. These measurements represent process VRAM after warmup. They exclude CubeCL's temporary startup allocations used to measure the device's memory bandwidth and guide autotuning.

For convolutional models, improvements cover both training and inference. Convolution gradients have more efficient execution paths, transposed convolution has a dedicated direct kernel, and BatchNorm training uses a dedicated operation with a closed-form backward pass.

Smarter Autotuning and Runtime Execution

In the previous release post, we discussed using hardware throughput to decide when to stop autotuning. The adaptive autotuning scheduler uses measured compute and memory throughput to estimate a performance target for an operation. When a candidate reaches the expected performance ,the tuner can stop exploring. It also eliminates weaker candidates early and adjusts how much measurement each candidate needs.

We also reduced overhead for repeated execution by adding graph capture and replay support on CUDA and HIP. This lets CubeCL reuse a sequence of kernel launches without repeating much of the launch preparation. For WGPU, replay is supported via a software graph that retains resolved pipelines and bindings.

CubeCL Environment and Deployment Caches

CubeCL already cached compiled kernels and autotuning results between runs, but we've reworked that caching around the new CubeCL Environment, which groups those together in a named environment. You can populate an environment's caches ahead of time, export them together, and ship them with your application. The environment also provides a common API for runtime configuration and cache storage across native applications, WebAssembly, and no_std.

We also expanded profiling support across Burn and CubeCL, including device timing on CUDA and HIP to help identify where execution time is spent.

New Compiler and Kernel Foundations

The work on custom operations and portable execution also required changes deeper in the stack. CubeCL has migrated its compiler infrastructure to Pliron, giving us a common foundation for representing, transforming, and lowering kernels. Its LLVM compiler now includes targets for AMD and NVIDIA GPUs, extending the LLVM path beyond CPU execution.

CubeK's tile abstraction addresses another challenge: complex kernels often need the same pieces of logic for loading data, staging it in memory, computing on it, and writing results. Making those pieces reusable lets kernel authors compose implementations for different layouts, precisions, and hardware capabilities.

That work is already used in tiled matrix multiplication, attention, and interpolation implementations. It also gives us a better foundation for extending and optimizing those kernels. Together with the backend extension API, it connects application-level customization to reusable compute building blocks further down the stack.

Fine-Tuning with LoRA and QLoRA

Burn 0.22 adds support for adapting existing models to new tasks with LoRA and QLoRA. LoRA keeps the base weights frozen and trains low-rank adapters, reducing the number of parameters that need gradients and optimizer state. QLoRA also keeps the base weights quantized to reduce their storage requirements.

You can attach adapters to an existing model without changing its layers or forward pass. Initialize the model on an autodiff device, select the parameters to adapt, and apply LoRA:

use burn::module::{Lora, Module, ParamGroup};

let attention = ParamGroup::from_predicate("attention");
let model = model.apply_lora(
    Lora::new(8, 16.0).set_param_group(attention),
);

This freezes the base floating-point parameters and attaches trainable adapters to matching two-dimensional weights. LoRA is built on a generalized weight reparameterization mechanism: a parameter keeps its stored base and an attached module computes its effective value. Adapter parameters participate in autodiff, optimization, device transfers, and checkpointing through the existing module infrastructure.

Controlling What Gets Trained

Parameter groups also connect fine-tuning to the optimizer. You can keep the base weights frozen, enable a task-specific head, and assign different optimizers or learning-rate policies to the head and adapters. For example, after attaching adapters to a model with a classifier head:

use burn::optim::AdamWConfig;

let classifier = ParamGroup::from_predicate("classifier");
let adapters = ParamGroup::from_regex(r"(^|\.)lora\.(a|b)$")?;

let model = model.unfreeze_group(classifier);
let optimizer = AdamWConfig::new().init().with_group(
    adapters,
    AdamWConfig::new().with_weight_decay(0.0).init(),
    None,
);

The classifier uses the default optimizer, while the adapters use their own configuration. Learning-rate schedulers support the same grouping model, so optimizer choice and learning-rate policy can be controlled independently.

Quantized Base Weights

QLoRA extends that workflow by keeping the frozen base weights quantized while training the adapters. The broader quantization work in this release supports it: two-level scales combine block-level scales with a tensor-wide scale, lookup-table quantization reconstructs values from a codebook, and software FP8 conversion makes low-precision conversion available across more hardware. We use native conversion when available, but supporting a format doesn't mean every operation can use accelerated arithmetic.

QLoRA currently reduces the memory used to store frozen base weights, but training still builds dense effective weights when combining them with the adapters. When input gradients are needed, those weights are retained for backward, which also computes full weight-shaped gradients before propagating into the adapter factors. Fusion can avoid some intermediate allocations, but it doesn't eliminate all of this overhead. Peak training memory can therefore be much higher than the packed weights and adapters alone suggest. See the implementation notes when sizing a training run.

Remote Compute

Runtime device selection also makes remote execution easier to integrate. The model keeps its usual tensor operations, while a remote device sends those operations to another machine's compute resources.

Burn Remote now uses Iroh as its primary transport. Peers prefer direct connections and can use relays when needed, and every connection is authenticated and encrypted. A server can expose multiple devices. With the fusion feature, a repeated operation graph is sent once and then replayed on the server by its id, which reduces communication. Tensor transfers between Iroh servers go directly between them, avoiding a round trip through the client. WebSocket remains available for trusted networks.

On the machine with the GPU, a server is the list of devices it hosts and a transport:

use burn::server::{
    IrohIdentity, IrohTransport, RemoteServer, TokenAuthorizer,
};
use burn::tensor::Device;

let identity = IrohIdentity::load_or_create("server.key")?;
let transport = IrohTransport::new(identity);
println!("server id: {}", transport.id());

RemoteServer::new([Device::cuda(0)])
    .with_authorizer(TokenAuthorizer::new(token)?)
    .serve(transport)?;

A client describes that server with a RemoteHost, from its id and the token its authorizer checks, and connects one of its devices. An existing model then moves to the remote device for inference, with its inputs created there:

use burn::module::Module;
use burn::remote::RemoteHost;
use burn::tensor::{Device, Tensor};

let host = RemoteHost::iroh(server_id).with_credential(token);
let device = Device::remote_options(&host).init()?;

let model = model.valid().to_device(&device);
let input = Tensor::<2>::ones([32, 128], &device);
let output = model.forward(input);

.device_index(1) picks another of the server's devices, and host.devices() lists all of them. The remote MNIST example demonstrates both training and inference on another machine's GPU. The remote guide covers server setup, relay configuration, and authorization.

Export Burn Models to ONNX

Burn's ONNX tooling has focused on importing models into readable Rust code. This release adds the other direction: exporting a Burn model to ONNX through burn-onnx.

The new computational graph capture infrastructure records the operations performed by a model's forward pass. The exporter lowers that representation into an ONNX model, including its weights. For the model used at the start of this post, initialized with a 128-feature linear layer:

use burn::module::Module;
use burn::tensor::Tensor;
use burn_onnx::export::OnnxExporter;

let model = model.valid();
let input = Tensor::<2>::zeros([1, 128], &device);

OnnxExporter::new()
    .export(&model, input, Model::forward)?
    .save("model.onnx")?;

This capture records the model's operations for export. The runtime graph capture and replay described earlier reuses an execution sequence on a device.

Export support is still early, and operator coverage remains limited. A model supported by Burn or by the ONNX importer is not necessarily exportable yet. The initial implementation includes vision-model validation, and coverage will expand incrementally. See the exporter for current capabilities.

Upgrading to Burn 0.22

This release includes substantial API changes. The 0.22 migration guide covers removing backend generics, selecting devices, and updating autodiff and training code.

Checkpointing becomes simpler as well. Burnpack is now the common record format for models, optimizers, and learning-rate schedulers, replacing the recorder and precision-settings machinery. Saving and loading model weights looks like this:

use burn::module::Module;

model.clone().save_file("model.bpk")?;
let model = model.try_load_file("model.bpk")?;

Legacy MessagePack, binary, and JSON recorder checkpoints require conversion before Burn 0.22 can load them. Renaming a file does not convert it. The migration guide explains how to export model weights from an older Burn project and load them into the corresponding 0.22 model. Legacy optimizer and scheduler state need to be handled separately.

Autodiff is now runtime state, and dataset access and more training APIs are fallible. Training and evaluation results report failures separately from user interruption, giving applications more control over how a run ends. The guide also covers Cargo feature changes and the Rust 1.95 requirement. Candle has been removed, while NdArray and LibTorch are deprecated.

What's Next

The new AI engine mentioned earlier has been driving several of these improvements, and we plan to release it publicly soon. Building applications with it has helped us identify gaps in Burn's APIs and bottlenecks in the runtime.

The engine uses CubeCL Environment to prepare compilation and autotuning caches before generation starts. A command-line tool, a native chat application, and an OpenAI-compatible server all use the same Rust library. In the release announcement, we'll introduce the supported models and show how to use that library in your own applications.

In the meantime, give Burn 0.22 a try and let us know how it works for your projects. We'd love to hear what you're building and where you're running into limitations.