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

Query–key dot product

×

no parameterstorch.matmul(xq, keys.transpose(2, 3))

Measure the feature agreement of rotated queries and their grouped keys.

S=QrKGT{\htmlData{light=similarities similarity}{S}}={\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}^{\mathsf T}}}}}
in
QrQ_{\mathrm r} [B batch, 32 q heads, T positions, 128 d_h]KGK_{\mathrm G} [B batch, 32 q heads, T positions, 128 d_h]
out
SS [B batch, 32 q heads, T queries, T keys]
Codemodel.py 179–184
model.py 179–184
        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
        keys = keys.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        values = values.transpose(
            1, 2
        )  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
Sources1

Direct evidence

Queries multiply transposed keys over their common head-feature axis.

attn-dot · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 179–184

Read code excerpt
        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
        keys = keys.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        values = values.transpose(
            1, 2
        )  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)