One OLMo 3 Block, at Toy Size

The problem

Write the block that OLMo 3 stacks 32 or 64 times, exactly as the released implementation computes it, but at toy size. It is the LLaMA (Large Language Model Meta AI) block of section 7 with four changes: QK-norm, a different RoPE (rotary position embedding) layout, grouped-query attention, and the norms moved to the output of each sub-layer.

For x of shape $(B, T, D)$, with $H$ query heads, $G$ key/value heads and head size $d = D/H$, the block does seven things.

The seven steps

1. Projections with QK-norm. $q = \mathrm{RMSNorm}_q(xW_q)$, $k = \mathrm{RMSNorm}_k(xW_k)$, $v = xW_v$. The norms (root mean square layer normalization) act on the whole projection ($Hd$ or $Gd$ features), before it is split into heads.
2. Heads. Reshape $q$ to $(B, H, T, d)$ and $k$, $v$ to $(B, G, T, d)$ (section 7: reshape the last axis first, then transpose).
3. RoPE on $q$ and $k$ (the given apply_rope).
4. GQA (grouped-query attention). Each key/value head serves $H/G$ consecutive query heads: repeat the key and value heads with np.repeat (section 7).
5. Attention with the given boolean mask (causal, or sliding-window), then merge the heads back to $(B, T, D)$ and multiply by $W_o$.
6. Reordered norm and first residual. $h = x + \mathrm{RMSNorm}_{\text{attn}}(\text{attention output})$.
7. Feed-forward part and second residual. $y = h + \mathrm{RMSNorm}_{\text{mlp}}\big((\mathrm{SiLU}(hW_g) \odot hW_u)\,W_d\big)$, where the feed-forward part is the SwiGLU (Swish-gated linear unit) MLP (multilayer perceptron) with gate $W_g$, up projection $W_u$ and down projection $W_d$, and $\odot$ is the elementwise product.

Example. What changes compared with the block of section 7, with the same input $x$:

• In the LLaMA block: $h = x + \mathrm{Attention}(\mathrm{Norm}(x))$. The norm sees the stream $x$ and the attention output is added as it is, with whatever size it has.
• OLMo 3: $h = x + \mathrm{Norm}(\mathrm{Attention}(x))$. The attention reads the raw stream, and its output is normalized before it is added: an attention output of $(300, 400)$ adds about $(0.85, 1.13)$ times the gain.

What to implement

olmo3_block(x, p, n_heads, n_kv_heads, mask). p is a dict with:

• "Wq" of shape $(D, Hd)$, "Wk" and "Wv" of shape $(D, Gd)$, "Wo" of shape $(Hd, D)$;
• "q_norm" of shape $(Hd,)$ and "k_norm" of shape $(Gd,)$: the QK-norm gains;
• "attn_norm" and "mlp_norm" of shape $(D,)$;
• "Wg" and "Wu" of shape $(D, F)$, "Wd" of shape $(F, D)$.

The helpers rmsnorm, apply_rope, attention and silu are given.

Requirements

• The norms of steps 6 and 7 go on the output of each sub-layer, not on its input: this is the "reordered norm" of OLMo 2 and 3, the main difference from the LLaMA block.
• No loops over heads or positions.