3. Metal性能优化原理

3.1 GPU 并行计算凭什么快？

Transformer 推理，说白了就是一堆矩阵乘法加注意力机制。CPU 擅长“串行干活”，但遇到同样操作重复几百万遍时就抓瞎了。GPU 正好相反——它有一堆小核心，每个核心运算能力不强，但架不住核心多啊。我记得刚接触 Metal 时，脑子还是 CPU 串行思维，写了个 for 循环让 GPU 逐个元素算……后来一看性能统计，直接傻眼，比 CPU 还慢。你想想看，让 GPU 干的活必须是“并行”的，而且数据要能拆成独立小块。

举个具体例子：一个 QKV 投影矩阵是 [4096, 4096] × [4096, 8192]，这个乘法可以拆成 4096 个独立行向量与整个矩阵的乘加。Metal 会把每一行分配给一个 threadgroup，每个 threadgroup 内部再做 tile 分块。这就是 GPU 并行计算的核心——数据并行 + 任务并行。我做的一个 7B 模型推理优化里，光是把矩阵乘法改成 Metal 并行实现，就比 CPU 版本快了 3 倍，但还没达到理论上限。

关键认知：GPU 不是“更快的 CPU”，而是“大量弱核的集合”。Transformer 推理中，90% 的计算量是矩阵乘法和注意力，天然适合 SIMD（单指令多数据）模式。Metal 的线程调度粒度是 thread（32 或 64 个 wave），理解这个才能写出高效的 shader。

3.2 Transformer 推理中 GPU 扮演什么角色？

我们拆开一个典型的 transformer 层看：

  • Self-attention 计算：Q、K、V 投影（矩阵乘法），然后 Q×K^T（又一个大矩阵乘法），Softmax，再乘 V。这几个步骤全是矩阵运算，GPU 最拿手。
  • 前馈网络 FFN：两个全连接层，中间加一个激活。还是矩阵乘法。
  • LayerNorm / RMSNorm：逐行 reduce 再 broadcast，GPU 可以用并行规约（parallel reduction）搞定。
  • 位置编码：RoPE 的旋转矩阵操作，也是按 并行的。

所以 GPU 就像一个高效的矩阵工厂，把 transformer 里这些“计算密集 + 高度并行”的步骤一股脑扔给它。但也别忘了——GPU 的显存带宽是瓶颈。记不记得我们第一章讲苹果统一内存？M 系列芯片的 GPU 和 CPU 共享内存，免去了 PCIe 拷贝，这是用 Metal 做推理的巨大优势。我踩过坑：在 Intel Mac 上写 OpenCL，每次 kernel 启动前都要从 CPU 拷到 GPU 显存，延迟生生吃掉一半性能。换成 M 系列后，直接用 shared memory 指针，零拷贝，爽多了。

操作 计算模式 Metal 并行策略
矩阵乘法（MatMul） 高度并行，数据复用 tile 分块 + 共享内存 + 向量化加载
Softmax（注意力） 行内 reduce，行列独立 threadgroup 内 warp-level reduce
LayerNorm 行内 reduce + broadcast 并行规约 + 共享内存缓存
RoPE 逐元素操作，无数据依赖 直接逐线程计算
SiLU / GELU 激活 逐元素 简单逐线程

记得：注意力层里 Q×K^T 的结果矩阵很大（batch × head × seq_len²），但 Softmax 是对每一行单独做的。这时候 GPU 的并行优势最明显——可以同时 softmax 所有行。我曾经手写一个 naive 版本，每行用一个 for 循环串行做 exp+sum，结果 GPU 利用率不到 20%。改成 threadgroup reduce 后，利用率直接拉到 90%。

3.3 Metal Shader 优化要点

好了，现在你知道 GPU 能干这活，但怎么写 Metal Shader 才能榨干 M 系列芯片？我总结几个关键点：

3.3.1 线程调度与工作组配置

Metal 的线程调度单位是 threadgroup（类似 CUDA 的 block），每个 threadgroup 内有 threads_per_threadgroup。M 系列 GPU 每个 threadgroup 最大支持 1024 个线程，但实际最好用 32 或 64 的倍数，因为 GPU 的 warp（wavefront）大小是 32（M1）或 64（某些后续芯片）。我用 M1 Pro 测试过，threadgroup 大小设为 64 比 96 提升了 15% 的利用率。为什么呢？因为 64 恰好是 wave 的整数倍，避免线程束部分空闲。

// 一个典型 threadgroup 配置示例
// dispatch = MTLSizeMake(numGroups, 1, 1);
// threadsPerGroup = MTLSizeMake(64, 1, 1);

kernel void matmul_kernel(
    device const float* A [[buffer(0)]],
    device const float* B [[buffer(1)]],
    device float* C [[buffer(2)]],
    constant uint& M [[buffer(3)]],
    constant uint& N [[buffer(4)]],
    constant uint& K [[buffer(5)]],
    uint2 gid [[thread_position_in_grid]]
) {
    // 确保不越界
    if (gid.x >= M || gid.y >= N) return;

    float sum = 0.0;
    for (uint k = 0; k < K; k++) {
        sum += A[gid.x * K + k] * B[k * N + gid.y];
    }
    C[gid.x * N + gid.y] = sum;
}

这个 naive 版本能跑，但很慢。因为每个线程都要从显存反复加载 A 和 B。我们接下来优化。

3.3.2 共享内存与 tile 分块

