Rust in Everything
基于 candle-core 0.11.0

Metal 后端

主线位置：算子主线在 Metal 上的终点和读回。Storage::matmul 匹配到 (Metal, Metal)（见 Storage 与后端 trait），MetalStorage::matmul 分配输出缓冲区、取编码器、调 candle-metal-kernels 的 call_mlx_gemm 把一次 GEMM 编码进当前命令缓冲区，然后立刻返回，GPU 这时还没算。to_vec2 走到 MetalStorage::to_cpu 时才提交并等待。本页讲这中间的每一步。

sequenceDiagram
  participant T as Tensor
  participant S as MetalStorage
  participant K as Kernels
  participant C as Commands
  participant G as GPU
  T->>S: matmul
  S->>S: 从缓冲池取输出
  S->>C: command_encoder
  S->>K: load_pipeline_with_constants
  K-->>S: 首次要编译 .metal
  S->>C: dispatch 进共享编码器
  S-->>T: 新 Tensor 立即返回
  C->>G: 每 50 次 commit
  T->>S: to_vec2
  S->>C: blit 到共享内存
  S->>C: flush_and_wait_current
  C->>G: commit 并等待
  G-->>T: 读回 Vec

代码分两个 crate：candle-core/src/metal_backend/（mod.rs 2378 行是 BackendStorage 的实现，device.rs 486 行是 MetalDevice 和缓冲池）；candle-metal-kernels（0.83 万行 Rust 加 1.7 万行 .metal）负责命令缓冲区、编码器、内核加载和每个内核的启动参数。它直接用 objc2-metal 调 Metal，不经过 wgpu。

本页的实验 crate 是 exp-candle/backend-ops（candle-core 0.11.0，features metal accelerate；M2 Max、macOS 27、rustc 1.99.0；全部 release）。计时前都调 device.synchronize()。

MetalDevice 有什么

MetalDevice 可以 Clone，字段除了 id 和设备句柄全是 Arc，克隆出来的还是同一个设备：

#[derive(Clone)]
pub struct MetalDevice {
    pub(crate) id: DeviceId,
    pub(crate) device: Device,
    pub(crate) commands: Arc<Commands>,
    pub(crate) buffers: Arc<RwLock<BufferMap>>,
    pub(crate) private_buffers: Arc<RwLock<BufferMap>>,
    pub(crate) kernels: Arc<Kernels>,
    pub(crate) seed: Arc<Mutex<Buffer>>,
    pub(crate) seed_value: Arc<RwLock<u64>>,
    pub(crate) residency_set: Arc<ResidencySet>,
}

（引用时删去了字段上的文档注释。）

字段是什么
id进程内自增的编号，判断「同一设备」用它，不用 GPU 的 registryID
deviceMTLDevice 的包装
commands命令队列、当前命令缓冲区、当前编码器，见下一节
buffers / private_buffers两个缓冲池，按大小分桶：HashMap<usize, Vec<Arc<Buffer>>>
kernels编译好的库和管线的缓存
seed / seed_value随机数内核用的种子，放在一个 8 字节的缓冲区里
residency_set所有分配过的缓冲区登记在这里，挂到命令队列上，让它们常驻 GPU 可见的内存

BackendDevice::new 按序号取 Device::all() 里的一个设备，建一个命令队列和一个空的 Kernels。实测 Device::new_metal(0) 花 8 到 9 ms；这时一个内核都没有编译。

MetalStorage 只有四个字段：buffer: Arc<Buffer>、device: MetalDevice、count、dtype。缓冲区不带类型，所以要另存 dtype。

一个 op 怎么编码

