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

Burn 深度学习框架自定义 CubeCL Kernel 实战从零手写融合算子并接入后端扩展【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesnt 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封装各后端的底层张量实现为了避免关联类型带来的冗长消歧代码中直接使用FloatTensorSelf别名。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: FloatTensorSelf, rhs: FloatTensorSelf, bias: FloatTensorSelf, ) - FloatTensorSelf; }其中#[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: Tensor3, rhs: Tensor3, bias: Tensor3) - Tensor3 { 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: Tensor3, rhs: Tensor3, bias: Tensor3) - Tensor3 { let x lhs.matmul(rhs) bias; activation::relu(x) }提供参考实现matmul_add_relu_reference并非强制要求但它能让数值验证变得非常简单——凡能用基础张量操作写出参考语义的场景都强烈建议保留一份这是后续所有测试的「黄金标准」。第二步编写前向 Kernel接下来用cubecl编译器前端编写融合 Kernel。为保持示例简洁这里不使用分块、向量化等高级技巧而是写一个朴素的 matmulrow、col、batch三个维度分别映射到 GPU 的绝对位置ABSOLUTE_POS_X / Y / Z每个线程负责输出矩阵中的一个元素。use cubecl::{cube, prelude::*}; #[cube(launch)] pub fn fused_matmul_add_relu_kernelF: Float( lhs: TensorF, rhs: TensorF, bias: TensorF, output: mut TensorF, ) { 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: FloatTensorSelf, rhs: FloatTensorSelf, bias: FloatTensorSelf, ) - FloatTensorSelf { // 定义 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::FusionCubeBackend提供了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是泛型的在AutodiffB, C上实现自定义 traituse burn::backend::autodiff::{ Autodiff, NodeId, checkpoint::{base::Checkpointer, strategy::CheckpointStrategy}, grads::Gradients, ops::{Backward, Ops, OpsKind, broadcast_shape}, }; // 为所有实现了我们自定义 trait 的后端实现 Autodiff 版本。 implB: Backend, C: CheckpointStrategy Backend for AutodiffB, C { fn fused_matmul_add_relu( lhs: FloatTensorSelf, rhs: FloatTensorSelf, bias: FloatTensorSelf, ) - FloatTensorSelf { // 创建零尺寸类型用于实现 Backward trait。 #[derive(Debug)] struct FusedMatmulAddReluBackward; // 为该操作实现 Backward需要计算三个梯度lhs、rhs、bias。 implB: Backend BackwardB, 3 for FusedMatmulAddReluBackward { // 前向阶段保存的状态供反向阶段使用。 type State (NodeId, NodeId, FloatTensorB, Shape); fn backward( self, ops: OpsSelf::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: FloatTensorB checkpointer.retrieve_node_output(lhs_state); let rhs: FloatTensorB 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) } } } }这段代码的逻辑可以总结为前向仍委托给内层后端在AutodiffB, C中真正的计算通过B::fused_matmul_add_relu转发给实现了自定义 trait 的CubeBackend状态State驱动反向Tracked分支把反向需要的全部信息——两个父节点的 checkpoint id、前向输出、bias 形状——打包成状态UnTracked分支图中该操作不需要反向则直接计算并返回不保存任何状态梯度推导反向阶段从计算图前序节点取出梯度grad依次计算——ReLU 导数用现成的B::relu_backwardmatmul 导数是对转置输入再做一次 matmulgrad ⊗ rhsᵀ与lhsᵀ ⊗ gradadd 导数恒为 1只做广播归约得到grad_bias条件注册ops.parents中的节点是Option只有当变量被require_grad跟踪tracked时才把对应梯度注册进Gradients避免无谓计算。还需要注意compute_bound()的语义它标记该操作为计算密集型因此在 checkpoint 机制下选择保存状态而非反向时重算这会影响显存与算力的权衡策略。此外state中保存了前向输出output因为反向的relu_backward需要它。第五步为 JIT 后端补上 Autodiff 声明最后一块拼图是为AutodiffCubeBackend实现自动微分后端 trait。由于上面的实现已经泛型覆盖了所有B: Backend这里的实现体可以是空的impl AutodiffBackend for AutodiffCubeBackend { }至此AutodiffCubeBackend同时满足自定义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 BoxesB: Backend { pub coords: FloatTensorB, pub scores: FloatTensorB, pub count: usize, // 非张量字段原样传递。 } #[backend_extension(Cube, Autodiff)] pub trait Backend: burn::backend::Backend { // 结构体作为输出。 fn detect(image: FloatTensorSelf) - BoxesSelf; // 结构体/枚举作为输入。 fn nms(#[extension_type] boxes: BoxesSelf, iou_threshold: f32) - BoxesSelf; }这类输入可以自由地与普通张量参数混用后端选择依据输入中的路由张量决定若没有任何输入携带张量则无法解析后端并会 panic。结构体/枚举输入同样支持Autodiff——浮点字段携带梯度其余字段不参与反向impl ... for AutodiffB, 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 doesnt compromise on flexibility, efficiency and portability.项目地址: https://gitcode.com/GitHub_Trending/bu/burn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考