SGEMM Optimization: From Naive to cuBLAS-Level on a Single GPU

This notebook walks through writing a single-precision matrix multiply (SGEMM) kernel on an NVIDIA GPU and getting it within a few percent of cuBLAS — the same exercise that produces every dense layer in every transformer at peak throughput. We follow the progression in Salykova’s Beating cuBLAS in Single-Precision General Matrix Multiplication (code), with cross-references to Simon Boehm’s tutorial where the framing differs. Both are listed in PAPERS.md.

By the end you should be able to:

• explain why a naive matmul gets ~1% of peak FLOPS no matter how big the GPU is,
• read the kernel in sgemm/05_async_pipeline.cu line by line and explain every choice,
• predict what each optimization buys you (and what it costs in code complexity),
• run the bench yourself on a cloud GPU and produce a results table you can defend.

The companion source files live in sgemm/. Each kernel is its own .cu file so the diff between versions is the lesson.

You don’t need a local NVIDIA GPU to read this notebook — the kernels are explained in prose and source. To run them, build on a Lambda Cloud / Colab / RunPod box with the CUDA toolkit; see docs/CLOUD_SETUP.md.

Tip📺 Video companion

How CUDA Programming Works (Stephen Jones, GTC 2022) explains the memory-bandwidth mental model behind this notebook’s roofline analysis, and GPU MODE Lecture 50 recounts a similar CUDA learning journey — more in Videos.

Note🎯 Goal

Optimize a CUDA matmul from naive to near-cuBLAS, one memory-hierarchy trick at a time.

🏗️ Chapter milestone: M12 — benchmark your model’s own decode GEMM shapes — the exercise at the end of this chapter; see The Project.

Tip⏱️ Sessions S5.7–S5.8

Two sessions (Colab): S5.7 climbs from naive to near-cuBLAS; S5.8 is 🏗️ M12 — your model’s GEMM shapes. Cards, prerequisites, and done-when tests: the Session Guide.

This notebook needs an NVIDIA GPU: on Colab pick Runtime → Change runtime type → T4 GPU.

# Setup — on Colab, fetch the repo so the CUDA sources in sgemm/ are available
import os
if not os.path.exists('sgemm'):
    !git clone -q https://github.com/ggreg/ai_playground.git
    %cd ai_playground/notebooks/05_gpu_nvidia_tools

import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

SGEMM_DIR = Path('sgemm').resolve()
assert SGEMM_DIR.is_dir(), f'expected {SGEMM_DIR} to exist'

# Try to detect a local CUDA GPU. The notebook still works without one - the
# kernel explanations are static; only the live `make && ./bench` cells skip.
try:
    import torch
    HAS_CUDA = torch.cuda.is_available()
    if HAS_CUDA:
        print(f'CUDA GPU: {torch.cuda.get_device_name(0)}')
        print(f'Compute capability: {torch.cuda.get_device_capability(0)}')
    else:
        print('No CUDA GPU - this notebook will show source and pre-computed numbers.')
except ImportError:
    HAS_CUDA = False
    print('PyTorch not installed; kernel cells will be illustrative.')

1. Why SGEMM is the canonical CUDA exercise

Every linear layer in every transformer — QKV projection, attention output, MLP up- and down-projections — is a GEMM. On a well-fed dense model the GEMMs dominate FLOPs and a few percent of GEMM throughput is a few percent of training throughput. cuBLAS achieves roughly 93% of theoretical peak FP32 on Ampere for square matmuls; understanding how it gets there is the foundation for understanding FlashAttention, CUTLASS, Triton kernels, and the Tensor Core variants.

Three reasons SGEMM is a great teaching exercise:

1. The roofline is favorable. SGEMM has arithmetic intensity proportional to the tile size — you can keep growing it until you’re far above the bandwidth ridge. Unlike, say, attention’s softmax, nothing fundamental keeps you bandwidth-bound.
2. The optimizations stack cleanly. Each step (tile in shared memory, accumulate in registers, vectorize, async-pipeline) targets a distinct bottleneck and gives a measurable jump.
3. You can validate against cuBLAS. Correctness is unambiguous: max |C_yours - C_cublas| should be within FP32 rounding (~1e-4 for K=4096).

2. The roofline calculation

For square M = N = K:

• FLOPs = 2 · M · N · K (one multiply + one add per inner-loop step).
• Minimum HBM bytes = (M·K + K·N + M·N) · 4 (read A and B once, write C once).
• Arithmetic intensity = FLOPs / bytes.

For M=N=K=4096 that’s 137.4 GFLOPs against 192 MB — about 715 FLOPs/byte. The ridge point on an RTX 3090 (35.6 TFLOPS FP32 peak / 936 GB/s HBM) sits at ~38 FLOPs/byte. SGEMM at 4K×4K is almost twenty times above the ridge: there is no excuse to be memory-bound.

But a naive kernel ignores reuse: each thread reads K floats from A and K from B straight out of HBM, doing one FMA per pair. That collapses the effective AI to ~1, putting you back below the ridge by a factor of 700. The whole optimization story is shrinking HBM traffic until the kernel can’t be anything but compute-bound.

# Roofline plot for an RTX 3090 (Ampere, sm_86), with our 5 kernel variants annotated
# at their effective arithmetic intensities (back-of-envelope estimates).
peak_tflops = 35.6      # RTX 3090 FP32 peak (FMA)
peak_bw_tbs = 0.936     # HBM2 bandwidth (TB/s)
ridge = peak_tflops / peak_bw_tbs

ai = np.logspace(-1, 4, 500)
perf = np.minimum(ai * peak_bw_tbs * 1e3, peak_tflops * 1e3)  # in GFLOPS for plot

# Approximate effective AI per kernel (FLOPs per byte read from HBM):
#   - naive: ~1, every FMA reloads from HBM
#   - smem (32x32 tile): each tile of A reused 32x, B reused 32x -> AI ~ 16
#   - block_tile (128x128, 8x8 thread tile, BK=8): A reused 128x, B reused 128x -> AI ~ 64
#   - vectorized: same data movement as block_tile, but compute is faster -> AI ~ 64
#   - async: same AI; the win is hiding latency, not reducing traffic
points = [
    ('01 naive',         1.0,    0.28e3, 'tab:red'),
    ('02 smem',          16.0,   4.13e3, 'tab:orange'),
    ('03 block_tile',    64.0,  13.49e3, 'tab:olive'),
    ('04 vectorized',    64.0,  16.72e3, 'tab:green'),
    ('05 async',         64.0,  18.35e3, 'tab:blue'),
    ('cuBLAS',           64.0,  19.24e3, 'tab:purple'),
]

fig, ax = plt.subplots(figsize=(9, 5.2))
ax.loglog(ai, perf, color='black', lw=1.4, label='RTX 3090 roofline (FP32)')
ax.axvline(ridge, ls='--', color='gray', lw=0.8)
ax.text(ridge*1.05, 1, f'ridge = {ridge:.0f} FLOPs/byte', color='gray', fontsize=9)

for name, x, y, c in points:
    ax.scatter([x], [y], s=70, color=c, zorder=3, label=name)
    ax.annotate(name, (x, y), xytext=(7, 5), textcoords='offset points', fontsize=9, color=c)

ax.set_xlabel('Arithmetic intensity (FLOPs / byte from HBM)')
ax.set_ylabel('Achieved throughput (GFLOPS)')
ax.set_title('SGEMM @ M=N=K=4096 on RTX 3090: each kernel buys back arithmetic intensity')
ax.set_xlim(0.3, 1e3)
ax.set_ylim(100, 5e4)
ax.grid(True, which='both', alpha=0.3)
fig.tight_layout()
plt.show()

Read this plot as a journey. Kernel 1 sits in the bandwidth-bound region (left of the ridge) at ~280 GFLOPS — the SM is starving for data. Kernels 2 through 5 walk progressively to the right and up, each one reducing redundant HBM reads until the kernel finally bumps against the FP32 FMA throughput ceiling, where cuBLAS lives.

3. What we’re building

Five kernels in sgemm/, each a small change over the previous one:

File	Big idea	Expected on RTX 3090 (4096³)
01_naive.cu	One thread = one output element. No reuse.	~280 GFLOPS — ~1.5% of peak
02_smem.cu	Block-level tile staged in shared memory.	~4 TFLOPS — ~22% of cuBLAS
03_block_tile.cu	128×128 block, 8×8 register accumulators per thread.	~13 TFLOPS — ~70%
04_vectorized.cu	float4 (LDG.128) + transposed As for vectorized inner-loop reads.	~17 TFLOPS — ~87%
05_async_pipeline.cu	cp.async + double-buffered software pipeline (Ampere+).	~18-20 TFLOPS — ~95%

