LeetGPU | 02_Matrix_Multiplication

12617 字
63 分钟
LeetGPU | 02_Matrix_Multiplication
2026-07-25
LeetGPU
AI Infra
/
Kernal
/
CUDA
/
Triton
/
Pytorch

关注 LeetGPU repository，了解更多内容。

Problem Description#

Difficulty: Easy

LeetGPU Challenge

Write a program that multiplies two matrices of 32-bit floating point numbers on a GPU. Given matrix AA of dimensions M×NM \times N and matrix BB of dimensions N×KN \times K, compute the product matrix C=A×BC = A \times B, which will have dimensions M×KM \times K. All matrices are stored in row-major format.

Implementation Requirements#

  • Use only native features (external libraries are not permitted).
  • The solve function signature must remain unchanged.
  • The final result must be stored in matrix C.

Example 1#

Input:

Matrix AA (2×22 \times 2):

[1.02.03.04.0]\begin{bmatrix} 1.0 & 2.0 \\ 3.0 & 4.0 \end{bmatrix}

Matrix BB (2×22 \times 2):

[5.06.07.08.0]\begin{bmatrix} 5.0 & 6.0 \\ 7.0 & 8.0 \end{bmatrix}

Output:

Matrix CC (2×22 \times 2):

[19.022.043.050.0]\begin{bmatrix} 19.0 & 22.0 \\ 43.0 & 50.0 \end{bmatrix}

Example 2#

Input:

Matrix AA (1×31 \times 3):

[1.02.03.0]\begin{bmatrix} 1.0 & 2.0 & 3.0 \end{bmatrix}

Matrix BB (3×13 \times 1):

[4.05.06.0]\begin{bmatrix} 4.0 \\ 5.0 \\ 6.0 \end{bmatrix}

Output:

Matrix CC (1×11 \times 1):

[32.0]\begin{bmatrix} 32.0 \end{bmatrix}

Constraints#

  • 1 <= M, N, K <= 8192.
  • Performance is measured with M = 8192, N = 6144, K = 4096.

CUDA#

Approach#

本题需要计算 A[M, N] 与 B[N, K] 的乘积，并将结果写入 C[M, K]。每个输出元素都是 A 的一行与 B 的一列所形成的内积（inner product）：

C[row,col]=∑inner=0N−1A[row,inner]×B[inner,col].C[\text{row},\text{col}] = \sum_{\text{inner}=0}^{N-1} A[\text{row},\text{inner}] \times B[\text{inner},\text{col}].

由于三个矩阵均采用行主序（row-major）存储，它们的二维坐标可分别转换为以下线性地址：

A[row, inner] -> A[row * N + inner]
B[inner, col] -> B[inner * K + col]
C[row, col]   -> C[row * K + col]

不同 (row, col) 坐标对应的输出彼此独立，因此 CUDA 可以并行处理输出矩阵的两个维度；沿 N 维度进行的归约（reduction），则仍由单个线程（thread）或该线程负责的输出分块（tile）完成。

传统 CPU 实现通常使用三重嵌套循环：外侧两个循环遍历 M×KM \times K 个输出位置，最内层循环完成长度为 NN 的归约，整体算术复杂度为 O(MNK)O(MNK)。GPU 优化并不会减少这些必要的算术运算，而是重新安排循环的执行位置、同时计算的输出数量，以及数据从各级内存中被读取的频率。

矩阵乘法的内积视角
矩阵乘法的内积视角

一个输出元素由 A 的一行和 B 的一列计算得到。

这些 CUDA 文件展示了一条逐步演进的优化路径。所有版本都保持相同的 solve(...) 接口和数学结果，主要区别在于每个 thread 负责多少输出，以及数据如何在全局内存（global memory）、共享内存（shared memory）与寄存器（register）之间流动。

整个优化过程围绕两个相互关联的目标展开：

  1. 提高复用率。 从 global memory 读取的数据，应在被替换前尽可能参与多个输出元素的计算。共享内存分块（shared-memory tiling）可以在同一个 block 内复用数据，线程分块（thread tiling）则能进一步在 register 中复用数据。
  2. 让数据移动与计算重叠，或提高二者的执行效率。 向量化传输可以减少 load/store 指令开销；软件流水线和硬件流水线可以让相邻 reduction tile 的数据移动与计算交叠；如果能够接受较低的 TF32 输入精度，还可以使用 Tensor Core 取代标量 FP32 乘加指令。
VersionMain changeWork owned by one thread or warp
V0Direct global-memory dot productsOne output element per thread
V1 sharedCooperative shared-memory tilingOne output element per thread
V1 1DRegister reuse along the row dimension16 x 1 outputs per thread
V1 2DRegister-level outer products8 x 4 outputs per thread
V2float4 cooperative transfers8 x 4 outputs per thread
V3Register prefetch plus two shared-memory buffers8 x 4 outputs per thread
V4Hardware-asynchronous cp.async transfers8 x 4 outputs per thread
V5WMMA TF32 Tensor Core operationsOne 16 x 16 output tile per warp

V0: Naive Global-Memory Kernel#

Approach#

V0 启动一个二维网格（grid），其中每个线程块（block）包含 16 x 16 个 thread。每个 thread 按如下方式确定自己负责的输出坐标：

row = blockIdx.y * blockDim.y + threadIdx.y
col = blockIdx.x * blockDim.x + threadIdx.x

grid 在两个输出维度上都使用向上取整除法计算大小。每个 thread 首先检查 row < M && col < K；只有坐标有效的 thread 才会计算完整的点积，并将结果写入 C 中对应的一个元素。

这种映射用二维 CUDA grid 取代了 CPU 实现最外侧的两个循环，但每个 thread 仍需串行完成整个内部归约。该实现结构直接，因而也很容易验证正确性。

V0 的主要弱点在于数据移动。分别计算 C[row, col0] 和 C[row, col1] 的两个 thread 都会遍历 A 的同一行，却各自重新加载这行数据；同样，计算同一输出列中不同行元素的 thread，也会反复读取 B 同一列中的值。硬件缓存（cache）或许能够捕获一部分复用机会，但该 kernel 并未显式地把这些数据保留在片上。因此，V0 虽然暴露了充足的输出并行性，却没有改善底层点积计算的数据局部性。

Solution#

#include <cuda_runtime.h>


__global__ void matrix_multiplication_v0(const float* A, const float* B, float* C,
                                         int M, int N, int K) {
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M || col >= K) return;


    float sum = 0.0f;
    for (int inner = 0; inner < N; ++inner) {
        sum += A[row * N + inner] * B[inner * K + col];
    }
    C[row * K + col] = sum;
}


extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
    const dim3 block(16, 16);
    const dim3 grid((K + block.x - 1) / block.x, (M + block.y - 1) / block.y);
    matrix_multiplication_v0<<<grid, block>>>(A, B, C, M, N, K);
}

V1: Original 16 x 16 Shared-Memory Tiling#

Approach#

这个原始版本仍让每个 thread 负责一个输出元素，但把归约维度划分为长度为 16 的 tile。每个 16 x 16 block 负责生成 C 中的一个 16 x 16 tile，从而形成一条明确的数据流动路径：

global memory -> shared memory -> thread registers -> C

global memory 保存完整矩阵，shared memory 充当 block 内部的暂存区，用于保存当前参与计算的一对输入 tile；每个 thread 则在 register 中维护自己的标量累加器。

处理 reduction tile tile 时，thread (ty, tx) 负责加载：

A[row, tile * 16 + tx] -> A_shared[ty][tx]
B[tile * 16 + ty, col] -> B_shared[ty][tx]

全部 256 个 thread 完成加载后，每个 A 值都可由负责不同输出列的 16 个 thread 共享，每个 B 值也可由负责不同输出行的 16 个 thread 共享。随后，每个 thread 执行 16 次乘加运算：

S += A_shared[ty][inner] * B_shared[inner][tx]

这样一来，每个从 global memory 读取一次的输入值，在同一个 block 内最多可以参与 16 次乘加运算。对于一个完整的 reduction tile，block 分别从两个输入矩阵加载 16×1616 \times 16 个值，随后利用这些数据完成 16×16×1616 \times 16 \times 16 次乘加运算。其算术计算与 V0 完全相同，但大量重复的 global-memory 流量被延迟更低的 shared-memory 访问取代。

每个 reduction tile 都需要两次同步屏障（barrier）。第一次 barrier 保证协作式加载（cooperative load）全部完成后才开始计算；第二次 barrier 则保证所有 thread 都用完当前 tile 后，下一个 tile 才能覆盖 shared memory 中的现有数据。

最后一个 reduction tile 可能不足 16 个元素。遇到无效的输入坐标时，程序不会提前返回，而是向 shared memory 写入零。零填充（zero padding）不会改变点积结果，同时能让所有 thread 都到达相同的 barrier，避免在同步位置发生分支发散。最终写回输出时，再通过 row < M && col < K 单独进行边界保护。

Solution#

#include <cuda_runtime.h>


#define BLOCK_SIZE  16
__global__ void matrix_multiplication_kernel(const float* A, const float* B, float* C, int M, int N,
                                             int K) {
    __shared__ float A_shared[BLOCK_SIZE][BLOCK_SIZE];
    __shared__ float B_shared[BLOCK_SIZE][BLOCK_SIZE];
    int row = threadIdx.y + blockIdx.y * blockDim.y;
    int col = threadIdx.x + blockIdx.x * blockDim.x;
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    float S = 0.0f;
    for (int i = 0 ; i < (N + BLOCK_SIZE - 1) / BLOCK_SIZE ; i++){
        int col2 = tx + i * BLOCK_SIZE;
        int row2 = ty + i * BLOCK_SIZE;
        A_shared[ty][tx] = (row < M && col2 < N) ? A[row * N + col2] : 0.0f;
        B_shared[ty][tx] = (row2 < N && col < K) ? B[row2 * K + col] : 0.0f;
        __syncthreads();
        for (int j = 0 ; j <  BLOCK_SIZE ; j++){
            S += A_shared[ty][j] * B_shared[j][tx];
        }
        __syncthreads();
    }
    if(row < M && col < K){
        C[row * K + col] = S ;
    }
}


extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
    dim3 threadsPerBlock(16, 16);
    dim3 blocksPerGrid((K + threadsPerBlock.x - 1) / threadsPerBlock.x,
                       (M + threadsPerBlock.y - 1) / threadsPerBlock.y);


    matrix_multiplication_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, M, N, K);
}

V1: One-Dimensional Thread Tiling#

Approach#

一维线程分块（1D Thread Tiling）版本使用以下参数：

BM = 64
BN = 64
BK = 16
TM = 16

每个 block 生成一个 64 x 64 输出 tile。block 的维度为 (64, 4)，共包含 256 个 thread；每个 thread 负责同一输出列中的 16 行：

row = blockIdx.y * 64 + ty * 16 + local_row
col = blockIdx.x * 64 + tx

对应的 16 个部分和保存在 float sum[16] 中。在一次归约步骤（reduction step）内，thread 从 shared memory 读取一个 B 值，并在全部 16 个累加器之间复用它：

sum[r]+=SA[ty×TM+r,inner]×b,0≤r<16.\text{sum}[r] \mathrel{+}= SA[ty \times TM+r,\text{inner}] \times b, \quad 0 \le r < 16.

在每个 thread 只负责一个输出的 kernel 中，每到一个归约位置（reduction position），thread 都要从 shared memory 各读取一个 A 值和一个 B 值，才能执行一次乘加运算。这里，同一个保存 B 值的 register 可以连续支持 16 次乘加。虽然额外的累加器会消耗更多 register，但它们增加了每次 shared-memory 访问所对应的算术工作量，也让 block 级同步的成本分摊到每个 thread 的更多计算上。

block 仍以协作方式填充 SA[64][16] 和 SB[16][64]。将 thread index 展平后，输入加载便不再依赖 thread 与输出位置之间的归属映射。协作式加载和最终写回都会进行边界检查，因此位于矩阵边缘的不完整 tile 也能得到正确结果。

Solution#

#include <cuda_runtime.h>


template <int BM = 64, int BN = 64, int BK = 16, int TM = 16>
__global__ void matrix_multiplication_v1_1d(const float* __restrict__ A,
                                            const float* __restrict__ B,
                                            float* __restrict__ C,
                                            int M, int N, int K) {
    constexpr int NT = (BM / TM) * BN;
    __shared__ float SA[BM][BK];
    __shared__ float SB[BK][BN];
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;
    const int col = blockIdx.x * BN + tx;
    float sum[TM] = {0.0f};


    for (int tile = 0; tile < N; tile += BK) {
        #pragma unroll
        for (int i = tid; i < BM * BK; i += NT) {
            const int local_row = i / BK;
            const int local_col = i % BK;
            const int row = blockIdx.y * BM + local_row;
            const int inner = tile + local_col;
            SA[local_row][local_col] =
                (row < M && inner < N) ? A[row * N + inner] : 0.0f;
        }
        #pragma unroll
        for (int i = tid; i < BK * BN; i += NT) {
            const int local_row = i / BN;
            const int local_col = i % BN;
            const int inner = tile + local_row;
            const int output_col = blockIdx.x * BN + local_col;
            SB[local_row][local_col] =
                (inner < N && output_col < K) ? B[inner * K + output_col] : 0.0f;
        }
        __syncthreads();


        #pragma unroll
        for (int inner = 0; inner < BK; ++inner) {
            const float b = SB[inner][tx];
            #pragma unroll
            for (int row = 0; row < TM; ++row) {
                sum[row] += SA[ty * TM + row][inner] * b;
            }
        }
        __syncthreads();
    }


    #pragma unroll
    for (int local_row = 0; local_row < TM; ++local_row) {
        const int row = blockIdx.y * BM + ty * TM + local_row;
        if (row < M && col < K) C[row * K + col] = sum[local_row];
    }
}


extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
    constexpr int BM = 64, BN = 64, BK = 16, TM = 16;
    const dim3 block(BN, BM / TM);
    const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
    matrix_multiplication_v1_1d<BM, BN, BK, TM><<<grid, block>>>(A, B, C, M, N, K);
}

V1: Two-Dimensional Thread Tiling#

Approach#

二维线程分块（2D Thread Tiling）版本使用以下参数：

BM = 64
BN = 64
BK = 16
TM = 8
TN = 4

block 的维度为 (16, 8)，即总共 128 个 thread。每个 thread 负责一个 8 x 4 微分块（micro-tile），其中包含 32 个输出元素：

row = blockIdx.y * 64 + ty * 8 + local_row
col = blockIdx.x * 64 + tx * 4 + local_col

在一个 reduction position 上，thread 会把 8 个 A 值和 4 个 B 值从 shared memory 读入 register，并通过二者的外积（outer product）更新全部 32 个累加器：

Cmicro+=[a0a1⋮a7][b0b1b2b3].C_{\text{micro}} \mathrel{+}= \begin{bmatrix} a_0\\a_1\\\vdots\\a_7 \end{bmatrix} \begin{bmatrix} b_0&b_1&b_2&b_3 \end{bmatrix}.

这提供了观察矩阵乘法的另一种视角：单个标量输出可以自然地描述为内积，而整个输出 tile 则可以看作沿 reduction 维度不断累加一系列秩一外积。在某个 reduction position 上，由 8 个元素组成的 A 列片段与由 4 个元素组成的 B 行片段，会共同更新该 thread 负责的整个 8 x 4 输出 tile。

矩阵乘法的外积视角
矩阵乘法的外积视角

每个 reduction position 都会为输出 tile 贡献一个秩一外积。

每个保存 A 值的 register 都会参与 4 个输出列的计算，每个保存 B 值的 register 则会参与 8 个输出行的计算。因此，thread 只需从 shared memory 读取 12 个标量，就可以执行 32 次乘加运算。

这正是二维 thread tiling 的核心优势：必要的数学工作量没有改变，但两个操作数都能先在 register 中得到充分复用，之后才需要再次访问 shared memory。

Solution#

#include <cuda_runtime.h>


template <int BM = 64, int BN = 64, int BK = 16, int TM = 8, int TN = 4>
__global__ void matrix_multiplication_v1_2d(const float* __restrict__ A,
                                            const float* __restrict__ B,
                                            float* __restrict__ C,
                                            int M, int N, int K) {
    constexpr int NT = (BM / TM) * (BN / TN);
    __shared__ float SA[BM][BK];
    __shared__ float SB[BK][BN];
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;
    float sum[TM][TN] = {{0.0f}};


    for (int tile = 0; tile < (N + BK - 1) / BK; ++tile) {
        const int base_inner = tile * BK;
        #pragma unroll
        for (int i = tid; i < BM * BK; i += NT) {
            const int local_row = i / BK;
            const int local_col = i % BK;
            const int row = blockIdx.y * BM + local_row;
            const int inner = base_inner + local_col;
            SA[local_row][local_col] =
                (row < M && inner < N) ? A[row * N + inner] : 0.0f;
        }
        #pragma unroll
        for (int i = tid; i < BK * BN; i += NT) {
            const int local_row = i / BN;
            const int local_col = i % BN;
            const int inner = base_inner + local_row;
            const int col = blockIdx.x * BN + local_col;
            SB[local_row][local_col] =
                (inner < N && col < K) ? B[inner * K + col] : 0.0f;
        }
        __syncthreads();


        #pragma unroll
        for (int inner = 0; inner < BK; ++inner) {
            float a_reg[TM], b_reg[TN];
            #pragma unroll
            for (int row = 0; row < TM; ++row) a_reg[row] = SA[ty * TM + row][inner];
            #pragma unroll
            for (int col = 0; col < TN; ++col) b_reg[col] = SB[inner][tx * TN + col];
            #pragma unroll
            for (int row = 0; row < TM; ++row)
                #pragma unroll
                for (int col = 0; col < TN; ++col)
                    sum[row][col] += a_reg[row] * b_reg[col];
        }
        __syncthreads();
    }


    #pragma unroll
    for (int local_row = 0; local_row < TM; ++local_row) {
        const int row = blockIdx.y * BM + ty * TM + local_row;
        if (row >= M) continue;
        #pragma unroll
        for (int local_col = 0; local_col < TN; ++local_col) {
            const int col = blockIdx.x * BN + tx * TN + local_col;
            if (col < K) C[row * K + col] = sum[local_row][local_col];
        }
    }
}


extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
    constexpr int BM = 64, BN = 64, BK = 16, TM = 8, TN = 4;
    const dim3 block(BN / TN, BM / TM);
    const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
    matrix_multiplication_v1_2d<BM, BN, BK, TM, TN><<<grid, block>>>(A, B, C, M, N, K);
}

V2: Vectorized float4 Transfers (BK = 32)#

Approach#

V2 保留了 8 x 4 register micro-tile，并把 reduction tile 的深度改为 BK = 32。一个 float4 包含 4 个连续的 FP32 值，共占 16 字节。kernel 在以下数据传输位置使用它，一次处理 4 个相邻值：

  • 从 global memory 加载 A 和 B；
  • 将数据写入 shared memory；
  • 从 shared memory 读取相邻的 B 值；
  • 最终写入 C 的相邻列。

