1. 先把“矩阵运算”这件事在PyTorch张量体系里说清楚
PyTorch张量的数学运算里，矩阵运算大概是被问得最多的一块。原因很朴素：深度学习里几乎所有核心算子——全连接层、注意力里的QKV打分、卷积展开后的乘加、归一化里的协方差矩阵——最后都能归结成几次矩阵乘法。很多人调用API没问题，但一遇到shape对不上、速度突然掉一半、反向传播报in-place错误，就抓瞎了。问题不在API记不熟，而在于脑子里没建立起“形状怎么对齐、内存怎么排布、梯度怎么回传”这三条并行的线。
这篇内容面向的是已经能写几行PyTorch、但没系统梳理过矩阵运算的人。不管你是刚跑通MNIST想弄明白@到底做了什么，还是已经在调模型、想搞清楚为什么加一个transpose就慢了30%，下面这些拆解应该都能对上你的痒处。我不会只列API清单，重点是解释每个选择背后的代价——为什么这里用mm不用matmul，为什么那里宁可多拷一次内存也要contiguous()，为什么用inv解方程是新手最容易犯的错之一。
1.1 张量不是数学符号，是带步长描述的内存视图
初学者最容易把张量想成一张二维表格，实际上它是一块连续内存加一份“步长（stride）元数据”。一个(3, 4)的float32张量，内存里就是12个连续浮点数，步长是(4, 1)，意思是“行方向每跳一步跨4个元素，列方向每跳一步跨1个元素”。真正决定形状的是这份元数据，而不是内存的物理排布。
这个认知一旦建立，很多诡异行为立刻就不诡异了。比如a.T根本不搬运任何数据，它只是把步长从(4, 1)改成(1, 4)，告诉后续算子“请按列优先来读”。所以转置是零成本的，但也因此转置后的张量不是内存连续的。后面所有关于性能的坑，根源都在这里。
我习惯在调试时随手打一行print(x.shape, x.stride(), x.is_contiguous())，比看什么都管用。尤其是接手别人代码的时候，一个看起来平平无奇的(B, S, H, D)四维张量，如果stride显示它是从某个permute过来的，那它十有八九不是连续的，后面接矩阵乘法就要留意开销。
注意：
stride()在排查“为什么这段代码比预期慢”时，信息量往往比shape()更大。
1.2 元素级运算和矩阵级运算的分界线
torch.mul、*、torch.add这类是逐元素运算，两个张量的形状要么完全相同，要么满足广播规则；而torch.mm、torch.matmul、@是矩阵乘法，遵循的是(m, k) @ (k, n) -> (m, n)这条内维收缩规则。这两套规则长得很像但完全不同，*和@只差一个符号，犯错的成本却很高。
一个典型的混淆场景：你想算两个向量的内积，写了a * b，得到的是逐元素乘积，还得再.sum()一次；写a @ b一步到位。两者在(n,)形状下的结果一个是张量、一个是标量，后续梯度形状也不同。更隐蔽的是，a * b在广播下可能悄悄把你的(B, 1)扩成(B, B)，代码不报错，显存却突然炸了。
我的经验是：凡是涉及“内维求和”的语义，一律用矩阵乘法API；凡是涉及“同位对齐”的语义，才用逐元素。写代码时强迫自己在注释里标一句“这里做的是内维收缩”，能挡掉一大半形状事故。
1.3 为什么必须区分mm、bmm、matmul这三个名字
这三个API的功能有重叠，但约束条件不同，选错了轻则报错，重则性能打折。
| API | 输入维度要求 | 是否支持广播 | 典型用途 |
|---|---|---|---|
torch.mm | 严格2D | 否 | 单次矩阵乘，形状已确定 |
torch.bmm | 严格3D，batch数必须一致 | 否 | 批量矩阵乘，batch对齐 |
torch.matmul/@ | 任意维，按后两维做矩阵乘 | 是 | 通用场景，含注意力等4D情况 |
torch.einsum | 任意维，按表达式指定 | 视表达式而定 | 复杂收缩、多张量连乘 |
mm和bmm是“我说清楚了我就要这个形状，你别帮我猜”，所以它们的前置检查更严格、开销更低；matmul是“你按广播规则帮我推”，灵活但需要多做一轮形状匹配。日常写模型，matmul和@是主力，但在数据加载、损失计算这类形状固定的热点路径上，用mm更容易暴露维度错误，也少一层推导。
2. 矩阵乘法API选型：从形状推演到实际写法
2.1 用shape对齐法手推广播规则
广播规则官方说法是“从尾部维度开始对齐，缺失的补1，为1的维度可扩展”。这话没错，但落到四维张量上容易绕晕。我自己的土办法叫“右对齐补1法”：把两个shape右对齐写下来，左边不够的补1，然后逐位比对，某一位要么相等、要么有一个是1，否则报错。
举个例子，注意力打分里最常见的写法：
import torch B, H, S, D = 2, 8, 128, 64 q = torch.randn(B, H, S, D) k = torch.randn(B, H, S, D) # 右对齐： # q: (2, 8, 128, 64) # kT: (2, 8, 64, 128) scores = q @ k.transpose(-2, -1) # -> (2, 8, 128, 128)transpose(-2, -1)只交换最后两维，前面(2, 8)原封不动。这里之所以能算，是因为matmul只对“最后两维”执行矩阵乘语义，前面的维度全部走广播。这个规则一定要刻进肌肉记忆：matmul的全部广播发生在batch维度，最后两维永远是(m, k) @ (k, n)。
再看一个容易出事的例子：
a = torch.randn(4, 1, 3, 5) b = torch.randn(4, 6, 5, 2) c = a @ b # -> (4, 6, 3, 2)结果是(4, 6, 3, 2)，因为a的第二维1被扩成了6。这看起来“很智能”，但如果你的本意是4个batch、每个batch一组，实际上多出了6倍计算量和6倍中间缓存。我就见过有人在这里显存翻了好几倍，排查了半天以为是模型结构问题。
实操心得：写四维矩阵乘之前，先把两个shape写在纸上右对齐比一遍，比在代码里试错快得多。如果推出来某一维是1、结果却是别的数，先问自己“这个扩展是我要的吗”。
2.2 einsum：把推导过程写成可读表达式
einsum是我认为最值得花20分钟专门学的API。它的价值不在于性能（小规模时略慢，和matmul打平），而在于它把形状推导变成了显式声明。你写'bij,bjk->bik'，读代码的人一眼就知道谁是内维、谁被保留、谁被求和消掉。
A = torch.randn(4, 16, 32) B = torch.randn(4, 32, 8) # 三种等价写法 c1 = torch.bmm(A, B) c2 = A @ B c3 = torch.einsum('bij,bjk->bik', A, B)三者在数值上是一致的（浮点误差范围内），但可读性和调试体验完全不同。当表达式超过两个张量时，einsum的优势就压倒性了：
# 双线性打分：q (B,H,S,D), k (B,H,T,D), w (H,D,D) score = torch.einsum('bhsd,bhtd,hde->bhst', q, k, w)这一行如果拆成matmul要写三四步，中间还要transpose和unsqueeze，很容易在某一步把维度搞反。当然代价是einsum对内存布局不敏感，它会自己做必要的重排，所以在大张量上不一定最优。我的实际做法是：先用einsum把逻辑跑通跑对，确认shape和数值无误，如果这段进了profiler的热点，再手工改写成matmul + 提前contiguous的版本。
顺带提一句，einsum还有两个实用变体：torch.einsum里加...可以做“省略前导维度”的写法，比如'...ij,...jk->...ik'会自动匹配任意前缀batch维度，适合写通用工具函数；而'ii->i'这类表达式能直接取对角线。
2.3 multi_dot与加法融合：两个容易被忽略的API
如果你要连乘三个以上矩阵，torch.linalg.multi_dot值得知道。矩阵连乘的结合顺序会显著影响计算量：(AB)C的代价是m*n*p + m*p*q，而A(BC)是n*p*q + m*n*q，当维度悬殊时能差一个数量级。multi_dot会用动态规划找较优的结合顺序。
A = torch.randn(1000, 2) B = torch.randn(2, 1000) C = torch.randn(1000, 2) naive = (A @ B) @ C # 先算出 1000x1000 的大矩阵 smart = torch.linalg.multi_dot([A, B, C]) # 走 1000x2 的低秩路径这个例子里，A @ B会产生一个百万元素的中间结果，而multi_dot会先算B @ C，中间量只有2*2。两者最终结果一样，但显存和耗时的差距是数量级的。这类低秩结构在手写模型的计算图里其实很常见。
另一个是torch.addmm，把beta * input + alpha * (mat1 @ mat2)融合成一个kernel。全连接层y = x @ W.T + b用它比先matmul再add少一次中间张量读写。在小模型上看不出差别，但在大batch的推理路径上是有意义的。同类还有torch.baddbmm处理批量情况。
3. 转置与内存布局：矩阵运算里最隐蔽的性能暗坑
3.1 transpose之后到底发生了什么
前面说了transpose只改步长不搬数据。这个设计的直接后果是：转置后的张量交给矩阵乘法，底层库可能要先做一次隐式拷贝。很多底层实现要求输入是行优先连续的，遇到非连续输入会先规整化再计算，这次拷贝你看不见，但它实实在在吃带宽。
a = torch.randn(1024, 1024) b = torch.randn(1024, 1024) b_t = b.T print(b_t.is_contiguous()) # False # 下面这行的实际开销，可能比 c1 高出一截 c = a @ b_t在CPU上这个差距有时不明显，在GPU上大矩阵时能看出差别。当然，底层库也不傻，很多情况下它能识别转置模式直接按转置方式读，这就是所谓“转置友好的GEMM”。我的建议是：别为了这点性能提前优化，但要知道这个变量的存在。如果profiler显示某个matmul异常慢，而它恰好接在一个permute后面，那基本就是这里的问题，加一个.contiguous()再测一次，对比一下。
顺带区分三个容易混的操作：transpose交换任意两维；permute一次性重排所有维度；t()只针对2D转置。三者都不拷贝数据，都返回非连续视图。
3.2 view为什么报错，reshape为什么能救
view要求源张量在“按新形状解读时”内存仍然连续，否则直接抛错。reshape则聪明一点：能view就view，不行就悄悄拷贝一份再给你。这就是为什么a.T.view(-1)必报错，而a.T.reshape(-1)能过。
a = torch.arange(12).reshape(3, 4) a.T.view(-1) # RuntimeError a.T.reshape(-1) # 可以，但发生了拷贝这个区别在debug时很关键。我见过有人在循环里对转置结果反复reshape，每次都在偷偷拷贝，循环一跑就是几百次，速度自然上不去。正确做法是提前contiguous()一次，后面想怎么view怎么view。
a_t = a.T.contiguous() # 一次拷贝 for _ in range(100): x = a_t.view(-1) # 后续零成本3.3 批次维度放在哪一位更划算
同样一个四维张量，维度顺序不同，性能可能差出可观的比例。核心原则是：让最内层维度（最后一维）承载连续访问，让需要一起做矩阵乘的两个维度尽量相邻。
以多头注意力为例，两种常见排布：
# 排布A：batch优先 q = torch.randn(B, H, S, D) scores = q @ q.transpose(-2, -1) # 排布B：序列优先 q2 = q.transpose(1, 2) # (B, S, H, D)排布A在做matmul时，(S, D) @ (D, S)是标准形状，transpose(-2,-1)只影响最后两维，前两维连续；排布B虽然语义上更贴近“每个token看所有头”，但transpose之后中间维度变味了，后续再和其他张量做广播时容易出岔子。
我踩过的具体坑是：把(B, H, S, D)先permute成(B, S, H, D)做了一层运算，再permute回去做第二个矩阵乘，来回折腾两次，每个permute都制造了非连续视图，两个matmul都多吃了拷贝开销。后来把维度顺序统一掉，从头到尾只在最后两维做文章，速度快了一点，代码也短了一截。
注意：不要孤立地看单个shape。整条链路里所有矩阵运算的维度排布如果能保持一致，省下的是成串的隐式拷贝。
4. 线性代数求解：inv、solve、分解怎么选
4.1 用inv解方程的精度代价
“解Ax = b”这件事，新手最直觉的写法是x = torch.inverse(A) @ b。这个写法能用，但有两个问题：一是inverse里藏着一次分解，你算出A^{-1}之后再乘b，等于把分解的功夫白扔了一半；二是显式求逆会放大条件数，A越接近奇异，误差越夸张。
A = torch.randn(64, 64) b = torch.randn(64) x_bad = torch.inverse(A) @ b x_good = torch.linalg.solve(A, b)torch.linalg.solve内部用的是LU分解加前后代入，不做显式求逆，数值上更稳，速度也更快。大批量情况下还有solve_ex可以顺便拿到分解结果复用，或者lu_factor+lu_solve分离两步。
我拿一个故意构造的病态矩阵试过：让A的最小奇异值和最大奇异值差出好几个数量级，inv路线的相对误差比solve高出不少。练手数据上看不出问题，实际模型里一旦遇到接近共线的特征，这个误差是会累积的。
实操心得：代码里出现
torch.inverse，基本都可以问一句“我是不是在解方程”。如果确实只是要矩阵本身（比如某些理论推导、协方差白化），那另说；如果目的是拿x，一律换solve。
4.2 SVD、QR、Cholesky的适用场景
不同分解对应不同前提，选错要么报错要么白算：
| 分解 | 前提条件 | 典型用途 | API |
|---|---|---|---|
| LU | 方阵、可逆（不严格要求） | 解线性方程组 | linalg.lu_factor |
| QR | 任意形状 | 最小二乘、正交化 | linalg.qr |
| Cholesky | 对称正定 | 协方差、高斯过程 | linalg.cholesky |
| SVD | 任意形状 | 降维、伪逆、低秩近似 | linalg.svd |
| 特征值 | 方阵 | 谱分析、稳定性判断 | linalg.eig/eigh |
几个实务要点。eigh专门处理对称矩阵，比通用eig快而且返回的特征值一定是实数；如果你的矩阵理论上对称，但浮点累加导致不对称，先用(A + A.T) / 2对称化再传进去，稳定性会好很多。cholesky在遇到数值上不完全正定的矩阵时会失败，实践中常加一个极小的对角扰动A + eps * I，eps取1e-6量级，具体要看数据尺度。
svd有个容易忽略的参数full_matrices，默认True会返回完整的U和Vh，对于高瘦或矮胖矩阵来说是巨大的浪费。算伪逆、做低秩近似时把full_matrices=False设上，省下的显存很可观：
U, S, Vh = torch.linalg.svd(X, full_matrices=False) k = 32 X_lowrank = (U[:, :k] * S[:k]) @ Vh[:k, :]4.3 行列式、迹、秩这类标量量的坑
torch.linalg.det在矩阵维度变大时数值上很脆弱，因为它本质上是所有特征值之积，量级很容易溢出或下溢。真要算对数行列式（变分推断里常见），用torch.linalg.slogdet返回符号和对数绝对值，稳定性好得多：
sign, logabsdet = torch.linalg.slogdet(A)torch.trace只能取主对角线之和，想取批次内每个矩阵的迹得用torch.diagonal(A, dim1=-2, dim2=-1).sum(-1)，直接trace会报维度错误。
至于秩，torch.linalg.matrix_rank内部靠SVD加阈值判断，比你手写(A != 0).sum()靠谱得多，但也要理解它是数值意义下的秩——浮点误差会让本该线性相关的行“看起来”独立。检查共线性时，我一般直接看奇异值的比值，比看秩更直观。
5. 矩阵运算与自动求导的配合
5.1 矩阵乘的梯度到底怎么回传
C = A @ B的反向传播，对A的梯度是grad_C @ B.T，对B的梯度是A.T @ grad_C。这两条公式本身简单，但里面藏着一个训练时经常被忽略的事实：只要A或B有一个不需要梯度，对应的那次矩阵乘就可以跳过。
W = torch.randn(128, 64, requires_grad=True) x = torch.randn(32, 128) # 不需要梯度 y = x @ W # 反向时不需要算 dA这个机制在冻结骨干网络、做特征提取时非常有用。但要注意，torch.no_grad()和requires_grad=False是两回事：前者连中间结果都不记录，后者只是标记某个张量不需要梯度，图还会照常构建。推理阶段一律套torch.no_grad()，训练里冻结参数用后者，别混用。
另外一个小细节：A.T在反向里也会用到，如果A本来就是某个转置来的非连续张量，两次非连续读会让梯度计算也变慢。这种情况下提前contiguous一次，省的是两处开销。
5.2 梯度检查：小规模数值差分
手写自定义算子或者怀疑某段矩阵运算的梯度对不对，最直接的办法是数值差分。用double精度、取小维度、中心差分，误差能压到1e-6量级：
import torch def grad_check(fn, x, eps=1e-6): x = x.double().requires_grad_(True) y = fn(x) y.backward(torch.ones_like(y)) analytical = x.grad.clone() numerical = torch.zeros_like(x) with torch.no_grad(): it = torch.ndindex if False else None for idx in torch.cartesian_prod(*[torch.arange(s) for s in x.shape]): idx = tuple(idx.tolist()) xp = x.clone(); xp[idx] += eps xm = x.clone(); xm[idx] -= eps numerical[idx] = (fn(xp).sum() - fn(xm).sum()) / (2 * eps) return (analytical - numerical).abs().max()这里一定要用double，float32的差分误差会盖过真实差异。还有一点：差分检查只在fn是确定性函数时有效，如果里面混了dropout、随机采样，得先切到eval模式。
我在排查一个问题时发现，某个自定义的批量矩阵乘在batch维度上做了广播，前向结果完全正确，反向却是错的——因为广播维度上的梯度需要sum归并，手写实现漏了这一步。数值差分一跑就暴露了，最大误差是1e-1量级，非常明显。这类bug只靠前向数值对拍是发现不了的。
5.3 in-place改动带来的报错复盘
矩阵运算和原地操作放在一起，最典型的报错是“a leaf Variable that requires grad is being used in an in-place operation”或者反向时的“one of the variables needed for gradient computation has been modified by an inplace operation”。
前一种好理解：叶子张量不能直接原地改，得在no_grad里改或者用.data（不推荐）。后一种更隐蔽：你改的可能不是叶子，而是某个中间结果，但那个中间结果恰好被反向公式需要。
a = torch.randn(4, 4, requires_grad=True) b = torch.randn(4, 4, requires_grad=True) c = a @ b # 反向需要 a 和 b a += 1 # 原地修改 a，反向会报错正确的改法是a = a + 1，建立一个新节点。如果确实需要在优化器更新里原地改参数，那要在torch.no_grad()上下文里做，这也是各种优化器内部的标准做法。
总结一下我的判断口诀：凡是被反向公式引用的张量，一律不原地改；只有纯前向的临时缓冲，才考虑原地省内存。像relu_、add_这种就地版本，在激活函数位置用问题不大，但在矩阵运算的输入上用，风险很高。
6. 常见问题与排查速查
6.1 形状类报错逐条对照
| 报错信息关键词 | 大概率原因 | 处理办法 |
|---|---|---|
mat1 and mat2 shapes cannot be multiplied | 内维不相等 | 打印两个shape右对齐比对内维 |
size mismatch | 广播不成立，某维既不等也不为1 | 补unsqueeze或squeeze |
Expected 2D or 3D input | 给mm/bmm传了错误维度 | 换matmul或先整理维度 |
view size is not compatible | 源张量非连续 | 改reshape或先contiguous |
expanded tensor ... in-place | 广播后的视图被原地改 | 先clone再改 |
CUDA out of memory（矩阵乘前后） | 广播把计算量放大了 | 逐维检查是否有意外的1维扩展 |
其中最后一条我想多说两句。显存爆炸不一定是模型太大，很多时候是一次不经意的广播。比如(B, 1, S, D) @ (B, H, D, S)，你以为在算H个头，结果一个批次里把(B,1,...)扩成了(B,H,...)，中间激活直接乘以H。这类问题用torch.cuda.max_memory_allocated()在可疑代码前后各打一次，差值能立刻定位到具体哪一行。
6.2 NaN与数值异常排查
矩阵运算出NaN，常见来源就那么几个：
- 求逆/求解遇到奇异矩阵：检查矩阵条件数，用
solve替代inv，必要时加eps正则。 - Cholesky对非正定矩阵：加
eps * I，或改用eigh并裁掉负特征值。 - softmax前打分过大：注意力里
q @ k.T没除以sqrt(D)，数值直接溢出。这一步缩放不是可选项，是必需的。 - log里出现0或负数：行列式、概率密度计算时用
slogdet和log_softmax的稳定版本。 - 半精度下的累加：
float16的表示范围窄，大矩阵乘法内维很长时，用torch.matmul的累加器通常还是float32，但如果你手工分块累加并保持半精度，就容易下溢。
排查手法上，我一般先定位到具体算子：用torch.autograd.set_detect_anomaly(True)跑一次，它会告诉你反向里哪一步先出现NaN。然后看这个算子的输入，torch.isfinite(x).all()逐层往前扫，找到第一个出现非有限值的位置。
注意：
set_detect_anomaly会让训练变慢好几倍，只在定位问题时开，别留在生产代码里。
6.3 精度与速度的开关清单
矩阵运算的性能，除了算法层面，还有一组全局开关值得知道：
# 允许 TF32（较新GPU上的张量核心精度模式，默认可能已开） torch.backends.cuda.matmul.allow_tf32 = True # 半精度矩阵乘的累加精度控制 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True # 确定性模式（复现实验用，会变慢） torch.use_deterministic_algorithms(True)这组开关的取舍要看你当前在干什么。做研究要复现结果，确定性模式该开就开；做推理要吞吐，TF32和半精度该开就开。但千万别在同一个项目里一会儿开一会儿关，否则你会得到“同一份代码两次跑出不同loss”的灵异现象，排查起来极其消耗精力。
另外，torch.compile对矩阵运算链的融合效果相当明显，尤其是连续几个matmul加逐元素操作的组合。我试过一个包含5次中小规模矩阵乘的模块，编译后整体耗时下降明显，主要收益来自kernel融合和少了中间张量的显存往返。代价是首次编译有等待时间，输入shape如果频繁变化还会触发重编译，所以更适合形状固定的推理路径。
最后分享一个我一直在用的调试习惯：任何涉及三次以上矩阵运算的函数，写完后先用小维度（比如B=2, S=4, D=8）跑一遍，把每一步的shape和dtype打出来，确认整条链路符合预期，再换成真实规模。小规模下报错信息更短、更好读，而且能在几秒内跑完，比在真实数据上等半分钟看一个size mismatch高效得多。这个习惯帮我省下的时间，大概比学任何单个API都多。
