The Anatomy of an LLM

The Anatomy of an LLM

Haifeng Jin, September 2026

In this article, we walk through every tensor operation in a modern large language model (LLM). For an ML systems engineer, mastering the exact compute graph is a must-have for optimizing training and inference performance, whereas the statistical theory and modeling assumptions are often secondary. When both are taught together, the abstract math often obscures the mechanical reality.

To reduce cognitive load, this article keeps the math to a strict minimum, using visual diagrams and clean Python code to trace how every tensor moves through the model. All code is available in the open-source readable-llm repository, which implements every operation in pure Python with zero dependencies. At the end of the article, we also put everything together in a complete computation graph so you can visually explore the end-to-end flow.

DISCLAIMER: This implementation does not optimize for lines of code or runtime performance, nor does it explain mathematical theory. Instead, it optimizes for readability on the compute level, while remaining mathematically equivalent to heavily optimized production systems.

LLMs have evolved

Over the past few years, the Transformer architecture has evolved significantly beyond the original "Attention Is All You Need" paper¹. Several key techniques have been introduced and widely adopted as the modern standard, including the decoder-only architecture², Mixture of Experts (MoE)³, Grouped-Query Attention (GQA)⁴, Root Mean Square Normalization (RMSNorm)⁵, and Rotary Position Embedding (RoPE)⁶.

In this article, we walk through an LLM architecture that adopts all these modern techniques, breaking down every tensor operation with clean Python code and visual diagrams.

pipeline

Let's start with an end-to-end abstraction of an LLM pipeline in the following figure. We input a question, and the pipeline appends the answer to the string. Note that we simplified the text by omitting special tokens and chat template markers here (such as <bos>, <start_of_turn>, and <eos>), which models use to mark the start and end of questions and responses.

Throughout the article, we use a single input prompt instance with no batch dimension for simplicity. Most operations process this sequence directly, except for a few cases in the MoE module where we further split that down to a single token.

Notice that pipeline is drawn as a capsule. Throughout this article, a capsule represents a module with internal structure and states that we will further break down.

As shown in the figure below, the pipeline breaks down into three parts: tokenizer.encode, model.generate, and tokenizer.decode. Let's walk through them one by one. Notice the small capsule in the top-right corner of the diagram, which shows which module we are currently breaking down.

Here is what pipeline looks like in Python:

def pipeline(prompt, tokenizer, model):
    # input_ids: [seq_len]
    input_ids = tokenizer.encode(prompt)
    # output_ids: [total_seq_len]
    output_ids = model.generate(input_ids)
    # output_text: string
    output_text = tokenizer.decode(output_ids)
    return output_text

tokenizer.encode

The job of a tokenizer is to convert between text strings and lists of integer token IDs. Let's see how that works in practice.

In the figure below, tokenizer.encode turns the input text into a list of token IDs. To keep things simple, we can imagine that each word or punctuation mark is its own token, though modern tokenizers usually split text into subword pieces. The tokenizer first chops the string into individual tokens (like ['What', 'is', '1', '+', '1', '?']), and then looks each one up in a dictionary that maps tokens to unique integers.

The lookup table (or vocabulary) maps each token string to its assigned ID. It contains all the possible tokens that the LLM can speak:

vocab = {
    "+": 10,
    "1": 16,
    "?": 30,
    "is": 318,
    " 1": 352,
    "What": 1867,
    # ...
}

Notice why the two 1s in our prompt get different IDs (352 and 16): tokenizers often treat whitespace as part of the token, so ' 1' (with a leading space) and '1' (without) get distinct entries.

Here is what tokenizer.encode looks like in Python:

class Tokenizer:
    def __init__(self, vocab):
        self.vocab = vocab
        self.inv_vocab = {token_id: token for token, token_id in vocab.items()}

    def encode(self, text):
        tokens = split_tokens(text)
        # input_ids: [seq_len]
        input_ids = [self.vocab[token] for token in tokens]
        return input_ids

We use seq_len (sequence length) to describe the number of tokens in this sequence. This dimension will appear in tensor shapes throughout the rest of the model.

tokenizer.decode

Similarly, tokenizer.decode takes the full ID sequence (the original prompt plus any newly generated tokens) and converts it back into text using an inverted lookup table:

class Tokenizer:
    # ...
    def decode(self, token_ids):
        # token_ids: [seq_len]
        return "".join(self.inv_vocab[token_id] for token_id in token_ids)

model.generate

model.generate is the module that generates new token IDs and appends them to the sequence, as shown in the figure below.

Those new tokens are generated one by one by greedy_sampler. With each step, the input sequence grows by one token as the newly predicted ID is appended:

At each step, model.generate calls greedy_sampler to predict the next token, appends that token to the sequence, and feeds the longer sequence back in. This loop continues until the model produces an end-of-sequence token or reaches a maximum length limit.

Here is what the generation loop looks like in Python:

MAX_NEW_TOKENS = 128
EOS_TOKEN_ID = 2

class Model:
    # ...
    def generate(self, input_ids):
        # input_ids: [seq_len]
        for _ in range(MAX_NEW_TOKENS):
            # next_token_id: int
            next_token_id = greedy_sampler(self, input_ids)
            # input_ids: [seq_len + 1]
            input_ids = input_ids + [next_token_id]
            if next_token_id == EOS_TOKEN_ID:
                break
        # input_ids: [total_seq_len]
        return input_ids

greedy_sampler

We use greedy sampling here for simplicity. In practice, modern LLMs typically rely on stochastic sampling techniques like temperature scaling, top-$p$ (nucleus) sampling, and top-$k$ sampling to generate more diverse and creative responses instead of always picking the single most probable token.

In a single call to the greedy_sampler, model.predict turns input_ids into logits, and then argmax (marked as A in the diagram) picks the highest-scoring token ID. While capsules represent modules with internal structure and states, circles in our diagrams represent pure functions with no internal states:

Each number in the logits vector represents the score for the token associated with that index. So the size of the vector is vocab_size, exactly one score for each token in the vocabulary. For example, logits[318] is the score for token ID 318, which corresponds to the string 'is' in our vocabulary. argmax simply picks whichever token ID has the highest score.

Here is what argmax looks like in Python:

def argmax(logits):
    # logits: [vocab_size]
    # max_idx: int
    max_idx = 0
    for i in range(len(logits)):
        if logits[i] > logits[max_idx]:
            max_idx = i
    return max_idx

Putting it together, here is what greedy_sampler looks like in Python:

def greedy_sampler(model, input_ids):
    # input_ids: [seq_len]
    # logits: [vocab_size]
    logits = model.predict(input_ids)
    # next_token_id: int
    next_token_id = argmax(logits)
    return next_token_id

model.predict

Now, we are ready to further break down model.predict into smaller components. model.predict is the step that actually calls the neural network model. It runs one forward pass of it.

It turns the input_ids into the logits for the next token across the vocabulary. It breaks down into three parts: embedding, decoder, and lm_head, as shown in the figure below.

Here is what model.predict looks like in Python:

class Model:
    # ...
    def predict(self, input_ids):
        # input_ids: [seq_len]
        # embed_out: [seq_len, hidden_size]
        embed_out = self.embedding(input_ids)
        # decoder_out: [seq_len, hidden_size]
        decoder_out = self.decoder(embed_out)
        # logits: [vocab_size]
        logits = self.lm_head(decoder_out)
        return logits

Now, we will tackle the three components one by one. We will explore embedding and lm_head first before diving into the more complex decoder.

embedding

Zooming into the embedding module, it takes input_ids as input and maps each integer ID to a vector of floating-point numbers, producing the matrix embed_out:

Note that embed_out is the end of the backward pass during training. Because token IDs are discrete integers, gradients cannot backpropagate any further into the input text. Gradients update the weights in embedding_table, and backpropagation stops there.

embedding is a token-wise operation, which means it processes each token in isolation. No token depends on the value or position of other tokens to pass through this module. As shown in the figure below, we can conceptually split the sequence into individual tokens, process each one independently, and combine the vectors back together to form embed_out:

A few quick notes here: 1. Splitting into individual tokens is purely for illustration. Real hardware does not loop through tokens one by one in Python; modern frameworks look up the entire sequence at once in parallel via batched tensor indexing (embedding_table[input_ids]). 2. Most operations in an LLM are token-wise. The attention layer in the decoder blocks is the sole exception where tokens actually interact with one another. 3. The blue box represents learnable weights. Here, embedding_table holds the model parameters updated during training. We will use this same blue style for all weights throughout the rest of the diagrams to distinguish them from intermediate activations (shown with white backgrounds).

During the lookup, each token ID indexes into embedding_table of shape [vocab_size, hidden_size], where each row is an embedding vector corresponding to one specific token ID.

How large is vocab_size in practice? Early models (such as LLaMA 1 and 2) used smaller vocabularies of 32K tokens. Modern open-weights models have shifted toward much larger vocabularies to represent multilingual text and code more compactly. Interestingly, model parameter count does not dictate vocab_size: Qwen 3.6 27B uses a vocab_size of 248K, while Kimi K3 2.8T uses 164K. A larger vocabulary compresses text into fewer tokens, but it also means embedding_table has 150k to 250k rows.

Each embedding vector has a length of hidden_size, an important dimension used throughout the model. Looking up every token ID in the sequence turns the [seq_len] list of IDs into a [seq_len, hidden_size] matrix.

Here is what lookup and the embedding module look like in Python:

def lookup(token_id, embedding_table):
    # token_id: int
    # embedding_table: [vocab_size, hidden_size]
    # token_vec: [hidden_size]
    token_vec = embedding_table[token_id]
    return token_vec

def embedding(input_ids, embedding_table):
    # input_ids: [seq_len]
    # embedding_table: [vocab_size, hidden_size]
    # embed_out: [seq_len, hidden_size]
    embed_out = [lookup(token_id, embedding_table) for token_id in input_ids]
    return embed_out

lm_head

After learning about embedding, whose output is fed into decoder, we turn to lm_head. It takes the output of the decoder stack (decoder_out) and converts it into logits, as shown in the figure below:

Just as embed_out marks the end of backpropagation, the logits mark its start. During training, these logits are compared against the actual next token (the ground truth) to compute the loss. Backpropagation begins at this loss, flowing gradients backward through lm_head, across all decoder blocks, and finally into embedding_table to update all the learnable weights.

During inference, the model does not compute a loss; instead, greedy_sampler uses the logits directly to pick the next token.

You might notice that the input shape of decoder is identical to its output shape. In other words, the tensor shape [seq_len, hidden_size] remains the same after passing through all the decoder blocks.

How large is hidden_size in practice? Unlike vocab_size, which is largely independent of model scale, hidden_size grows directly with parameter count. Smaller models around 7B or 8B typically use a hidden_size of 4,096. Moving up the scale, Qwen 3.6 27B uses 5,120, while Kimi K3 2.8T reaches 7,168. Wider vectors give the model more capacity to hold rich representations for each token as they flow through the stack.

Looking inside lm_head, it breaks down into three sequential parts: rms_norm, logits_matmul, and slice:

Here is what lm_head looks like in Python:

def lm_head(decoder_out, gamma, embedding_table_T):
    # decoder_out: [seq_len, hidden_size]
    # gamma: [hidden_size]
    # embedding_table_T: [hidden_size, vocab_size]

    # rms_out: [seq_len, hidden_size]
    rms_out = rms_norm(decoder_out, gamma)
    # all_logits: [seq_len, vocab_size]
    all_logits = logits_matmul(rms_out, embedding_table_T)
    # logits: [vocab_size]
    logits = slice_last(all_logits)
    return logits

Let's walk through these three operations one by one.

rms_norm

Before projecting decoder_out to vocabulary logits, lm_head first normalizes the input using rms_norm (short for Root Mean Square Normalization, commonly written as RMSNorm). Because normalization does not change tensor dimensions, the input and output shapes are identical:

Zooming into rms_norm: it is a token-wise operation, meaning each row of the input matrix is normalized independently:

Here, gamma is the only learnable weight. It has a shape of [hidden_size], matching the dimension of each token vector. It is multiplied element-wise with each token vector.

Here is the Python code for norm_token, which takes a single token vector and the learnable weight gamma:

EPS = 1e-6

def norm_token(token_vec, gamma):
    # token_vec: [hidden_size]
    sum_of_squares = 0
    for v in token_vec:
        sum_of_squares += v ** 2
    rms = (sum_of_squares / len(token_vec) + EPS) ** 0.5

    # output: [hidden_size]
    output = []
    for i in range(len(token_vec)):
        output.append(token_vec[i] / rms * gamma[i])
    return output

def rms_norm(tensor, gamma):
    # tensor: [seq_len, hidden_size]
    # rms_out: [seq_len, hidden_size]
    rms_out = [norm_token(token_vec, gamma) for token_vec in tensor]
    return rms_out

This rms_norm module is used throughout the model; we will see it again soon in the attention and feed-forward blocks.

logits_matmul

The output of rms_norm feeds directly into logits_matmul, which performs a matrix multiplication to produce all_logits:

You may notice that the output tensor still has the same number of tokens (seq_len), but the length of each token vector has expanded from hidden_size to vocab_size. Each row in all_logits is now a vector of logits, containing an unnormalized prediction score for every token in the vocabulary (as we introduced in greedy_sampler).

Zooming into logits_matmul: like embedding and rms_norm, this is a token-wise operation. The input matrix splits into individual token vectors. Each token vector multiplies the weight matrix embedding_table.T, and the resulting vectors are combined back into all_logits:

Notice that the learnable weight matrix is the transpose of embedding_table. In the embedding step, the table mapped discrete token IDs into dense vectors of length hidden_size. Multiplying by its transpose does the reverse, projecting each token vector from shape [hidden_size] back to [vocab_size]. Reusing the embedding weights here is called weight tying, which saves a significant number of parameters. When weights are untied, models use a separate projection matrix of shape [hidden_size, vocab_size] instead, but the computation is identical.

Because multiplying a vector by a matrix is used in several places throughout the model, we write a general matmul helper in plain Python. Here is what matmul and logits_matmul look like in Python:

def matmul(vec, matrix):
    # vec: [in_dim]
    # matrix: [in_dim, out_dim]
    in_dim = len(vec)
    out_dim = len(matrix[0])
    output = []
    for col in range(out_dim):
        dot_product = sum(vec[k] * matrix[k][col] for k in range(in_dim))
        output.append(dot_product)
    return output

def logits_matmul(rms_out, embedding_table_T):
    # rms_out: [seq_len, hidden_size]
    # embedding_table_T: [hidden_size, vocab_size]

    # all_logits: [seq_len, vocab_size]
    all_logits = [matmul(token_vec, embedding_table_T) for token_vec in rms_out]
    return all_logits

slice

slice is the last step in lm_head, as well as the final operation inside model.predict.

As we saw earlier, the output of model.predict is a single 1D logits vector of shape [vocab_size]. But logits_matmul produces all_logits, a 2D matrix containing seq_len logit vectors. The job of slice is simple: it extracts the very last vector from the matrix and returns it, as shown in the figure below:

That final vector provides the scores for the sampler to pick the next token. Because slicing the last row of a 2D list in Python only requires index -1, the code is straightforward:

def slice_last(all_logits):
    # all_logits: [seq_len, vocab_size]
    # logits: [vocab_size]
    logits = all_logits[-1]
    return logits

decoder

So far, we have explored everything before and after decoder: embedding at the input and lm_head at the output. Now, we are ready to dive into the core and most complex part of the Transformer architecture.

As shown in the figure below, the input and output shapes of the module are identical, [seq_len, hidden_size]:

Zooming into decoder: it consists of a stack of identical decoder blocks chained in series, with each block's output feeding directly into the next:

As a reference, Qwen 3.6 27B has 64 decoder blocks, and Kimi K3 2.8T has 93 decoder blocks.

Because each individual block preserves tensor dimensions, the sequence retains its [seq_len, hidden_size] shape from the first block to the last.

Here is what decoder looks like in Python:

def decoder(embed_out, decoder_blocks):
    # embed_out: [seq_len, hidden_size]
    # decoder_blocks: list of decoder_block layers

    # decoder_out: [seq_len, hidden_size]
    decoder_out = embed_out
    for decoder_block in decoder_blocks:
        decoder_out = decoder_block(decoder_out)
    return decoder_out

decoder_block

Zooming into a single decoder_block: each block contains two sub-layers with residual connections: grouped query attention (gqa_block) and mixture of experts (moe_block). The input passes through gqa_block and is added back to itself element-wise via a residual connection, and that result passes through moe_block and is added back to itself element-wise:

Because the residual additions are element-wise, the tensor shape [seq_len, hidden_size] remains unchanged through both gqa_block and moe_block.

Here is what a single decoder_block looks like in Python:

def add(tensor_a, tensor_b):
    # tensor_a: [seq_len, hidden_size]
    # tensor_b: [seq_len, hidden_size]
    seq_len = len(tensor_a)
    hidden_size = len(tensor_a[0])

    # output: [seq_len, hidden_size]
    output = [[0] * hidden_size for _ in range(seq_len)]
    for i in range(seq_len):
        for j in range(hidden_size):
            output[i][j] = tensor_a[i][j] + tensor_b[i][j]
    return output

def decoder_block(decoder_in, gqa_block, moe_block):
    # decoder_in: [seq_len, hidden_size]

    # gqa_out: [seq_len, hidden_size]
    gqa_out = gqa_block(decoder_in)
    residual_1 = add(decoder_in, gqa_out)

    # moe_out: [seq_len, hidden_size]
    moe_out = moe_block(residual_1)

    # decoder_out: [seq_len, hidden_size]
    decoder_out = add(residual_1, moe_out)
    return decoder_out

moe_block

Let's dive into moe_block first. As shown in the figure below, the input and output shapes of the module are identical, [seq_len, hidden_size]:

Zooming into moe_block: it consists of three steps in sequence: normalizing tokens with rms_norm (which we introduced earlier in lm_head), computing routing weights with router, and passing them to moe to compute the weighted expert outputs:

Here is what moe_block looks like in Python:

def moe_block(moe_in, rms_norm, router, moe):
    # moe_in: [seq_len, hidden_size]

    # rms_out: [seq_len, hidden_size]
    rms_out = rms_norm(moe_in)

    # top_weights: [seq_len, num_experts]
    top_weights = router(rms_out)

    # moe_out: [seq_len, hidden_size]
    moe_out = moe(rms_out, top_weights)
    return moe_out

A new dimension shows up in top_weights: num_experts. This represents the total number of expert networks available in the layer. We will dive into how tokens route to these experts in the next few sections.

Notice that moe_block normalizes its input with rms_norm right at the start, before computing routing weights and expert outputs and after the start of branching out the residual connection. This arrangement is known as Pre-LN (pre-layer normalization).

This technique works hand in hand with the residual connection. Without Pre-LN, the tensor between the end of one residual connection and the start of the next would be normalized. Every token had to pass through that normalization at each choke point, with no way around it.

With Pre-LN, normalization happens inside the branch between the start and end of the residual connection. This moves normalization off the model's main choke points and places it onto a parallel path, leaving the residual stream untouched.

router

The normalized matrix is fed into router, which outputs a matrix called top_weights with shape [seq_len, num_experts] as shown in the figure below:

num_experts is an important hyperparameter here: it sets the number of subnetworks, also known as experts, in each MoE module. As a reference, Qwen 3.6 35B has 256 experts, while Kimi K3 2.8T uses 896 experts per MoE module.

Zooming into router: it is again a token-wise operation. The input matrix splits into individual token vectors. Each token vector passes through route_token (circled as $R$) along with the weight matrix w_router to compute its expert routing weights. The output row vectors are then stacked back into the matrix top_weights:

Inside route_token, we see another important hyperparameter: TOP_K. In an MoE module, we do not activate all num_experts for every token. Instead, we only select the top ones. TOP_K is the number of experts we select for each token; we will see how they compute outputs in the next few sections.

As a reference, Qwen 3.6 35B activates 9 out of 256 experts per token, while Kimi K3 2.8T activates 16 out of 896 experts per token.

In route_token, we first multiply the token vector by the weight matrix w_router to produce a logits vector, where each value corresponds to one expert.

Then, we construct the top_weights vector. We pick the TOP_K highest scores from the logits vector, set the ones not in TOP_K to 0, and normalize the selected scores with a softmax function so they sum to 1.

For example, with TOP_K = 2, suppose a token produces raw logits [1.2, 1.6, 0.3] for 3 experts. The top two scores are 1.2 (expert 0) and 1.6 (expert 1). Normalizing them with softmax yields approximately .4 and .6. Setting the unselected expert 2 to 0 gives the final vector [.4, .6, 0], matching the first row of top_weights in the diagram above.

Here is what route_token and router look like in Python:

import math

TOP_K = 2

def softmax(scores):
    # scores: list of scalars
    max_score = max(scores)
    exp_scores = [math.exp(s - max_score) for s in scores]
    sum_exp = sum(exp_scores)
    return [s / sum_exp for s in exp_scores]

def route_token(token_vec, w_router):
    # token_vec: [hidden_size]
    # w_router:  [hidden_size, num_experts]

    # logits: [num_experts]
    # example value: [1.2, 1.6, 0.3]
    logits = matmul(token_vec, w_router)

    num_experts = len(logits)

    # indices: [num_experts]
    # example value: [0, 1, 2]
    indices = [i for i in range(num_experts)]

    # sorted_indices: [num_experts]
    # example value: [1, 0, 2]
    sorted_indices = sorted(indices, key=lambda i: logits[i], reverse=True)

    # top_indices: [TOP_K]
    # example value: [1, 0]
    top_indices = sorted_indices[:TOP_K]

    # top_logits: [TOP_K]
    # example value: [1.6, 1.2]
    top_logits = [logits[i] for i in top_indices]

    # top_probs: [TOP_K]
    # example value: [0.6, 0.4]
    top_probs = softmax(top_logits)

    # top_weights: [num_experts]
    # example value: [0.4, 0.6, 0]
    top_weights = [0.0] * num_experts
    for k in range(TOP_K):
        top_weights[top_indices[k]] = top_probs[k]
    return top_weights

def router(rms_out, w_router):
    # rms_out: [seq_len, hidden_size]
    # top_weights: [seq_len, num_experts]
    top_weights = [route_token(token_vec, w_router) for token_vec in rms_out]
    return top_weights

moe

moe takes the normalized tensor rms_out and the router output top_weights as inputs, producing moe_out with the exact same shape as rms_out:

To make this easy to follow, we split rms_out and top_weights into individual token rows. Taking token 0 as an example, rms_out_0 and top_weights_0 are the corresponding rows. We also split top_weights_0 into three scalar weights, one for each expert. Here is how they work together to produce moe_out_0:

First, rms_out_0 feeds into all three experts (expert_0, expert_1, and expert_2), and each expert produces an output vector of shape [hidden_size]. Next, each output vector is multiplied by its corresponding scalar from top_weights_0. This is an element-wise scalar-vector multiplication, scaling every element in the vector. Finally, an element-wise add combines the three scaled vectors into moe_out_0.

In this example, expert_2 is not activated. Its output is multiplied by 0, producing an all-zero vector that contributes nothing to the element-wise add. Its weight was set to 0 during routing because it was not among the TOP_K experts for this token.

Here is what moe_token and moe look like in Python:

def moe_token(token_vec, top_weights, experts):
    # token_vec: [hidden_size]
    # top_weights: [num_experts]
    hidden_size = len(token_vec)

    # moe_out: [hidden_size]
    moe_out = [0.0] * hidden_size
    for i in range(len(experts)):
        # expert_out: [hidden_size]
        expert_out = experts[i](token_vec)
        for j in range(hidden_size):
            moe_out[j] += top_weights[i] * expert_out[j]
    return moe_out

def moe(rms_out, top_weights, experts):
    # rms_out: [seq_len, hidden_size]
    # top_weights: [seq_len, num_experts]

    # moe_out: [seq_len, hidden_size]
    moe_out = [
        moe_token(token_vec, weights, experts)
        for token_vec, weights in zip(rms_out, top_weights)
    ]
    return moe_out

We can use the following mental model to understand how the MoE module works across the full sequence. Each token activates its own top-$k$ experts, and the selected experts can differ from token to token. Tokens are only routed to these activated experts, and their outputs are weighted and added together element-wise to produce moe_out:

Note that there are also dense models that do not use the MoE architecture, such as Llama 3 or Gemma 2. You can think of a dense model as using one big expert that is always activated to process all the tokens.

expert

Now, we are ready to peek into the last module in MoE: expert. Each expert is a SwiGLU (Swish Gated Linear Unit) feed-forward network. The architecture of an expert is shown in the following figure. Again, for better illustration, we only use a single token as input here:

Each expert uses three weight matrices (w_gate, w_up, and w_down) to do the matmuls.

silu (marked as S in the diagram) is the SiLU (Swish) activation function. Both silu and the element-wise multiplication that follows operate on vectors element by element without changing their dimensions. Whenever we omit intermediate output tensors between operations in our diagrams, it means their shape stays exactly the same as the incoming input tensor, which is [inter_size] here.

Here is the Python code for silu:

import math

def silu(x):
    return x / (1.0 + math.exp(-x))

Using our matmul helper, here is what expert_token and expert look like in Python:

def expert_token(token_vec, w_gate, w_up, w_down):
    # token_vec: [hidden_size]
    # w_gate, w_up: [hidden_size, inter_size]
    # w_down: [inter_size, hidden_size]

    # x_gate: [inter_size]
    x_gate = matmul(token_vec, w_gate)

    # x_up: [inter_size]
    x_up = matmul(token_vec, w_up)

    # x_act: [inter_size]
    x_act = [silu(x) for x in x_gate]

    # x_inter: [inter_size]
    inter_size = len(x_act)
    x_inter = [x_act[i] * x_up[i] for i in range(inter_size)]

    # x_down: [hidden_size]
    x_down = matmul(x_inter, w_down)
    return x_down

def expert(tensor, w_gate, w_up, w_down):
    # tensor: [seq_len, hidden_size]
    # output: [seq_len, hidden_size]
    output = [expert_token(token_vec, w_gate, w_up, w_down) for token_vec in tensor]
    return output

As a reference for the intermediate dimension (inter_size), Qwen 3.6 35B uses an inter_size of 512 per expert, while Kimi K3 2.8T uses 3,584. In MoE models with many experts, this inner dimension is often kept relatively narrow for each individual expert so that activating multiple experts per token remains computationally efficient.

gqa_block

So far, we have covered moe_block, which forms the second half of a single decoder_block. The first half is gqa_block. Here, gqa stands for Grouped-Query Attention (GQA), a modern evolution of the multi-head attention mechanism introduced in the original Transformer paper. We will see how it works in detail in the next few sections. As shown in the following figure, the input and output shapes are identical ([seq_len, hidden_size]):

Zooming into gqa_block: it consists of three steps in sequence: rms_norm, gqa, and out_matmul:

Here is what gqa_block looks like in Python:

def gqa_block(gqa_block_in, rms_norm, gqa, out_matmul):
    # gqa_block_in: [seq_len, hidden_size]

    # rms_out: [seq_len, hidden_size]
    rms_out = rms_norm(gqa_block_in)

    # gqa_out: [seq_len, hidden_size]
    gqa_out = gqa(rms_out)

    # gqa_block_out: [seq_len, hidden_size]
    gqa_block_out = out_matmul(gqa_out)
    return gqa_block_out

Notice that each sub-module inside gqa_block shares the exact same input and output shape ([seq_len, hidden_size]).

Just like moe_block, gqa_block follows the Pre-LN design. It normalizes its input with rms_norm after the residual connection has branched out. This keeps normalization off the choke point, leaving the main residual stream unnormalized as it passes through the block.

Because rms_norm is identical to what we explored in lm_head and moe_block, we can jump straight into out_matmul and gqa.

out_matmul

Let's look at out_matmul first. It is the simpler of the two remaining components: just a standard projection that maps the concatenated attention output back to the residual stream. out_matmul takes the output of gqa (gqa_out) and produces gqa_block_out with the exact same shape:

Zooming into out_matmul: like rms_norm and logits_matmul, this is a token-wise operation. The input matrix splits into individual token vectors, each multiplying the weight matrix w_matmul, and the resulting vectors combine back into gqa_block_out:

Because each token vector simply multiplies w_matmul, we can reuse the matmul helper subroutine from the expert section:

def out_matmul(gqa_out, w_matmul):
    # gqa_out: [seq_len, hidden_size]
    # w_matmul: [hidden_size, hidden_size]

    # gqa_block_out: [seq_len, hidden_size]
    gqa_block_out = [matmul(token_vec, w_matmul) for token_vec in gqa_out]
    return gqa_block_out

gqa

gqa takes the normalized tokens rms_out and produces gqa_out, which feeds directly into out_matmul. Its input and output share the exact same shape, [seq_len, hidden_size]:

Zooming into gqa: the attention computation is divided into independent head groups (group_0, group_1, and group_2). To clearly illustrate how the dimensions and projections work in this module and the upcoming group_0 section, we expand the visual representation of hidden_size from 3 cells to 12 cells.

Each group receives the full rms_out matrix of shape [seq_len, hidden_size]. Inside each group, the tokens are projected into Query, Key, and Value representations, and the group produces an output matrix of shape [seq_len, head_dim] (where head_dim = hidden_size // num_groups). Finally, the outputs from all groups are concatenated back along the column dimension to form gqa_out:

Here is what gqa looks like in Python:

def gqa(rms_out, groups):
    # rms_out: [seq_len, hidden_size]
    seq_len = len(rms_out)

    # Each group computes attention on the full rms_out input
    # group_out: [seq_len, head_dim]
    group_outs = [group(rms_out) for group in groups]

    # Concatenate group outputs back along the column dimension
    # gqa_out: [seq_len, hidden_size]
    gqa_out = []
    # Iterate over the rows of gqa_out
    for t in range(seq_len):
        row = []
        # Iterate over groups to concatenate the outputs
        for out in group_outs:
            row.extend(out[t])
        gqa_out.append(row)
    return gqa_out

Here we introduce a new hyperparameter, num_groups, as well as head_dim, which depends on it: head_dim = hidden_size // num_groups. As a reference, Qwen 3.6 27B has 4 groups, and Kimi K3 2.8T does not use this architecture.

group_0

Zooming into group_0: it takes the full rms_out matrix and projects it into query (q), key (k), and value (v) tensors with matmul using the learned weights w_q, w_k, and w_v. Note that these matmuls are also token-wise operations.

Next, q and k pass through rope (circled as P), while v remains unchanged. Then, all three branches feed into attention_head (circled as H) to produce the final output, group_0_out of shape [seq_len, head_dim]:

As you can see, w_q has an extra depth axis (the z-axis in our diagram), so its output also has a z-axis. We call each matrix along this depth axis a head. In group_0, we have one k head and one v head, but multiple q heads. The length of this z-axis is the number of query heads, q_heads.

This introduces another hyperparameter: d_head, the length of each token vector in q, k, and v. Once you choose d_head, the number of query heads q_heads is decided as well. That is because attention_head concatenates all q_heads outputs along their channel dimension into the group output head_dim:

$$ \text{head_dim} = \text{q_heads} \times \text{d_head} $$

In production configurations, architects usually define this in reverse: they pick the total number of attention heads first and derive d_head = hidden_size // num_attention_heads.

For comparison, Qwen 3.6 27B uses 6 query heads per group, while models like Kimi K3 2.8T use a different attention architecture.

Here is what group_0 looks like in Python:

def group_0(rms_out, w_q, w_k, w_v):
    # rms_out: [seq_len, hidden_size]
    # w_q: [q_heads, hidden_size, d_head]
    # w_k: [hidden_size, d_head]
    # w_v: [hidden_size, d_head]

    # Project inputs to q, k, v token-wise using our matmul helper
    # q: [q_heads, seq_len, d_head]
    # k: [seq_len, d_head]
    # v: [seq_len, d_head]
    q = [[matmul(token_vec, w) for token_vec in rms_out] for w in w_q]
    k = [matmul(token_vec, w_k) for token_vec in rms_out]
    v = [matmul(token_vec, w_v) for token_vec in rms_out]

    # Encode positions with rotary embedding
    # rope is token-wise and preserves tensor shape
    q = rope(q)
    k = rope(k)

    # Compute multi-head attention and concatenate heads
    # group_0_out: [seq_len, head_dim]
    group_0_out = attention_head(q, k, v)
    return group_0_out

Even though rope and attention_head are stateless functions, both are central to how attention works. Let's dive into them one by one.

rope

rope (circled as P) is short for Rotary Position Embedding. Earlier transformer architectures injected position information right at the start of the model. Modern LLMs instead move position encoding directly into the attention layer. Notice that only q and k go through rope, while v skips rope entirely.

rope is a token-wise operation: each token vector is processed independently based on its position index, without mixing information across different tokens in the sequence.

In group_0, k is a 2D matrix ([seq_len, d_head]), while q is a 3D tensor with multiple query heads ([q_heads, seq_len, d_head]). The entry-point rope function inspects the input shape and routes the tensor to either rope_2d or rope_3d:

def rope(x):
    # x: [seq_len, d_head] or [q_heads, seq_len, d_head]
    if isinstance(x[0][0], list):
        return rope_3d(x)
    return rope_2d(x)

Both functions handle tensor dimensions by calling a token-level subroutine, rope_token. rope_2d processes each token vector across the sequence, while rope_3d passes each head matrix into rope_2d:

def rope_2d(x):
    # x: [seq_len, d_head]
    # return: [seq_len, d_head]
    return [rope_token(token_vec, pos) for pos, token_vec in enumerate(x)]

def rope_3d(x):
    # x: [q_heads, seq_len, d_head]
    # return: [q_heads, seq_len, d_head]
    return [rope_2d(head) for head in x]

Inside rope_token, the vector is processed two numbers at a time. It steps through consecutive pairs along the token vector, passing each pair together with its token position to rope_pair:

def rope_token(token_vec, pos):
    # token_vec: [d_head]
    # pos: scalar
    d_head = len(token_vec)
    # output: [d_head]
    output = []
    # Iterate over consecutive coordinate pairs: (x0, x1), (x2, x3), ...
    for i in range(0, d_head, 2):
        r0, r1 = rope_pair(token_vec[i], token_vec[i + 1], pos, i, d_head)
        output.append(r0)
        output.append(r1)
    return output

Finally, rope_pair operates on two individual numbers. It updates the pair based on the coordinates of x0 and x1 on the matrix (token position pos and the dimension index i), returning two modified values. The operations here are all on scalars. Feel free to dive deeper on your own.

import math

def rope_pair(x0, x1, pos, i, d_head):
    # The variables in this function are all scalars.
    freq = 1.0 / (10000 ** (i / d_head))
    angle = pos * freq
    cos_val = math.cos(angle)
    sin_val = math.sin(angle)
    return x0 * cos_val - x1 * sin_val, x0 * sin_val + x1 * cos_val

Because rope preserves tensor shapes, q and k emerge with the exact same dimensions they had before entering rope, ready for attention_head.

attention_head

The attention head circled as H takes the query tensor q, key k, and value v. This is the only operation in the LLM that is not a token-wise operation. We have to pass in all the tokens at once instead of one by one. This is the most complex function we need to explain in this article.

Remember that the q tensor has a z-axis of length q_heads? In the attention_head function, we loop over that dimension, calling the subroutine single_attention_head, which processes a single 2D query matrix alongside the shared key and value matrices. We then concatenate the outputs from single_attention_head back to one output tensor of shape [seq_len, head_dim]:

def attention_head(q, k, v):
    # q: [q_heads, seq_len, d_head]
    # k: [seq_len, d_head]
    # v: [seq_len, d_head]
    seq_len = len(k)

    # Compute attention for each query head
    # head_outs: [q_heads, seq_len, d_head]
    head_outs = []
    for single_q in q:
        # head_out: [seq_len, d_head]
        head_out = single_attention_head(single_q, k, v)
        head_outs.append(head_out)

    # Concatenate all head outputs
    # out: [seq_len, head_dim]
    out = []
    for t in range(seq_len):
        # token_out: [head_dim]
        # head_dim == q_heads * d_head
        token_out = []
        for head_out in head_outs:
            # head_out: [seq_len, d_head]
            token_out.extend(head_out[t])
        out.append(token_out)
    return out

Zooming into single_attention_head: this subroutine was also used by traditional multi-head attention (MHA) before GQA was introduced. It takes care of one Q, K, and V head. The function first transposes k into k_T so that we can compute dot products between single_q and keys with matmul. Because we have subsequent operations following the dot products, we iterate over the tokens (rows) in single_q to process them one by one. For each token, it delegates the computation to a token-level subroutine, attention_token, and appends the resulting vector to head_out:

def single_attention_head(single_q, k, v):
    # single_q: [seq_len, d_head]
    # k: [seq_len, d_head]
    # v: [seq_len, d_head]
    seq_len = len(k)
    d_head = len(k[0])

    # Transpose k to compute dot products with q: [d_head, seq_len]
    k_T = [[k[row][col] for row in range(seq_len)] for col in range(d_head)]

    # head_out: [seq_len, d_head]
    head_out = []
    for i, q_token in enumerate(single_q):
        # token_out: [d_head]
        token_out = attention_token(q_token, k_T, v, i)
        head_out.append(token_out)
    return head_out

Finally, attention_token performs scaled dot-product attention for a single token. It first computes dot products with k_T to produce a vector of shape [seq_len]. Those scores are then scaled and converted to attention weights, weights, before multiplying with v.

To keep generation autoregressive, only the first i + 1 values in weights (from position 0 up to and including current position i) are non-zero. The remaining future positions are padded with zeros:

import math

def attention_token(q_token, k_T, v, i):
    # q_token: [d_head]
    # k_T: [d_head, seq_len]
    # v: [seq_len, d_head]
    d_head = len(q_token)
    seq_len = len(v)

    # Compute raw dot products against all keys: [seq_len]
    dot_products = matmul(q_token, k_T)

    # Scale scores up to the current token position i: [i + 1]
    scores = []
    for j in range(i + 1):
        scores.append(dot_products[j] / math.sqrt(d_head))

    # weights: [seq_len]
    # Pad zeros for future tokens to make length [seq_len]
    weights = softmax(scores) + [0.0] * (seq_len - (i + 1))

    # Compute weighted sum of values: [d_head]
    token_out = matmul(weights, v)
    return token_out

Putting It All Together

We put everything together in this figure to show the end-to-end workflow of model.predict. You can hover over any tensor to see its name and shape, or any operation to see its Python code. In the code popups, we used NumPy operations without explicit token-wise iterations to make them concise and easier to read. You can bookmark the page and come back to it anytime you need a quick reference.

You can also explore the complete, runnable Python implementation in the readable-llm repository on GitHub.

References

1. Vaswani, A., & others. (2017). Attention is all you need. NeurIPS. ↩
2. Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). Improving language understanding by generative pre-training. OpenAI. ↩
3. Shazeer, N., & others. (2017). Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. ICLR. ↩
4. Ainslie, J., & others. (2023). GQA: Training generalized multi-query transformer models from multi-head checkpoints. EMNLP. ↩
5. Zhang, B., & Sennrich, R. (2019). Root mean square layer normalization. NeurIPS. ↩
6. Su, J., & others. (2024). RoFormer: Enhanced transformer with rotary position embedding. Neurocomputing. ↩