Day 99Module 10
in-technical-review

Assembling a forward and a backward pass

A forward pass tells you what the model predicted. Training asks a second question: how would every input, weight and bias have to move to lower the loss? That question changes the shapes of the CUDA work. The same linear layer that reads a weight transposed on the way forward needs a transposed GEMM for its weight gradient, a column reduction for its bias gradient and a mask for the ReLU gradient.

Today is one deliberately small classifier:

Z = X W^T + b
H = ReLU(Z)
L = mean(cross_entropy(softmax(H), labels))

X is [batch, inputs], W is stored [classes, inputs], and H is [batch, classes]. There is no optimiser and no performance claim. The goal is to make one complete forward and backward pass visible before day 100 asks you to repeat it across two layers and many batches.

The forward pass keeps what backward will need

The linear kernel writes both Z and H. Keeping both looks wasteful until backward reaches ReLU:

H[r,c] = max(Z[r,c], 0)

The derivative is one where Z was positive and zero everywhere else. Discarding Z means either reconstructing the mask or storing it in another form. This program keeps the obvious tensor because day 99 is about the dependency, not the storage optimisation.

The forward GEMM is X W^T, not X W. A weight stored as [classes, inputs] puts one class's coefficients next to each other, so one output dot product walks one row of W. CUDA sees only an address. Swapping the two interpretations can stay in bounds and still compute the wrong model, which is why the dimensions live in every formula on this page.

Stable softmax and cross-entropy are one stage

A direct exp(logit) can overflow even when the final probability is representable. Subtracting the row maximum changes neither probability:

p[c] = exp(H[c] - max(H)) / sum_j exp(H[j] - max(H))
loss = log(sum_j exp(H[j] - max(H))) - (H[label] - max(H))

The CUDA kernel assigns one block to one row, reduces the maximum, then reduces the denominator. It also writes

dH = (p - one_hot(label)) / batch

from those same probabilities. Combining loss and dH matters for correctness: two separate softmax implementations could each pass a loose test while disagreeing with each other. PyTorch documents cross-entropy as log-softmax followed by negative log likelihood, the same stable algebra used here: https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html (checked 2026-09-02).

Four lines carry the gradient backward

The calculus is supplied because deriving it is not the CUDA exercise:

dZ = dH * (Z > 0)
dW = dZ^T X
db = column_sum(dZ)
dX = dZ W

The shapes explain the kernels.

Tensor Shape CUDA work
dZ [batch, classes] elementwise ReLU mask
dW [classes, inputs] GEMM with dZ transposed
db [classes] reduce each dZ column over the batch
dX [batch, inputs] GEMM using stored W

The new GEMM orientation is easiest to see in the output indexing:

__global__ void weightGradient(const float* __restrict__ dZ,
                               const float* __restrict__ x,
                               float* __restrict__ dW) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    const int c = blockIdx.y * blockDim.y + threadIdx.y;
    if (i >= kInputs || c >= kClasses) return;
    float sum = 0.0f;
    for (int row = 0; row < kBatch; ++row) {
        sum += dZ[row * kClasses + c] * x[row * kInputs + i];
    }
    dW[c * kInputs + i] = sum;
}

One thread computes one dW[c,i] and walks the batch. That is intentionally naive. Day 44 supplies the tiling work once the operand orientation is correct. db is similarly one thread per class walking a column. Day 26 supplies the parallel reduction after the simple version passes.

Assembly is the launch order

Full source: code/day99-forward-pass/forward_pass.cu.

    linearReluForward<<<dim3(1, 2), tile>>>(dX, dW, dBias, dZ, dH);
    stableSoftmaxCrossEntropy<<<kBatch, kSoftmaxThreads>>>(dH, dLabels, dLoss, dDH);
    reluBackward<<<(nAct + 255) / 256, 256>>>(dZ, dDH, dDZ);
    weightGradient<<<dim3(1, 1), tile>>>(dDZ, dX, dDW);
    biasGradient<<<1, 32>>>(dDZ, dDb);
    inputGradient<<<dim3(1, 2), tile>>>(dDZ, dW, dDX);

The dependency chain is the lesson. dZ cannot start before dH and Z exist. All three parameter/input gradients can start after dZ; this single-stream program keeps them ordered so the dataflow is obvious. A later version could place the independent tails on separate streams, but only a measurement could say whether these tiny kernels repay the launch and synchronisation cost.

The gates cover values and gradients

The host writes the same pass independently in double precision. It does not call the CUDA kernels' helper code. Eight arrays are copied back and checked:

  1. Z, H, per-row loss and dH cover the forward and stable loss stage.
  2. dZ proves that the ReLU mask used the pre-activation sign.
  3. dW proves the transposed GEMM indexing.
  4. db proves the column reduction.
  5. dX proves the gradient reached the layer input.

Every element must satisfy |got - reference| <= 2e-5 + 2e-4 * |reference|. The mixed absolute and relative bound follows day 66: a pure relative check fails near zero, while a pure absolute check grows too loose with magnitude. The remaining difference is ordinary floating-point ordering, because the GPU accumulates in float and the reference accumulates in double.

The program exits with failure on the first tensor whose gate fails, after printing its first bad element. A final PASS therefore means the complete chain agreed with the independent reference, not merely that every launch returned success.

Results

Re-verified on the same Tesla T4 with driver 580.173.02 and CUDA 13.0 (V13.0.88). All eight errors and the mean loss reproduced exactly; only the printed runtime version changed from 12060 to 13000. Both transcripts are listed in front matter. All predictions held:

Gate Max abs error Verdict
Z 2.653e-07 PASS
H 2.292e-07 PASS
loss 2.693e-07 PASS and finite
dH 2.288e-09 PASS
dZ 2.288e-09 PASS
dW 9.460e-09 PASS
db 1.036e-08 PASS
dX 2.779e-09 PASS

Mean cross-entropy was 2.228264. The program printed final PASS and exited 0, so the forward layout, stable loss, ReLU mask, transposed weight gradient, column reduction and input gradient all matched the independent double reference. This run settles correctness only; the program makes no timing or profiler claim.

Exercise

Change weightGradient so it reads dZ[c * kBatch + row], as though dZ were stored transposed. Predict which gates still pass, then build and run. The forward gates, dZ, db and dX do not depend on dW, so only the dW row should fail. That is the useful property of stage-level gradient checks: a wrong parameter gradient does not get hidden behind a later optimiser step.

Pitfalls

The loss is finite on ordinary inputs, so softmax must be stable. Small logits do not exercise overflow. The subtraction is part of the algorithm, not an optimisation to add after a NaN appears.

Backward uses H > 0 instead of Z > 0. That happens to give the same mask for ReLU, but it hides which forward value the derivative depends on and does not generalise to activations whose output loses information. Keep Z until the pass is correct.

db reduces rows. Bias broadcasts across the batch in forward, so backward sums the gradient across that same broadcast dimension. Reducing classes produces one number per example, the wrong shape.

A transposed GEMM is treated as a transposed buffer. dZ^T X describes the indexing of the multiplication. This program never materialises dZ^T. Changing the read pattern is enough.

Go deeper

Next

Day 100 turns these shapes into training: two forward GEMMs, stable loss, the backward kernels, an update and repeated batches. Nothing in that capstone is first contact now. The work is assembly, measurement and proving that the loss actually falls.