各个 op 在 metal_backend/mod.rs 里的写法一样：从缓冲池取输出缓冲区，device.command_encoder() 取编码器，按 dtype 和布局拼出内核名，调 candle_metal_kernels 里对应的 call_* 函数。以 matmul 为例：

        let buffer = self
            .device
            .new_buffer_builder()
            .with_size_for(b * m * n, self.dtype)
            .with_label("matmul")
            .build()?;
        let encoder = self.device.command_encoder()?;
        let dtype = match self.dtype {
            DType::F32 => candle_metal_kernels::GemmDType::F32,
            DType::F16 => candle_metal_kernels::GemmDType::F16,
            DType::BF16 => candle_metal_kernels::GemmDType::BF16,
            dtype => {
                return Err(
                    MetalError::Message(format!("mlx matmul doesn't support {dtype:?}")).into(),
                )
            }
        };
        candle_metal_kernels::call_mlx_gemm(
            &self.device.device,
            &encoder,
            &self.device.kernels,
            dtype,
            (b, m, n, k),
            // ...
        )
        .map_err(MetalError::from)?;

call_mlx_gemm 是从 MLX 移植的 GEMM：m == 1 或 n == 1 时改走 gemv；否则按矩阵大小、dtype、转置方式和设备类型选一组分块参数，内核名形如 gemm_nn_f32_f32_32_32_16_2_2，再用函数常量（function constants）传「m、n、k 是否对齐到分块」这类开关。最后 set_params! 绑定缓冲区、dispatch_thread_groups 发出去，函数返回。没有任何等待。

批维的处理是这个函数的一个缺口。CPU 的 ab_skip 要求所有批维能合成一个步长，合不成就报 MatMulUnexpectedStriding（见 CPU 后端）；call_mlx_gemm 不做这个检查，批步长只取倒数第三维的步长：

        let a_stride = if lhs_stride.len() > 2 {
            lhs_stride[lhs_stride.len() - 3] as isize
        } else {
            (m * k) as isize
        };

b 个批次一律按这一个步长往后跳。对一个在第二维上 narrow 过的四维张量，这个假设不成立，结果就错了，而且不报错。batch_mm 程序（最小复现）：

let x = Tensor::randn(0f32, 1., (2, 4, 3, 5), &cpu)?;
let y = Tensor::randn(0f32, 1., (2, 2, 5, 6), &cpu)?;
let xn = x.narrow(1, 0, 2)?; // 形状 (2, 2, 3, 5)，步长 (60, 15, 5, 1)
let reference = xn.contiguous()?.matmul(&y)?;
let xm = x.to_device(&metal)?.narrow(1, 0, 2)?;
let got = xm.matmul(&y.to_device(&metal)?)?.to_device(&cpu)?;

xn dims [2, 2, 3, 5] stride [60, 15, 5, 1]
cpu strided: Err(MatMulUnexpectedStriding { lhs_l: Layout { shape: [2, 2, 3, 5], stride: [60, 15, 5, 1], start_offset: 0 }, rhs_l: Layout { shape: [2, 2, 5, 6], stride: [60, 30, 6, 1], start_offset: 0 }, bmnk: (4, 3, 6, 5), msg: "non-contiguous lhs" })
metal batch (0,0) max |diff| = 0.00000008940697
metal batch (0,1) max |diff| = 0.00000023841858
metal batch (1,0) max |diff| = 5.5542965
metal batch (1,1) max |diff| = 7.631793
metal after contiguous(): max |diff| = 0.00000023841858

四个批次按步长 15 排，第三个批次从偏移 30 读，正确的位置是 60，后两个批次读到的是 x 里被 narrow 掉的那部分。这是 0.11.0 的上游 bug；先 contiguous() 就对了。布局与视图一页从布局的角度讲了同一个问题。

