Writing a CUDA matmul that catches cuBLAS — then training a GPT on it

The hardware is an RTX 4070 Laptop (Ada, sm_89): 36 SMs, 256 GB/s of memory bandwidth, 8 GB, and a 55 W power budget. Two numbers from that spec sheet explain everything that follows.

At the clock these benchmarks run at, the card can do about 11.1 TFLOP/s of fp32 and move 256 GB/s. Divide them and you get the ridge point: 43 FLOP/byte. Every byte read from memory has to feed 43 floating-point operations before the arithmetic units stop waiting on memory. A naive matmul manages 0.25.

That gap — a factor of ~170 — is the entire project. Almost every optimization below is the same move applied at a different level of the memory hierarchy: load a value once, then spend it on as much arithmetic as possible before letting it go.

The nine kernels

GFLOP/s at N=4096, SM clock pinned to 1200 MHz. The cuBLAS baseline is fp32, which is what cuBLAS does by default. Kernel 9 uses tensor cores and is therefore not computing the same thing — see below.

Kernel 1 → 2: one index, 8× faster

The naive kernel maps threadIdx.x to the row of the output. That is the natural thing to write and exactly wrong. Lanes 0–31 of a warp then read addresses K floats apart, so each lane needs its own memory transaction — 32 transactions for what could be one.

Swapping the mapping so threadIdx.x is the column makes the warp read 32 consecutive floats: one 128-byte transaction. The arithmetic is byte-for-byte identical. Only the address pattern moved, and throughput went from 83 to 648 GFLOP/s.

Kernels 3–5: reuse, at three levels

Shared memory tiling gets each value loaded from global memory once per block instead of once per thread (0.25 → 8 FLOP/byte). Then register tiling attacks the next bottleneck: every fused multiply-add in the shared-memory kernel needs two SMEM loads, and shared memory cannot feed the FMA pipes that fast. Giving each thread an 8×8 patch of output means 16 loads serve 64 FMAs instead of 128.

The measurement was nearly wrong — in my favor

The fix is to pin the SM clock. 1500 MHz does not hold — the power cap drags it to ~1320 and it wanders. 1200 MHz holds rock steady through a sustained dense-GEMM load (39–47 W, 68–83 °C, clock never moved). With clocks pinned, best-vs-median across runs agrees to under 1%.

Correctness needed the same care. Elementwise relative error is the wrong metric for a matmul: each output is a sum of K products, so where cancellation drives an entry near zero the rounding noise of the other terms remains. An entry of magnitude 1e-3 can carry 1e-4 of absolute error while every input was computed correctly. Comparing worst absolute error against max|ref| separates cleanly — reordered fp32 summation lands around 1e-6, a real bug lands at 1e-1.

The profiler found what reading could not

At kernel 6 I was at 88% and out of ideas. Nsight Compute gave the answer in three lines:

DRAM Throughput          15.1%   <- global memory is not the problem
L1/TEX Cache Throughput  81.4%   <- saturated
Compute (SM) Throughput  54.2%

Shared memory and L1 share the same LSU/MIO datapath on NVIDIA hardware. 81% L1 against 15% DRAM means the kernel was bound on shared-memory-to-register traffic. The FMA units were idle waiting for operands. No further work on global memory access — the thing I had spent five kernels on — would have bought anything.

Warp tiling fixes exactly that, by adding a blocking level between the block and the thread so each value pulled from SMEM feeds more arithmetic. Loads per FMA drop from 0.25 to 0.1875. Measured after the change:

Kernel 8: the loop was shaped wrong

Kernel 7 left the kernel in a genuinely different regime — neither memory (45%) nor compute (55%) saturated. It was latency bound, and the reason was visible in the loop structure rather than in any counter:

load global -> shared        <- ~500 cycle latency
__syncthreads()              <- every warp blocks until it lands
compute on the tile
__syncthreads()              <- and blocks again before overwriting

The load is immediately followed by a barrier, so nothing overlaps it. Every K-chunk pays a full round trip to DRAM with the arithmetic units idle. Double buffering keeps two tiles in shared memory: while the warps compute on one, the next chunk’s loads are already in flight. It also halves the barriers, because reading one buffer and writing the other cannot conflict.

