Decoder Block 5: Causal GQA Attention, Correct First

Lesson

Everything this path has built converges here. The packed layout supplies Q, K and V by column offset; the GQA mapping picks KV head h / 4 with one integer division; the reduction trees you wrote in the Reduction path compute the row maximum and the denominator; and the stable-softmax rule from that path's softmax stages becomes a graded requirement.

The contract per (query row i, query head h), with blockIdx.x encoding i * 8 + h:

g = h / 4
score_j = dot(Q[i,h], K[j,g]) / 8 for j <= i only
m = max_j(score_j)
p_j = exp(score_j - m) / sum(exp(score_k - m))
out[i,h] = sum_j p_j * V[j,g]

Materializing the score row is deliberate at this stage. The starter stages it in a shared array bounded at the 2048 profile length, and keeping it inspectable is what lets you tell a causal-mask bug from a stability bug: a masked row you can print answers questions an online accumulator cannot yet. Stage 6 removes the materialization; this stage earns the right to.

Two hidden hazards are planted rather than left to chance. Tokens 0 and 40 carry full-magnitude Q and K values whose matched score is 128, and float exp overflows past roughly 88, so a softmax that skips the max subtraction returns inf or NaN on those rows and fails even at one token, where the single causal probability must be exactly 1. And the causal mask is checked from row 0: the first row may attend only to itself, which an unmasked kernel gets wrong immediately.

The starter is correct and serial: thread 0 does everything while 255 threads idle. There is no speed gate to force the cooperative version, but the block reductions are the natural next step, and the timeline will show you exactly what serial rows cost at 2048 tokens.

The fixed driver specifies tensor dimensions and tolerance. For query row i, only keys j <= i participate; the one-token case must normalize to probability 1.

What you will learn

• Compose packed-layout addressing, the GQA head map, and causal masking in one kernel.
• Grade numerical stability as correctness: subtract the row maximum before exponentiating.
• Reuse block-reduction trees for a max and a sum over the same shared score row.
• Explain why materializing the score row is the right correctness-stage trade.

Starter code

    // One block per (query row, query head); this starter leaves 255 of its
    // 256 threads idle while thread 0 walks the row serially. Correct, slow,
    // and the shape the cooperative version keeps: the shared score row is
    // bounded at the 2048 profile length.
    __shared__ float scores[2048];
    int i = blockIdx.x / HQ;
    int h = blockIdx.x % HQ;
    int g = h / (HQ / HKV);
    if (i >= T) return;
    if (threadIdx.x == 0) {
        float m = -3.402823466e+38f;
        for (int j = 0; j <= i; ++j) {
            float dot = 0.0f;
            for (int d = 0; d < DH; ++d)
                dot += qkv[i * QKV_OUT + h * DH + d] * qkv[j * QKV_OUT + D + g * DH + d];
            scores[j] = dot * 0.125f;
            if (scores[j] > m) m = scores[j];
        }
        float denom = 0.0f;
        for (int j = 0; j <= i; ++j) denom += expf(scores[j] - m);
        for (int j = 0; j <= i; ++j) scores[j] = expf(scores[j] - m) / denom;
        for (int d = 0; d < DH; ++d) {
            float acc = 0.0f;
            for (int j = 0; j <= i; ++j) acc += scores[j] * qkv[j * QKV_OUT + D + KV + g * DH + d];
            attnOut[i * D + h * DH + d] = acc;
        }
    }