22. GPU Engineering: Correctness, Throughput, and Attention

Chapters 12–13 explain GPU mapping and shared-memory kernels. This chapter turns those ideas into a hardware investigation. Every performance result here is an assignment to measure on your GPU; there are no implied benchmark results. Start with the architecture sections of the official CUDA Programming Guide and record the GPU model, supported compute capability, compiler, and driver used for your experiment.

22.1 Execution, memory, and synchronization

A block contains threads that execute in warps. Blocks may execute in any order; a conventional kernel cannot depend on a different block reaching a software barrier. Use separate launches or a supported cooperative mechanism when global coordination is required.

A block barrier requires participating threads to reach it consistently. This is unsafe when the last block is partial:

if global_index < n:
    shared[local_index] = input[global_index]
    block_barrier()

Move the barrier outside the divergent condition and initialize invalid lanes to the operation's identity. For a sum use zero; for a maximum use negative infinity. Ensure consumers cannot read an uninitialized slot. Warp-level primitives have their own active-mask and participation requirements; lockstep intuition is not a substitute for the specified synchronization contract.

Lab: implement vector addition, a block reduction, and a multi-block reduction. Test lengths 1, 31, 32, 33, and 1003. Run the available memory, race, and synchronization checking tools on real hardware. Keep a deliberately broken version to demonstrate which check catches the error.

22.2 Coalescing and shared-memory banks

For a row-major matrix, adjacent threads loading adjacent columns typically access nearby addresses. If adjacent threads instead stride across rows, the warp may need many more memory transactions. Logical bytes read and physical bytes transferred can therefore differ substantially.

A tiled transpose loads contiguous rows into shared memory and stores contiguous rows of the transposed result. Shared-memory layout can create bank conflicts. Under the common teaching model of 32 banks and 4-byte words, bank index is (byte_address/4) mod 32; a 32×32 float tile accessed by column maps those accesses to the same bank. Padding its row stride to 33 spreads those accesses. Check the actual architecture's rules, access width, and broadcast behavior before generalizing.

Lab: compare naive transpose, tiled transpose, and padded tiled transpose. Handle rectangular sizes and edge tiles. Report useful bandwidth 2*M*N*sizeof(element)/time and relevant profiler transaction and bank-conflict evidence. Explain any mismatch between predicted and measured behavior.

22.3 The GEMM optimization ladder

For a block output tile BM × BN and reduction chunk BK, load BM*BK + BK*BN elements into shared memory, then perform approximately 2*BM*BN*BK operations. With f32 inputs, the idealized input arithmetic intensity per chunk is:

2*BM*BN*BK / (4*(BM*BK + BK*BN)) FLOPs per byte

This ignores output traffic, cache effects, and repeated block loads, so it is a model to check. Each thread can retain several output accumulators in registers. That increases reuse but also register demand. Shared memory per block, registers per thread, thread count, and architectural limits jointly constrain residency. High occupancy alone does not imply high throughput.

Build and retain each rung: naive global loads; shared-memory tiling; per-thread register tiles; vectorized loads where alignment permits; tensor-core operations; pipelined memory transfer. For each rung, measure correctness, time, register use, shared-memory use, spills, and occupancy. A larger tile that spills can lose to a smaller tile with less nominal reuse.

For tensor cores, match supported tile shapes, storage and accumulator types, layouts, alignment, and warp participation rules. WMMA fragments are implementation-defined distributed register representations; never assume each thread owns a simple contiguous subsection of the matrix. Compare f16-input/f32-accumulate output against a reference using the same quantized inputs.

Lab: benchmark square and rectangular GEMMs with aligned and irregular dimensions. For each optimization predict its benefit, then inspect generated code and profiler evidence. Demonstrate at least one configuration where an apparently stronger optimization loses and explain why.

22.4 Asynchronous pipelines

A sequential tile loop loads tile t, synchronizes, computes it, then starts t+1. A double-buffered pipeline can overlap loading t+1 with computing t. It needs separate storage for tiles in flight, a completion mechanism before consumption, and protection against reusing storage too early.

Write a timeline before code:

prologue: begin load tile 0
iteration t: wait tile t; begin load t+1 into other buffer; compute tile t
boundary: ensure all readers finished before reusing tile t's buffer

Exact wait and barrier operations depend on the target instruction and architecture. A block barrier does not automatically replace an asynchronous-copy completion protocol. More stages consume more shared memory, potentially reducing residency enough to remove the benefit.

Lab: add a supported asynchronous-copy pipeline to your tiled GEMM. Validate the final partial K tile, buffer reuse, and all synchronization paths. Compare two and three stages across several K sizes. If the hardware lacks the feature, retain the design and mark the device experiment incomplete.

22.5 FlashAttention beyond the formula