只有当访问地址满足对齐要求，并且对应的 4 个元素全部有效时，向量加载才是合法的。因此，仅当 N 和 K 都能被 4 整除时，程序才会选择向量化 kernel；任一条件不满足，solve(...) 都会启动能够安全处理边界的标量实现。通用的 LeetGPU 接口必须保留这一回退路径（fallback），不能为了向量化而牺牲不规则矩阵形状（shape）的正确性。

在向量化路径中，最后一个 block 的输出行仍需单独进行边界保护。由于此时 K 是 4 的倍数，一个由 4 列组成的分组要么完整落在矩阵范围内，要么完整位于矩阵范围外，不会只剩部分列有效。

行主序布局让同一行中的连续列在内存中相邻，因此适合沿 A 的归约坐标、B 的输出列坐标以及 C 的输出列坐标进行向量化。在这些传输位置，一条向量指令可以替代 4 条标量指令。实际传输的总字节数，以及矩阵乘法所需的 2MNK2MNK 次浮点运算都没有改变；向量化主要减少的是指令数量和地址生成开销。因此，实际收益取决于地址对齐、指令发射效率，以及数据移动与算术计算之间的平衡，并不会仅仅因为使用了更宽的数据类型就自动出现。

Solution#

#include <cuda_runtime.h>


#ifndef MATMUL_V2_BM
#define MATMUL_V2_BM 64
#endif
#ifndef MATMUL_V2_BN
#define MATMUL_V2_BN 64
#endif
#ifndef MATMUL_V2_BK
#define MATMUL_V2_BK 32
#endif
#ifndef MATMUL_V2_TM
#define MATMUL_V2_TM 8
#endif
#ifndef MATMUL_V2_TN
#define MATMUL_V2_TN 4
#endif


__global__ void scalar_fallback(const float* A, const float* B, float* C,
                                int M, int N, int K) {
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M || col >= K) return;
    float sum = 0.0f;
    for (int inner = 0; inner < N; ++inner) {
        sum += A[row * N + inner] * B[inner * K + col];
    }
    C[row * K + col] = sum;
}


template <int BM = MATMUL_V2_BM, int BN = MATMUL_V2_BN,
          int BK = MATMUL_V2_BK, int TM = MATMUL_V2_TM,
          int TN = MATMUL_V2_TN>
__global__ void matrix_multiplication_v2(const float* __restrict__ A,
                                         const float* __restrict__ B,
                                         float* __restrict__ C,
                                         int M, int N, int K) {
    static_assert(BK % 4 == 0 && BN % 4 == 0 && TN % 4 == 0,
                  "Vectorized dimensions must be multiples of four");
    constexpr int NT = (BM / TM) * (BN / TN);
    __shared__ float SA[BM][BK];
    __shared__ float SB[BK][BN];
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;
    float sum[TM][TN] = {{0.0f}};


    for (int tile = 0; tile < (N + BK - 1) / BK; ++tile) {
        const int base_inner = tile * BK;
        #pragma unroll
        for (int i = tid; i < BM * BK / 4; i += NT) {
            const int local_row = i / (BK / 4);
            const int local_col = (i % (BK / 4)) * 4;
            const int row = blockIdx.y * BM + local_row;
            const int inner = base_inner + local_col;
            float4 value = make_float4(0, 0, 0, 0);
            if (row < M && inner + 3 < N) {
                value = *reinterpret_cast<const float4*>(&A[row * N + inner]);
            }
            *reinterpret_cast<float4*>(&SA[local_row][local_col]) = value;
        }
        #pragma unroll
        for (int i = tid; i < BK * BN / 4; i += NT) {
            const int local_row = i / (BN / 4);
            const int local_col = (i % (BN / 4)) * 4;
            const int inner = base_inner + local_row;
            const int col = blockIdx.x * BN + local_col;
            float4 value = make_float4(0, 0, 0, 0);
            if (inner < N && col + 3 < K) {
                value = *reinterpret_cast<const float4*>(&B[inner * K + col]);
            }
            *reinterpret_cast<float4*>(&SB[local_row][local_col]) = value;
        }
        __syncthreads();


        #pragma unroll
        for (int inner = 0; inner < BK; ++inner) {
            const float4 b = *reinterpret_cast<const float4*>(&SB[inner][tx * TN]);
            #pragma unroll
            for (int row = 0; row < TM; ++row) {
                const float a = SA[ty * TM + row][inner];
                sum[row][0] += a * b.x;
                sum[row][1] += a * b.y;
                sum[row][2] += a * b.z;
                sum[row][3] += a * b.w;
            }
        }
        __syncthreads();
    }


    #pragma unroll
    for (int local_row = 0; local_row < TM; ++local_row) {
        const int row = blockIdx.y * BM + ty * TM + local_row;
        const int col = blockIdx.x * BN + tx * TN;
        if (row < M && col + 3 < K) {
            *reinterpret_cast<float4*>(&C[row * K + col]) =
                make_float4(sum[local_row][0], sum[local_row][1],
                            sum[local_row][2], sum[local_row][3]);
        }
    }
}


extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
    if ((N & 3) != 0 || (K & 3) != 0) {
        const dim3 block(16, 16);
        const dim3 grid((K + 15) / 16, (M + 15) / 16);
        scalar_fallback<<<grid, block>>>(A, B, C, M, N, K);
    } else {
        constexpr int BM = MATMUL_V2_BM, BN = MATMUL_V2_BN;
        constexpr int BK = MATMUL_V2_BK, TM = MATMUL_V2_TM;
        constexpr int TN = MATMUL_V2_TN;
        const dim3 block(BN / TN, BM / TM);
        const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
        matrix_multiplication_v2<BM, BN, BK, TM, TN>
            <<<grid, block>>>(A, B, C, M, N, K);
    }
}

V2 Control: Vectorized Single Buffer (BK = 16)#

Approach#

这个对照版本仍编译同一个单缓冲（single-buffered）V2 kernel，但将 BK = 16，以匹配 V3 的归约分块深度（reduction-tile depth）。V2 的其余设计均保持不变，包括每个 thread 负责的输出、向量化传输、register micro-tile、边界 fallback 和启动布局（launch geometry）。

这一参数匹配的配置，可以把 reduction-tile depth 的影响与双缓冲（double buffering）的影响区分开来。在两款受测 GPU 上，single-buffered BK = 16 对照版本仍然快于 V3。因此，不能仅凭 V2 原本使用 BK = 32、而 V3 使用 BK = 16 来解释 V3 较低的实测性能；在当前实现中，额外的预取 register、两套 shared-memory buffer、buffer 切换与同步操作，同样没有带来足以抵消自身成本的收益。

Solution#

#define MATMUL_V2_BK 16


#include "02_Matrix_Multiplication_v2_vectorized.cu"

V2 Control: Vectorized Single Buffer (128 x 128)#

Approach#

这个对照版本保留 V2 的 single-buffered 算法，只把编译期 tile 参数改为 BM = 128、BN = 128 和 BK = 16。每个 thread 负责的 8 x 4 micro-tile、float4 传输、标量边界 fallback 以及累加逻辑均保持不变。

block 的规模从 128 个 thread 增加到 512 个 thread，生成的输出元素数量也变为原来的 4 倍；然而，更大的 tile 并未提升这个 single-buffered kernel 的性能。在两款受测 GPU 上，它都慢于参数匹配的 64 x 64、BK = 16 对照版本。这说明额外的 block 级数据复用本身还不足以抵消更大 block 带来的调度成本和资源开销。

Solution#

#define MATMUL_V2_BM 128
#define MATMUL_V2_BN 128
#define MATMUL_V2_BK 16
#define MATMUL_V2_TM 8
#define MATMUL_V2_TN 4


#include "02_Matrix_Multiplication_v2_vectorized.cu"

V3: Register Prefetch and Double-Buffered Shared Memory#

Base 64 x 64 Configuration#

Approach#

V3 使用以下参数：

BM = 64
BN = 64
BK = 16
TM = 8
TN = 4

它为每个 shared-memory tile 分配两个副本：

SA[2][64][16]
SB[2][16][64]

第一个 tile 在主循环之前加载。这一阶段称为流水线序言（pipeline prologue）：只有一对完整的输入 tile 已进入 shared memory 后，计算才能开始。

对于后续的每个 tile，thread 首先发出 global load，将数据加载到私有的预取 register。随后，它基于当前 shared-memory buffer 进行计算，等待所有 thread 都用完该 buffer，再将预取值写入另一个 buffer，最后在交换读写 buffer index 之前完成同步。这一阶段称为稳态（steady state），其中 reduction tile ii 的计算与 tile i+1i+1 的数据移动交错进行。

最后一个预取 tile 成为当前 tile 后，kernel 会完成其计算，但不再请求下一个 tile。这一阶段称为流水线尾声（pipeline epilogue）。

预期的流水线如下：

tile i + 1: global memory -> prefetch registers
tile i:     shared memory -> accumulator registers
tile i + 1: prefetch registers -> alternate shared-memory buffer

这是带 register prefetch 的软件双缓冲（software double buffering），而不是异步 cp.async 流水线。global load 会在当前 tile 被使用之前发出，从而使编译器和 GPU 有机会将尚未完成的内存操作与独立的算术计算重叠执行。不过，当 shared buffer 切换角色时，该实现仍需要 block 范围的 barrier。

实际源码使用直接的行主序 tile 布局保存 SA[buffer][row][inner] 和 SB[buffer][inner][col]，并未暗中加入转置 shared-memory 布局或其他未实现的访问变换。这一点很重要，因为评价一条流水线时，不能脱离其实际加载模式、register footprint、shared-memory layout 和同步调度。

