1. 先把“矩阵运算”这件事的边界划清楚

我第一次被人问“PyTorch张量怎么做矩阵运算”，对方的真实困惑其实不是不会写@，而是不知道自己写的那个@到底调到了哪一层，为什么有时候形状对得上、有时候对不上，为什么同样的数学式子换成另一种写法速度快了一倍。这个问题在项目里被放大得很厉害：小规模 demo 里for循环累加完全跑得动，一旦数据量上来，同样的逻辑换成批量矩阵运算，耗时能从几秒掉到几十毫秒，差的不是机器性能，是你有没有用对 API。

所以这篇内容不讲“张量是什么”这类从零起步的概念，而是聚焦在PyTorch 张量上的矩阵运算这一个点上，把它拆成三块：乘法族的 API 到底各管什么、非乘法的线性代数操作（求逆、解方程、分解、范数）怎么用、以及形状/精度/性能这些真正会在项目里咬人的细节。适合已经能写出简单模型、但在写自定义损失、实现注意力、做批量相似度计算时被形状报错和性能问题卡住的人；也适合从 NumPy 转过来、发现两边行为不完全一致的同学。以下所有结论都是我在实际项目和反复 benchmark 中得到的，凡是涉及参数选择的，我都会把推演过程写出来，方便你自己判断。

1.1 标量、向量、矩阵、张量在代码里到底差在哪

数学上矩阵是二维的，张量是高维的，但在 PyTorch 里，它们全都是torch.Tensor，区别只体现在shape和dim()上。一个shape=(3,4)的张量是矩阵，shape=(3,)的是向量，shape=(2,3,4)的是三维张量。这个“统一表示”带来一个副作用：你没法靠类型系统区分“这是个矩阵”还是“这是个批量矩阵”，只能靠 API 的约定去约束。

关键就在于batch 维度的位置约定。PyTorch 里几乎所有批量矩阵 API 都默认批量维在前，矩阵的两个维度在后。比如torch.bmm(A, B)要求A的形状是(b, n, m)、B是(b, m, p)，输出(b, n, p)。这个约定不是随便定的，它和内存布局有关：连续张量的最后一个维度步长最小，批量维在最前面意味着每个 batch 的矩阵是连续存放的一整块，GPU 上做批量 GEMM 时可以一次性调度，效率最高。

很多人踩的坑是把 batch 维放到中间，比如(n, b, m)，然后硬套bmm，结果报错或者拿到错误结果。这时候要么permute一下变成(b, n, m)，要么换matmul让它走广播。permute回来的是非连续张量，后续.view()会失败，得用.reshape()或者先.contiguous()——这是我在处理变长序列时被教育过好几次的地方。

还有一个容易被忽略的点：PyTorch 里没有“行向量”和“列向量”的区分。shape=(3,)的张量在参与matmul时，会根据另一侧的形状自动决定是当行向量还是列向量用。这个规则在torch.matmul文档里写得很清楚，但实际读代码时经常看不出来，尤其是当两边都是一维的时候——torch.matmul(a, b)对两个一维张量做的是点积，返回标量；而torch.outer(a, b)返回的是外积矩阵。同样的两个输入，函数不同结果维度完全不同，这是新手最容易混淆的地方之一。

1.2 为什么 PyTorch 要搞出那么多乘法函数

torch.mm、torch.mv、torch.dot、torch.bmm、torch.matmul、torch.einsum、torch.outer，光乘法就有七八个入口。乍看是历史包袱，实际上每一个都对应一种明确的形状契约，用对了能少写很多判断，用错了报错信息也更直接。

torch.mm只接受二维张量，不接受广播，不做任何形状兜底。这个“严格”恰恰是它的价值：在代码 review 时看到mm，你就能确定两个操作数都是矩阵，不会出现“某个分支下传进来的是一维张量”的隐藏风险。torch.bmm同理，严格三维，不广播。把这种严格 API 用在核心计算路径上，相当于给自己加了一道类型检查。

torch.matmul是“万能版”，它内部会做形状规约和广播，能处理一维、二维、多维的任意组合。方便是方便，代价是错误被推迟了：形状不匹配的时候，报错信息往往指向matmul内部，提示“size mismatch”，但真正的 bug 可能在十行之前就已经埋下。我自己的习惯是，核心路径优先用严格 API，只有在确实需要广播或者维度不固定的地方才用matmul，这样出错时定位成本低得多。

torch.einsum是另一类东西，它不针对某一种形状，而是用爱因斯坦求和约定描述“哪些维度参与缩并、哪些保留”。一个einsum表达式可以替代多重循环、permute + bmm、diagonal、trace等一大堆操作。它的缺点是字母标记容易写错、初学者读起来费劲，而且早期版本的einsum在某些形状下性能不如手写bmm。不过近几年优化已经好很多，尤其是配合optimize='optimal'参数时，它会自动搜索最优的缩并顺序。

2. 乘法族 API 的能力边界与形状推演

这一节把乘法相关的函数放在一起对照，重点不是背函数签名，而是建立一个判断流程：看到两个张量要相乘，先问自己三个问题——各是几维、需不需要广播、批量维在哪。这三个问题回答完，用哪个函数基本就定了。

2.1 mm、mv、dot、bmm、matmul 的适用条件对照

把常用乘法 API 整理成一张表，方便随时查。这里列的是“能不能用”，不是“推荐用”，实际选择还要结合可读性和性能。

torch.dot有一个限制经常被忽略：它只接受一维张量，且要求dtype一致。如果你拿一个shape=(3,)和一个shape=(3,1)去做dot，会直接报错，因为后者是二维。这类错误在 batch size 为 1 的时候特别隐蔽，因为(3,1)在肉眼上和(3,)看起来差不多，但shape完全不同。我在调试损失函数时遇到过一次，原因就是某次unsqueeze多了一个维度，导致内积变成了矩阵乘，输出形状从标量变成(3,1)，广播到后面全部算错，但因为有广播，程序不报错，只是 loss 曲线莫名其妙地不对。

torch.mv的第二个参数必须是一维。如果你手上是(n, 1)的列向量，得先squeeze(-1)。反过来matmul就没这个限制，(n, m) @ (m, 1)得到(n, 1)，(n, m) @ (m,)得到(n,)，两种都能跑但结果形状不同。这个差异在下游要做broadcast的时候会有实际影响，比如后面接一个(n,)的 bias 相加，前者(n,1)+(n,)会广播成(n,n)，后者(n,)+(n,)是正确的。这类“静默广播错误”是矩阵运算里最难查的一类 bug，因为程序不报错、结果也不崩，就是数值不对。

torch.bmm的批量维必须严格相等，不做广播。如果两个张量的批量维一个是 1、一个是 32，bmm会报错，这时候要么expand成 32，要么换matmul。expand出来的是共享内存的视图，不占额外显存，但它要求被扩展的维度大小是 1，且后续如果是 in-place 操作会出问题。我一般直接用matmul，让它自己处理广播，省得纠结。

torch.matmul的一维规约规则值得单独记一下：如果第一个参数是一维，会在它前面补一个 1 变成二维；如果第二个参数是一维，会在它后面补一个 1；算完之后再把补的维度去掉。所以(3,) @ (3,)其实是(1,3) @ (3,1) = (1,1)，去掉补的维变成标量。(3,) @ (3,4)是(1,3) @ (3,4) = (1,4)，去掉前面的 1 得到(4,)。这条规则理解了，很多“为什么结果维度少了”的困惑就解决了。

2.2 广播机制：matmul 真正强的地方

广播这个词听起来抽象，其实就是“维度大小对不上时，如果有一边是 1，就自动复制成另一边的大小”。规则是：从最后一个维度开始往前对齐，每一维要么相等、要么其中一个是 1、要么其中一方没有这一维，否则报错。