For one query row, process scores in tiles and maintain running maximum m, normalizer l, and unnormalized weighted-value accumulator o. Starting with m=-infinity, l=0, o=0, a new score tile S updates:

m_new = max(m, max(S))
p = exp(S - m_new)
alpha = exp(m - m_new)
l_new = alpha*l + sum(p)
o_new = alpha*o + p @ V_tile
output = o / l

Rescaling old state is what makes the tiled result match a stable full softmax in real arithmetic. Floating-point reduction order still changes rounding. Causal masking must exclude future positions before they contribute to maxima and sums. Fully masked rows require explicit semantics to avoid -infinity - -infinity producing NaN.

For backward, with P=softmax(S) and output O=P*V, compute dV=P^T*dO, dP=dO*V^T, and dS=P*(dP-row_sum(dP*P)). Then apply the score scale when forming dQ and dK. A memory-efficient backward recomputes needed probabilities from scores and saved row statistics instead of storing all of P. Its scheduling and race-free gradient accumulation are substantial implementation work.

Lab: implement forward and backward against an unfused reference. Test causal and noncausal attention, irregular sequence lengths, several head dimensions, and extreme logits. Run finite differences on tiny double-precision cases. Measure peak memory and latency as sequence length grows; report both forward and backward. Add dropout only after defining reproducible masks.

22.6 Framework comparison and backend portability

Use the official Triton tutorials as a second implementation route for vector operations, reductions, matrix multiplication, and attention. Compare the schedule expressed by a block program with the thread and shared-memory schedule your CUDA renderer emits. Read the intermediate representations for one kernel.

For benchmarks, warm up compilation and caches, measure repeated device work with synchronization or device events, and report median plus spread. Host wall time without waiting often measures enqueue latency. Keep transfers and cold compilation as separate end-to-end measurements. Match PyTorch's precision settings and verify that the compared operation performs equivalent work.

For HIP or another backend, list the portable semantics and the architecture-specific assumptions: subgroup width, matrix instructions, memory spaces, barriers, launch limits, and runtime APIs. A textual rename of CUDA builtins is insufficient. Port a reduction before attempting tensor cores.

Exit gate: a correctness suite on real hardware, a transpose and GEMM optimization notebook, an attention forward/backward implementation, and a workload-wide comparison. Competitive performance is earned separately for each tested workload and device.

22.7 Worked lab: derive and execute tiled attention backward

Let scores be S=Q*K^T/sqrt(d) and probabilities P=softmax(S). For a loss gradient G on O=P*V, the chain rule gives dV=P^T*G and dP=G*V^T. A single softmax row has Jacobian diag(P)-P*P^T, so its vector-Jacobian product is dS=P*(dP-sum(dP*P)).

The row correction can be computed without materializing dP or P:

sum_j dP[i,j]*P[i,j]
= sum_j sum_c G[i,c]*V[j,c]*P[i,j]
= sum_c G[i,c]*O[i,c]

That identity lets backward save output O and one log-normalizer per query row, then recompute each probability tile as exp(scores-log_normalizer). For a score tile of width T, temporary score and probability storage is proportional to queries*T, rather than queries*keys. Q, K, V, output, and gradients still occupy memory: the algorithm removes quadratic score storage, not all storage.

For a scalar online-softmax example, process scores 1000 and 1001 in separate tiles. After the first tile, m=1000 and l=1. After the second, m=1001 and l=exp(-1)+1. The first contribution to the weighted-value sum is also rescaled by exp(-1). Forgetting that rescaling produces wrong results even though the normalizer looks stable.

Run the implementation and its numerical checks:

.venv/bin/python -m labs.attention
.venv/bin/python -m pytest tests/test_advanced_labs.py -q -k attention

"""Tiled online attention and recomputing backward, using NumPy as a CPU lab.

Single head, Q:(queries,dim), K:(keys,dim), V:(keys,value_dim).
Causal queries use absolute positions starting at query_offset. Fully masked rows
produce zero output, -inf log-normalizer, and zero gradients by explicit policy.
Run: python -m labs.attention
"""
import numpy as np


def _inputs(q, k, v, tile, query_offset):
  if q.ndim != 2 or k.ndim != 2 or v.ndim != 2 or q.shape[1] != k.shape[1] or k.shape[0] != v.shape[0]:
    raise ValueError("expected compatible rank-two Q, K, V")
  if min(*q.shape, *k.shape, *v.shape) < 1 or tile < 1 or query_offset < 0:
    raise ValueError("dimensions and tile must be positive; query_offset nonnegative")
  if not all(np.issubdtype(x.dtype, np.floating) and np.isfinite(x).all() for x in (q, k, v)):
    raise ValueError("this lab requires finite floating-point inputs")


