A3 — CUDA Kernels

Files: llmc/encoder.cuh, llmc/layernorm.cuh, llmc/attention.cuh, llmc/matmul.cuh, llmc/gelu.cuh, llmc/fused_classifier.cuh, llmc/adamw.cuh, train_gpt2.cu
Needs: 1 NVIDIA GPU
Builds on: A0–A2
Check with: ./grade A3

---

Overview

Here you port the forward pass, backward pass, and optimizer from A1/A2 to CUDA. The goal is not just correctness — you must write kernels that saturate GPU memory bandwidth. A correct but slow kernel is insufficient.

By the end of A3, ./train_gpt2cu will train GPT-2 124M at hardware-limited speed on a single GPU. Multi-GPU is A4.

This one requires an NVIDIA GPU. A0–A2 did not; everything from here does. ./grade A3 reports every check as SKIP when nvcc is absent, so you can still read through it without one.

To list your work:

grep -rn "TODO (A3)" llmc/ train_gpt2.cu

When these notes and the in-file HINT disagree, the code is authoritative — please report the discrepancy so it can be fixed.

Key techniques you will use:

• 128-bit vectorized loads/stores (x128 / Packed128<T>) to reach peak memory bandwidth
• Warp-level reductions (warpReduceSum, warpReduceMax) to avoid atomicAdd bottlenecks
• Block-level reductions (blockReduce) using shared memory and warp shuffles
• Online softmax to compute softmax in a single pass without materializing all values
• Stochastic rounding to correctly convert FP32 accumulations back to BF16/FP16
• cuBLASLt for matrix multiplications with fused epilogue operations

---

Where this fits