(Numbers are illustrative; actual results depend on driver, clocks, and thermal state. The bench in sgemm/bench.cu validates correctness against cuBLAS and reports your machine’s real numbers.)

4. Compile and run

The kernels need an NVIDIA GPU and the CUDA toolkit. On a fresh Lambda Cloud / Colab / RunPod box:

cd sgemm
make SM_ARCH=sm_86         # set to match your GPU; sm_80 = A100, sm_86 = RTX 3090, sm_89 = RTX 4090, sm_90 = H100
./bench --size 4096 --iters 50 --json results.json

For locked-clock reproducible numbers (recommended for any serious comparison):

sudo nvidia-smi --persistence-mode=1
sudo nvidia-smi -lgc 1395    # base core clock for RTX 3090; check `nvidia-smi -q -d SUPPORTED_CLOCKS`
./bench --size 4096 --iters 100 --json results.json
sudo nvidia-smi -rgc         # release lock when done

The bench writes results.json which we’ll load further down to plot the speedup ladder. Until then, the next sections walk through each kernel’s source.

5. Kernel 1: naive (sgemm/01_naive.cu)

The simplest correct implementation. Each thread computes exactly one output element of C by iterating over the K dimension and accumulating into a scalar.

const int col = blockIdx.x * BLOCK + threadIdx.x;   // varies fastest in a warp
const int row = blockIdx.y * BLOCK + threadIdx.y;
float acc = 0.0f;
for (int k = 0; k < K; ++k)
    acc += A[row * K + k] * B[k * N + col];
C[row * N + col] = alpha * acc + beta * C[row * N + col];

Where the time goes:

• Threads in a warp share row and have consecutive col. Reads of B[k*N + col] are coalesced into a single 128-byte transaction per warp — good.
• Reads of A[row*K + k] are broadcast within a warp (all threads read the same address) — also fine.
• But across the whole grid, every value of A is read N times (once per output column it participates in) and every value of B is read M times. For M=N=K=4096 that’s 4096 redundant reads of every element. Effective arithmetic intensity is ~1 FLOP/byte — firmly in the bandwidth-bound region.

Ceiling = HBM bandwidth / 4 bytes per FMA ≈ 936 GB/s · (1 FMA / 4 bytes) = ~234 GFLOPS. That’s exactly the kernel’s measured speed. The kernel is doing as well as it can given how much data it pulls from HBM. Optimization is about pulling less.

# Print the full source so you can read it inline.
print((SGEMM_DIR / '01_naive.cu').read_text())

6. Kernel 2: shared-memory tiling (sgemm/02_smem.cu)

Shared memory is on-chip SRAM with ~10 TB/s effective bandwidth on Ampere — about 10x faster than HBM. The classic GEMM trick: each block stages a BM × BK slab of A and a BK × BN slab of B into shared memory, then all BM · BN threads in the block read from shared instead of HBM.

For BM = BN = BK = 32 and a 32×32 thread block:

• Each block loads 32×32 of A and 32×32 of B = 8 KB total per K-tile.
• Within the block, the 32×32 = 1024 outputs share that data: each A row is reused by 32 threads (one per output column), each B column by 32 threads.
• HBM traffic per output element drops from O(K) (kernel 1) to O(K / BM) — 32× less bandwidth pressure.

Inner loop:

As[ty][tx] = A[row * K + (k0 + tx)];
Bs[ty][tx] = B[(k0 + ty) * N + col];
__syncthreads();
for (int k = 0; k < BK; ++k)
    acc += As[ty][k] * Bs[k][tx];
__syncthreads();

Two __syncthreads() per K-tile: one to ensure the staging is complete before any thread reads from shared, one to ensure all threads finished reading before the next iteration overwrites the staging.

What’s still wrong: each thread issues one shared-memory load per FMA. Shared is fast but not free — we’re effectively bottlenecked on shared-memory bandwidth instead of HBM bandwidth. Kernel 3 fixes this by holding many partial sums in registers.

print((SGEMM_DIR / '02_smem.cu').read_text())

7. Kernel 3: 2D thread tile + register accumulators (sgemm/03_block_tile.cu)

Registers are ~10x faster than shared memory again. Instead of one thread per output, give each thread an 8×8 patch of outputs held entirely in 64 registers, and reuse each loaded value 8 times.

Tile geometry:

	value
Block tile (BM × BN)	128 × 128
K-step (BK)	8
Per-thread output (TM × TN)	8 × 8
Threads per block	(BM·BN) / (TM·TN) = 256
Threads laid out as	16 × 16

Inner loop (the hot path):

for (int k = 0; k < BK; ++k) {
    for (int i = 0; i < TM; ++i) regM[i] = As[(threadRow*TM + i) * BK + k];
    for (int j = 0; j < TN; ++j) regN[j] = Bs[k * BN + threadCol*TN + j];
    for (int i = 0; i < TM; ++i)
        for (int j = 0; j < TN; ++j)
            threadResults[i*TN + j] += regM[i] * regN[j];
}

That’s 64 FMAs for 16 register reads — arithmetic intensity from shared memory is now 4 FLOPs/byte instead of the 0.25 it was in kernel 2. Big jump in throughput.

Two cooperative-load details to look at in the source:

• Loading A: 256 threads need to fill a 128×8 tile (1024 floats) — 4 elements per thread. Threads vary innerColA = tid % BK within a warp, so consecutive tids read consecutive columns of A — coalesced.
• Loading B: 256 threads fill a 8×128 tile (1024 floats). Threads vary innerColB = tid % BN within a warp — consecutive columns of B — one 128-byte coalesced transaction per warp.

This kernel, with default Ampere register allocation, gets ~3 blocks resident per SM — enough warps in flight to hide most shared-memory latency. Most of the heavy lifting is done by the time you write this kernel correctly.

print((SGEMM_DIR / '03_block_tile.cu').read_text())

8. Kernel 4: vectorized 128-bit traffic + transposed As (sgemm/04_vectorized.cu)

Two layered changes on top of kernel 3.

8.1 float4 loads and stores (LDG.128 / STG.128)

A single 128-bit (float4) load amortizes the address-calculation and instruction-issue overhead over four floats. SASS issues fewer LD/ST instructions per FMA, the load/store unit gets one wide transaction instead of four narrow ones, and HBM throughput goes up. On Ampere this is the difference between ~75% and ~90% of peak.

*reinterpret_cast<float4*>(&regM[0]) =
    *reinterpret_cast<float4*>(&As[k * BM + threadRow * TM + 0]);

8.2 Transpose A on load

In kernel 3, As is row-major (As[m][k]). For a fixed k in the inner loop, the 8 regM elements one thread needs are at addresses As[m..m+7][k] — strided by BK floats, which means scalar reads (8 LDS instructions) and a 4-way bank conflict on each.

The fix: store A transposed in shared memory as As[k][m]. The 8 regM values become contiguous in memory, so the inner loop reads them as 2 × float4. The transpose happens once per tile, on the load:

float4 a = *reinterpret_cast<const float4*>(&A[innerRowA * K + innerColA * 4]);
As[(innerColA*4 + 0) * BM + innerRowA] = a.x;
As[(innerColA*4 + 1) * BM + innerRowA] = a.y;
As[(innerColA*4 + 2) * BM + innerRowA] = a.z;
As[(innerColA*4 + 3) * BM + innerRowA] = a.w;

One vectorized HBM load per thread, four scalar shared stores. The cost is paid up front; the inner loop — the hot path — sees only fast vectorized shared reads.

8.3 What’s left to optimize?

The compute path is now efficient. What remains is the memory pipeline: HBM loads still serialize against compute. While the SM is loading the next tile, the FMA units sit idle. Kernel 5 fixes this with cp.async.

print((SGEMM_DIR / '04_vectorized.cu').read_text())

9. Kernel 5: cp.async + double-buffered software pipeline (sgemm/05_async_pipeline.cu)

The last big lever, and the headline of Salykova’s article. Available on Ampere (SM_80) and newer.

9.1 What cp.async does

