Home › Neural network training & inference › GPU kernels
algorithm

Tiled matrix multiplication

Problem Developers writing GPU or CPU kernels for matrix multiplication, the main workload of neural networks, find naive code runs at a tiny fraction of hardware speed because memory cannot keep up with the arithmetic.

Block GEMM into tiles held in shared memory and registers, so each byte loaded from DRAM feeds many FMAs; the core of fast GEMM.

2·M·N·K FLOPs; DRAM traffic about M·N·K·(1/BN + 1/BM) elements with BM×BN block tiles, instead of 2·M·N·K for the naive loop
gemmcudashared-memoryregister-blockingtensor-coresroofline
Overview

A naive GEMM loads two operands per FMA and runs at about 1% of peak. Tiling creates reuse: a thread block stages a BM×BK slice of A and a BK×BN slice of B in shared memory, and each staged element feeds BN (or BM) FMAs. Inside the block, each thread keeps a TM×TN micro-tile of C in registers and reuses every shared-memory load TN (or TM) times.

Key facts

Use when

  • Writing a GEMM, conv-as-GEMM or attention kernel, or fusing epilogues that libraries can't.
  • Dense linear algebra in scientific computing: BLAS-level GEMM underlies LAPACK-style solvers (Goto and van de Geijn), alongside the neural-network use.

Avoid when

  • Standard shapes: cuBLAS, CUTLASS and Triton are already tuned.
  • Decode GEMV (M ≤ 16): bandwidth-bound, so cut bytes (weight quantization).

Deep dive

How it works

In: A (M×K), B (K×N) in DRAM, tile sizes BM, BN, BK, TM, TN · Out: C = A·B (M×N) · State: none between calls; per block a shared-memory tile pair and per thread a TM×TN register accumulator · Cost: 2·M·N·K FLOPs; DRAM traffic about M·N·K·(1/BN + 1/BM) elements

C (M×N) is split into BM×BN block tiles, one per thread block.
for k0 in 0..K step BK:
    cooperatively load A[BM×BK] and B[BK×BN] → shared memory
    __syncthreads
    for k in 0..BK:                         ← register level
        a[TM] = As[k][my rows], b[TN] = Bs[k][my cols]
        acc[TM][TN] += a ⊗ b                (TM·TN FMAs per TM+TN shared loads)
    __syncthreads
write acc → C

Hierarchy on a modern GPU:

  • DRAM/L2 → shared memory: block tile, e.g. 128×128×32. The reuse factor sets DRAM traffic.
  • Shared → registers: warp tile, then thread tile (8×8 fp32 = 64 accumulators per thread). Reuse here sets shared-memory bandwidth, which is about 32 lanes × 4 B per clock per SM.
  • Tensor cores: mma.sync.m16n8k16 (Ampere, one warp) or wgmma 64×N×16 (Hopper, a 4-warp warpgroup reading operands straight from shared memory).
  • Pipelining: cp.async (Ampere) or TMA (Hopper) streams the next k-tile into a second or third shared buffer while the current one computes.

Code

CUDA, fp32, row-major, 2D block tiling plus register tiling. 256 threads compute a 64×64 tile, 4×4 outputs each.

#include <cuda_runtime.h>

template <int BM, int BN, int BK, int TM, int TN>
__global__ void __launch_bounds__((BM / TM) * (BN / TN))
sgemm_tiled(int M, int N, int K, const float* __restrict__ A,
            const float* __restrict__ B, float* __restrict__ C) {
    constexpr int THREADS = (BM / TM) * (BN / TN);
    __shared__ float As[BK][BM + 4];   // A stored transposed; +4 pad kills store bank conflicts
    __shared__ float Bs[BK][BN];

    const int tid = threadIdx.x;
    const int tr = tid / (BN / TN), tc = tid % (BN / TN);  // this thread's micro-tile
    const int row0 = blockIdx.y * BM, col0 = blockIdx.x * BN;

    float acc[TM][TN] = {};
    float a[TM], b[TN];

    for (int k0 = 0; k0 < K; k0 += BK) {
        for (int i = tid; i < BM * BK; i += THREADS) {     // load A tile, zero-fill edges
            int r = i / BK, c = i % BK, gr = row0 + r, gc = k0 + c;
            As[c][r] = (gr < M && gc < K) ? A[(size_t)gr * K + gc] : 0.f;
        }
        for (int i = tid; i < BK * BN; i += THREADS) {     // load B tile (coalesced along N)
            int r = i / BN, c = i % BN, gr = k0 + r, gc = col0 + c;
            Bs[r][c] = (gr < K && gc < N) ? B[(size_t)gr * N + gc] : 0.f;
        }
        __syncthreads();
#pragma unroll
        for (int k = 0; k < BK; ++k) {
#pragma unroll
            for (int i = 0; i < TM; ++i) a[i] = As[k][tr * TM + i];
#pragma unroll
            for (int j = 0; j < TN; ++j) b[j] = Bs[k][tc * TN + j];
#pragma unroll
            for (int i = 0; i < TM; ++i)
#pragma unroll
                for (int j = 0; j < TN; ++j) acc[i][j] = fmaf(a[i], b[j], acc[i][j]);
        }
        __syncthreads();                                   // don't overwrite tiles still being read
    }
#pragma unroll
    for (int i = 0; i < TM; ++i)
#pragma unroll
        for (int j = 0; j < TN; ++j) {
            int r = row0 + tr * TM + i, c = col0 + tc * TN + j;
            if (r < M && c < N) C[(size_t)r * N + c] = acc[i][j];  // epilogue: fuse bias/act here
        }
}

