企业官网建设流程全解析

Burn 深度学习框架自定义 CubeCL Kernel 实战：从零手写融合算子并接入后端扩展

【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesn't compromise on flexibility, efficiency and portability.项目地址: https://gitcode.com/GitHub_Trending/bu/burn

导读

本篇技术指南围绕 Burn 框架的后端扩展（Backend Extension）机制展开，讲解如何借助 CubeCL 编译器前端编写自定义 GPU Kernel，实现「矩阵乘 + 加法 + ReLU」三算子融合的自定义操作。读者将掌握：如何声明自定义后端 trait、如何用#[cube]宏编写并启动 Kernel、如何为 JIT 后端（CubeBackend）与自动微分后端（Autodiff）分别实现前向与反向传播，最终得到一个可在任意 CubeCL 运行时（CUDA / WGPU / 融合后端）上运行、且能与参考实现做数值校验的完整扩展。

文中所有代码均可在仓库 examples/custom-cubecl-kernel 目录下找到（src/lib.rs、src/kernel.rs、src/forward.rs、src/backward.rs、examples/custom-cubecl-kernel.rs），建议边读边对照源码。

注意：CubeCL 仍在活跃开发中，本节内容可能随时间演进有所出入，请以当前仓库源码为准。

为什么需要手动编写融合 Kernel

在深度学习实践中，matmul → add(bias) → relu是各类模型中极其常见的一段计算序列。Burn 本身已经具备自动算子融合能力，能够把多个连续操作合并执行，减少 Kernel 启动次数与中间张量的内存读写。但在某些场景下，手工实现的融合 Kernel 可能更加高效——例如需要精确控制寄存器占用、缓存复用策略，或对特定形状做针对性优化时。

本节选择这个例子，正是因为它既能展示 Kernel 编写的基本流程，又贴近真实模型结构，具备典型的工程价值。

第一步：定义自定义 Backend Trait

扩展后端的起点，是为新操作确定类型签名。Burn 的Backendtrait 通过关联类型TensorPrimitive封装各后端的底层张量实现；为了避免关联类型带来的冗长消歧，代码中直接使用FloatTensor<Self>别名。

use burn::backend::{backend_extension, tensor::FloatTensor}; /// 我们创建自己的 Backend trait，扩展 Burn 的 backend trait。 #[backend_extension(Autodiff, Cube)] pub trait Backend: burn::backend::Backend { fn fused_matmul_add_relu( lhs: FloatTensor<Self>, rhs: FloatTensor<Self>, bias: FloatTensor<Self>, ) -> FloatTensor<Self>; }

其中#[backend_extension(Autodiff, Cube)]属性宏会为 trait 生成配套的调度与自动微分支持代码，其展开逻辑实现在 burn-backend-extension crate 中（catalog.rs、dispatch.rs、extension.rs、routing.rs等），它负责把Tensor → Bridge → Dispatch → Backend这条调用链串联起来。感兴趣的读者可以在该 crate 中查看宏的底层展开细节。

定义好 trait 后，项目代码中就可以用这个自定义 trait 替代burn::backend::Backend使用。由于 Burn 的用户级 API 通常操作Tensor结构体而非底层 primitive 张量，我们还需要用普通函数封装新操作，保持 API 一致：

use burn::tensor::{Tensor, activation}; /// 使用自定义后端 trait 上新增的函数实现自定义操作。 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) } /// 仅使用基础张量操作实现的参考版本，用于测试校验。 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) }

提供参考实现（matmul_add_relu_reference）并非强制要求，但它能让数值验证变得非常简单——凡能用基础张量操作写出参考语义的场景，都强烈建议保留一份，这是后续所有测试的「黄金标准」。

第二步：编写前向 Kernel

接下来用cubecl编译器前端编写融合 Kernel。为保持示例简洁，这里不使用分块、向量化等高级技巧，而是写一个朴素的 matmul：row、col、batch三个维度分别映射到 GPU 的绝对位置ABSOLUTE_POS_X / Y / Z，每个线程负责输出矩阵中的一个元素。

use cubecl::{cube, prelude::*}; #[cube(launch)] pub fn fused_matmul_add_relu_kernel<F: Float>( lhs: &Tensor<F>, rhs: &Tensor<F>, bias: &Tensor<F>, output: &mut Tensor<F>, ) { let row = ABSOLUTE_POS_X; let col = ABSOLUTE_POS_Y; let batch = ABSOLUTE_POS_Z; let n_rows = output.shape(output.rank() - 2); let n_cols = output.shape(output.rank() - 1); let dim_k = rhs.shape(rhs.rank() - 1); // 越界线程直接返回，保证 grid 大小不整除时依然正确。 if row >= n_rows || col >= n_cols { return; } // 计算 batch 维的偏移，并支持任意维度的 batch 前缀。 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); } // 单线程串行完成 K 维累加。 let mut sum = F::new(0.0); 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]; } // 在同一个 Kernel 内完成 bias 加法与 ReLU 激活。 let out_index = row * n_cols + col; let index = offset_output + out_index; output[index] = F::max(sum + bias[index], F::new(0.0)); }

对照仓库中 examples/custom-cubecl-kernel/src/kernel.rs 可以看到，实际示例代码在细节上有两处演进：

• 越界处理使用terminate!()主动终止线程，而非return；
• 通过#[define(F)] _dtype: ElemType把元素类型作为编译期常量传入，配合launch参数中的dtype_to_storage_type(dtype)使用，从而支持f32/f16等不同精度。

代码中那段「prelude」看似繁琐，但它的作用至关重要：把每个计算单元正确地映射到它负责的数据上。offset_lhs/offset_rhs的推导利用了张量的stride与shape信息，使 Kernel 能够支持带 batch 前缀（如[B1, B2, M, N]甚至更高维）的输入，而不仅仅局限于三维张量。dim_k取自rhs的最后一维，n_rows/n_cols取自输出张量的倒数两维，这正对应[..., M, K] × [..., K, N] → [..., M, N]的 matmul 语义。

关于#[cube]宏的完整语法（Tensor、ABSOLUTE_POS_X、terminate!等内置项），可以继续阅读 CubeCL 官方 Book，这里不再展开。

第三步：为 JIT 后端实现前向

Kernel 写好后，还需要编写「启动代码」，实现自定义 trait。这里的关键洞察是：CubeCL 的 JIT 后端只有一个——CubeBackend，它覆盖了所有 CubeCL 运行时；张量运行在哪个设备上由device决定。因此这一份impl同时覆盖了burn-cuda、burn-wgpu以及融合后端。

use cubecl::{CubeCount, CubeDim}; use burn_cubecl::{CubeBackend, kernel::into_contiguous, tensor::CubeTensor}; /// 为 CubeBackend 实现我们的自定义后端 trait。 impl Backend for CubeBackend { fn fused_matmul_add_relu( lhs: FloatTensor<Self>, rhs: FloatTensor<Self>, bias: FloatTensor<Self>, ) -> FloatTensor<Self> { // 定义 cube 维度，这里为简单起见硬编码。 let cube_dim = CubeDim { x: 16, y: 16, z: 1 }; // 设备一致性检查。 lhs.assert_is_on_same_device(&rhs); lhs.assert_is_on_same_device(&bias); // 为简单起见，确保每个张量都是连续的。 let lhs = into_contiguous(lhs); let rhs = into_contiguous(rhs); let bias = into_contiguous(bias); // 提取 matmul 相关形状。 let ndims = lhs.shape.num_dims(); let num_rows = lhs.shape[ndims - 2]; let num_cols = rhs.shape[ndims - 1]; // 计算输出形状，同时统计 batch 数量。 let mut num_batches = 1; let mut shape_out = vec![0; ndims]; for i in shape_out.clone().into_iter().take(ndims - 2) { shape_out[i] = usize::max(lhs.shape[i], rhs.shape[i]); num_batches *= shape_out[i]; } shape_out[ndims - 2] = num_rows; shape_out[ndims - 1] = num_cols; let shape_out = Shape::from(shape_out); // 为输出张量分配缓冲区。 let buffer = lhs .client .empty(shape_out.num_elements() * core::mem::size_of::<F>()); // 构造输出张量 primitive。 let output = CubeTensor::new_contiguous( lhs.client.clone(), lhs.device.clone(), shape_out, buffer, F::dtype(), ); // 根据形状计算三个方向需要的 cube 数量。 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); // 惰性执行 Kernel，传入启动信息与各缓冲区。 fused_matmul_add_relu_kernel::launch::<F, R>( &lhs.client, cube_count, cube_dim, lhs.into_tensor_arg(), rhs.into_tensor_arg(), bias.into_tensor_arg(), output.clone().into_tensor_arg(), ); // 返回输出张量。 output } }

