Rust in Everything
基于 candle-core 0.11.0

Invalid matmul arguments

主线位置：算子主线的后端一步。Tensor::matmul 只看形状，把两边的 Layout（形状、步长、偏移）原样交给后端；CpuStorage::matmul 和 MetalStorage::matmul 各自决定接受哪些步长。布局和视图的来龙去脉见布局与视图。

flowchart TB
  v["narrow 或 transpose"] --> lay["非连续 Layout"]
  lay --> be{"哪个后端"}
  be -->|"CPU gemm"| ok["能算"]
  be -->|"CPU accelerate"| e1["报 Striding 错"]
  be -->|Metal| e2["Invalid matmul"]
  fix["contiguous"] --> cl["连续 Layout"]
  cl --> ok2["三个后端都能算"]

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

现象

一个融合的 QKV 投影输出 (seq, 3 * d)，沿最后一维切出 Q 再乘一个矩阵：

let qkv = Tensor::randn(0f32, 1., (8, 3 * 16), &dev)?;
let w = Tensor::randn(0f32, 1., (16, 16), &dev)?;
let q = qkv.narrow(1, 0, 16)?;
show("q.matmul(&w)", q.matmul(&w));

开 accelerate 的 CPU 和 Metal：

== Cpu
q dims [8, 16] stride [48, 1] contiguous false
q.matmul(&w): Err(MatMulUnexpectedStriding { lhs_l: Layout { shape: [8, 16], stride: [48, 1], start_offset: 0 }, rhs_l: Layout { shape: [16, 16], stride: [16, 1], start_offset: 0 }, bmnk: (1, 8, 16, 16), msg: "non-contiguous lhs" })
q.contiguous()?.matmul(&w): Ok [8, 16]
w.matmul(&q.t()): Err(MatMulUnexpectedStriding { lhs_l: Layout { shape: [16, 16], stride: [16, 1], start_offset: 0 }, rhs_l: Layout { shape: [16, 8], stride: [1, 48], start_offset: 0 }, bmnk: (1, 16, 8, 16), msg: "non-contiguous rhs" })
q.t().matmul(&q): Err(MatMulUnexpectedStriding { lhs_l: Layout { shape: [16, 8], stride: [1, 48], start_offset: 0 }, rhs_l: Layout { shape: [8, 16], stride: [48, 1], start_offset: 0 }, bmnk: (1, 16, 16, 8), msg: "non-contiguous rhs" })
...
== Metal { gpu_id: 4294968354 }
q dims [8, 16] stride [48, 1] contiguous false
q.matmul(&w): Err(Metal error Invalid matmul arguments [48, 1] [16, 1] (8, 16, 16))
q.contiguous()?.matmul(&w): Ok [8, 16]
w.matmul(&q.t()): Err(Metal error Invalid matmul arguments [16, 1] [1, 48] (16, 8, 16))
q.t().matmul(&q): Err(Metal error Invalid matmul arguments [1, 48] [48, 1] (16, 16, 8))

不开 accelerate，同一段代码在 CPU 上全部成功：

== Cpu
q dims [8, 16] stride [48, 1] contiguous false
q.matmul(&w): Ok [8, 16]
q.contiguous()?.matmul(&w): Ok [8, 16]
w.matmul(&q.t()): Ok [16, 8]
q.t().matmul(&q): Ok [16, 16]

在 Linux 上（gemm）开发、到 Mac 上开 accelerate 或换 Metal，才第一次看到这个错误。

不只是 matmul。同一个程序里另外两个常用算子：

softmax_last_dim(s.t()): Err(input has to be contiguous)
softmax(s.t(), D::Minus1): Ok [6, 4]
softmax_last_dim(s.t().contiguous()): Ok [6, 4]
cache.append(k) [5 tokens]: Err(slice-set only supports contiguous tensors)
1 token transposed: contiguous true
cache.append(k1) [1 token]: Ok [1, 3, 1, 64]
cache.append(k.contiguous()): Ok [1, 3, 5, 64]

（这是 CPU 的输出；Metal 上 softmax_last_dim 报的是 Non contiguous softmax-last-dim is not implemented，append 的两行相同。）k 是常见的 (batch, seq, heads, d) 转成 (batch, heads, seq, d)：

