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 DecodeP1,000

0 127 1,000 4,096 8,191

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

Prepare Q, K, VprojectRoPE Open →

Q_r[B, 32, 1, 128]K_r[B, 8, 1, 128]V[B, 8, 1, 128]KV cacheK8 headsV8 headsreads 0…1,000 Repeat heads8 → 328 → 32 K_G[B, 32, L, 128]V_G[B, 32, L, 128]

{\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}}

\htmlClass{vl-absent}{+\,{\htmlClass{vl-stream-context}{M}}}

\Big)

\cdot

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

H[B, 32, 1, 128]

Merge and projectmerge headsoutput projection Open →

Δ_attn[B, 1, 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, 1 queries, 1001 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)