torch.bmm适用于标准三维批量矩阵乘法（B×N×M × B×M×P），高效但要求batch size严格对齐；torch.einsum更灵活，支持不规则维度、广播和掩码，语义清晰但语法稍重，适合复杂场景。
PyTorch中torch.bmm和torch.einsum怎么选？
批量矩阵乘法的核心是避免用for循环调用torch.mm，否则GPU利用率低、显存碎片多、速度慢。最常用的是torch.bmm，但它只支持三维输入（B×N×M × B×M×P），且要求batch size严格对齐。如果维度不规整或需要更灵活的索引（比如带mask、跨轴广播），torch.einsum更稳妥，虽然语法稍重，但可读性和扩展性更好。
常见错误现象：RuntimeError: Expected 3-D tensor, but got 2-D for argument #1 'input'——这是误把单个矩阵当批量传给torch.bmm；或者torch.bmm两个张量batch size不一致，直接报错而非静默截断。
-
torch.bmm适合标准场景：所有样本都含相同数量的子矩阵，如一批句子的注意力权重计算（batch_size × seq_len × seq_len） -
torch.einsum('bik,bkj->bij', A, B)语义清晰，支持自动广播、省略维度、甚至混合标量运算，调试时更容易定位维度错位 - 性能上，
torch.bmm在纯三维同构场景下通常快5–10%，但差别随batch size增大而收窄；若需条件逻辑（如只对部分样本计算），torch.einsum配合布尔索引更自然
如何处理变长序列导致的批量矩阵乘法失败？
真实NLP或图神经网络中，序列长度常不一致，强行pad成统一shape会导致大量无效计算。直接喂给torch.bmm会出错，因为padding后虽满足三维，但实际有效区域稀疏。
正确做法不是“先pad再乘”，而是用torch.masked_fill或torch.where屏蔽无效位置，让梯度只流经有效区域。注意：不能只靠mask * result，必须在乘法前就阻断无效路径，否则NaN风险高。
立即学习“Python免费学习笔记（深入）”；
- 先构造布尔mask：
mask = torch.arange(max_len)[None, :] （<code>lengths是每条样本的真实长度） - 将mask扩展为三维用于点乘：
mask_3d = mask.unsqueeze(-1) & mask.unsqueeze(-2)，得到B×L×L的注意力掩码 - 在
torch.bmm后立即应用：attn_weights = torch.bmm(Q, K.transpose(-2, -1)) / scale，再attn_weights = attn_weights.masked_fill(~mask_3d, float('-inf')) - 避免在
torch.einsum里硬编码mask维度——它不自动广播，得显式写出'bik,bkj,bij->bij'这类三输入形式，易错
torch.matmul能替代torch.bmm吗？有什么坑？
能，但要小心广播规则。在PyTorch中，torch.matmul对输入维度有智能推导：当输入是二维，行为同torch.mm；当至少一个是三维，就按batch semantics广播，等价于torch.bmm（前提是其他维度可对齐）。但它不校验batch size是否一致——如果A是2×4×5、B是3×5×6，torch.matmul会尝试广播成6×4×6，结果完全不对，且无警告。
- 显式检查batch维度：
assert A.shape[0] == B.shape[0], "Batch size mismatch"，别依赖matmul自动兜底 - 遇到
torch.Size([1, 8, 16])×torch.Size([8, 16, 32])这种非对齐情况，torch.matmul会静默广播第一个张量的batch维到8，结果变成8×8×32，极易被忽略 - 在JIT或Triton kernel中，
torch.bmm的符号更稳定，torch.matmul可能触发额外的reshape开销
为什么torch.compile对批量矩阵乘法加速有限？
PyTorch 2.0+ 的torch.compile对torch.bmm这类底层cuBLAS调用优化空间很小——它已经是最优实现。真正能被加速的是周围控制流：比如动态生成mask、条件分支下的不同乘法路径、或嵌套在nn.Module.forward中与激活函数/归一化混合的复杂计算图。
实测显示，单纯包装torch.bmm加@torch.compile，训练吞吐量变化通常在±3%以内；但若把mask构建+softmax+dropout+bmm整个块编译，端到端可提速12–18%（尤其在A100上）。
- 不要单独编译单个
torch.bmm调用，浪费编译时间 - 把
bmm嵌入更大粒度的forward函数中再编译，确保输入shape在多次迭代中稳定（否则触发recompilation） - 启用
fullgraph=True，否则control flow（如if batch_size > 32）会导致fallback，失去加速效果
批量矩阵乘法的瓶颈从来不在乘法本身，而在数据搬运、mask同步、以及前后算子的内存布局是否连续。写高效代码的关键，是让bmm的输入已经处于contiguous()状态，并复用中间张量的storage，而不是反复拷贝或view。
