expand 之后 matmul，后几批算错
主线位置：张量主线的「计算」一步。一组共享的查询向量扩到整个 batch，再和每个样本的键做批量矩阵乘。
版本：PyTorch 2.15.0（CPU）、candle 0.11.0（CPU 开和不开 accelerate；Metal）
flowchart TB q["queries 1 m k"] --> ex["扩成 b 批"] ex --> mm["matmul keys_t"] mm -->|"gemm CPU"| bad["30 个错，15 个是 0"] mm -->|"accelerate Metal"| ok["正确"] q -->|"broadcast_matmul"| ok ex -->|"contiguous 再乘"| ok
本页的对照：PyTorch 侧是实验代码里的 p10.py（torch.set_num_threads(1)），candle 侧是实验 crate pita-t 的 p10_expand_mm（features metal accelerate）和 pita-noacc 的 p10_expand_mm_noacc（只开 metal，CPU 走 gemm crate），都是 release。平台 M2 Max、macOS 27、rustc 1.99.0。Linux 和 Windows 上 candle 默认就是 gemm 后端。
PyTorch 的写法
注意力池化、Perceiver 的潜变量、Q-Former 的查询，都是「一组可学的查询，对每个样本的键算分数」：
b, m, k, n = 4, 2, 3, 5
queries = torch.arange(m * k, dtype=torch.float32).reshape(1, m, k)
keys_t = torch.arange(b * k * n, dtype=torch.float32).reshape(b, k, n) * 0.01
q = queries.expand(b, -1, -1)
s = q @ keys_t
q stride (0, 3, 1)
expand @ vs contiguous: max|diff| = 0.0
batch 3 row 0: [1.600000023841858, 1.6299999952316284, 1.659999966621399, 1.6899999380111694, 1.71999990940094]
expand 不复制，批维的步长是 0；@ 对这种视图和对连续张量结果相同。
搬到 candle 的现象
expand 写成 broadcast_as，形状检查能过（秩相同，批维相同）：
let queries = Tensor::arange(0f32, (m * k) as f32, &dev)?.reshape((1, m, k))?;
let keys_t = (Tensor::arange(0f32, (b * k * n) as f32, &dev)?.reshape((b, k, n))? * 0.01)?;
let q = queries.broadcast_as((b, m, k))?; // queries[None].expand(b, -1, -1)
let scores = q.matmul(&keys_t)?;
let reference = q.contiguous()?.matmul(&keys_t)?;
不开 accelerate 的 CPU：
== Cpu (gemm)
q stride [0, 3, 1]
expand+matmul vs contiguous: max|diff| = 6.5799994
elements 40: exact zeros 15, wrong 30
batch 0: [0.24999999, 0.28, 0.31, 0.33999997, 0.37]
batch 1: [1.3, 1.42, 1.54, 1.66, 1.78]
batch 2: [2.5, 2.62, 2.7399998, 2.86, 2.98]
batch 3: [0.0, 0.0, 0.0, 0.0, 0.0]
开 accelerate 的 CPU 和 Metal：
== Cpu (accelerate)
q stride [0, 3, 1]
expand+matmul vs contiguous: max|diff| = 0
elements 40: exact zeros 0, wrong 0
batch 0: [0.24999999, 0.28, 0.31, 0.33999997, 0.37]
batch 1: [0.7, 0.72999996, 0.76, 0.78999996, 0.82]
batch 2: [1.15, 1.18, 1.21, 1.24, 1.27]
batch 3: [1.6, 1.63, 1.66, 1.6899999, 1.7199999]
（Metal 的输出除第一行外与此逐字相同。）gemm 后端上 40 个输出只有第 0 批的 10 个对，后面的批要么是别的数，要么是 0，没有任何报错。在 macOS 上开着 accelerate 开发、到 Linux 上部署，正是会撞上这一条的组合。
原因
gemm 那一支的 MatMul 在进入逐批循环之前，尝试把批维合并掉：
let (b, m, n, k) = if b_skip == 0 && a_skip == m * k {
// a_skip and c_skip should be updated but step is always 0 so
// it wouldn't matter.
(1, b * m, n, k)
} else if a_skip == 0 && b_skip == n * k {
(1, m, b * n, k)
} else {
(b, m, n, k)
};
第一支（右边批步长为 0）把批并进行，是对的。第二支（左边批步长为 0，expand 出来的正是这种）把批并进列，按一个 m × (b·n) 的矩阵算一次；但右边是 [b, k, n] 按批排的，各批的列在内存里不相邻，输出的行步长也还是 n，算出来的数互相覆盖，没写到的位置保持初始的 0。修复它的提交说明给出了没写到的个数：只写了 (m - 1) * n + b * n 个元素。这里 m = 2, n = 5, b = 4，写了 25 个，剩下 15 个是 0，和实测一致。
accelerate 和 Metal 的 matmul 实现里没有这条合并，所以不受影响。这条路径在形状与广播的「expand 之后直接 matmul」里已经拆过一遍。Metal 上另有一个批维步长的问题，现象不同，见源码卷的 Metal 上批量 matmul 结果不对。
origin/main 上已经修了：#3960「Fix silently incorrect CPU matmul for a stride-zero broadcast batch」删掉了第二支，让左边批步长为 0 的情况走下面的逐批循环，并加了测试。截至 2026-10-29 还没有包含它的发布版，0.11.0 仍是最新版。
两条合并来自 #2884「Optimize the batched matmul for the cpu backend.」（首次进入 0.9.0-alpha.2，第一个正式版是 0.9.0）。那次提交只在 cpu_backend/mod.rs 里加了 9 行，就是上面这段 if，两支对称地写出来，提交说明只有标题。一年半以后，#3960 的提交说明解释了为什么不能对称：「That fold is not valid: rhs is batch-major, so its batches are not adjacent columns, and dst_rs stays n while each merged row is b * n wide.」修复后的代码在原来第二支的位置留了一段注释，说明这里「deliberately」没有镜像的合并。
改法
在 0.11.0 上，左边要广播批维时，让它真正有那些批：
let fixed = queries.broadcast_matmul(&keys_t)?;
let fixed2 = q.contiguous()?.matmul(&keys_t)?;
fix broadcast_matmul vs contiguous: max|diff| = 0
fix q.contiguous().matmul vs contiguous: max|diff| = 0
三种后端上都是这两行。broadcast_matmul 内部先 broadcast_as 再 contiguous()，复制出 b 份，所以安全；代价是这一次复制。和 PyTorch 的结果相同（批 3 第 0 行都是 [1.6, 1.63, 1.66, 1.69, 1.72]）。
怎么检查
- 搜代码里
broadcast_as/expand之后紧跟matmul的地方，左边是被扩展的那一个就改成broadcast_matmul。 - CI 在 Linux 上跑一次不开 accelerate 的数值测试。只在 macOS 开 accelerate 跑，这个 bug 永远不会出现。
- 对批量运算的结果逐批和参考比较，不要只看第 0 批：本页第 0 批是对的。
小结
queries.broadcast_as((b, m, k))?.matmul(&keys)在 candle 0.11.0 的 gemm CPU 后端上静默算错，40 个输出错 30 个；accelerate 和 Metal 正确，PyTorch 正确。- 原因是
MatMul把「左边批步长为 0」并进列的那条合并；main 上 #3960 已删掉，尚未发布。 - 改法：
broadcast_matmul，或先contiguous()再matmul。