The PTX instruction cp.async.ca.shared.global issues an asynchronous global → shared memory copy that does not pass through registers and does not block the issuing thread. The SM continues executing instructions while the HBM transaction is in flight. You commit a batch of in-flight copies with cp.async.commit_group and wait for them with cp.async.wait_all or the more granular cp.async.wait_group N.

__device__ __forceinline__ void cp_async_16(uint32_t smem_int, const void* gmem) {
    asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n"
                 :: "r"(smem_int), "l"(gmem));
}
__device__ __forceinline__ void cp_async_commit() {
    asm volatile("cp.async.commit_group;\n");
}
__device__ __forceinline__ void cp_async_wait_all() {
    asm volatile("cp.async.wait_all;\n");
}

Three benefits:

1. No register pressure for in-flight loads. In kernel 4, each load occupied registers from issue until commit-to-shared; kernel 5 frees those for accumulators.
2. No instruction blocking. The SM keeps issuing FMAs past the load.
3. Overlap with compute. This is the real win and needs double buffering to materialize.

9.2 The pipeline structure

We allocate two stages of shared memory and arrange the loop so the next tile is loading while the current tile is being multiplied:

prologue:
    cp.async tile 0 -> stage 0
    wait_all; transpose A; sync
    cp.async tile 1 -> stage 1   (in-flight)

main loop t = 0 .. num_tiles - 2:
    compute tile t                <-- runs WHILE tile (t+1) loads from HBM
    wait_all; transpose A; sync
    cp.async tile (t+2) -> stage (t % 2)   (reuse the just-finished stage)

epilogue:
    compute tile (num_tiles - 1)

9.3 The transpose scratchpad

cp.async preserves the global-memory layout, so we can’t directly cp.async into the transposed As from kernel 4 with a single 16-byte instruction. The kernel allocates a small scratchpad A_raw and does a shared→shared transpose pass between the wait and the compute. This is a tutorial-level simplification; Salykova’s article goes further and uses four 4-byte cp.async operations into transposed positions, eliminating the scratchpad. The trade-off is more PTX-level complexity for a modest extra speedup.

Take a moment to read the full source — the prologue/main-loop/epilogue structure is the canonical software-pipeline pattern you’ll see in every CUTLASS kernel and every flash-attention implementation.

print((SGEMM_DIR / '05_async_pipeline.cu').read_text())

10. The benchmark harness (sgemm/bench.cu)

Three things to call out in the bench source.

10.1 cuBLAS row-major trick

cuBLAS is column-major. We get a row-major result by computing C^T = B^T @ A^T in the column-major view, which uses the same bytes as C row-major:

cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, N, A, K, &beta, C, N);

10.2 L2 flush between replays

Cached A or B in L2 makes the second call faster than the first by ~20% on a fresh tile. To get steady-state numbers we cudaMemsetAsync a 64 MB buffer between replays, evicting the GEMM working set:

size_t flush_bytes = std::max<size_t>(2 * l2_size, 64 * 1024 * 1024);
cudaMemsetAsync(dFlush, 0, flush_bytes);

10.3 Median, not mean

Thermal throttling produces long-tail outliers. We sort and pick the median of the timed runs — a robust point estimate that won’t be skewed by one slow iteration when the GPU just hit its power limit.

11. Run the bench and plot the ladder

The cell below tries to load sgemm/results.json from a previous ./bench --json results.json run. If it doesn’t exist, it falls back to the example numbers from sgemm/README.md so the plot still renders.

import json

example_results = {
    'gpu': 'NVIDIA GeForce RTX 3090 (illustrative)',
    'sm': '8.6',
    'size': 4096,
    'iters': 50,
    'results': [
        {'kernel': 'cublas',            'median_ms':  7.142, 'tflops': 19.24, 'pct_cublas': 100.0, 'max_err': 0.0},
        {'kernel': '01_naive',          'median_ms': 482.31, 'tflops':  0.28, 'pct_cublas':   1.5, 'max_err': 1e-5},
        {'kernel': '02_smem',           'median_ms':  33.27, 'tflops':  4.13, 'pct_cublas':  21.5, 'max_err': 1e-5},
        {'kernel': '03_block_tile',     'median_ms':  10.19, 'tflops': 13.49, 'pct_cublas':  70.1, 'max_err': 1e-5},
        {'kernel': '04_vectorized',     'median_ms':   8.22, 'tflops': 16.72, 'pct_cublas':  86.9, 'max_err': 1e-5},
        {'kernel': '05_async_pipeline', 'median_ms':   7.49, 'tflops': 18.35, 'pct_cublas':  95.4, 'max_err': 1e-5},
    ],
}

