Attention block param budget = 4·C²

Background

nanoGPT's MultiHeadAttention has no biases and fuses the three projections into one Linear(C, 3C), plus an output projection Linear(C, C). So the parameter count is:

$$\underbrace{3C^2}_{W_Q,W_K,W_V \text{ fused}} + \underbrace{C^2}_{W_O} = 4C^2$$

Crucially this does not depend on n_head — splitting C into more or fewer heads just reorganises the same weights. For GPT-2 small (C=768): 4·768² = 2,359,296.

Problem statement

Implement attention_params(n_embd, n_head=12) returning the attention block's parameter count.

Input

• n_embd — embedding dim C.
• n_head — number of heads (does not affect the count).

Output

Returns an int: 4 * n_embd**2.

Examples

Example 1 — GPT-2 small

Input:  n_embd = 768, n_head = 12
Output: 2359296

Constraints

• Return 4 * n_embd**2.
• The result is independent of n_head.

Notes

• Three fused QKV projections + one output projection, all C×C, no biases — exactly 4C².