The same thing happens on the implementation side. Squeezing code into as few lines as possible makes it cryptic, while piling on hardware-specific performance optimizations quickly buries the core architecture under complex engineering tricks. That is why our implementation does not optimize for lines of code or raw runtime performance. Instead, it optimizes strictly for readability at the compute level, while remaining mathematically equivalent to heavily optimized production systems.

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.

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)⁶. The architecture we walk through here adopts all of these modern techniques.

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 treat whitespace as part of the token, so ' 1' (with a leading space) and '1' (without) get distinct entries. That is also why ' is' carries its space. Decoding is then a plain string join, with no need to guess where the spaces went.

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

Notice that __init__ also builds self.inv_vocab. This is an inverted dictionary that flips the keys and values, mapping each integer token ID back to its string token. We don't need it inside encode, but it will be used in tokenizer.decode to turn generated token IDs back into text.

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 the inverted lookup table inv_vocab we built earlier:

Here is what tokenizer.decode looks like in Python:

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 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.

Throughout the rest of the article, we quote real numbers from two reference points: Qwen 3.6 and the much larger Kimi K3 2.8T. Qwen 3.6 ships in two flavors, and the difference matters for the numbers we quote. The dense 27B runs every one of its parameters for every single token. The sparse 35B is a Mixture of Experts model: it carries more parameters in total, but each token only activates a small slice of them, so the compute per token stays close to that of a much smaller model.

We use the dense 27B wherever the number does not depend on MoE, and switch to the sparse 35B once we reach the expert-specific dimensions.

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.

RMSNorm is an improved version of the earlier LayerNorm. It keeps the rescaling and drops the mean subtraction, which turned out to matter little in practice.

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

gqa_block runs first inside the block, but moe_block is the simpler of the two, so let's start there. 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 would have 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. This is where we switch to the sparse Qwen 3.6 35B, which 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:

This also gives us a neat way to think about the dense models we mentioned earlier: a dense model is just one big expert that is always activated for every token.

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.

