相同FLOPS下numpy单一大矩阵乘法比多次小矩阵乘法慢的原因

架构老林

2026-10-5

测试代码

import numpy as np
import time

C = np.random.rand(4, 100000)

# TEST A
times_A = []
for _ in range(10):
    t0 = time.perf_counter()
    A = np.random.rand(100, 3, 4)
    X = A.dot(C)
    times_A.append(time.perf_counter() - t0)

# TEST B
times_B = []
for _ in range(10):
    t0 = time.perf_counter()
    for _ in range(100):
        B = np.random.rand(3, 4)
        X = B.dot(C)
    times_B.append(time.perf_counter() - t0)

print('TIME A: ', np.mean(times_A))
print('TIME B: ', np.mean(times_B))

运行输出

TIME A:  1.002193902921863
TIME B:  0.0581539266044274

性能差异原因

输出内存开销差异是核心因素
TEST A中三维矩阵A与C点乘的输出形状为(100, 3, 100000)，按双精度浮点数计算总大小约240MB，远超普通CPU的L2/L3缓存容量。每次计算都需要申请大块连续内存，且计算结果必须全部写入主存，内存分配、访问、写入的开销占比极高。
TEST B每次点乘的输出形状为(3, 100000)，总大小仅约2.4MB，可完全放入CPU缓存，且循环中每次生成的X会被覆盖回收，内存可以复用，几乎没有大块内存操作的额外开销。
BLAS库计算适配效率差异
主流优化BLAS库（MKL、OpenBLAS等）的GEMM矩阵乘接口对「小左矩阵 × 大右矩阵」的场景做了极致优化：TEST B中的左矩阵B只有3行，可以完全放入L1缓存，计算时C的每行只需要读取一次，缓存命中率接近100%。
TEST A的批量三维矩阵乘在numpy底层实现中，虽然也是拆分批量维度调用GEMM，但输出的三维数组内存布局会导致写入时缓存命中率极低，反而放大了计算开销。
Python API调用开销的影响被掩盖
你预期的循环调用Python-C API的开销每次只有微秒级，和几百毫秒级的内存访问、大数组操作开销相比完全可以忽略，不会影响最终耗时排名。