Inside NVIDIA GPUs: Anatomy of high performance matmul kernels

From GPU architecture and PTX/SASS to warp-tiling and deep asynchronous tensor core pipelines

October 10, 2023

This is a sample blog post to visualize the components. It includes blockquotes, images, code blocks, and references.

This is a blockquote inside the blog post. It is used to highlight important information or quotes in a visually distinct way.

Measure before optimizing

Wall-clock timing around a launch measures the launch, not the kernel. Use CUDA events, synchronize on both ends, and convert to TFLOP/s so numbers stay comparable across shapes.

benchmark.py

import torch


def benchmark(fn, *args, warmup: int = 10, iters: int = 100) -> float:
    """Time a CUDA op in milliseconds, ignoring warm-up and launch noise."""
    for _ in range(warmup):
        fn(*args)

    start, end = torch.cuda.Event(True), torch.cuda.Event(True)
    torch.cuda.synchronize()
    start.record()
    for _ in range(iters):
        fn(*args)
    end.record()
    torch.cuda.synchronize()

    return start.elapsed_time(end) / iters


if __name__ == "__main__":
    a = torch.randn(4096, 4096, device="cuda", dtype=torch.float16)
    b = torch.randn(4096, 4096, device="cuda", dtype=torch.float16)

    ms = benchmark(torch.matmul, a, b)
    tflops = (2 * 4096 ** 3) / (ms * 1e-3) / 1e12
    print(f"matmul: {ms:.3f} ms | {tflops:.1f} TFLOP/s")

The CPU baseline

A blocked CPU version is worth writing first. The blocking is the same idea the kernel uses later: keep the working set small enough that it stays in fast memory.

matmul.cpp

#include <algorithm>
#include <cstddef>
#include <vector>

// Cache-blocked reference matmul: C = A * B, all row-major.
// The j-loop stays innermost so B is walked along rows, never down columns.
template <std::size_t BLOCK = 64>
void matmul(const std::vector<float>& a,
            const std::vector<float>& b,
            std::vector<float>& c,
            std::size_t n) {
    for (std::size_t ii = 0; ii < n; ii += BLOCK) {
        for (std::size_t kk = 0; kk < n; kk += BLOCK) {
            for (std::size_t jj = 0; jj < n; jj += BLOCK) {
                const std::size_t i_max = std::min(ii + BLOCK, n);
                const std::size_t k_max = std::min(kk + BLOCK, n);
                const std::size_t j_max = std::min(jj + BLOCK, n);

                for (std::size_t i = ii; i < i_max; ++i) {
                    for (std::size_t k = kk; k < k_max; ++k) {
                        const float a_ik = a[i * n + k];
                        for (std::size_t j = jj; j < j_max; ++j) {
                            c[i * n + j] += a_ik * b[k * n + j];
                        }
                    }
                }
            }
        }
    }
}

Here is an image of Roman Reigns:

Tribal Chief Roman Reigns

Shared-memory tiling on the GPU

Each block stages a TILE × TILE slice of both operands in shared memory, so every element pulled from global memory is reused TILE times.

matmul.cu

#include <cstdio>
#include <cuda_runtime.h>

constexpr int TILE = 32;

// One 32x32 output tile per block; each thread owns a single element of C.
__global__ void matmul_tiled(const float* __restrict__ A,
                             const float* __restrict__ B,
                             float* __restrict__ C,
                             int N) {
    __shared__ float As[TILE][TILE];
    __shared__ float Bs[TILE][TILE + 1];  // pad by one to dodge bank conflicts

    const int row = blockIdx.y * TILE + threadIdx.y;
    const int col = blockIdx.x * TILE + threadIdx.x;

    float acc = 0.0f;

    for (int t = 0; t < N / TILE; ++t) {
        As[threadIdx.y][threadIdx.x] = A[row * N + t * TILE + threadIdx.x];
        Bs[threadIdx.y][threadIdx.x] = B[(t * TILE + threadIdx.y) * N + col];
        __syncthreads();

        #pragma unroll
        for (int k = 0; k < TILE; ++k) {
            acc = fmaf(As[threadIdx.y][k], Bs[k][threadIdx.x], acc);
        }
        __syncthreads();  // don't refill the tile while peers are still reading
    }

    C[row * N + col] = acc;
}

void launch(const float* A, const float* B, float* C, int N, cudaStream_t stream) {
    const dim3 block(TILE, TILE);
    const dim3 grid(N / TILE, N / TILE);

    matmul_tiled<<<grid, block, 0, stream>>>(A, B, C, N);

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        std::fprintf(stderr, "launch failed: %s\n", cudaGetErrorString(err));
    }
}

The padding in Bs[TILE][TILE + 1] is not cosmetic — without it, every thread in a warp hits the same shared-memory bank on a column read and the access serializes.

Dropping to warp primitives

Once a reduction fits inside a single warp, shared memory and barriers are pure overhead.__shfl_down_sync moves the partial sums directly between registers.

warp_reduce.cuh

// Warp-level reduction: no shared memory, no barriers, just register shuffles.
__device__ __forceinline__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;  // lane 0 holds the total
}

This is a blockquote inside the blog post. It is used to highlight important information or quotes in a visually distinct way.

Get notified when I publish a new post.

References

1. Next.js Documentation: https://nextjs.org/docs

2. Tailwind CSS Documentation: https://tailwindcss.com/docs

3. Next Themes Documentation: https://next-themes.com/docs

4. CUDA C++ Programming Guide: https://docs.nvidia.com/cuda/cuda-c-programming-guide/