举个具体例子，A形状(32, 1, 8, 16)，B形状(16, 32)。先对齐：B是二维，A是四维，B前面补两个 1 变成(1, 1, 16, 32)。逐维比较：最后一维16vs32，不相等且都不是 1，报错。如果B是(16, 16)，补完是(1, 1, 16, 16)，最后一维16 vs 16通过，倒数第二维8 vs 16又不通过。所以要让这个乘法成立，B的后两个维度得是(8, x)或者(1, x)配合A的8。

实际编码中，广播最常见的用法是“单头 vs 多头”。比如你有一个(heads, head_dim, head_dim)的旋转矩阵，要作用到(batch, heads, seq, head_dim)的张量上。直接matmul，PyTorch 会把旋转矩阵前面补一个 1 变成(1, heads, head_dim, head_dim)，然后和(batch, heads, seq, head_dim)做批量乘。这里1和batch广播、heads对齐、seq和head_dim对齐，最后一个维度做矩阵乘，输出(batch, heads, seq, head_dim)。整个式子写出来就是一行x @ R.transpose(-1, -2)，比手写循环清爽太多。

需要注意的是，广播出来的张量是虚拟扩展，不真的分配内存，但一旦参与计算，硬件层面可能还是会按扩展后的大小去读。expand之后连续做多次运算，性能不一定比repeat好，因为访存模式是重复读取同一块。我实测过(1, 64, 128, 128)扩展成(32, 64, 128, 128)再matmul的情况，和直接repeat再算，在小尺寸下差别不明显，在seq_len上千的时候，expand版本反而略慢，原因是缓存局部性差。所以“expand 省显存”是对的，但别默认它一定更快。

还有一类坑是“意外的标量广播”。(n, 1) + (n,)会得到(n, n)，这个前面提过。更隐蔽的是(n, 1) * (1, m)得到(n, m)，如果本意是逐元素对齐，就会算出一个完全错误的矩阵。防御办法很简单：在关键位置assert一下形状，尤其是损失函数、mask 处理这类地方。一行assert x.shape == y.shape能省掉半夜调参的时间。

2.3 用 einsum 统一表达所有乘法形态

einsum的核心是下标字符串：输入张量各用一个字母序列标注维度，箭头右边写输出保留的维度，没出现在右边的字母就是被缩并掉的。'ij,jk->ik'是矩阵乘，'i,i->'是内积，'i,j->ij'是外积，'ii->'是迹，'ij->ji'是转置。

它最大的好处是把“形状变换”和“计算”合成一句，代码的数学含义直接可读。比如实现多头注意力的打分：Q形状(b, h, s, d)，K形状(b, h, t, d)，要算(b, h, s, t)的分数矩阵。用matmul的写法是Q @ K.transpose(-1, -2)，用einsum是torch.einsum('bhsd,bhtd->bhst', Q, K)。后者一眼就能看出d被缩并掉了，输出的四个维度分别来自哪里，不需要在脑子里做转置推演。

再看一个matmul不太好表达的：批量加权求和。feat形状(b, n, d)，weight形状(b, n)，要对每个样本按权重把n个向量加权平均成(b, d)。matmul写法是(weight.unsqueeze(1) @ feat).squeeze(1)，中间要插维度、再去掉。einsum写法是torch.einsum('bn,bnd->bd', weight, feat)，直接表达“沿 n 缩并”。可读性和出错概率都低不少。

性能上，einsum内部会把它翻译成bmm、permute、sum的组合。默认的缩并顺序是按表达式从左到右，不一定是全局最优。加optimize='optimal'会让它枚举所有缩并顺序，选计算量最小的那个——代价是第一次调用有额外的规划开销，适合形状固定的场景，比如固定的模型结构。如果是动态形状，每次调用形状都变，规划开销会抵消收益，这时候用默认或者optimize='greedy'更稳妥。

有一个einsum的坑要说清楚：它不做广播。所有出现在表达式里的维度，输入张量必须有，大小必须严格匹配。'bhsd,bhtd->bhst'里如果某个输入的h是 1、另一个是 12，einsum会报错，而matmul会自动广播。所以在写可变头数的代码时，要自己处理单头扩展的情况，不能指望einsum兜底。

3. 非乘法类矩阵运算：分解、求逆与解方程

乘法只解决“乘”的问题，实际项目里更多时候要解方程、求逆、做分解。PyTorch 把这部分统一放在torch.linalg命名空间下，接口和 NumPy 的linalg基本对齐，行为上有一些差异需要注意。

3.1 求逆是最不该用的操作

torch.linalg.inv能用，但绝大多数情况下不该用。原因有三层：一是数值稳定性，显式求逆再相乘的误差通常比直接解方程大一个量级；二是性能，解方程是O(n^3)的分解加O(n^2)的回代，求逆是O(n^3)加一次完整的矩阵乘，常数更大；三是显存，n×n的逆矩阵要额外占一块空间，而解方程只需要一个n×n的分解缓存。

正确的写法是torch.linalg.solve(A, B)，直接解AX = B。它在内部做 LU 分解（对称正定时走 Cholesky，速度更快），然后回代，不显式构造逆矩阵。如果确实需要逆矩阵本身（比如要拿它做后续的多次乘法），可以用torch.linalg.inv_ex或者先lu_factor再lu_solve，把分解结果复用起来。

我在一个做高斯过程回归的模块里踩过这个坑。原代码是K_inv = torch.linalg.inv(K)然后alpha = K_inv @ y，在小数据上没问题，n到 2000 的时候开始出现loss = nan。排查下来是K的条件数接近 1e12，求逆过程中小奇异值被放大，数值彻底崩掉。改成torch.linalg.solve(K, y)并加了一点点jitter（对角线加1e-6），问题消失。这里的关键不是solve有多神，而是它避免了显式构造一个数值上很脆弱的中间量。

torch.linalg.pinv用在非方阵或者秩亏的情况下，内部走 SVD，速度比solve慢不少，但鲁棒。它的rcond参数控制小奇异值的截断阈值，默认值是max(m, n) * eps，实际用的时候通常要手动调大，比如1e-5，否则噪声方向会被当成有效信息保留下来，反而有害。

3.2 解线性方程组、最小二乘与矩阵分解

除了solve，还有几个必须知道的入口。

torch.linalg.lstsq(A, B)直接解最小二乘min ||AX - B||，不用你手动写正规方程。这一点很重要，因为正规方程(A^T A)^{-1} A^T B会把条件数平方，精度损失严重。lstsq内部走 QR 或者 SVD，稳定性好很多。它的返回值和 PyTorch 版本有关，新版返回一个 namedtuple，包含solution、residuals、rank、singular_values，取结果时注意用属性名而不是位置索引，避免版本升级后出错。

torch.linalg.qr做 QR 分解，torch.linalg.svd做奇异值分解，torch.linalg.eigh做对称矩阵的特征分解，torch.linalg.cholesky做 Cholesky 分解。这几个是构建更复杂算法的基石，比如 PCA 就是中心化之后取 SVD 的前 k 个奇异向量。

svd有一个参数full_matrices，默认True，返回完整的U和Vh，对于(m, n)且m >> n的矩阵，U是(m, m)，显存开销巨大。做降维的时候设成False，只返回经济型分解，U是(m, n)，能省下大量显存。这个参数我在处理图像特征矩阵时吃过亏，一个(50000, 128)的矩阵，full_matrices=True直接 OOM，改成False之后轻松跑完。

cholesky只对对称正定矩阵有效，但它比solve快将近一倍，而且能顺便判断正定性——如果矩阵不正定，分解会直接报错。做协方差矩阵相关的运算时，先试cholesky，失败再退回solve，是一个很实用的策略。

3.3 范数、条件数与秩：判断矩阵“好不好用”

