Vectorized Memory Access: Pushing CUDA SGEMM to 29 TFLOPS on H100

After getting the 2D block-tiling kernel running and benchmarking around 18.9 TFLOPS on an NVIDIA H100, I was happy with the progress, but I knew there was still substantial performance left on the table. In that kernel, each thread computed an 8 by 8 patch of the output matrix using register reuse, which gave us roughly 36% of cuBLAS speed.

However, when I looked closely at how data moved into shared memory, an obvious inefficiency stood out: we were still fetching memory from global DRAM into shared memory one 32-bit float at a time.

Modern GPU memory controllers have native 128-bit memory buses. Asking the hardware to load 32 bits per instruction is like hiring a large cargo truck to transport a single small box. You pay the full overhead of issuing a memory instruction, but you only utilize a fraction of the available memory bus width.

To fix this, I rewrote the loading pipeline in my gemm-kernels repository to use vectorized 128-bit loads (float4), expanded the block tile to 128 by 128, and benchmarked it on an NVIDIA H100 via Modal.

The result was an immediate leap: performance climbed from 18.9 TFLOPS to 29.12 TFLOPS, reaching over 57% of cuBLAS speed on raw FP32 CUDA cores with TF32 disabled.

Here is the story of how vectorized memory access works, why tile geometry changes everything, and what the code looks like.

---

1. Why Scalar Loads Waste Instruction Bandwidth

When a thread loads a standard float from memory, the GPU compiler emits a 32-bit load instruction.

Even though multiple threads in a warp can coalesce their individual 32-bit reads into cache-line transactions, each thread is still issuing an instruction that requests only 4 bytes of data. Every memory instruction must pass through the GPU instruction issue pipelines, consuming scoreboard entries and instruction queue slots.

CUDA provides built-in vector types like float2 (64 bits, 8 bytes) and float4 (128 bits, 16 bytes).

When you cast a float pointer to const float4* and load four contiguous floats together:

const float4 tmp_a = reinterpret_cast<const float4*>(&block_a[inner_row_a * num_cols_a + a_col])[0];

The compiler translates that single C++ expression into a single 128-bit vector load instruction (LD.E.128).

Instead of issuing four separate 32-bit load instructions to fetch 16 bytes of data, the thread issues a single instruction that pulls all 16 bytes across the bus at once. This cuts the number of load instructions issued by a factor of 4, leaving the warp scheduler with far less instruction pressure and more room to dispatch math operations.

---

2. Upsizing the Tile: 128 by 128 Geometry

In the earlier 2D block-tiling kernel, our block output tile was 64 by 64 (BM = 64, BN = 64). Because each thread handled an 8 by 8 output patch, each thread block only needed (64 / 8) * (64 / 8) = 64 threads.

While 64 threads worked fine on smaller GPUs like the T4, on an NVIDIA H100 with 132 large Streaming Multiprocessors, running only 64 threads (2 warps) per block leaves the multiprocessor under-occupied, making it harder for the hardware scheduler to hide execution latency.

For the vectorized kernel, I expanded the tile dimensions:

constexpr int BM = 128; // Output tile height
constexpr int BN = 128; // Output tile width
constexpr int BK = 8;   // Step size along the inner K dimension
constexpr int TM = 8;   // Rows per thread
constexpr int TN = 8;   // Columns per thread

Now, each thread block computes a 128 by 128 patch of matrix C, which equals 16,384 output elements.

With each thread computing an 8 by 8 patch (64 elements), the number of threads per block becomes:

threads_per_block = (BM / TM) * (BN / TN) = (128 / 8) * (128 / 8) = 16 * 16 = 256 threads

A block of 256 threads equals 8 full warps. This is a well-known sweet spot for CUDA kernels: it provides plenty of warps for latency hiding without creating excessive register pressure that would throttle active block occupancy.

---

3. The Perfect Match: 256 Threads Meet Float4

One of the most satisfying parts of this design is how cleanly the thread count aligns with the shared memory tile size:

• Shared memory tile_a has size BM * BK = 128 * 8 = 1024 floats.
• Shared memory tile_b has size BK * BN = 8 * 128 = 1024 floats.

Because each float4 holds 4 floats, 1024 floats is exactly 1024 / 4 = 256 vector loads!

Since our block contains exactly 256 threads, every single thread loads exactly one float4 for tile A and one float4 for tile B.

There is no need for awkward strided loops or complex index math just to fill shared memory. In a single synchronized loading phase, all 256 threads execute in parallel:

const uint inner_row_a = threadIdx.x / (BK / 4);
const uint inner_col_a = threadIdx.x % (BK / 4);
const uint inner_row_b = threadIdx.x / (BN / 4);
const uint inner_col_b = threadIdx.x % (BN / 4);

Every thread calculates its coordinates, issues its 128-bit load, and writes the 4 floats directly into the shared memory buffer.

---

4. Transposing Tile A to Protect Shared Memory Bandwidth

If you look closely at how data is written into tile_a, there is an intentional layout choice:

const float4 tmp_a = reinterpret_cast<const float4*>(&block_a[inner_row_a * num_cols_a + a_col])[0];

tile_a[(inner_col_a * 4 + 0) * BM + inner_row_a] = tmp_a.x;
tile_a[(inner_col_a * 4 + 1) * BM + inner_row_a] = tmp_a.y;
tile_a[(inner_col_a * 4 + 2) * BM + inner_row_a] = tmp_a.z;
tile_a[(inner_col_a * 4 + 3) * BM + inner_row_a] = tmp_a.w;

