跳到主要内容

GEMM/GEMV 优化

GEMM/GEMV 优化核心原则
  1. GEMV（矩阵向量乘法）是典型的内存瓶颈负载，优化重点是带宽利用率
  2. GEMM（矩阵矩阵乘法）是计算瓶颈负载，优化重点是计算单元利用率
  3. 分块（Tiling）是核心策略，平衡容量与计算密度
  4. 充分利用张量核心（Tensor Core）等专用硬件单元

阅读本文时，建议先判断算子属于 GEMV 还是 GEMM，再选择对应优化路径。GEMV 优先检查访存是否合并；GEMM 优先检查分块复用、寄存器累加和张量核心使用。

场景主要瓶颈优先检查项
GEMV内存带宽合并访存、向量化加载、向量缓存
GEMM计算吞吐分块复用、寄存器累加、张量核心

基础示例：从基础实现（Naive）到共享内存优化​

下面的示例先给出直接实现，再给出共享内存分块实现。两者计算同一个矩阵乘法，差异在于是否把 A/B 的局部数据缓存到共享内存中复用。

基础实现（Naive，无共享内存）​

每个线程直接从全局内存读取数据，计算 C 矩阵的一个元素：

// 矩阵存储结构
typedef struct {
    int width;
    int height;
    float* elements;
} Matrix;

#define BLOCK_SIZE 16

__global__ void MatMulNaive(Matrix A, Matrix B, Matrix C)
{
    // 每个线程计算 C 的一个元素
    float Cvalue = 0;
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < C.height && col < C.width) {
        for (int e = 0; e < A.width; ++e) {
            Cvalue += A.elements[row * A.width + e] * B.elements[e * B.width + col];
        }
        C.elements[row * C.width + col] = Cvalue;
    }
}

// 启动配置
dim3 dimBlock(BLOCK_SIZE, BLOCK_SIZE);
dim3 dimGrid((B.width + dimBlock.x - 1) / dimBlock.x,
             (A.height + dimBlock.y - 1) / dimBlock.y);
MatMulNaive<<<dimGrid, dimBlock>>>(A, B, C);
查看问题分析

问题分析：

矩阵全局内存读取次数原因
AB.width 次每个线程独立读取 A 的行元素
BA.height 次每个线程独立读取 B 的列元素

示例：对于 1024×1024 矩阵，BLOCK_SIZE=16

  • 全局内存读取：约 2 × 1024 × 1024 × 16 = 33M 次
  • 带宽利用率极低

访存趋势对比​

对于 1024×1024 矩阵乘法：

实现全局内存读取趋势说明
基础实现重复读取 A/B 元素实现简单，但全局内存访问压力高
共享内存分块每个块内复用 A/B 元素减少重复读取，实际收益需结合矩阵形状和硬件实测

GEMV 优化（矩阵×向量）​

GEMV 定义与特征​

GEMV（矩阵向量乘法）常出现在小批量推理、归约后的线性变换、向量投影等场景。它的计算量相对有限，但需要持续读取矩阵和向量数据，因此通常更受内存带宽限制，而不是受计算单元峰值能力限制。

GEMV（General Matrix-Vector Multiplication） 操作：

y=α⋅A⋅x+β⋅by = \alpha \cdot A \cdot x + \beta \cdot b

其中：

  • AA 是 m×km \times k 矩阵
  • xx 是 kk 维向量
  • yy 是 mm 维输出向量

算术强度（以 float 为例）：

Arithmetic Intensity=2mk FLOPs(mk+k+m)×4 Bytes≈0.5 FLOP/Byte\text{Arithmetic Intensity} = \frac{2mk \text{ FLOPs}}{(mk + k + m) \times 4 \text{ Bytes}} \approx 0.5 \text{ FLOP/Byte}

特征：

  • 算存比远低于 1 → 典型的内存瓶颈负载
  • 性能受限于内存带宽，而非计算能力

GEMV 逐步优化策略​

GEMV 优化应按顺序推进：先保证合并访存，再做向量化加载，最后根据向量复用情况考虑共享内存缓存。下面三个标签页展示同一优化路径的不同阶段。

阶段解决的问题适用判断
基础并行化建立行级并行需要先把单线程点积改为线程束协作
向量化加载减少访存指令压力已满足连续地址访问和对齐要求
缓存向量 x降低 x 的重复读取同一线程块内多个输出行复用同一段 x

步骤 1: 基础并行化（线程束级）​

