AI Engineer Exhaustive
MediumPlusgpt-oss

GQA Attention

gpt-oss

Medium

Forward pass combining grouped-query attention, optional sliding window, and learned softmax sinks.

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Every gpt-oss layer uses grouped-query attention. Section 2.2 fixes 64 query heads of width 64 and 8 key-value heads. Eight query heads share each KV head, so the KV cache is (8/64=1/8) the size of full MHA while the residual width stays 2880. The same attention block also applies YaRN-RoPE to (Q) and (K), an optional 128-token sliding window on even layers, and a learned sink logit per query head.

The model card lists these pieces side by side. The composition lives in AttentionBlock plus sdpa in the reference PyTorch file. Pre-norm RMSNorm and the residual add wrap the block; they are part of the layer but not part of the attention scores themselves.

How it works

Let the residual stream be (x\in\mathbb{R}^{T\times d}) with (d=2880). After RMSNorm, a fused linear qkv produces a single tensor of width

d_{\mathrm{qkv}} = d_h\bigl(n_q + 2n_{\mathrm{kv}}\bigr) = 64\cdot(64+16)=5120.

Split that into (Q\in\mathbb{R}^{T\times n_q\times d_h}), (K,V\in\mathbb{R}^{T\times n_{\mathrm{kv}}\times d_h}). The official view then groups queries as

Q \in \mathbb{R}^{T \times n_{\mathrm{kv}} \times q_{\mathrm{mult}} \times d_h},\qquad q_{\mathrm{mult}}=n_q/n_{\mathrm{kv}}=8,

and leaves (K,V) as (\mathbb{R}^{T\times n_{\mathrm{kv}}\times d_h}). Rotary embeddings run on (Q) and (K) only. sdpa expands (K) and (V) along the (q_{\mathrm{mult}}) axis so each of the eight queries in a group dots with the same key and mixes the same value:

s_{t,g,m,j} = \frac{\langle q_{t,g,m},\,k_{j,g}\rangle}{\sqrt{d_h}} + M_{tj}.

(M) is (0) on allowed keys and (-\infty) elsewhere. Even layers add the sliding-window band (i-W<j\le i) with (W=128); odd layers use a full causal triangle. After masking, a per-query-head sink (\sigma_{g,m}) is concatenated as a (T+1)-st logit (see the sinks note). Softmax over keys-plus-sink, drop the sink weight, and mix (V):

o_{t,g,m} = \sum_{j} \alpha_{t,g,m,j}\, v_{j,g}.

Flatten the ((g,m)) axes back to 64 heads and project with out: (\mathbb{R}^{T\times 4096}\to\mathbb{R}^{T\times 2880}). Add the residual: (x\leftarrow x+\mathrm{out}(o)).

Scale is (1/\sqrt{64}=1/8), stored as sm_scale. There is no QK-norm and no attention dropout in the public graph. Both qkv and out are bfloat16 linears with bias (PyTorch nn.Linear default). The paper does not mention those biases; they are an official-code detail.

GQA sharing is only on (K) and (V). Sinks, output projection rows, and the eight queries inside a group are distinct. Repeating KV is an expand, not a copy of learned weights: there are 8 key projections, not 64.

Official code

gpt_oss/torch/model.py: AttentionBlock.forward does RMSNorm, the fused qkv split, the GQA reshape, RotaryEmbedding, then sdpa. sdpa expands KV, applies the causal / window mask, concatenates sinks, softmaxes, and mixes values. Config defaults in ModelConfig match the card: 64 / 8 / 64 and sliding_window=128.

Watch-outs

Sources