encoder 是什么？Commands::command_encoder：

    pub fn command_encoder(&self) -> Result<CommandsGuard<'_>, MetalKernelError> {
        let mut state_guard = self.state.lock().unwrap();
        let count = self.compute_count.fetch_add(1, Ordering::Relaxed);
        let flush = count >= self.compute_per_buffer;

        if flush {
            self.commit_swap_locked(&mut state_guard, 1)?;
        }

        if state_guard.current_encoder.is_none() {
            let fence = Arc::new(Fence::new(&self.device));
            let enc = state_guard.current.compute_command_encoder(&fence);
            // Wait for all prior encoder fences before the first dispatch.
            // ...
            state_guard.current_encoder = Some(enc);
        }

        Ok(CommandsGuard { guard: state_guard })
    }
  • 返回的 CommandsGuard 持有 Commands 的互斥锁，整个 op 编码期间别的线程进不来；guard 在 op 函数结束时释放。
  • 计算编码器是共享的：当前命令缓冲区上只有一个 current_encoder，第一次用到时创建，之后每个 op 都往同一个编码器里追加 dispatch，直到提交或者有人要 blit。
  • 这个编码器是并发的：compute_command_encoder 用 MTLDispatchType::Concurrent 创建，GPU 可以让相邻的 dispatch 重叠执行。缓冲区又都用 HazardTrackingModeUntracked 分配（RESOURCE_OPTIONS），Metal 不替你跟踪读写依赖。

依赖由 candle 自己跟踪。内核绑定缓冲区时分两种：set_input_buffer 和 set_output_buffer，按缓冲区指针记进上一批 dispatch 的输入、输出集合。这次要读的缓冲区上一批写过（写后读，RAW），或者这次要写的缓冲区上一批读过、写过（WAR、WAW），就标记 needs_barrier；下一次 dispatch 前 auto_barrier 插一道 memoryBarrierWithScope(Buffers)。没有依赖的 dispatch 之间没有屏障，可以并行。跨编码器、跨命令缓冲区的依赖用 MTLFence：每个编码器结束时 updateFence，把自己写过的缓冲区登记到 prev_ce_outputs，下一个编码器开始时对表里的每个 fence waitForFence。

这套「共享并发编码器 + 自动屏障」是 PR #3511（「[Metal] Concurrent dispatching」，首次进入 0.11.0）引入的。之前每个 op 各开一个编码器，编码器之间天然串行。也因此，candle-metal-kernels 的 set_params! 里写输出要包一层 Output::new(buffer)：裸的 &Buffer 会被当成输入登记，自定义 Metal 算子同样要遵守（见自定义算子）。

什么时候提交

上面的 count >= self.compute_per_buffer 决定提交时机。Commands::new 读环境变量：

const DEFAULT_CANDLE_METAL_COMPUTE_PER_BUFFER: usize = 50;
// ...
        let compute_per_buffer = match std::env::var("CANDLE_METAL_COMPUTE_PER_BUFFER") {
            Ok(val) => val
                .parse()
                .unwrap_or(DEFAULT_CANDLE_METAL_COMPUTE_PER_BUFFER),
            _ => DEFAULT_CANDLE_METAL_COMPUTE_PER_BUFFER,
        };

字段的文档注释说它是「每个命令缓冲区最多几个计算编码器」，但编码器已经是共享的了，compute_count 实际数的是 command_encoder() 和 blit_command_encoder() 被调用的次数，大致等于 op 的个数。数到 50，commit_swap_locked 结束当前编码器、commit 当前命令缓冲区、放进 in_flight、换一个新的。提交之后 GPU 才开始执行这一批；不等它执行完，CPU 继续编码下一批。环境变量解析失败（比如 abc）退回 50；设成 0 等于每次取编码器都提交。

metal_queue 程序做 500 次 x = x.add(&y)（x、y 是 1048576 个 f32），分三段计时：编码循环本身、synchronize、to_vec1；另一组是 1024 个元素的张量连加 5000 次。先用一次小链预热，编译好要用的内核：

let t0 = Instant::now();
for _ in 0..steps {
  x = x.add(&y)?;
}
let encode = ms(t0);
let t1 = Instant::now();
if sync_first {
  dev.synchronize()?;
}
let sync = ms(t1);
let t2 = Instant::now();
let v = x.to_vec1::<f32>()?;
let read = ms(t2);

