Skip to content

Latest commit

 

History

6 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 

Repository files navigation

GEMM Optimization on H100: From Naive CUDA to Triton, Plus a Strassen Design

Target GPU: NVIDIA H100 (sm_90a)

This README walks through how I actually got to the final numbers for both parts — what I tried first, where it broke, and why I ended up where I did. If you just want the headline results, they're in the "Where I landed" sections. Everything else here is the reasoning behind them. The raw code, logs, and profiler screenshots referenced below are in the Task_1/ and Task_2/ folders (see the map at the bottom).


Part 1: Getting a GEMM close to cuBLAS on an H100

Starting point

I began the simplest way possible: a naive CUDA kernel where every thread computes exactly one output element, reading straight from global memory with no reuse at all (v0_naive_kernel in Task_1/code/v4_tuned.cu). It works, but it's slow, because every single multiply-add re-reads a full row and column from global memory instead of reusing anything.

Chasing the memory wall — v0 → v1 → v4

The obvious next step was shared-memory tiling: load a tile of A and a tile of B into shared memory once, and let every thread in the block reuse them (v1_tiled_kernel, TILE = 16). That alone gave a solid speedup — at N=512, runtime dropped from about 71 µs to about 51 µs, and at N=4096 it dropped from roughly 31 ms to well under half that.

But profiling in Nsight Compute (screenshots in 512_profiler_v0_v1_v4/, 1024_.../, 2048_.../, 4096_.../) showed there was still a lot left on the table, especially in how much time was going into memory movement versus actual math. So I pushed further into a templated, register-blocked kernel (gemm_tiled_kernel, the "v4" version): bigger shared-memory tiles, each thread now computing a small 4×4 microtile of output instead of one element, and vectorized float4 loads/stores so each thread instruction moves 128 bits at a time instead of 32. This is also where I deliberately kept the tile size small enough that N=512 still spawns 256 thread blocks instead of just 16 — more on why that matters below.

That last version (v4) was the biggest jump: at N=512, duration went from ~71 µs (v0) → ~51 µs (v1) → ~29 µs (v4). At N=4096, it went from ~31 ms (v0) down to ~4.4 ms (v4) — about a 7x improvement just from tiling and register blocking, before touching Tensor Cores at all.

The ceiling with hand-written CUDA

Even after all that tuning, the hand-written CUDA kernel topped out around 72–78% of cuBLAS's speed across matrix sizes. Good, but not good enough, and pushing it further by hand (moving to wmma Tensor Core APIs) ran into a wall: the compiler was spilling registers into local memory because of how opaque the wmma fragment structs compile down, and no amount of manual tuning got that fully under control.

Why I switched to Triton

Rather than keep fighting register spilling by hand, I rewrote the GEMM in Triton. On Hopper, Triton's backend emits wgmma.mma_async directly and manages the TMA (Tensor Memory Accelerator) descriptors and shared-memory swizzling itself — the exact thing I was fighting to get right by hand in CUDA. The FP32 kernel (python_gemm_fp32_triton.py) forces pure IEEE FP32 FMAs (input_precision="ieee", no silent TF32 truncation) and is autotuned across five block-size configs. The BF16 kernel (triton_gemm_tma_bf16.py) goes further and builds actual TMA tensor descriptors for A, B, and C, with autotune configs explicitly split into large/medium/small-N buckets — I kept extra small-block configs specifically for small N, again to avoid starving the GPU of work.

Where I landed

  • FP32 (pure IEEE, no TF32): beat cuBLAS at N=1024 (104.8% of cuBLAS), and stayed strong at larger sizes (88.8% of cuBLAS at N=4096).
  • BF16 (TMA + WGMMA): reached 79.6% of cuBLAS at N=4096, around 599 TFLOP/s.

The profiling insight that explained the small-N numbers

Across every version, the worst relative performance always showed up at small N (512, sometimes 1024). Nsight Compute made the reason obvious: at N=512 with a naive tiling scheme, you only get 8–16 CTAs total — nowhere near enough to fill the H100's 132 SMs, so most of the chip just sits idle (SM Busy time as low as ~1.86% in the worst case). That single observation is what drove the register-blocked tile sizing in v4 and the small-N-specific autotune configs in the Triton BF16 kernel — both are direct fixes for GPU starvation, not just "faster math."


Part 2: Is Strassen's algorithm actually worth it on an H100?

Working it out by hand first

Before writing anything about GPUs, I worked a small 2×2 block example by hand (Task_2/Strassen's Matrix.pdf) to make sure I actually understood the trade Strassen's algorithm makes. Standard block matrix multiply on a 2×2 split costs 8 multiplications and 4 additions. Strassen's method computes 7 products (M1 through M7) instead of 8, but needs 18 additions to combine them into the final result. That's the whole trade in one sentence: fewer multiplications, a lot more additions.

Doing the math for a real GPU

I then worked out the FLOP count formula for recursion depth L (Task_2/theorotically_total_flop.pdf): total GEMM FLOPs come out to 2·N³·(7/8)^L — so yes, FLOPs really do shrink as you recurse deeper, exactly as advertised. But that's only half the story on a GPU. An H100's Tensor Cores do roughly 1000 TFLOP/s of compute, while memory bandwidth is only about 3 TB/s. Matrix multiplication runs on the fast compute path; the 18 extra additions per level run on the slow memory-bound path. So every level of Strassen recursion is trading work off the path you have the most of (compute) onto the path you have the least of (memory bandwidth) — and as L increases, memory traffic actually grows by a factor of (7/4)^L. Net effect: Strassen pushes the whole computation further into memory-bound territory on the roofline model, which is the opposite of what you want on this hardware.

Three problems stacked on top of each other

Beyond the bandwidth trade itself, two more issues compound it, and one of them is a direct callback to Part 1:

  • GPU starvation, again. Deeper recursion means smaller sub-matrices (N / 2^L). I'd already seen exactly this failure mode in Part 1 at N=512 — tiny grids that leave most of the 132 SMs idle. Strassen recursion walks straight into the same trap on purpose.
  • Kernel launch overhead. At recursion level L you're issuing 7^L sub-GEMMs plus 18 addition kernels per level — that's hundreds of individual CPU-to-GPU launches, and that launch latency adds up fast unless it's fused or captured in a CUDA graph.

My proposed fix

This part stays at the design level rather than a full benchmarked implementation, so here's the approach I'd take to make Strassen viable on H100 instead of just explaining why it's currently a bad idea:

  1. Kill the memory traffic with fusion. Fold the 18 additions directly into the GEMM pipeline instead of materializing intermediate matrices in DRAM at all. Pre-additions like A11 + A22 happen during the TMA load stage; post-additions like M1 + M4 − M5 + M7 get accumulated directly in the WGMMA accumulator registers right before the final store. Nothing extra ever touches global memory.
  2. Stop recursing before you hit starvation. Set a base-case threshold (I used N_base = 2048) — once sub-matrices shrink below that, stop recursing and just run a normal dense GEMM instead of splitting further.
  3. Fill the chip with concurrency. Launch the 7^L sub-GEMMs for a given level concurrently via CUDA streams (or MPS) instead of serially, so all 132 SMs stay saturated even while each individual sub-GEMM is smaller.

Together, those three address the memory-bandwidth trade, the starvation problem, and most of the launch-overhead problem at once — the actual write-up is in Task_2/Tradeoff.txt.


Repo map

  • Task_1/code/v4_tuned.cu — the hand-written CUDA path: v0 (naive) → v1 (shared-memory tiled) → the templated register-blocked kernel referred to as v4.
  • Task_1/code/python_gemm_fp32_triton.py — the Triton FP32 GEMM (autotuned, pure IEEE FMA).
  • Task_1/code/triton_gemm_tma_bf16.py — the Triton BF16 GEMM using TMA tensor descriptors.
  • Task_1/code/fp32_sol.txt, Task_1/code/tma_sol.txt — raw Nsight Compute output backing the numbers above.
  • Task_1/{512,1024,2048,4096}_profiler_v0_v1_v4/ — Nsight Compute screenshots for each kernel version at each matrix size.
  • Task_1/fp32.png, Task_1/task1_perf_summary.png — the summary charts.
  • Task_2/Strassen's Matrix.pdf — the hand-worked 2×2 numeric example.
  • Task_2/theorotically_total_flop.pdf — the FLOP-count-vs-recursion-depth derivation.
  • Task_2/Tradeoff.txt — the full written analysis and proposed fusion design.

FP32 and BF16 Performance

About

what I tried first, where it broke, and why I ended up where I did. If you just want the headline results, they're in the "Where I landed" sections. Everything else here is the reasoning behind them. The raw code, logs, and profiler screenshots referenced below are in the Task_1/ and Task_2/ folders (see the map at the bottom).

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages