Split-K GEMM: Trade Extra Reduction for More Parallelism
REFERENCE CONTEXT · CUDA 12.2; Turing sm_75 and later concepts
Core Concept
Split a long K reduction into P partial GEMMs, then merge the partial output. Split-K is useful when M and N are small, K is large, and the ordinary GEMM launches too few output tiles to occupy the GPU.
Why It Affects Performance
A conventional GEMM assigns one output tile to one block. For small M/N, the grid can be too small even when K is huge. Split-K creates additional independent blocks along K, buying parallelism at the cost of workspace traffic and a reduction.
Example or Pseudocode
// C_partial[p, m, n] = sum over K slice p
partial_gemm<<<grid_mn_times_p, block>>>(A, B, workspace, split_k);
reduce_split_k<<<grid_mn, block>>>(workspace, C, split_k);Nonlinear epilogues such as ReLU must run after the partial results are merged, not on each partial independently.
Profiling and Performance Judgment
Compare ordinary and Split-K kernels with Nsight Systems/Compute. Verify that the original kernel is grid-limited, then account for partial-output writes, reduction reads, synchronization, and workspace size. Stop increasing P when reduction/workspace cost exceeds the gain from added waves.
Five-Minute Exercise
For M=N=128, K=8192, output tiles 64×64, and P=16: how many ordinary blocks vs Split-K blocks? What is each split K? What FP32 workspace is required? If ordinary GEMM is 80 μs and Split-K partial+merge is 28+12 μs, what is the speedup? Can ReLU run before the merge?
Show answer
Ordinary GEMM launches 4 output-tile blocks. Split-K launches 4×16=64 partial blocks. Each split covers K=512. Workspace is 16×128×128×4 bytes = 1 MiB; a write plus merge read moves roughly 2 MiB. Runtime becomes 40 μs vs 80 μs, or 2× faster. ReLU cannot be applied before the merge because ReLU(sum partials) is not equal to sum(ReLU(partials)).