double buffering 并不会自动提升速度。它使 shared-memory 存储量翻倍，并增加预取 register、控制流、pipeline prologue 和 epilogue。这些开销可能降低 occupancy，或抵消隐藏延迟带来的收益。在实测工作负载中，基础 64 x 64 V3 在两款受测 GPU 上都略慢于参数匹配的 single-buffer V2 对照版本。后面的 128 x 128 V3 配置更快，是因为更大的输出 tile 同时改变了复用和调度；保留各个独立对照版本，可以避免将该结果错误地完全归因于 double buffering。

Solution#
#include <cuda_runtime.h>


#ifndef MATMUL_BM
#define MATMUL_BM 64
#define MATMUL_BN 64
#define MATMUL_BK 16
#define MATMUL_TM 8
#define MATMUL_TN 4
#endif


__global__ void scalar_fallback_v3(const float* A, const float* B, float* C,
                                   int M, int N, int K) {
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M || col >= K) return;
    float sum = 0.0f;
    for (int inner = 0; inner < N; ++inner) {
        sum += A[row * N + inner] * B[inner * K + col];
    }
    C[row * K + col] = sum;
}


template <int BM = 64, int BN = 64, int BK = 16, int TM = 8, int TN = 4>
__global__ void matrix_multiplication_v3(const float* __restrict__ A,
                                         const float* __restrict__ B,
                                         float* __restrict__ C,
                                         int M, int N, int K) {
    constexpr int NT = (BM / TM) * (BN / TN);
    constexpr int A_FLOAT4_PER_THREAD = (BM * BK / 4 + NT - 1) / NT;
    constexpr int B_FLOAT4_PER_THREAD = (BK * BN / 4 + NT - 1) / NT;
    __shared__ float SA[2][BM][BK];
    __shared__ float SB[2][BK][BN];


    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;
    float a_prefetch[A_FLOAT4_PER_THREAD][4];
    float b_prefetch[B_FLOAT4_PER_THREAD][4];
    float sum[TM][TN] = {{0.0f}};


    auto load_tile = [&](int tile) {
        const int base_inner = tile * BK;
        #pragma unroll
        for (int i = 0; i < A_FLOAT4_PER_THREAD; ++i) {
            const int index = (i * NT + tid) * 4;
            float4 value = make_float4(0, 0, 0, 0);
            if (index < BM * BK) {
                const int local_row = index / BK;
                const int local_col = index % BK;
                const int row = blockIdx.y * BM + local_row;
                const int inner = base_inner + local_col;
                if (row < M && inner + 3 < N) {
                    value = *reinterpret_cast<const float4*>(&A[row * N + inner]);
                }
            }
            a_prefetch[i][0] = value.x; a_prefetch[i][1] = value.y;
            a_prefetch[i][2] = value.z; a_prefetch[i][3] = value.w;
        }
        #pragma unroll
        for (int i = 0; i < B_FLOAT4_PER_THREAD; ++i) {
            const int index = (i * NT + tid) * 4;
            float4 value = make_float4(0, 0, 0, 0);
            if (index < BK * BN) {
                const int local_row = index / BN;
                const int local_col = index % BN;
                const int inner = base_inner + local_row;
                const int col = blockIdx.x * BN + local_col;
                if (inner < N && col + 3 < K) {
                    value = *reinterpret_cast<const float4*>(&B[inner * K + col]);
                }
            }
            b_prefetch[i][0] = value.x; b_prefetch[i][1] = value.y;
            b_prefetch[i][2] = value.z; b_prefetch[i][3] = value.w;
        }
    };


    auto store_tile = [&](int buffer) {
        #pragma unroll
        for (int i = 0; i < A_FLOAT4_PER_THREAD; ++i) {
            const int index = (i * NT + tid) * 4;
            if (index < BM * BK) {
                const int local_row = index / BK;
                const int local_col = index % BK;
                *reinterpret_cast<float4*>(&SA[buffer][local_row][local_col]) =
                    make_float4(a_prefetch[i][0], a_prefetch[i][1],
                                a_prefetch[i][2], a_prefetch[i][3]);
            }
        }
        #pragma unroll
        for (int i = 0; i < B_FLOAT4_PER_THREAD; ++i) {
            const int index = (i * NT + tid) * 4;
            if (index < BK * BN) {
                const int local_row = index / BN;
                const int local_col = index % BN;
                *reinterpret_cast<float4*>(&SB[buffer][local_row][local_col]) =
                    make_float4(b_prefetch[i][0], b_prefetch[i][1],
                                b_prefetch[i][2], b_prefetch[i][3]);
            }
        }
    };


    const int tiles = (N + BK - 1) / BK;
    load_tile(0);
    store_tile(0);
    __syncthreads();


    int read_buffer = 0;
    for (int tile = 0; tile < tiles; ++tile) {
        if (tile + 1 < tiles) load_tile(tile + 1);


        #pragma unroll
        for (int inner = 0; inner < BK; ++inner) {
            #pragma unroll
            for (int local_col = 0; local_col < TN; local_col += 4) {
                const float4 b = *reinterpret_cast<const float4*>(
                    &SB[read_buffer][inner][tx * TN + local_col]);
                #pragma unroll
                for (int row = 0; row < TM; ++row) {
                    const float a = SA[read_buffer][ty * TM + row][inner];
                    sum[row][local_col + 0] += a * b.x;
                    sum[row][local_col + 1] += a * b.y;
                    sum[row][local_col + 2] += a * b.z;
                    sum[row][local_col + 3] += a * b.w;
                }
            }
        }
        __syncthreads();


        if (tile + 1 < tiles) {
            const int write_buffer = read_buffer ^ 1;
            store_tile(write_buffer);
            __syncthreads();
            read_buffer = write_buffer;
        }
    }


    #pragma unroll
    for (int local_row = 0; local_row < TM; ++local_row) {
        const int row = blockIdx.y * BM + ty * TM + local_row;
        #pragma unroll
        for (int local_col = 0; local_col < TN; local_col += 4) {
            const int col = blockIdx.x * BN + tx * TN + local_col;
            if (row < M && col + 3 < K) {
                *reinterpret_cast<float4*>(&C[row * K + col]) =
                    make_float4(sum[local_row][local_col + 0],
                                sum[local_row][local_col + 1],
                                sum[local_row][local_col + 2],
                                sum[local_row][local_col + 3]);
            }
        }
    }
}


#ifndef MATMUL_NO_SOLVE
extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K) {
    if ((N & 3) != 0 || (K & 3) != 0) {
        const dim3 block(16, 16);
        const dim3 grid((K + 15) / 16, (M + 15) / 16);
        scalar_fallback_v3<<<grid, block>>>(A, B, C, M, N, K);
    } else {
        constexpr int BM = MATMUL_BM, BN = MATMUL_BN, BK = MATMUL_BK;
        constexpr int TM = MATMUL_TM, TN = MATMUL_TN;
        const dim3 block(BN / TN, BM / TM);
        const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
        matrix_multiplication_v3<BM, BN, BK, TM, TN>
            <<<grid, block>>>(A, B, C, M, N, K);
    }
}
#endif

Large-Tile 128 x 128 Configuration#

Approach#

large-tile 配置将 V3 的编译期参数改为：

BM = 128
BN = 128
BK = 16
TM = 8
TN = 4

block 包含 (128 / 4) x (128 / 8) = 32 x 16 = 512 个 thread。每个 thread 仍负责 32 个输出元素，而整个 block 共同生成一个 128 x 128 tile。其 double-buffered 输入存储占用 32 KiB：

2 * (128 * 16 + 16 * 128) * sizeof(float)

更大的输出 tile 提高了 block 级复用，并将 input-tile 加载开销分摊到更多输出元素上，同时完整保留基础 V3 配置的 register-prefetch 和 double-buffer 机制。它也增加了 block 的 thread 数量，并改变 register 需求、occupancy 和调度，因此它是 V3 的一种参数配置，而不是独立的优化原理。

在修订后的 benchmark 中，该配置是两款受测 GPU 上本项目最快的 CUDA 实现。然而，参数匹配的 single-buffered 实验在扩大到相同的 128 x 128 tile 后反而变慢。因此，收益并非单独来自 tile 大小：只有当更大的输出 tile 与这条 double-buffered 数据路径结合时，收益才会出现。实测结果表明这两种选择存在相互作用，而不能由此得出更大 tile 普遍更快的结论。

Solution#
#define MATMUL_BM 128
#define MATMUL_BN 128
#define MATMUL_BK 16
#define MATMUL_TM 8
#define MATMUL_TN 4


#include "02_Matrix_Multiplication_v3_double_buffered.cu"

Controlled Comparison#

下表 4 行均使用 BK = 16、每个 thread 一个 8 x 4 micro-tile、确定性的非零输入、5 次 warm-up iteration 和 10 次 measured iteration。

BufferingOutput TileRTX 5070 Ti Laptop GPURTX 4090
Single buffer64 x 6438.761 ms11.314 ms
Single buffer128 x 12838.842 ms13.146 ms
Double buffer64 x 6439.642 ms11.514 ms
Double buffer128 x 12834.590 ms10.052 ms

在 64 x 64 配置下，double buffering 的影响很小且取决于平台：在 RTX 4090 上略快，在 RTX 5070 Ti Laptop GPU 上略慢。将 single-buffered kernel 扩大到 128 x 128 后，两款 GPU 上的速度都更慢。相比之下，128 x 128 double-buffered 配置明显快于 single-buffered large tile，也是表中最快的通用配置。因此，实测加速来自更大输出 tile 与 double-buffered 数据路径的相互作用，不能由此得出任一选择单独使用时普遍更快的结论。

以下编译器资源报告由面向 RTX 4090 target 的 nvcc -O2 -arch=sm_89 -Xptxas=-v 生成：

BufferingOutput TileThreads per BlockRegisters per ThreadShared Memory per BlockSpills
Single buffer64 x 64128868 KiB0
Single buffer128 x 1285128616 KiB0
Double buffer64 x 6412812216 KiB0
Double buffer128 x 12851210032 KiB0

该报告证实了资源方面的权衡。single-buffered large tile 仍为每个 thread 使用 86 个 register，但 thread 数量增加到 4 倍，因此每个 block 分配的 register 从 11,008 个增加到 44,032 个。基础 double-buffered kernel 因预取状态而将每个 thread 的 register 用量提高到 122 个。对于 large double-buffered 实例，编译器为每个 thread 使用 100 个 register，且没有 spill；block 在保持相同 per-thread micro-tile 的同时，完成了 4 倍的输出工作。这些事实解释了为何 occupancy 和复用程度会随配置发生显著变化，但仅凭 register 与 shared-memory 数量并不能确定实际运行时的 stall 行为。

V4: Hardware-Asynchronous cp.async Pipeline#

Approach#

V4 保留 large V3 配置的 128 x 128 输出 tile、BK = 16 和每个 thread 一个 8 x 4 micro-tile，但改变了下一个输入 tile 进入 shared memory 的方式。

V3 首先将 global-memory 中的值加载到普通 thread register，随后再把这些 register 中的值存入另一个 shared-memory buffer。V4 则改为发出 Ampere 或更新架构支持的 PTX 指令：

cp.async.cg.shared.global

每条指令直接从 global memory 向 shared memory 复制 16 字节。cp.async.commit_group 提交一个 thread 发出的复制操作，cp.async.wait_group 0 则等待 block 所需的全部未完成 group 执行完毕。

该流水线包含 prologue、steady state 和 epilogue：

  1. prologue 异步填充 shared buffer 0，并在首次使用前等待填充完成。
  2. 在一次 steady-state iteration 中，block 会先将 tile i + 1 的复制操作发往另一个 buffer，再计算 tile i。
  3. 算术循环只读取当前 buffer，因此 GPU 可将 fused multiply-add 工作与尚未完成的 global-to-shared 复制重叠执行。
  4. buffer 交换角色前，cp.async.wait_group 0 保证复制完成，__syncthreads() 则使新 tile 对所有 thread 可见。
  5. 最后一个 tile 没有后继，因此循环结束时不会再发出另一组复制操作。

越界的 16 字节 vector 会使用有效字节数为零的 source-size operand。硬件将目标 shared-memory vector 填零，而不会解引用无效的 global address。reduction 维度或 output-column 维度不适合 4 个 float 向量传输的输入会使用标量 fallback。compute capability 低于 8.0 的设备因不支持 cp.async，也会使用 fallback。

与 V3 的 register-prefetch 软件流水线不同，这是一条真正的硬件异步流水线（hardware-asynchronous pipeline）。不过，它仍不保证更快：两个版本在 shared buffer 切换角色时都需要同步，而包含 512 个 thread 的 block、地址生成、copy-group 管理和算术指令流都会争用执行资源。

Solution#

#include <cuda_runtime.h>


namespace {


constexpr int BM = 128;
constexpr int BN = 128;
constexpr int BK = 16;
constexpr int TM = 8;
constexpr int TN = 4;
constexpr int THREADS_X = BN / TN;
constexpr int THREADS_Y = BM / TM;
constexpr int THREAD_COUNT = THREADS_X * THREADS_Y;


__global__ void scalar_fallback_v4(const float* A, const float* B, float* C,
                                   int M, int N, int K) {
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M || col >= K) return;


    float sum = 0.0f;
    for (int inner = 0; inner < N; ++inner) {
        sum += A[row * N + inner] * B[inner * K + col];
    }
    C[row * K + col] = sum;
}


#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
__device__ __forceinline__ void cp_async_16(void* shared_destination,
                                            const void* global_source,
                                            int valid_bytes) {
    const unsigned int shared_address =
        static_cast<unsigned int>(__cvta_generic_to_shared(shared_destination));
    asm volatile(
        "cp.async.cg.shared.global [%0], [%1], 16, %2;\n"
        :
        : "r"(shared_address), "l"(global_source), "r"(valid_bytes));
}


__device__ __forceinline__ void cp_async_commit() {
    asm volatile("cp.async.commit_group;\n" : :);
}


__device__ __forceinline__ void cp_async_wait_all() {
    asm volatile("cp.async.wait_group 0;\n" : :);
}
#endif


__global__ void matrix_multiplication_v4_cp_async(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    __shared__ __align__(16) float shared_a[2][BM][BK];
    __shared__ __align__(16) float shared_b[2][BK][BN];


    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * blockDim.x + tx;
    float accumulators[TM][TN] = {{0.0f}};


    auto issue_tile = [&](int tile, int buffer) {
        constexpr int A_VECTORS = BM * BK / 4;
        constexpr int B_VECTORS = BK * BN / 4;
        const int inner_base = tile * BK;


        for (int vector_index = tid;
             vector_index < A_VECTORS;
             vector_index += THREAD_COUNT) {
            const int element_index = vector_index * 4;
            const int local_row = element_index / BK;
            const int local_inner = element_index % BK;
            const int row = blockIdx.y * BM + local_row;
            const int inner = inner_base + local_inner;
            const int valid_bytes = (row < M && inner + 3 < N) ? 16 : 0;
            const float* source = A + static_cast<size_t>(row < M ? row : 0) * N +
                                  (inner + 3 < N ? inner : 0);
            cp_async_16(&shared_a[buffer][local_row][local_inner],
                        source, valid_bytes);
        }


        for (int vector_index = tid;
             vector_index < B_VECTORS;
             vector_index += THREAD_COUNT) {
            const int element_index = vector_index * 4;
            const int local_inner = element_index / BN;
            const int local_col = element_index % BN;
            const int inner = inner_base + local_inner;
            const int col = blockIdx.x * BN + local_col;
            const int valid_bytes = (inner < N && col + 3 < K) ? 16 : 0;
            const float* source = B + static_cast<size_t>(inner < N ? inner : 0) * K +
                                  (col + 3 < K ? col : 0);
            cp_async_16(&shared_b[buffer][local_inner][local_col],
                        source, valid_bytes);
        }
        cp_async_commit();
    };


    const int tile_count = (N + BK - 1) / BK;
    issue_tile(0, 0);
    cp_async_wait_all();
    __syncthreads();


    int read_buffer = 0;
    for (int tile = 0; tile < tile_count; ++tile) {
        const bool has_next_tile = tile + 1 < tile_count;
        if (has_next_tile) {
            issue_tile(tile + 1, read_buffer ^ 1);
        }


        #pragma unroll
        for (int inner = 0; inner < BK; ++inner) {
            const float4 b = *reinterpret_cast<const float4*>(
                &shared_b[read_buffer][inner][tx * TN]);
            #pragma unroll
            for (int local_row = 0; local_row < TM; ++local_row) {
                const float a =
                    shared_a[read_buffer][ty * TM + local_row][inner];
                accumulators[local_row][0] += a * b.x;
                accumulators[local_row][1] += a * b.y;
                accumulators[local_row][2] += a * b.z;
                accumulators[local_row][3] += a * b.w;
            }
        }


        if (has_next_tile) {
            cp_async_wait_all();
            __syncthreads();
            read_buffer ^= 1;
        }
    }


    #pragma unroll
    for (int local_row = 0; local_row < TM; ++local_row) {
        const int row = blockIdx.y * BM + ty * TM + local_row;
        const int col = blockIdx.x * BN + tx * TN;
        if (row < M && col + 3 < K) {
            *reinterpret_cast<float4*>(&C[static_cast<size_t>(row) * K + col]) =
                make_float4(accumulators[local_row][0],
                            accumulators[local_row][1],
                            accumulators[local_row][2],
                            accumulators[local_row][3]);
        }
    }
#endif
}


}  // namespace


extern "C" void solve(const float* A, const float* B, float* C,
                      int M, int N, int K) {
    static const bool supports_cp_async = [] {
        int device = 0;
        cudaDeviceProp properties{};
        cudaGetDevice(&device);
        cudaGetDeviceProperties(&properties, device);
        return properties.major >= 8;
    }();


    if (!supports_cp_async || (N & 3) != 0 || (K & 3) != 0) {
        const dim3 block(16, 16);
        const dim3 grid((K + 15) / 16, (M + 15) / 16);
        scalar_fallback_v4<<<grid, block>>>(A, B, C, M, N, K);
        return;
    }


    const dim3 block(THREADS_X, THREADS_Y);
    const dim3 grid((K + BN - 1) / BN, (M + BM - 1) / BM);
    matrix_multiplication_v4_cp_async<<<grid, block>>>(A, B, C, M, N, K);
}

V5: WMMA TF32 Tensor Core Kernel#

Approach#

V5 同时改变了执行单元和数值格式。公共接口仍然接收和返回 FP32 矩阵，但 fast path 会将输入值显式舍入为 TensorFloat-32，并在 Tensor Core 上使用 FP32 累加器执行矩阵乘加运算。

TF32 保留 FP32 的 8 位指数范围，但使用 10 位显式尾数。转换通过以下指令完成：

cvt.rna.tf32.f32

这意味着 V5 在数值上与严格的 FP32 CUDA Core 版本并不完全相同。它以降低输入尾数精度为代价，换取使用 Tensor Core 矩阵指令的能力。因此，测试采用 5e-3 的缩放容差，同时仍会拒绝所有非有限结果，并将每个输出元素与 CPU reference 进行比较。