let k = Tensor::ones((1, 5, 3, 64), DType::F32, &dev)?.transpose(1, 2)?;
let mut cache = KvCache::new(2, 16);
cache.append(&k, &k)

prefill（5 个 token）报错，decode（1 个 token）却没事：seq 为 1 时，长度为 1 的维度不参与连续性判断，转置后仍算连续。于是只测了单 token 的代码，上线后一个多 token 的 prompt 就失败。

出处

三个 matmul 实现对最后两维的要求（布局与视图一页有完整的表）：

  • gemm 那一支把行步长、列步长都传给 gemm crate，任意步长都行。
  • accelerate 那一支要映射到 BLAS 的 N / T：要么行步长等于列数、列步长为 1，要么反过来；[48, 1] 的 8 × 16 矩阵两样都不是（行步长 48 ≠ 16），报 MatMulUnexpectedStriding，msg 说明是哪一边。
  • Metal 的 call_mlx_gemm 做同样的判断：
    let (lda, a_trans) = if (lhs_m1 == 1 || k == 1) && (lhs_m2 == k || m == 1) {
        (k as i32, false)
    } else if (lhs_m1 == m || k == 1) && (lhs_m2 == 1 || m == 1) {
        (m as i32, true)
    } else {
        return Err(MetalKernelError::MatMulNonContiguous {
            lhs_stride: lhs_stride.to_vec(),
            rhs_stride: rhs_stride.to_vec(),
            mnk: (m, n, k),
        }
        .bt())?;
    };

MatMulNonContiguous 的 Display 就是 Invalid matmul arguments {lhs_stride:?} {rhs_stride:?} {mnk:?}，三项依次是左边的步长、右边的步长、(m, n, k)。

softmax_last_dim 是融合算子 SoftmaxLastDim，CPU 实现取 layout.contiguous_offsets()，拿不到就 bail!("input has to be contiguous")；Metal 实现检查 is_contiguous()。拼出来的 ops::softmax 没有这个要求（但慢，见常用层）。

KvCache::append 把新的 k、v 用 slice_set 写进预分配的缓冲区，slice_set 的第一项检查就是 if !self.is_contiguous() || !src.is_contiguous()（缓存本身见 KV 缓存与注意力）。

修法

在交给这几个算子之前 contiguous()：

let q = qkv.narrow(1, 0, 16)?.contiguous()?;
let k = k.transpose(1, 2)?.contiguous()?;

contiguous() 对已经连续的张量只是 clone（不复制数据），所以放在可能不连续的地方没有额外代价；真正不连续时它复制一次，这正是后端需要的。

几个经验：

  • 转置本身不用 contiguous()：w.t() 是标准的 T 形式，accelerate 和 Metal 都接受（上面三个后端里都能直接 a.matmul(&w.t()?)）。要处理的是 narrow 出来的列、permute 过的批维、转置后又 narrow 的组合。
  • 喂 KvCache::append 之前固定写 .contiguous()?，不要依赖「decode 时碰巧连续」。
  • 要在 Linux 和 Mac 上都能跑的代码，在 Mac 上开 accelerate 和 Metal 各测一遍；只在 gemm 上测过的步长组合不说明问题。

怎么确认没中招

  • 报错里有 non-contiguous、Invalid matmul arguments、has to be contiguous、only supports contiguous 之一：打印输入的 stride() 和 is_contiguous()，找出产生它的 narrow / transpose / permute。
  • KV 缓存的测试至少包含一次多 token 的 prefill。
  • 批维不连续是另一种情况：Metal 不报错而是算错，见 Metal 上批量 matmul 结果不对。

小结

  • matmul 对最后两维的步长：gemm 全收，accelerate 和 Metal 只收行主序和转置两种；narrow 出来的列两样都不是。
  • softmax_last_dim、KvCache::append（经 slice_set）要求连续输入；单 token 时转置后的张量碰巧连续，问题只在 prefill 暴露。
  • 修法都是 contiguous()；已经连续时它不复制。

💬 关联讨论

发起讨论
加载中...