results_path = SGEMM_DIR / 'results.json'
if results_path.exists():
    data = json.loads(results_path.read_text())
    print(f'Loaded results from {results_path}')
else:
    data = example_results
    print(f'No {results_path.name} found; using example numbers from README. '
          f'Run `cd sgemm && make && ./bench --json results.json` on a CUDA box to get your own.')

print(f"\nGPU: {data['gpu']}  (SM {data['sm']}, size={data['size']}, iters={data['iters']})\n")
print(f'{"kernel":<22}{"median ms":>12}{"TFLOPS":>10}{"% cuBLAS":>11}{"max |err|":>14}')
for r in data['results']:
    print(f"{r['kernel']:<22}{r['median_ms']:>12.3f}{r['tflops']:>10.2f}{r['pct_cublas']:>10.1f}%{r['max_err']:>14.2g}")

# Plot the speedup ladder.
results = data['results']
# Re-order so cublas appears at the right (the target).
ordered = [r for r in results if r['kernel'] != 'cublas'] + [r for r in results if r['kernel'] == 'cublas']
labels = [r['kernel'].replace('_', ' ') for r in ordered]
tflops = [r['tflops'] for r in ordered]
pct = [r['pct_cublas'] for r in ordered]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4.5))

colors = ['tab:red', 'tab:orange', 'tab:olive', 'tab:green', 'tab:blue', 'tab:purple']
ax1.barh(labels, tflops, color=colors[:len(labels)])
ax1.set_xlabel('Throughput (TFLOPS, FP32)')
ax1.set_title(f"Each kernel's measured throughput\n({data['gpu']}, M=N=K={data['size']})")
ax1.invert_yaxis()
for i, t in enumerate(tflops):
    ax1.text(t + 0.2, i, f'{t:.2f}', va='center', fontsize=9)

ax2.barh(labels, pct, color=colors[:len(labels)])
ax2.axvline(100, color='black', ls='--', lw=0.8)
ax2.set_xlabel('% of cuBLAS')
ax2.set_title('Same data as % of cuBLAS reference')
ax2.invert_yaxis()
for i, p in enumerate(pct):
    ax2.text(p + 1, i, f'{p:.1f}%', va='center', fontsize=9)

fig.tight_layout()
plt.show()

What to take away from the ladder. The biggest single jump is kernel 2 → 3 (smem to block-tile + register accumulators) — over 3× on this hardware. Each subsequent step is a smaller percentage but increasingly involves PTX-level reasoning. cuBLAS is sitting just past kernel 5 because its inner loop uses warp-tile decomposition and bank-conflict-free shared layouts that we left as exercises (see §13 below).

12. Things we deliberately did not do

To keep the source readable, the tutorial stops at kernel 5. Each of the following would be a worthwhile follow-up:

1. Warp-tile decomposition. Group the 256 threads into 8 warps of 32, give each warp a WM × WN sub-tile, and have threads within a warp share regM/regN reads more cleverly. This is the [BlockTile, WarpTile, ThreadTile] hierarchy used by CUTLASS. Adds ~5-10% on top of kernel 5.

2. Bank-conflict-free Bs. With Bs[k][n] and 16 columns of threads reading float4 at offsets threadCol * 8 + j, threads 0/4/8/12 hit the same banks. Padding Bs to width 132 fixes it; Salykova does this. Adds a few percent.

3. Tensor cores (mma.sync / WMMA). FP32 → TF32 mode on Ampere gives ~10× over FP32 FMA at the cost of mantissa precision. mma.sync.aligned.m16n8k8 is the relevant Ampere instruction. The CUTLASS GEMM templates use this by default.

4. Hopper TMA + warp-specialized kernels. SM_90 (H100) replaces cp.async with the Tensor Memory Accelerator, which streams whole tiles into shared memory in a single instruction issued by one warp. The producer/consumer warp split is a different programming model; CUTLASS 3.x supports it.