一个 block 生成一个 128 x 128 输出 tile，并以 32 为单位分块归约 K。该 block 包含 16 个 warp。每个 warp 负责一条 16 x 64 输出带，由 4 个 16 x 16 accumulator fragment 表示。两个 warp 分别覆盖同一个 16 行输出带的左半部分和右半部分。

对于每个 K=32 shared-memory tile：

  1. 全部 512 个 thread 协作加载 A[128, 32] 和 B[32, 128]，并将每个值转换为 TF32。
  2. 一次 block 同步使两个 tile 均可见。
  3. 由于 TF32 WMMA 使用 16 x 16 x 8 fragment，该 tile 被划分为 4 个 WMMA reduction step。
  4. 每个 warp 在每个 step 加载一个 A fragment 和 4 个 B fragment。
  5. wmma::mma_sync 更新 4 个 FP32 accumulator fragment。
  6. 所有 reduction tile 完成后，wmma::store_matrix_sync 将 fragment 写入行主序 FP32 output memory。

fast path 要求 M 和 K 与 128 x 128 输出 tile 对齐，N 与 BK=32 对齐。其他 shape、较旧架构和边界情况会使用标量 FP32 fallback。这样既保留了题目接口和一般情况下的正确性，也使 WMMA kernel 无需处理 partial-fragment store。

Tensor Core 能够提供很高的峰值吞吐率，但该实现也需要承担显式 TF32 转换、shared-memory staging、频繁 fragment load、block 同步和 WMMA fragment 归属管理等开销。因此，这些测量结果证明的是一个正确的底层 WMMA 实现，并不意味着它必然优于经过精细调优的 CUDA Core kernel。

Solution#

#include <cuda_runtime.h>
#include <mma.h>


namespace {


constexpr int BM = 128;
constexpr int BN = 128;
constexpr int BK = 32;
constexpr int WARPS_PER_BLOCK = 16;
constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * 32;
constexpr int WMMA_K = 8;
constexpr int OUTPUT_FRAGMENTS_PER_WARP = BN / 32;


__device__ __forceinline__ float to_tf32(float value) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    unsigned int tf32_bits;
    asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(tf32_bits) : "f"(value));
    return __uint_as_float(tf32_bits);
#else
    return value;
#endif
}


__global__ void scalar_fallback_v5(const float* A, const float* B, float* C,
                                   int M, int N, int K) {
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M || col >= K) return;


    float sum = 0.0f;
    for (int inner = 0; inner < N; ++inner) {
        sum += A[row * N + inner] * B[inner * K + col];
    }
    C[row * K + col] = sum;
}


__global__ void matrix_multiplication_v5_wmma_tf32(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    using namespace nvcuda;


    __shared__ __align__(16) float shared_a[BM][BK];
    __shared__ __align__(16) float shared_b[BK][BN];


    const int tid = threadIdx.x;
    const int warp_id = tid / 32;
    const int warp_row = warp_id / 2;
    const int warp_col_group = warp_id % 2;


    wmma::fragment<wmma::accumulator, 16, 16, WMMA_K, float>
        accumulators[OUTPUT_FRAGMENTS_PER_WARP];
    #pragma unroll
    for (int fragment = 0; fragment < OUTPUT_FRAGMENTS_PER_WARP; ++fragment) {
        wmma::fill_fragment(accumulators[fragment], 0.0f);
    }


    for (int inner_base = 0; inner_base < N; inner_base += BK) {
        for (int index = tid; index < BM * BK;
             index += THREADS_PER_BLOCK) {
            const int local_row = index / BK;
            const int local_inner = index % BK;
            const int row = blockIdx.y * BM + local_row;
            const float value = A[static_cast<size_t>(row) * N +
                                  inner_base + local_inner];
            shared_a[local_row][local_inner] = to_tf32(value);
        }
        for (int index = tid; index < BK * BN;
             index += THREADS_PER_BLOCK) {
            const int local_inner = index / BN;
            const int local_col = index % BN;
            const int col = blockIdx.x * BN + local_col;
            const float value = B[static_cast<size_t>(inner_base + local_inner) * K +
                                  col];
            shared_b[local_inner][local_col] = to_tf32(value);
        }
        __syncthreads();


        const int local_row = warp_row * 16;
        #pragma unroll
        for (int inner = 0; inner < BK; inner += WMMA_K) {
            wmma::fragment<wmma::matrix_a, 16, 16, WMMA_K,
                           wmma::precision::tf32, wmma::row_major> a_fragment;
            wmma::load_matrix_sync(
                a_fragment, &shared_a[local_row][inner], BK);


            #pragma unroll
            for (int fragment = 0;
                 fragment < OUTPUT_FRAGMENTS_PER_WARP;
                 ++fragment) {
                wmma::fragment<wmma::matrix_b, 16, 16, WMMA_K,
                               wmma::precision::tf32,
                               wmma::row_major> b_fragment;
                wmma::load_matrix_sync(
                    b_fragment,
                    &shared_b[inner][warp_col_group * 64 + fragment * 16],
                    BN);
                wmma::mma_sync(accumulators[fragment], a_fragment,
                               b_fragment, accumulators[fragment]);
            }
        }
        __syncthreads();
    }


    const int output_row = blockIdx.y * BM + warp_row * 16;
    const int output_col = blockIdx.x * BN + warp_col_group * 64;
    #pragma unroll
    for (int fragment = 0; fragment < OUTPUT_FRAGMENTS_PER_WARP; ++fragment) {
        wmma::store_matrix_sync(
            &C[static_cast<size_t>(output_row) * K +
               output_col + fragment * 16],
            accumulators[fragment], K, wmma::mem_row_major);
    }
#endif
}


}  // namespace


extern "C" void solve(const float* A, const float* B, float* C,
                      int M, int N, int K) {
    static const bool supports_wmma_tf32 = [] {
        int device = 0;
        cudaDeviceProp properties{};
        cudaGetDevice(&device);
        cudaGetDeviceProperties(&properties, device);
        return properties.major >= 8;
    }();


    if (!supports_wmma_tf32 ||
        M % BM != 0 || N % BK != 0 || K % BN != 0) {
        const dim3 block(16, 16);
        const dim3 grid((K + 15) / 16, (M + 15) / 16);
        scalar_fallback_v5<<<grid, block>>>(A, B, C, M, N, K);
        return;
    }


    const dim3 block(THREADS_PER_BLOCK);
    const dim3 grid(K / BN, M / BM);
    matrix_multiplication_v5_wmma_tf32<<<grid, block>>>(A, B, C, M, N, K);
}

Test Methodology#

所有 CUDA 方法都使用同一个测试文件，并且必须通过 9 组正确性 shape，其中包括标量、矩形、对齐、未对齐、partial-tile 和 multi-block 维度。检查器会拒绝非有限输出值，并使用 5e-3 的缩放容差将每个输出元素与 CPU reference 比较。该容差覆盖 V5 有意进行的 TF32 输入舍入；严格的 FP32 方法也使用同一测试评估，从而确保对比中的每一行都遵循统一流程。129 x 96 乘 96 x 132 的 case 用于覆盖 V4 处理不完整输出 tile 的 asynchronous fast path，而 256 x 96 乘 96 x 256 的 case 则用于覆盖多个 WMMA output block。

性能测试 case 使用 M=8192、N=6144 和 K=4096。输入被初始化为确定性的非零值，以免全零分配产生缺乏代表性的内存行为。每种方法都针对远程 GPU 的原生架构编译。在 10 次基于 CUDA event 的测量前，先执行 5 次 warm-up iteration。下表给出这 10 次测量的平均值、最小值、最大值和计算得到的浮点吞吐率。

Test Code#

#include <cuda_runtime.h>


#include <array>
#include <cmath>
#include <iostream>


extern "C" void solve(const float* a, const float* b, float* c,
                      int m, int n, int k);


int main() {
    constexpr int m = 2;
    constexpr int n = 2;
    constexpr int k = 2;
    const std::array<float, 4> matrix_a = {1.0f, 2.0f, 3.0f, 4.0f};
    const std::array<float, 4> matrix_b = {5.0f, 6.0f, 7.0f, 8.0f};
    const std::array<float, 4> expected = {19.0f, 22.0f, 43.0f, 50.0f};
    std::array<float, 4> actual = {};
    constexpr size_t bytes = 4 * sizeof(float);


    float* device_a = nullptr;
    float* device_b = nullptr;
    float* device_c = nullptr;


    auto check_cuda = [](cudaError_t status, const char* operation) {
        if (status == cudaSuccess) {
            return true;
        }
        std::cerr << "Test failed. CUDA error in " << operation << ": "
                  << cudaGetErrorString(status) << '\n';
        return false;
    };


    bool success = check_cuda(cudaMalloc(&device_a, bytes), "cudaMalloc(device_a)") &&
                   check_cuda(cudaMalloc(&device_b, bytes), "cudaMalloc(device_b)") &&
                   check_cuda(cudaMalloc(&device_c, bytes), "cudaMalloc(device_c)") &&
                   check_cuda(cudaMemcpy(device_a, matrix_a.data(), bytes,
                                         cudaMemcpyHostToDevice), "copy matrix_a") &&
                   check_cuda(cudaMemcpy(device_b, matrix_b.data(), bytes,
                                         cudaMemcpyHostToDevice), "copy matrix_b");


    if (success) {
        solve(device_a, device_b, device_c, m, n, k);
        success = check_cuda(cudaGetLastError(), "solve") &&
                  check_cuda(cudaDeviceSynchronize(), "synchronize") &&
                  check_cuda(cudaMemcpy(actual.data(), device_c, bytes,
                                        cudaMemcpyDeviceToHost), "copy result");
    }


    cudaFree(device_a);
    cudaFree(device_b);
    cudaFree(device_c);


    if (!success) {
        return 1;
    }


    for (size_t index = 0; index < actual.size(); ++index) {
        if (std::fabs(actual[index] - expected[index]) > 1e-5f) {
            std::cerr << "Test failed at index " << index
                      << ". Expected: " << expected[index]
                      << ", Actual: " << actual[index] << '\n';
            return 1;
        }
    }


    std::cout << "Test passed.\n";
    return 0;
}