在动手算之前，先花一行代码看看矩阵的“体质”，能省掉大量调试时间。

torch.linalg.norm(A, ord=...)支持各种范数，矩阵场景下常用ord=2（谱范数，即最大奇异值）、ord='fro'（F 范数，所有元素平方和开根）、ord=1（列和最大值）、ord=torch.inf（行和最大值）。做梯度裁剪的时候用的是向量范数，ord=2就是把所有元素平方求和开根，这个和矩阵的谱范数不是一回事，别混。

torch.linalg.cond(A)算条件数，等于最大奇异值除以最小奇异值。经验值：小于 1e3 很安全，1e3 到 1e6 要留意，超过 1e8 基本可以认为在 float32 下不可逆，超过 1e12 在 float64 下也悬。我一般会在涉及求逆或解方程的模块里加一个可选的 debug 分支，训练前跑一次cond，数值大了就报警，比等到 loss 变 nan 再回头查高效得多。

torch.linalg.matrix_rank(A)算数值秩，内部也是走 SVD，阈值同样由rcond控制。判断一个矩阵是不是满秩、要不要加正则项，用它比看det靠谱——行列式接近 0 不代表一定奇异，行列式很大也不代表条件数好，这两个指标经常误导人。

4. 从零跑一遍：线性回归的正规方程与向量化改造

理论讲完，跑一段能直接复现的代码。这一段的目标不是训一个多好的模型，而是把“同一个数学式子有四种写法，性能和稳定性各不相同”这件事演示清楚。

4.1 环境与数据构造

不需要 GPU 也能跑，CPU 上足够看出差异。

import torch import time torch.manual_seed(42) torch.set_default_dtype(torch.float32) n_samples, n_features = 5000, 64 X = torch.randn(n_samples, n_features) true_w = torch.randn(n_features, 1) y = X @ true_w + 0.01 * torch.randn(n_samples, 1) print(X.shape, y.shape) # torch.Size([5000, 64]) torch.Size([5000, 1])

数据构造里有一个细节值得说：y的形状是(n_samples, 1)而不是(n_samples,)。这个选择会影响后面所有形状推演。如果y是一维的，X.t() @ y得到(64,)，而X.t() @ X是(64, 64)，解出来的w是(64,)；如果y是二维的，w是(64, 1)。两种都能用，但混着用就会出问题，比如预测值X @ w的形状在一维w下是(n,)，二维w下是(n, 1)，和y相减时如果y是(n, 1)、预测是(n,)，会广播成(n, n)，loss 直接爆炸。我的习惯是回归任务统一用二维(n, 1)，全程形状一致，少一类 bug。

4.2 用 inv 硬算与用 solve 稳算的对比

正规方程的数学形式是w = (X^T X)^{-1} X^T y。三种实现，从“最直译”到“最推荐”。

XTX = X.t() @ X # (64, 64) XTy = X.t() @ y # (64, 1) # 写法一：显式求逆，数学上最直译，实际最不该用 w_inv = torch.linalg.inv(XTX) @ XTy # 写法二：解方程，稳定性好，速度快 w_solve = torch.linalg.solve(XTX, XTy) # 写法三：直接最小二乘，不走正规方程，条件数不被平方 w_lstsq = torch.linalg.lstsq(X, y).solution

三者结果应该非常接近，因为数据是良态的。为了看出差异，给X加一点共线性：把第 2 列设成第 1 列加一点噪声，人为把条件数拉高。

X_bad = X.clone() X_bad[:, 1] = X_bad[:, 0] + 1e-4 * torch.randn(n_samples) print("cond =", torch.linalg.cond(X_bad.t() @ X_bad).item())

条件数到 1e8 以上时，inv版本的结果会明显偏离，solve稍微好一点，lstsq最稳。这个实验我建议你自己跑一遍，把绝对误差打出来看，比看任何文字描述都有说服力。原因在于X^T X的条件数是X的条件数平方，本来 1e4 的条件数经过平方变成 1e8，float32 只有大约 7 位有效十进制数字，误差直接吃掉了大部分精度。

