Ritesh Yadav
Ritesh Yadav
/device-software/2d-block-tiling-gemm
GEMM

2D Block Tiling: Scaling Matrix Multiplication on NVIDIA H100

In the last few days, after writing the 1D kernel in CUDA, I started learning about 2D block tiling. In the 1D approach, each thread calculated a vertical strip of output values. That was already an exciting step forward because caching values from matrix B inside registers saved a lot of shared memory traffic, taking performance up to around 42% of cuBLAS on an NVIDIA T4.

Still, looking at how the hardware actually works, 1D tiling leaves an obvious inefficiency on the table. Each thread only reuses data along one direction. While matrix B values sit comfortably in registers, values from matrix A still have to be fetched repeatedly from shared memory across neighboring threads.

That led me to explore 2D block tiling in my gemm-kernels repository, benchmarking it on an NVIDIA H100 GPU using Modal. The kernel reached close to 19 TFLOPS on standard FP32 math (roughly 36% of cuBLAS with TF32 disabled). Here is the story of how 2D block tiling works, what I observed along the way, and why it changes both thread counts and memory access patterns so dramatically.


1. The Core Shift: From a 1D Strip to a 2D Patch

In 1D block tiling, a single thread is responsible for computing TM elements (say, 8 rows by 1 column). In 2D block tiling, each thread takes ownership of a two-dimensional patch of output values: TM rows by TN columns.

In my kernel, both TM and TN are set to 8. This means every thread computes an 8 by 8 patch, producing 64 output values of matrix C all by itself.

Why does this matter? It comes down to arithmetic intensity inside the thread:

  1. In 1D tiling (TM = 8, TN = 1), to compute 8 results, the thread fetches 8 values from tile A and 1 value from tile B. That is 9 shared memory reads to perform 8 multiply-accumulate operations.
  2. In 2D tiling (TM = 8, TN = 8), the thread fetches 8 values from tile A and 8 values from tile B into registers. With those 16 values in hand, the thread executes an 8 by 8 outer product, performing 64 multiply-accumulate operations.

Instead of needing roughly 1.1 shared memory reads per computation, we only need 16 reads for 64 computations: a mere 0.25 reads per multiply-accumulate. That is a 4x reduction in shared memory bandwidth pressure.


2. Block Dimensions: Why We Only Need 64 Threads

One of the most interesting surprises when moving from 1D to 2D tiling is how dramatically the block size shrinks.

Here are the constants used in the kernel:

constexpr int BM = 64; // Block output tile height
constexpr int BN = 64; // Block output tile width
constexpr int BK = 8;  // Step size along the inner K dimension
constexpr int TM = 8;  // Number of output rows per thread
constexpr int TN = 8;  // Number of output columns per thread

A block covers a BM x BN (64 by 64) tile of matrix C, which equals 4096 elements.

  • In the 1D kernel, each thread handled 8 elements. Covering 4096 elements required 4096 / 8 = 512 threads per block.
  • In the 2D kernel, each thread handles 8 * 8 = 64 elements. Covering 4096 elements requires only 4096 / 64 = 64 threads per block.

A block of 64 threads is just two warps. That small thread count introduces an immediate new question: how do 64 threads load 512 elements into shared memory?


3. Cooperative Strided Loading

Because BM * BK = 64 * 8 = 512 elements for tile_a and BK * BN = 8 * 64 = 512 elements for tile_b, 64 threads cannot load the tiles in a single instruction.

Instead, each thread must load 512 / 64 = 8 values cooperatively using a strided loop:

#pragma unroll
for (uint load_offset = 0; load_offset < BM * BK; load_offset += num_threads) {
    uint load_idx = threadIdx.x + load_offset;
    uint a_row = load_idx / BK;
    uint a_col = load_idx % BK;
    tile_a[load_idx] = matrix_a[a_row * num_cols_a + a_col];
}

Every thread steps through the shared memory array with a stride of num_threads = 64. By the time the loop completes, the entire tile is safely staged in shared memory, ready for the compute stage.


4. The Inner Compute Loop: Register Outer Product

Once the shared memory tiles are filled and synchronized with __syncthreads(), the thread steps through the BK = 8 depth.

At each step, it pulls 8 floats of tile_a into a private register array register_m, and 8 floats of tile_b into register_n. Then it updates its 64 private accumulators:

for (uint dot_idx = 0; dot_idx < BK; ++dot_idx) {
    for (uint i = 0; i < TM; ++i) {
        register_m[i] = tile_a[(thread_row * TM + i) * BK + dot_idx];
    }

    for (uint i = 0; i < TN; ++i) {
        register_n[i] = tile_b[dot_idx * BN + thread_col * TN + i];
    }

    for (uint res_idx_m = 0; res_idx_m < TM; ++res_idx_m) {
        for (uint res_idx_n = 0; res_idx_n < TN; ++res_idx_n) {
            thread_results[res_idx_m * TN + res_idx_n] +=
                register_m[res_idx_m] * register_n[res_idx_n];
        }
    }
}

Because TM and TN are compile-time constants, the compiler unrolls these loops completely. All 64 values in thread_results live in hardware registers throughout the entire K loop, never touching cache or DRAM until final write-back.


5. Handling Non-Multiple Matrix Shapes

Real-world matrices do not always divide cleanly into blocks of 64 by 64. If you have a matrix of size 1000 by 1000, dividing by 64 leaves 15 full blocks (960 rows and columns) and a leftover tail of 40 elements.

Rather than cluttering the high-speed main kernel with per-element bounds checks on every iteration, I split the execution into two paths:

  1. Main kernel: Handles all full 64 x 64 blocks with zero branch checks.
  2. Edge kernel: Launched specifically for the remaining right edge, bottom edge, and corner blocks.

In the edge kernel, data loads are zero-padded when indices exceed matrix boundaries:

tile_a[load_idx] = (global_row_a < num_rows_a && global_col_a < num_cols_a)
                       ? matrix_a[a_row * num_cols_a + a_col]
                       : 0.0f;

When storing results back to matrix C, only valid elements within bounds are written:

if (global_row < num_rows_a && global_col < num_cols_b) {
    const uint c_idx = (thread_row * TM + res_idx_m) * num_cols_b +
                       (thread_col * TN + res_idx_n);
    matrix_c[c_idx] = alpha * thread_results[res_idx_m * TN + res_idx_n] +
                      beta * matrix_c[c_idx];
}

This keeps the fast path completely unhindered while guaranteeing numerical correctness across arbitrary matrix dimensions.


6. Benchmarks on NVIDIA H100

To benchmark this kernel on modern hardware, I ran it on an NVIDIA H100 80GB HBM3 GPU via Modal. The test environment disabled TF32 to ensure an apples-to-apples comparison against raw FP32 cuBLAS:

Matrix Size (M = N = K)Custom Kernel TimeCustom TFLOPScuBLAS TimecuBLAS TFLOPS% cuBLAS
5120.0779 ms3.440.0163 ms16.4620.9%
10240.1560 ms13.760.0562 ms38.1936.0%
20480.9721 ms17.670.3412 ms50.3635.1%
40967.4775 ms18.382.6448 ms51.9735.4%
819258.0961 ms18.9321.0749 ms52.1736.3%

Key Observations

  • Throughput scaling: On smaller matrices (size 512), the kernel does not generate enough blocks to saturate all 132 streaming multiprocessors on the H100. By size 2048 and above, occupancy stabilizes, and throughput settles near 18 to 19 TFLOPS.
  • Consistent cuBLAS percentage: Across 1024 to 8192, the kernel reliably delivers between 35% and 36.3% of cuBLAS throughput.
  • Pure CUDA cores: Reaching almost 19 TFLOPS without Tensor Cores or compiler intrinsics highlights how much raw performance is unlocked purely by restructuring memory hierarchy accesses.

What Is Next

While 2D block tiling is a major leap over 1D tiling, there is still headroom before approaching cuBLAS performance:

  1. Vectorized loads: Currently, global loads fetch 32-bit floats individually. Switching to 128-bit float4 instructions will saturate memory bus bandwidth much faster (see our deep dive on Vectorized Float4 GEMM).
  2. Double buffering: Overlapping global memory loads for the next tile with compute on the current tile will hide memory latency.
  3. Shared memory bank conflicts: Ensuring access patterns across the 8x8 micro-tile avoid bank collisions during inner loop register loads (see our deep dive on Shared Memory Bank Conflicts).
  4. Warp tiling and Tensor Cores: Scaling further to leverage modern Hopper Tensor Cores via WMMA.

About the author

Ritesh Yadav works as an AI/ML Engineer. He writes independent research notes on ML performance, infrastructure, and systems, covering CUDA, low-latency inference, generative AI, distributed training, Kubernetes, and LLMOps.