Inference Attention 2: Flash Prefill Against a Cached Prefix

Lesson

An inference server rarely runs prefill as one giant matrix. vLLM and its peers cut a long prompt into chunks and run each chunk against everything already in the KV cache, so a long prompt never blocks the other users' next tokens. This lab is that kernel: a chunk of C new tokens at positions P to P + C - 1 attends to the cached keys and values at positions 0 to P + C - 1, causally, so query row i sees keys j <= P + i.

The tensors use the attention-v1 contract shared with the course's later decode labs: Llama-3.2-1B attention per layer, 32 query heads, 8 KV heads, head dimension 64, query head h reading KV head h / 4. q is the chunk, [C][32][64]. kCache and vCache hold the whole visible sequence, [P + C][8][64]. out is [C][32][64]. The scale is 1/8 and is exact. There is no workspace argument, so there is nowhere to put a score matrix.

The starter is correct and already online: one warp per (query row, head), a running max m, a running denominator l and a 64-value accumulator, rescaled by exp(m_old - m_new) whenever the max moves. Its cost is architectural. Every query row walks its whole K/V prefix through the memory system on its own, so at C = 2048 and P = 2048 each KV head is read about 2,048 times per query head.

The tiled version changes that multiplier. Keep a tile of query rows resident in one block, stage each K/V tile in shared memory once, and let every resident row consume it before it is replaced. Per row you carry the same three things as the starter: m, l and the accumulator. When a tile arrives with tile maximum m_tile:

m_new = max(m_old, m_tile)
alpha = exp(m_old - m_new)
l_new = alpha * l_old + sum_j exp(s_j - m_new)
acc = alpha * acc + sum_j exp(s_j - m_new) * V[j]

One division by l at the very end normalizes the row. If you finished Reduction 6 or Decoder Block 6, this is the rescale you already wrote; what is new is the cached prefix, the causal edge at P + i, and the serving shapes.

Why the regime matters. Each K/V byte staged for a tile of Br query rows feeds Br rows, about Br / 2 floating-point operations per byte issued, so Br = 64 gives about 32, and K/V tiles often hit L2 on top of that. Prefill attention should therefore be compute-bound, and the profile checks it: the Range pass reports SM throughput and DRAM bytes for your kernel. Shrink the chunk to one token and the same formula gives 0.5 operations per byte: that is decode, the subject of the next lab, and it is bandwidth-bound.

Hidden hazards are planted rather than left to chance. Planted positions carry 4.0 in query head 0 and KV head 0, so a planted query against a planted key scores 128, past float expf's overflow near 88: a softmax that skips the max subtraction returns inf or NaN. One case plants only the last key of the sequence, so a row's maximum first appears in its final tile, and a missing rescale on either l or the accumulator fails there. Chunks of 1, 17 and 33 rows and prefixes of 31 and 1,000 tokens keep every tile edge honest. The output buffer is poisoned, so a row your kernel never writes fails by itself.

The driver times your kernel against the starter's exact kernel in the same run, at C = 2048 and P = 2048, and prints both times and SPEEDUP. This lab is Experimental: correctness is graded now, and a speed gate follows calibration on each GPU. Every value is checked against the attention-v1 oracle (double sums, tolerance 1e-4 plus 1e-3 relative), which the driver first proves against the contract's golden case.

What you will learn

• Write causal attention for a chunk at a prefix offset: query row i sees keys up to position P + i.
• Stage K/V tiles in shared memory once and reuse them across a resident tile of query rows.
• Carry the online softmax state across tiles, rescaling the denominator and the accumulator by exp(m_old - m_new).
• Read SM throughput and DRAM bytes to decide whether a kernel is compute-bound or bandwidth-bound.

Starter code

    // One block per 64 query rows of one head. Each warp walks its 8 rows one at a time,
    // and every row re-reads its whole K/V prefix from global memory: correct, online, slow.
    const int h = blockIdx.y;
    const int g = h / GROUP;
    const int warp = threadIdx.x / 32;
    const int lane = threadIdx.x % 32;
    for (int r = 0; r < 8; ++r) {
        const int row = blockIdx.x * 64 + warp * 8 + r;
        if (row >= chunkLen) return;
        const int pos = prefixLen + row;
        const float* qRow = q + (size_t)row * HQ * DH + h * DH;
        const float q0 = qRow[lane];
        const float q1 = qRow[lane + 32];
        float m = -INFINITY, l = 0.0f, acc0 = 0.0f, acc1 = 0.0f;
        for (int j = 0; j <= pos; ++j) {
            const float* kRow = kCache + (size_t)j * HKV * DH + g * DH;
            const float* vRow = vCache + (size_t)j * HKV * DH + g * DH;
            float dot = q0 * kRow[lane] + q1 * kRow[lane + 32];
            for (int off = 16; off > 0; off >>= 1) dot += __shfl_xor_sync(0xffffffffu, dot, off);
            const float s = dot * 0.125f;
            const float mNew = fmaxf(m, s);
            const float alpha = expf(m - mNew);
            const float p = expf(s - mNew);
            l = alpha * l + p;
            acc0 = alpha * acc0 + p * vRow[lane];
            acc1 = alpha * acc1 + p * vRow[lane + 32];
            m = mNew;
        }
        float* outRow = out + (size_t)row * HQ * DH + h * DH;
        outRow[lane] = acc0 / l;
        outRow[lane + 32] = acc1 / l;
    }