Llama 3.1 8B · Repeat heads

32 independent transformer blocks › Attention

Repeat heads

no parameterskeys = repeat_kv(keys, self.n_rep); values = repeat_kv(values, self.n_rep)

Copy each key head and each value head to every query head in its group.

{\htmlData{light=grouped-keys repeat-kv}{{\htmlClass{vl-stream-k}{K_{\mathrm G}}}}}[h]={\htmlData{light=rotated-keys k-rope}{{\htmlClass{vl-stream-k}{K_{\mathrm r}}}}}\big[\lfloor h/4\rfloor\big]

{\htmlData{light=grouped-values repeat-kv}{{\htmlClass{vl-stream-v}{V_{\mathrm G}}}}}[h]={\htmlData{light=values v-heads}{{\htmlClass{vl-stream-v}{V}}}}\big[\lfloor h/4\rfloor\big]

in: $K_{\mathrm r}$ [B batch, 8 kv heads, T positions, 128 d_h] $V$ [B batch, 8 kv heads, T positions, 128 d_h]
out: $K_{\mathrm G},\,V_{\mathrm G}$ [B batch, 32 q heads, T positions, 128 d_h]

repeat_kv copies each key and value head to the 4 query heads of its group, with expand and reshape. The scores and the weighted sum then run over all 32 heads; what grouping saves is the cache and the K and V projections.

Explanation

Heads are counted from zero. The repeat runs after the cache is written and read (model.py 165–177), so the cache holds 8 heads.

One of 8 groups: query heads 0–3 share the same key/value pair. The association repeats across 32 query heads.

Codemodel.py 78–177 · sku_list.py 19, 231–248

model.py 78–87

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
    bs, slen, n_kv_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
    )

model.py 129–144

        self.cache_k = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()

model.py 165–177

        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

        # repeat k/v heads if n_kv_heads < n_heads
        keys = repeat_kv(
            keys, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)
        values = repeat_kv(
            values, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, 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,
        ),

Sources3

Direct evidence

repeat_kv copies each KV head n_rep times (model.py 78–87).

repeat-kv · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 78–87

Read code excerpt

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
    bs, slen, n_kv_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
    )

Direct evidence

The cache is allocated at n_kv heads (129–144), written and read at that size, and the repeat follows the read (165–177).

kv-cache · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 129–177

Read code excerpt

        self.cache_k = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()

… lines 145–164 omitted …

        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

        # repeat k/v heads if n_kv_heads < n_heads
        keys = repeat_kv(
            keys, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)
        values = repeat_kv(
            values, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, 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,
        ),