Burn 运行时后端分发（Dispatch）：一个程序驱动多后端 Tensor 的架构核心

• 人工智能
• 深度学习
• 机器学习
• 本地部署

导读：本文围绕burn-dispatch这一 crate 展开，它是 Burn 深度学习框架中“运行时后端选择”的枢纽——每个burn::tensor::Tensor背后的全局执行后端。读完本文你将掌握：Dispatch 如何把同一份模型代码路由到 CPU / CUDA / wgpu / Flex / Remote 等多个后端并让它们并存；Device构造函数与BURN_DEVICE环境变量如何决定运行时用哪个后端；autodiff、梯度检查点与 kernel fusion 如何作为“装饰器”叠加；以及#[backend_extension]如何把自定义算子接入 dispatch 层。全文以 crates/burn-dispatch/README.md 为骨架，并结合仓库源码给出实现级证据。

一、Dispatch 是什么：Tensor 的“全局后端”

Burn 的定位是“不牺牲灵活性、效率与可移植性的下一代张量库与深度学习框架”。要实现“同一份模型代码，多种后端可跑”，就需要一个能管理所有底层后端的统一执行层。burn-dispatch正是这个执行层。

官方文档对该 crate 的定位是一句话点明：

Dispatchis the backend behind everyburn::tensor::Tensor.

也就是说，Burn 用户代码中写到的Tensor，其底层的“后端类型”既不是Cpu也不是Cuda，而是统一的Dispatch。它在编译期持有每一个被编译进应用的 backend 的张量 primitive（张量句柄），并在运行时把每个算子路由到拥有该张量的那个后端上执行。

由此带来两个直接结果（也是 Burn 区别于传统泛型后端框架的关键设计）：

• 模型代码没有后端类型参数：Tensor::<2>、Tensor::<3>这样的类型不再携带B: Backend泛型，写模型时不需要Tensor<B, 2>这种参数化签名；
• 一个程序可以同时使用多个后端：只要把不同的Device传给不同的 Tensor，Dispatch 就能让它们并行共存，甚至在它们之间搬运数据。

这一设计在 crates/burn/src/lib.rs 的 crate 文档中同样被强调——“Every enabled backend is available at runtime through aDeviceconstructor, and several can be used side by side”。

二、一条 Tensor 操作的路由链路

2.1 核心调用链：Tensor → BridgeTensor → DispatchTensor → backend primitive

原文档给出了 Dispatch 的数据流骨架：

Tensor -> BridgeTensor -> DispatchTensor -> backend primitive

逐层解读：

1. Tensor：用户可见的高层 API（位于 crates/burn-tensor），携带形状、dtype 与Device；
2. BridgeTensor：桥接层，把高层 API 的调用转换为后端无关的操作描述；
3. DispatchTensor：分发层张量，定义在 crates/burn-dispatch/src/tensor.rs。它由两部分组成——kind: DispatchTensorKind（真正指向哪个后端张量）与autodiff: DispatchAutodiffContext（携带该张量的自动微分 / 梯度检查点上下文）；
4. backend primitive：最终落到burn-cubecl、burn-flex、burn-remote等具体后端的张量原语上执行。

关键数据结构是DispatchTensorKind——一个“每个已启用后端一个 variant”的枚举：

pub enum DispatchTensorKind { Cube(BackendTensor<Cube>), // 所有 CubeCL 运行时共享这一个 variant Flex(BackendTensor<Flex>), NdArray(BackendTensor<NdArray>), LibTorch(BackendTensor<LibTorch>), Remote(BackendTensor<Remote>), Capture(BackendTensor<Capture>), Autodiff(Box<DispatchTensorKind>), }

其中BackendTensor<B>是对该后端 float / int / bool / quantized / autodiff 五种张量句柄的封装（见 crates/burn-dispatch/src/tensor.rs）。注意：所有 CubeCL 运行时（CUDA、ROCm、Metal、Vulkan、WebGPU、wgpu、CPU）共用同一个Cubevariant，运行时具体跑在哪个设备上由张量携带的Device决定——这正是后面“一个后端、多种运行时”设计的落点。

2.2 算子如何被“转发”：backend_dispatch过程宏

在 crates/burn-dispatch/src/ops/tensor.rs 中，FloatTensorOps、IntTensorOps、BoolTensorOps、QTensorOps、ModuleOps等 trait 的Dispatch实现，几乎全部由#[backend_dispatch]属性宏生成，例如：

#[backend_dispatch] impl FloatTensorOps<Self> for Dispatch { fn float_add(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> { B::float_add(lhs, rhs) } // ... 其余算子同理 }

这里的B是宏生成的“路由结果”：它根据输入张量的运行时后端 tag，把B绑定到具体后端类型（如Cube、Flex），然后把B::float_add(..)转发到该后端。这一生成流水线定义在 crates/burn-backend-extension/src/lib.rs：

#[backend_dispatch] ─┐ ├─> ir::Operation ─> routing ─> generated enum dispatch #[backend_extension] ─┘

• ir：描述张量输入、输出与后端调用，与具体前端（built-in dispatch 或用户 extension）无关；
• routing：负责共享的后端选择、输入提取、调用与输出包装；
• catalog：唯一的“运行时后端清单”。

需要特殊路由的方法（如float_to_device的跨后端搬运）用#[backend_dispatch(skip)]标记，改由手写的float_to_device!宏生成完整的分发矩阵（见 crates/burn-dispatch/src/ops/tensor.rs 与 crates/burn-dispatch/src/macros.rs）。

2.3 后端清单：单一权威来源

Dispatch 的路由宏（backend_list!、distributed_backend_list!、backend_matrix!）并不自行维护后端列表，而是委托给burn-backend-extension的backend_catalog!过程宏（见 crates/burn-dispatch/src/macros.rs）。权威清单定义在 crates/burn-backend-extension/src/catalog.rs：

这份 catalog 同时派生出“分布式后端子集”与“转移矩阵”，供所有生成路径与手写路径共用，避免后端列表在多处维护导致漂移。

三、支持的后端变体与 feature 组合

3.1 后端总表

原文档列出的后端变体如下（Feature 列为在burn/burn-dispatch上启用的 Cargo feature）：

两处要点（README 与 crates/burn-dispatch/src/lib.rs 的模块文档一致确认）：

• 所有 cubecl 后端共享同一个Cube后端：feature 只决定“把哪些运行时编译进去”，具体用哪个运行时由张量携带的 device 决定。可以同时启用多个运行时 feature，它们共享唯一的DispatchDevice::Cubevariant。
• autodiff与fusion是装饰器：autodiff可以把上面任意一个后端包装成 burn-autodiff；fusion则为 CubeCL 与 remote 后端开启 kernel fusion。feature 可以自由组合（如flex+autodiff+fusion）。

3.2 从源码看 feature 门控

在 crates/burn-dispatch/src/backend.rs 中，Dispatch的每个 Backend 方法都按DispatchDevice的 variant 展开 match 分支，例如：

fn name(device: &Self::Device) -> String { let inner = dispatch_device!(device, |device| B::name(device)); format!("dispatch<{inner}>") }

dispatch_device!宏会为每个cfg满足条件的后端生成一个分支，把B绑定为具体后端类型后执行同一段 body（见 crates/burn-dispatch/src/macros.rs）。若某个 variant 的 feature 未启用，对应 match 分支直接通过#[cfg]消失——这就是“编译期裁剪后端、运行时选择后端”的实现机制。

四、Device：运行时的“后端选择器”

4.1 Device 与 DispatchDevice 的关系

用户面向的Device（定义在 crates/burn-tensor/src/device.rs）实际上是对DispatchDevice的类型擦除封装：

pub struct Device { blob: device_opaque::Opaque, // 类型擦除后的 DispatchDevice }

而DispatchDevice（定义在 crates/burn-dispatch/src/device.rs）是一个“每个已启用后端一个 variant”的枚举：

pub enum DispatchDevice { Cube(CubeDevice), // 所有 CubeCL 运行时 Flex(FlexDevice), NdArray(NdArrayDevice), LibTorch(LibTorchDevice), Remote(RemoteDevice), Capture(CaptureDevice), Autodiff(AutodiffDevice), }

Device提供了一组按 feature 门控的工厂方法（见 crates/burn-tensor/src/device.rs），每个方法把具体后端设备转换为DispatchDevice：

底层转换关系在 crates/burn-dispatch/src/device.rs 中通过一组From<XxxDevice> for DispatchDevice实现：例如From<CudaDevice> for DispatchDevice将其包装为DispatchDevice::Cube(CubeDevice::Cuda(..))——即 CUDA 设备最终也是落到Cubevariant 上，印证了“一个 Cube 后端覆盖所有运行时”。

4.2Device::default()的选择顺序与BURN_DEVICE

DispatchDevice::default()的实现（crates/burn-dispatch/src/device.rs）遵循以下行为：

1. 在std构建下，若设置了环境变量BURN_DEVICE，则按其值选择后端，可接受值包括cuda、rocm、metal、vulkan、webgpu、wgpu、cpu、tch、remote、flex、ndarray；
2. 未设置时，按固定优先级选择第一个已启用的后端：CUDA → Metal → ROCm → Vulkan → WebGPU → wgpu → CPU → LibTorch → Flex → Remote → NdArray；
3. 若BURN_DEVICE指定了未知名称或未启用的 feature，会直接panic!并给出配置提示；若完全没有启用任何执行后端，也会panic!，提示启用flex、wgpu或cuda等后端 feature（仅有capture时可通过Device::capture()记录图而不执行）。

注意一个细节：默认顺序之所以“每个 feature 单独写死一个return”，是为了避免CubeDevice::default()被 cargo feature 统一（unification）干扰——例如工作区同时构建了burn-cuda时，即使本 crate 只启用了wgpu，CubeDevice::default()也可能默认到 CUDA。逐个 feature 显式返回保证了“调用者未选择时的默认值”不随依赖图漂移（源码注释对此有明确说明）。

4.3 设备枚举与运行时过滤

Dispatch::enumerate(type_id)（crates/burn-dispatch/src/backend.rs）可以列出某类设备。对Cube而言，枚举结果会经过cube_runtime_enabled过滤：cargo 的 feature 统一可能让 cubecl 编译进比本 crate 请求更多的运行时（比如 workspace 里有人启用了burn-cuda），因此枚举必须按本 crate 自己的 feature过滤，只交出当前构建真正启用的运行时设备（见 crates/burn-dispatch/src/backend.rs）。Dispatch::enumerate_cube(runtime)则进一步只列出指定 runtime 的设备。

五、autodiff 与梯度检查点：运行时上下文而非类型装饰

Burn 0.22 的一个关键变化是：autodiff 不再通过泛型Autodiff<B>出现在模型代码类型中，而是作为DispatchDevice的运行时属性与DispatchTensor携带的上下文存在。

5.1 Device 层：.autodiff()与.gradient_checkpointing()

Device::autodiff()把设备包装为DispatchDevice::Autodiff(AutodiffDevice)（crates/burn-tensor/src/device.rs），AutodiffDevice内含两个字段（crates/burn-dispatch/src/device.rs）：

pub struct AutodiffDevice { inner: Box<DispatchDevice>, checkpointing: GradientCheckpointingStrategy, }

GradientCheckpointingStrategy只有两个取值（crates/burn-dispatch/src/device.rs）：

用户链式调用即可开启：Device::default().autodiff().gradient_checkpointing()。如果对未开启 autodiff 的设备调用gradient_checkpointing()会 panic（见 crates/burn-tensor/src/device.rs）。Device::without_autodiff()则反向移除 autodiff 关联，且操作幂等。

5.2 Tensor 层：DispatchAutodiffContext的合并规则

DispatchTensor携带autodiff: DispatchAutodiffContext（crates/burn-dispatch/src/tensor.rs），取值Disabled或Enabled(strategy)。一次算子调用中，多个输入张量的上下文按以下规则合并（源码中的merge方法）：

• Disabled与Enabled相遇：disabled 张量被当作常数（constant）参与运算，操作与输出使用 enabled 上下文；
• 两个Enabled相遇：必须使用相同的梯度检查点策略，否则 panic（提示 “Gradient checkpointing strategy mismatch”）。

这一规则有专门的单元测试覆盖，见 crates/burn-dispatch/src/backend.rs 底部的autodiff_context_tests模块：例如disabled_and_enabled_float_contexts_merge_in_both_orders验证两种顺序合并结果一致；fixed_arity_inputs_reject_mismatched_checkpointing_strategies验证策略不一致时 panic；q_matmul_rejects_mismatched_checkpointing_strategies验证量化算子同样遵守该规则。

5.3 运行时策略到编译期泛型的映射

with_autodiff_backend!宏（crates/burn-dispatch/src/macros.rs）把运行时的GradientCheckpointingStrategy映射为编译期Autodiff<B, C>的具体泛型参数：Balanced→BalancedCheckpointing，Disabled→NoCheckpointing。这样 dispatch 层既能在运行时选择策略，又能复用 burn-autodiff 的编译期优化。

5.4 AutodiffBackend 的实现

Dispatch实现了AutodiffBackend（crates/burn-dispatch/src/backend.rs），backward、grad、grad_remove、grad_replace、inner、from_inner等均按DispatchTensorKind的 variant 分发到对应后端的 autodiff 实现。其中：

• Capture张量不支持 autodiff，调用即 panic（“Capture tensors do not support autodiff”）；
• inner/from_inner是“剥离 / 恢复 autodiff 上下文”的转换，测试模块inner_transitions_clear_and_restore_context验证了上下文在转换中被正确清除与恢复；
• grad返回的梯度张量上下文为Disabled（“gradients are inner backend tensors”，见gradients_are_inner_backend_tensors测试）——即梯度本身是可脱离 autodiff 图的张量。

六、跨后端张量转移（to_device）

Dispatch 的另一个核心职责是“管理跨后端张量转移”。在 crates/burn-dispatch/src/macros.rs 中，to_device!与float_to_device!通过backend_matrix!生成所有后端两两组合的转移矩阵，行为要点：

• 同后端转移走快路径：直接把张量转移到该后端的新设备；
• 跨后端转移（如Flex → Cube）通过宿主内存中介：float_transfer把源张量同步读回TensorData，再用Dst::float_from_data在目标设备重建（见 crates/burn-dispatch/src/ops/transfer.rs）；
• 带 autodiff 的跨后端转移通过HostTransfer适配器同时记录前向与反向（DifferentiableTransfer），使梯度能沿转移路径反向传播（见 crates/burn-dispatch/src/ops/transfer.rs）；
• Capture 是单向的：普通后端张量可以被“物化”到 capture 设备（成为捕获图中的 initializer），但已捕获的张量没有真实数据，不能搬回其他后端——转移矩阵中 capture 被刻意排除在交叉转移之外，违规则 panic（见 crates/burn-dispatch/src/macros.rs 及多处注释）。

转移相关的测试在 crates/burn-dispatch/src/ops/transfer.rs 底部，例如recorded_transfer_uses_adapter_in_both_directions_without_replay验证Autodiff<Flex> → Autodiff<NdArray>转移时前向/反向适配器各被调用一次且梯度正确回到源设备。

七、后端扩展：#[backend_extension]自定义算子

原文档指出：后端扩展通过 burn-backend-extension 的#[backend_extension]宏为 dispatch 增加算子。仓库给出了完整示例：

• 示例项目：examples/custom-cubecl-kernel：演示如何编写自定义 CubeCL kernel，并通过扩展 trait 接入 Dispatch。示例入口 examples/custom-cubecl-kernel/examples/custom-cubecl-kernel.rs 展示了使用方式——Device::default()取设备，matmul_add_relu_custom调用自定义融合算子，并与参考实现对比输出和梯度；
• 扩展 trait 的实现：examples/custom-cubecl-kernel/src/lib.rs 中use burn::backend::{Dispatch, backend_extension, tensor::FloatTensor}，说明自定义算子的签名直接写在Dispatch上。

从实现原理看（crates/burn-backend-extension/src/lib.rs）：

• #[backend_extension]为扩展 trait 生成Dispatch实现；
• 后端选择来自一个 routing tensor（优先 float）；所有张量输入的 autodiff 上下文被合并，disabled 输入视为常数，enabled 输入必须共享同一梯度检查点策略；
• struct / enum 输入通过ExtensionTypetrait（定义在 crates/burn-core/src/backend.rs）跨 dispatch 边界映射；
• 执行后端选择器（selector）包括Cube、Flex、NdArray、LibTorch、Remote；注意Wgpu、Cuda等不是选择器——它们只是 CubeCL 运行时，统一归入Cube；
• 自定义算子的 autodiff 支持需要为Autodiff<B, C>手写扩展 trait 的实现，宏不自动生成反向传播。

八、如何在自己的项目中使用 Dispatch（实用清单）

虽然应用层不直接依赖burn-dispatch（它是通过burn的 backend features 间接启用的），但理解其工作机制有助于正确地配置你的项目：

1. 在Cargo.toml中启用一个或多个后端 feature，例如：

   [dependencies] burn = { version = "0.22", features = ["wgpu"] } # 或 ["flex"]、["cuda"]、["cpu"] 等

   Burn 默认不携带任何执行后端，必须显式选择（见 crates/burn/src/lib.rs 的 Feature Flags 说明）。

2. 创建Device并在运行时选择后端：

   use burn::prelude::*; let device = Device::default(); // 按优先级自动选择 let cuda = Device::cuda(0); // 显式指定（需 cuda feature） let flex = Device::flex(); // 纯 Rust CPU（需 flex feature） let ad_device = device.autodiff(); // 开启自动微分

3. 利用BURN_DEVICE环境变量在运行时切换默认后端，便于同一份二进制在不同机器上跑不同硬件。

4. 让不同张量落在不同后端上并存：Dispatch 允许Tensor::zeros([128, 128], &cuda)与Tensor::zeros([128, 128], &flex)共存于一个程序，跨后端运算会自动触发to_device转移。

5. 需要自定义算子时参考 examples/custom-cubecl-kernel 与 crates/burn-backend-extension/README.md，使用#[backend_extension]接入 Dispatch 路由。

九、测试与验证：Dispatch 正确性的保障

仓库为 Dispatch 提供了多层测试验证：

• 上下文合并语义测试：crates/burn-dispatch/src/backend.rs 的autodiff_context_tests（需autodiff+flexfeature），覆盖策略传播、混用合并、不匹配 panic、int/bool/float 转换保持关联、量化算子上下文合并等；
• 设备 ID 往返测试：crates/burn-dispatch/src/device.rs，验证DispatchDevice::from_id(device.to_id())对 Remote 与 Capture 设备无损往返；
• 跨后端梯度传递测试：crates/burn-dispatch/src/ops/transfer.rs，验证Autodiff<Flex> → Autodiff<NdArray>前向/反向转移、分布式参数在不同后端上的拒绝与回环；
• 高层 API 集成测试：crates/burn/tests 下的backend_extension_remote.rs、backend_extension_runtime.rs等验证扩展算子在不同运行时上的行为；
• 公开 API 测试：crates/burn/tests/public_api_linalg.rs 等确保 dispatch 层的公开 API 形状稳定。

十、小结

burn-dispatch是 Burn 多后端架构的“中枢神经系统”：它用DispatchDevice在运行时承载后端选择，用DispatchTensor/DispatchTensorKind在运行时携带张量的后端归属与 autodiff 上下文，用backend_dispatch过程宏把每个算子路由到正确后端，用转移矩阵支撑跨后端数据搬运，并用#[backend_extension]让第三方算子无缝接入。理解这一层，就理解了为什么 Burn 的模型代码可以不带任何后端泛型参数，却能在 CPU、GPU、Web 与远程设备之间自由切换与并存。相关进一步阅读：后端扩展机制见 crates/burn-backend-extension/README.md，后端实现见 crates/burn-cubecl/README.md 与 crates/burn-flex/README.md，完整 API 见 crates/burn/src/lib.rs。

• 人工智能
• 深度学习
• 机器学习
• 本地部署