AI Engineer Exhaustive
MediumPlusGLM-4.5

GQA Attention

GLM-4.5

Medium

Forward pass combining GQA, partial RoPE, and QK-norm before scaled dot-product attention.

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

GLM-4.5 attention is grouped-query attention with two extra twists named in §2.1: partial RoPE, and (on the 355B model only) QK-Norm. The authors also over-provision heads — 96 query heads on a 5120-wide model, about (2.5\times) the usual (D / d_{\mathrm{head}}) count — because that width helped MMLU / BBH even though it did not lower training loss.

GQA is the KV-cache compression. There are 96 query heads and only 8 key/value heads, so each KV head is reused by 12 query heads. Partial RoPE rotates the first half of each head and leaves the rest position-free. QK-Norm is a per-head RMSNorm on (Q) and (K) that keeps attention logits from drifting as the network deepens.

The official Glm4MoeAttention applies those steps in a fixed order: project, reshape to heads, optional QK-Norm, transpose, partial RoPE, then scaled dot-product with KV repeat. Air uses the same graph with use_qk_norm=false.

How it works

Input (x \in \mathbb{R}^{B \times T \times D}) is already RMSNorm-ed by the block. Head width is (d=128). Let (n_q=96), (n_{kv}=8), so the group size is (n_q / n_{kv} = 12).

Q = x W_Q + b_Q, \quad K = x W_K + b_K, \quad V = x W_V + b_V.

(W_Q \in \mathbb{R}^{D \times n_q d}), (W_K,W_V \in \mathbb{R}^{D \times n_{kv} d}). The released configs set attention_bias: true; the output projection (W_O) has no bias.

Reshape to ((B,T,n_{\cdot},d)). When QK-Norm is on (GLM-4.5, not Air), RMSNorm runs on the last axis of (Q) and of (K), independently, with its own (\gamma_Q,\gamma_K \in \mathbb{R}^{d}):

Q \leftarrow \mathrm{RMSNorm}(Q),\qquad K \leftarrow \mathrm{RMSNorm}(K).

Then transpose to ((B,n_{\cdot},T,d)) and apply partial RoPE. Cosine/sine tables have length (d_{\mathrm{rot}} = \lfloor \rho, d\rfloor) with (\rho=0.5), so (d_{\mathrm{rot}}=64). Split each head, rotate the prefix, pass the suffix through:

Q = \big[\,R_{\Theta,t}\,Q_{:d_{\mathrm{rot}}}\;;\; Q_{d_{\mathrm{rot}}:}\,\big],

and the same for (K). (V) is never rotated.

Grouped-query mixing repeats each KV head (12) times along the head axis so (K,V) match (n_q), then runs ordinary causal attention:

\mathrm{Attn}(Q,K,V) = \mathrm{softmax}\!\left(\frac{Q\,\mathrm{repeat}(K)^\top}{\sqrt{d}} + M\right)\mathrm{repeat}(V).

The scale is (d^{-1/2} = 128^{-1/2}), not (D^{-1/2}). Flatten heads and project:

\mathrm{out} = \mathrm{Attn}\,W_O, \qquad W_O \in \mathbb{R}^{n_q d \times D}.

The paper does not write this product. The order QK-Norm → RoPE → GQA is inferred from Transformers / SGLang, not from an equation in §2.1.

Official code

Glm4MoeAttention and apply_rotary_pos_emb in modeling_glm4_moe.py. apply_rotary_pos_emb splits on cos.shape[-1], which is already the partial rotary width. repeat_kv implements the GQA broadcast. SGLang’s Glm4MoeAttention in glm4_moe.py does the same with partial_rotary_factor=0.5.

zai-org/GLM-4.5 does not contain a separate attention kernel; it defers to those implementations.

Watch-outs

Sources