Custom CubeCL Kernel

This example fuses matrix multiplication, bias addition, and ReLU into one custom operation. The example project contains the source and dependencies. The code below is included directly from that project so that changes to its API also update this chapter.

The CubeCL frontend compiles the kernel for the selected runtime. The example uses WGPU; its CubeBackend implementation is shared by CubeCL runtimes.

Custom Backend Trait

#[backend_extension(Autodiff, Cube, Fusion)] registers a trait with runtime dispatch. The low-level signature uses FloatTensor<Self> primitives, while the public wrapper accepts Tensor<3>. into_dispatch() and from_dispatch() cross that boundary without requiring a backend generic in the caller's code. The reference implementation composes ordinary tensor operations for comparison.

mod backward;
mod forward;
mod kernel;

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

/// We create our own Backend trait that extends the Burn backend trait.
#[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>;
}

/// We define our custom implementation using the added function on our custom backend.
pub fn matmul_add_relu_custom(lhs: Tensor<3>, rhs: Tensor<3>, bias: Tensor<3>) -> Tensor<3> {
    let output = Dispatch::fused_matmul_add_relu(
        lhs.into_dispatch(),
        rhs.into_dispatch(),
        bias.into_dispatch(),
    );

    Tensor::from_dispatch(output)
}

/// We define a reference implementation using basic tensor operations.
pub fn matmul_add_relu_reference(lhs: Tensor<3>, rhs: Tensor<3>, bias: Tensor<3>) -> Tensor<3> {
    let x = lhs.matmul(rhs) + bias;

    activation::relu(x)
}

fn output_shape(lhs: &Shape, rhs: &Shape, bias: &Shape) -> Shape {
    assert!(lhs.num_dims() >= 2, "matmul needs at least two dimensions");
    let shape = burn::backend::calculate_matmul_output(lhs, rhs).expect("compatible matmul shapes");
    assert_eq!(
        &shape, bias,
        "kernel requires bias to match the output shape"
    );
    shape
}

The Autodiff selector enables routing to the handwritten autodiff implementation below. There is no additional user-defined AutodiffBackend marker trait to implement.

Lazy Fusion registration

Enable Burn's fusion feature and list Fusion on #[backend_extension] as shown above. The annotation copies lhs's dtype and passes borrowed Shape values to output_shape. Fusion uses that metadata to register a lazy output, then calls the forward implementation when the operation executes.

Sharing output_shape with execution keeps both paths consistent: it checks matrix dimensions, batch broadcasting, and the requirement that bias match the output shape. The wrapper does not compare output shapes with backend results, so the metadata calculation must be correct.

The custom kernel remains opaque to the Fusion optimizer. The wrapper does not combine it with neighboring kernels or generate gradients; the handwritten Autodiff<B, C> implementation below supplies the backward pass.

For structured outputs, use #[fusion(meta = callable)]. The callback must describe the same result as direct backend execution, including the actual values of any non-tensor fields: Fusion returns those values immediately and discards the backend's later values without comparison. Use #[fusion(default)] to inherit an existing trait body, or omit Fusion from #[backend_extension] to write the Fusion implementation yourself.

Forward Kernel

The kernel computes each output from the corresponding row and column, then adds the bias and applies ReLU. This is an instructional implementation; production matmul kernels use more advanced tiling and hardware-specific optimizations.

use cubecl::{cube, prelude::*};

/// Declare a custom kernel that gets compiled to `wgpu`/`CUDA`
#[cube(launch)]
pub fn fused_matmul_add_relu_kernel<F: Float>(
    lhs: &Tensor<F>,
    rhs: &Tensor<F>,
    bias: &Tensor<F>,
    output: &mut Tensor<F>,
    #[define(F)] _dtype: ElemType,
) {
    let row = ABSOLUTE_POS_X as usize;
    let col = ABSOLUTE_POS_Y as usize;
    let batch = ABSOLUTE_POS_Z as usize;

    let n_rows = output.shape(output.rank() - 2);
    let n_cols = output.shape(output.rank() - 1);
    let dim_k = rhs.shape(rhs.rank() - 2);

    if row >= n_rows || col >= n_cols {
        terminate!();
    }

    let offset_output = batch * n_rows * n_cols;
    let mut offset_lhs = 0;
    let mut offset_rhs = 0;

    let batch_dims = output.rank() - 2;
    for dim in 0..batch_dims {
        offset_lhs += offset_output / output.stride(dim) % lhs.shape(dim) * lhs.stride(dim);
        offset_rhs += offset_output / output.stride(dim) % rhs.shape(dim) * rhs.stride(dim);
    }

    let mut sum = F::new(0.0_f32);
    for k in 0..dim_k {
        let lhs_index = row * dim_k + k;
        let rhs_index = k * n_cols + col;

        sum += lhs[offset_lhs + lhs_index] * rhs[offset_rhs + rhs_index];
    }

    let out_index = row * n_cols + col;
    let index = offset_output + out_index;

    output[index] = F::max(sum + bias[index], F::new(0.0_f32));
}

The forward implementation checks that operands share a device, makes them contiguous, allocates the output, and launches the kernel with the required grid and element type.

use crate::{FloatTensor, kernel::fused_matmul_add_relu_kernel};

use super::Backend;
use burn::backend::cubecl::dtype_to_storage_type;
use burn_cubecl::{CubeBackend, kernel::into_contiguous, tensor::CubeTensor};
use cubecl::{CubeCount, CubeDim};

/// Implement our custom backend trait for the cubecl backend. One impl covers every runtime:
/// a tensor's device says which one it runs on.
impl Backend for CubeBackend {
    fn fused_matmul_add_relu(
        lhs: FloatTensor<Self>,
        rhs: FloatTensor<Self>,
        bias: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
        let dtype = lhs.dtype;
        // Define cube dim, hardcoded for simplicity.
        let cube_dim = CubeDim::new_3d(16, 16, 1);

        lhs.assert_is_on_same_device(&rhs);
        lhs.assert_is_on_same_device(&bias);

        // For simplicity, make sure each tensor is continuous.
        let lhs = into_contiguous(lhs);
        let rhs = into_contiguous(rhs);
        let bias = into_contiguous(bias);

        assert_eq!(lhs.dtype, rhs.dtype, "matrix dtypes must match");
        assert_eq!(lhs.dtype, bias.dtype, "bias dtype must match");
        let shape_out = crate::output_shape(lhs.meta.shape(), rhs.meta.shape(), bias.meta.shape());

        // Get the matmul relevant shapes after validating the inputs.
        let ndims = lhs.meta.num_dims();
        let num_rows = lhs.meta.shape()[ndims - 2];
        let num_cols = rhs.meta.shape()[ndims - 1];
        let num_batches: usize = (0..ndims - 2).map(|i| shape_out[i]).product();

        // Create a buffer for the output tensor.
        let buffer = lhs.client.empty(shape_out.num_elements() * dtype.size());

        // Create the output tensor primitive.
        let output = CubeTensor::new_contiguous(
            lhs.client.clone(),
            lhs.device.clone(),
            shape_out,
            buffer,
            dtype,
        );

        // Declare the wgsl workgroup with the number of cubes in x, y and z.
        let cubes_needed_in_x = f32::ceil(num_rows as f32 / cube_dim.x as f32) as u32;
        let cubes_needed_in_y = f32::ceil(num_cols as f32 / cube_dim.y as f32) as u32;
        let cube_count =
            CubeCount::Static(cubes_needed_in_x, cubes_needed_in_y, num_batches as u32);

        // Execute lazily the kernel with the launch information and the given buffers. For
        // simplicity, no vectorization is performed
        fused_matmul_add_relu_kernel::launch(
            &output.client,
            cube_count,
            cube_dim,
            lhs.into_tensor_arg(),
            rhs.into_tensor_arg(),
            bias.into_tensor_arg(),
            output.clone().into_tensor_arg(),
            dtype_to_storage_type(dtype),
        );

        // Return the output tensor.
        output
    }
}

Backward

The custom trait is implemented for Autodiff<B, C> where the inner B implements the same trait. The backward state saves the operands and shape information needed to compute the derivatives. Tracked operations register that state with the graph; untracked operations only execute the forward computation.

use crate::FloatTensor;

use super::Backend;
use burn::{
    backend::{
        TensorMetadata,
        autodiff::{
            Autodiff, NodeId,
            checkpoint::{base::Checkpointer, strategy::CheckpointStrategy},
            grads::Gradients,
            ops::{Backward, Ops, OpsKind, broadcast_shape},
        },
    },
    tensor::Shape,
};

