Llama 3.1 8B · Grouped-query attention

32 independent transformer blocks › Attention

Prepare Q, K, V

learnedAttention.forward · model.py 146–190

Projects the normalized input into query heads and fewer key/value heads, and rotates queries and keys by position.

\begin{aligned}&{\htmlData{light=q}{{\htmlClass{vl-stream-q}{Q}}}}={\htmlData{light=attention-input}{x}}W_Q\\ &{\htmlData{light=k}{{\htmlClass{vl-stream-k}{K}}}}={\htmlData{light=attention-input}{x}}W_K\\ &{\htmlData{light=v}{{\htmlClass{vl-stream-v}{V}}}}={\htmlData{light=attention-input}{x}}W_V\\ &{\htmlData{light=rotated-queries q-rope}{{\htmlClass{vl-stream-q}{Q_{\mathrm r}}}}}=\operatorname{RoPE}({\htmlData{light=queries q-heads}{{\htmlClass{vl-stream-q}{Q}}}},{\htmlData{light=attention-rotations}{{\htmlClass{vl-stream-context}{R}}}})\\ &{\htmlData{light=rotated-keys k-rope}{{\htmlClass{vl-stream-k}{K_{\mathrm r}}}}}=\operatorname{RoPE}({\htmlData{light=keys k-heads}{{\htmlClass{vl-stream-k}{K}}}},{\htmlData{light=attention-rotations}{{\htmlClass{vl-stream-context}{R}}}})\end{aligned}

in: $x$ [B batch, T positions, 4096 d] $R$ [T positions, 64 pairs]
out: $Q_{\mathrm r}$ [B batch, 32 q heads, T positions, 128 d_h] $K_{\mathrm r}$ [B batch, 8 kv heads, T positions, 128 d_h] $V$ [B batch, 8 kv heads, T positions, 128 d_h]
parameters: 25,165,824 · × 32 blocks = 805,306,368

No biases (model.py 100–127). K and V leave at the key/value head count and are cached at that size; repeat_kv runs only after the cache is read (165–177).

Parameters

25,165,824 parameters per block, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Query projection	$4{,}096 \times 4{,}096$	16,777,216
Key projection	4,096 × (8 × 128)	4,194,304
Value projection	4,096 × (8 × 128)	4,194,304
Per-instance total		25,165,824

Across 32 independent blocks: $25{,}165{,}824 \times 32 = 805{,}306{,}368$ parameters.

The group references these owners; it adds no weights of its own.

Explanation

Three separate learned projections produce queries of 4,096 features and keys and values of 1,024 each: $Q=xW_Q$ , $K=xW_K$ , $V=xW_V$ . Queries are reshaped into 32 heads and keys and values into 8, each of 128 features. RoPE rotates query and key feature pairs by position with the supplied coefficients: $Q_{\mathrm r}=\operatorname{RoPE}(Q,R)$ . Keys and values leave with 8 heads, and the reference caches them at that size; repeat_kv repeats each across its group of 4 query heads only after the cache is read. Only the three projections own parameters; reshaping and rotation reorganize values without learned weights.

Operations: Query projection · Key projection · Value projection · Reshape queries · Reshape keys · Reshape values · RoPE queries · RoPE keys
Intermediate states: Query heads · Key heads

Codelayers.py 258–380 · 2 more files

layers.py 258–262


        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        self.weight = Parameter(torch.Tensor(self.output_size_per_partition, self.in_features))

layers.py 286–290

    def forward(self, input_: torch.Tensor) -> torch.Tensor:  # type: ignore
        # Set up backprop all-reduce.
        input_parallel = copy_to_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = F.linear(input_parallel, self.weight, self.bias)

layers.py 345–349


        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        self.weight = Parameter(torch.Tensor(self.out_features, self.input_size_per_partition))

layers.py 374–380

        # Set up backprop all-reduce.
        if self.input_is_parallel:
            input_parallel = input_
        else:
            input_parallel = scatter_to_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = F.linear(input_parallel, self.weight)

model.py 49–54

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    return freqs_cis

model.py 271–275

        self.freqs_cis = precompute_freqs_cis(
            params.dim // params.n_heads,
            params.max_seq_len * 2,
            params.rope_theta,
        )

model.py 277–282

    @torch.inference_mode()
    def forward(self, tokens: torch.Tensor, start_pos: int):
        _bsz, seqlen = tokens.shape
        h = self.tok_embeddings(tokens)
        self.freqs_cis = self.freqs_cis.to(h.device)
        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]

model.py 45–69

def apply_scaling(freqs: torch.Tensor) -> torch.Tensor:
    # Values obtained from grid search
    scale_factor = 8
    low_freq_factor = 1
    high_freq_factor = 4
    old_context_len = 8192  # original llama3 length

    low_freq_wavelen = old_context_len / low_freq_factor
    high_freq_wavelen = old_context_len / high_freq_factor

    wavelen = 2 * torch.pi / freqs
    new_freqs = torch.where(wavelen > low_freq_wavelen, freqs / scale_factor, freqs)
    smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
    return torch.where(
        (wavelen >= high_freq_wavelen) & (wavelen <= low_freq_wavelen),
        (1 - smooth) * new_freqs / scale_factor + smooth * new_freqs,
        new_freqs,
    )


