Custom WGPU 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.

This example supplies WGSL source to the WGPU runtime. It requires a compiler accepting WGSL; a WGPU build selecting SPIR-V or MSL rejects this source. Registering Cube does not make a WGSL kernel portable to CUDA, ROCm, or the CubeCL CPU runtime.

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;

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>;
}

fn output_shape(lhs: &Shape, rhs: &Shape, bias: &Shape) -> Shape {
    let rank = lhs.num_dims();
    assert!(rank >= 2, "matmul needs at least two dimensions");
    assert_eq!(rank, rhs.num_dims(), "matrix ranks must match");
    assert_eq!(
        lhs[rank - 1],
        rhs[rank - 2],
        "contraction dimensions must match"
    );
    let mut shape = lhs.clone();
    for i in 0..rank - 2 {
        assert!(
            lhs[i] == rhs[i] || lhs[i] == 1 || rhs[i] == 1,
            "batch dimensions must broadcast"
        );
        shape[i] = if lhs[i] == 1 { rhs[i] } else { lhs[i] };
    }
    shape[rank - 1] = rhs[rank - 1];
    assert_eq!(
        &shape, bias,
        "kernel requires bias to match the output shape"
    );
    shape
}

/// 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)
}

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.

@group(0)
@binding(0)
var<storage, read_write> lhs: array<{{ elem }}>;

@group(0)
@binding(1)
var<storage, read_write> rhs: array<{{ elem }}>;

@group(0)
@binding(2)
var<storage, read_write> bias: array<{{ elem }}>;

@group(0)
@binding(3)
var<storage, read_write> output: array<{{ elem }}>;

@group(0)
@binding(4)
var<storage, read_write> info: array<u32>;

const BLOCK_SIZE = {{ workgroup_size_x }}u;

@compute
@workgroup_size({{ workgroup_size_x }}, {{ workgroup_size_y }}, 1)
fn main(
    @builtin(global_invocation_id) global_id: vec3<u32>,
    @builtin(local_invocation_index) local_idx: u32,
    @builtin(workgroup_id) workgroup_id: vec3<u32>,
) {
    // Indices
    let row = workgroup_id.x * BLOCK_SIZE + (local_idx / BLOCK_SIZE);
    let col = workgroup_id.y * BLOCK_SIZE + (local_idx % BLOCK_SIZE);
    let batch = global_id.z;

    // Basic information
    let dim = info[0];
    let n_rows = info[6u * dim - 1u];
    let n_cols = info[6u * dim];
    let K = info[5u * dim - 1u];

    // Returns if outside the output dimension
    if row >= n_rows || col >= n_cols {
        return;
    }

    // Calculate the corresponding offsets with support for broadcasting.
    let offset_output = batch * n_rows * n_cols;
    var offset_lhs: u32 = 0u;
    var offset_rhs: u32 = 0u;

    let batch_dims = dim - 2u;
    for (var b: u32 = 1u; b <= batch_dims; b++) {
        let stride_lhs = info[b];
        let stride_rhs = info[b + dim];
        let stride_output = info[b + 2u * dim];
        let shape_lhs = info[b + 3u * dim];
        let shape_rhs = info[b + 4u * dim];

        offset_lhs += offset_output / stride_output % shape_lhs * stride_lhs;
        offset_rhs += offset_output / stride_output % shape_rhs * stride_rhs;
    }

    // Basic matmul implementation
    var sum = 0.0;
    for (var k: u32 = 0u; k < K; k++) {
        let lhs_index = row * K + k;
        let rhs_index = k * n_cols + col;

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

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

    output[index] = max(sum + bias[index], 0.0);
}

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;

use super::Backend;
use burn::{
    backend::wgpu::{
        CubeBackend, CubeTensor, KernelSource, SourceKernel, SourceTemplate, build_info,
        into_contiguous, kernel_source,
    },
    tensor::DType,
};
use cubecl::{CubeCount, CubeDim, prelude::KernelId, server::KernelArguments};
use derive_new::new;

// Source the kernel written in WGSL.
kernel_source!(FusedMatmulAddReluRaw, "./kernel.wgsl");

// Define our kernel type with cube information.
#[derive(new, Debug)]
struct FusedMatmulAddRelu {
    cube_dim: CubeDim,
    dtype: DType,
}

// Implement the dynamic kernel trait for our kernel type.
impl KernelSource for FusedMatmulAddRelu {
    fn source(&self) -> SourceTemplate {
        // Extend our raw kernel with cube size information using the
        // `SourceTemplate` trait.
        FusedMatmulAddReluRaw::new()
            .source()
            .register("workgroup_size_x", self.cube_dim.x.to_string())
            .register("workgroup_size_y", self.cube_dim.y.to_string())
            .register("elem", self.dtype.name())
            .register("int", "i32")
    }

    fn id(&self) -> KernelId {
        KernelId::new::<Self>().info(self.cube_dim)
    }

    // `kernel.wgsl` is WGSL, so this runs on a wgpu build that compiles to WGSL. A build whose
    // `AutoCompiler` picks SPIR-V or MSL instead rejects the kernel rather than mis-reading it.
    fn lang(&self) -> &'static str {
        "wgsl"
    }
}

/// Implement our custom backend trait for the cubecl backend. The WGSL source below only
/// compiles on the wgpu runtime, which is the one this example's device selects.
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);

        // Get the matmul relevant shapes.
        let ndims = lhs.meta.shape().num_dims();
        let num_rows = lhs.meta.shape()[ndims - 2];
        let num_cols = rhs.meta.shape()[ndims - 1];

        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());
        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,
        );

        // Create the kernel.
        let kernel = FusedMatmulAddRelu::new(cube_dim, dtype);

        // Build info buffer with tensor information needed by the kernel, such as shapes and strides.
        let info = build_info(&[&lhs, &rhs, &output]);
        let info_handle = lhs.client.create_from_slice(bytemuck::cast_slice(&info));

        // 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.
        lhs.client.launch(
            Box::new(SourceKernel::new(kernel, cube_dim)),
            cube_count,
            KernelArguments::new().with_buffers(vec![
                lhs.handle.binding(),
                rhs.handle.binding(),
                bias.handle.binding(),
                output.handle.clone().binding(),
                info_handle.binding(),
            ]),
        );

        // 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.
//
// Note that we could implement the backend trait only for the Wgpu backend instead of any backend that
// also implements our own API. This would allow us to call any function only implemented for Wgpu
// and potentially call a custom kernel crafted only for this task.
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. 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.