AI Engineer Exhaustive
HardPlusGemma 3

Gemma 3 Attention Block

Gemma 3

Hard

Full Gemma 3 attention with routing and QK-Norm

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

Gemma 3 attention is grouped-query attention with a per-layer choice of local or global span, RMSNorm on (Q) and (K), and RoPE whose base frequency depends on that choice. The technical report lists those pieces in §2: GQA, pre- and post-RMSNorm around the block, QK-norm in place of Gemma 2’s logit soft-cap, and a 5:1 interleaving of sliding-window layers with global layers. The report does not write the forward equations. The official Attention module in gemma/gm/nn/_modules.py is the concrete graph.

The design is aimed at 128K context without a KV cache that grows in every layer. Local layers attend inside a short window (1024 tokens on 4B–27B; 512 on 270M/1B). Global layers see the full sequence, use a RoPE base of (10^6), and on the long-context sizes apply a RoPE scale of 8 after the 32K→128K extension. QK-norm keeps the logits stable once the residual stream is no longer soft-capped.

How it works

Input (x) has shape [B, T, D]. Positions segment_pos are [B, T]. The incoming attn_mask is [B, T, S] (causal and padding; (S=T) without cache).

Projections. If num_kv_heads == num_heads, a fused qkv_einsum produces (Q,K,V) each [B, T, H, d_h]. Otherwise a q_einsum yields (Q) as [B, T, H, d_h] and a kv_einsum yields (K,V) as [B, T, H_{kv}, d_h]. Released Gemma 3 sizes are GQA except where (H_{kv}=1) (270M, 1B), which is multi-query.

QK-norm, then RoPE, then scale. When use_qk_norm is set (true for all Gemma 3 configs),

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

with last-axis RMS and ((1+\gamma)) (see the QK-norm note). RoPE is applied to (Q) and (K) using the layer’s rope_base_frequency and rope_scale_factor. Queries are then scaled:

Q \leftarrow Q\cdot \alpha,\qquad \alpha = \begin{cases} d_h^{-1/2} & \text{1B--12B}\\ (D/H)^{-1/2} & \text{27B}. \end{cases}

The 27B exception is QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_EMBED_DIM_DIV_NUM_HEADS in _config.py: (D=5376), (H=32), so (\alpha=168^{-1/2}), not (128^{-1/2}). Soft-cap is None; the tanh branch in Attention is unused.

Cache. If a cache is present, (K) and (V) (and positions) are written into a left-aligned buffer of length (S) with dynamic_update_slice. Sliding-window layers need those stored positions later.

Scores. Full attention is

\mathrm{logits}_{b,t,h,s} = \sum_{c=1}^{d_h} Q_{b,t,h,c}\,K_{b,s,h,c}.

GQA reshapes (Q) to [B, T, H_{kv}, G, d_h] with (G=H/H_{kv}), einsums against (K), then flattens heads back. Local layers multiply attn_mask by the sliding mask

M_{t,s}=\mathbf{1}[s>t-W]\,\mathbf{1}[s<t+W]

(create_sliding_mask). Global layers leave attn_mask as the causal mask. Blocked entries are filled with K_MASK = -2.3819763e38. Softmax over (s) produces weights; they mix (V) and the output einsum attn_vec_einsum maps [B, T, H, d_h] back to [B, T, D].

Routing. Transformer.setup stamps each layer with attn_type from GEMMA3_ATTENTION_PATTERN tiled to depth: five LOCAL_SLIDING then one GLOBAL. That type selects both the mask and the RoPE pair ((\theta,\mathrm{scale})). Local: (\theta=10^4), scale 1. Global: (\theta=10^6), scale 8 on 4B/12B/27B.

The paper specifies GQA, QK-norm instead of soft-cap, 5:1 local/global, window 1024, and the two RoPE bases. Projection einsums, the 27B query scale, K_MASK, cache-position masking, and the 512-token small-model window are from the official module.

Official code

Watch-outs

Sources