void sgemm(int M, int N, int K, const float* A, const float* B, float* C, cudaStream_t s) {
    constexpr int BM = 64, BN = 64, BK = 8, TM = 4, TN = 4;     // 256 threads, 4.2 KB smem
    dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
    sgemm_tiled<BM, BN, BK, TM, TN><<<grid, (BM / TM) * (BN / TN), 0, s>>>(M, N, K, A, B, C);
}

Next steps toward cuBLAS speed, in order of payoff: float4 vectorized global loads; a larger thread tile (8×8) with warp tiling; double-buffered shared memory with cp.async; then tensor cores via mma.sync or CUTLASS/CuTe.

Complexity and constants

  • FLOPs = 2MNK. At M=N=K=4096: 137 GFLOP, about 0.14 ms at 989 TFLOPS. Operand bytes (bf16) are only 100 MB, so a good kernel is compute-bound.
  • DRAM traffic with block tiles: each A element is loaded N/BN times and each B element M/BM times. L2 catches part of that if blocks are rasterized (swizzled) so neighbors share panels.
  • Shared memory per SM: up to 164 KB (A100) or 228 KB (H100). Registers: 64K × 32-bit per SM, at most 255 per thread. Occupancy is traded against tile size.
  • Wave quantization: H100 SXM has 132 SMs. A GEMM with 133 tiles needs 2 waves, so about half the second wave sits idle. Pick tile sizes (or use stream-K) so the tile count is just under a multiple of the SM count.

Efficiency

  • Cost: 2·M·N·K FLOPs: 137 GFLOP at M = N = K = 4096, about 0.14 ms at 989 TFLOPS bf16 on an H100 SXM, with only 100 MB of operands, so a good kernel is compute-bound. A 128×128 fp32 block tile gives 32 FLOP/byte against a ridge point of about 300 FLOP/byte (989 TFLOPS / 3.35 TB/s). A naive kernel sits at 0.25 FLOP/byte and about 1% of peak (order of). Bound: DRAM and shared-memory bandwidth until tiled, then tensor-core throughput.
  • Levers:
    • block and register tiling — one stage cuts DRAM traffic by the tile size; Boehm's fp32 worklog reached about 94% of cuBLAS after ten steps — every GEMM
    • tensor cores (mma, wgmma) with async copy (cp.async, TMA) pipelines — overlap loads with math; needed to reach bf16/fp8 peak — Ampere and Hopper
    • swizzled block order for L2 reuse, stream-K for wave quantization — keeps all 132 SMs busy (Osama et al.) — shapes just over a multiple of the SM count
    • epilogue fusion (bias, activation, quantize) in registers — saves a full read and write of C — GEMM followed by elementwise ops
    • pad M, N, K to tile and 8/16-element multiples — no tail waves or misaligned loads — model-sizing time
  • Naive → standard → state of the art:
    • triple loop, two loads per FMA · about 1% of peak
    • shared-memory block tiles with register micro-tiles · tens of percent of cuBLAS fp32
    • warp tiling or wgmma with TMA pipelines and stream-K (CUTLASS, Triton) · near cuBLAS and peak tensor-core throughput

Implementation notes

  • Store A transposed in shared memory so the inner loop reads contiguous As[k][...]. Pad rows (+4 floats) to spread the transposing stores across all 32 banks.
  • Unroll fully. The accumulators must be compile-time indexed or they spill to local memory (check -Xptxas -v for spills).
  • Make M, N, K multiples of the tile and of 8/16 elements (tensor-core alignment), or pad at the model level. Hidden sizes that are multiples of 64/128 are not an accident.
  • Fuse the epilogue: bias, activation, residual, scaling, fp8/int8 quantization while the tile is still in registers. That saves a full read and write of C.
  • CPUs use the same idea (Goto/BLIS): pack panels so they fit in L2 and L1, with a register microkernel of mr×nr accumulators (e.g. 6×16 fp32 with AVX2 = 12 ymm registers).