M 系列 GPU 每个 threadgroup 有 32KB ~ 64KB 的共享内存（threadgroup memory）。写矩阵乘法时，把 A 和 B 的 tile 加载到共享内存，然后让 threadgroup 内所有线程共享。这样每次加载一块到共享内存，计算这块内的所有结果，然后再加载下一块。这就是经典的 tiled matmul。我个人习惯 tile 大小取 16×16 或 32×32，配合 32 线程的 threadgroup 正好。

避坑指南：共享内存空间有限，tile 太大塞不下，太小又浪费带宽。我在 M1 Max 上做过扫描：tile=32×32 在矩阵维度 4096 时达到最好性能，带宽利用率约 85%；tile=16×16 时利用率只有 60%。如果你也优化 attention，注意 q×k^T 的 tile 受 seq_len 影响，不一定能选最优 tile。

3.3.3 向量化内存访问

Metal 支持 half2、float2、float4 等向量类型。用 float4 一次加载 16 字节，比单浮点加载快很多。M 系列芯片的内存控制器对连续 128-bit 访问有专门的优化。我记得有个项目，把输入从 float 改成 packed_float4（对齐到 16 字节），kernel 运行时间直接砍掉 30%。但要小心：向量化要求数据对齐和连续性，乱序时反而更慢。

// 向量化加载示例：用 float4 一次读 4 个 float
kernel void vec_matmul(
    device const float4* A [[buffer(0)]],
    device const float4* B [[buffer(1)]],
    device float* C [[buffer(2)]],
    ...
) {
    // 假设 A 和 B 按 float4 对齐
    float4 a = A[gid.x * (K/4) + tileId];
    float4 b = B[tileId * (K/4) + gid.y];
    // 做点积
    sum += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
}

3.3.4 减少分支与线程束分歧

Metal shader 里 if-else 会导致同一 wavefront 里的线程走不同分支，形成 divergence。轻则浪费周期，重则让并行度下降 50% 以上。在 transformer 推理里，比如 attention mask、padding 的处理，容易引发分支。我的经验是：把 batch 处理写成“掩码乘法”而非条件判断。比如 masked softmax，用一个很大的负值乘到需要 mask 的位置上，这样既不用分支，结果也是一样的。

3.3.5 利用 Metal Performance Shaders (MPS)

其实你未必需要自己从头写 shader。Apple 提供了 MPS 框架，内有优化的矩阵乘法（MPSMatrixMultiplication）和卷积操作。我自己做实验时，用 MPSMatrixMultiplication 跑 4096×4096 的矩阵乘法，性能比手写 naive 版本好了 5 倍，比自己写 tiled 版本好了 2 倍。但 MPS 也不是万能的——它对非标准形状（比如不对称矩阵、小矩阵）的优化有限。我建议：标准尺寸的大矩阵用 MPS，特殊尺寸（比如注意力中的 small K）自己写定制 shader。

警告：MPS 内部可能依赖 Metal 运行时的版本。macOS 13.0 以下对某些 MPS 操作支持不完整。如果你部署到旧系统，最好后备一个手写 shader。我遇到过客户在 macOS 12.6 上跑 MPS 矩阵乘法直接闪退，排查两小时才发现是系统版本问题。

3.4 GPU 并行在 Transformer 中的瓶颈与应对

你以为只要堆并行度和带宽就能无限提速？不是。Transformer 推理里有个麻烦——注意力计算的 内存带宽瓶颈。Q×K^T 的结果需要写回显存，然后 Softmax 又要重新读回来，再乘 V，产生大量 IO。在 M 系列芯片上，内存带宽虽然高（M1 Max 约 400 GB/s，M2 Ultra 约 800 GB/s），但跟计算能力一比，还是容易成为瓶颈。

怎么破？两种思路：

  1. 算子融合：把 Q×K^T、Softmax、×V 三个 kernel 合并成一个。计算完 Q×K^T 的 tile 后，直接把 Softmax 结果存到共享内存，然后立即跟 V 的 tile 相乘，最后累加写回。这样大幅减少显存读写。
  2. Flash Attention 理念：用分块算法避免整个注意力矩阵的显存分配。Metal 里可以写一个多级 tiling 的 kernel，让注意力的计算完全在共享内存和寄存器里完成，只把最终结果写回去。我按照 flash attention 论文在 M1 Max 上实现了一遍，Batch=1, seq_len=4096 时，推理延迟比原始 naive 版本降低了约 40%。

下面这张图帮你理清 GPU 并行计算的完整链路：

GPU 并行计算在 Transformer 推理中的核心链路 输入 Tokens + Weights 数据分块 / tile Threadgroup 调度 共享内存 MatMul / Attention 计算 算子融合 / Flash Attention 结果写回统一内存 优化关键 • 向量化加载（float4 / half2） • 共享内存 tile 复用 • 避免线程束分歧

3.5 我的一些实践经验总结

讲真，Metal 优化是个不断迭代的过程。我刚开始做 M1 优化时，习惯把 CUDA 那套直接搬过来，结果发现 Metal 的线程模型、内存模型和编译器行为跟 CUDA 有细微差别。比如 Metal 的 threadgroup memory 不能用 dynamic indexing 随意访问，否则编译器会生成很差的代码。后来我改成用 static tile 地址计算，性能就上来了。

另外，M 系列每个芯片的 GPU 核心数和内存带宽都不同。M1 8 核 GPU 和 M2 Ultra 76 核 GPU，优化策略不能一刀切。我建议你做 profiling 时固定一个目标芯片，比如 M1 Pro 或 M2 Max，先把这个型号榨到极致，再考虑向下兼容。Metal 的 GPU 计数器可以看 occupancy、memory 带宽利用率、ALU 利用率，用这些数据指导优化，比盲目猜要靠谱得多。

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