Llama 3

Meta · Apr 2024Dense decoder

Grouped-query attention

Reference code meta-llama/llama3 @ a0940f9, configuration meta-llama/llama-models @ 0e0b8c5 · checkpoint, tokenizer and tensors not verified · How to read

Full sequence Prefill Decode

x[B, T, 4096]RoPE coefficientsR[T, 64]

Prepare Q, K, VprojectRoPE Open →

Q_r[B, 32, T, 128]K_r[B, 8, T, 128]V[B, 8, T, 128]Repeat heads8 → 328 → 32 K_G[B, 32, T, 128]V_G[B, 32, T, 128]Causal maskM[T, T]

{\htmlData{light=head-output retrieve merged merge}{H}}

=

\underset{\text{over keys}}{\mathrm{softmax}}

\Big(

{\htmlData{light=rotated-queries q-rope}{{\htmlClass{vl-stream-q}{Q_{\mathrm r}}}}}{\htmlData{light=grouped-keys repeat-kv}{{\htmlClass{vl-stream-k}{K_{\mathrm G}^{\top}}}}}

/\mathrlap{\sqrt{128}}\hphantom{\sqrt{128}}

+\,{\htmlData{light=attention-mask}{{\htmlClass{vl-stream-context}{M}}}}

\Big)

\cdot

{\htmlData{light=grouped-values repeat-kv}{{\htmlClass{vl-stream-v}{V_{\mathrm G}}}}}

H[B, 32, T, 128]

Merge and projectmerge headsoutput projection Open →

Δ_attn[B, T, 4096]

32 independent transformer blocks › Attention

Attention weights

tensorscores · 187

Normalized weights over the allowed key positions for each query.

shape: $A$ [B batch, 32 q heads, T queries, T keys]

Codemodel.py 187

model.py 187

        scores = F.softmax(scores.float(), dim=-1).type_as(xq)

Sources1

Contextual evidence

Softmax normalizes the last score axis in float32 and restores the query dtype.

attn-softmax · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 187–187

Read code excerpt

        scores = F.softmax(scores.float(), dim=-1).type_as(xq)