Test Result#

以下表格特意将同一款 GPU 的结果集中展示。它们用于比较同一平台上的不同方法，而不是为两款 GPU 排名。Default 当前选择 V3 large-tile 实现，因此该行是对同一源码路径的重复运行，并非另一种算法。

NVIDIA GeForce RTX 5070 Ti Laptop GPU#

MethodArithmeticStatusAverage TimeMinimum TimeMaximum TimePerformance
V0 naiveFP32PASS383.370 ms348.679 ms417.879 ms1075.506 GFLOPS
V1 shared memoryFP32PASS250.884 ms224.301 ms302.463 ms1643.454 GFLOPS
V1 1D thread tilingFP32PASS135.891 ms128.005 ms150.389 ms3034.181 GFLOPS
V1 2D thread tilingFP32PASS123.423 ms117.553 ms135.046 ms3340.669 GFLOPS
V2 vectorized, BK=32FP32PASS46.376 ms44.035 ms48.737 ms8890.723 GFLOPS
V2 vectorized, BK=16FP32PASS38.761 ms36.343 ms40.281 ms10637.490 GFLOPS
V2 vectorized, 128 x 128FP32PASS38.842 ms37.423 ms41.546 ms10615.226 GFLOPS
V3 software double buffer, 64 x 64FP32PASS39.642 ms34.580 ms42.673 ms10400.993 GFLOPS
V3 software double buffer, 128 x 128FP32PASS34.590 ms32.101 ms37.409 ms11920.091 GFLOPS
V4 hardware cp.asyncFP32PASS34.572 ms31.259 ms37.045 ms11926.402 GFLOPS
V5 WMMA Tensor CoreTF32 input, FP32 accumulationPASS57.805 ms52.967 ms66.704 ms7132.925 GFLOPS
Default (V3 128 x 128)FP32PASS34.486 ms31.874 ms37.132 ms11956.000 GFLOPS

在这款 GPU 上，V4 在最终统一测试中的平均时间最低，但考虑到观察到的运行间波动，其 34.572 ms 结果实际上与匹配的 large V3 的 34.590 ms 持平。V4 多次运行的时间大约在 32 至 35 ms 之间。因此，证据表明 hardware-asynchronous copy 在该平台上具有竞争力，而不是 0.018 ms 的最终差异具有显著意义。V5 在文档所述的 TF32 容差下结果正确，但它慢于经过调优的 CUDA Core kernel，因为这个手写 WMMA 映射需要承担大量转换、shared-memory、同步和 fragment 管理开销。

NVIDIA GeForce RTX 4090#

MethodArithmeticStatusAverage TimeMinimum TimeMaximum TimePerformance
V0 naiveFP32PASS83.140 ms82.831 ms83.375 ms4959.321 GFLOPS
V1 shared memoryFP32PASS59.379 ms58.985 ms59.865 ms6943.818 GFLOPS
V1 1D thread tilingFP32PASS19.484 ms19.393 ms19.505 ms21161.508 GFLOPS
V1 2D thread tilingFP32PASS13.446 ms13.420 ms13.465 ms30665.525 GFLOPS
V2 vectorized, BK=32FP32PASS11.646 ms11.098 ms12.106 ms35404.086 GFLOPS
V2 vectorized, BK=16FP32PASS11.314 ms11.213 ms11.345 ms36441.761 GFLOPS
V2 vectorized, 128 x 128FP32PASS13.146 ms13.135 ms13.173 ms31363.839 GFLOPS
V3 software double buffer, 64 x 64FP32PASS11.514 ms11.380 ms11.654 ms35809.707 GFLOPS
V3 software double buffer, 128 x 128FP32PASS10.052 ms9.939 ms10.213 ms41019.869 GFLOPS
V4 hardware cp.asyncFP32PASS10.293 ms10.230 ms10.415 ms40056.186 GFLOPS
V5 WMMA Tensor CoreTF32 input, FP32 accumulationPASS12.675 ms12.107 ms13.133 ms32528.826 GFLOPS
Default (V3 128 x 128)FP32PASS10.096 ms10.010 ms10.144 ms40840.572 GFLOPS

在这款 GPU 上，large V3 软件流水线仍是实测最快的实现。V4 将复制与算术计算重叠执行，但未能抵消额外的 copy-group 和同步开销。V5 再次证明 Tensor Core 执行功能正常，却没有超越已有的 FP32 路径。这一结果说明了一个有用的边界：当数据准备和 fragment 调度尚未达到生产级 GEMM 库的成熟程度时，仅凭 Tensor Core 的峰值吞吐率并不能保证端到端 kernel 更快。

Triton#

Approach#

Triton 实现计算与 CUDA 相同的行主序矩阵乘积：

C[m,k]=∑n=0N−1A[m,n]B[n,k].C[m,k] = \sum_{n=0}^{N-1} A[m,n]B[n,k].

一个 Triton program instance 负责一个 64 x 64 输出 tile。与 CUDA 源码不同，该实现不会将单个标量输出显式分配给 threadIdx 坐标，而是描述完整的 index block 和 value block；Triton 编译器再将这些操作映射到 GPU thread 和 warp。reduction 维度以 32 为单位分块处理：

BLOCK_SIZE_M = 64
BLOCK_SIZE_N = 32
BLOCK_SIZE_K = 64
GROUP_SIZE_M = 8

输出 grid 被展平为一维。tl.program_id(0) 用于标识当前 program instance。kernel 根据这个 linear ID 推导 PID_M 和 PID_K，二者分别标识输出行 tile 和输出列 tile。完整 grid 包含：

⌈M64⌉×⌈K64⌉\left\lceil \frac{M}{64} \right\rceil \times \left\lceil \frac{K}{64} \right\rceil

个 program instance。

program 按最多 8 个 M tile 分组排序，之后才沿 K 继续推进。这只会改变执行顺序，不会改变每个 program 所负责的输出 tile。相邻 program 更有可能复用 A 和 B 的 cache 区域；相比之下，简单的行主序排列可能会沿一个输出维度前进较远后，才重新访问可复用的输入数据。

对于一个 program，行列 offset 分别为：

a_m_offset = PID_M * 64 + [0, ..., 63]
b_k_offset = PID_K * 64 + [0, ..., 63]

在每个 reduction step 中，tl.arange 创建行、reduction 和列 offset。使用 [:, None] 和 [None, :] 添加 singleton dimension 后，这些 vector 会被广播为来自 A 的 64 x 32 pointer tile 和来自 B 的 32 x 64 pointer tile。A、B 和 C 的行主序 stride 分别以 (N, 1)、(K, 1) 和 (K, 1) 显式传入。

mask 在正确性方面所起的作用，与 CUDA kernel 中的边界检查和零填充相同。被 mask 的输入位置会按零加载，因此不会对 reduction 产生贡献；被 mask 的输出位置则不会被存储。因此，同一个 kernel 即可处理不是 64 或 32 的倍数的维度，无需启动单独的边界 kernel。

tl.dot 将两个输入 tile 相乘，并累加到一个 float32 64 x 64 block 中。allow_tf32=False 使计算保持在该实现采用的 non-TF32 路径上。reduction loop 结束后，最终的 mask 会防止向 C 范围外存储。

launch 为每个 program 使用 4 个 warp 和 3 个 pipeline stage。与此前固定的 32 x 32 输出 tile 和未分组的二维 grid 相比，更大的 tile 使每个 program 完成更多工作，分组排序则改善了局部性。在统一 benchmark 下，这使 RTX 4090 上的平均时间从 16.907 ms 降至 9.803 ms，并使 RTX 5070 Ti Laptop GPU 上的平均时间从 78.710 ms 降至 32.518 ms。

Solution#

import torch
import triton
import triton.language as tl




@triton.jit
def matrix_multiplication_kernel(
    a, b, c, M, N, K, stride_am, stride_an, stride_bn, stride_bk, stride_cm, stride_ck,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr
):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_k = tl.cdiv(K, BLOCK_SIZE_K)
    num_pid_in_group = GROUP_SIZE_M * num_pid_k
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_SIZE_M)
    PID_M = first_pid_m + (pid % num_pid_in_group) % group_size_m
    PID_K = (pid % num_pid_in_group) // group_size_m
    MAX_N = tl.cdiv(N, BLOCK_SIZE_N)


    accumulated_block = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K), dtype=tl.float32)


    start_a_m = PID_M * BLOCK_SIZE_M
    a_m_offset = start_a_m + tl.arange(0, BLOCK_SIZE_M)


    start_b_k = PID_K * BLOCK_SIZE_K
    b_k_offset = start_b_k + tl.arange(0, BLOCK_SIZE_K)


    for n in tl.range(MAX_N):
        start_a_n = n * BLOCK_SIZE_N
        a_n_offset = start_a_n + tl.arange(0, BLOCK_SIZE_N)


        a_mn_mask = (a_m_offset[:, None] < M) & (a_n_offset[None, :] < N)
        a_mn_ptrs = a + a_m_offset[:, None] * stride_am + a_n_offset[None, :] * stride_an
        block_a_mn = tl.load(a_mn_ptrs, mask=a_mn_mask, other=0.0)


        start_b_n = n * BLOCK_SIZE_N
        b_n_offset = start_b_n + tl.arange(0, BLOCK_SIZE_N)




        b_nk_mask = (b_n_offset[:, None] < N) & (b_k_offset[None, :] < K)
        b_nk_ptrs = b + b_n_offset[:, None] * stride_bn + b_k_offset[None, :] * stride_bk


        block_b_nk = tl.load(b_nk_ptrs, mask=b_nk_mask, other=0.0)


        accumulated_block = tl.dot(block_a_mn, block_b_nk, accumulated_block, allow_tf32=False)
        # block_ab = tl.dot(block_a_mn, block_b_nk)
        # accumulated_block += block_ab


    c_mk_ptrs = c + a_m_offset[:, None] * stride_cm + b_k_offset[None, :] * stride_ck
    c_mk_mask = (a_m_offset[:, None] < M) & (b_k_offset[None, :] < K)
    tl.store(c_mk_ptrs, accumulated_block, mask=c_mk_mask)