flowchart LR
    A0["A0
tokens, batches, PRNG"] --> A1["A1
forward pass (CPU)"]
    A1 --> A2["A2
backward + AdamW (CPU)"]
    A2 --> A3["A3
port to CUDA (1 GPU)"]
    A3 --> A4["A4
scale to N GPUs"]
    A3 --> A5["A5
schedules, clipping, checkpoints"]
    style A3 fill:#2563eb,stroke:#1e40af,color:#fff

The maths does not change in this assignment — A1 and A2 already defined it. What changes is that you are now writing for a machine with tens of thousands of threads and a memory system that is the actual bottleneck. Most of these kernels do very little arithmetic per byte loaded, so the goal is not “compute faster”, it is stop wasting memory bandwidth.

flowchart TB
    subgraph gpu["one GPU"]
        direction TB
        hbm["HBM / global memory
~2 TB/s, ~80 GB
this is the bottleneck"]
        l2["L2 cache"]
        smem["shared memory
per block, ~100 TB/s"]
        regs["registers
per thread"]
        hbm <--> l2 <--> smem <--> regs
    end
    hbm -.->|"load128cs
bypasses L1"| regs
    regs -.->|"__shfl_xor_sync
warpReduceSum"| regs

Three ideas do most of the work:

Technique	What it buys you
x128 vectorized load/store	One 128-bit transaction instead of 8 scalar ones. Memory-bound kernels get most of their speed here.
warp/block reductions	A sum over C without atomicAdd serialisation — and deterministic, which A4 needs.
kernel fusion	fused_residual_forward5 and fused_classifier_kernel5 avoid writing a [B,T,C] or [B,T,Vp] tensor to HBM only to read it straight back.

Why the classifier is fused. The naive path materialises probs[B,T,Vp]. At B=4, T=1024, Vp=50304 that is ~820 MB in BF16 — written, then read, then written again for the gradient. Fusing softmax + loss + dlogits into one kernel removes both round trips, which is why it is the single biggest kernel in the codebase and worth 10 points on its own.

Why stochastic rounding. Gradients accumulate in FP32 but weights live in BF16, which has 8 mantissa bits. Round-to-nearest turns every update smaller than half an ULP into a no-op, so small consistent gradients never move the weight at all. Rounding probabilistically makes the update correct in expectation. This is the classic silent A3 bug: everything runs, nothing errors, and the model converges to a visibly worse loss.

flowchart LR
    subgraph pat["the three kernel shapes in this codebase"]
        direction TB
        pa["A: elementwise
one thread per x128 chunk
encoder, gelu, residual"]
        pb["B: row reduction over C
one block per (b,t)
layernorm fwd/bwd"]
        pc["C: row reduction over V
one block per (b,t), online softmax
fused classifier"]
    end

---

Background Reading

Source	Read	Why it matters here
llmc/cuda_utils.cuh and llmc/cuda_common.h	Packed128, load128/store128 + cs variants, warpReduceSum, blockReduce, stochastic_rounding	Provided, and you will use them in nearly every kernel. Read them before starting — most of A3 is composing these four primitives.
NVIDIA, CUDA C++ Programming Guide	Ch. 5 (memory hierarchy), §5.3 (shared memory), §5.4 (warp shuffles)	The mental model for why x128 and warp reductions matter. §5.4 in particular explains the shuffle intrinsics warpReduceSum is built from.
Harris, Optimizing Parallel Reduction in CUDA	all 38 slides	The canonical walkthrough: seven successive versions of a sum reduction, each faster than the last, with the reason for each change. This is the single most useful document for A3.
Milakov & Gimelshein 2018, Online normalizer calculation for softmax	§3, Algorithm 3	The one-pass max-and-sum recurrence used by softmax_forward_kernel5 and prepare_softmax_blockwide3.
Micikevicius et al. 2018, Mixed Precision Training	§3	Why FP32 master weights exist alongside BF16 working weights, and what goes wrong without them. Directly explains the use_master_weights path in Part 7.
Gupta et al. 2015, Deep Learning with Limited Numerical Precision	§2.2	Stochastic rounding, and the experiment showing round-to-nearest stalls training at low precision. This is the justification for the warning in Part 7.
NVIDIA, cuBLASLt docs	skim	Background only — matmul_cublaslt is provided. Useful for understanding what a “fused epilogue” is when you meet gelu_fusion.

---

Key Data Types

// floatX is the training precision: BF16 by default, FP16 or FP32 with flags
typedef __nv_bfloat16 floatX;   // default

// x128: 128-bit packed load of 8 x BF16 (or 4 x FP32)
typedef Packed128<floatX> x128;

// f128: 128-bit packed load of 4 x float (always FP32)
typedef Packed128<float> f128;

Note x128::size is precision-dependent: 8 in BF16/FP16 mode, 4 in FP32 mode. Never hard-code it.

Always load and store in x128 chunks. The load128cs (cache-streaming) and store128cs variants bypass L1 for data that won’t be reused. Use load128cs for large sequential reads and store128cs for write-only outputs.

---

Parallelism Strategy

Most kernels follow one of three patterns:

Pattern A — One thread per output element:

idx = blockIdx.x * blockDim.x + threadIdx.x
out[idx] = f(inp[idx])

Used for: encoder_forward_kernel3, gelu_forward_kernel2, residual_forward_kernel

Pattern B — One warp/block per row (reduction over C):

row = blockIdx.x
each thread processes C/blockDim.x elements of that row
warpReduceSum / blockReduce to aggregate

Used for: layernorm_forward_kernel3, layernorm_backward_kernel10

Pattern C — One block per (b,t) position (reduction over vocab):

position = blockIdx.x  // indexes into B*T
each thread handles part of the V dimension
online max/sum reduction via blockReduce

Used for: fused_classifier_kernel5, prepare_softmax_blockwide3

---

Part 1 — Encoder (llmc/encoder.cuh)

encoder_forward_kernel3

Each thread processes x128::size contiguous channels:

tid = blockIdx.x * blockDim.x + threadIdx.x
c   = tid % (C / x128::size)   // which x128-chunk of channels
bt  = tid / (C / x128::size)   // which (b,t) position
b   = bt / T,  t = bt % T

// load x128 from wte[inp[b,t]*C + c*x128::size]
// load x128 from wpe[t*C + c*x128::size]
// add element-wise and store128cs to out[bt*C + c*x128::size]

Use load128cs for both embedding reads and store128cs for the output.

wte_backward_kernel

The backward pass over wte is complicated because multiple sequence positions may reference the same token. A naive atomicAdd would introduce non-determinism — which A4 then makes visible as a bit-exactness failure across GPU counts. The codebase pre-sorts tokens into buckets by vocabulary index on the CPU, then assigns one CUDA block per bucket. Within the block, warps cooperate over bucket items, and warp 0 accumulates the final gradient using stochastic_rounding.

You do not implement the bucket-sorting logic (it is in the launcher). Focus on the kernel: given bucket_info (start index, size) for this block’s token, accumulate the gradient from all positions in the bucket into a warp-level accumulator, then write back to dwte using stochastic rounding.

wpe_backward_kernel

Simpler: for each position t, accumulate dout[b, t, c] over the batch dimension:

for b in 0..B:
    acc += load128cs(dout + (b*T + t)*C + c)
stochastic_rounding(acc, dwpe + t*C + c, seed ^ t ^ c)

---

Part 2 — LayerNorm (llmc/layernorm.cuh)

layernorm_forward_kernel3

One block per (b, t) row. Each thread computes a partial sum over C/blockDim.x elements. After a blockReduce for mean and variance:

mean = blockReduceSum(local_sum) / C
rstd = rsqrtf(blockReduceSum(local_sq_diff) / C + 1e-5f)
// write mean, rstd to output arrays
// apply normalization with weight and bias

Use load128cs for reading inp, store128 for writing out.

fused_residual_forward_kernel5

Fuses the residual addition with the subsequent LayerNorm. In the training loop the pattern residual_forward + layernorm_forward appears twice per transformer block; fusing saves a full read of a [B,T,C] tensor. The kernel computes both outputs in one launch:

out_residual = inp1 + inp2
out_ln       = layernorm(out_residual, weight, bias)

layernorm_backward_kernel10

One block per (b,t) row. Accumulates dweight and dbias across rows via shared-memory scratch plus a final atomic flush, and computes dinp using the two-reduction formula from A2 Part 2.

---

Part 3 — Attention (llmc/attention.cuh)

The GPU attention path transposes Q, K, V into a layout amenable to batched GEMM, then calls cuBLASLt for the QK^T and att·V products. You implement the surrounding permutation and softmax kernels.

permute_kernel / permute_kernel_backward

Rearrange inp[B, T, 3C] → Q[B, NH, T, HS], K[...], V[...]:

// forward:
q[b, nh, t, hs] = inp[b, t, nh*HS + hs]        (first C channels)
k[b, nh, t, hs] = inp[b, t, C + nh*HS + hs]
v[b, nh, t, hs] = inp[b, t, 2C + nh*HS + hs]

// backward: reverse the mapping, accumulate into dinp

unpermute_kernel / unpermute_kernel_backward

Inverse of permute: merge per-head outputs back into a single [B, T, C] tensor.

softmax_forward_kernel5

Row-wise causal softmax over [B, NH, T, T]. Each block handles one row (b, nh, t1) and only the first t1+1 entries are live. Use the online max-then-sum algorithm:

maxval = -INFINITY, sumval = 0
for each chunk: x = inp[...]
    maxval_new = max(maxval, x)
    sumval = sumval * exp(maxval - maxval_new) + exp(x - maxval_new)
    maxval = maxval_new
// warpReduceMax then warpReduceSum for block-level aggregation
// normalize and write out

softmax_autoregressive_backward_inplace_kernel

Given datt and the post-softmax att, compute dpreatt in place:

dpreatt[t1,t2] = att[t1,t2] * (datt[t1,t2] - dot(att[t1,:], datt[t1,:]))

Use a block reduction for the per-row dot product.

---

Part 4 — Matrix Multiplication (llmc/matmul.cuh)

You do not write a GEMM kernel — no tiling, no tensor-core intrinsics, no software pipelining. cuBLASLt does that, and it is very hard to beat. What you do write is everything around it, which is more than it sounds:

matmul_cublaslt

The wrapper every matmul in the model goes through. Build a cublasLtMatmulDesc and the three cublasLtMatrixLayouts, set TRANSA/TRANSB for the operand order you want, attach the optional epilogue (CUBLASLT_EPILOGUE_BIAS, or the fused ..._GELU_AUX variants), then call cublasLtMatmulAlgoGetHeuristic and cublasLtMatmul with a workspace.

This is where fused epilogues come from: asking cuBLASLt to add the bias and apply GELU while the output tile is still in registers, instead of launching two more kernels that each re-read a [B,T,4C] tensor from HBM. That is what the gelu_fusion flag toggles, and it is a real speedup for free.

matmul_forward_cublaslt

Thin dispatcher: pick the argument order, and if gelu_fusion < 1, run gelu_forward as a separate kernel afterwards instead of in the epilogue.

matmul_backward_bias_kernel9

Reduces dout[B*T, OC] along B*T to get dbias[OC], using warp-level shuffle reductions for the inner sum.

reduce_add_sum_kernel

Sums partial results across thread blocks when the bias gradient doesn’t fit in a single block reduction.

matmul_backward

The launcher. Calls:

1. matmul_cublaslt for dinp = dout @ weight
2. matmul_cublaslt for dweight = dout.T @ inp
3. If bias != NULL: matmul_backward_bias_kernel9 for dbias

---

Part 5 — GeLU (llmc/gelu.cuh)

gelu_forward_kernel2 / gelu_forward — same formula as A1, vectorized over x128::size elements per thread.

gelu_backward_inplace_kernel / gelu_backward_inplace — same derivative as A2, vectorized. Note the signature: it takes a single floatX* d_in_out and works in place, reading the incoming gradient and overwriting it with the GeLU-adjusted one, rather than taking separate dinp/dout buffers as the CPU version did. Launch with block_size = 128.

---

Part 6 — Fused Classifier (llmc/fused_classifier.cuh)

The most complex kernel in the codebase. It fuses softmax, cross-entropy, and the logit gradient into a single kernel, avoiding materializing the [B*T, Vp] probability tensor.

prepare_softmax_blockwide3 (device function)

Computes the softmax normalizer for one (b,t) row using the online algorithm, processing the vocabulary in x128 chunks with each thread holding a running (max, sum) pair. After a block-wide reduction:

return SoftmaxParams{1.0f / sumval, maxval}

fused_classifier_kernel5

Template kernel over <WriteDLogits, WriteProbs>. For each (b,t) position (one block):

1. prepare_softmax_blockwide3 → {scale, offset}
2. loss = -log(exp(logits[target] - offset) * scale)
3. Write to losses[b*T+t]
4. If WriteDLogits: logits[b,t,v] = dloss * (expf(logits[b,t,v] - offset) * scale - (v == target)) — note this writes the gradient in place over the logits buffer
5. If WriteProbs: write expf(...) * scale to probs

The template parameters are compile-time so the branches vanish. The training loop calls with WriteDLogits=True; validation with False.

As in A1/A2, the reduction and the target lookup run over [0, V) while the row stride is Vp.

---

Part 7 — AdamW on the GPU (llmc/adamw.cuh)

The same optimizer you wrote for the CPU in A2, now with three additions.

lerp (device function)

start + weight * (end - start), expressed as two FMAs so the result is exact at the endpoints. Use it for the m and v updates.

adamw_update (device function)

The A2 update, plus:

• a grad_scale multiplier applied to the gradient (this is how A5’s gradient clipping is applied — the gradient buffer itself is never modified)
• selective weight decay: the caller passes weight_decay = 0 for 1D tensors (biases, LayerNorm scales/shifts) and the configured value for 2D tensors. Unlike the CPU version in A2, decay is not uniform.
• optional master weights: when present, the FP32 master copy is the one being updated, and the working-precision copy is written with stochastic_rounding.

adamw_kernel3 — dispatches the device-side adamw_update for the right slice, offsetting each pointer by blockIdx.y * stride first.

init_from_master_kernel — stochastically rounds master_params (float) down into params_memory for one slice. A5’s checkpoint resume uses this to rebuild working-precision weights from the master copy.

adamw_update (host overload) / init_from_master (host) — the launchers. Both launch over a 2D grid, dim3(num_blocks, num_slices), with block_size = 512. Note adamw_update is overloaded: a __device__ function that updates one parameter, and a host function that launches adamw_kernel3. The launcher precomputes beta1_correction = 1 - beta1^t and beta2_correction = 1 - beta2^t so the kernel does not recompute powf per element.

Plain (floatX)value instead of stochastic_rounding introduces a systematic bias that compounds over thousands of steps. This is the single most common silent A3 bug: training appears to work and converges to a visibly worse loss.

---

Part 8 — Wiring it together (train_gpt2.cu)

gpt2_build_from_checkpoint

Same job as A1’s, with two differences: the version must be 3 (fp32) or 5 (bf16) and must be cross-checked against the compiled PRECISION_MODE (loading a bf16 checkpoint into an fp32 build is a fatal error, and vice versa); and weights are streamed to the device with file_to_device rather than fread.

gpt2_forward

Same structure as A1 Part 8, but:

• encoder_forward calls your kernel launcher
• QKV projection and all other GEMMs: matmul_forward_cublaslt (provided)
• Attention: your kernels
• After the attention projection: fused_residual_forward5 — one call replaces residual_forward + layernorm_forward
• MLP: matmul_forward_cublaslt with the gelu_fusion flag, which folds GeLU into the GEMM epilogue
• After each layer except the last: fused_residual_forward5 produces both residual3 and the next layer’s ln1

gpt2_backward_and_reduce

Same structure as A2 Part 3, with CUDA launchers. It also takes grad_accum_steps / micro_step: zero the gradient buffers only on micro_step == 0, and use dloss = 1.0f / (B * T * grad_accum_steps).

This function also carries a TODO (A4) for the NCCL gradient reduction. For A3, leave that part out — with one GPU multi_gpu_async_reduce_gradient returns immediately anyway.

gpt2_update

Drives the host-side adamw_update launcher over all 16 parameter tensors, passing weight_decay only for the 2D ones. It also carries a TODO (A4) for ZeRO-1 sharding and the ncclAllGather; for A3 implement the single-GPU path only.

---

Testing

make train_gpt2cu
make test_gpt2cu
./test_gpt2cu

test_gpt2cu is the CUDA sibling of test_gpt2 and checks against the same gpt2_124M_debug_state.bin reference. It accepts three flags:

Flag	Meaning
-w 0\|1	use master weights
-r 0\|1\|2	recompute level (2 also recomputes LayerNorm)
-ge 0\|1	GeLU fusion in the matmul epilogue

Run all of them — the recompute and fusion paths take different routes through your kernels:

./test_gpt2cu
./test_gpt2cu -w 0
./test_gpt2cu -r 2
./test_gpt2cu -ge 0

Then train for real:

./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin \
               -j dev/data/tinyshakespeare/tiny_shakespeare_val.bin \
               -b 4 -t 1024 -v 20 -x 100

These are single-character flags (-i, -b, -x, …), not GNU long options. --input-bin will not work — the parser reads argv[i][1]. Run ./train_gpt2cu with no arguments to see the full list. The ones you need most in A3:

Flag	Meaning
-i / -j	train / val data shard glob
-b / -t	micro-batch size B / sequence length T
-d	total batch size in tokens (sets gradient accumulation)
-x	number of steps
-v	run validation every N steps
-w	use master weights (default 1)
-r	recompute level
-ge	GeLU fusion
-f	override TF32

MFU is reported in the per-step log line, but only once gpt2_estimate_mfu is implemented in A5. Until then it prints -1. For A3 we measure performance from the reported step time; see below.

---

Performance Targets

GPU	Expected MFU (BF16)	Naive baseline
A100 80GB	≥ 50%	~15%
H100 80GB	≥ 55%	~12%
RTX 4090	≥ 40%	~10%
RTX 3090	≥ 35%	~8%

On a 4090, step time for B=4, T=1024 GPT-2 124M should be under 200 ms. MFU below 20% means something is leaving a lot on the table — profile with Nsight Compute. These numbers depend heavily on which GPU you are on, so treat them as a target rather than a threshold.

---

Common Mistakes

• Hard-coding x128::size as 8. It is 4 in FP32 builds, and the FP32 test configuration will fail.
• Shared memory bank conflicts. If threads in a warp hit the same bank, throughput drops 32×. Pad shared arrays by one element.
• Inconsistent __syncthreads(). Every thread in a block must reach it. Don’t put syncs inside divergent conditionals.
• Misaligned x128 reads. load128 requires 16-byte alignment. C is a multiple of 64 in GPT-2, so c * sizeof(floatX) is aligned whenever c is a multiple of x128::size.
• Forgetting cudaCheck(cudaGetLastError()) after launches. Launch failures are otherwise silent until the next sync.
• Stochastic rounding omitted. See the warning in Part 7.
• Writing gradient into the [V, Vp) padding in the fused classifier.

---

Checking your work

./grade A3

Builds the CUDA trainer and runs ./test_gpt2cu in all four configurations, scoring the kernel checks by how many pass. The MFU and compute-sanitizer lines are reported as SKIP: both depend on your specific GPU, so run them yourself with nsys/ncu and compute-sanitizer ./test_gpt2cu.

---

Scoring

What ./grade A3 checks, and what each part is worth.

Component	Points
encoder_forward_kernel3 correctness	7
layernorm_forward_kernel3 and fused_residual_forward_kernel5 correctness	9
layernorm_backward_kernel10 correctness	7
permute / unpermute kernels correctness	7
softmax_forward_kernel5 and the autoregressive softmax backward correctness	8
fused_classifier_kernel5 correctness	10
matmul_cublaslt, matmul_backward_bias_kernel9, matmul_backward correctness	7
wte_backward_kernel and wpe_backward_kernel correctness	7
adamw.cuh: lerp, adamw_update, selective decay, stochastic rounding	10
gpt2_build_from_checkpoint, gpt2_forward, gpt2_backward_and_reduce, gpt2_update wire the kernels correctly	8
test_gpt2cu passes in all four configurations (-w 0, -r 2, -ge 0, default)	5
Performance: MFU ≥ target on the reference GPU	10
No memory errors (compute-sanitizer)	5
Total	100