That took kernel 8 to 95.0% of cuBLAS at N=4096, and past it at sizes that divide evenly into the 128×128 block tile — 104.0% at N=1536 and 104.2% at N=6144, reproducible across runs. The cost was 219 registers against 186, and 32 KiB of shared memory against 16.

Kernel 9: the silicon I had not touched

Everything to this point runs on the SM’s fp32 FMA pipes. The card also has tensor cores, which had been idle for eight kernels. The honest way to see how much that matters is to let cuBLAS use them too:

So the tensor-core kernel beats fp32 cuBLAS by 18%, and trails cuBLAS’s own tensor-core path by about 21%. Both of those are worth saying out loud; quoting only the first would be the flattering half.

Tuning it was also a reminder that intuition is not a substitute for measurement. The first version gave each thread a 4×4 grid of accumulator fragments and hit 255 registers — the hardware ceiling — which throttled occupancy. Spreading the same block tile over 8 warps instead of 4 cut that to 128 registers. And of two arrangements with identical register counts and identical instruction counts, one was 27% faster than the other, purely from how the warps’ fragment loads land in shared memory.

Four ways not to speed up a tensor-core kernel

Kernel 9 trails cuBLAS’s own TF32 path by about 21%. The obvious place to look is the profiler, which says: DRAM 64.6%, compute 36.9%, L2 hit rate 58%, occupancy 32.8%. The tempting reading is bandwidth bound — and three of the four attempts below came from taking that at face value. All four made it slower. Three runs each, N=4096:

The bigger tile is the one that settles it. A block tile’s arithmetic intensity is BM·BN / (2(BM+BN)) — 32 FLOP/byte at 128×128, and 42.7 at 256×128, which clears this card’s 43 ridge point. If the kernel were really bandwidth bound, that is the fix. It cost 16%. So it is not bandwidth bound, however much the DRAM counter looks like it.

And double buffering is kernel 8’s cure for exactly this neither-counter-is-saturated signature, which earlier in this project was worth 6%. Here it does nothing, and not for want of registers — 128 either way, twelve bytes of spill. So the latency is not on the global-memory path either.

Kernel 10: the hypothesis was wrong and the kernel got faster anyway

Shared memory does not have to hold a matrix. It only has to hold whatever makes the next read cheap. The mma.m16n8k8 TF32 instruction requires each of the 32 lanes to hold four particular elements of A — a0=(g,t) a1=(g+8,t) a2=(g,t+4) a3=(g+8,t+4), where g=laneid>>2 and t=laneid&3. In a row-major tile those sit at four unrelated addresses. So kernel 10 stores each 16×8 tile lane-major — lane L’s four elements at L*4 — and the fragment load becomes one 128-bit access at base + laneid*16, which is also the ideal shared-memory pattern: 32 lanes over 512 contiguous bytes, no bank conflict possible. Per warp per k-step, 24 shared-load instructions become 6, moving exactly the same 3072 bytes.

It bought 2.7%. Instruction issue on the shared path was never the constraint. The bytes were — and WMMA moves exactly as many.

The other 8% came from somewhere kernel 9 could not reach. What this kernel is actually limited by is shared-memory reuse: bytes read from shared per mma issued is 4096·(WM+WN)/(WM·WN), which falls only when both warp-tile dimensions grow — and the accumulator costs WM·WN/32 registers per thread, so reuse is a register-budget problem in disguise. Going from a 32×64 warp tile to 64×64 halves it, and is worth 8.7%, three times what hand-written PTX was worth on its own.

And that exact shape is already in kernel 9’s tuning table, at 7783 GF/s — slower than the shape it settled on. WMMA at a 64×64 warp tile needs 255 registers and spills; the hand-written version spills four bytes. So raw PTX did matter — indirectly, by making the tile affordable rather than by making the loads cheaper. The stated hypothesis was wrong and the conclusion drawn from it was right, which is not the same thing as being right.

Two failures on the way, both worth more than the result

A fragment layout is a fact to measure, not to recall. I wrote the A register order from memory and swapped a1 and a2 — correct for the analogous f16 shape, wrong for TF32, which concatenates two k4 chunks rather than two row halves. The kernel compiled, ran at full speed, and returned garbage. Guessing produced a silent wrong answer, so there is now a one-hot probe that discovers the mapping on the hardware: set one element of A to 1, make B the identity, and see which lane and register light up.

Fewer instructions is worth nothing if the bytes arrive four at a time. The first correct version ran 8% slower than the WMMA kernel it was meant to beat, on a layout whose entire justification was cheaper shared access. Lane-major staging is 8-way bank conflicted on A and 16-way on B, because a warp of staging threads varies only in bits the slot index scales by 4. The fix is an XOR swizzle keyed on the tile index: staging sees it vary across the warp and spreads out, while a fragment load reads one whole tile per warp so it is uniform there and the permutation is invisible. 8-way and 16-way become 2-way, and the load stays perfectly conflict-free.

That is index arithmetic, not a hardware mystery, so it can be settled without a GPU. tools/smem_banks.py derives both layouts, proves they are bijections, checks that a fragment load picks up exactly the elements the instruction demands in exactly the right register slots, and simulates the bank pattern of every access.

Putting it in the model, where it behaves differently

The ladder kernel only does NN; the model needs all four transpose cases. The layout ports cleanly, because it is a function of the logical element (m,k) rather than of how the operand is stored — all four cases share one map, and only the axis the global float4 runs along changes. 67.7 → 65.2 ms per training step, 3.7%, replicated three times each way at a pinned clock, with identical loss to four decimals.

But the swizzle has to key on the axis staging walks. It must be uniform across a warp doing a fragment load and varying across a warp doing a staging store, and only one tile coordinate is both — the one the staging warp walks, which the transpose flag decides. I ported the untransposed key to all four cases. For the transposed ones that is warp-uniform during staging, so the swizzle does nothing and the stores go back to 16-way conflicted. It is not a correctness bug, so every test passed. It showed up only as the transposed cases running 20% slower than the WMMA path they replaced: geomean over the model’s twenty shape/transpose combinations was 0.887×, and 1.044× with the key fixed. What identified it was the pattern — NN and NT near parity, TN and TT at 0.80–0.84, which is exactly the half where transA is true.

And the best tile on a square benchmark is not the best tile in the model. The 64×64 warp tile that is worth 8.7% at N=4096 loses in situ, 1.034× against the narrower shape’s 1.044×. The model’s GEMMs are 4096×384×384 and friends, so a 128×128 block tile gives 96 blocks against 36 SMs — the machine is not full, and a 128-thread block brings half as many warps per SM to hide latency with. The extra reuse is real and there is nothing to spend it on. The ladder and the model deliberately run different tiles.

The compiler optimized the thing it could see

Kernel 9 sat at 8064 GF/s until a question about a version number. The writeup said CUDA 12.5; 13.3 was also installed. Building the same source with both — same machine, same pinned clock, N=4096 — eight of the nine kernels came out indistinguishable, and one did not:

A fifth of the throughput, reproducibly, at identical numerical error — so it was still genuinely TF32, just slower. ptxas -v gives the whole story in two lines:

12.5:  128 registers, 12 bytes spill stores
13.3:  142 registers,  0 spills

nvcc 13.3 spent 14 more registers to eliminate a 12-byte spill. In isolation that is a good trade. Here it crosses an occupancy cliff, because registers are allocated per warp in multiples of eight and this kernel runs 256 threads per block:

• 128 regs → 4096/warp → 65536/4096 = 16 warps per SM → 2 resident blocks
• 142 regs → rounds to 144 → 4608/warp → 65536/4608 = 14 warps per SM → 1 block

Halving the resident blocks to avoid twelve bytes of spill. The compiler optimized what it could see — the spill — and could not see what it cost.

The fix is to state what the kernel needs instead of hoping the register allocator infers it. The second argument to __launch_bounds__ is minimum blocks per SM:

__global__ __launch_bounds__(NUM_THREADS, 2) void tensorcore_kernel(...)