对照仓库 examples/custom-cubecl-kernel/src/forward.rs 可以看到，实际示例还额外为burn_fusion::Fusion<CubeBackend>提供了impl Backend（当前以todo!()占位），并在CubeDim::new_3d(16, 16, 1)、dtype_to_storage_type(dtype)等细节上做了与最新 API 对齐的调整。

这段启动代码中有几个值得注意的要点：

• into_contiguous预处理：Kernel 内部用连续内存地址 + 手工偏移计算索引，因此输入必须先转换为连续张量，否则可能读到错误数据；
• 输出形状的广播语义：batch 维度取max(lhs.shape[i], rhs.shape[i])，即沿 batch 维做广播，这与 Burn 常规 matmul 的行为一致；
• CubeCount::Static(x, y, z)：分别代表三个方向的 grid 规模，通过ceil(尺寸 / cube_dim)计算得到，保证覆盖所有输出元素；
• 缓冲区的可变性边界：CubeCL 的上下文（client）有能力对任意缓冲区执行任意可变操作，Rust 的借用规则在这里并不生效。上面的例子只修改新建的输出缓冲区，因此没有问题，但编写更复杂的 Kernel 时需要格外留意这一点。

第四步：实现反向传播

完成 JIT 后端的前向实现后，matmul_add_relu_custom已可用于推理。但如果需要计算梯度，还必须为自动微分后端实现反向。

反向传播的实现思路非常巧妙：不为 backward 另写 CubeCL Kernel，而是复用前向融合 Kernel，梯度用基础张量操作计算。做法是利用burn-autodiff的通用后端实现（它对内层后端B是泛型的），在Autodiff<B, C>上实现自定义 trait：

use burn::backend::autodiff::{ Autodiff, NodeId, checkpoint::{base::Checkpointer, strategy::CheckpointStrategy}, grads::Gradients, ops::{Backward, Ops, OpsKind, broadcast_shape}, }; // 为所有实现了我们自定义 trait 的后端实现 Autodiff 版本。 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> { // 创建零尺寸类型，用于实现 Backward trait。 #[derive(Debug)] struct FusedMatmulAddReluBackward; // 为该操作实现 Backward，需要计算三个梯度（lhs、rhs、bias）。 impl<B: Backend> Backward<B, 3> for FusedMatmulAddReluBackward { // 前向阶段保存的状态，供反向阶段使用。 type State = (NodeId, NodeId, FloatTensor<B>, Shape); fn backward( self, ops: Ops<Self::State, 3>, grads: &mut Gradients, checkpointer: &mut Checkpointer, ) { // 取出每个变量的计算图节点。 let [node_lhs, node_rhs, node_bias] = ops.parents; // 取出当前节点传下来的梯度。 let grad = grads.consume::<B>(&ops.node); // 从状态中恢复前向信息。 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); // 取出张量形状以支持广播。 let shape_lhs = lhs.shape(); let shape_rhs = rhs.shape(); // 复用 Burn 基础后端 trait 中已有的 relu_backward。 let grad_output = B::relu_backward(output, grad); // lhs 梯度 = d(matmul)/d(lhs)，支持广播。 let grad_lhs = broadcast_shape::<B>( B::float_matmul(grad_output.clone(), B::float_transpose(rhs)), &shape_lhs, ); // rhs 梯度 = d(matmul)/d(rhs)，支持广播。 let grad_rhs = broadcast_shape::<B>( B::float_matmul(B::float_transpose(lhs), grad_output.clone()), &shape_rhs, ); // add 的导数为 1，只需广播即可得到 bias 梯度。 let grad_bias = broadcast_shape::<B>(grad_output, &shape_bias); // 根据节点是否被 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); } } } // 准备一个有状态操作：传入各变量节点及对应计算图。 match FusedMatmulAddReluBackward .prepare::<C>([lhs.node.clone(), rhs.node.clone(), bias.node.clone()]) // 标记为计算密集型操作：checkpoint 时保存状态而非重算。 .compute_bound() .stateful() { OpsKind::Tracked(mut prep) => { // 至少有一个节点被 tracked 时，注册反向步骤。 // 需要父节点输出，因此 checkpoint 它们的 id；同时保存 bias 形状等 // 辅助数据。本操作输出也存入状态（因为是计算密集型）。 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.primitive, ); let state = (lhs_state, rhs_state, output.clone(), bias_shape); prep.finish(state, output) } OpsKind::UnTracked(prep) => { // 没有任何节点被 tracked 时，无需保存状态，直接做前向计算。 let output = B::fused_matmul_add_relu(lhs.primitive, rhs.primitive, bias.primitive); prep.finish(output) } } } }