# a, b, c are tensors on the GPU
def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, M: int, N: int, K: int):
    stride_am, stride_an = N, 1
    stride_bn, stride_bk = K, 1
    stride_cm, stride_ck = K, 1


    BLOCK_SIZE_M = 64
    BLOCK_SIZE_N = 32
    BLOCK_SIZE_K = 64
    GROUP_SIZE_M = 8


    grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(K, BLOCK_SIZE_K),)
    matrix_multiplication_kernel[grid](
        a, b, c, M, N, K, stride_am, stride_an, stride_bn, stride_bk, stride_cm, stride_ck,
        BLOCK_SIZE_M = BLOCK_SIZE_M,
        BLOCK_SIZE_N = BLOCK_SIZE_N,
        BLOCK_SIZE_K = BLOCK_SIZE_K,
        GROUP_SIZE_M = GROUP_SIZE_M,
        num_warps=4,
        num_stages=3,
    )

Test Code#

import importlib.util
import sys
from pathlib import Path


import torch




def load_implementation():
    source_path = Path(__file__).resolve().parents[2] / "src" / "triton" / "02_Matrix_Multiplication.py"
    spec = importlib.util.spec_from_file_location("matrix_multiplication_triton_minimal", source_path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Unable to load implementation: {source_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module




def main():
    matrix_a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="cuda")
    matrix_b = torch.tensor([[5.0, 6.0], [7.0, 8.0]], device="cuda")
    expected = torch.tensor([[19.0, 22.0], [43.0, 50.0]], device="cuda")
    actual = torch.empty_like(expected)


    load_implementation().solve(matrix_a, matrix_b, actual, 2, 2, 2)
    torch.cuda.synchronize()
    torch.testing.assert_close(actual, expected)
    print("Test passed.")




if __name__ == "__main__":
    try:
        main()
    except Exception as error:
        print(f"Test failed: {error}", file=sys.stderr)
        raise

Test Result#

PlatformStatusProblem SizeIterationsAverage TimeMinimum TimeMaximum TimePerformance
NVIDIA GeForce RTX 5070 Ti Laptop GPUPASSM=8192, N=6144, K=40961032.518 ms30.066 ms34.157 ms12679.673 GFLOPS
NVIDIA GeForce RTX 4090PASSM=8192, N=6144, K=4096109.803 ms8.994 ms10.005 ms42061.461 GFLOPS

PyTorch#

Approach#

PyTorch 实现刻意保持简短，因为该框架已经以内置 tensor operation 的形式提供了矩阵乘法：

torch.matmul(A, B, out=C)

在本题中，A 和 B 是 shape 分别为 M x N 和 N x K 的二维 tensor。torch.matmul 会检查内部维度是否兼容，计算 M x K 的矩阵乘积，并将结果直接写入传入的输出 tensor C。使用 out=C 很重要，因为题目要求最终结果必须存放在调用方提供的输出 tensor 中。

PyTorch 提供了多种乘法接口，但它们的含义并不完全相同：

  • A * B 执行逐元素乘法，并非矩阵乘法。
  • A @ B 是矩阵乘法的 Python 运算符语法，遵循 torch.matmul 语义。
  • torch.mm(A, B) 专用于二维矩阵输入。
  • torch.bmm(A, B) 用于处理成批的三维矩阵，不支持通用 broadcasting。
  • torch.matmul(A, B) 支持 vector、matrix、更高维 batch 和 broadcasting。
  • torch.einsum(...) 通过 index notation 显式表达 contraction，适用于更复杂的 tensor 关系。

torch.mm 和 A @ B 都可以表达本题的二维矩阵运算，但实际源码使用的是带 out argument 的 torch.matmul，因此本文档说明和测试的也是这一操作。

该实现中没有显式的 grid、thread index、tile size、mask 或同步。PyTorch 根据 tensor 的 device 和 dtype 分派操作，并由其 CUDA backend 选择底层 GPU 实现。所需的 solve signature 仍保留 M、N 和 K argument，但不会直接读取它们；实际执行的矩阵乘法由 tensor shape 决定。

Solution#

import torch


# A, B, C are tensors on the GPU
def solve(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, M: int, N: int, K: int):
    torch.matmul(A,B,out=C)

Test Code#

import importlib.util
import sys
from pathlib import Path


import torch




def load_implementation():
    source_path = Path(__file__).resolve().parents[2] / "src" / "pytorch" / "02_Matrix_Multiplication.py"
    spec = importlib.util.spec_from_file_location("matrix_multiplication_pytorch_minimal", source_path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Unable to load implementation: {source_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module




def main():
    matrix_a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="cuda")
    matrix_b = torch.tensor([[5.0, 6.0], [7.0, 8.0]], device="cuda")
    expected = torch.tensor([[19.0, 22.0], [43.0, 50.0]], device="cuda")
    actual = torch.empty_like(expected)


    load_implementation().solve(matrix_a, matrix_b, actual, 2, 2, 2)
    torch.cuda.synchronize()
    torch.testing.assert_close(actual, expected)
    print("Test passed.")




if __name__ == "__main__":
    try:
        main()
    except Exception as error:
        print(f"Test failed: {error}", file=sys.stderr)
        raise

Test Result#

PlatformStatusProblem SizeIterationsAverage TimeMinimum TimeMaximum TimePerformance
NVIDIA GeForce RTX 5070 Ti Laptop GPUPASSM=8192, N=6144, K=40961029.305 ms25.729 ms31.689 ms14069.752 GFLOPS
NVIDIA GeForce RTX 4090PASSM=8192, N=6144, K=4096107.469 ms7.188 ms7.559 ms55204.880 GFLOPS

References#

Acknowledgements#

本文档参考了 Du Ziyuan 的文章 CUDA Learning Journey [11] — A Detailed Explanation of Matrix Multiplication 中的讲解材料、图示和动画。原作采用 CC BY-NC-SA 4.0 许可；改编后的文档和在本地重新制作的媒体内容也按相同许可发布。改动包括围绕 LeetGPU 接口重新组织讲解，使讨论和代码与本项目的实现相匹配，并加入项目专用测试和实测结果。

关注 LeetGPU repository，了解更多内容。

文章分享

如果这篇文章对你有帮助，欢迎分享给更多人！

LeetGPU | 02_Matrix_Multiplication
https://github.com/HaoyangPing0324/LeetGPU
作者
平昊阳
发布于
2026-07-25
许可协议
CC BY-NC-SA 4.0
LeetCode | 两数之和
LeetGPU | 01_Vector_Addition
相关文章智能推荐
1
LeetGPU | 01_Vector_Addition
LeetGPULeetGPU Vector Addition 题解：CUDA、Triton 与 PyTorch 三种实现方式。
2
LeetCode | 合并两个有序链表
LeetCodeLeetCode 合并两个有序链表题解：双指针遍历比较节点值，dummy 虚拟头节点串联结果链表。
3
LeetCode | 两数相加
LeetCodeLeetCode 两数相加题解：模拟竖式逐位相加，dummy 虚拟头节点串联结果链表，维护 carry 进位。
4
LeetCode | 最长回文子串
LeetCodeLeetCode 最长回文子串题解：中心扩展法枚举每个奇偶中心向两侧扩展，时间复杂度 O(n²)，空间复杂度 O(1)。
5
llm-algo-leetcode 推理优化笔记
llm-algo-leetcode 推理优化Datawhale llm-algo-leetcode 推理优化课程的笔记导航页，汇总本系列各 Task 的学习笔记（GPU 架构与内存、KV Cache 优化技术等）。
随机文章随机推荐

评论区

分享你的想法，与大家交流讨论

Profile Image of the Author
平昊阳
乘长风，破巨浪， 展鸿图于未央！
--
总访问量
--
访客数
公告
欢迎来到我的个人博客！欢迎关注交流吖！
更多相关公告，见
「社交-留言」。
音乐
封面

音乐

暂未播放

0:000:00
暂无歌词
站点统计
文章
66
分类
16
标签
93
总字数
477,284
运行时长
0 天
最后活动
0 天前
苏公网安备32058102003077号 苏ICP备2026047448号
萌ICP备20266603号 喵ICP备20268044号 魔ICP备2026821142号-1
茶ICP备2026070376号 幸ICP备20260324号
酷AIA盟2026070324号 梵AIA盟2026070324号
三年之约 三年之约2026072103000139号
Copyright © 2026平昊阳. All Rights Reserved.RSSSitemapPowered byAstroFirefly

文章目录