def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, use_scaled: bool = False):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device, dtype=torch.float32)
    if use_scaled:
        freqs = apply_scaling(freqs)

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

Sources4

Evidence for the operations this group contains. The group itself is a presentation and has no implementation of its own.

Contextual evidence from Query projection

The query projection is bias-free and produces query-head count times head width features.

attn-q · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 93–154

Read code excerpt

        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

        self.wq = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )

… lines 107–152 omitted …

        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

Contextual evidence from Query projection

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

Contextual evidence from Query projection

FairScale 0.4.13 stores linear weights with output features first and calls F.linear. This is operation-context evidence; the Meta requirements leave the FairScale version unpinned.

linear-storage · facebookresearch/fairscale / fairscale/nn/model_parallel/layers.py
Pinned revision 50b06d251730 ↗ · lines 258–380

Read code excerpt


        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        self.weight = Parameter(torch.Tensor(self.output_size_per_partition, self.in_features))

… lines 263–285 omitted …

    def forward(self, input_: torch.Tensor) -> torch.Tensor:  # type: ignore
        # Set up backprop all-reduce.
        input_parallel = copy_to_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = F.linear(input_parallel, self.weight, self.bias)

… lines 291–344 omitted …


        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        self.weight = Parameter(torch.Tensor(self.out_features, self.input_size_per_partition))

… lines 350–373 omitted …

        # Set up backprop all-reduce.
        if self.input_is_parallel:
            input_parallel = input_
        else:
            input_parallel = scatter_to_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = F.linear(input_parallel, self.weight)

Contextual evidence from Key projection

The key projection is bias-free and produces KV-head count times head width features.

attn-k · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 93–154

Read code excerpt

        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

… lines 99–106 omitted …

        self.wk = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )

… lines 114–152 omitted …

        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

Contextual evidence from Value projection

The value projection is bias-free and produces KV-head count times head width features.

attn-v · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 93–154

Read code excerpt

        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

… lines 99–113 omitted …

        self.wv = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )

… lines 121–152 omitted …

        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

Contextual evidence from Reshape queries

Queries reshape into query heads; keys and values reshape into KV heads. The reference transposes into head-first order for score computation.

attn-reshape · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 93–183

Read code excerpt

        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        model_parallel_size = fs_init.get_model_parallel_world_size()
        self.n_local_heads = args.n_heads // model_parallel_size
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

… lines 99–152 omitted …

        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

… lines 159–178 omitted …

        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
        keys = keys.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        values = values.transpose(
            1, 2
        )  # (bs, n_local_heads, cache_len + seqlen, head_dim)

Contextual evidence from RoPE queries

Complex-pair rotations apply to queries and keys, with their original dtype restored.

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

Read code excerpt

def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

… lines 76–155 omitted …

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

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

Contextual evidence from RoPE queries

RoPE coefficients depend on position and configured frequency base; the model slices them by the forward-pass position range.

positions · meta-llama/llama3 / llama/model.py
Pinned revision a0940f9cf706 ↗ · lines 49–282

Read code excerpt

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    return freqs_cis

… lines 55–270 omitted …

        self.freqs_cis = precompute_freqs_cis(
            params.dim // params.n_heads,
            params.max_seq_len * 2,
            params.rope_theta,
        )

… lines 276–276 omitted …

    @torch.inference_mode()
    def forward(self, tokens: torch.Tensor, start_pos: int):
        _bsz, seqlen = tokens.shape
        h = self.tok_embeddings(tokens)
        self.freqs_cis = self.freqs_cis.to(h.device)
        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]

Contextual evidence from RoPE queries

apply_scaling: scale_factor 8, low_freq_factor 1, high_freq_factor 4, old_context_len 8192; precompute_freqs_cis applies it when use_scaled is set (68–69). Apart from this and device placement, the file computes as the pinned Llama 3 model.py.

scaling · meta-llama/llama-models / models/llama3/model.py
Pinned revision 0e0b8c519242 ↗ · lines 45–69

Read code excerpt

def apply_scaling(freqs: torch.Tensor) -> torch.Tensor:
    # Values obtained from grid search
    scale_factor = 8
    low_freq_factor = 1
    high_freq_factor = 4
    old_context_len = 8192  # original llama3 length

    low_freq_wavelen = old_context_len / low_freq_factor
    high_freq_wavelen = old_context_len / high_freq_factor

    wavelen = 2 * torch.pi / freqs
    new_freqs = torch.where(wavelen > low_freq_wavelen, freqs / scale_factor, freqs)
    smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
    return torch.where(
        (wavelen >= high_freq_wavelen) & (wavelen <= low_freq_wavelen),
        (1 - smooth) * new_freqs / scale_factor + smooth * new_freqs,
        new_freqs,
    )


def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, use_scaled: bool = False):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device, dtype=torch.float32)
    if use_scaled:
        freqs = apply_scaling(freqs)