5. Split-K / persistent grids. For tall-skinny matrices where M·N doesn’t produce enough blocks to fill the GPU, split the K reduction across multiple blocks and atomic-add their partial results. Persistent kernels keep blocks resident across many tiles, amortizing launch overhead.

6. Auto-tuning the tile shape. Different problem sizes prefer different [BM, BN, BK, TM, TN]. Salykova’s article ends up using two kernels (128x128x8 and 128x256x8) and dispatches based on size. CUTLASS exposes hundreds of pre-compiled shapes.

13. Suggested exercises

Pick one and budget a couple of evenings for each. Record numbers in docs/PROGRESS_TEMPLATE.md.

1. Sweep the problem size. Run ./bench --size N for N in {1024, 2048, 4096, 6144, 8192} and plot TFLOPS vs N for each kernel. Where do the curves diverge from peak, and why?
2. Lock the clock and re-bench. Compare locked vs unlocked numbers. By how much does kernel 5 drop at unlocked clocks because of thermal throttling? (Salykova’s article spends real effort on this distinction.)
3. Add bank-conflict-free Bs. Pad the second dimension of Bs from 128 to 132 and adjust the indexing. Bench and report the delta.
4. Add an mma.sync (TF32) variant. Use nvcuda::wmma or inline PTX. Expect ~3-4× over kernel 5 in TFLOPS but ~10-bit mantissa precision.
5. Profile kernel 5 with Nsight Compute. Confirm that the SM is busy during cp.async loads (smsp__inst_executed.avg.per_cycle_active should stay high while smsp__average_warp_latency_per_inst_issued.ratio shows few stalls). See notebooks/05_gpu_nvidia_tools/03_nsight_profiling.ipynb.

🏗️ Project milestone M12 — your GEMM shapes

Every decode step of your model is a handful of small GEMMs: list them for batch B tokens — (B x dim) @ (dim x dim) attention projections, (B x dim) @ (dim x hidden) FFN, and the (B x dim) @ (dim x vocab) head. On Colab’s T4, time torch.matmul at those shapes for B ∈ {1, 8, 32} and record GFLOPs vs this chapter’s kernels in metrics.json['m12_gemms'] (CPU fallback: record the shape table only). Note how far batch-1 GEMMs sit below peak — that’s M9’s roofline conclusion, measured.

Accept when: the shape table (and timings, if GPU) is in metrics.json. Part of The Project: Serve Your Own LLM.

Session S5.8 — verify: uv run pytest tests/milestones/test_m12_gemms.py.

Key Takeaways

1. The optimization sequence is universal. Block-level shared-memory tiles → register accumulators → vectorized 128-bit traffic → async pipelined loads. Read FlashAttention, Mamba, or any modern Triton kernel and you’ll see the same hierarchy.
2. Naive is bandwidth-bound; optimized is compute-bound. The whole game is increasing arithmetic intensity until the kernel can’t be anything but compute-bound. Once you cross the roofline ridge, the remaining work is squeezing the FMA pipeline.
3. Each technique targets one bottleneck. Shared memory removes redundant HBM reads. Register tiles remove redundant shared reads. Vectorization removes redundant LD/ST instructions. Async copies remove load-compute serialization. Stack them.
4. Benchmarking is its own skill. Locked clocks, L2 flush, median over many replays, validation against a reference. Without this discipline your “X% faster” numbers are noise — the article’s emphasis on benchmark methodology is half the lesson.
5. You can match cuBLAS for FP32. The remaining gap to “beat cuBLAS” requires problem-specific dispatch (small vs large, square vs skinny), and that’s exactly where Salykova’s article goes next.

Next: 03_nsight_profiling.ipynb — profile kernel 5 with Nsight Compute and see exactly why each optimization showed up in the SASS counters. Or read PAPERS.md for the production-grade implementation in CUTLASS.

---

🔨 From-scratch project: p5 — A decode roofline model from a blank file

M9 and M11 computed one roofline for one model. The phase’s from-scratch project generalizes it: a calculator that takes any config and any GPU spec and predicts decode FLOPs, bytes, arithmetic intensity, and tokens/sec — no GPU required. Its punchline test: discover for yourself that at real context lengths no batch size makes decode compute-bound.

Brief and rules: projects/p5_roofline/README.md

uv run pytest projects/p5_roofline/ -v