Both toolkits then allocate 128 registers and accept the spill. This is not a 13.3 workaround — it is faster on both, and 8290 GF/s is the fastest this kernel has ever run:

A spilled byte is cheap; a resident block is not. __launch_bounds__ is how you tell the compiler which one you are buying. With the fix in place the two toolkits agree across the board, so the build now takes whichever is newest and prints which one it used.

Then: a language model on top of it

A matmul is only interesting if something uses it. So the second half was a GPT — 6 layers, 6 heads, 384 embedding, 256 context, weight-tied head, 10.8M parameters — trained on character-level Shakespeare. Layernorm, GELU, causal multi-head attention, softmax, cross-entropy and AdamW are all hand-written; every matmul routes through the kernels above.

The backward pass needs dX = dY·Wᵀ and dW = dYᵀ·X. Rather than materializing transposed copies — an extra bandwidth-bound pass over the data — the transpose folds into the shared-memory staging. The A tile was already stored transposed, because that is what makes the per-thread register reads contiguous, so each of the four transpose cases is just a different index map.

A falling loss does not verify a backward pass

This is the part most from-scratch projects skip. A dropped correction term in layernorm still produces a descending loss curve — just a worse model. So the gradients are checked against finite differences directly.

The naive check perturbs one scalar and watches the loss, but in fp32 that barely works: individual gradients here are ~1e-4, so the loss moves by about as much as the forward pass’s own rounding noise. Instead each parameter tensor is stepped along u = g/‖g‖, so every element contributes coherently and the predicted change is exactly ‖g‖ — three orders of magnitude above the noise floor.

All 16 parameter tensors agree to 1e-5…2e-3. And the check is sensitive: deliberately deleting the x̂·mean(dx̂·x̂) term from layernorm backward makes 14 of 16 tensors fail immediately. (The two that still pass are the final layernorm’s weight and bias — their gradients do not flow through the path that was broken, which is exactly right.)

What it writes

Sampled from the best checkpoint at temperature 0.8:

What the profiler found that I never would have

This repo had the tooling to profile a training step for two sessions before it ever ran one, because reading GPU performance counters needs administrator on Windows and it never seemed worth the interruption. It was worth the interruption. One click, thirty seconds, and it said something I had not guessed — twice.

67.7 → 59.9 ms per step, 11.5%. Everything that grew as a share grew because the total shrank.

The bias add did not need to exist

Second biggest thing in the step, for one add per element. As its own kernel it reads the entire output tensor and writes it back to do that; in the GEMM epilogue — which is already holding the value in a register, about to store it — the same add costs two floats per lane out of L1 and no global traffic at all.

There is a reason it was a separate kernel, and it is not laziness. Adding a per-column value in an epilogue requires knowing which accumulator register holds which column, and that is exactly what WMMA’s fragment type hides. Kernel 9 could not have done this; the raw-PTX kernel can, because the register mapping is the thing it was built around. The abstraction that turned out not to be the bottleneck for arithmetic turned out to be a real constraint on what could be fused — a better argument for writing the PTX than the 2.7% was.

The column reduction read memory the wrong way round

What was left after fusing the forward was the bias backward, at 5.4% of a step for an operation whose floor is a single streaming read. It ran one block per column, striding down the rows — under a comment of mine asserting the reads were coalesced because consecutive blocks own consecutive columns. They were not. Coalescing happens within a warp, and in that arrangement a warp’s 32 threads read 32 different rows at the same column: 32 addresses C floats apart, so 32 separate transactions fetching 32 bytes each to use 4.

Neighbouring blocks do re-use those sectors out of L2, which is why it was bad rather than catastrophic — and why it survived. Nothing about the source looks wrong. It has the shape of a coalesced kernel and a comment explaining why it is one. The fix is the first thing anyone learns about CUDA.

The residual and GELU go the same way, and cost three detours

Same waste, same fix — and two activation buffers that then had nothing left to hold, since attproj and fcproj are never read again, not even by the backward. fp32 70.1 → 69.1 ms, TF32 59.8 → 58.9 ms, 0.91 → 0.84 GB resident. About 1.5% — and getting to an honest 1.5% took three detours that are worth more than the number.

It measured as 37% first. The profile said total kernel time had barely moved, which is the only reason I looked: ncu resets the application clock when it detaches, so an earlier clock lock had been silently undone and the card was boosting. 59.9 → 38 ms is 1.58×; it was a clock ratio wearing a speedup’s clothes. This project already had a rule about never quoting a ratio on an unpinned clock. What it did not have was a way to notice the pin coming undone mid-session.

Then the loss started varying in the fourth decimal, which looked exactly like a race I had just introduced. It is not mine and it is not new: two backward kernels accumulate through global atomics, so floating-point summation order varies between runs. The parent commit does it too — 1 run in 14. I was one plausible story away from attributing a pre-existing property of the model to my own change, and what prevented it was building the parent and running it fourteen times.

And the fp32 path came out 6% slower. Not register pressure — 219 against 221, essentially unchanged. The epilogue tested if (ep.gelu_out) at runtime, and tanhf expands to a substantial block of code that sits in the kernel whether or not the branch is taken. Making the feature set a template parameter and testing it with if constexpr turned a 6% regression into a 1.4% gain. Code you do not execute is not free.

One kernel was never going to be enough

Before starting on the attention backward I noticed the GEMM benchmark had never covered three of the model’s own shapes. Adding them showed the attention-projection weight gradient running at 1846 GF/s where every other shape reaches 7000–8000. A 128×128 tile cuts 384×384 into nine blocks on a 36-SM card; three quarters of the machine is idle and no amount of inner-loop work fixes that.

Two fixes, both shape-dependent — which is the point, and is why cuBLAS ships dozens of kernels rather than one good one. A second, half-height tile (64×128, 4 blocks/SM) for grids that cannot fill the machine once: +25% on the nine-block shape, −3% on the 96-block ones, so it is chosen per shape. And split-K when the output is small and K is long.

I had costed split-K before, for 4096×384×384, and correctly rejected it: the output there is 6.3 MB, so every extra partial is another 6.3 MB of traffic to recover a third of a wave. The weight gradients are the opposite shape — 590 KB of output against K=4096 — so the partials are nearly free. Same technique, opposite verdict, and the deciding quantity is K against M·N, not the wave count I had been staring at.

58.9 → 54.0 ms per step. The part worth dwelling on is that three of the four largest matmuls in the backward pass had never been measured individually, so the one running at a quarter speed was invisible — averaged into a GEMM category that looked healthy at 65%.

The clock lock does not stay locked

Twice now a change has measured as a large speedup and been a clock ratio. ncu resets the application clock when it detaches; the lock has also lapsed on its own. Both times the number looked plausible — 59.9 → 38 ms is 1.58×, and so is 1900/1200. This project already had a rule saying never to quote a ratio on an unpinned clock, and that rule does not survive a pin coming undone silently, between the pinning and the measurement. Every timing run now reads the clock on both sides of itself and labels the output UNPINNED if either reading is wrong. A discipline that depends on remembering to check is not a discipline.

Fusing attention: the score matrix never exists

The table above says attention costs about 21% of a training step across three kernels. None of them is badly written — the batched GEMM is the same code that reaches 90% of cuBLAS. The cost is structural. A (B, NH, T, T) score matrix gets written to global memory and read back, and the softmax over it does roughly 5 flops per 8 bytes on a card whose ridge point is 43 FLOP/byte. That is 0.6% of peak no matter how good the kernel is. The only fix is to not have the intermediate.

A softmax normally needs the whole row before it can emit anything, because it needs the row max and the row sum. But both are running statistics. After seeing part of a row you hold a partial max m and partial sum l, and when a later block raises the max to m′, everything computed so far is corrected by one factor — exp(m - m′):

