前兩天我們一直在優化 GEMM：

Naive GEMM
    ↓
Shared Memory Tiling
    ↓
Register / Thread Tiling

做的事情基本上都是同一件事：增加 Data Reuse，減少資料搬移。

但就算資料搬得很好，最後矩陣乘法還是要真的算。

而 NVIDIA GPU 上有一組特別為矩陣運算設計的硬體：

Tensor Core

CUDA Core 和 Tensor Core 差在哪？

前面的 GEMM kernel，本質上一直在做很多次：

a * b + c

也就是 FMA（Fused Multiply-Add）。

可以先很粗略地理解成：

CUDA Core

a × b + c

Tensor Core 則是把這件事情放大成：

Tensor Core

D = A × B + C

這裡的 A、B、C 已經不是單一數字，而是一小塊 Matrix。

所以 Tensor Core 並不是單純讓「一個乘法」變快，而是：

直接用專門的硬體處理 Matrix Multiply-Accumulate。

這類操作通常就叫做 MMA（Matrix Multiply-Accumulate）。

一個 Warp 一起算 Matrix Tile

Tensor Core 的使用方式和之前「一個 Thread 算自己的東西」也不太一樣。

以 CUDA 的 WMMA API 為例，一個 Warp，也就是 32 個 Threads，會一起合作處理一個 Matrix Tile。

例如常見的：

A：16 × 16
B：16 × 16
C：16 × 16

概念上就是：

        16
     ┌──────┐
  16 │  A   │
     └──────┘
         ×
     ┌──────┐
     │  B   │
     └──────┘
         ↓
     ┌──────┐
     │  C   │
     └──────┘

      1 Warp
    共同完成 MMA

【圖 1：一個 Warp 使用 Tensor Core 計算 Matrix Tile】

和昨天相比：

昨天：

Thread 0 → 一小塊 C
Thread 1 → 一小塊 C
Thread 2 → 一小塊 C
...

今天則是：

一整個 Warp
    ↓
共同操作 Matrix Fragments
    ↓
Tensor Core
    ↓
完成 Matrix Multiply-Accumulate

CUDA 裡怎麼用 Tensor Core？

CUDA 提供了 WMMA（Warp Matrix Multiply-Accumulate）API。

一個最簡化的例子大概長這樣：

#include <cuda_fp16.h>
#include <mma.h>

using namespace nvcuda;

__global__ void tensor_core_gemm(
    const half* A,
    const half* B,
    float* C
) {
    wmma::fragment<
        wmma::matrix_a,
        16, 16, 16,
        half,
        wmma::row_major
    > a_frag;

    wmma::fragment<
        wmma::matrix_b,
        16, 16, 16,
        half,
        wmma::col_major
    > b_frag;

    wmma::fragment<
        wmma::accumulator,
        16, 16, 16,
        float
    > c_frag;

    wmma::fill_fragment(c_frag, 0.0f);

    wmma::load_matrix_sync(a_frag, A, 16);
    wmma::load_matrix_sync(b_frag, B, 16);

    wmma::mma_sync(
        c_frag,
        a_frag,
        b_frag,
        c_frag
    );

    wmma::store_matrix_sync(
        C,
        c_frag,
        16,
        wmma::mem_row_major
    );
}

裡面最重要的是：

wmma::mma_sync(
    c_frag,
    a_frag,
    b_frag,
    c_frag
);

也就是：

C = A × B + C

只是這次是一整個 Warp 共同操作 Matrix fragments，底層可以使用 Tensor Core 來完成計算。

為什麼 Tensor Core 對 LLM 這麼重要？

因為 LLM 裡到處都是 GEMM。

例如：

Hidden States
    ↓
Q / K / V Projection
    ↓
Attention
    ↓
Output Projection
    ↓
MLP

其中很多 Linear Layer 最後都會變成大型矩陣乘法。

所以我們前幾天做的事情可以分成兩個方向：

讓資料餵得更有效率

Global Memory
    ↓
Shared Memory
    ↓
Register

靠 Tiling 和 Data Reuse 減少 Memory Traffic。

讓矩陣乘法本身算得更快

Matrix Tile
    ↓
MMA
    ↓
Tensor Core

兩個其實缺一不可。

Tensor Core 很會算，但如果資料餵不進來，一樣只能等。

這也是為什麼真正高效能的 GEMM kernel，不只是「用了 Tensor Core」就結束了。

那 H100 呢？

這篇先用 WMMA 理解最基本的 Tensor Core 概念。

但我們現在用的 LeetGPU 是 H100（Hopper）。

Hopper 上還有更進一步的 WGMMA（Warpgroup Matrix Multiply-Accumulate）：不再只是一個 Warp，而是一個 Warpgroup，也就是 4 個 Warps / 128 Threads 合作進行更大的 Matrix Multiply-Accumulate。

而且 Hopper 還有另一個之後很重要的東西：

TMA（Tensor Memory Accelerator）

所以接下來會開始碰到另一個問題：

Tensor Core 已經算得這麼快了，要怎麼確保資料來得及送進去？

也就是：

搬下一批資料
        ↘
          同時進行
        ↗
計算這一批資料

下一篇來看看：

TMA / Async Pipeline：GPU 可以一邊搬資料，一邊算嗎？

Reference

• NVIDIA CUDA Programming Guide — WMMA
• NVIDIA CUTLASS — Warpgroup MMA Programming Guide