GPU算子性能优化与显存控制实战 - 第1节

一、算子优化概述：到底什么是算子？

大家好，我是你们这门课的老朋友。今天咱们正式开篇，聊聊算子优化。

先问个问题：你写神经网络的时候，调用的那些 conv2d、relu、matmul，在GPU眼里到底是什么？

说白了，算子（Operator）就是计算图里的一个基本操作单元。你可以把它理解成乐高积木里的一块小砖头。一个卷积层是一个算子，一个激活函数也是一个算子。整个深度学习模型，就是由这些算子拼装起来的。

我个人习惯把算子分成三类：

  • 计算密集型：比如卷积、矩阵乘法。这类算子大部分时间都在做数学运算。
  • 访存密集型：比如ReLU、加法、Reshape。这类算子大部分时间在读写数据，计算本身很快。
  • 通信密集型：比如AllReduce。多卡训练时才会遇到，数据搬来搬去很耗时。

嗯，这里要注意：不同类型的算子，优化策略完全不同。你拿优化卷积的思路去优化ReLU，大概率会翻车。我在项目中就见过有人把访存算子的循环展开搞得很复杂，结果性能反而下降了——因为GPU的显存带宽成了瓶颈，计算单元都在空等。

二、算子优化的目标：我们到底在追求什么？

你可能会想：优化算子，不就是让它跑得更快吗？

对，但不全对。我个人认为，算子优化的核心目标有三个，缺一不可：

1. 计算速度

这是最直观的。同样的模型，同样的batch size，你的算子跑一遍只要5毫秒，别人要10毫秒，这就是优势。但速度不是唯一的追求。

2. 显存占用

这个很多人容易忽略。我记得有一次做视频理解模型，模型结构没问题，计算量也不大，但一跑就OOM（显存溢出）。查了半天，发现是某个中间算子的临时缓冲区开得太大了。优化之后，显存占用从12GB降到了4GB，模型才跑起来。

所以，显存控制和计算速度同样重要。有时候为了省显存，我们甚至愿意牺牲一点速度——这叫“以空间换时间”的反面，“以时间换空间”。

3. 数值精度与稳定性

这个容易被新手忽略。你优化了半天，速度上去了，但模型精度掉了两个点，那等于白干。我曾经在优化一个LayerNorm算子时，为了省显存用了半精度，结果梯度爆炸了。后来发现是数值范围没处理好。

所以，优化的前提是保证精度不变。如果精度有损失，必须能解释清楚，并且控制在可接受范围内。

一句话总结算子优化的目标：在保证数值精度的前提下，尽可能降低计算延迟和显存占用。

三、性能度量指标：怎么量化优化效果？

光说“优化了”不行，你得拿出数据来。我常用的指标有三个，咱们一个一个说。

1. FLOPs（浮点运算次数）

FLOPs 是 Floating Point Operations 的缩写。它衡量的是计算量。比如一个矩阵乘法 C = A × B，如果A是m×n，B是n×p，那FLOPs大约是 2×m×n×p（乘法和加法各算一次）。

但注意：FLOPs 低不代表速度快。为什么？因为GPU的利用率也很重要。你想想看，一个算子FLOPs只有1G，但数据在显存和计算单元之间搬来搬去，搬了10次，那实际耗时可能比一个FLOPs是2G但数据只搬1次的算子还要长。

我刚开始做优化时，就犯过这个错。我费尽心思把一个算子的FLOPs降了30%，结果跑起来反而慢了。后来分析发现，我用的算法虽然计算量小了，但数据访问模式变得很不规则，导致缓存命中率暴跌。

小技巧：FLOPs 适合用来评估计算密集型算子的理论下限。但对于访存密集型算子，更该关注的是带宽利用率。

2. 吞吐量（Throughput）

吞吐量指的是单位时间内处理的数据量。常用单位是 samples/s（每秒处理多少个样本）或者 tokens/s（每秒处理多少个token）。

对于推理场景，吞吐量特别重要。比如你部署一个在线服务，每秒能处理1000个请求，那就是1000 QPS（Queries Per Second）。

吞吐量和延迟是什么关系？我举个例子：

  • 如果一次推理需要10ms，那理论上单线程的吞吐量是100 samples/s。
  • 但如果用batch处理，一次推理处理8个样本，耗时可能只增加到20ms，那吞吐量就变成了 8/0.02 = 400 samples/s。

