Meta has recently announced the release of LLaMA 3.3, the latest iteration in their state-of-the-art open-source foundational large language models (LLMs).

This new version, particularly the LLaMA 3.3 70B model, demonstrates remarkable performance that either matches or surpasses the capabilities of OpenAI's closed-source GPT-4o across a wide range of benchmarks. LLaMA 3 represents a significant leap in large language model development, building upon Meta's ongoing commitment to open-source AI innovation. In this blog post, we will look at the architectural advancements and optimizations that set LLaMA 3 apart from earlier models and the original transformer framework.

Prerequisites

Before diving into the model's code, we will first explore the most significant architectural innovation introduced in Llama 3.

Rotary positional Encodings (RoPE)

Rotary Position Embeddings (RoPE) are an advanced method for encoding positional information in transformer models, designed to enhance their ability to process sequential data. Unlike fixed or learned positional embeddings, RoPE introduces a rotational transformation in the query and key vectors of the self-attention mechanism. This approach embeds positional relationships directly into the angular components of these vectors, allowing the model to generalize better across different sequence lengths and maintain relative positional information.

Given the vectors $q_m$ and $k_n$, which represent the embedded vectors of a word at positions $m$ and $n$ in a sequence, the self-attention mechanism between these tokens at different positions is defined by $q_m^\top k_n$ or $q_m k_n^\top$, depending on whether the configuration is row-wise or column-wise.

To incorporate relative positional information, we aim to express the inner product of the query vector $q_m$ and the key vector $k_n$ as a function $g$. This function should depend solely on the word embeddings $x_m$ and $x_n$, as well as their relative position $m - n$. In other words, the inner product should encode positional information solely in a relative manner:

$\langle f_q(x_m, m), f_k(x_n, n) \rangle = g(x_m, x_n, m - n).$

The primary objective is to devise an encoding mechanism such that the functions $f_q(x_m, m)$ and $f_k(x_n, n)$ align with and satisfy this relationship.

This is illustrated in the exemplary figure below. The later the word "dog" appears in the sentence, the more it is rotated. However, the most important observation is that the rotation angle between two word vectors (e.g., between "pig" and "dog") remains consistent if their relative distance stays the same, regardless of their absolute position or the length of the sentence. This is what the existence of such a function $g$ implies.

To derive RoPE, we will focus on the case where $x_q$ and $x_k$ are 2-dimensional. We define their positionally encoded counterparts as:

where $m$ and $n$ denote their positions, such that:

Additionally, we define the initial conditions for $m = n = 0$ as:

To find solutions for $f_q$ and $f_k$, we note that our vectors $x_q$ and $x_k$ are 2-dimensional. This allows us to represent them as a complex number $z_{(q,k)}$ using Euler's formula:

Using this representation, we can define $f_q$, $f_k$, and $g$ as follows:

Here, $R$ represents the magnitude (scaling factor) and $\Theta$ represents the phase (angle of rotation) for the respective functions. This complex representation elegantly encodes both magnitude and phase, enabling us to incorporate relative positional information. As we want to satisfy $\langle f_q(x_m, m), f_k(x_n, n) \rangle = g(x_m, x_n, m - n).$, we can derive the following eqautions for the radial and phase functions $R$ and $\Theta$:

The derivation of the first one is indeed trival. The $\Theta$ equality can be derived as follows:

In this representation, our initial conditions can be defined as:

With the following relationships:

and

Now, setting $m = n$ and leveraging our knowledge about the initial conditions, we obtain:

From this, we can observe that a possible straight forward solution for $R_q$, $R_k$, and $R_g$ can be derived as follows:

From this, we can see that these solutions for $R_q$, $R_k$, and $R_g$ are, in fact, independent of any positional information.

Additionally, we observe that $\Theta$ in $\Theta_q(x_q, m) - \Theta_k(x_k, m) = \theta_q - \theta_k$ is independent of $x_q$ or $x_k$. Since $\Theta_q(x_q, m) - \Theta_k(x_k, m) = \theta_q - \theta_k$, there must exist some function $\psi(m)$ such that:

Now let’s set $n=m+1$, which means we get

As $\Theta_g(x_q, x_k, 1) + \theta_q - \theta_k$ is constant, we can write $\psi$ as

with $\theta = \Theta_g(x_q, x_k, 1) + \theta_q - \theta_k \neq 0$ and $\gamma \in \mathcal{R}$.