Notice that instead of storing tile_a as BM x BK (128 rows by 8 columns), the code transposes the data as it lands in shared memory, storing it as BK x BM (8 rows by 128 columns).

Why transpose tile A during the load?

Look at the compute loop where threads read from tile_a:

for (uint dot_idx = 0; dot_idx < BK; ++dot_idx) {
    #pragma unroll
    for (uint i = 0; i < TM; ++i) {
        register_m[i] = tile_a[dot_idx * BM + thread_row * TM + i];
    }
    // ...
}

Because tile_a is transposed, adjacent threads in a warp (which have sequential thread_row values) access consecutive addresses along the BM dimension. This produces sequential, unit-stride memory access across the threads, ensuring conflict-free shared memory reads throughout the inner compute loop.

---

5. Handling Unaligned Tails and Boundaries

In production code, matrix dimensions are not always perfect multiples of 128 or 4. If an arbitrary dimension causes a load to cross a boundary, issuing a blind 16-byte float4 read would cause out-of-bounds memory faults.

To handle this cleanly, the kernel checks whether the dimension is aligned and within bounds:

if (num_cols_a % 4 == 0 && inner_row_a < BM && a_col + 3 < static_cast<uint>(num_cols_a)) {
    // Fast path: aligned 128-bit vector load
    const float4 tmp_a = reinterpret_cast<const float4*>(
        &block_a[inner_row_a * num_cols_a + a_col])[0];
    tile_a[(inner_col_a * 4 + 0) * BM + inner_row_a] = tmp_a.x;
    tile_a[(inner_col_a * 4 + 1) * BM + inner_row_a] = tmp_a.y;
    tile_a[(inner_col_a * 4 + 2) * BM + inner_row_a] = tmp_a.z;
    tile_a[(inner_col_a * 4 + 3) * BM + inner_row_a] = tmp_a.w;
} else {
    // Safe fallback: scalar loop with zero-padding
    #pragma unroll
    for (uint i = 0; i < 4; ++i) {
        tile_a[(inner_col_a * 4 + i) * BM + inner_row_a] =
            (a_col + i < static_cast<uint>(num_cols_a))
                ? block_a[inner_row_a * num_cols_a + a_col + i]
                : 0.0f;
    }
}

Full interior tiles take the high-speed 128-bit path, while boundary tiles safely fall back to bounds-checked scalar reads padded with zeros.

Similarly, when writing final accumulated results back to global matrix C, each thread checks row and column boundaries before committing its values:

if (row < static_cast<uint>(num_rows_a) && col < static_cast<uint>(num_cols_b)) {
    const uint c_idx = row * num_cols_b + col;
    matrix_c[c_idx] = alpha * thread_results[res_idx_m * TN + res_idx_n] +
                      beta * matrix_c[c_idx];
}

This ensures mathematical correctness across non-square shapes, arbitrary sizes, and arbitrary alpha/beta parameters.

---

6. Benchmark Results on NVIDIA H100

To measure the real-world impact, I benchmarked this kernel on an NVIDIA H100 80GB HBM3 GPU via Modal, testing square matrices from size 512 up to 8192 with TF32 disabled:

Matrix Size (M = N = K)	Custom Kernel Time	Custom TFLOPS	cuBLAS Time	cuBLAS TFLOPS	% of cuBLAS
512	0.1118 ms	2.40	0.0152 ms	17.68	13.6%
1024	0.2062 ms	10.41	0.0542 ms	39.64	26.3%
2048	0.5900 ms	29.12	0.3386 ms	50.74	57.4%
4096	4.7672 ms	28.83	2.6394 ms	52.07	55.4%
8192	38.1850 ms	28.79	21.0683 ms	52.19	55.2%

Comparison Against Earlier Kernels

Here is how the performance has evolved across kernel iterations on the H100 for size 4096:

• 2D Block Tiling (64x64 tile): 18.38 TFLOPS (35.4% cuBLAS)
• Vectorized Float4 Tiling (128x128 tile): 28.83 TFLOPS (55.4% cuBLAS)

That represents an immediate 56.8% jump in raw throughput.

On larger matrices (size 2048 to 8192), the kernel consistently sustains 28.8 to 29.1 TFLOPS, capturing more than 57% of what NVIDIA cuBLAS achieves on pure FP32 CUDA cores.

---

Takeaways and What Comes Next

Moving to 128-bit loads teaches a simple but fundamental lesson in GPU engineering: raw arithmetic speed means nothing if the memory instruction pipeline is constantly throttled by high instruction counts.

By grouping 4 floats into a single instruction and sizing the block to 256 threads, we simultaneously maximized memory bus utilization, achieved conflict-free shared memory access, and provided enough warps to keep the SM pipelines saturated.

Our optimization roadmap continues from here:

1. Double Buffering: Currently, threads load data into shared memory, synchronize with __syncthreads(), compute, and synchronize again. Using double buffering allows global memory loads for the next tile to occur in parallel while the current tile is computing.
2. Asynchronous Memory Copies (cuda::memcpy_async): Bypassing registers altogether during global-to-shared transfers using hardware asynchronous copy engines.
3. Warp Tiling and Tensor Cores: Structuring register micro-tiles at the warp level and utilizing Hopper Tensor Cores for FP16 and FP8 execution.