Rust in Everything
基于 candle-core 0.11.0

Metal 上批量 matmul 结果不对

主线位置：算子主线的 MetalStorage::matmul → call_mlx_gemm。所有批维被压成一个批维 b，第 i 个矩阵的起点按一个步长算；这个步长怎么取、CPU 怎么检查，见布局与视图的「matmul 接受哪些步长」和 Metal 后端。

flowchart TB
  q["qkv.narrow 出 q"] --> lay["批步长 480 和 40"]
  lay --> be{"后端"}
  be -->|CPU| err["报 Striding 错"]
  be -->|Metal| one["只用倒数第三维 40"]
  one --> bad["batch 1 读错位置"]
  fix["contiguous"] --> ok["两边都对"]

实验 crate exp-candle/pa-pit（features metal accelerate）和 exp-candle/pa-noacc（只开 metal），bin p12_bmm、example linear4d、rank5，debug；M2 Max、macOS 27、rustc 1.99.0。

现象

注意力里一个常见写法：Q、K、V 沿头的那一维叠在一起，形状 (batch, 3 * heads, seq, head_dim)，用 narrow 切出 Q 和 K，算 Q Kᵀ：

let (b, h, s, d) = (2, 4, 5, 8);
let qkv = qkv_cpu.to_device(dev)?;
let q = qkv.narrow(1, 0, h)?;
let k = qkv.narrow(1, h, h)?;
let att = q.matmul(&k.t()?)?;

和先 contiguous() 再乘的结果逐个头比较：

== Cpu
q dims [2, 4, 5, 8] stride [480, 40, 8, 1]
q.matmul(&k.t()): Err(MatMulUnexpectedStriding { lhs_l: Layout { shape: [2, 4, 5, 8], stride: [480, 40, 8, 1], start_offset: 0 }, rhs_l: Layout { shape: [2, 4, 8, 5], stride: [480, 40, 1, 8], start_offset: 160 }, bmnk: (8, 5, 5, 8), msg: "non-contiguous lhs" })
q.contiguous().matmul(k.contiguous().t()): max |diff| 9.54e-7
== Metal { gpu_id: 4294968354 }
q dims [2, 4, 5, 8] stride [480, 40, 8, 1]
q.matmul(&k.t()): Ok [2, 4, 5, 5]
  batch 0, max |diff| per head: 0.00e0 0.00e0 0.00e0 0.00e0
  batch 1, max |diff| per head: 9.55e0 9.63e0 8.00e0 1.07e1
q.contiguous().matmul(k.contiguous().t()): max |diff| 0.00e0

CPU（开不开 accelerate 都一样）拒绝；Metal 返回一个形状正确的结果，第一个 batch 全对，第二个 batch 全错，没有任何报错。

candle-nn 的 Linear 也会碰到。它对非连续的四维输入走 broadcast_left + matmul（linear4d，输入是 (2, 4, 3, 5) 在第 1 维 narrow(1, 0, 2) 的结果）：

cpu Linear on narrowed 4D: 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: [0, 0, 1, 5], start_offset: 0 }, bmnk: (4, 3, 6, 5), msg: \"non-contiguous lhs\" }")
metal Linear batch (0,0) max |diff| = 2.38e-7
metal Linear batch (0,1) max |diff| = 4.77e-7
metal Linear batch (1,0) max |diff| = 4.65e0
metal Linear batch (1,1) max |diff| = 5.75e0

出处

q 的批维是 (2, 4)，步长 (480, 40)：同一个 batch 里相邻两个头差 40 个元素，相邻两个 batch 差 480，而不是 4 × 40 = 160。这两个批维合不成一个等距的批。

CPU 的 MatMul::ab_skip 把这种情况挡下来：

        let a_skip: usize = match lhs_stride[..rank - 2] {
            [s1, stride] if s1 == stride * lhs_l.dims()[1] => stride,
            [_, stride] if lhs_l.dims()[0] == 1 => stride,
            [stride, _] if lhs_l.dims()[1] == 1 => stride,
            [stride] => stride,
            [] => m * k,
            _ => Err(self.striding_error(lhs_l, rhs_l, "non-contiguous lhs"))?,
        };

Metal 的 call_mlx_gemm 不检查，直接取倒数第三维的步长：

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