这也是为什么在特征维度高、特征之间相关性强的场景下，宁可直接用lstsq或者加 L2 正则（把XTX换成XTX + λI），也不要硬解正规方程。加正则之后的矩阵条件数会被λ拉低，λ取1e-4量级就能把 1e8 降到 1e4 左右，效果立竿见影。

4.3 把三重循环改写成 einsum

前一个例子是一步矩阵乘，体现不出向量化的价值。换一个更真实的场景：计算一批样本两两之间的加权相似度。

假设有b个批次，每批n个向量，每个向量d维，另有一个d×d的度量矩阵M，要算S[i,j] = x_i^T M x_j。朴素写法是三重循环：

def pairwise_naive(Xb, M): b, n, d = Xb.shape out = torch.zeros(b, n, n) for bi in range(b): for i in range(n): for j in range(n): out[bi, i, j] = Xb[bi, i] @ M @ Xb[bi, j] return out

n=200、b=8的时候这个循环要跑 32 万次内层迭代，Python 层面开销巨大。改写分两步：先把M吸收进X，即Z = X @ M，形状(b, n, d)；然后S = Z @ X.transpose(-1, -2)，形状(b, n, n)。用einsum表达就是：

def pairwise_einsum(Xb, M): Z = torch.einsum('bnd,de->bne', Xb, M) return torch.einsum('bnd,bmd->bnm', Z, Xb)

注意einsum的第二个表达式里，两个输入的d是同一个字母，表示缩并；输出的n和m分别来自两个输入，所以得到(b, n, n)。这里用不同字母n和m而不是都用n，是为了避免einsum误解成“对角”或者“逐元素”——同一个字母出现在两个输入的非缩并位置，语义是“这两个维度必须相等且逐元素对应”，不会产生两两组合。这个细节很多人第一次写会搞错，写成'bnd,bnd->bnn'得到的结果完全不同。

实测下来，n=200, b=8, d=64，朴素循环大约 1.8 秒，einsum版本 3 毫秒左右，差了两个数量级。而且循环版本还有一个隐患：out[bi, i, j] = ...是原地赋值，如果out参与自动求导，这种逐元素写入的方式在某些情况下会导致梯度计算变慢或者报错，向量化版本完全没有这个问题。

5. 形状、dtype、device 与数值精度的高频塌方点

前面讲的是“怎么写对”，这一节讲“为什么明明写对了还是出错”。矩阵运算的报错信息通常很简短，size mismatch、expected scalar type、device mismatch，但根因经常在很远的地方。

5.1 形状报错的定位套路

遇到形状报错，别急着改代码，先做三件事。

第一，把报错行涉及的所有张量形状打出来，包括中间变量，不要只打最终的两个操作数。很多时候根因是某个unsqueeze加错了位置，或者某个sum少了dim参数，导致维度凭空多出来一个。

第二，用matmul的规则手动推一遍广播。从最后一维往前对齐，逐维检查，把每一维的来源写下来。这个过程慢，但能定位到具体哪一维不匹配。比如报错说Expected size for argument #2 'mat2' (64) to be 32，那就是第二个操作数的倒数第二维是 64，第一个操作数的最后一维是 32，它们要做矩阵乘必须相等。

第三，检查有没有“隐式 squeeze”。torch.sum(x, dim=0)会真正去掉那一维，而torch.sum(x, dim=0, keepdim=True)保留成 1。代码里混用这两种写法，形状就会飘。我现在的习惯是，凡是在流程中间做 reduction，一律keepdim=True，到最后确定要输出的时候再显式squeeze，这样形状变化都有迹可循。

