Split the fused QKV projection

Background

nanoGPT projects to Q, K, V in one Linear(C, 3C) for speed, then splits the result back into three (.., C) tensors:

q, k, v = self.qkv(x).split(n_embd, dim=-1)

The fused output has its last axis of size 3C laid out as [Q | K | V].

Problem statement

Implement split_qkv(qkv) returning (q, k, v) by splitting the last axis into three equal parts.

Input

• qkv — array-like whose last axis has size 3C (e.g. shape (T, 3C) or (B, T, 3C)).

Output

Returns (q, k, v), each an np.ndarray with last axis C (the first, middle, last thirds).

Examples

Example 1

Input:  qkv = [[1, 2, 3, 4, 5, 6]]      # C = 2
Output: q = [[1, 2]], k = [[3, 4]], v = [[5, 6]]

Constraints

• C = last_axis // 3.
• q is the first C, k the next C, v the last C.
• Works for any leading dimensions.

Notes

• One big matmul + a cheap split beats three separate Linear(C, C) calls — same math, better memory locality.