This means when now have a way to write our encoding functions $f_q$and $f_k$:

As we dont make assumtions about $f_q(x_q,0)$ and $f_k(x_k,0)$ we choose

Therefore we can derive the following solution:

Instead of using Euler's formula with complex numbers, we can express rotations in the attention mechanism using rotation matrices. Let's define a rotation matrix $R^d_{\Theta,m}$ which rotates 2D vectors by $m\theta$:

We can now rewrite $f_q$ and $f_k$ as follows:

Using this matrix representation, the standard attention mechanism can be expressed as:

the last equality highlights a crucial point: the applied rotation depends solely on the distance between tokens ($m - n$) and not on the sentence length or absolute token positions

Currently, we have only considered 2D vectors for $x_q$ and $x_k$. To generalize this to $x_{q,k} \in \mathbb{R}^d$, we divide the $d$-dimensional vector into $d/2$ two-dimensional subspaces. Each subspace is rotated based on the position of the token in the sequence.

Now given a $x_{m} \in \mathcal{R}^d$, we formulate a $R^d\_{\Theta,m}$, with:

with pre-defined parameters:

This approach takes every pair of $(W_{q,k}x_m)$ and independently rotates it based on the position $m$ and the index $i$ of the pair in $x_m$ ($\theta_i$).

The above figure illustrates this process. For each token embedding, we take successive 2D subvectors and rotate them by a predefined angle $\theta_i$, where $i$ is the index of the subvector within the token embedding, multiplied by $m$, the position of the token in the sequence.

Note that the above formulation of $\mathbf{R}^d_{\Theta,m}$ is highly inefficient, as it constructs a mostly empty $d \times d$ matrix. Therefore, the paper introduces very simple elementwise vector operations, which are equivalent to multiplying $W_{q,k}x_m$with the very sparse $\mathbf{R}^d_{\Theta,m}$. It defines

Now this is the operation that we will implement later in PyTorch.

Grouped-Query Attention (GQA)

Grouped Query Attention (GQA) is an optimization technique in attention mechanisms, aimed at reducing the computational complexity of standard multi-head attention. Instead of computing a separate query for each head, GQA groups multiple heads together to share a single query representation.

The above figure illustrates this clearly. Each rectangle in the figure represents one attention head. In standard attention, each head in the queries attends to the corresponding head in the keys. However, in this approach, multiple query heads are grouped together to attend to a single key head. The standart attention mechanism can be expressed as:

where:

• $X \in \mathbb{R}^{n \times d_{\text{model}}}$ is the input sequence of token embeddings,
• $W^Q, W^K, W^V \in \mathbb{R}^{d_{\text{model}} \times d_{hs}}$ are the weight matrices for the query, key, and value, respectively,
• $Q, K, V \in \mathbb{R}^{n \times d_{hs}}$ are the computed query, key, and value matrices,
• $n$ is the sequence length, and $d_\text{model}$ is the embedding dimension.

Here, $d_{hs} = \text{heads} \times d_h$, where $d_h$ is the head dimension. In practice, $d_\text{model} = \text{heads} \times d_h$, as we later reshape the $Q$, $K$, and $V$ tensors into the shape:

The fundamental difference in GQA lies in the shape of the weight matrices. In standard attention, the weight matrices are of size:

where $d_{hs} = \text{heads} \times d_h$.

In GQA, the weight matrices for $K$ and $V$ are of size:

where $\text{heads}_{\text{kv}}$ refers to the number of key-value heads. For a group size of $\text{group\_size} = 2$, as shown in the figure, the number of key-value heads is reduced:

This will result in the shape $(\text{b}, \text{seq\_len},\text{heads}_{\text{kv}}* d_h)$ for our $K,V$. Such that we can reshape them into:

Now, since $\text{heads}_{\text{kv}} \leq \text{heads}$, we observe that by sharing attention heads between queries, we can drastically reduce the memory footprint.

Preparing the Dataset

We will train out Llama model on the TinyStories datasets. The TinyStories dataset is a dataset of short, narrative-style stories designed to be used for training natural language processing (NLP) models, particularly language models like Llama. These dataset consist of small, self-contained stories that capture a variety of writing styles, including simple plots, dialogues, and events, making them suitable for fine-tuning language models on tasks related to story generation. This is a small example from the TinyStories dataset:

Once upon a time, there was a little car named Beep. Beep loved to go fast and play in the sun. Beep was a healthy car because he always had good fuel. Good fuel made Beep happy and strong. One day, Beep was driving in the park when he saw a big tree. The tree had many leaves that were falling. Beep liked how the leaves fall and wanted to play with them. Beep drove under the tree and watched the leaves fall on him. He laughed and beeped his horn. Beep played with the falling leaves all day. When it was time to go home, Beep knew he needed more fuel. He went to the fuel place and got more healthy fuel. Now, Beep was ready to go fast and play again the next day. And Beep lived happily ever after.

Now let’s look our code that prepares this dataset for training:

import torch
from datasets import load_dataset

torch.set_printoptions(profile="full")
from datasets import load_dataset
from transformers import GPT2TokenizerFast
from model import Llama3, ModelArgs
from transformers import DataCollatorForLanguageModeling
from torch.utils.data import DataLoader
import logging

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s: %(message)s",
    handlers=[logging.StreamHandler()],  # Logs to stdout
)
logger = logging.getLogger()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

max_seq_len = 256
batch_size = 70
dataset = load_dataset("roneneldan/TinyStories", split="train").select(range(1000))
tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
tokenizer.add_special_tokens({"pad_token": "[PAD]"})


def tokenize_function(examples):
    return tokenizer(examples["text"], max_length=max_seq_len, truncation=True)


dataset = dataset.filter(
    lambda example: example["text"] is not None and example["text"].strip() != ""
)

tokenized_dataset = dataset.map(
    tokenize_function, batched=True, remove_columns=["text"]
)

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

dataloader = DataLoader(
    tokenized_dataset,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=lambda feature: Llama3.gen_labels(feature, tokenizer, data_collator),
)

This loads the TinyStories dataset and prepares it for training a language model. It tokenizes the text using GPT2TokenizerFast, ensuring that the sequences are of a fixed length (max_seq_len). The dataset is filtered to remove empty or invalid texts. A DataLoader is created using a custom collate function (Llama3.gen_labels) to generate labels for training, while efficiently handling batch processing with DataCollatorForLanguageModeling, which pads sequences and generates attention masks.

class Llama3
  (...)
  @staticmethod
  def __build_masks(seq_len, attention_mask, device, position=0, training=False):
      causal = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool)).unsqueeze(0)
      (...)
      if attention_mask == None:
          return causal
      attention_mask = attention_mask.unsqueeze(1).repeat(1, seq_len, 1).int()

      return (causal & attention_mask).int()

  @staticmethod
  def gen_labels(labels, tokenizer, data_collator, ignore_index=-100):
      batch = data_collator(labels)
      labels = batch["labels"]
      attention_mask = batch["attention_mask"]
      for i in range(labels.shape[0]):
          l = torch.roll(labels[i], -1)
          target_indices = (l == ignore_index).nonzero(as_tuple=True)[0]
          if len(target_indices) == 0:
              l[-1] = tokenizer.eos_token_id
          else:
              seq_len = len(l) - len(target_indices)
              l[-1] = ignore_index
              l[seq_len - 1] = tokenizer.eos_token_id

          labels[i] = l

      batch["labels"] = labels
      batch["attention_mask"] = Llama3.__build_masks(
          labels.shape[1], attention_mask, device, training=True
      )
  (...)

The gen_labels function is designed to prepare the labels for training by shifting them one token to the left and appending an EOS (End of Sequence) token at the end. At the start of the function, the labels are just a copy of the input. The challenge is to place the EOS token at the correct position, which requires determining the length of the unpadded sequence. If padding is present, the function places the EOS token at the end of the actual data, not on top of padding tokens. If no padding exists, the EOS token is placed at the very end of the sequence.

Once the labels are shifted and the EOS token is placed correctly, the function must handle the attention masks. Since DataCollatorForLanguageModeling provides an array of attention values for a single sequence, the function needs to convert this into a matrix that can be used for causal language modeling. This involves combining the attention mask with a causal mask, which is a lower triangular matrix, ensuring that the model only attends to previous tokens in the sequence and not to any future tokens.

Llama3 - Model Architecture

Llama3 Architecture Overview

Now that we have everything in place, let's look at how the Llama3 architecture is structured. Below is an overview of its design:

As we can see, Llama3 is conceptually very similar to the vanilla transformer architecture. Most importantly, the Llama3 and other GPT-style models only consist of a decoder-only transformer. This means that we won't use the encoder and the cross-attention, which we covered in the previous blog post. Another notable change that stands out is the use of RMSNorm instead of LayerNorm for normalization.