还有一种情况是 reshape 顺序问题。(b, n, d)直接reshape(b*n, d)是对的，但如果中间做过transpose变成(b, d, n)，再reshape(b*n, d)就会把内存顺序搞乱，得到错误的数据（不报错，但结果全错）。判断方法是用.is_contiguous()检查，或者干脆用permute之后再.reshape()，reshape会在必要时自动拷贝一份保证语义正确，代价是多一次内存分配。这个坑在实现 Transformer 的多头拆分时特别常见：(b, s, h*d)拆成(b, s, h, d)再转成(b, h, s, d)，如果不用reshape而用view，或者转置顺序写错，后面所有注意力计算全是垃圾数据。

5.2 dtype 与 device 不一致的连锁反应

float32和float64混用，矩阵乘会直接报错，PyTorch 不做隐式升级。这一点和 NumPy 不同，NumPy 会自动提升到更高精度。PyTorch 的选择是明确的：混精度必须显式转换，避免不知不觉中损失性能。

排查套路是，在报错处打印x.dtype和y.dtype。常见来源有三个：一是数据加载时numpy默认float64，torch.from_numpy之后没转；二是某个常数写成了 Pythonfloat，参与运算时被当成float32，但如果另一边是float64就报错；三是torch.zeros没指定dtype，用了全局默认值，而全局默认被某处set_default_dtype改过。

device 不一致的报错更直接，Expected all tensors to be on the same device。麻烦的是多卡场景，模型在某张卡上，数据在另一张卡上，或者某个 buffer 没有跟着.to(device)走。我的做法是写一个工具函数，在训练循环的第一步检查所有关键张量的 device 和 dtype，一次性打印出来，而不是等报错出现再逐个查。

混合精度训练时还要注意，autocast上下文里的矩阵乘会被自动转成半精度，但如果手动调用了torch.linalg下的函数，有些操作在半精度下不支持，会报错或者悄悄降级。linalg.inv、linalg.svd这类操作在半精度下精度损失严重，建议显式float()之后再算，算完再转回来。

5.3 浮点精度、TF32 与混合精度

float32有 24 位有效尾数，大约 7 位十进制精度。矩阵乘是累加操作，d维内积累加d次，误差大致按sqrt(d)增长。d=1024时，相对误差可能在 1e-5 量级，对大多数任务够用，但对需要高精度对齐的场景（比如和参考实现做数值比对）就会出问题。这时候要么用float64，要么把比对阈值放宽到 1e-4。

GPU 上还有一个额外变量：TF32。这是较新架构的 GPU 在矩阵乘上默认开启的一种折中精度格式，尾数只有 10 位，比float32少很多，但比float16多。它能显著提升矩阵乘吞吐，代价是精度下降。PyTorch 里由torch.backends.cuda.matmul.allow_tf32控制，不同版本默认值变过，做数值复现实验时要显式设成False，否则你会发现同样的代码在两张不同卡上结果不一致。

混合精度里，float16的动态范围小，容易出现下溢。做矩阵乘之前如果输入的量级差异大，建议用bfloat16，它的指数位和float32一样宽，不容易溢出，代价是尾数更短、精度更低。实际项目里通常是bfloat16更省心，尤其是在做注意力里的 softmax 之前，float16经常需要把分数先缩放一下防溢出，bfloat16基本不用管。

5.4 原地操作与自动求导的冲突

矩阵运算里最容易触发自动求导报错的是 in-place 操作。比如x.t_()、x.add_()、x[:, 0] = 0，如果x是叶子张量或者被其他计算依赖，autograd 会报a leaf Variable that requires grad is being used in an in-place operation，或者更隐蔽的版本号不匹配错误。

t_()和transpose_()在较新版本里已经不推荐用了，因为它们改变的是 stride 而不是数据，和后续的视图、求导混在一起容易出问题。统一用t()和transpose()返回新视图，虽然多了一个 Python 对象，但没有副作用。

填充 mask 的时候，x.masked_fill_(mask, 0)是原地操作，如果x需要梯度，会报错。改成x.masked_fill(mask, 0)返回新张量就行，或者在torch.no_grad()上下文里做原地填充。后者用在推理阶段更省显存，训练阶段老老实实用非原地版本。