8 个矩阵按步长 40 排开：前 4 个落在 0、40、80、120，正好是 batch 0 的 4 个头；后 4 个落在 160 到 280，读到的是 batch 0 里 K 的头，而 batch 1 的 Q 从 480 开始。所以 batch 0 全对、batch 1 全错。rhs 的批步长同样只取一个。这是 0.11.0 的上游 bug。

Linear 的情况相同：非连续的四维输入走 Linear::forward 的 else 分支，self.weight.broadcast_left((b1, b2))?.t()? 再 x.matmul(&w)，输入的批步长 (60, 15) 合不成一个。

📝NOTE
冷知识：唯一一处检查批步长的代码只是为了优化

call_mlx_gemm 里批步长取 stride[len - 3] 的写法，是 PR #2468（「Integrate the MLX gemm kernels」，首次进入 0.7.0）接入 MLX GEMM 时就有的。PR #3313（「Metal GEMM Dynamic Tile Selection and Batch Collapse Optimization」，首次进入 0.10.0）之后，函数里多了一个 check_batch_collapse，注释说它移植自 MLX 的「batch collapse optimization」：当 A 的批步长等于 M * K、B 的批步长为 0（二维权重被广播）时，把 [batch, M, K] @ [K, N] 折成 [batch*M, K] @ [K, N] 一次算完。这是 mlx_gemm.rs 里唯一看批步长是否「连续」的地方，但它同样只看倒数第三维，条件不满足就退回原来的路径，并不拒绝批维合不成一个步长的输入。上面 Linear 的例子正好满足折叠条件：输入倒数第三维的步长 15 等于 M * K = 3 × 5，广播出来的权重批步长是 0，于是 4 个批被当成连续的 12 行（每行 5 个元素）一次读完，后两个批读的是偏移 30 到 59，正确的位置是 60 到 89。这个 PR 的提交说明只有「mlx gemm opt init」「fix bug」「opt」这类条目，没有提到批维布局。

修法

批量 matmul 之前，凡是来自 narrow、permute、transpose 批维的输入都 contiguous()：

let att = q.contiguous()?.matmul(&k.contiguous()?.t()?)?;

上面的输出里，修过之后 Metal 和参考结果完全一致（0.00e0），CPU 也能算。k.contiguous()?.t()? 的转置不用再 contiguous()，最后两维的转置 Metal 和 CPU 都接受（见 Invalid matmul arguments）。

Linear 也在调用前 contiguous()：连续的输入走 reshape 成二维的那条路径，不经过批维。

切头时如果沿批维 narrow、transpose、permute，结果都要 contiguous() 一次再进 matmul，不管下一步看起来多简单。

怎么确认没中招

  • 批量 matmul 的两个输入，检查 is_contiguous()，或者至少检查批维的步长满足 stride[i] == stride[i + 1] * dims[i + 1]。只在 Metal 上跑的代码最需要这一步，CPU 会替你报错，Metal 不会。
  • 新写的注意力代码在 Metal 上先和 CPU（或先 contiguous() 的版本）逐 batch 比一次；只比第一个 batch 发现不了，本页的例子里 batch 0 恰好全对。 另一个方向的问题：CPU 的 ab_skip 只处理最多两个批维，秩为 5 的连续输入也被拒绝，报的同样是「non-contiguous lhs」（rank5）：
Cpu rank5 contiguous: Err("MatMulUnexpectedStriding { lhs_l: Layout { shape: [2, 2, 2, 3, 4], stride: [48, 24, 12, 4, 1], start_offset: 0 }, rhs_l: Layout { shape: [2, 2, 2, 4, 5], stride: [80, 40, 20, 5, 1], start_offset: 0 }, bmnk: (8, 3, 5, 4), msg: \"non-contiguous lhs\" }")
Metal { gpu_id: 4294968354 } rank5 contiguous: Ok([2, 2, 2, 3, 5])

遇到时把批维 reshape 成一维再乘。

小结

  • Metal 的 call_mlx_gemm 只用倒数第三维的步长当批步长，批维合不成一个步长时静默算错（0.11.0 上游 bug）；CPU 同样的输入报 MatMulUnexpectedStriding。
  • 典型来源：沿头的维度 narrow 出来的 Q / K，Linear 的非连续四维输入。
  • 修法：批量 matmul 前 contiguous()；在 Metal 上逐 batch 对照验证。
  • CPU 反过来对秩 5 的连续输入也报「non-contiguous」。

💬 关联讨论

发起讨论
加载中...