GEMM Optimization — infra.bea.ai

GEMM (General Matrix Multiply) is the most performance-critical kernel in deep learning — every linear layer, every attention score, every FFN computation is a GEMM. Understanding GEMM optimization reveals every key GPU programming technique: tiling, data reuse, memory hierarchy exploitation, and Tensor Core utilization.

1. Why GEMM Matters

For a model with $N$ parameters, a single forward pass does roughly $2N$ FLOPs of GEMM. For LLaMA-2-7B with a 2K sequence: ~28 TFLOP per forward pass, ~95% of which is GEMM.

GEMM is THE kernel that determines GPU utilization for training and prefill. If GEMM is slow, everything is slow.

The Problem

Compute $C = A \times B$ where $A \in \mathbb{R}^{M \times K}$ , $B \in \mathbb{R}^{K \times N}$ , $C \in \mathbb{R}^{M \times N}$ :

C_{ij} = \sum_{k=0}^{K-1} A_{ik} \cdot B_{kj}

FLOPs: $2MKN$
Data: $2(MK + KN + MN)$ bytes (FP16)
Arithmetic intensity: $\frac{2MKN}{2(MK + KN + MN)}$

For square matrices ( $M=N=K=4096$ ): AI ≈ 1365 → compute-bound on H100 (ridge point = 295). GEMM is one of the few operations that can actually saturate GPU compute.

2. V0: Naive Implementation

Each thread computes one output element:

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

    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;
    }
}

Performance: ~2% of cuBLAS. Each output element reads an entire row of A and column of B from HBM — massive redundancy.

3. V1: Shared Memory Tiling

Load tiles of A and B into shared memory, compute, iterate along K:

#define TILE 32

__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 row = blockIdx.y * TILE + threadIdx.y;
    int col = blockIdx.x * TILE + threadIdx.x;
    float sum = 0.0f;

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

        // Compute from shared memory
        for (int k = 0; k < TILE; k++) {
            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
        }
        __syncthreads();
    }

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

Performance: ~20-30% of cuBLAS. Tile size 32 → 32× fewer HBM reads. But still limited by shared memory bandwidth.

4. V2: Register-Level Tiling (Thread Tile)

Each thread computes not one output element but a small tile (e.g., 8×8) using registers:

Thread block tile: 128×128 (in shared memory)
Thread tile: 8×8 (in registers)
Each thread:
  - Loads a column of As and row of Bs into registers
  - Computes 8×8 outer product
  - Accumulates across K iterations

// Each thread maintains an 8×8 accumulator in registers
float acc[8][8] = {0};

for (int t = 0; t < K; t += TK) {
    // Load tile to shared memory (cooperative)
    load_tile_to_shared(As, Bs, t);
    __syncthreads();

    for (int k = 0; k < TK; k++) {
        // Load thread's column of A and row of B from shared memory
        float a_frag[8], b_frag[8];
        for (int i = 0; i < 8; i++) a_frag[i] = As[threadRow*8 + i][k];
        for (int j = 0; j < 8; j++) b_frag[j] = Bs[k][threadCol*8 + j];

        // 8×8 outer product in registers
        for (int i = 0; i < 8; i++)
            for (int j = 0; j < 8; j++)
                acc[i][j] += a_frag[i] * b_frag[j];
    }
    __syncthreads();
}

Key insight: The 8×8 accumulator lives in registers — each element of a_frag is reused 8 times and each element of b_frag is reused 8 times. This is register-level data reuse on top of shared memory-level reuse.

Performance: ~60-70% of cuBLAS.

5. V3: Double Buffering

Overlap loading the next tile with computing the current tile:

// Two sets of shared memory buffers
__shared__ float As[2][BM][BK];
__shared__ float Bs[2][BK][BN];

// Load first tile
load_tile(As[0], Bs[0], k=0);
__syncthreads();

for (int t = BK; t < K; t += BK) {
    int curr = (t/BK) % 2;
    int next = 1 - curr;

    // Async load next tile while computing current
    load_tile_async(As[next], Bs[next], t);
    compute_tile(As[curr], Bs[curr], acc);
    __syncthreads();
}