// Implement our custom backend trait for any backend that also implements our custom backend trait.
impl<B: Backend, C: CheckpointStrategy> Backend for Autodiff<B, C> {
    fn fused_matmul_add_relu(
        lhs: FloatTensor<Self>,
        rhs: FloatTensor<Self>,
        bias: FloatTensor<Self>,
    ) -> FloatTensor<Self> {
        // Create our zero-sized type that will implement the Backward trait.
        #[derive(Debug)]
        struct FusedMatmulAddReluBackward;

        // Implement the backward trait for the given backend B, the node gradient
        // with three other gradients to calculate (lhs, rhs, and bias).
        impl<B: Backend> Backward<B, 3> for FusedMatmulAddReluBackward {
            // Our state that we must build during the forward pass to compute the backward pass.
            //
            // Note that we could improve the performance further by only keeping the state of
            // tensors that are tracked, improving memory management, but for simplicity, we avoid
            // that part.
            type State = (NodeId, NodeId, FloatTensor<B>, Shape);

            fn backward(
                self,
                ops: Ops<Self::State, 3>,
                grads: &mut Gradients,
                checkpointer: &mut Checkpointer,
            ) {
                // Get the nodes of each variable.
                let [node_lhs, node_rhs, node_bias] = ops.parents;
                // Fetch the gradient for the current node.
                let grad = grads.consume::<B>(&ops.node);

                // Set our state.
                let (lhs_state, rhs_state, output, shape_bias) = ops.state;
                let lhs: FloatTensor<B> = checkpointer.retrieve_node_output(lhs_state);
                let rhs: FloatTensor<B> = checkpointer.retrieve_node_output(rhs_state);

                // Fetch shapes of our tensor to support broadcasting.
                let shape_lhs = lhs.shape();
                let shape_rhs = rhs.shape();

                // Compute the gradient of the output using the already existing `relu_backward`
                // function in the basic Burn backend trait.
                let grad_output = B::relu_backward(output, grad);

                // Compute the lhs gradient, which is the derivative of matmul with support for
                // broadcasting.
                let grad_lhs = broadcast_shape::<B>(
                    B::float_matmul(grad_output.clone(), B::float_transpose(rhs)),
                    &shape_lhs,
                );
                // Compute the rhs gradient, which is the derivative of matmul with support for
                // broadcasting.
                let grad_rhs = broadcast_shape::<B>(
                    B::float_matmul(B::float_transpose(lhs), grad_output.clone()),
                    &shape_rhs,
                );
                // The add derivative is only 1, so we just need to support broadcasting to
                // compute the bias gradient.
                let grad_bias = broadcast_shape::<B>(grad_output, &shape_bias);

                // Register the gradient for each variable based on whether they are marked as
                // `tracked`.
                if let Some(node) = node_bias {
                    grads.register::<B>(node.id, grad_bias);
                }
                if let Some(node) = node_lhs {
                    grads.register::<B>(node.id, grad_lhs);
                }
                if let Some(node) = node_rhs {
                    grads.register::<B>(node.id, grad_rhs);
                }
            }
        }

        // Prepare a stateful operation, keeping its inputs alive through registration.
        // `node()` returns a guard that `prep.finish` releases after registering the output.
        //
        // Each node can be fetched with `ops.parents` in the same order as defined here.
        match FusedMatmulAddReluBackward
            .prepare::<C>([lhs.node(), rhs.node(), bias.node()])
            // Marks the operation as compute bound, meaning it will save its
            // state instead of recomputing itself during checkpointing
            .compute_bound()
            .stateful()
        {
            OpsKind::Tracked(mut prep) => {
                // When at least one node is tracked, we should register our backward step.

                // The state consists of what will be needed for this operation's backward pass.
                // Since we need the parents' outputs, we must checkpoint their ids to retrieve
                // their node output at the beginning of the backward pass. We can also save
                // utility data such as the bias shape. If we also need this operation's output,
                // we can either save it in the state or recompute it.
                // during the backward pass. Here we choose to save it in the state because it's a
                // compute bound operation.
                let lhs_state = prep.checkpoint(&lhs);
                let rhs_state = prep.checkpoint(&rhs);
                let bias_shape = bias.primitive().shape();

                let output = B::fused_matmul_add_relu(
                    lhs.primitive().clone(),
                    rhs.primitive().clone(),
                    bias.into_primitive(),
                );

                let state = (lhs_state, rhs_state, output.clone(), bias_shape);

                prep.finish(state, output)
            }
            OpsKind::UnTracked(prep) => {
                // When no node is tracked, we can just compute the original operation without
                // keeping any state.
                let output = B::fused_matmul_add_relu(
                    lhs.into_primitive(),
                    rhs.into_primitive(),
                    bias.into_primitive(),
                );
                prep.finish(output)
            }
        }
    }
}

The backward pass masks the incoming gradient for ReLU, multiplies by transposed operands for the matmul gradients, and reduces broadcast dimensions to match the original input and bias shapes. Verify all three input gradients against the reference implementation, as well as forward values.

For application code, initialize inputs on an autodiff device and mark only source leaves whose gradients you need. Backend generics stay inside these primitive implementations; callers use the same matmul_add_relu_custom function for training and inference.