def _scores(q, k_tile, start, scale, causal, query_offset):
  scores = (q @ k_tile.T) * scale
  if causal:
    mask = np.arange(start, start + len(k_tile))[None, :] > (np.arange(len(q)) + query_offset)[:, None]
    scores = np.where(mask, -np.inf, scores)
  return scores


def online_attention(q, k, v, tile=32, causal=False, query_offset=0):
  _inputs(q, k, v, tile, query_offset)
  dtype = np.result_type(q, k, v, np.float32)
  q, k, v = (x.astype(dtype, copy=False) for x in (q, k, v))
  maximum = np.full(len(q), -np.inf, dtype=dtype)
  normalizer = np.zeros(len(q), dtype=dtype)
  output = np.zeros((len(q), v.shape[1]), dtype=dtype)
  scale = 1 / np.sqrt(q.shape[1])
  for start in range(0, len(k), tile):
    scores = _scores(q, k[start:start+tile], start, scale, causal, query_offset)
    next_maximum = np.maximum(maximum, scores.max(axis=1))
    # Zero is a harmless exponent reference for a still-fully-masked row.
    safe_maximum = np.where(np.isfinite(next_maximum), next_maximum, 0)
    alpha = np.exp(maximum - safe_maximum)
    probability = np.exp(scores - safe_maximum[:, None])
    output = alpha[:, None] * output + probability @ v[start:start+tile]
    normalizer = alpha * normalizer + probability.sum(axis=1)
    maximum = next_maximum
  output = np.divide(output, normalizer[:, None], out=np.zeros_like(output), where=normalizer[:, None] != 0)
  log_normalizer = np.full_like(maximum, -np.inf)
  valid = normalizer > 0
  log_normalizer[valid] = maximum[valid] + np.log(normalizer[valid])
  return output, log_normalizer


def attention_backward(q, k, v, output, log_normalizer, grad_output, tile=32, causal=False, query_offset=0):
  """Recompute probability tiles; never allocate a full queries-by-keys matrix."""
  _inputs(q, k, v, tile, query_offset)
  dtype = np.result_type(q, k, v, grad_output, np.float32)
  q, k, v = (x.astype(dtype, copy=False) for x in (q, k, v))
  if output.shape != (len(q), v.shape[1]) or grad_output.shape != output.shape or log_normalizer.shape != (len(q),):
    raise ValueError("invalid saved state or output gradient shape")
  dq, dk, dv = np.zeros_like(q), np.zeros_like(k), np.zeros_like(v)
  correction = (grad_output * output).sum(axis=1)
  scale = 1 / np.sqrt(q.shape[1])
  valid = np.isfinite(log_normalizer)
  safe_log = np.where(valid, log_normalizer, 0)
  for start in range(0, len(k), tile):
    scores = _scores(q, k[start:start+tile], start, scale, causal, query_offset)
    probability = np.exp(scores - safe_log[:, None])
    probability[~valid] = 0
    dp = grad_output @ v[start:start+tile].T
    ds = probability * (dp - correction[:, None])
    dq += (ds @ k[start:start+tile]) * scale
    dk[start:start+tile] += (ds.T @ q) * scale
    dv[start:start+tile] += probability.T @ grad_output
  return dq, dk, dv


def main():
  rng = np.random.default_rng(0)
  q, k, v = (rng.normal(size=shape) for shape in ((7, 5), (11, 5), (11, 3)))
  output, log_normalizer = online_attention(q, k, v, tile=4, causal=True)
  gradients = attention_backward(q, k, v, output, log_normalizer, np.ones_like(output), tile=4, causal=True)
  print("output:", output.shape, "saved row statistics:", log_normalizer.shape)
  print("gradient shapes:", *(x.shape for x in gradients))
  print("CPU algorithm lab; no GPU performance claim")


if __name__ == "__main__": main()

This lab is a CPU implementation of the algorithm. It supports one head, finite floating-point inputs, causal masking, and a nonnegative query-position offset for decoding-style subproblems. Its loops do not establish GPU synchronization correctness or performance. Finite inputs can still overflow a dot product at extreme magnitudes; the tested numerical domain must be reported.

Experiment: vary tile size through 1, 4, and larger than the key length. Forward results should agree within rounding error. Perturb Q, K, and V separately in random directions, comparing the dense loss finite difference with the analytic gradient's directional inner product. The test suite uses an independent dense computation, so a shared error in tiled forward and backward cannot validate itself.

Next implementation: lower the tile loop into a real GPU kernel, then handle parallel dK/dV accumulation without races. Compare memory and device time with the dense baseline at matched precision. The serial CPU loop's safe accumulation order is not an automatic GPU scheduling solution.

22.8 Continue into measured register tiling

Chapter 35 implements a 64×64 shared-memory output tile with per-thread register accumulators, masks partial dimensions, and measures the actual CUDA kernels on an RTX 3050. It includes the kernel, event-timing harness, independent numerical checks, and raw recorded results.