How I Built a CUDA GEMM Kernel that Hits 932.9 GFLOPS

I'm currently an Engineering student at Pikes Peak State College hoping to transfer to the Colorado School of Mines to study Computer Science with a minor in Computer Engineering. This summer instead of taking classes I decided that I wanted to build the type of code that I wanted to write for a living: GPU kernels.

Specifically, I built a CUDA implementation of GEMM also known as General Matrix Multiplication which is the operation that is behind every neutral network. I first started making a naive version and profiled it using Nsight Compute to help me identify the bottleneck and rewrote the code with shared memory tiling.The final verison of my kernel hits 932.9 GFLOPS at 4096 x 4096 on a Tesla T4 GPU which was about 11.5% of the GPU's FP32 peak and about 26% of what cuBLAS achieves with the same hardware

This post will walk through the entire process including why naive GEMM is slow, what does tiling do with the data from the profiler supporting every decision I make along the way.

Why GEMM?

GEMM (General Matrix Multiplication) computes C = A * B, where A is an M * K matrix and B is a K * N matrix which results in a M * K matrix. While this may seem like a simple linear algebra operation it is the most common computation in deep learning which every linear layer, attention projection and fully connected layer ties to.

This is also why frameworks like PyTorch don't do their own matrix multiplication but instead call libraries like cuBLAS which has been personall been hand-tuned by engineers over years. I learned that understanding why cuBLAS is fast by making a slower version and figuring out how to close the gap step by step is one of the best ways to learn GPU programming.

Naive GEMM Demonstration

One of the first obvious GPU implementations is to assign one output to one threads.

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

    if (row < M && col < N) {
        float sum = 0.0f;

        for (int k = 0; k < K; ++k) {
            sum += A[row * K + k] * B[k * N + col];
        }

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

Each entire thread reads an entire row of A and an entire column from B from global memory which is the GPU largest memory but also its slowest. For specifically the K dimension that is 2 * K global memory reads per output element. Even though this work it becomes a problem since the same threads that are near each other re-read the same data with zero sharing between them. This is where tiling comes in.

Tiling

Tiling breaks matrix A and matrix B into smaller square tiles loading one tile into shared memory (a small, fast, per thread block memory space) which lets each thread within that thread block use the same tile as many times as it wants before moving on to the next tile.

Here is a visualization for context:

And here is the actual kernel:

#define TILE_WIDTH 32

__global__ void tiled_GEMM(const float *A, const float *B, float *C, int M, int K, int N){
    __shared__ float sharedA[TILE_WIDTH][TILE_WIDTH];
    __shared__ float sharedB[TILE_WIDTH][TILE_WIDTH];

    int row = blockIdx.y * TILE_WIDTH + threadIdx.y;
    int col = blockIdx.x * TILE_WIDTH + threadIdx.x;
    float sum = 0.0f;
    int numTiles = (K + TILE_WIDTH - 1) / TILE_WIDTH;

    for (int t = 0; t < numTiles; ++t) {
        int aCol = t * TILE_WIDTH + threadIdx.x;
        sharedA[threadIdx.y][threadIdx.x] = (row < M && aCol < K) ? A[row * K + aCol] : 0.0f;

        int bRow = t * TILE_WIDTH + threadIdx.y;
        sharedB[threadIdx.y][threadIdx.x] = (bRow < K && col < N) ? B[bRow * N + col] : 0.0f;

        __syncthreads();
        for (int i = 0; i < TILE_WIDTH; ++i)
            sum += sharedA[threadIdx.y][i] * sharedB[i][threadIdx.x];
        __syncthreads();
    }

    if (row < M && col < N) C[row * N + col] = sum;
}

The trade here is to pay the slow global memory cost once per tile the read from shared memory using TILE_WIDTH times. The two __syncthreads I put here especially matter because the first one makes sure each thread has finished writing to shared memory before anyone reads it and the second makes sure a thread doesn't overwrite the tile while a slower thread is still using it.
This is also the same idea that cuBLAS is built on although mine isn't as tuned as it.

Results

The tile size I used wasn't a guess. I benchmarker TILE_WIDTH at 16 vs TILE_WIDTH = 32 using Nsight Compute at 2048 * 2048 and these were the results:

Tile Size Comparison

Metric	TILE=16	TILE=32
Kernel duration (Nsight Compute)	46.26 ms	41.89 ms
Achieved occupancy	99.63%	99.97%
DRAM read bandwidth	92.49 GB/s	39.43 GB/s

The TILE= 32 won with it being about 9.4% faster showing high occupancy. The lower DRAM bandwidth number showed less total data movement since each element is reused more times per tile at a larger scale. I also checked for shared bank conflicts just in case and found none.

Full Benchmark Table Validated

Performance Comparison

Matrix Size	CPU (ms)	Naive CUDA (ms)	Tiled CUDA (ms)	PyTorch CUDA (ms)	Tiled GFLOPS
256×256	20.647	0.218	0.207	0.057	162.1
512×512	203.569	1.005	0.599	0.116	448.1
1024×1024	3509	5.574	5.355	0.608	401.0
2048×2048	65316.9	69.062	28.7*	4.203	598.6
4096×4096	813709	330.641	147.31*	38.039	932.9

At 4096 x 4096 the tiled kernel reaches 932.9 GFLOPS which is about 11.5% of the T4's FP32 hardware peak which is 8.1 TFLOPS and about 26% of what PyTorch cuBLAS backend hits on the same hardware. This gap is due to cuBLAS using register blocking, double buffering and hand-tuned instruction schedule that my kernel doesn't have. Being able to close that gap is the exact kind of work I want to be able to do next.

What is Next?

My kernel is currently callable from Python since I wrapped it in pybind11 so tiled_gemm runs the CUDA kernel on NumPy arrays.

One project planned:

GPU Kernel Library: Flash attention varient,SoftMax,LayerNorm and multi-GPU GEMM benmarked against cuDNN planned for next summer

If you made it this far thank you for reading. The Full Code is on Github: https://github.com/cwilliams-systems/cuda-ml-accelerator

DEV Community