这段代码的逻辑可以总结为：

1. 前向仍委托给内层后端：在Autodiff<B, C>中，真正的计算通过B::fused_matmul_add_relu转发给实现了自定义 trait 的CubeBackend；
2. 状态（State）驱动反向：Tracked分支把反向需要的全部信息——两个父节点的 checkpoint id、前向输出、bias 形状——打包成状态；UnTracked分支（图中该操作不需要反向）则直接计算并返回，不保存任何状态；
3. 梯度推导：反向阶段从计算图前序节点取出梯度grad，依次计算——ReLU 导数用现成的B::relu_backward；matmul 导数是对转置输入再做一次 matmul（grad ⊗ rhsᵀ与lhsᵀ ⊗ grad）；add 导数恒为 1，只做广播归约得到grad_bias；
4. 条件注册：ops.parents中的节点是Option，只有当变量被require_grad跟踪（tracked）时才把对应梯度注册进Gradients，避免无谓计算。

还需要注意compute_bound()的语义：它标记该操作为计算密集型，因此在 checkpoint 机制下选择保存状态而非反向时重算，这会影响显存与算力的权衡策略。此外，state中保存了前向输出output，因为反向的relu_backward需要它。

第五步：为 JIT 后端补上 Autodiff 声明

最后一块拼图，是为Autodiff<CubeBackend>实现自动微分后端 trait。由于上面的实现已经泛型覆盖了所有B: Backend，这里的实现体可以是空的：

impl AutodiffBackend for Autodiff<CubeBackend> { }

至此，Autodiff<CubeBackend>同时满足自定义Backendtrait 与 Burn 的AutodiffBackendtrait，训练链路完整打通。

验证：与参考实现做数值对比

示例仓库的 examples/custom-cubecl-kernel.rs 提供了两个验证函数：

推理验证——随机初始化形状为[1, 32, 32]、[32, 32, 32]、[32, 32, 32]的三个张量，分别用参考实现与自定义 Kernel 计算输出，然后通过assert_approx_eq::<f32>(&custom, Tolerance::default())断言近似相等：

let reference = matmul_add_relu_reference(lhs.clone(), rhs.clone(), bias.clone()) .into_data() .convert::<f32>(); let custom = matmul_add_relu_custom(lhs, rhs, bias) .into_data() .convert::<f32>(); reference.assert_approx_eq::<f32>(&custom, Tolerance::default());

反向验证——对三个输入张量调用require_grad()，先对参考实现做backward()取出三个梯度，再detach()后对自定义实现做同样的反向，逐一用assert_approx_eq对比lhs、rhs、bias三个梯度：

let lhs = Tensor::<3>::random([1, 32, 32], Distribution::Default, device).require_grad(); // ... rhs、bias 同理 let reference = matmul_add_relu_reference(lhs.clone(), rhs.clone(), bias.clone()); let mut gradients = reference.backward(); let lhs_grad_ref = lhs.grad_remove(&mut gradients).unwrap(); // ... 对比自定义实现的梯度

