Decoder Block 4: RoPE and GQA Are Layout, Not Copies

Lesson

Position information enters this block as a rotation. RoPE splits each 64-value Q and K head into two 32-value halves and rotates the pair (f, f + 32) by an angle that depends on the token position t and the frequency index f. The tables are precomputed: ropeCos and ropeSin hold T rows of 32 values each, and your kernel reads row t, never recomputes it.

The tensor you rotate is stage 3's packed output, and the layout does the work: Q head h lives at columns h * 64, K head h at columns 512 + h * 64, and V at columns 640..767, which you must not touch. Rotation happens in place; there is no separate rotated tensor, and there is no expansion.

That last point is the GQA half of this stage. Eight query heads share two KV heads, and high-level code sometimes makes that literal by repeating K and V four times into an eight-head copy. Here that implementation cannot even be written: the runner allocates no expansion buffer, because the mapping kv_head = q_head / 4 is an index calculation the next stage performs at read time. Logical repetition is not a memory operation, and the workspace map in your evidence shows exactly one (T, 768) tensor.

One thread per (token, head, frequency): T x 10 x 32 threads rotate everything. The hidden cases are precise about the classic mistakes: rotating V, reading another token's table row, rotating adjacent pairs (2f, 2f + 1) instead of the half-split, and mis-addressing the last Q head. Each fails naming the head or region it corrupted. Token 0's rotation is the identity, so a bug that misreads the table shows up from token 1 onward, and one hidden case is exactly that boundary.

With 8 query heads and 2 KV heads, q heads 0,1,2,3 read KV head 0; q heads 4,5,6,7 read KV head 1. RoPE rotates Q/K pairs and must leave V unchanged.

What you will learn

• Apply RoPE's half-split rotation in place inside a packed tensor, by column offset.
• Read the per-position table row for the token being rotated, never a neighbor's.
• Leave V untouched: position information belongs to Q and K only.
• State why GQA sharing is an index mapping and what an expansion copy would cost.

Starter code

    // One thread per (token, head, frequency). Ten heads rotate: Q heads 0..7
    // live at columns h * 64, K heads 0..1 at columns 512 + h * 64. V never
    // rotates. Decode the flat index, then rotate the pair (f, f + 32).
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int half = DH / 2;
    if (idx >= T * (HQ + HKV) * half) return;
    // TODO: f = idx % half; h = (idx / half) % (HQ + HKV); t = idx / (half * (HQ + HKV));
    // TODO: base column = h < HQ ? h * DH : D + (h - HQ) * DH;
    // TODO: rotate qkv[t * QKV_OUT + base + f] and qkv[t * QKV_OUT + base + f + half]
    //       by ropeCos/ropeSin[t * half + f].