CANDLE_METAL_COMPUTE_PER_BUFFER=None
len=1048576  steps=500   encode     0.78 ms  synchronize     0.00 ms  to_vec1     8.01 ms  total     8.79 ms  x[0]=500
len=1048576  steps=500   encode     0.82 ms  synchronize     7.44 ms  to_vec1     0.71 ms  total     8.97 ms  x[0]=500
len=1024     steps=5000  encode     6.90 ms  synchronize     0.00 ms  to_vec1     3.42 ms  total    10.31 ms  x[0]=5000
CANDLE_METAL_COMPUTE_PER_BUFFER=Some("1")
len=1048576  steps=500   encode     8.88 ms  synchronize     0.00 ms  to_vec1     1.58 ms  total    10.46 ms  x[0]=500
len=1048576  steps=500   encode     8.64 ms  synchronize     1.12 ms  to_vec1     0.71 ms  total    10.47 ms  x[0]=500
len=1024     steps=5000  encode    84.83 ms  synchronize     0.00 ms  to_vec1     1.68 ms  total    86.51 ms  x[0]=5000

（每种设置跑两轮，这里是第二轮；第一轮数字相近。）

设置500 次大加法：编码 / 总计5000 次小加法：总计
默认（50）0.8 ms / 8.8 到 9.0 ms10.0 到 10.3 ms
18.6 到 8.9 ms / 10.5 ms85.5 到 86.5 ms
08.6 ms / 10.0 到 10.8 ms82.3 到 88.3 ms
101.2 ms / 8.7 到 8.9 ms13.4 到 13.7 ms
10000.6 ms / 9.1 到 9.4 ms11.9 到 12.3 ms
abc0.8 到 0.9 ms / 8.8 到 9.1 ms10.0 到 11.1 ms
  • 默认设置下，编码 500 个 op 只要 0.8 ms，平均每个 1.6 µs。GPU 的活（约 8 ms，4 MB 读两份写一份，500 次）全部落在后面那一次等待里：不先 synchronize 就算在 to_vec1 头上，先 synchronize 就算在它头上，to_vec1 本身（blit 4 MB 再拷进 Vec）只剩 0.7 ms。
  • 每个 op 提交一次（1 或 0）时，提交本身的开销在每个 op 十几微秒，小张量上总时间变成 8 倍。大张量上 GPU 时间占大头，差别只有 1 到 2 ms。
  • 50 和 1000 差不多，abc 就是 50。这台机器上默认值已经在平台区，往大调没什么收益。

读回：等的是谁

MetalStorage::to_cpu 是所有 to_vec*、to_scalar、to_device(&Cpu) 的终点：

    pub(crate) fn to_cpu<T: Clone>(&self) -> Result<Vec<T>> {
        let size = self.count * self.dtype.size_in_bytes();
        let buffer = self
            .device
            .new_buffer_builder()
            .with_size(size)
            .with_label("blit_to_cpu_dst")
            .build()?;
        {
            let mut blit = self.device.blit_command_encoder()?;
            blit.set_label("blit_to_cpu");
            blit.copy_from_buffer(&self.buffer, 0, &buffer, 0, size);
        }
        self.device.flush_and_wait_current()?;
        Ok(read_to_vec(&buffer, self.count))
    }

中间结果默认放在 StorageModePrivate 的缓冲区里（只有 GPU 能读，见下一节），所以读回要先用 blit 拷进一个 StorageModeShared 的缓冲区，再从 contents() 指针拷进 Vec。拷的是整块 count 个元素，和张量的 Layout 无关，按布局挑元素是之后在 CPU 上做的。

flush_and_wait_current 在持锁时提交当前命令缓冲区并记下它，放锁后等这一个缓冲区完成。命令队列按提交顺序执行，它完成了，之前提交的也都完成了。

📝NOTE
冷知识：读回曾经可能读到还没写的内存

