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

Retrieve values

×

no parameterstorch.matmul(scores, values)

Combine the grouped value vectors using each query’s attention weights.

H=AVG{\htmlData{light=head-output retrieve}{H}}={\htmlData{light=weights softmax}{A}}{\htmlData{light=grouped-values repeat-kv}{{\htmlClass{vl-stream-v}{V_{\mathrm G}}}}}
in
AA [B batch, 32 q heads, T queries, T keys]VGV_{\mathrm G} [B batch, 32 q heads, T positions, 128 d_h]
out
HH [B batch, 32 q heads, T positions, 128 d_h]
Codemodel.py 175–188
model.py 175–188
        values = repeat_kv(
            values, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)

        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)
        if mask is not None:
            scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)
Sources1

Direct evidence

Attention weights multiply the matching grouped values to produce a field of head outputs.

attn-retrieve · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 175–188

Read code excerpt
        values = repeat_kv(
            values, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)

        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)
        if mask is not None:
            scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)