Llama 3.1

All models, with previews →

Meta · Jul 2024Dense decoder · follows Llama 3

Grouped-query attention

Reference code meta-llama/llama3 @ a0940f9, which computes as Llama 3.1 but for its scaled RoPE (meta-llama/llama-models @ 0e0b8c5, model.py 45–69); configuration meta-llama/llama-models @ 0e0b8c5 · checkpoint, tokenizer and tensors not verified · How to read

Full sequencePrefillDecode

32 independent transformer blocks › Attention

Row-wise softmax

×

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

Convert each query row into normalized attention weights.

A[b,i,t,s]=exp⁡(Sm[b,i,t,s])∑u=0T−1exp⁡(Sm[b,i,t,u]){\htmlData{light=weights softmax}{A}}[b,i,t,s]=\frac{\exp({\htmlData{light=causal}{S_{\mathrm m}}}[b,i,t,s])}{\sum_{u=0}^{T-1}\exp({\htmlData{light=causal}{S_{\mathrm m}}}[b,i,t,u])}
in
SmS_{\mathrm m} [B batch, 32 q heads, T queries, T keys]
out
AA [B batch, 32 q heads, T queries, T keys]

In float32, then cast back to the query type.

Explanation

bb, ii, tt and ss index batch, head, query and key; uu runs over the keys.

Codemodel.py 187
model.py 187
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
Sources1

Direct 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)