最直接的并行方式是让一个线程束协作计算一个输出元素。线程束内每个线程负责一部分 K 维度，最后通过线程束内归约得到该行的点积结果。

__device__ float warpReduceSum(float value) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        value += __shfl_down_sync(0xffffffff, value, offset);
    }
    return value;
}

// 一个线程束计算一个输出元素，每个线程负责部分 K 维度
__global__ void gemvStep1(float* A, float* x, float* y, int m, int k) {
    int row = blockIdx.x;
    if (row >= m) return;

    int lane = threadIdx.x % 32;
    float partial = 0.0f;

    for (int col = lane; col < k; col += 32) {
        partial += A[row * k + col] * x[col];
    }

    // 线程束内归约，得到该行输出
    partial = warpReduceSum(partial);

    if (lane == 0) {
        y[row] = partial;
    }
}

// 启动配置
int threadsPerBlock = 32;
int blocksPerGrid = m;
gemvStep1<<<blocksPerGrid, threadsPerBlock>>>(A, x, y, m, k);
查看该阶段的限制

问题：

  • 每个线程每次只加载 32-bit 数据，访存指令数量较多
  • x 会被不同行重复读取，缓存和外部内存带宽压力较高
  • 当 kk 较大时，可能需要多个线程束协作计算一个输出元素，并通过共享内存完成跨线程束归约

GEMV 优化检查清单​

  • 访存是否合并？

    • 一个线程束协作计算一行，线程块内可包含多个线程束
    • 同一线程束内相邻线程访问矩阵 A 的连续列，避免跨步访问

  • 是否使用向量化加载？

    • 在合并访存基础上使用 128-bit 或更宽的向量化加载
    • 确保数据 16B 对齐

  • 是否缓存向量 x？

    • 使用共享内存分段缓存
    • TILE_V 根据容量调整

GEMM 优化（矩阵×矩阵）​

GEMM 定义与特征​

GEMM（通用矩阵乘法）是深度学习、科学计算和图形计算中的核心算子。在 LLM、卷积、全连接层等场景中，GEMM 通常占据大量计算时间，因此 GEMM 优化是获得端到端性能收益的关键路径。

GEMM（General Matrix Multiplication） 操作：

C=A×BC = A \times B

其中：

  • AA 是 m×km \times k 矩阵
  • BB 是 k×nk \times n 矩阵
  • CC 是 m×nm \times n 矩阵

通常把输入矩阵记为 A/B，把输出矩阵记为 C。优化 GEMM 时，需要同时关注计算密度、片上存储容量、共享内存带宽和寄存器占用。

算术强度（以 float 为例）：

Arithmetic Intensity=2mnk FLOPs(mn+mk+nk)×4 Bytes\text{Arithmetic Intensity} = \frac{2mnk \text{ FLOPs}}{(mn + mk + nk) \times 4 \text{ Bytes}}

当 m,n,km, n, k 较大时，算术强度远高于 1 → 计算瓶颈负载

GEMM 分块策略​

分块不会减少总计算量，但可以把局部数据保存在共享内存和寄存器中，让同一批 A/B 数据被多次复用。这样可以提高计算访存比，避免计算单元长时间等待全局内存数据。

分块的核心思想：

  • 将大矩阵分解为小块
  • 每块数据加载到高速存储器（共享内存）
  • 在高速存储器中重复使用数据

分块参数选择：

参数含义典型值限制因素
TILE_MC 矩阵 M 维度分块64-256共享内存容量
TILE_NC 矩阵 N 维度分块64-256共享内存容量
TILE_KK 维度分块8-32共享内存容量
THREAD_M每线程计算 M 维度8-16寄存器数量
THREAD_N每线程计算 N 维度8-16寄存器数量

经典分块实现​

// 分块 GEMM 实现：每个线程计算一个 C 元素
template<int TILE>
__global__ void gemmTiled(float* A, float* B, float* C, int M, int N, int K) {
    __shared__ float As[TILE][TILE];
    __shared__ float Bs[TILE][TILE];

    int tx = threadIdx.x;
    int ty = threadIdx.y;

    int row = blockIdx.y * TILE + ty;
    int col = blockIdx.x * TILE + tx;

    float accum = 0.0f;

    // 分块迭代
    for (int t = 0; t < (K + TILE - 1) / TILE; t++) {
        // 协作加载 A 块到共享内存（合并访问）
        int a_col = t * TILE + tx;
        if (row < M && a_col < K) {
            As[ty][tx] = A[row * K + a_col];
        } else {
            As[ty][tx] = 0.0f;
        }

        // 协作加载 B 块到共享内存（合并访问）
        int b_row = t * TILE + ty;
        if (b_row < K && col < N) {
            Bs[ty][tx] = B[b_row * N + col];
        } else {
            Bs[ty][tx] = 0.0f;
        }
        __syncthreads();

        // 在共享内存中计算
        #pragma unroll
        for (int k = 0; k < TILE; k++) {
            accum += As[ty][k] * Bs[k][tx];
        }
        __syncthreads();
    }

    // 写回结果
    if (row < M && col < N) {
        C[row * N + col] = accum;
    }
}

// 启动配置
dim3 blockDim(16, 16);  // 256 线程/Block
dim3 gridDim((N + 15) / 16, (M + 15) / 16);
gemmTiled<16><<<gridDim, blockDim>>>(A, B, C, M, N, K);

优化要点​

1. 合并访存：

// A 矩阵加载：行优先，threadIdx.x 映射到连续列
As[ty][tx] = A[c_row * K + (t * TILE_K + tx)];

// B 矩阵加载：列优先，threadIdx.y 映射到连续行
Bs[ty][tx] = B[(t * TILE_K + ty) * N + c_col];

2. 寄存器优化：

// 每线程计算 8×8 子块，64 个累加器存储在寄存器
float accum[8][8];

// 优势：
// - 减少共享内存访问
// - 提高计算/访存比
// - 充分利用寄存器带宽

3. 循环展开：

// #pragma unroll 自动展开内层循环
// 增加指令级并行
// 隐藏访存延迟

张量核心（Tensor Core）优化​

张量核心（Tensor Core）简介​

张量核心（Tensor Core）是专为矩阵乘法加速的硬件单元：

  • 单指令完成 D=A×B+CD = A \times B + C
  • A、B：16 位浮点（FP16/BF16）
  • C、D：16 位或 32 位浮点
  • MT GPU S5000 及后续架构支持

MT GPU 张量核心特性：

  • 支持 FP16、BF16、INT8 精度
  • 每指令 256 FLOPs（4×4×4 矩阵乘）
  • 需对齐到 16 字节边界

张量核心（Tensor Core）GEMM 实现​

推荐：使用 MUTLASS 库

对于生产环境，建议使用 MUTLASS（MUSA Tensor Core Library，MUSA 张量核心库）库来实现张量核心（Tensor Core）加速的 GEMM。MUTLASS 提供了经过深度优化的张量核心（Tensor Core）实现，支持 FP16/BF16/TF32/FP8 等多种精度。

详见：MUTLASS 库文档

MUTLASS 主循环（Mainloop）示例​

下面两个标签页展示 MUTLASS 两阶段主循环的核心片段。第一个片段强调基础拷贝流程，第二个片段展示边界谓词拷贝。两者都省略了张量视图、迭代器、谓词和类型定义等上下文，不能作为独立 kernel 直接编译。

该片段用于说明 MUTLASS_PRAGMA_NO_UNROLL 如何配合静态循环使用：

#include <mutlass/mutlass.h>
#include <mute/algorithm/gemm.hpp>

// MUTLASS mainloop 两阶段实现
template <class... Args>
__global__ void gemm_mainloop_kernel(Args... args) {
  // ... 初始化代码省略 ...

  // K 维度上的外积（outer product）数量
  auto K_BLOCK_MAX = size<2>(tCrA);

  // 使用 MUTLASS_PRAGMA_NO_UNROLL 禁用循环展开
  // 静态 for 循环需要手动控制展开行为
  MUTLASS_PRAGMA_NO_UNROLL
  while (k_tile_count > -1)
  {
    // 使用 for_each 进行静态循环流水线化
    for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block)
    {
      if (k_block == K_BLOCK_MAX - 1)
      {
        __syncthreads();

        // 将寄存器数据写回共享内存（双缓冲）
        copy(tArA, tAsA);
        copy(tBrB, tBsB);
        __syncthreads();
      }

      // 预加载下一个 K 块的数据（共享内存 smem → 寄存器内存 rmem）
      int k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX;  // 静态取模
      copy(tCsA(_,_,k_block_next), tCrA_copy_view(_,_,k_block_next));
      copy(tCsB(_,_,k_block_next), tCrB_copy_view(_,_,k_block_next));

      if (k_block == 0)
      {
        // 从全局内存（gmem）加载下一个 K 块（全局内存 gmem → 寄存器内存 rmem）
        copy(gmem_tiled_copy_a, tAgA(_,_,_,*k_tile_iter), tArA);
        copy(gmem_tiled_copy_b, tBgB(_,_,_,*k_tile_iter), tBrB);
        ++k_tile_iter;
        --k_tile_count;
      }

      // 计算前变换数据
      mute::transform(tCrA(_,_,k_block), TransformA{});
      mute::transform(tCrB(_,_,k_block), TransformB{});

      // 线程级寄存器矩阵乘累加
      mute::gemm(tiled_mma, accum, tCrA(_,_,k_block), tCrB(_,_,k_block), src_accum);
    });
  }
}

关键实现要点​

要点说明
双缓冲（Double Buffering）使用 K_BLOCK_MAX 块进行数据流水线，避免等待
静态循环使用 for_each(make_int_sequence<K_BLOCK_MAX>{}) 替代普通 for 循环
MUTLASS_PRAGMA_NO_UNROLL禁用 while 循环展开，保持静态循环的确定性
寄存器复用数据在共享内存（smem）→ 寄存器内存（rmem）→ 寄存器中流动，最大化复用

张量核心（Tensor Core）优化要点​

优化项要求理由
精度选择FP16/BF16张量核心（Tensor Core）原生支持
对齐要求16 字节边界避免非对齐访问开销
分块大小16 的倍数匹配张量核心（Tensor Core）指令形状
累加器精度FP32保持数值稳定性

性能收益趋势​

GEMV 优化收益趋势​

优化阶段主要收益说明
基础并行建立行级并行一个线程束或多个线程束协作计算输出元素
+ 合并访存/向量化提高全局内存带宽利用率连续线程访问连续地址，向量化减少访存指令数量
+ 共享内存缓存降低向量重复读取压力将向量分段缓存到共享内存，在多个线程束间复用

GEMM 优化收益趋势​

优化阶段主要收益说明
基础实现（Naive）实现简单计算访存比较低，难以充分利用计算单元
共享内存分块提升数据复用将 A/B 分块缓存到共享内存，降低全局内存访问压力
+ 寄存器优化提升计算密度每线程计算多个输出元素，累加器存放在寄存器中
+ 张量核心（Tensor Core）提升矩阵乘吞吐适用于 FP16/BF16/INT8 等张量核心支持的数据类型

实际性能收益与矩阵形状、数据类型、硬件架构、访存对齐和 kernel 实现有关，应以目标环境实测为准。

性能优化检查清单​

GEMV 优化​

  • 访存是否合并？

    • 一个线程束协作计算一行，线程块内可包含多个线程束
    • 同一线程束内相邻线程访问矩阵 A 的连续列，避免跨步访问

  • 是否使用向量化加载？

    • 在合并访存基础上使用 128-bit 或更宽的向量化加载
    • 确保 16B 对齐

  • 是否缓存向量？

    • 使用共享内存分段缓存
    • TILE_V 根据容量调整（512-2048）

GEMM 优化​

  • 分块参数是否合理？

    • TILE_M, TILE_N, TILE_K 平衡容量与计算密度
    • 典型值：128×64×32

  • 是否使用寄存器优化？

    • 每线程计算 8×8 或 16×16 子块
    • 累加器存储在寄存器

  • 访存是否合并？

    • A 矩阵：行优先，threadIdx.x 映射到列
    • B 矩阵：列优先，threadIdx.y 映射到行

  • 是否利用张量核心（Tensor Core）

    • FP16/BF16 精度
    • 分块大小为 16 的倍数

常见问题​

Q1：GEMV 和 GEMM 哪个更难优化？

答：GEMV 主要受限于内存带宽，优化重点是提高带宽利用率；GEMM 可以通过分块、寄存器累加和张量核心（Tensor Core）提高计算吞吐，通常具备更大的优化空间。实际收益应以目标硬件和矩阵形状实测为准。

Q2：如何选择分块大小？

答：

  1. 从典型值开始（128×64×32）
  2. 使用 Moore Perf Compute 分析 占用率 和共享内存使用
  3. 微调参数，找到最优组合

相关文档​

版权所有 2024 摩尔线程，本文档受国际版权法保护。

摩尔线程和摩尔线程徽标是摩尔线程智能科技 (北京)有限责任公司的注册商标。

免责声明

本文档提供有关摩尔线程产品的信息。本文档并未授权任何知识产权的许可，并未以明示或暗示，或以禁止反言或其他方式授予任何知识产权许可。