While LayerNorm normalizes the input by using the mean and standard deviation of the tensor, RMSNorm uses the root mean square of the elements, making it simpler and computationally more efficient. Both methods have learnable scaling $\gamma$ and bias $\beta$ parameters, but RMSNorm avoids the need for computing the mean and standard deviation, which can be more efficient in practice.

Additionally, we observe that the positional encoding (RoPE) is applied directly to the queries and keys in the attention layer, rather than at the embedding level. The values, however, are not position encoded.

After that the outputs of the attention layer are again feed into a FFN. This time however we use the so called SwiGLU activation function. The SwiGLU (Swish-Gated Linear Unit) activation function is a variation of the standard activation functions like ReLU and GELU. It is designed to combine the benefits of the Swish activation function and the Gated Linear Unit (GLU). SwiGLU introduces a gating mechanism using a linear function combined with the Swish activation, which helps improve model performance, particularly in transformer-based architectures.

The SwiGLU activation function is mathematically defined as:

Where:

• $Swish(x)$ is the Swish activation function defined as:

  $$\text{Swish}(x) = x \cdot \sigma(x)$$

  with $\sigma(x)$ being the sigmoid function:

  $$\sigma(x) = \frac{1}{1 + e^{-x}}$$

• $h(X)$ is the Gated Linear Unit with: $h(X)=(XW+b) \otimes \sigma(XV+c) \quad b,c \in \mathbb{R}$ with $\otimes$ being the element-wise product

Now that we have an overview of the Llama3 architecture, let's quickly discuss the implementation of the rotary positional encodings in our model. As mentioned earlier, the direct formulation of a large, mostly empty $R^d_{\Theta,m}$ matrix for positional encodings is inefficient and wasteful in terms of memory and computation. Therefore, we pursue a different approach to implement positional encodings.

Below we have the RoPE-class, which implents the positional encoding.

class RoPE:
    @staticmethod
    def compute_freq(head_dim: int, seq_len: int, base: int = 10000):
        exp = -2 * torch.arange(0, head_dim, 2).float() / head_dim
        thetas = torch.pow(base, exp)
        m = torch.arange(0, seq_len).float()
        freq = torch.outer(m, thetas).float()
        freq_comp = torch.polar(torch.ones_like(freq), freq)
        return freq_comp

    @staticmethod
    def apply_rotary_embedding(x: torch.Tensor, freq_cis):
        # batch,seq_len,heads,d_k
        x_comp = torch.view_as_complex(x.float().reshape(*(x.shape[:-1]), -1, 2))
        freq_com = freq_cis.unsqueeze(0).unsqueeze(2)

        x_out = torch.view_as_real(x_comp * freq_com).reshape(*x.shape)
        return x_out.float()

The compute_freq function precomputes the $m\theta_i$ for $m \leq \text{max\_len}$ and $i \leq \text{d\_model}$.

While the first two lines implement the following equation:

the third line implements the outer product between the array of positions and the $\theta_i$.

Lastly, we use torch.polar, which converts the number into a complex number using the matrix entries as angular components. Mathematically, this whole process can be expressed as:

After the torch.polar operation, which converts the matrix entries into their polar representation, it looks like this:

As we can see the radial components are $1$.

We will now shift our focus on the apply_rotary_embedding function, which given queries $Q$ and keys $K$, will encode them based on their position in the sentence.

let’s demonstrate this problems bei focusing on $polar(\psi)_j$, which contains the angular components for position $j$. Let $x \in \{q,k\} \sub \mathbb{R}^{1 \times \text{d\_model}}$ be the queries ($Q$) and keys($K$). We will transfrom them in the following way:

Now we take this vector of complex entries and multiply it element-wise with $polar(\psi)_j$:

We can now view this result again as a real matrix and reshape it into its original dimension:

As we can see, we arrived at our end result of $R^d_{\Theta,j}x$, with the operations performed in our PyTorch code in the function apply_rotation.

Now that we have covered the most important details that distinguish the Llama3 architecture from a vanilla decoder-only transformer, we can actually take a look at the code of the Llama class.

class FFN(nn.Module):
    def __init__(self, d_model=256, multiple_of=256):
        super(FFN, self).__init__()
        hidden = 4 * d_model
        hidden = int(2 * hidden / 3)

        hidden = multiple_of * ((hidden + multiple_of - 1) // multiple_of)
        self.w1 = nn.Linear(d_model, hidden, bias=False)
        self.v = nn.Linear(d_model, hidden, bias=False)
        self.w2 = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x):
        return self.w2(nn.functional.silu(self.w1(x)) * self.v(x))

class MultiHeadGQAttention(nn.Module):
    flash = False

    def __init__(self, heads=4, d_model=256, group_size=2, max_seq_len=256):
        super(MultiHeadGQAttention, self).__init__()
        self.heads = heads
        self.d_model = d_model
        self.group_size = group_size

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model // group_size, bias=False)
        self.W_v = nn.Linear(d_model, d_model // group_size, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(
        self,
        q,
        k,
        v,
        mask,
        freq_cis,
        position=-1,
    ):
        d_k = self.d_model // self.heads
        bs, seq_len = q.shape[:2]
        q, k, v = self.W_q(q), self.W_k(k), self.W_v(v)

        q = q.view(
            q.shape[0], q.shape[1], self.heads, -1
        )  # (batch,seq_len,heads,d_k) or (batch,1,heads,d_k)
        k = k.view(
            k.shape[0], k.shape[1], self.heads // self.group_size, -1
        )  # (batch,seq_len,heads,d_k) or (batch,1,heads,d_k)
        v = v.view(
            v.shape[0], v.shape[1], self.heads // self.group_size, -1
        )  # (batch,seq_len,heads,d_k) or (batch,1,heads,d_k)
        q = RoPE.apply_rotary_embedding(q, freq_cis)
        k = RoPE.apply_rotary_embedding(k, freq_cis)
        if not self.training:
            k, v = self.kv_cache.update(k, v, position)


        q, k, v = (x.transpose(1, 2) for x in (q, k, v))

        if MultiHeadGQAttention.flash:
            q, k, v = (x.contiguous() for x in (q, k, v))
            output = (
                F.scaled_dot_product_attention(
                    q, k, v, attn_mask=(mask == 1).unsqueeze(1)
                )
                .transpose(1, 2)
                .reshape(bs, seq_len, -1)
            )
        else:
            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
            if mask is not None:
                mask = mask.unsqueeze(1)
                scores = torch.masked_fill(scores, mask == 0, float("-inf"))
                attention = nn.functional.softmax(scores, dim=-1)
                output = (
                    torch.matmul(attention, v)
                    .transpose(1, 2)
                    .contiguous()
                    .view(bs, seq_len, -1)
                )

        output = self.W_o(output)
        return output

class DecoderLayer(nn.Module):
    def __init__(self, d_model=256, heads=4, group_size=2, max_seq_len=256):
        super(DecoderLayer, self).__init__()
        self.norm1 = nn.RMSNorm(d_model, eps=1e-6)
        self.norm2 = nn.RMSNorm(d_model, eps=1e-6)
        self.ffn = FFN(d_model=d_model)
        self.attention = MultiHeadGQAttention(
            heads=heads, d_model=d_model, group_size=group_size, max_seq_len=max_seq_len
        )

    def forward(self, x, tgt_causal_mask, pos, freqs_cis):
        x_norm = self.norm1(x)
        x = x + self.attention(
            x_norm, x_norm, x_norm, tgt_causal_mask, freqs_cis, position=pos
        )
        return x + self.ffn(self.norm2(x))

class Llama3(nn.Module):
    def __init__(self, params: ModelArgs):
        super(Llama3, self).__init__()
        self.tokenizer = params.tokenizer
        self.max_seq_len = params.max_seq_len
        self.ignore_index = params.ignore_index
        self.num_layers = params.num_layers
        self.layers = nn.ModuleList(
            [
                DecoderLayer(
                    d_model=params.d_model,
                    heads=params.heads,
                    group_size=params.group_size,
                    max_seq_len=params.max_seq_len,
                )
                for _ in range(params.num_layers)
            ]
        )
        self.embedding = nn.Embedding(params.vocab_size, params.d_model)
        self.norm = nn.RMSNorm(params.d_model, eps=1e-6)
        self.ffn = nn.Linear(params.d_model, params.vocab_size, bias=False)

        self.d_k = params.d_model // params.heads
        self.freqs_cis = RoPE.compute_freq(
            head_dim=params.d_model // params.heads, seq_len=params.max_seq_len
        )
        self.d_model = params.d_model

        MultiHeadGQAttention.flash = params.use_flash
    (...)
    def calc_loss(self, logits, labels):
        loss = nn.functional.cross_entropy(
            logits.view(-1, logits.shape[-1]),
            labels.view(-1),
            ignore_index=self.ignore_index,
        )
        return loss

    def __run_model(self, tgt, attention_mask, position=0):

        tgt_embed = self.embedding(tgt)
        freqs_cis = self.freqs_cis[position : position + tgt_embed.shape[1]].to(
            tgt.device
        )

        for i in range(self.num_layers):
            tgt_embed = self.layers[i](tgt_embed, attention_mask, position, freqs_cis)
        return self.ffn(self.norm(tgt_embed))

    def forward(self, tgt, attention_mask, labels):
        logits = self.__run_model(tgt, attention_mask)
        return self.calc_loss(logits, labels)

