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

32 independent transformer blocks › Attention

Scale similarities

×

no parameters/ math.sqrt(self.head_dim)

Control the magnitude of similarities entering softmax.

Ss=S/128{\htmlData{light=scale}{S_{\mathrm s}}}={\htmlData{light=similarities similarity}{S}}/\sqrt{128}
in
SS [B batch, 32 q heads, 1 queries, 1001 keys]
out
SsS_{\mathrm s} [B batch, 32 q heads, 1 queries, 1001 keys]
Explanation

For independent, zero-mean, unit-variance query and key features, a 128-term dot product has variance 128; dividing by 128\sqrt{128} returns it to 1, so the scores keep the same scale at any head width. Learned features need not meet these assumptions.

Codemodel.py 184 · sku_list.py 19, 231–248
model.py 184
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
sku_list.py 19
LLAMA3_VOCAB_SIZE = 128256
sku_list.py 231–248
        Model(
            core_model_id=CoreModelId.llama3_1_8b,
            description="Llama 3.1 8b model",
            huggingface_repo="meta-llama/Llama-3.1-8B",
            arch_args={
                "dim": 4096,
                "n_layers": 32,
                "n_heads": 32,
                "n_kv_heads": 8,
                "vocab_size": LLAMA3_VOCAB_SIZE,
                "ffn_dim_multiplier": 1.3,
                "multiple_of": 1024,
                "norm_eps": 1e-05,
                "rope_theta": 500000.0,
                "use_scaled_rope": True,
            },
            pth_file_count=1,
        ),
Sources2

Direct evidence

Raw query–key similarities are divided by the square root of head width.

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

Read code excerpt
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)

Direct evidence

Meta's Llama 3.1 8B base catalogue entry supplies dimensions, head counts, the vocabulary (LLAMA3_VOCAB_SIZE, line 19), MLP rounding, epsilon, rope_theta and use_scaled_rope: True.

config · meta-llama/llama-models / models/sku_list.py
Pinned revision 0e0b8c519242 ↗ · lines 19–248

Read code excerpt
LLAMA3_VOCAB_SIZE = 128256

… lines 20–230 omitted …

        Model(
            core_model_id=CoreModelId.llama3_1_8b,
            description="Llama 3.1 8b model",
            huggingface_repo="meta-llama/Llama-3.1-8B",
            arch_args={
                "dim": 4096,
                "n_layers": 32,
                "n_heads": 32,
                "n_kv_heads": 8,
                "vocab_size": LLAMA3_VOCAB_SIZE,
                "ffn_dim_multiplier": 1.3,
                "multiple_of": 1024,
                "norm_eps": 1e-05,
                "rope_theta": 500000.0,
                "use_scaled_rope": True,
            },
            pth_file_count=1,
        ),