The gate is what makes this a SwiGLU rather than a plain feed-forward network, with w_up producing the values and w_gate one multiplier for each. Why that beats the simpler alternatives is still not well understood, and the paper that introduced these variants is refreshingly blunt about having no explanation for it⁷.

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, 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, while Kimi K3 2.8T uses a different attention architecture entirely.

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.

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. Because the angle is proportional to position, the two rotations cancel down to a difference in the later dot product, so attention sees relative position. 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 dot_product_attention (marked as D in the figure below), which processes a single 2D query matrix alongside the shared key and value matrices. Notice in the figure that q is the only tensor that splits. Every head reads the same k and v. We then concatenate the outputs back into one tensor of shape [seq_len, head_dim], marked as C:

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 q_head in q:
        # head_out: [seq_len, d_head]
        head_out = dot_product_attention(q_head, 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 dot_product_attention: 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 q_head and the keys with matmul. Because more operations follow the dot products, we iterate over the tokens (rows) in q_head and handle 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.

The figure below unrolls that loop and shows all three tokens at once. Every row of q_head takes its own trip through the same steps, and the shared k_T and v feed all of them. Stack the resulting attention weights and you get a triangle, because a token is only allowed to look at itself and the tokens before it:

def dot_product_attention(q_head, k, v):
    # q_head: [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(q_head):
        # 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. To check that the architecture actually works, we trained a 0.000005B model (4,596 parameters) on a single sentence. It can only answer the one prompt "What is 1+1?", which is all we needed. The training code is in the repo too, if you want to see how those weights were produced.

Where to go next

What we built here is a foundation. Most models you will run into are variations on the same shape, and they get much easier to read once you know the baseline. From here you can branch off in whatever direction interests you: other architectures, the theory behind why any of this works, or the performance engineering that separates code written to be read from code written to be fast.

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. ↩
7. Shazeer, N. (2020). GLU variants improve transformer. arXiv Preprint arXiv:2002.05202. ↩

To measure the intelligence scale of artificial intelligence (AI), various benchmark datasets have been designed, analogous to tests for humans. For example, the MMLU dataset⁵ consists of multiple-choice questions to test multi-language understanding, the GPQA dataset⁶ consists of graduate-level questions on a variety of subjects, and the MATH dataset⁷ comprises competition math problems.

However, when experiments are conducted to test AI on these datasets, unlike tests for humans, only the quality of work is emphasized, while the time aspect of the test is typically ignored. Such an approach does not satisfy the requirements of AI in real-world applications.

From a pragmatic perspective, many tasks are time-sensitive, such as autonomous driving and customer service. In these applications, the output of intelligence is required within a limited time window. Even in the early days of deep learning, there was an implicit assumption of a time limit. When AlphaGo⁸ defeated the best human Go player in 2016, it adhered to the same rules, including time constraints, as the human player.

Implications of test-time scaling

With the advent of the test-time scaling law⁹, the scores achieved on these AI benchmark datasets alone can no longer fully measure intelligence.

According to the test-time scaling law, a small model can solve harder problems by spending more tokens "thinking", while a larger model can solve the same problem with fewer tokens. By analogy to a human IQ test, it is as if one person spent more time and used a long sheet of scratch paper to finish the test, while another person used very little time and no scratch paper at all. If both achieve the same number of correct answers, it is more rational to conclude that the person who completed the test faster possesses higher intelligence, likely employing a more efficient method.

That smaller and larger models can answer the same question correctly does not imply they have the same intelligence level; rather, they approached the problem differently. Therefore, without any constraint on the output, whether in terms of the number of tokens or output speed, it is impossible to measure the actual intelligence level of AI solely based on correctness.

From a user's perspective, estimating the time required for AI to complete tasks has become increasingly challenging. Prior to the advent of test-time scaling, models typically used a similar number of tokens to solve each problem. With test-time scaling, models are incentivized to generate more tokens for certain problems, leading to greater variability. Users now have access to metrics such as tokens per second for the APIs they utilize, but lack information about the total number of tokens needed for a given task. As a result, they can no longer reliably estimate the time required for task completion.

Why speed was ignored

If processing speed is so important to intelligence, why was it historically ignored? The reasons are twofold. First, AI was not sufficiently advanced. Early systems could not pass the Turing test¹⁰, solve math or coding problems, or perform in-depth reading or writing. Given that AI was in a primitive stage, the focus was on enabling AI to perform new tasks rather than on the speed of task completion. Second, AI applications were all task-specific. The AI for X-ray image processing¹¹ was entirely different from the AI used for recommendation systems¹². The applicability of AI was measured case by case by application builders, without the need for a unified processing speed metric.

However, with the advent of large language models (LLMs)¹³, AI has become substantially more capable and generalizable, with applications spanning from medical diagnosis¹⁴ to recommendation systems¹⁵. Consequently, it is now both rational and timely to establish a unified metric for AI processing speed.

A mental shift for evaluating AI

Additionally, there has been a significant shift in how application builders interface with models as summarized in Table 1. Initially, developers focused on training application-specific deep learning models, beginning with the advent of AlexNet¹⁶. With the emergence of ChatGPT¹³, the paradigm shifted toward pretraining large language models (LLMs). The public release of Llama¹⁷ enabled post-training and fine-tuning of open models. As the capabilities of open models advanced, such as the Qwen series¹⁸, it became increasingly feasible to serve an open-weight model "as is" without additional fine-tuning. More recently, the introduction of DeepSeek¹⁹, which substantially reduced token costs, has made it more cost-effective to utilize hosted APIs from major AI service providers rather than maintaining proprietary infrastructure.

Table 1. Paradigm shifts of AI usages

Therefore, when evaluating the intelligence of AI, a shift is required from evaluating static models, pure mathematical constructs defined by neural architectures and parameters, to evaluating hosted AI services or APIs. It is necessary to assess their processing speed for applicability to time-sensitive tasks.

Existing metrics

There are established metrics to evaluate the speed of LLMs, such as time to first token (TTFT) and tokens per second (TPS)²⁰. These are suitable metrics for serving language models, but there are two major limitations when considering them as general metrics for AI processing speed.

Unlike the time-limit for an IQ test, these metrics focus on the number of tokens rather than the actual useful information produced by intelligence. For example, a simple program without any intelligence can output random tokens at a very high speed.

Therefore, a new metric is needed that measures the intelligent portion of the output rather than the total volume of tokens produced per unit time.

Intelligence goodput

I propose intelligence goodput as a metric for measuring the processing speed of AI.

Intelligence goodput (IG) is a measurement indicating the maximum amount of intelligent information that an AI service can produce in a given amount of time, formally expressed as:

$$ G=\frac{I}{t} $$

where $G$ is the intelligence goodput, $I$ is the amount of intelligent information, and $t$ is the total time spent.

It is important to note that intelligence goodput primarily measures the output speed of AI, rather than input, for two reasons. First, the main impact of intelligence goodput is in human-AI interaction, which will be discussed further in the next section. Second, AI's processing speed is mainly bounded by output, not input. Thus, it is more meaningful to track AI progress with a metric that measures its bottleneck.

Despite the formal definition of intelligence goodput, certain ambiguities remain in practical measurement. First, "intelligent information" is not an unambiguous term. Defining the amount of intelligence contained in the output of an AI model remains a challenging task. Second, time is also subject to multiple interpretations in computer science.

Regarding time measurement, the two popular choices are CPU time and wall time. Wall time is preferred, as intelligence goodput is primarily relevant for human-AI interactions and applications. End-to-end latency is more informative than technical details such as CPU time.

To address the ambiguity in "intelligent information" in the definition, a method for calculating intelligence goodput are proposed. We divide the score achieved by AI on benchmark datasets by the time used to produce the answer. In this way, we delegate the challenging task of measuring intelligence to the existing benchmarks. It can be formally expressed as follows.

Let $S = \{s_1, s_2, \dots, s_n\}$ be a set of scores from $n$ individual AI benchmarks normalized to the same range, and $W = \{w_1, w_2, \dots, w_n\}$ be a corresponding set of weights, where $w_i$ represents the relative importance of each benchmark. The total time expended across all assessments is denoted by $t$. Intelligence goodput $G$ can be defined as:

$$ G = \frac{\sum_\limits{i=1}^{n} w_i s_i}{\sum_\limits{i=1}^{n} w_i \cdot t} $$

Experiments

Several high-quality models from leading AI providers were evaluated. Source data, including benchmark scores, tokens per second (TPS), and total tokens generated while running the benchmarks, were collected from Artificial Analysis²¹.

We aim to answer three questions: * Do models use very different numbers of tokens to complete the same tasks? * Can users roughly estimate per-task time from TPS? * Does intelligence goodput (IG) provide a better estimate of per-task time? To answer these questions, we conduct three measurements.

First, we compare the total tokens used by different models on the same benchmark suite. The results are shown in Figure 1. The bar length represents the number of tokens generated to complete the tasks by each model.

Figure 1: Model output tokens running benchmarks.

Token usage varies substantially across models. The highest (Grok 4) uses nearly three times the tokens of the lowest (Claude 4.5 Sonnet). Therefore, models can require very different token budgets for the same workload.

Second, we compare model speed measured in TPS with the actual wall-clock time to complete the same benchmark suite. The results are shown in Figure 2, where each point is a model; the x-axis is TPS, and the y-axis is days to finish all tasks.

Figure 2: Total time running all benchmarks vs. TPS.

From Figure 2, we observe weak correlation between TPS and total time. Therefore, TPS alone is an unreliable predictor of per-task or per-suite wall-clock time.

Third, we compute intelligence goodput for each model and analyze the results. The experimental results are shown in Table 2, which reports the intelligence index (I; normalized and averaged benchmark scores), Time (days to run all benchmarks), and intelligence goodput (IG). For details on the methodology for computing the intelligence index, please visit Artificial Analysis²¹.

Table 2. Intelligence Goodput Results

From the results, we make two observations:

• Among models with similar intelligence, IG differentiates by time-to-complete. For example, Grok 4 Fast and Gemini 2.5 Pro both have I = 60, but their times differ a lot (2.7d vs 7.5d). Consequently, Grok 4 Fast has much higher IG and completes tasks significantly faster.

• Among models with similar total time, IG reflects differences in correctness. GPT-5 High and Claude 4.5 Sonnet both finish in about 8 days, but GPT-5 High has a higher I (68 vs 63), meaning it completed more tasks correctly in the same time, which is reflected in their IGs (100.00 vs 91.75).

In both cases, no matter it is different time-to-complete for similar correct answers or number of correct answers in a fixed time window, IG gives a better estimation on the per-task time.

Discussions

Benchmark datasets score only the final answer, disregarding intermediate outputs such as the reasoning process. This approach effectively isolates the intelligent portion of the total output, aligning with the definition of intelligence goodput.

For AI application developers working with time-bounded tasks, the tokens-per-second metric provides limited insight, as the number of tokens required to complete a task remains unknown beforehand. In contrast, intelligence goodput offers a more informative measure of how much useful work can be accomplished per unit time.

Furthermore, incorporating intelligence goodput as an optimization target during model training may help mitigate the verbosity problem commonly observed in LLMs. This problem manifests as unnecessarily long chains of thought that repeatedly revisit the same logical steps. Since longer reasoning processes result in slower time-to-answer and consequently lower intelligence goodput, optimizing for this metric naturally incentivizes more concise and efficient reasoning.

Limitations

The proposed intelligence goodput metric has two primary limitations.

First, it is currently limited to text-based outputs. While AI models are increasingly multimodal, benchmark datasets for evaluating the intelligence of outputs in other modalities, such as images, remain underdeveloped. Such outputs are typically assessed for real-world fidelity and artistic merit rather than intelligence.

Second, the metric is computationally expensive to measure. Calculating intelligence goodput requires significant engineering effort to build evaluation infrastructure capable of running APIs through comprehensive benchmark datasets. Additionally, the cost of API token consumption for executing these benchmarks can be substantial.

Third, tokens spent on incorrect answers are wasted. This property may bias intelligence goodput toward higher-accuracy models because lower-accuracy models consume tokens on wrong answers that do not increase the intelligence index, thereby lowering their IG.

However, in some settings this penalty is desirable. For example, in customer service, an incorrect answer may prompt the customer to re-prompt the AI before receiving a correct answer. From this perspective, the tokens and time spent on wrong attempts are part of the true time to a correct answer. Intelligence goodput appropriately includes this cost in the denominator.

Conclusions

This article introduces the concept of intelligence goodput a metric that measures the processing speed of AI services by quantifying the rate at which they produce intelligent information. The article advocates for a new evaluation paradigm that focuses on dynamic, served AI systems, emphasizing the crucial interplay between hardware, software, and models.

Ultimately, the most important takeaway is that when evaluating AI for performing tasks traditionally performed by humans, the evaluation criteria should mirror those used for humans. The distinction between assessing human and AI performance for a given task will become increasingly blurred. Innovative approaches are required to integrate processing speed into comprehensive intelligence assessments for AI.

References

1. Woodcock, R. W., & others. (1989). Woodcock-Johnson tests of cognitive ability. DLM Teaching Resources. ↩
2. Fry, A. F., & others. (1996). Processing speed, working memory, and fluid intelligence: Evidence for a developmental cascade. Psychological Science. ↩
3. Lichtenberger, E. O., & others. (2012). Essentials of WAIS-IV assessment. John Wiley & Sons. ↩
4. Flanagan, D. P., & others. (2017). Essentials of WISC-V assessment. John Wiley & Sons. ↩
5. Hendrycks, D., & others. (2021). Measuring Massive Multitask Language Understanding. ICLR. ↩
6. Rein, D., & others. (2024). GPQA: A Graduate-Level Google-Proof Q&A Benchmark. COLM. ↩
7. Hendrycks, D., & others. (2021). Measuring Mathematical Problem Solving With the MATH Dataset. NeurIPS. ↩
8. Silver, D., & others. (2016). Mastering the game of Go with deep neural networks and tree search. Nature. ↩
9. Snell, C., & others. (2025). Scaling LLM test-time compute optimally can be more effective than scaling model parameters. ICLR. ↩
10. Turing, A. M. (2009). Computing machinery and intelligence. Springer. ↩
11. Çallı, E., & others. (2021). Deep learning for chest X-ray analysis: A survey. Medical Image Analysis. ↩
12. He, X., & others. (2017). Neural collaborative filtering. The Web Conference. ↩
13. Brown, T., & others. (2020). Language models are few-shot learners. NeurIPS. ↩ ↩
14. Ullah, E., & others. (2024). Challenges and barriers of using large language models (LLM) such as ChatGPT for diagnostic medicine with a focus on digital pathology–a recent scoping review. Diagnostic Pathology. ↩
15. Wu, X., & others. (2024). Could small language models serve as recommenders? Towards data-centric cold-start recommendation. The Web Conference. ↩
16. Krizhevsky, A., & others. (2012). Imagenet classification with deep convolutional neural networks. NeurIPS. ↩
17. Touvron, H., & others. (2023). Llama: Open and efficient foundation language models. arXiv Preprint arXiv:2302.13971. ↩
18. Bai, J., & others. (2023). Qwen technical report. arXiv Preprint arXiv:2309.16609. ↩
19. Guo, D., & others. (2025). DeepSeek-R1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv Preprint arXiv:2501.12948. ↩
20. Zhong, Y., & others. (2024). DistServe: disaggregating prefill and decoding for goodput-optimized large language model serving. OSDI. ↩
21. Artificial Analysis, Inc. (2024). AI Model & API Providers Analysis. ↗ ↩ ↩

With the emergence of multimodal AI, such as image generation^1,2,3 and video generation^4,5,6, it’s reasonable to expect that both human-AI interaction⁷ and AI-AI interaction⁸ will become multimodal. Images may even appear in a model’s chain of thought, approximating aspects of visual reasoning in the human brain.

Analogy to the internet

This invites a comparison between AI's trajectory and the internet's. Viewed through the lens of network bandwidth, increases in throughput enabled richer content formats⁹, evolving from text-forward sites like early Twitter¹⁰ to video-centric platforms like YouTube¹¹. This transformation fundamentally changed how people use the internet.

Nielsen's law¹² states that internet bandwidth increases by 50% annually. This predictive framework enabled anticipation of when the internet would become fast enough for users to easily download images or stream videos, thereby forecasting the emergence and popularization of platforms such as Pinterest¹³ and YouTube¹¹.

If bandwidth served as a progress proxy for the web, can we identify an analogous metric to track the rate of AI output? Moreover, if we can observe a growth pattern in such a metric, we may be able to predict the timeline for new modalities of human-AI interaction in the future.

Intelligence bandwidth

I propose intelligence bandwidth as a metric to measure the output rate of AI services. It quantifies the number of bits in the raw outputs of an AI model, with the unit expressed as KB/s.

Intelligence bandwidth is formally expressed as:

$$ B=\frac{D}{t} $$

where $B$ is the intelligence bandwidth, $D$ is the total size of output data from AI measured in KB, and $t$ is the total time spent.

The measurement methodology is straightforward. For any generative AI service, whether it produces text, images, or videos, the output can be saved to disk. Intelligence bandwidth is then calculated by dividing the storage size of the output by the time the service took to generate that content.

Measuring models across modalities

With intelligence bandwith, we can project all the GenAI models to the same axis, KB/s. All historically significant generative AI models are measured and plotted in a single figure. The models covered include large language models, image generators, and video generators. Most of the data presented in this section is collected from Artificial Analysis¹⁴.

The experimental results are shown in Figure 1, where the X-axis is the release date of the models, and the Y-axis is the intelligence bandwidth of the models measured in kilobytes per second. The modality of the models is indicated by different colors as shown in the legend.

Figure 1: Intelligence bandwidth (KB/s) over time.

Key observations from the experimental results include: 1. Most language models are between 0 KB/s and 3 KB/s.

1. Image generators exhibit an exponential growth rate.

2. The video generator, Veo3⁵, currently exhibits an even lower intelligence bandwidth than the state-of-the-art image generators. This is primarily attributable to less mature serving technologies for video models compared to those for large language models and image generators. As serving efficiency for video generators improves, substantial growth in their intelligence bandwidth is anticipated in the near future.

3. The Gemini 2.5 Flash¹⁵ image generator is an outlier, primarily because it is optimized for low latency and usability rather than best quality and fidelity.

Jin's law

Effective metrics can help reveal new laws to predict the future, analogous to the number of transistors and Moore's law¹⁶, FLOPS and Huang's Law, or the internet bandwidth and Nielsen's law¹². Similarly, a robust metric for AI output rate can facilitate the discovery of macro-level trends in AI development.

We now assess the validity of intelligence bandwidth as such a metric by examining whether it supports the formulation of a predictive law for future AI growth.

The dotted curve in Figure 1 represents the estimated growth of intelligence bandwidth. The prediction of the growth rate is based primarily on Imagen 4, the state-of-the-art high-quality image generator, rather than a model balanced between speed and quality, such as Gemini 2.5 Flash¹⁵. The growth rate of intelligence bandwidth is summarized in a simple law named after the author's surname, presented as follows.

Jin's law: The intelligence bandwidth (KB/s) of the best generative AI service available to the public doubles every year.

The formal definition of this law is as follows. Let $B(t)$ be the intelligence bandwidth of the best hosted AI model available to the public at time $t$ and $B_0$ ​be the initial intelligence bandwidth at time $t_0$. The law can be expressed as:

$$ B(t)=B_0 \cdot 2^{(t−t_0 ​)/T} $$

where $T$ is the doubling period. The best estimate is $T=1$, the unit of which is year.

In Jin's law, intelligence bandwidth is defined by the modality exhibiting the highest KB/s measurement. Currently, image generators are at the forefront of this growth. As advancements in models and serving technologies for image generation reach a plateau, it is anticipated that video generators will become the primary drivers of further increases in intelligence bandwidth.

We only measure the best GenAI services. Any model or service that is not well-received by the public is not considered by Jin's law.

We only measure the public available GenAI services. Any service that is not public accessible is not considered by Jin's law.

Impact on human-AI interaction

To use Jin's law to predict the future of human-AI interaction, we first need to recap a few basics on this subject.

Current human-AI interactions are still predominantly text-based. This is possible because the output speed of AI has exceeded human reading speed, which is approximately 238 words per minute¹⁷. This is enabled by state-of-the-art serving technologies, which can generate 14,000 words per minute. Similarly, speech generation speed is far beyond human listening speed¹⁸.

To clarify the prerequisites for enabling real-time human-AI interaction within a given modality, it is helpful to examine the specific conditions that must be satisfied. For self-paced media formats such as text and images, individuals typically consume content at their maximum perceptual speed. Once the intelligence bandwidth exceeds the human perceptual threshold, real-time interaction in that modality becomes feasible. In contrast, for fixed-speed media formats such as audio and video, users generally adhere to the inherent playback speed. Thus, as long as the AI generation speed surpasses the fixed playback rate of these media formats, real-time interaction in those modalities is achievable.

Research in multi-modal AI continues to address the bottlenecks of other modalities. With the increase of the intelligence bandwidth of image generation, visual illustrations will become integrated into AI responses. AI may also be able to perform visual reasoning akin to humans on a whiteboard, and iteratively refine graphical designs as designers do on paper. As the intelligence bandwidth of leading AI models continues to increase, it is expected that video illustrations and real-time generated environment interactions, such as those demonstrated by Genie 3¹⁹, will become feasible. In a speculative future, AI could generate entire worlds that users can interact with in real time, an idealistic scenario enabled by the cognitive capabilities of AI.

Achieving such increases in intelligence bandwidth requires not only advances in AI models as static collections of neural architectures and parameters, but also significant improvements in AI hardware and machine learning systems. In this vision, hardware, software, and models are no longer orthogonal, but are deeply integrated to enable superintelligence.

Predictions

Based on Jin's law, two predictions about human-AI interaction in the near future are made as an example of how to use the law to predict the exponential growth of AI in the future:

1. Images will soon be used in AI interactions. The latency of the Gemini 2.5 Flash¹⁵ image generator is lower than that of large language model responses. Consequently, large language models may soon incorporate images to provide enhanced illustrations in their outputs. Currently, it takes only 4.6 seconds to generate an image. If this speed doubles within a year, it is likely that applications will emerge in which images become a primary mode of interaction and illustration.

2. Real-time video interaction will be widely available in three years. Given the intelligence bandwidth of models in 2025, it is currently possible to generate 8 seconds of video in 50 to 60 seconds. If generation speed increases by a factor of 7 to 8, real-time video generation will become feasible. Achieving an 8-fold increase corresponds to approximately $\log_2(8) = 3$ years. While some uncertainty remains in this prediction, video generators are presently below the projected growth curve and possess significant potential for accelerated improvement as serving technologies advance.

There are many other implications that can be derived from Jin's law. It is hoped that this law will guide AI application developers in identifying optimal time windows to bring products to market, and policymakers in enacting regulations at appropriate times to maximize development while minimizing harm.

Limitations

There are two main limitations to this work.

First, the accuracy of the estimated doubling period $T$ in Jin's law is uncertain. As the field of AI is still in its early stages, the current estimation is based on a limited number of available data points.

Second, the exponential growth described by Jin's law represents an idealized scenario. Just like the idea of "the end of Moore's Law", Jin's law would not live for ever. In practice, growth may be constrained by factors such as energy supply limitations or economic pressures, particularly if the AI sector experiences a market correction.

Conclusions

The article introduces the concept of intelligence bandwidth, a metric that measures the processing speed of AI services. The article observes an exponential growth trend in this metric across various modalities and formulates Jin's law, which states that the intelligence bandwidth of the best publicly available AI doubles annually. This law provides a predictive framework, forecasting the imminent integration of real-time images into human-AI text interactions and the widespread availability of real-time video generation within three years.

Further reading

Intelligence bandwidth measures only the total size of AI output, whereas the proportion of intelligent information contained within that output can vary significantly. To address this limitation, it is also important to measure the rate at which an AI service produces genuinely intelligent information. I propose another metric to evaluate this aspect, named intelligence goodput, which is discussed in detail in this companion article.

References

1. Kingma, D. P., & others. (2014). Auto-encoding variational Bayes. ICLR. ↩
2. Goodfellow, I. J., & others. (2014). Generative adversarial nets. NeurIPS. ↩
3. Rombach, R., & others. (2022). High-resolution image synthesis with latent diffusion models. CVPR. ↩
4. Oord, Aäron van den, & others. (2024). State-of-the-art video and image generation with Veo 2 and Imagen 3. ↗ ↩
5. Hume, T., & others. (2025). Meet Flow: AI-powered filmmaking with Veo 3. ↗ ↩ ↩
6. OpenAI. (2024). Sora: System Card. ↗ ↩
7. Amershi, S., & others. (2019). Guidelines for human-AI interaction. CHI. ↩
8. Surapaneni, R., & others. (2025). Announcing the Agent2Agent Protocol (A2A). ↗ ↩
9. Coffman, K. G., & others. (2002). Internet growth: Is there a “Moore’s Law” for data traffic? Handbook of Massive Data Sets. ↩
10. Murthy, D. (2018). Twitter. Polity Press Cambridge. ↩
11. Snickars, P., & Vonderau, P. (2009). The YouTube reader. Kungliga biblioteket. ↩ ↩
12. Nielsen, J. (1998). Nielsen’s law of internet bandwidth. ↗ ↩ ↩
13. Gilbert, E., & others. (2013). “ I need to try this”? A statistical overview of Pinterest. CHI. ↩
14. Artificial Analysis, Inc. (2024). AI Model & API Providers Analysis. ↗ ↩
15. Team, G. (2025). Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv Preprint arXiv:2507.06261. ↩ ↩ ↩
16. Moore, G. E., & others. (1965). Cramming more components onto integrated circuits. ↩
17. Brysbaert, M. (2019). How many words do we read per minute? A review and meta-analysis of reading rate. Journal of Memory and Language. ↩
18. Kuperman, V., & others. (2021). A lingering question addressed: Reading rate and most efficient listening rate are highly similar. Journal of Experimental Psychology: Human Perception and Performance. ↩
19. Parker-Holder, J., & others. (2025). Genie 3: A new frontier for world models. ↗ ↩

It only took about 800 lines of Python and C++ code altogether to build the framework. The following is a basic example of what it is capable of, but you are also welcome to check out the end-to-end classification example.

import framework

x = framework.Tensor.from_numpy(np.array([[2.0, 3.0]], dtype=np.float32))
y = framework.Tensor.from_numpy(np.array([[4.0], [5.0]], dtype=np.float32))
z = framework.ops.matmul(x, y)  # Expected: [[2*4 + 3*5]] = [[23.0]]
loss = framework.ops.sum(z)

loss.backward()  # Compute the gradients of loss w.r.t. tensor x, y, and z.

print("x.grad:", x.grad.numpy())  # Expected: [[4.0, 5.0]]
print("y.grad:", y.grad.numpy())  # Expected: [[2.0], [3.0]]

The codebase is thoroughly commented, which brings it up to around 2000 lines of code. Everyone can easily read the source code to understand how an ML framework works and feel more confident about building a new one. This project is named the "Readable ML Framework" and is available on GitHub.

This article dives into all the core concepts in ML frameworks citing the implementations from the repo. We will also discuss other important elements of ML frameworks beyond the project.

Why care about ML frameworks?

ML framework lies in the center of the ML infra. Understanding ML framework gives you the bigger picture by connecting all the dots, including CUDA/Triton kernels, GPU/TPU, data pipelining, ML compilers, distributed training and so on.

How did I come up with this idea?

It all started with a question from a friend. He asked me if I knew how to build a project like TensorFlow from scratch. The answer was absolutely no at the time. It planted a seed in my mind that grew my curiosity for ML frameworks through the years. Five years later, I built this project to learn more about ML frameworks myself.

Having spent years working on machine learning infrastructure, many people would think that I knew ML frameworks very well, but I did not. Then, I started to realize that if I did not know much about it, others would not either unless they worked on building an ML framework directly.

The next question I asked myself was how hard could it be to build a basic ML framework. To find out, I read the source code of PyTorch. Here are my notes. I found it was not that hard if we got rid of all the complex features.

So, I decided to build one with almost no features except tensor operations and automatic gradient computation. This "bare-bones" style for educational projects is largely inspired by the super tiny compiler project, which helped me learn what a compiler is.

Python C++ interfacing

Let's start with the number one issue that turns the code readers away from the ML frameworks, Python C++ interfacing. Whenever I wanted to track down a tensor operation (We call it op/ops in the rest of the article for short) from its Python interface to its actual implementation in C++, I got lost. It could take a super complex system to bind all the C++ functions to Python functions. We call this one-to-one mapping between Python functions and C++ functions bindings.

In PyTorch, it indeed developed a system called torchgen to generate the Python C++ bindings. TensorFlow and JAX uses Pybind11, which is a third-party library to simplify the code for Python C++ bindings. For simplicity and the readability of the code, we use Pybind11 for our framework.

It works as follows. * Write some C++ functions/classes. * Map the functions/classes to Python objects with Pybind11. Here is an example from python_bind.py of using Pybind11 to present the C++ functions as Python functions.

// "core" is the Python submodule name.
// Its Python import path is framework.core
// "m" is the module object.
PYBIND11_MODULE(core, m) {
    ...

    // Create a Python submodule named ops under framework.core.
    auto ops_module = m.def_submodule("ops");

    // Bind C++ function ops::matmul to
    // Python function framework.core.ops.matmul().
    ops_module.def(
        "matmul",                  // Function name in Python.
        &ops::matmul,              // Pointer to the C++ function.
        "(m, k), (k, n) -> (m, n)" // Docstring of the function in Python.
    );
    ...
}

• Define the Python module as a C++ extension in setup.py.

# Define the extension module
ext_modules = [
    Pybind11Extension(
        # The Python import path to the C++ module.
        "framework.core",
        # Pass in all the C++ files needed to compile the module. We put all
        # the .cpp and .h files under framework/core dir.
        glob.glob("framework/core/**/*.cpp", recursive=True),
        # Tell setuptools where to find the pybind11 header files while
        # compiling the C++ extension module.
        include_dirs=[pybind11.get_include()],
    ),
]

# Setup function
setuptools.setup(
    ...
    # Add the C++ extension module to the package.
    ext_modules=ext_modules,
    ...
)

That is it! With the above, we can call the C++ function as a Python function from Python code.

When we build the project, for example, with pip install, the setuptools will compile the C++ code into a .so file in Linux, which is binary file. It is also known as a shared object or shared library. When we call the corresponding Python function in Python, it will automatically convert the arguments of Python types to the parameters of C++ types, and call the C++ function.

In our framework, it is easy to track down an op. Just look at python_bind.py and see which C++ function it is bound to.

Why mix Python with C++?

The simple answer is we want the speed of C++ and the simplicity of Python programming. There are parts of an ML framework that have to be Python and parts that have to be C++/CUDA.

The user-facing APIs have to be in Python. Python is the programming language of machine learning and AI. All the research scientists and machine learning engineers love Python. It has an amazing ecosystem of machine learning libraries. For the optimal user experience, the frameworks have to have Python APIs.

The underlying data structure of tensors and the tensor ops have to be implemented in C++/CUDA for speed performance. Python is slow, especially with its Global Interpreter Lock (GIL) to prevent parallel programming. So we need C++/CUDA to write efficient data structures for tensors and parallel tensor ops.

Note there are also ongoing efforts to include the merits of both languages in one programming language, like the 🔥 Mojo programming language.

Tensor

Tensor is the most essential data structure used in an ML framework. To be more efficient and easier to explicitly manage the memory allocation, they are usually implemented in C++/CUDA.

In modern ML frameworks, tensors are split into the data and the view. The data contains the actual values in the tensor occupying a large amount of memory. The view is how these data are presented to the ops or users.

The data can be a contiguous chunk of memory of floats or integers. The view can contain the shape of the tensor, the order of the axis, the step size to iterate the data, and so on.

This data-view split allows many ops to be super efficient. For example, if you want to transpose a tensor or matrix, you only need to create a new view indexing the values differently. No copy or modification to the data is needed.

For simplicity, we chose a simple split strategy to implement the Tensor class in our framework. The data part is always a flattened C++ vector. The view only contains the shape of the tensor. Data and shape are the two members of the Tensor class.

For example, the following 2x2 matrix can be represented as shape=[2,2] and data = [1,2,3,4]. $$ \begin{bmatrix} 1 & 2 \\ 3 & 4 \end{bmatrix} $$

Here is our Tensor class.

class Tensor {
public:
    /**
     * @brief The shape of the tensor, represented as a vector of size_t.
     *
     * Each element in the vector corresponds to the size of a dimension.
     * For example, a tensor with shape {2, 3} would be a 2x3 matrix.
     */
    std::vector<size_t> shape;

    /**
     * @brief The underlying data of the tensor, stored as a contiguous vector
     * of floats.
     *
     * The elements are typically ordered according to a row-major (C-style)
     * layout, though this class itself doesn't enforce a specific layout.
     */
    std::vector<float> data;

    ...
}

Read tensor.h and tensor.cpp for more details.

Tensor operations

Implementing the tensor operations is the most challenging part of building our ML framework. We had to make sure they were parallelized correctly. Let's take a look at matrix multiplication, one of the most important ops we implemented, and how backward functions work in general.

Matrix multiplication

Matrix multiplication is important because it is so widely used in neural networks. Optimizing it has a high ROI. The technique we used here is called tiling.

Let's say we have two tensors (matrices) x and y. Also, we have allocated memory for the output tensor. Then, output[i,j] is the inner product of the ith row of x and the jth column of y.

Tiling is a technique that chops large matrices into small blocks that can be processed in parallel. You split x into row blocks and y into column blocks, and hand each row-block/column-block pair to its own task. A task assigned x[a:b,:] and y[:,c:d] produces output[a:b,c:d]. If the shared dimension is also long, you can chop it too, computing partial inner products on shorter slices and summing them up afterward. The sub-blocks produced by all the tasks are then combined (tiled) into the final output matrix.

There is even fully optimized hardware just for matrix multiplication, such as the Tensor Cores in NVIDIA GPUs and the MXUs in Google's TPUs. To use them, you chop the matrix, feed the blocks to the ASIC, and tile the output back together.

Backward functions

A quick recap of forward and backward functions. The forward function is the function of the op itself, for example, softmax, matmul, and relu. It computes the output of the op given the input to the op. The backward function is to compute the gradients w.r.t. the input given the gradients w.r.t. the output. The backward function may also access the input and output of the forward function.

The backward function is useful for the neural network optimizers during backward propagation to compute the gradients of loss w.r.t all the weights of the neural network.

The following figure shows how gradients are computed. In the forward function, we compute z using x and y. In the backward function, we compute the x.grad, y.grad using x, y, and z.grad, where x.grad, y.grad, z.grad are the gradients of loss w.r.t. x, y, z.

To put it in math, if the forward function is $\mathbf{z}=f(\mathbf{x}, \mathbf{y})$. The backward function is to compute $\partial L / \partial \mathbf{x}$ and $\partial L / \partial \mathbf{y}$ where $L$ is the loss. We can compute them using the chain rule as follows.

$$ \frac{\partial L}{\partial \mathbf{x}} = \frac{\partial L}{\partial \mathbf{z}} \cdot \frac{\partial \mathbf{z}}{\partial \mathbf{x}} $$ $$ \frac{\partial L}{\partial \mathbf{y}} = \frac{\partial L}{\partial \mathbf{z}} \cdot \frac{\partial \mathbf{z}}{\partial \mathbf{y}} $$

In practice, $\mathbf{x}$, $\mathbf{y}$, $\mathbf{z}$, and $\frac{\partial L}{\partial \mathbf{z}}$ are known at the time we call the backward function. We need to compute $\frac{\partial L}{\partial \mathbf{x}}$ using them. As a straightforward way to compute it, we only need to compute $\frac{\partial \mathbf{z}}{\partial \mathbf{x}}$, and $\frac{\partial \mathbf{z}}{\partial \mathbf{y}}$, which are computed using $\mathbf{x}$ and $\mathbf{y}$.

For a simplified example, given $z = f(x, y) = x \times y$, $\partial L / \partial z = 3$, $x=4$, $y=5$, $z=20$ are scalars, we have $$ \frac{\partial z}{\partial x} = y, \\ \frac{\partial z}{\partial y} = x. $$ So, we can compute the gradients as follows. $$ \frac{\partial L}{\partial x} = \frac{\partial L}{\partial z} \cdot \frac{\partial z}{\partial x} = \frac{\partial L}{\partial z} \cdot y = 3\cdot 5 = 15 $$ $$ \frac{\partial L}{\partial y} = \frac{\partial L}{\partial z} \cdot \frac{\partial z}{\partial y} = \frac{\partial L}{\partial z} \cdot x = 3\cdot 4 = 12 $$

Note that in the example above, the output of the forward function, $z$, is not directly involved. However, we may still need the output in some cases as a trick for performance (running speed) optimization, which is where recomputing comes in.

Recomputing

Some backward functions need the forward function's output, not just its input, to compute gradients. That leaves two ways to implement them: store the output when the forward function runs and reuse it in the backward function, or discard it and recompute it from the input when the backward function needs it.

Why recompute the output if we already have it saved somewhere? The memory it's saved in could be slow to load from. Recomputing the output and writing it to a local fast-but-small memory can sometimes even be faster than loading it from a big-and-slow memory.

The same trade-off shows up at a much larger scale in training deep neural networks. Normally, every layer's outputs, called activations, are kept in memory during the forward pass so the backward pass can load them instead of recomputing them. If we feel recomputing is actually faster for some of the layers, we do not save those tensors during the forward pass, and recompute them on demand during the backward pass instead. This technique of only keeping a subset of activations and recomputing the rest on demand is known as activation checkpointing. It substantially cuts peak memory usage, helping fit large model training workloads into limited memory.

Please refer to ops.h and ops.cpp for the ops implementations.

Automatic differentiation

Automatic differentiation, also known as autodiff or autograd, is a mechanism to compute the gradients of the loss w.r.t. the weights of the neural network automatically. The user only defines the forward pass. The framework records the ops performed on each tensor and computes the gradients. This is a core feature of an ML framework.

Before we dive into the implementation, There are a few things we need to clarify.

Compute graph

The following figure shows the compute graphs of the forward pass and the backward pass of a super simple neural network, whose inputs are a, b, d, and the output is e.

In the graph, each node is a tensor, and each edge or two edges pointing to the same output tensor is an op.

To put the compute shown in the figure in code, it should look like the following.

def forward(a, b, d):
    c = ops.add(a, b)
    e = ops.multiply(c, d)
    return e

# Run if we call e.backward().
def backward(a, b, c, d, e):
    c.grad, d.grad = ops.multiply_backward(output_grad=e.grad, inputs=(c, d))
    a.grad, b.grad = ops.add_backward(output_grad=c.grad, inputs=(a, b))

In the forward pass, we compute the output e from the inputs. In the backward pass, we compute the gradients of e w.r.t. each tensor (i.e. a.grad, b.grad, c.grad, d.grad).

The user defines the forward pass of the neural network. We need to automatically figure out the backward pass when the user calls e.backward(), or, in most cases, loss.backward().

Tracing

Tracing refers to the process of extracting the compute graph from the user-defined forward pass. It is called tracing because it is the process of tracing the paths of tensors flowing in the compute graph.

There are many different ways to do this, but we implemented the simplest one. We just record the op history on every tensor, namely OpRecord. We use this information to reconstruct the compute graph from the very last output tensor, which, in most cases, is the loss.

What information do we store in an OpRecord? In the example above, we just record c and d in e.op_record, and record a and b in c.op_record. So, we record the input tensors of the op in the OpRecord of the output tensor. This information should be good enough for us to reconstruct the entire compute graph by starting a breadth-first search (BFS) from the last output tensor.

Backward propagation

During backward propagation, we want to compute the gradients w.r.t. all the tensors. So, every time, we just call the backward function of the op with the gradients and the input tensors of the op. So, we need one more piece of information in the OpRecord, the backward function of the op to call.

Following is our implementation of the OpRecord.

class OpRecord:
    def __init__(self, func_backward, input_tensors, output_tensor):
        self.func_backward = func_backward
        self.input_tensors = input_tensors
        self.output_tensor = output_tensor

So, it only has 3 attributes:

• func_backward. The backward function to call. We need it because we need to know which function to call during backward propagation.
• input_tensors. The input tensors to the forward operation. We need it to trace the compute graph and for the backward function call.
• output_tensor. The output tensor of the operation. We need it for the backward function call because we need the gradients w.r.t. the output tensor, i.e., output_tensor.grad.

Here are the steps to do it.

• During the forward pass, assign an OpRecord object to the .op_record attribute of the Tensor when we produce the Tensor with an op.
• When loss.backward() is called, use the OpRecords to trace back the entire compute graph.
• Sort the nodes in the compute graph in reverse topological order.
• Call the backward function stored in the .op_record of every tensor following the sorted order, which starts from the loss node.

Please refer to op_record.py, forward.py, backward.py, autograd.py for more details, for example, why we need topological sorting.

End-to-end workflow

So far, we have reviewed all the major components of the framework. Let's connect the dots by describing an end-to-end workflow.

Let's look at the code example at the beginning of this article.

import framework

x = framework.Tensor.from_numpy(np.array([[2.0, 3.0]], dtype=np.float32))
y = framework.Tensor.from_numpy(np.array([[4.0], [5.0]], dtype=np.float32))
z = framework.ops.matmul(x, y)  # Expected: [[2*4 + 3*5]] = [[23.0]]
loss = framework.ops.sum(z)

loss.backward()  # Compute the gradients of loss w.r.t. tensor x, y, and z.

print("x.grad:", x.grad.numpy())  # Expected: [[4.0, 5.0]]
print("y.grad:", y.grad.numpy())  # Expected: [[2.0], [3.0]]

First, we constructed two tensors x and y from NumPy arrays. The framework.Tensor is a Python class, it uses the C++ implementation of the tensor, framework.core.Tensor, as its attribute. The NumPy array will be passed to the constructor of the C++ class. The argument type conversion from Python types to C++ types is handled by Pybind11.

Then, we call two ops matmul and sum on them to get the single value tensor, the loss. This defines the forward pass of the model. framework.ops.matmul and framework.ops.sum are just Python wrappers of the C++ implementation, framework.core.ops. Before calling the C++ code, the Python wrapper stores x, y, and matmul_backward() in z.op_record. It also does something similar to loss.op_record.

Finally, when loss.backward() is called, the autograd module trace and build the compute graph using the op records, and compute the gradients calling the backward functions of the ops.

So far, we have gone through all the features we have implemented in our ML framework. The main challenges we faced were the Python C++ interfacing and ops implementations.

If you follow everything we introduced, you can explore the heavily commented codebase and try to understand all the code by yourself. You will feel confident to implement a framework by yourself. I encourage you to do so to better understand all the details.

Other features

In the rest of the article, we will explore the important features that are not included in our framework. We will explain what they are, why they are important, and how they can be added to our framework.

GPU backend

A GPU backend of an ML framework should allow you to run your models on GPU for faster speed. The implementation requires adding a new C++ extension module to the Python library. This new module should implement a new Tensor class and all the ops using CUDA, a programming language similar to C++ but for GPUs specifically. The Python code of the framework should dynamically route to different backends based on the user configuration.

It is not hard to implement the GPU backend, but extremely challenging if you want it to run fast. There are a lot of tricks in CUDA kernel implementation. A CUDA kernel is a function that is executed on a GPU making use of its parallel processing capabilities.

More precisions

In our framework, we only support float32 as the data type for all the tensors. However, people may want to use a different precision like bfloat16, int8, and so on. This allows the user to control the compute precision on a finer-grained level to better address their speed and power requirements.

The implementation of this feature is mainly about abstraction. To reduce code duplication, we need to make code for the ops general enough for more data types and also leave space for custom optimization for certain data types.

More tensor types

In many applications, people may find sparse tensors and ragged tensors useful. For example, sparse tensors are widely used in the embedding lookup tables of recommendation systems. Ragged tensors, which means a dimension of the model may have different lengths, for example, the length of the input sentences to a natural language processing model.

To implement these, we need to make the Tensor class of the framework extensible, and make the ops more supportive for different types of tensors. We may even need to introduce new hardware feature support to better process sparse tensors, for example, making use of the sparse core of Google TPUs.

Data pipelining

If not careful, the data pipeline can become the performance bottleneck of the overall ML process. The data pipelining module needs to load the data in advance to feed the model.

The implementation is rather orthogonal to what we have introduced above. It should prefetch the data on the compute device, for example, the HBM of GPUs before the model actually requests that data for computation.

ML compiler

ML compiler is a big topic. We will dig more into the ML compilers in a future blog post. The ML compiler is different from a traditional programming language compiler, it does not compile any programming language, but optimizes the compute graph and lowers it down to execute on a GPU/TPU.

Here is an example of an optimization that an ML compiler can do. We use softmax to convert logits $\mathbf{z}$ into probability $\mathbf{p}$ and use categorical cross-entropy as the loss function to get the loss value $L$, shown as follows.

The softmax function: $$ \mathbf{p} = \text{softmax}(\mathbf{z}) $$ $$ p_i = \text{softmax}(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} $$

The categorical cross-entropy loss: $$ L = - \sum_{i=1}^{K} y_i \log(p_i) $$

where $\mathbf{y} = (y_1, y_2, \cdots, y_K)$ is the one-hot encoded label.

It would be time-consuming to call the backward function of all the ops involved (exp, log, divide, multiply, sum) one by one to get the gradients of the loss $L$ w.r.t. the logits $\mathbf{z}$ denoted as $\partial L/\partial \mathbf{z}$.

However, if you derive it by hand, you can find that $\partial L/\partial \mathbf{z}$ can be simplified to $\mathbf{p} - \mathbf{y}$.

$$ \frac{\partial L}{\partial \mathbf{z}} = \mathbf{p} - \mathbf{y} $$

The ML compiler should have a built-in rule to capture this pattern of a long chain of backward function in the compute graph and optimize it to $\mathbf{p}-\mathbf{y}$.

Regarding the implementation of an ML compiler, it takes tremendous effort to build a good one. It is far beyond what a solo project can do. It should capture the compute graph through tracing, and serialize the compute graph to some intermediate representations (IR). The compiler takes the IR as input and optimizes it to output another IR. After multiple levels of such optimization, the IR should be lowered down to the hardware to execute the computation.

Distributed training

Finally, here is the killer feature of an ML framework. It is almost impossible to do distributed training without an ML framework. It distributes the compute and memory usage across multiple machines and multiple GPUs. It allows the users to have an experience similar to running the model on a single machine.

There are many different distribution techniques in machine learning. We will also have a future post dedicated to this topic.

Conclusions

So far, we have gone through the basics of building an ML framework from scratch including Python and C++ interfacing, tensor data structures, tensor operations, and automatic differentiation.

As you can see, it is not difficult to build a simple working ML framework. The challenge is to make it efficient enough for all the different scenarios. It may include running on different hardware, on different numbers of machines, with different data types and tensor types. To support these scenarios, it may require a good ML compiler, distribution features, optimized op implementations, and so on.

After reading this article, hope you are more confident about your knowledge on ML frameworks now.

This article analyzes Meta's GenAI and large model strategy to understand the considerations of open-sourcing their large models. We also discuss how this wave of open-source models is similar to and different from traditional open-source software.

DISCLAIMER: Whether the Llama models are genuinely open-source falls outside the scope of this article. All information is from public sources.

The illusion of proprietary models

If Meta made its models open-source, how could they make money? People would just build their own services instead of paying for the service provided by Meta, for example, the chatbot on meta.ai, an API based on Llama, or services for fine-tuning and serving the model.

However, you can never prevent people from building their own solutions. Keeping your models proprietary would not help you achieve this goal. Regardless of whether you make your models open-source, others, like Mistral AI, Alibaba, and even Google, made their models open-source.

The users are building their models based on the best open-source models. Therefore, it doesn't matter whether your model is open-source or not unless it is always better than the best open-source models in the long run. I don't think many companies are so confident about that.

So, now, you only have two choices: to be the first and the leader of open-source models, or to be a follower by releasing your models later.

Why be the leader of open-source models?

Being the leader of open-source models has many benefits, but the most important is attracting talent.

The war of GenAI is a talent competition bottlenecked by computing power. How much computing power you get largely depends on two things: your cash flow, and your relationship with Nvidia. However, how many talents you have is another story.

According to Elon Musk¹, Google had two-thirds of the AI talent pool in the early 2010s. To counter Google's power, they founded OpenAI. Then, some of the best people left OpenAI and founded Anthropic to focus on AI safety. So, these three companies have the best and the most AI experts in the job market. Everyone else is super hungry for more AI experts.

Being the leader of open-source models would help Meta bridge this gap of AI experts. Open-source models attract talent in two different ways.

First, the AI experts want to work for Meta. It is super cool to have the whole world use the model you built. It gives you so much exposure for your work, amplifies your professional impact, and benefits your future career. So, many talented people would like to work for them.

Second, the AI experts in the community do the work for Meta for free. Right after the release of Llama, people started to experiment with it. They help you develop new serving technologies to reduce costs, fine-tune your models to discover new applications and scrutinize your model to discover vulnerabilities to make it safer. For example, according to this article, they did instruction tuning, quantization, quality improvements, human evals, multimodality, and RLHF for Llama within a month after its initial release. Delegating this work to the community saves Meta huge amounts of computing and human resources.

Iterate fast with the community.

With open-source models, Meta can iterate quickly with the community by directly incorporating their newly developed methods.

How much would it cost Google to adopt a new method from the community? The process consists of two phases: implementation and evaluation. First, they need to reimplement the method for Gemini. This involves rewriting the code in JAX, which requires a fair amount of engineering resources. During the evaluation, they need to run a list of benchmarks on it, which requires a lot of computing power. Most importantly, it takes time. It stopped them from iterating on the latest technologies when they were first available.

Conversely, if Meta wants to adopt a new method from the community, it will cost them nothing. The community has done the experiments and benchmarks on the Llama model directly, so not much further evaluation is needed. The code is written in PyTorch. They can just copy and paste it into their system.

Llama built a flywheel between Meta and the community. Meta brings in the latest technology from the community and rolls out its next-generation model to the community. PyTorch is the common language they speak.

Can they still make money?

The model is open-source. Wouldn't people just build their own service? Why would they want to pay Meta for a service built on an open-source model? Of course, they will. The service is difficult to build even with an open-source model.

How do you fine-tune and align the model to your specific application? How do you balance between the service cost and the model quality? Are you aware of all the tricks to fully utilize your GPUs?

The people who know the answers to these questions are expensive to hire. Even with enough people, it is hard to get the computing power to fine-tune and serve the model. Imagine how hard it is to build Meta AI from the open-source Llama model. I would expect hundreds of employees and GPUs to be involved.

So, it is likely that people will still pay for Meta's GenAI service if they have any in the future.

It's just like open-source software, but not quite.

The situation is very similar to traditional open-source software. The “free code paid service” framework still applies. The code or the model is free to attract more users to the ecosystem. With a larger ecosystem, the owner collects more benefits. The service built upon the free code is for profit.

However, it is also NOT like open-source software. The main difference can be summarized as low user retention and a new type of ecosystem.

Low user retention

Open-source models have lower user retention. Migrating to a new model is much easier than to new software.

It is hard to migrate software. PyTorch and HuggingFace have established a strong ecosystem for deep learning frameworks and model pools. Imagine how hard it would be to shift their dominance even slightly if you created a new deep learning framework or model pool to compete with them.

A good example is JAX. It has better support for large-scale distributed training, but it is hard to onboard users to JAX because it has a smaller ecosystem and community. It lacks a helpful community to support users with issues. Moreover, the engineering cost of migrating the entire infra to a new framework is too high for most companies.

Open-source models do not have these problems. They are easy to migrate and require almost no user support. Therefore, it is easy for people to shift to the latest and best models. To maintain your leadership in open-source models, you must constantly release new models at the top of the leaderboard. This is also noted as a downside or challenge to be the leader in open-source models.

A new type of ecosystem

Open-source models create a new type of ecosystem. Unlike open-source software, which creates ecosystems of contributors and new software built upon them, open-source models create ecosystems of fine-tuned and quantized models, which can be seen as forks of the original model.

As a result, an open-source foundational model doesn't have to be super good at every specific task because users would fine-tune it for their applications with domain-specific data. The most important feature of a foundational model is to meet the deployment requirements of the users, such as low latency in inferencing or being small enough to fit an end device.

This is why Llama has multiple sizes for each version. For example, Llama-3 has three versions: 8B, 70B, and 400B. They want to ensure they cover all deployment scenarios.

Summary

Even if Meta did not open-source their model, others would. So, it would be wise for Meta to open-source it early and lead the open-source models. Then, Meta can iterate quickly with the community to improve its models and catch up with OpenAI and Google.

When open-sourcing your model, there is no need to worry about people not using your service since there is still a huge gap between the foundational model and a well-built service.

Open-source models are similar to open-source software in that they all follow the “free code paid service” framework but differ in user retention rate and the type of ecosystem they create.

In the future, I would expect to see more open-source models from more companies. Unlike the deep learning frameworks converged on PyTorch, open-source models will remain diverse and competitive for a long time.

References

1. Musk, E. (2023). Interview on Google and OpenAI AI talent. YouTube. ↗ ↩

DISCLAIMER: The views expressed in this article are solely those of the author and do not necessarily reflect the opinions or viewpoints of Google or its affiliates. The entirety of the information presented within this article is sourced exclusively from publicly available materials.

A popular theory

One popular theory I heard is that Google wants to attract more customers to its cloud services. If they sell it to other cloud service providers, they are less competitive in the cloud service market.

According to cloud service customers, this theory does not make much sense. No corporate-level customer wants to be locked to one specific cloud service provider. They want to be flexible enough to move to another whenever needed. Otherwise, if the provider increases the price, they can do nothing about it.

If they are locked to Google Cloud for using TPUs, they would rather not use it. This is why many customers don't want to use TPUs. They only started to feel less locked in recently when OpenXLA, an intermediate software to access TPUs, supported more frameworks like PyTorch.

So, using TPUs to attract customers to Google Cloud is not a valid reason for not selling them. Then, what is the real reason? To answer this question better, we must look into how Google started the TPU project.

Why did Google start the TPU project?

The short answer is for proprietary usage. There was a time when GPUs could not meet the computing requirements for AI hardware.

Let's try to estimate when the TPU project was started. Given it was first announced to the public in 2016¹, it would be a fair guess that it started around 2011. If that is true, they started the project pretty early since we did not see a significant improvement in computer vision until 2012 by AlexNet.

With this timeline, we know GPUs were less potent than today when the project started. Google saw this AI revolution early and wanted faster hardware for large-scale computing. Their only choice is to build a new solution for it.

That was why Google started this project, but there are more questions. Why were GPUs not good enough back in the day? What potential improvements did Google see that are significant enough to start their new hardware project?

The answer lies in the microarchitecture of GPUs and TPUs. Let's examine the design of the cores on GPUs and TPUs.

The design idea of GPUs

First, let's do a quick recap of the background knowledge of CPUs. When an instruction comes, it is decoded by the instruction decoder and fed into the arithmetic logic unit (ALU) together with data from the registers. The ALU does all the computing and returns the results to one of the registers. If you have multiple cores in the CPU, they can work in parallel.

What is a GPU? It is short for the graphics processing unit. It was designed for graphics computing and later discovered suitable for machine learning. Most of the operations in a GPU are matrix operations, which could run in parallel. This also means there are not many operations they need to support compared with a CPU.

The more specialized the chip is for a given task, the faster it is on the task.

The key idea of the GPU's initial design was to have a feature-reduced CPU with smaller but more cores for faster parallel computing. The number of instructions supported on a GPU is much less than on a CPU, which makes the area taken by a single core on a chip much smaller. This way, they can pack more cores onto the chip for large-scale parallel computing.

Why do fewer features mean a smaller area on the chip? In software, more features mean more code. In hardware, all features are implemented using logical circuits instead of code. More features mean the circuit is more complex. For example, a CPU must implement more instructions on the chip.

Smaller also means faster. A simpler design of the logic gates leads to a shorter cycle time.

The design idea of TPUs

TPU further developed this idea of specialized chips for deep learning. The defining feature of a TPU is its matrix-multiply unit (MXU). Since matrix multiplication is the most frequent operation in deep learning, TPU builds a specialized core for it, the MXU.

This is even more specialized than a GPU core, capable of many matrix operations, while the MXU only does one thing: matrix multiplication.

It works quite differently from a traditional CPU/GPU core. We usually refer to these processors with new architectures as application-specific integrated circuits (ASICs).

In an MXU, all the dynamics and generality are removed. It has a grid of nodes all connected together. Each node only does multiplication and addition in a predefined manner. The results are directly pushed to the next node for the next multiplication and addition. So, everything is predefined and fixed.

Here is a simple example to explain how it works. Every matrix multiplication is decomposed into a few vector inner-product computations. We have $\mathbf{x}$ and $\mathbf{y}$ as the input of the MXU, where

$$ \mathbf{x} = \begin{bmatrix} x_1,\ x_2,\ x_3\end{bmatrix},\ \mathbf{y} = \begin{bmatrix} y_1 \\ y_2 \\ y_3\end{bmatrix}. $$

We load the float values onto the registers and multiply the corresponding value pairs together to get

$$ \mathbf{z} = \begin{bmatrix} z_1,\ z_2,\ z_3\end{bmatrix}=\begin{bmatrix} x_1y_1,\ x_2y_2,\ x_3y_3\end{bmatrix}.\ $$

Then, we accumulate the values in $\mathbf{z}$ with

$$ \mathbf{s}=\begin{bmatrix} z_1,\ z_2,\ z_3\end{bmatrix}, $$ where $$s_i=s_{i-1} + z_i$$ and $$ s_1=z_1. $$

The following figure shows how the values flow in an MXU.

All the multiplications to compute $\mathbf{z}$ are independent, which can run in parallel. All the add operations to compute $\mathbf{s}$ are sequential. However, no operation will be blocked by this accumulation because the data of input vectors are streamed into the MXU. As soon as the multiplication of two floats are done, a new pair of $x$ and $y$ from the next pair of vectors are pushed in. Similarly, when we finish one add operaion to produce $s_2=s_1+z_2$, the $s_1$ and $z_2$ for the next vector pair is already there. It is like a factory pipeline, though the steps are sequential, they do not wait in idle, but stream the products through the pipeline.

This example only shows how to compute the vector inner-product of length 3. However, this small structure is tileable. If we stack more of such structure together, we can compute much longer vectors and matrices.

With MXU instead of GPU, we save time by removing the need for instruction decoding since it just multiplies and adds whatever it receives. There is no register for writing and reading since we already know where the results should go, and there is no need to store it for arbitrary operations that come next.

Besides the MXU, the TPU has also been designed for better scalability. It has dedicated ports for high-bandwidth inter-chip interconnection (ICI). It is designed to sit on the racks in Google's data centers and to be used in clusters. Since it is for proprietary usage only, they don't need to worry about selling single chips or the complexity of installing the chips on the racks.

Are TPUs still faster today?

It doesn't make sense others didn't come up with the same simple idea of building dedicated cores for tensor operations (matrix multiplication). Even if they did not, it doesn't make sense that they don't copy.

From the timeline, it seems Nvidia came up with the same idea at about the same time. A similar product from Nvidia, the Tensor Cores, was first announced to the public in 2017², one year after Google's TPU announcement.

It is unclear whether TPUs are still faster than GPUs today. I cannot find public benchmarks of the latest generations of TPUs and GPUs, and it is unclear to me which generation and metrics should be used for benchmarking.

However, we can use one universal application-oriented metric: dollars per epoch. I found an interesting benchmark from Google Cloud³ that aligns different hardware to the same axis: money. TPUs appear cheaper on Google Cloud if you have the same model, data, and number of epochs.

Large models, like Midjourney, Claude, and Gemini, are all very sensitive to the training cost because they consume too much computing power. As a result, many of them use TPUs on Google Cloud.

Why are TPUs cheaper?

One important reason is the software stack. You are using not only the hardware but also the software stack associated with it. Google has better vertical integration for its software stack and AI hardware than GPUs.

Google has dedicated engineering teams to build a whole software stack for it with strong vertical integration, from the model implementation (Vertex Model Garden) to the deep learning frameworks (Keras, JAX, and TensorFlow) to a compiler well-optimized for the TPUs (XLA).

The software stack for GPUs is very different. PyTorch is the most popular deep learning framework used with Nvidia GPUs, and it was mainly developed by Meta. The most widely used model pools with PyTorch are transformers and diffusers developed by HuggingFace. It is much harder to do perfect vertical integration for the software stack across all these companies.

One caveat is that fewer models are implemented with JAX and TensorFlow. Sometimes, you may need to implement the model yourself or use it from PyTorch on TPUs. Depending on the implementation, you may experience some friction when using PyTorch on TPUs. So, there might be extra engineering costs besides the hardware cost itself.

Why not start selling TPUs?

We understand the project was started for proprietary usage and acquired a pretty good user base on Google Cloud because of its lower price. Why did not Google just start to sell it to customers directly, just like Nvidia's GPUs?

The short answer is to stay focused. Google is in fierce competition with OpenAI for generative AI. At the same time, it is in the middle of multiple waves of tech layoffs to lower its cost. A wise strategy now would be to focus its limited resources on the most important projects.

If Google ever wants to start selling its TPUs, it will be competing with two strong opponents, Nvidia and OpenAI, at the same time, which may not be a wise move at the moment.

The huge overhead of selling hardware

Selling hardware directly to customers creates huge overheads for the company. Conversely, renting TPUs on their cloud services is much more manageable.

When TPUs are only served on the cloud, they can have a centralized way to install all the TPUs and related software. There is no need to deal with various installation environments or the difficulty of deploying a TPU cluster.

They know exactly how many TPUs to make. The demands are all internal, so there is no uncertainty. Thus, managing the supply chain is much easier.

Sales also become much easier since it is just selling the cloud service. There is no need to build a new team experienced in selling hardware.

The advantages of the TPU approach

Without all the overhead of selling hardware directly to the customers, Google got a few advantages in return.

First, they can have a more aggressive TPU architecture design. The TPUs have a unique way of connecting the chips. Unlike multiple GPUs that connect to the same board, TPUs are organized in cubes. They arranged 64 TPUs in a 4 by 4 by 4 cube to interconnect them with each other for faster inter-chip communication. There are 8960 chips in a single v5p Pod. They can be easily used together. This is the advantage of fully control your hardware installation environment.

Second, they can iterate faster to push out new generations. Since they only need to support a small set of use cases for proprietary usages, it drastically reduces their research and development cycle for every generation of the chips. I wonder if Nvidia came up with the TensorCore idea earlier than Google, but because of the overhead of selling hardware to external customers, they could only announce it one year later than Google.

From the perspective of serving its most important purpose, competing in GenAI, these advantages put Google in a very good position. Most importantly, with this in-house hardware solution, Google saved huge money by not buying GPUs from Nvidia at a monopoly price.

The downside of the TPU approach

So far, we have discussed many advantages of Google's AI hardware approach, but is there any downside? Indeed, there is a big one. Google became a tech island.

Every pioneer in tech will become an island isolated from the rest of the world, at least for a while. This is because they started early when the corresponding infrastructure was not ready. They need to build everything from scratch. Due to the migration cost, they will stick with their solution even if everyone else uses something else.

This is exactly what Google is experiencing right now. The rest of the world is innovating with models from HuggingFace and PyTorch. Everyone is quickly tweaking each other's models to develop better ones. However, Google cannot join this process easily since its infra is largely built around TensorFlow and JAX. When putting a model from external into production, it must be re-implemented with Google's framework.

This “tech island” problem slows Google down in taking good solutions from the external world and further isolates it from others. Google will either start bringing more external solutions like HuggingFace, PyTorch, and GPUs or always ensure its in-house solutions are the best in the world.

What does the future of AI hardware look like?

Finally, let's peek into the future of AI hardware. What would the future AI hardware look like? I have two observations: mode collapse as the hardware becomes more specialized, scale-out rather than scale-up.

Mode collapse

Hardware will be further coupled with the applications. For example, support more precision formats for better language model serving. Like with bfloat16 and TF32, they may better support int8 and int4. Nvidia announced their second generation of the Transformer Engine, which works with Blackwell GPU. This made optimizing their hardware for transformer models easier without changing the user code. A lot of codesign is happening.

On the other hand, software cannot easily jump out of the transformer realm. If they do, they will be slow due to a lack of hardware support. On the contrary, they implement their models with the hardware in mind. For example, the FlashAttention algorithm⁴ is designed to leverage the memory hierarchy of GPUs for better performance.

We see a big mode collapse coming soon. The hardware and software are so well optimized for each other for the current models. Neither of them can easily leave the current design or algorithm. If there is a new model completely different from the transformers, it needs to be 10x better to get widely adopted. It must incentivize people to make new hardware as fast and cheap as transformers.

Scaling out rather than scaling up

Today, AI hardware cares more about scaling out than scaling up. They prefer to optimize the overall performance of a large cluster over improving the performance of a single chip. This is mainly for two reasons.

First, it is hard to improve single-chip performance anymore. Both compute and high bandwidth memory sizes are facing challenges for further improvement. It is also hard to build a chip with a smaller feature size while balancing its cost and heat dissipation.

Second, the single-chip performance is no longer the bottleneck. In the era of large language models, no single chip can handle the entire training job. It requires hundreds or even thousands of GPUs to work together to train a model. The communication bandwidth between chips, machines, rooms, buildings, or even data centers becomes the bottleneck. NVidia is inventing new technologies like NVlink and InfiniBand, to replace the traditional PCIe. Google is also optimizing its TPU's inter-chip interconnection (ICI) for faster large-scale training. These technologies significantly increase the communication speed. We expect to see many more innovations in this direction.

Summary

In conclusion, the TPU project started for proprietary usage when the GPU's computing power was insufficient. Google wants to focus on GenAI instead of competing in the AI hardware market to avoid slowing the iteration speed and sacrificing its innovative design. Faster computing at a lower cost helped Google significantly in doing AI research and developing AI applications. However, it also made Google a tech island.

In the future, AI hardware will be even more optimized for specific applications, like the transformer models. Neither the hardware nor the models could easily jump out of this mode collapse. Also, the hardware will work better on a larger scale with more bandwidth between them.

References

1. Jouppi, N. (2016). Google supercharges machine learning tasks with custom chip. Google Cloud Blog. ↗ ↩
2. Nvidia. (2017). Inside Volta: The world’s most advanced data center GPU. Nvidia Developer Blog. ↗ ↩
3. Google Cloud Platform. (2023). JAX ViT benchmarking report on TPU vs GPU. Vertex AI Model Garden Samples. ↗ ↩
4. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. NeurIPS. ↗ ↩

Through all these years, Keras has survived as a high-level user-facing library supporting different backends, including TensorFlow, PyTorch, and JAX. Studying Keras closely, I learned how much the team cares about user experience for the software and how they ensured a good user experience by following a few simple yet powerful principles in their design process.

In this article, I will share the 3 most important software design principles I learned from studying Keras over the past few years, which may be generalizable to all types of software and help you make an impact in the open-source community with yours.

Why user experience is important for open-source software

Before we dive into the main content, let's quickly discuss why user experience is so important. We can learn this through the PyTorch vs. TensorFlow case.

They were developed by two tech giants, Meta and Google, and have quite different cultural strengths. Meta is good at product, while Google is good at engineering. As a result, Google's frameworks like TensorFlow and JAX are the fastest to run and technically superior to PyTorch, as they support sparse tensors and distributed training well. However, PyTorch still took away half of the market share from TensorFlow because it prioritizes user experience over other aspects of the software.

Better user experience wins for the research scientists who build the models and propagate them to the engineers, who take models from them since they don't always want to convert the models they receive from the research scientists to another framework. They will build new software around PyTorch to smooth their workflow, which will establish a software ecosystem around PyTorch.

TensorFlow also made a few blunders that caused its users to lose. TensorFlow's general user experience is good. However, its installation guide for GPU support was broken for years before it was fixed in 2022. TensorFlow 2 broke the backward compatibility, which cost its users millions of dollars to migrate.

So, the lesson we learned here is that despite technical superiority, user experience decides which software the open-source users would choose.

All deep learning frameworks invest heavily in user experience

All the deep learning frameworks, TensorFlow, PyTorch, and JAX, invest heavily in user experience. Good evidence is that they all have a relatively high Python percentage in their codebases.

All the core logic of deep learning frameworks, including tensor operations, automatic differentiation, compilation, and distribution are implemented in C++. Why would they want to expose a set of Python APIs to the users? It is just because the users love Python and they want to polish their user experience.

Investing in user experience is of high ROI

Imagine how much engineering effort it requires to make your deep learning framework a little bit faster than others. A lot.

However, for a better user experience, as long as you follow a certain design process and some principles, you can achieve it. For attracting more users, your user experience is as important as the computing efficiency of your framework. So, investing in user experience is of high return on investment (ROI).

The three principles

I will share the three important software design principles I learned by contributing to Keras, each with good and bad code examples from different frameworks.

Principle 1: Design end-to-end workflows

When we think of designing the APIs of a piece of software, you may look like this.

class Model:
    def __call__(self, input):
        """The forward call of the model.

        Args:
          input: A tensor. The input to the model.
        """
        pass

Define the class and add the documentation. Now, we know all the class names, method names, and arguments. However, this would not help us understand much about the user experience.

What we should do is something like this.

input = keras.Input(shape=(10,))
x = layers.Dense(32, activation='relu')(input)
output = layers.Dense(10, activation='softmax')(x)
model = keras.models.Model(inputs=input, outputs=output)
model.compile(
    optimizer='adam', loss='categorical_crossentropy'
)

We want to write out the entire user workflow of using the software. Ideally, it should be a tutorial on how to use the software. It provides much more information about the user experience. It may help us spot many more UX problems during the design phase compared with just writing out the class and methods.

Let's look at another example. This is how I discovered a user experience problem by following this principle when implementing a hyperparameter tuning library.

When using KerasTuner, users can use this RandomSearch class to select the best model. We have the metrics, and objectives in the arguments. By default, objective equals validation loss. So, it helps us find the model with the smallest validation loss.

class RandomSearch:
    def __init__(self, ..., metrics, objective="val_loss", ...):
        """The initializer.

        Args:
            metrics: A list of Keras metrics.
            objective: String or a custom metric function. The
                name of the metirc we want to minimize.
        """
        pass

Again, it doesn't provide much information about the user experience. So, everything looks OK for now.

However, if we write an end-to-end workflow like the following. It exposes many more problems. The user is trying to define a custom metric function named custom_metric. The objective is not so straightforward to use anymore. What should we pass to the objective argument now?

tuner = RandomSearch(
    ...,
    metrics=[custom_metric],
    objective="val_???",
)

It should be just "val_custom_metric”. Just use the prefix of "val_" and the name of the metric function. It is not intuitive enough. We want to make it better instead of forcing the user to learn this. We easily spotted a user experience problem by writing this workflow.

If you wrote the design more comprehensively by including the implementation of the custom_metric function, you will find you even need to learn how to write a Keras custom metric. You have to follow the function signature to make it work, as shown in the following code snippet.

def custom_metric(y_true, y_pred):
    squared_diff = ops.square(y_true - y_pred)
    return ops.mean(squared_diff, axis=-1)

After discovering this problem. We specially designed a better workflow for custom metrics. You only need to override HyperModel.fit() to compute your custom metric and return it. No strings to name the objective. No function signature to follow. Just a return value. The user experience is much better right now.

class MyHyperModel(HyperModel):
    def fit(self, trial, model, validation_data):
        x_val, y_true = validation_data
        y_pred = model(x_val)
        return custom_metric(y_true, y_pred)

tuner = RandomSearch(MyHyperModel(), max_trials=20)

One more thing to remember is we should always start from the user experience. The designed workflows backpropagate to the implementation.

Principle 2: Minimize cognitive load

Do not force the user to learn anything unless it is really necessary. Let's see some good examples.

The Keras modeling API is a good example shown in the following code snippet. The model builders already have these concepts in mind, for example, a model is a stack of layers. It needs a loss function. We can fit it with data or make it predict on data.

model = keras.Sequential([
    layers.Dense(10, activation="relu"),
    layers.Dense(num_classes, activation="softmax"),
])
model.compile(loss='categorical_crossentropy')
model.fit(...)
model.predict(...)

So basically, no new concepts were learned to use Keras.

Another good example is the PyTorch modeling. The code is executed just like Python code. All tensors are just real tensors with real values. You can depend on the value of a tensor to decide your path with plain Python code.

class MyModel(nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return self.path_a(x)
        return self.path_b(x)

You can also do this with Keras with TensorFlow or JAX backend but needs to be written differently. All the if conditions need to be written with this ops.cond function as shown in the following code snippet.

class MyModel(keras.Model):
    def call(self, inputs):
        return ops.cond(
            ops.sum(inputs) > 0,
            lambda : self.path_a(inputs),
            lambda : self.path_b(inputs),
        )