看到了吗？吞吐量和延迟不是简单的反比关系。通过增大batch size，我们可以用更少的延迟代价换来更高的吞吐量。这就是算子优化里常用的“批处理”技巧。

3. 延迟（Latency）

延迟就是一次推理从输入到输出所花的时间。单位通常是毫秒（ms）或微秒（μs）。

对于实时性要求高的场景，比如自动驾驶、语音助手，延迟是硬指标。你不能说“我吞吐量很高，但每次推理要100ms”——那车都撞上了。

延迟的测量要注意两点：

  • 要测多次取平均。GPU的第一次运行通常有预热（warm-up）开销，所以一般会先跑几次“热身”，再正式计时。
  • 要区分P50、P99。P50是中位数延迟，P99是99%分位的延迟。对于在线服务，P99比P50更有参考价值——因为用户遇到的那次慢请求，可能就是P99里的。

避坑指南：我曾经在优化一个语音模型时，只看了平均延迟，觉得优化效果很好。上线后用户投诉说“偶尔卡顿”。一查，P99延迟比优化前还高了30%。所以，永远不要只看平均值。

四、三个指标的关系：一张图看懂

这三个指标不是孤立的。我画了一张图，帮你理清它们的关系：

算子优化核心指标关系图 FLOPs 计算量 理论下限 延迟 实时性 P50 / P99 吞吐量 效率 samples/s 优化目标 速度 + 显存 + 精度 FLOPs 低 ≠ 延迟低 延迟 × 并发 ≈ 吞吐量 注：三个指标需综合权衡，单一指标优化可能带来其他指标恶化

这张图想表达的是：FLOPs、延迟、吞吐量三者相互影响，但并非线性关系。你降低FLOPs，延迟不一定降；你提高吞吐量，延迟可能上升。真正的优化，是在这三个指标之间找到平衡点。

五、一个简单的例子：从代码看指标

光说理论太枯燥。咱们写一段简单的PyTorch代码，看看这些指标怎么算：

import torch
import time

# 定义一个简单的矩阵乘法算子
def matmul_benchmark(M, N, K, num_iter=100):
    A = torch.randn(M, K, device='cuda')
    B = torch.randn(K, N, device='cuda')
    
    # 预热
    for _ in range(10):
        C = torch.matmul(A, B)
    
    # 正式计时
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(num_iter):
        C = torch.matmul(A, B)
    torch.cuda.synchronize()
    end = time.time()
    
    avg_latency = (end - start) / num_iter * 1000  # 毫秒
    flops = 2 * M * N * K  # 乘法和加法
    throughput = 1.0 / (avg_latency / 1000)  # 每秒多少次
    
    print(f"矩阵大小: {M}x{K} × {K}x{N}")
    print(f"平均延迟: {avg_latency:.3f} ms")
    print(f"理论FLOPs: {flops/1e9:.2f} GFLOPs")
    print(f"吞吐量: {throughput:.0f} ops/s")

# 跑一下
matmul_benchmark(1024, 1024, 1024)

这段代码很简单，但包含了我们刚才讲的所有要素：预热、同步、多次取平均。你跑一下就会发现，理论FLOPs和实际延迟之间，差了一个GPU利用率。利用率高，FLOPs才能转化为实际的吞吐量。

我的习惯：每次优化完一个算子，我都会跑一遍这个基准测试，记录三个指标的变化。如果FLOPs降了但延迟没变，说明优化方向可能错了——瓶颈不在计算量，而在访存或调度。

六、总结：这一章你该记住什么？

好，咱们把这一章的核心点捋一捋：

  • 算子是计算图的基本单元，分计算密集型、访存密集型、通信密集型三类。
  • 优化目标是：在保证精度的前提下，降低延迟和显存占用。
  • 三个指标：FLOPs（理论计算量）、延迟（实际耗时）、吞吐量（处理效率）。它们相互关联，但需要综合权衡。

嗯，这一章就到这里。内容不多，但都是基础中的基础。后面的章节，我们会深入到具体的算子优化技巧，比如算子融合、内存复用、循环优化等等。到时候你会发现，今天讲的这些指标，每一个都会反复用到。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →