编写第一个压测脚本：矩阵乘法算子压测

好，咱们直接动手。前面聊了那么多理论，现在该写点真东西了。

我个人习惯，学任何新东西都从最简单的开始。矩阵乘法，说白了就是深度学习里最基础的算子之一。你想想看，全连接层、注意力机制、卷积的底层实现，哪个离得开它？所以用它来做第一个压测脚本，再合适不过。

环境准备

先确认你的机器上有 PyTorch。我建议用 1.10 以上版本，CUDA 最好也配好。怎么检查？跑两行代码：

import torch
print(torch.__version__)
print(torch.cuda.is_available())

如果输出 True，恭喜你，可以开始了。如果输出 False，嗯，那你得先搞定 CUDA 环境。我在项目中遇到过有人折腾了一整天，结果发现 PyTorch 装的是 CPU 版本——这种坑，踩一次就够了。

核心思路

压测一个算子，我们需要记录两样东西：耗时和吞吐。

耗时：一次算子执行花了多少毫秒
吞吐：单位时间内处理了多少数据，通常用 TFLOPS 或 GB/s 表示

对于矩阵乘法，吞吐的计算公式很简单：

吞吐 = 2 * M * N * K / 耗时

为什么是 2？因为一次矩阵乘法包含乘法和加法两种操作。M、N、K 是矩阵的维度。这个公式我建议你记下来，后面所有算子的吞吐计算都大同小异。

写脚本

来，直接上代码。这是我写的一个最小可用的压测脚本：

import torch
import time

def benchmark_matmul(M, N, K, device='cuda', num_warmup=10, num_iter=100):
    """
    矩阵乘法压测函数
    M, N, K: 矩阵维度
    device: 运行设备
    num_warmup: 预热次数
    num_iter: 正式测试次数
    """
    # 创建随机矩阵
    A = torch.randn(M, K, device=device)
    B = torch.randn(K, N, device=device)
    
    # 预热：让 GPU 先跑起来
    for _ in range(num_warmup):
        C = torch.mm(A, B)
    
    # 同步，确保前面的操作都完成了
    torch.cuda.synchronize()
    
    # 开始计时
    start = time.time()
    for _ in range(num_iter):
        C = torch.mm(A, B)
    torch.cuda.synchronize()
    end = time.time()
    
    # 计算耗时和吞吐
    total_time = (end - start) * 1000  # 转成毫秒
    avg_time = total_time / num_iter
    
    # 吞吐：2 * M * N * K / 平均耗时(秒)
    flops = 2 * M * N * K
    throughput = flops / (avg_time / 1000) / 1e12  # 转成 TFLOPS
    
    return avg_time, throughput

# 测试一下
M, N, K = 1024, 1024, 1024
avg_time, throughput = benchmark_matmul(M, N, K)
print(f"平均耗时: {avg_time:.3f} ms")
print(f"吞吐: {throughput:.3f} TFLOPS")

关键点说明：

预热：GPU 第一次跑的时候很慢，因为要加载驱动、分配显存。不预热的话，测出来的数据会偏大很多。
同步：PyTorch 的 CUDA 操作是异步的。不加 torch.cuda.synchronize()，你测的只是 CPU 提交任务的时间，不是 GPU 真正执行的时间。
多次迭代：单次测量误差太大，跑 100 次取平均，结果才稳定。

避坑指南

我曾经犯过一个低级错误：忘记加预热，直接跑了一次就记录结果。结果测出来的吞吐只有理论值的 1/10。当时我还以为是显卡坏了，折腾了半天才发现是预热的问题。嗯，这种坑，你记住就好。

还有一个常见问题：矩阵维度选太小。比如 M=N=K=16，这种小矩阵在 GPU 上根本跑不满，测出来的吞吐会很低。我建议至少从 512 起步，越大越好，直到显存撑不住为止。

结果分析

跑一下上面的代码，你会得到类似这样的结果：

矩阵维度	平均耗时 (ms)	吞吐 (TFLOPS)
512x512x512	0.152	1.76
1024x1024x1024	1.021	2.10
2048x2048x2048	7.834	2.19
4096x4096x4096	62.112	2.21

看到没？随着矩阵变大，吞吐逐渐接近一个稳定值。这个稳定值就是你的 GPU 在矩阵乘法上的理论峰值。比如我用的这张卡，理论峰值是 2.4 TFLOPS，实测 2.21，利用率大概 92%。还不错。

小技巧：如果你想测出 GPU 的极限，可以试试把矩阵维度设成 2 的幂次，比如 1024、2048、4096。因为 GPU 的硬件设计对 2 的幂次有优化，测出来的结果更接近理论值。

核心逻辑流程图

下面这张图展示了整个压测流程，我画成 SVG 方便你理解：

扩展一下

上面的脚本只测了方阵。实际工作中，矩阵往往不是方的。比如 Transformer 里的 QKV 投影，M=1024, N=64, K=1024 这种。我建议你改一下参数，多测几组不同形状的矩阵，看看 GPU 的表现有什么差异。

另外，别忘了记录环境信息：GPU 型号、驱动版本、PyTorch 版本、CUDA 版本。这些信息在对比结果时非常重要。我习惯在脚本开头加一行：

print(f"GPU: {torch.cuda.get_device_name(0)}")

好了，第一个压测脚本就写完了。你跑一下试试，看看你的显卡能跑多少 TFLOPS？

警告：别一上来就测 4096x4096 的矩阵。显存占用 = M*K*4 + K*N*4 + M*N*4 字节。4096 的矩阵需要 192MB 显存，看起来不多，但如果你同时跑多个进程，显存可能爆掉。我建议从小矩阵开始，逐步增大，直到显存占用达到 80% 左右为止。

想要更多资料，联系可可

📚 课程目录

编写第一个压测脚本：矩阵乘法算子压测

环境准备

核心思路

写脚本

避坑指南

结果分析

核心逻辑流程图

扩展一下