PR #3595（「Metal: fix CPU readback race under concurrent command submission」，首次进入 0.11.0）之前，to_cpu 编码完 blit 调的是 wait_until_completed。那个函数提交当前缓冲区，然后从 in_flight 列表里取走最后一个去等。提交说明讲了问题：别的线程同时在提交、等待时，「the in-flight list can be taken by another thread's flush_and_wait between the blit encode and the wait, causing the reader to return before its blit has executed and to read stale or unwritten destination memory」。修法就是上面的 flush_and_wait_current：「commits the current command buffer while holding the state lock and waits on that specific buffer」，并把已完成的缓冲区从 in_flight 里清掉，免得读回频繁时列表无限变长。同一个 PR 也改了量化张量的两处读回；说明里写，新加的并发读回测试在旧代码上「fail on every thread within a few iterations」。单线程程序碰不到这个竞态。

device.synchronize() 是另一条路：MetalDevice::wait_until_completed → Commands::flush_and_wait，提交当前缓冲区，取走整个 in_flight 列表，只等最后一个（注释说逐个等每个要多付 1 到 2 ms 的系统通知延迟）。两条路最后都调 drop_unused_buffers。

缓冲区：按 2 的幂分桶复用

with_size_for(..) 最终调 MetalDevice::new_buffer：

    pub fn new_buffer(
        &self,
        element_count: usize,
        dtype: DType,
        _name: &str,
    ) -> Result<Arc<Buffer>> {
        let size = element_count * dtype.size_in_bytes();
        let mut buffers = self.private_buffers.write().map_err(MetalError::from)?;
        if let Some(b) = find_available_buffer(size, &buffers) {
            return Ok(b.clone());
        }
        let size = buf_size(size);
        let subbuffers = buffers.entry(size).or_insert(vec![]);

        let new_buffer = self
            .device
            .new_buffer(size, PRIVATE_RESOURCE_OPTIONS)
            .map_err(MetalError::from)?;
        let new_buffer = Arc::new(new_buffer);
        self.residency_set.insert(&new_buffer);
        subbuffers.push(new_buffer.clone());
        Ok(new_buffer)
    }
  • 池里每个缓冲区都由池自己持有一个 Arc。find_available_buffer 在不小于需求的桶里找 Arc::strong_count == 1 的缓冲区，也就是除了池没人再用的那个，取最小的桶。
  • 找不到就新分配，大小向上取到 2 的幂（buf_size 是 size.next_power_of_two()）。一个 1025 字节的张量占 2048 字节。
  • 中间结果用 private_buffers（StorageModePrivate，只有 GPU 访问，省掉 CPU 一致性开销）；需要 CPU 读写的（读回的目标、with_zeros、从 CPU 数据创建的张量）用 buffers（StorageModeShared）。
  • 池不会自己缩小。drop_unused_buffers 把 strong_count == 1 的缓冲区从池里删掉（真正释放），它只在 synchronize 和读回时被调用。MetalDevice 字段上的注释说「Whenever we actually allocate a new buffer, we make a full sweep to clean up unused buffers」，0.11.0 的代码里分配路径并没有做这件事。

池里的 Arc 计数还解释了另一个 0.11.0 的上游 bug。MetalStorage::try_clone 是 Ok(self.clone())，只多了一个 Arc<Buffer> 引用，没有新缓冲区；而 Tensor::copy 正是靠 try_clone「复制」数据。于是在 Metal 上 copy() 出来的张量和原张量共用显存，原地写其中一个（slice_set、zero_set、自定义的 InplaceOp），另一个跟着变。backend-ops 的 storage 程序（最小复现）：

let orig = Tensor::zeros(4, DType::F32, dev)?;
let copied = orig.copy()?;
let src = Tensor::ones(2, DType::F32, dev)?;
orig.slice_set(&src, 0, 0)?;

Cpu: orig=[1.0, 1.0, 0.0, 0.0] copy=[0.0, 0.0, 0.0, 0.0]
Metal { gpu_id: 4294968354 }: orig=[1.0, 1.0, 0.0, 0.0] copy=[1.0, 1.0, 0.0, 0.0]