m' = max(m, rowmax(S_j))
l' = l * exp(m - m') + rowsum(exp(S_j - m'))
O' = O * exp(m - m') + exp(S_j - m') @ V_j

with O divided by l only at the end. Each score tile lives in registers for the few instructions it takes to consume it and is then gone. This is FlashAttention. Note it does more arithmetic than the unfused version, not less — the whole win is in what never gets written.

Two more wins fall out of the structure, and on this hardware they matter as much as the famous one. Causality becomes a loop bound rather than a mask: the unfused path computes all T×T scores and discards the upper triangle inside the softmax, while a query block here simply never visits key blocks past its diagonal — so both attention matmuls do half the work. And the head permute disappears: the unfused path pays a bandwidth-bound pass over 3·B·T·C to make each head’s slice contiguous because a batched GEMM needs uniform strides, while the fused kernel indexes q/k/v straight out of the (B, T, 3C) projection — one block owns one head, so the head offset is a constant on the row pointer.

Backward never stores the score matrix either. It rebuilds the probabilities from the saved log-sum-exp, which is exact and needs no reductions at all: P[i,j] = exp(S[i,j] − lse[i]). Recomputing S costs one matmul; reading a stored score matrix back costs 25 MB of DRAM per layer. On this card that trade is not close. It is two kernels rather than one, because dQ reduces over keys while dK and dV reduce over queries — fusing them would push one of the three through global atomics on a tensor the size of the activations. Recompute beats communication, which is this project’s whole lesson restated one level up.

The backward is only 1.19×, and why

The forward result is most of the story; the backward is nearly a wash, and it is worth saying why rather than quoting the forward alone. My first answer was a guess: the fastest backward config ran 64 threads per block at 46 KB of shared memory — two blocks per SM, 8% occupancy. So I added a config with twice the threads. It gained 2.4%. The profiler explains why:

L1/TEX saturated near 70% with DRAM at ~21% is a signature this project has met before — kernel 6 showed 81% against 15%. Shared memory and L1 share the LSU/MIO datapath, so this is shared-memory→register traffic and the FMA pipes are starved. The same disease as kernel 6, one level up the hierarchy. Counting loads per FMA agrees: the forward spends 12 shared loads on 32 FMAs in its P@V loop, the backward accumulation spends 16 on 32. That is also why more threads did not help — a narrower column tile buys occupancy by loading more per unit of arithmetic, so the two effects nearly cancel.

What the memory actually buys: twice the context

The speedup is nice; the memory is the part that changes what the card can do. Unfused attention is quadratic in context length, fused is linear. Total resident memory at batch 16 — parameters, gradients, Adam state, activations and backward scratch:

This card has 8 GB. The unfused path tops out at context 1024; the fused path trains at context 2048, and the unfused path would need nearly three times the card to do the same. Fusing attention doubled the context this GPU can train.

Twenty percent that was hiding behind a square benchmark

The benchmark above measures square matrices. The model does not run square matrices — its matmuls are 4096×384×384, 4096×1536×384, 384×384×4096. I had written that sentence in the README months earlier as an explanation for why end-to-end throughput sits below the headline GFLOP/s, and never measured it. So ./bench/sgemm got --mnk M,N,K, and two pieces of free performance fell out immediately.

First: the model’s GEMM was built on kernel 7, and kernel 8 beats it at every one of these shapes. The transpose-aware GEMM was written against the kernel-7 structure, kernel 8 added double buffering afterwards, and the model simply never got the prefetch — on about 80% of a training step. Porting it across took 11.4% off the step, with the arithmetic untouched: loss and gradient norm identical to every printed digit.

Nothing was broken. No bug, no regression, every test passing the whole time. It was invisible because the benchmark measured a shape the model does not run.

The speedup that did not survive integration

Second: the tensor-core kernel is 1.22–1.58× faster than kernel 7 at these shapes — a wider margin than the 1.30× it manages on square matrices, because skinny matmuls starve the fp32 pipes harder than they starve the tensor cores. So I built a TF32 path for the model’s GEMM. Wired into training it was 2.4% slower.

A speedup that does not survive integration is a measurement that has not finished. Timing the same entry point both ways, per shape and per transpose case, gave a split so clean it named the cause by itself:

WMMA wants its A fragment stored m-major. Holding As[m][k] forces the transposed case to scatter four scalars for every float4 read, where the fp32 kernel stores As[k][m] and gets a straight vector copy. And the backward pass computes every weight gradient as TN — so the forward won, and the backward handed the winnings straight back.

The fix is not to pick a better tile. It is to stop insisting on one storage orientation: keep each operand in whichever layout makes staging a vector copy, and choose the fragment layout to match — col_major when the operand arrived transposed. Same bytes, same mma instruction, no scatter.

Does TF32 cost the model anything?

Full 5000-step runs, same seed, same configuration, clock pinned:

The validation losses differ by 0.004 nats. That is comfortably inside the run-to-run spread this setup shows from fp32 summation order alone — across configurations on the same data and seed I have measured 1.5035, 1.5056, 1.5138, 1.5147, 1.5178 and 1.5194. So TF32 buys 10% of training time and costs nothing I can distinguish from noise, which is why every framework defaults to it on Ampere and later.

Two GPUs, and the lesson that turned out to be wrong

The received wisdom about multi-GPU training is that communication becomes the bottleneck. Split the work across two cards and the interconnect, not the arithmetic, is what limits you. I rented two A40s to see it happen.

First, the collective itself. No NCCL — for the same reason there is no cuBLAS here. A ring all-reduce splits the gradient buffer into n chunks and runs two phases of n−1 steps, reduce-scatter then all-gather, with every device sending and receiving at once. Each device moves 2(n−1)/n · S bytes, which stops growing with n, and every link stays busy. That is why bandwidth-optimal collectives are rings.

The driver said peer-to-peer worked. It did not.

cudaDeviceCanAccessPeer returned true both directions. cudaDeviceEnablePeerAccess succeeded. Every cudaMemcpyPeerAsync returned success, every stream synchronise returned success — and the bytes never arrived. PCIe ACS or IOMMU misconfiguration on a virtualised host, which is most rented hardware.

It failed silently. The all-reduce reported error exactly 1.00 — the result was exactly zero — at 0.2 GB/s. And before the test caught it, the training log had already said so in a way I nearly missed: one rank reported a gradient norm of 15.023, two ranks reported 7.672. Exactly half, because each rank kept its own shard’s gradient and then divided by the rank count. No sum ever happened, and the loss curve looked entirely reasonable while it did not.

Communication was 6%. Two GPUs were still slower than one.

Two A40s deliver roughly half the throughput of one, and the interconnect accounts for six percent of it. The wire is not the problem. The host loop is: one process driving both devices, with blocking calls inside the step, so the two GPUs barely overlap. Giving each rank its own thread — CUDA’s current device is per-thread — moved it from 159 to 149 ms. Real, and nowhere near the ~85 ms two genuinely overlapped ranks should reach.

Chasing that further turned up a bug worth the trip on its own. The loss reduction cached its output scalar in a static device pointer, so every rank got the same pointer, allocated on whichever device called first — one rank reducing into another rank’s memory, both racing to read one scalar. It produced plausible losses the whole time.

Attention on the tensor cores, and a round trip that vanishes

Every matmul in the project ran on tensor cores under --tf32 except the ones inside attention, which stayed on fp32 FMAs with 4×2 and 4×4 register tiles while the model’s own GEMM ran 8×8. That gap was most of why the fused backward beat the unfused path by only 1.06×. Attention was the last consumer in the repo computing a matmul the slow way.

Porting it turns on one fact. S = Q@Kᵀ comes out of an mma as an accumulator, and P = softmax(S) then has to go back in as the A operand of P@V. If those layouts agree the fragment is reused in place; if they do not, P goes through shared memory — and that round trip is the structural reason the backward was slow.

But the disagreement is confined to the four lanes that share a row group, so eight __shfl_synces convert one to the other exactly, and P never reaches shared memory at all. One trap on the way, worth stating because it is silent: shuffling the already-selected register — __shfl_sync(m, par ? d[1] : d[0], src) — reads whichever register the source lane’s own par chose. It lands on the neighbouring column and stays entirely plausible. Both registers have to be shuffled and the selection applied afterwards.

The backward needed a transpose, and did not get one

The forward ports directly and is 1.81× the best fp32 tile. The backward is two kernels and only one of them is that easy: dQ’s second matmul reduces over keys, so its A operand is dS in the accumulator’s own orientation. But dK/dV reduces over queries — dV[c][j] += Σ P[i][c] dO[i][j] — so its A operand is P transposed, and transposing an accumulator is exactly the round trip being deleted.

So that kernel does not compute S. It computes Sᵀ = K@Qᵀ directly, key as M, and its accumulator is [key][query] from the start — the orientation the second matmul wants. The reorientation is free, and the reason is the part worth keeping: the backward needs no row reductions at all. The forward must reduce along a query’s keys for the running max and sum; here lse and D are already known, so P and dS are elementwise with a per-query scalar.

End to end, both arms with --tf32 matmuls and attention pinned each way: −6.9% of a training step at ctx 256, −18.5% at 1024, −24.3% at 2048. The win grows with context because attention’s share of the step does.

A stale constant, and two wrong explanations for it

With attention ported the matmuls are back to roughly 72% of a step, so the next thing is the tensor-core dispatch’s choice between a 128×128 tile and a 64×128 one. Its rule was TILE_SWITCH_BLOCKS = 72 — two blocks per SM across the 36 SMs of the 4070 it was written on. The card it now runs on has 46. Re-measuring is worth 7.6% of a training step, 43.6 → 40.2 ms, and lifts mlp up to 96.5% of the ladder’s own square-N peak.

The interesting part is that I explained it wrong twice, and both times building the explanation is what exposed it.

Wave quantisation was the first: 96 blocks against 46 SMs × 2 = 92 slots is one full wave plus four stragglers. The control that killed it was forcing the narrow tile on every shape — it wins at 384 blocks as decisively as at 96 — and the arithmetic agrees it was never the story, since 96-into-92 and 192-into-184 are the same 52% efficiency.

The roofline was the second, and wrong more interestingly. A BM×BN tile reads (BM+BN)·BK·4 bytes per 2·BM·BN·BK flops, so its arithmetic intensity is BM·BN / 2(BM+BN): 32.0 FLOP/byte wide, 21.3 narrow. This card’s ridge point is 21.0 — so the narrow tile clears it and the wide one does not, which explains the measurement and the 4070’s opposite result. A tidy story.

What survives is the ratio between machines, robust where the knife-edge was not: the 4070’s ridge is ~44 at base clock against this card’s 25.4. So the rule shipped is empirical with the roofline as motivation — take the narrow tile unless the machine’s ridge sits above the wide tile’s intensity, i.e. unless it is bandwidth-starved enough that intensity is the limit. Both cards clear by 20–40%, it is evaluated at base clock so it is deterministic, and it is calibrated on two data points and labelled as such rather than dressed up as a law.

What I’d do next

1. Raw mma.sync — done, above. 8318 → 9210 GF/s on the ladder. The hypothesis behind it was wrong, which is the part worth keeping.
2. cp.async — done. Worth 8.1% on the ladder, and the reasoning that promoted it to “the main event” accounted for only 2.1 of those 8.1 points; the asynchrony it treated as a bonus was the rest. Wiring it into the model’s GEMM was tried and lost at every shape, and the suspected cause — shared footprint costing a resident block — is not the reason either. It is BK: reuse per barrier.
3. Fuse attention — done, and then ported to tensor cores, above. The prescription this list used to carry — chunk the head dimension so bigger register tiles stop competing with more blocks per SM — was carried out, and it works structurally while buying only 0.4–3.5%. The sharpest datum from it is the config that frees a third resident block with nothing traded away and is not one percent faster. Attention is no longer computing anything on fp32 FMAs.
4. Re-profile before choosing the next thing. The step has gone 46.8 → 40.2 ms, and the two largest costs both moved a lot, so the old ranking cannot be trusted. This project is three-for-three on my intuitions losing to the profiler, and the honest move is to measure the new distribution rather than guess which of layernorm, the bias reductions or the remaining GEMM headroom is now on top.
5. Multi-GPU — started, above. The collective is correct and cheap; the data-parallel driver is not yet good enough to profit from it. Two ranks need to genuinely overlap, which means one process per GPU or a step with no blocking calls left in it.