As this is rather much to comprehend, we will start with the forward method of the Llama class. It is easy to see that this just calls __run_model, which precomputes the frequency components for the position encoding, based on the position of the current token. This is a result of the KV-Cache, which we will cover in a bit. The determined $m_j \theta_i$ are then passed to the decoder layers, together with the token embeddings and masks.

The DecoderLayer may be of more interest:

class DecoderLayer(nn.Module):
    def __init__(self, d_model=256, heads=4, group_size=2, max_seq_len=256):
        super(DecoderLayer, self).__init__()
        self.norm1 = nn.RMSNorm(d_model, eps=1e-6)
        self.norm2 = nn.RMSNorm(d_model, eps=1e-6)
        self.ffn = FFN(d_model=d_model)
        self.attention = MultiHeadGQAttention(
            heads=heads, d_model=d_model, group_size=group_size, max_seq_len=max_seq_len
        )

    def forward(self, x, tgt_causal_mask, pos, freqs_cis):
        x_norm = self.norm1(x)
        x = x + self.attention(
            x_norm, x_norm, x_norm, tgt_causal_mask, freqs_cis, position=pos
        )
        return x + self.ffn(self.norm2(x))

As we can see from our architecture overview from before, we first do an RMSNorm before we feed the output into the MultiHeadGQAttention layer. Applying the skip connection and the FFN is pretty self explanotroy. So the most interesting Layer is the MultiHeadGQAttention. So let’s look at that:

class MultiHeadGQAttention(nn.Module):
    flash = False

    def __init__(self, heads=4, d_model=256, group_size=2, max_seq_len=256):
        super(MultiHeadGQAttention, self).__init__()
        self.heads = heads
        self.d_model = d_model
        self.group_size = group_size

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model // group_size, bias=False)
        self.W_v = nn.Linear(d_model, d_model // group_size, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(
        self,
        q,
        k,
        v,
        mask,
        freq_cis,
        position=-1,
    ):
        d_k = self.d_model // self.heads
        bs, seq_len = q.shape[:2]
        q, k, v = self.W_q(q), self.W_k(k), self.W_v(v)

        q = q.view(
            q.shape[0], q.shape[1], self.heads, -1
        )  # (batch,seq_len,heads,d_k) or (batch,1,heads,d_k)
        k = k.view(
            k.shape[0], k.shape[1], self.heads // self.group_size, -1
        )  # (batch,seq_len,heads,d_k) or (batch,1,heads,d_k)
        v = v.view(
            v.shape[0], v.shape[1], self.heads // self.group_size, -1
        )  # (batch,seq_len,heads,d_k) or (batch,1,heads,d_k)
        q = RoPE.apply_rotary_embedding(q, freq_cis)
        k = RoPE.apply_rotary_embedding(k, freq_cis)
        if not self.training:
            k, v = self.kv_cache.update(k, v, position)


        q, k, v = (x.transpose(1, 2) for x in (q, k, v))

        if MultiHeadGQAttention.flash:
            q, k, v = (x.contiguous() for x in (q, k, v))
            output = (
                F.scaled_dot_product_attention(
                    q, k, v, attn_mask=(mask == 1).unsqueeze(1)
                )
                .transpose(1, 2)
                .reshape(bs, seq_len, -1)
            )
        else:
            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
            if mask is not None:
                mask = mask.unsqueeze(1)
                scores = torch.masked_fill(scores, mask == 0, float("-inf"))
                attention = nn.functional.softmax(scores, dim=-1)
                output = (
                    torch.matmul(attention, v)
                    .transpose(1, 2)
                    .contiguous()
                    .view(bs, seq_len, -1)
                )

        output = self.W_o(output)
        return output