需要真正独立的一份时，可以用一个总会分配新输出的运算绕过去，例如 t.affine(1.0, 0.0)?（注意 to_dtype 遇到同一 dtype 时直接返回 clone，不行）。Tensor 的结构一页有更多例子。

一个张量被 drop，它的缓冲区立刻就能被下一个 op 拿去当输出，而 GPU 上可能还有排着队、要读这个缓冲区的 dispatch。之所以不出错，靠的是上面的依赖跟踪：新 op 把这个缓冲区登记为输出，而它在上一批的输入集合里，于是插入屏障（WAR）；跨编码器时有 fence。注释里那句「in order for this to work, we need to guarantee the order of operation」说的就是这个前提。

📝NOTE
冷知识：2 字节的缓冲区曾经被分成 1 字节

PR #3394（「fix(metal): buffer pool rounds 2 down by 1」，首次进入 0.10.0）之前，buf_size 是 size.saturating_sub(1).next_power_of_two()。对多数大小这和 next_power_of_two() 一样，但 2 - 1 = 1，1.next_power_of_two() == 1：一个 bf16 或 f16 的标量要 2 字节，池给了 1 字节的缓冲区。这个 PR 改成现在的写法，并加了三个单元测试，其中 test_buf_size_bf16_f16_scalar 的注释写着「A scalar tensor requests a 2-byte buffer. This must not be rounded down to 1.」。

内核：运行时编译，两级缓存

.metal 源码不在构建时编译。source.rs 用 include_str! 把 16 个文件作为字符串嵌进二进制：

pub const AFFINE: &str = include_str!("metal_src/affine.metal");
pub const BINARY: &str = include_str!("metal_src/binary.metal");
// ...
pub const MLX_GEMM: &str = include_str!("metal_src/mlx_gemm.metal");
// ...
pub const QUANTIZED: &str = include_str!("metal_src/quantized.metal");
// ...
pub const SDPA: &str = include_str!("metal_src/scaled_dot_product_attention.metal");

目录里的第 17 个文件 utils.metal 不在其中，也没有任何文件 #include 它：newLibraryWithSource 拿到的只是一个字符串，没有包含路径，所以 get_strided_index 这样的辅助函数在 affine、binary、cast、indexing、ternary、unary 六个文件里各抄了一份。CUDA 后端正好相反，candle-kernels 的 build.rs 在构建时调 nvcc 生成 PTX；本卷没有实测 CUDA。

Kernels 有两张表：libraries: HashMap<Source, Library> 每个源文件一项，pipelines: HashMap<(KernelName, Option<ConstantValues>), ComputePipeline> 每个（内核名，函数常量）一项。load_library 第一次碰到某个 Source 时把整个文件交给 newLibraryWithSource 编译：

    pub fn load_library(
        &self,
        device: &Device,
        source: Source,
    ) -> Result<Library, MetalKernelError> {
        let mut libraries = self.libraries.write()?;
        if let Some(lib) = libraries.get(&source) {
            Ok(lib.clone())
        } else {
            let lib = {
                let source_content = self.get_library_source(source);
                let compile_options = get_compile_options();
                device
                    .new_library_with_source(source_content, Some(&compile_options))
                    .map_err(|e| MetalKernelError::LoadLibraryError(e.to_string()))?
            };
            libraries.insert(source, lib.clone());
            Ok(lib)
        }
    }

编译时拿着写锁，同一设备上别的线程要加载任何库都得等。get_compile_options 默认开快速数学（CANDLE_METAL_ENABLE_FAST_MATH，默认 true）。load_pipeline_with_constants 在库之上再按（名字，常量）建管线：同一个库里换一个内核，或者 GEMM 换一组对齐常量，还要再建一次管线，但不用重新编译整个文件。

Kernels 在 MetalDevice::new 里新建，每个设备一份。

实测：首次与之后

metal_first 按顺序对每个 op 计「第一次」（含编译）和之后 5 次的中位数，张量是 512×512 的 f32（matmul 500 用 500×500）。下面是一次缓存未命中时的运行：

report("add f32 [Binary lib]", timed(&dev, || x.add(&y))?);
report("mul f32 [same lib]", timed(&dev, || x.mul(&y))?);
report("exp f32 [Unary lib]", timed(&dev, || x.exp())?);
report("matmul 512 [Gemm lib]", timed(&dev, || x.matmul(&y))?);
report("matmul 500 [new consts]", timed(&dev, || z.matmul(&z))?);
// ...

new_metal: 8.40 ms
add f32 [Binary lib]         first    579.72 ms   later(median of 5)   0.551 ms
mul f32 [same lib]           first      0.54 ms   later(median of 5)   0.220 ms
exp f32 [Unary lib]          first    123.05 ms   later(median of 5)   0.203 ms
matmul 512 [Gemm lib]        first   4011.89 ms   later(median of 5)   0.702 ms
matmul 500 [new consts]      first     97.68 ms   later(median of 5)   0.712 ms
sum_keepdim [Reduce lib]     first   4490.36 ms   later(median of 5)   0.549 ms
softmax_last_dim [Reduce]    first      0.38 ms   later(median of 5)   0.207 ms
to_dtype f16 [Cast lib]      first     54.92 ms   later(median of 5)   0.210 ms
QMatMul Q8_0 [Quantized lib] first    819.26 ms   later(median of 5)   0.532 ms
dev2: add f32                first      0.59 ms   later(median of 5)   0.552 ms
dev2: matmul 512             first      1.04 ms   later(median of 5)   0.389 ms
  • 首次代价按库付：binary.metal 0.58 s，mlx_gemm.metal 4.0 s，reduce.metal 4.5 s，quantized.metal 0.8 s。和文件行数不成比例（quantized.metal 7741 行反而比 1483 行的 mlx_gemm.metal 快），起作用的是模板实例化出多少个内核。
  • 同一个库里的下一个内核（mul 之于 add、softmax 之于 sum）不到 1 ms；GEMM 换形状换了函数常量，要 98 ms 建新管线。
  • 第二个设备 dev2 有自己空的 Kernels，第一次 add、matmul 却只要 1 ms。再让 dev2 先跑一个 dev1 也没用过的排序内核：dev2 上 71 ms，之后 dev1 第一次跑它只要 1 ms。也就是说在 candle 自己的 Kernels 之下，进程内还有一层由 Metal 框架维护的缓存。

同一个程序再运行一次：

add f32 [Binary lib]         first      3.14 ms   later(median of 5)   0.561 ms
matmul 512 [Gemm lib]        first      3.00 ms   later(median of 5)   0.353 ms
sum_keepdim [Reduce lib]     first      1.92 ms   later(median of 5)   0.209 ms
QMatMul Q8_0 [Quantized lib] first      2.42 ms   later(median of 5)   0.540 ms

跨进程还有一层：系统把编译结果存在用户缓存目录下（$(getconf DARWIN_USER_CACHE_DIR)com.apple.metal/，非 app 进程共用一个目录，实测里面的库数据有 214 MB）。命中时整个库 2 到 3 ms。把 CANDLE_METAL_ENABLE_FAST_MATH=false 换一组编译选项，第一次又回到 0.57 s / 3.8 s / 4.4 s，第二次又是 2 到 3 ms，可见缓存的键里包含编译选项。

这层缓存不总能命中。实验刚开始时同一个二进制连续运行两次，两次都是约 4 s；把可执行文件复制到新路径后第一次运行也是冷的，而硬链接、改名后的副本、touch 过的原文件都是热的。没能找出稳定的失效条件；这几天机器上有多个写手同时在编译、运行别的 Metal 程序，系统缓存可能被挤掉，也可能是写入延迟。实际影响是：一个 candle Metal 程序第一次用到 GEMM 或归约时，有可能卡 4 秒以上，也可能只花几毫秒，取决于系统缓存。base 的 SmolLM2 生成程序在 Metal 上记录过 9.1 s（CPU 6.0 s）；在本页的副本 backend-gen 上重跑，缓存热时 Metal 87 到 91 ms，换成 CANDLE_METAL_ENABLE_FAST_MATH=false（部分库此前没以这组选项编译过）是 2.49 s，同一次运行 CPU 都在 0.11 到 0.18 s。所以计时必须把第一次单独列出来，在 Metal 上做基准时也要先跑一遍预热。

和 burn 对照：burn 的 metal feature 走 wgpu，内核由 cubecl 在运行期从 IR 生成 MSL 源码、经 wgpu 交给 Metal 建管线，同样有首次编译的代价（见 wgpu 与 Metal）；candle 手写 .metal、直接调 Metal，没有中间语言，也没有自动调优，GEMM 的分块参数是按尺寸查表选的。

❗IMPORTANT
动脑筋：把 CANDLE_METAL_COMPUTE_PER_BUFFER 设成 1，1024 个元素的张量连加 5000 次，总时间大约是默认设置的几倍？

用上面的 metal_queue，看 len=1024 steps=5000 那一行的 total。提示：GPU 上每次加法几乎不花时间。

动脑筋答案

CANDLE_METAL_COMPUTE_PER_BUFFER=None
len=1024     steps=5000  encode     7.69 ms  synchronize     0.00 ms  to_vec1     2.55 ms  total    10.24 ms  x[0]=5000
len=1024     steps=5000  encode     6.90 ms  synchronize     0.00 ms  to_vec1     3.42 ms  total    10.31 ms  x[0]=5000
CANDLE_METAL_COMPUTE_PER_BUFFER=Some("1")
len=1024     steps=5000  encode    83.98 ms  synchronize     0.00 ms  to_vec1     1.52 ms  total    85.50 ms  x[0]=5000
len=1024     steps=5000  encode    84.83 ms  synchronize     0.00 ms  to_vec1     1.68 ms  total    86.51 ms  x[0]=5000
  • 约 8.4 倍：10.3 ms 对 85.5 到 86.5 ms。
  • 默认设置下 5000 次只提交 100 个命令缓冲区，每次 op 的 CPU 开销约 1.4 µs（取缓冲区、取锁、绑定参数、dispatch）。设成 1 之后每个 op 都 commit 一次并新建命令缓冲区，每个 op 多出约 15 µs，全落在编码阶段。
  • to_vec1 反而变短：命令缓冲区早就陆续提交了，最后等的只剩一小段。小张量、op 多的场景（逐 token 的解码）对这个值最敏感。

小结

  • MetalDevice 可克隆，内部是共享的 Commands、两个缓冲池、Kernels 缓存；「同一设备」看进程内的 DeviceId。
  • 每个 op 取一个缓冲区、往共享的并发计算编码器里追加一次 dispatch，立即返回；读写依赖由 candle 按缓冲区指针跟踪，自动插入屏障和 fence（0.11.0 的 #3511）。
  • 每 50 次取编码器提交一次命令缓冲区（CANDLE_METAL_COMPUTE_PER_BUFFER）；设成 1 时小 op 的开销变成 8 倍。
  • 读回先 blit 到共享内存，再 flush_and_wait_current 等这一个命令缓冲区；GPU 的活常常算在第一次读回头上。#3595 修了多线程下读回过早返回的竞态。
  • 缓冲区按 2 的幂分桶，strong_count == 1 即可复用，只在等待时释放。
  • 两个 0.11.0 的上游 bug：try_clone 只复制 Arc，Tensor::copy 在 Metal 上不复制；call_mlx_gemm 只用倒数第三维的步长当批步长，批维合不成一个步长时静默算错，CPU 上同样的输入会报错。
  • .metal 源码嵌在二进制里，第一次用到某个库时整文件编译：GEMM 4.0 s、归约 4.5 s；进程内和跨进程（系统缓存）都有缓存，命中时只要几毫秒，但跨进程缓存不保证命中。

💬 关联讨论

发起讨论
加载中...