Pitfalls

  • Missing the second __syncthreads(): fast threads overwrite the tile while slow ones still read it. The result is wrong only intermittently.
  • Uncoalesced loads. Consecutive threads must touch consecutive addresses. Loading A along M instead of K drops bandwidth by up to 8-32x.
  • Bank conflicts: a 32-float row stride maps a whole column to one bank (32-way conflict).
  • Register spills from oversized thread tiles. They show up as local-memory traffic in Nsight Compute.
  • Benchmarking with one launch: the first call includes module load. Warm up, then use CUDA events over many iterations.

Variants and alternatives

Split-Kpartition K across blocks and reduce at the end, for small M·N with large K.
Stream-Kgive each SM an equal share of MAC iterations to remove wave quantization.
Strassen-like algorithmsfewer FLOPs, worse numerics and memory behavior. Not used for DL GEMMs.
Triton tl.dotwrite block-level code and let the compiler handle the warp, thread and tensor-core tiling.

In production

cuBLAS/cuBLASLt, CUTLASS/CuTe (the basis of many FlashAttention and fp8 kernels), rocBLAS/hipBLASLt, Triton, OpenBLAS/BLIS/MKL on CPU, Apple AMX/MPS.

Implementations

namelangrepoinstallverdictlicencehealthas_ofnotes
faerRustsarah-quinones/faer-rscargo add faer————dense linear algebra with its own blocked, SIMD GEMM
gemmRustsarah-quinones/gemmcargo add gemm————packed, SIMD GEMM kernels for f16/f32/f64/complex
matrixmultiplyRustbluss/matrixmultiplycargo add matrixmultiply————BLIS-style packed GEMM, backend of ndarray
CUTLASS / CuTeC++ / CUDANVIDIA/cutlassn/a (header-only; clone the repo)————tiled tensor-core GEMM templates, stream-K
TritonPython / MLIRtriton-lang/tritonpip install triton————block-level tl.dot; compiler does warp and tensor-core tiling

Sources

  • Goto and van de Geijn, Anatomy of High-Performance Matrix Multiplication, ACM TOMS, 2008. paper
  • Volkov and Demmel, Benchmarking GPUs to Tune Dense Linear Algebra, SC 2008. paper
  • Osama, Merrill, Cecka, Garland, Owens, Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU, PPoPP 2023. paper
  • NVIDIA, CUTLASS documentation and Matrix Multiplication Background User's Guide. CUTLASS · guide · code
  • Simon Boehm, How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog, 2022. post

Builds on

gpu memory coalescing banking · plannedCoalesced global loads, shared memory bank conflicts, L2 and HBM bandwidth on GPUs
CUDA and GPU languagesHow should you program a GPU, and when is it worth writing custom kernels rather than using existing libraries? This guide helps pick between libraries, Triton and CUDA C++ for numeric workloads.

Used by

FlashAttentionTraining or running a transformer on long sequences is slow and runs out of GPU memory, because the attention step builds and moves a huge intermediate table. Engineers want exact attention that is faster and leaner.
gpu health checks · plannedDetecting bad GPUs/links before and during jobs: ECC/XID errors, NCCL tests, throughput probes, draining
kernel autotuning · plannedSearching tile sizes and launch parameters per shape/hardware, caching the best configs
Kernel fusionGPU programs made of many small operations, such as normalisation, activations and optimizer steps, run well below hardware capability because each step re-reads and rewrites memory. Engineers want to speed these up.
tensor compiler scheduling · plannedTile, fuse and lay out tensor ops for a target (XLA, TVM, Triton schedules), plus autotuning
tensor network contraction · plannedContract tensor networks with good contraction orders; MPS and slicing for circuit simulation
Tensor parallelismA single layer of a very large neural network is too big or too slow for one GPU. Engineers need to divide the layer's work across several GPUs on a fast link.

Systems (15):

gpuai acceleratorcuda platformedge ai camerallm pretrainingllm inferenceon device inferenceml compilertensor runtimelinear solver suiteeigensolverarray computingquantum chemistryles dns solverearthquake sim

Alternatives and related

Mixed-precision trainingNeural network training is slow and memory-hungry in full 32-bit arithmetic. Engineers want to use faster, smaller number formats on modern GPUs without losing accuracy or hitting numerical failures.Roofline modelAn engineer optimising a compute-heavy kernel needs to know whether it is limited by arithmetic or by memory bandwidth, and how much faster it could possibly get, before deciding what to work on.