Llama 3 8B · Repeat heads · OpenModelAtlas

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:1001]\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:1001]\big[\lfloor h/4\rfloor\big]

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

Reads rows 0–1,000 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, 188–207

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 188–207

def llama3_base_models() -> list[Model]:
    return [
        Model(
            core_model_id=CoreModelId.llama3_8b,
            description="Llama 3 8b model",
            huggingface_repo="meta-llama/Llama-3-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": False,
            },
            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 original Llama 3 8B base catalogue entry supplies dimensions, head counts, vocabulary, MLP rounding settings, epsilon and unscaled RoPE settings.

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

Read code excerpt

LLAMA3_VOCAB_SIZE = 128256

… lines 20–187 omitted …

def llama3_base_models() -> list[Model]:
    return [
        Model(
            core_model_id=CoreModelId.llama3_8b,
            description="Llama 3 8b model",
            huggingface_repo="meta-llama/Llama-3-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": False,
            },
            pth_file_count=1,
        ),