Llama 3 8B · Retrieve values

32 independent transformer blocks › Attention

Retrieve values

no parameterstorch.matmul(scores, values)

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

{\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: $A$ [B batch, 32 q heads, 1 query, 1001 keys] $V_{\mathrm G}$ [B batch, 32 q heads, 1001 positions, 128 d_h]
out: $H$ [B batch, 32 q heads, 1 position, 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)