Decoder Block 3: Packed QKV Projection

Lesson

Stage 2's timeline showed three adjacent projection launches reading the SAME normalized input: Q, then K, then V. Their weight matrices sit contiguously in the packed buffer, all with input width 512, so together they are one 768-row matrix. This stage replaces the three launches with one T x 512 x 768 GEMM whose output columns are partitioned by offset: 0..511 are Q, 512..639 are K, 640..767 are V.

The harness hands you the slab directly: W points at the packed buffer's WQ row 0, and row o of the slab starts at W + o * 512 for all 768 rows. Your kernel fills QKV in the interleaved (T, 768) layout, and verification is slice-aware: land a K value in a V column and the error names the k slice, not a flat index.

The starter already computes this correctly, one thread per output element, and it is already about twice as fast as the three launches it replaced: packing alone removes two launches and two extra passes over X. The gate asks for at least 3.5x, and that gap is the Matrix Multiplication path's lesson applied here: tile X and the slab through shared memory so each 16 x 16 block loads its inputs once and reuses them sixteen times. 768 and 512 are exact tile multiples; only the T direction has a tail, it is graded at T = 17 and 33, and the input carries a poisoned guard region, so a kernel that reads past the last row fails loudly instead of passing by luck.

The SPEEDUP is measured against a trusted three-launch naive baseline timed in the same run on the same GPU, so the judgement travels across hardware. After Run, compare the single projection launch on the timeline with stage 2's three.

The current starter already packs QKV. Your immediate task is tiling and reuse. The historical three-launch baseline is a separate comparison and is not the current starter.

What you will learn

• Treat adjacent weight matrices in a packed buffer as one wider GEMM operand.
• Keep the Q/K/V column partition exact: packed storage is an offset contract.
• Apply shared-memory tiling to a non-square GEMM whose only tail is the token dimension.
• Separate the launch-count win of packing from the reuse win of tiling in measured evidence.

Starter code

    // One thread per output element. Correct at every T, and already about
    // twice as fast as the three separate launches it replaced: one launch,
    // one pass over X. The gate expects the tiled version's further gains.
    int o = blockIdx.x * blockDim.x + threadIdx.x;
    int t = blockIdx.y * blockDim.y + threadIdx.y;
    if (t >= T || o >= QKV_OUT) return;
    float acc = 0.0f;
    for (int i = 0; i < D; ++i) {
        acc += X[t * D + i] * W[o * D + i];
    }
    QKV[t * QKV_OUT + o] = acc;