Performance: ~75-85% of cuBLAS. Memory latency hidden behind compute.

L2 Cache and Wave Analysis

Beyond tiling and registers, production GEMM kernels also consider L2 cache behavior:

Wave model: When a GEMM launches, blocks are assigned to SMs in “waves.” If the grid has more blocks than SMs, the first wave runs, finishes, then the second wave starts. Between waves, L2 cache contents from the first wave may still be warm.

L2 cache reuse: If Block (0,0) and Block (0,1) read the same row of A, and they’re scheduled in the same wave on the same SM cluster, the second read may hit L2 instead of HBM. Production GEMM kernels (cuBLAS) choose block-scheduling order to maximize L2 hits.

Swizzle patterns: cuBLAS reorders block indices so that blocks reading similar data are scheduled on nearby SMs (which share L2 partitions). This “swizzled” scheduling can improve L2 hit rates by 20-40% for medium-sized GEMMs.

L2 optimization matters most for medium matrices (1K-4K). For very large matrices, tiles are too large to fit in L2 anyway. For very small matrices, the entire A/B may fit in L2. The sweet spot where L2 optimization helps is the range used by most Transformer FFN layers.

6. Tensor Core GEMM

Tensor Cores operate on small matrix tiles (16×8×16 on Hopper). Instead of scalar FMA, each Tensor Core MMA computes a matrix multiply-accumulate per cycle:

#include <mma.h>
using namespace nvcuda::wmma;

// Declare fragments (Tensor Core registers)
fragment<matrix_a, 16, 16, 16, half, row_major> a_frag;
fragment<matrix_b, 16, 16, 16, half, col_major> b_frag;
fragment<accumulator, 16, 16, 16, float> c_frag;

fill_fragment(c_frag, 0.0f);

for (int k = 0; k < K; k += 16) {
    load_matrix_sync(a_frag, A + row*K + k, K);
    load_matrix_sync(b_frag, B + k*N + col, N);
    mma_sync(c_frag, a_frag, b_frag, c_frag);
}

store_matrix_sync(C + row*N + col, c_frag, N, mem_row_major);

Performance: With proper tiling around Tensor Cores, ~90-95% of cuBLAS. In practice, use CUTLASS or cuBLAS directly — they handle all the complexity.

7. The Production Stack

In AI infrastructure, you almost never write GEMM from scratch:

Library	When to use	Performance
cuBLAS	Standard GEMM, always available	95-100% of peak
CUTLASS	Custom GEMM variants, epilogue fusion	90-98% of peak
Triton	Prototyping custom kernels in Python	80-95% of peak
torch.mm / torch.matmul	PyTorch (calls cuBLAS under the hood)	Same as cuBLAS

Why learn GEMM optimization if you’ll use cuBLAS? Because:

Every tiling, memory hierarchy, and warp-level optimization in GEMM applies to other kernels
Custom epilogues (GEMM + activation, GEMM + bias + residual) require understanding GEMM internals
Debugging performance issues requires understanding what cuBLAS is doing

8. Performance Summary

Version	Key technique	% of cuBLAS
V0 Naive	One element per thread	~2%
V1 Shared memory tiling	Block-level data reuse	~25%
V2 Register tiling	Thread-level outer product	~65%
V3 Double buffering	Overlap load and compute	~80%
+ Tensor Cores	WMMA / MMA instructions	~95%
cuBLAS	All of the above + autotuning	100%

Self-Check

Explain why GEMM is compute-bound for large matrices and calculate arithmetic intensity
Implement tiled GEMM with shared memory (V1)
Explain register tiling and why the outer product pattern enables register reuse
Explain double buffering and how it hides memory latency
Explain what Tensor Cores do at the hardware level (MMA instruction)
Explain when to use cuBLAS vs CUTLASS vs Triton vs hand-written CUDA

References

How to Optimize a CUDA Matmul Kernel (Simon Boehm) — excellent step-by-step guide
CUTLASS — NVIDIA CUDA Templates for Linear Algebra
cuBLAS Documentation
NVIDIA — Optimizing GEMM with Tensor Cores