运行入口main中，先对Device::default()执行inference()，再对device.autodiff()执行autodiff()，即可一次性验证前向与反向的正确性：

fn main() { let device = Device::default(); inference(&device); autodiff(&device.autodiff()); }

这种「自定义实现 vs 参考实现」的双实现验证模式，是整个示例最值得借鉴的工程实践：它把 Kernel 正确性验证完全自动化，任何索引偏移、广播处理或边界条件上的疏漏都会在断言处暴露。

依赖与特性配置

在 examples/custom-cubecl-kernel/Cargo.toml 中可以看到运行本示例所需的依赖与特性：

[dependencies] burn = { workspace = true, features = [ "autodiff", "wgpu", "autotune", "template", "extension" ] } burn-cubecl = { workspace = true, features = ["default"] } cubecl = { workspace = true, features = ["wgpu"] }

其中burn的extension特性启用#[backend_extension]属性宏，template特性启用#[cube]模板编译支持，autodiff与wgpu分别提供自动微分与 WGPU 运行时；burn-cubecl提供CubeBackend、into_contiguous等 JIT 后端基础设施，cubecl则是 Kernel 编写所需的编译器前端。若需要 CUDA 支持，可将wgpu替换为对应的 CUDA 运行时特性。

扩展：不止于张量参数

本节示例中的操作只接收/返回张量，但后端扩展的能力边界不止于此。仓库中 burn-book/src/advanced/backend-extension/README.md 展示了更强的用法：自定义操作可以接收和返回包含张量的结构体或枚举，只需为类型派生ExtensionType，并用#[extension_type]标注输入参数：

#[derive(ExtensionType)] pub struct Boxes<B: Backend> { pub coords: FloatTensor<B>, pub scores: FloatTensor<B>, pub count: usize, // 非张量字段原样传递。 } #[backend_extension(Cube, Autodiff)] pub trait Backend: burn::backend::Backend { // 结构体作为输出。 fn detect(image: FloatTensor<Self>) -> Boxes<Self>; // 结构体/枚举作为输入。 fn nms(#[extension_type] boxes: Boxes<Self>, iou_threshold: f32) -> Boxes<Self>; }

这类输入可以自由地与普通张量参数混用；后端选择依据输入中的路由张量决定，若没有任何输入携带张量则无法解析后端并会 panic。结构体/枚举输入同样支持Autodiff——浮点字段携带梯度，其余字段不参与反向，impl ... for Autodiff<B, C>中反向逻辑的编写方式与普通张量输入完全一致。这意味着你可以把自定义 Kernel 封装成高层次的「算子对象」，例如 flash attention 或自定义检测头，而无需把每个中间张量都暴露给用户。

总结

在本指南中，我们借助cubecl编译器前端实现了一个融合 Kernel，使其能够在任意 GPU 与任意 CubeCL 后端上执行；同时深入剖析了 JIT 后端与自动微分后端的内部工作方式：

• 自定义后端 trait用#[backend_extension(Autodiff, Cube)]声明新原语操作，并配套Dispatch封装与参考实现；
• 前向 Kernel用#[cube(launch)]编写，通过ABSOLUTE_POS_X/Y/Z做线程映射，支持 batch 前缀与广播；
• JIT 后端实现一份impl覆盖 CUDA / WGPU / Fusion 全部运行时，包含设备检查、连续性保证、形状推导、缓冲区分配与惰性启动；
• Autodiff 实现复用前向 Kernel，反向梯度用基础操作计算，通过OpsKind::Tracked / UnTracked分支管理状态与 checkpoint 策略；
• 数值验证通过参考实现 +assert_approx_eq自动校验前向输出与三个梯度。

虽然扩展后端比直接操作张量更具挑战性，但收益同样可观：你可以用对执行过程的精细控制来打造自定义模型，从而在特定负载上显著提升性能。希望这篇指南能帮助你理解 Burn 后端扩展的全貌，并在自己的项目中释放自定义 Kernel 的潜力。

【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesn't compromise on flexibility, efficiency and portability.项目地址: https://gitcode.com/GitHub_Trending/bu/burn