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=kv-cache}{C_K}}[0:T]\big[\lfloor h/4\rfloor\big]

{\htmlData{light=grouped-values repeat-kv}{{\htmlClass{vl-stream-v}{V_{\mathrm G}}}}}[h]={\htmlData{light=kv-cache}{C_V}}[0:T]\big[\lfloor h/4\rfloor\big]

in: $C_{K}$ [B batch, 8 kv heads, T positions, 128 d_h] $C_{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]

Reads rows 0 … T − 1 of the cache, after this call's write, and copies each key/value head to the query heads that share it.

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,
        ),