32 independent transformer blocks
Causal mask
×supplied · no parametersmask
Identify which key positions each query can attend to.
- shape
- [T queries, T keys]
Codemodel.py 185–186, 284–296
if mask is not None:
scores = scores + mask # (bs, n_local_heads, seqlen, cache_len + seqlen) mask = None
if seqlen > 1:
mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device)
mask = torch.triu(mask, diagonal=1)
# When performing key-value caching, we compute the attention scores
# only for the new sequence. Thus, the matrix of scores is of size
# (seqlen, cache_len + seqlen), and the only masked entries are (i, j) for
# j > cache_len + i, since row i corresponds to token cache_len + i.
mask = torch.hstack(
[torch.zeros((seqlen, start_pos), device=tokens.device), mask]
).type_as(h)Sources1
Contextual evidence
The model constructs an upper-triangular negative-infinity mask; attention adds it to the scaled scores.
attn-mask · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 185–296
Read code excerpt
if mask is not None:
scores = scores + mask # (bs, n_local_heads, seqlen, cache_len + seqlen)… lines 187–283 omitted …
mask = None
if seqlen > 1:
mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device)
mask = torch.triu(mask, diagonal=1)
# When performing key-value caching, we compute the attention scores
# only for the new sequence. Thus, the matrix of scores is of size
# (seqlen, cache_len + seqlen), and the only masked entries are (i, j) for
# j > cache_len + i, since row i corresponds to token cache_len + i.
mask = torch.hstack(
[torch.zeros((seqlen, start_pos), device=tokens.device), mask]
).type_as(h)