Llama 3 8B · KV cache · OpenModelAtlas

32 independent transformer blocks › Attention

KV cache

state · kept between callsself.cache_k, self.cache_v

The rotated keys and the values of every position so far, kept between forward calls.

{\htmlData{light=kv-cache}{C_K}}[1000]\gets {\htmlData{light=rotated-keys k-rope}{{\htmlClass{vl-stream-k}{K_{\mathrm r}}}}}

{\htmlData{light=kv-cache}{C_V}}[1000]\gets {\htmlData{light=values v-heads}{{\htmlClass{vl-stream-v}{V}}}}

{\htmlData{light=grouped-keys repeat-kv}{{\htmlClass{vl-stream-k}{K_{\mathrm G}}}}},{\htmlData{light=grouped-values repeat-kv}{{\htmlClass{vl-stream-v}{V_{\mathrm G}}}}}\gets\mathrm{repeat}({\htmlData{light=kv-cache}{C_K}}[0:1001],\,{\htmlData{light=kv-cache}{C_V}}[0:1001])

shape: $C_{K}$ cache_k [B batch, 8 kv heads, 1001 positions, 128 d_h] $C_{V}$ cache_v [B batch, 8 kv heads, 1001 positions, 128 d_h]

This call writes row 1,000 and reads rows 0–1,000 (the formula's $[0:1001]$ ends before row 1,001). Keys are stored after the rotation; the ×4 repeat happens after the read, so the cache holds only the 8 key/value heads, split across model-parallel ranks. It is allocated up front with max_seq_len rows, which the caller chooses up to 8,192; the atlas assumes the maximum (1 GiB per sequence across all 32 layers), of which 125.1 MiB is in use after this call. The code stores [B, positions, heads, $d_h$ ]; the drawing puts heads first, as the attention plates do.

Codemodel.py 129–302 · 2 more files

model.py 129–145

        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 160–177

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)

        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)

generation.py 105–109

        if torch.cuda.is_bf16_supported():
            torch.set_default_tensor_type(torch.cuda.BFloat16Tensor)
        else:
            torch.set_default_tensor_type(torch.cuda.HalfTensor)
        model = Transformer(model_args)

model.py 281–302

        self.freqs_cis = self.freqs_cis.to(h.device)
        freqs_cis = self.freqs_cis[start_pos : start_pos + 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)

        for layer in self.layers:
            h = layer(h, start_pos, freqs_cis, mask)
        h = self.norm(h)
        output = self.output(h).float()
        return output

generation.py 67–69

        assert 1 <= max_seq_len <= 8192, f"max_seq_len must be between 1 and 8192, got {max_seq_len}."
        assert os.path.isdir(ckpt_dir), f"Checkpoint directory '{ckpt_dir}' does not exist."
        assert os.path.isfile(tokenizer_path), f"Tokenizer file '{tokenizer_path}' does not exist."

generation.py 95–104

        with open(Path(ckpt_dir) / "params.json", "r") as f:
            params = json.loads(f.read())

        model_args: ModelArgs = ModelArgs(
            max_seq_len=max_seq_len,
            max_batch_size=max_batch_size,
            **params,
        )
        tokenizer = Tokenizer(model_path=tokenizer_path)
        assert model_args.vocab_size == tokenizer.n_words

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

Sources6

Contextual evidence

storage

storage · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 129–145

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

Contextual evidence

update

update · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 160–177

Read code excerpt

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)

        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)

Contextual evidence

dtype

dtype · meta-llama/llama3 / llama/generation.py
Pinned revision a0940f9cf706 ↗ · lines 105–109

Read code excerpt

        if torch.cuda.is_bf16_supported():
            torch.set_default_tensor_type(torch.cuda.BFloat16Tensor)
        else:
            torch.set_default_tensor_type(torch.cuda.HalfTensor)
        model = Transformer(model_args)

Contextual evidence

mask

mask · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 281–302

Read code excerpt

        self.freqs_cis = self.freqs_cis.to(h.device)
        freqs_cis = self.freqs_cis[start_pos : start_pos + 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)

        for layer in self.layers:
            h = layer(h, start_pos, freqs_cis, mask)
        h = self.norm(h)
        output = self.output(h).float()
        return output

Contextual evidence

loader

loader · meta-llama/llama3 / llama/generation.py
Pinned revision a0940f9cf706 ↗ · lines 67–104

Read code excerpt

        assert 1 <= max_seq_len <= 8192, f"max_seq_len must be between 1 and 8192, got {max_seq_len}."
        assert os.path.isdir(ckpt_dir), f"Checkpoint directory '{ckpt_dir}' does not exist."
        assert os.path.isfile(tokenizer_path), f"Tokenizer file '{tokenizer_path}' does not exist."

… lines 70–94 omitted …

        with open(Path(ckpt_dir) / "params.json", "r") as f:
            params = json.loads(f.read())

        model_args: ModelArgs = ModelArgs(
            max_seq_len=max_seq_len,
            max_batch_size=max_batch_size,
            **params,
        )
        tokenizer = Tokenizer(model_path=tokenizer_path)
        assert model_args.vocab_size == tokenizer.n_words

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