6. 性能实测：不同 API 的速度差异从哪来

讲性能不能空谈，得有数据。下面这段 benchmark 脚本可以直接复制运行，CPU 和 GPU 都适配。

6.1 一个简单的 benchmark 脚本

import torch import time def bench(fn, warmup=10, rep=50): for _ in range(warmup): fn() if torch.cuda.is_available(): torch.cuda.synchronize() t0 = time.perf_counter() for _ in range(rep): fn() if torch.cuda.is_available(): torch.cuda.synchronize() return (time.perf_counter() - t0) / rep * 1000 # ms device = 'cuda' if torch.cuda.is_available() else 'cpu' b, n, m, p = 32, 256, 256, 256 A = torch.randn(b, n, m, device=device) B = torch.randn(b, m, p, device=device) print("bmm :", bench(lambda: torch.bmm(A, B))) print("matmul :", bench(lambda: torch.matmul(A, B))) print("einsum :", bench(lambda: torch.einsum('bij,bjk->bik', A, B))) print("loop :", bench(lambda: torch.stack([A[i] @ B[i] for i in range(b)])))

warmup 不能省。第一次调用某个形状的矩阵乘时，底层库要做算法选择、分配缓存，耗时会明显偏高。不做 warmup 直接计时，测出来的数字完全没有参考价值。GPU 上还要加torch.cuda.synchronize()，因为核函数是异步提交的，不同步的话计时只量到了“提交任务”的时间，不是“算完”的时间。

6.2 结果解读

在 GPU 上，bmm和matmul的耗时基本一致，因为matmul在批量维严格相等时会直接转调bmm，没有额外开销。einsum稍慢一点，大概 5% 到 15%，主要是表达式解析和中间调度的开销。loop版本慢一个数量级以上，因为 32 次独立的核函数启动，每次都有调度延迟。

在 CPU 上结论不同：小矩阵（比如 32×32）时，loop版本可能反而更快，因为每次调用的矩阵小，多线程调度的固定开销占比高，批量 API 的优势体现不出来。矩阵尺寸超过 128 之后，批量 API 才开始明显领先。这个交叉点在不同 CPU 上不一样，你的机器上具体是多少，跑一下就知道，别照搬别人的数字。

还有一个容易被忽略的因素是非连续张量。如果A是permute出来的，.is_contiguous()是False，bmm在内部要先拷一份连续版本，耗时可能翻倍。做 benchmark 时要区分连续和非连续两种情况，否则会得出“某个 API 很慢”的错误结论，实际上是输入布局的问题。修法是提前.contiguous()，把拷贝成本单独承担，后续多次复用就摊薄了。

具体到工程选择，我的排序是：核心路径上的固定形状批量矩阵乘，用bmm，语义明确、性能稳定；维度不固定或者需要广播的，用matmul；涉及三个以上张量缩并、或者表达式本身就复杂（比如注意力里的QK^T之后再乘V），用einsum换可读性，必要时加optimize；至于手写循环，只在 batch 很小、调试阶段用，上生产前替换掉。

我个人的体会是，矩阵运算的坑大多不在“不会写”，而在“写错了但不报错”。广播把形状不匹配变成静默计算，reshape把内存顺序问题藏起来，float32把精度损失摊薄到你察觉不到，直到某一天 loss 曲线诡异地上飘，或者和参考实现比对时差了 1e-3，才开始逐个环节回查。所以我现在的习惯是，每个涉及矩阵运算的模块写完，先构造一个手算得出来的小例子（3×3 的矩阵足够），把期望结果和实际结果打出来对一遍，再放大到真实尺寸。这一步多花十分钟，能省掉后面几天的排查。另外，凡是形状变化的地方，加一行assert或者注释写清楚变换前后的形状，也是我自己强制执行的一条规矩——过两周再回来看代码，那行注释比任何记忆都可靠。