LLM speedrun: Architecture

AI Use

This entire article is written by me without any kind of text generated from AI. None of the python code is designed by AI either. I used AI to proofread the article to highlight gaps, and to translate the python code to Rust.

Audience

This blog series is me publishing my personal notes on the state of LLMs and scalably serving them. I’ll discuss the motivation for ideas specific to LLMs / transformers, but not the building blocks behind machine learning and language modeling. To that end, I expect the reader to understand these topics: normalization, standard feed forward networks, LSTMs and basics of linear algebra. To follow along in code, I expect you to either be comfortable with Python’s numpy or with Rust’s rayon crate (I actually familiarized myself with rayon while writing this blog).

Running code

Most sections have code attached with them that you can copy in your local and run directly. The full Qwen implementation is available in harness section.

For python, use:

uv run <file.py> <args>

For rust, use:

cargo +nightly -Zscript <file.rs> -- <args>

This is a first of three part series, deep diving into individual parts of LLMs, starting with architecture. The other 2 planned posts are: inference, and training.

A note on notation:

All scalers are represented as $1 \times 1$ matrix.
All vectors are represented as column vectors, like $d_{embed} \times 1$ .
A matrix of $X$ vectors is represented as $X$ concatenated column vectors, like $d_{embed} \times X$ .
$\times$ represents a matrix multiplication, while $\odot$ represents an element wise multiplication.

Overview

For a while, one popular way to model language has been to predict the next token.

That is the premise behind RNNs or their most successful derivative, LSTMs. But these models require all previous tokens to be processed before processing the current one. This makes the training process inherently serial.

This was a huge bottleneck as we couldn’t train over a large set of corpus (and apparently we needed to train over sum of all human knowledge for models to be able to speak fluently).

The transformers architecture solved that problem, with the core mechanism of “attention,” which was introduced in a 2014 paper. Folks at Google looked at it and said “maybe… maybe that’s all you need?” leading to the “Attention is all you need” paper in 2017 --- unironically for machine translation --- paving the way to parallelize the computation of language models.

Transformer architecture

Setting up the inputs

Similar to RNNs, the input text is converted into learned embeddings. Concretely, given a sequence of length $T$ tokens, the embedding for the sequence is represented by a matrix of $d_{embed} \times T$ . But what about the position? In RNNs the position is automatic, as the tokens are processed one at a time but that is not the case with transformers where each token is processed in a position-independent way. So, position is encoded as an embedding as well which is added to the embedding of a token (modern LLMs do not add positions, rather use “Rotary Position Embedding”).

Attention

Attention mechanism entails that each token asks a query “how relevant are you to me” to every token (including itself) that answers using its key. Then the weighted some of values of all tokens is the final value of the querying token. The inputs pass through the transformer block which is made up of many attention layers (as anything in deep learning).

In RNNs, the single state vector carries the information of every text that has been seen by the model up to a point (and it can be arbitrarily large). Attention allows each token to look back at all the past tokens, hence allowing it to change its context encoding dynamically based on its needs. So, being able to train in parallel and being able to look back allowed attention mechanism to be extremely powerful in representing language.

Concretely, each token (of size $d_{embed}$ ) has a meaning in some high dimensional space called a “value” space (of size $d_{model}$ ). Now each token wants to know how much other tokens matter to it (and other tokens will also ask the same of it). Or stated otherwise, how much of “value” of token $x$ should affect a token $y$ (where $y$ can be equal to $x$ ).

How to figure that out? What if each token asks a “query” to every other token, and all those other tokens provide a “key” which tells how relevant it is to the query token. This means that each token has 3 high dimensional representation $K$ , $Q$ and $V$ , which is calculated by 3 matrices $W^K$ , $W^Q$ , and $W^V$ .

Now the input sequence is transformed into 3 other sequence matrices:

K = W^K \times I

Q = W^Q \times I

V = W^V \times I

And each key query pair is resolved as:

Attention^{weights} = Softmax(\frac{K^T \times Q}{\sqrt{d_{model}}})

Size of $Attention^{weights}$ is $T \times T$ noting the attention each token is paying to every other token.

Finally the output of attention is attention score weighted sum of values:

Attention^{output} = V \times Attention^{weights}

QWhat's up with Softmax?

What’s the technique to output probabilistic contribution of a bunch of things based on some unbounded weights? You got it - Softmax.

QThen what's up with

\sqrt{d_{model}}

That’s just normalization to make sure that softmax does not misbehave and its variance is 1.0.

Every token with every other token, even the future ones?

I am no alien from Arrival, I read sequentially

We do look at things in parallel too. We do not read each pixel of an image sequentially. So, looking into the future has its advantages.

But for language that’s not the case. So a causal mask makes sure a token does not attend to future tokens.

M_{ij} = \begin{cases} 0 & \text{if } i \le j \\ -\infty & \text{if } i \gt j \end{cases}

Attention^{weights} = Softmax(\frac{K^T \times Q}{\sqrt{d_{model}}} + M)

where $i$ is the row index and $j$ is the column index.

I am going to omit this causal mask in future equations to keep math representations terse.

But how does it remember things?

Attention is not the end of things for a transformer, which also has a classic fully connected layer, a feed-forward network. This layer is huge, typically having $4 \cdot d_{model}$ the intermediate activations. This is also the bulk of the weights in an LLM (generally accounting for 70% of it). Modern LLMs use gated activations rather than ReLU, that too typically SwiGLU.

Output = W_{down} \times ((W_{up} \times Attention^{output}) \odot Swish(W_{gate} \times Attention^{output}))

This FFN layer is considered to be the “brain” of an LLM, however, interpretability of ML models is an active area of research. So Google around (ironic in an LLM blog).

Stacking

The output from one layer of transfomer is passed to the next one. One can stack as many of these layers as one wants (and have the capacity to train, something called scaling laws that I’ll write about in the training part of the series).

You might notice one discrepancy though. The input to the first layer is of size $d_{embed} \times T$ , while the output is of size $d_{model} \times T$ . So either the rest of the layers need to take input of $d_{model}$ size or $d_{embed}$ is set equal to $d_{model}$ and generally the latter is chosen but not for just this reason.

Notice how setting

d_{embed} = d_{model}

allows arbitrary stacking of the same layer. This is very powerful as this makes making deep networks trivial, all the while retaining the parallel training benefits of each individual layer.

Gradients be vanishing

or activations exploding

Notice that softmax is a part of the attention layer and FFN will also include an activation function. This makes it very hard for the model to learn or makes learning unstable as multiple transformer layers are stacked. So transformer borrows the same idea from “ResNet” and uses residual connections.

In transformer, residual connection is added at two places: After calculating $Attention^{output}$ and at the final output location $Output$ . Mathematically:

Intermediate = I + Attention^{output}

Output = I + FFN(Intermediate)

This does require $d_{model} = d_{embed}$

NNote on residual impact

With residual flow, the intent of the entire network has changed. Each layer of transformer now acts as a perturbation on the initial input. Conceptually, initial layers will act as a larger perturbation (the network is trying to figure out the language specifics) and the later layers as smaller perturbations (the network is finetuning the representation of the larger image). Later, I’ll discuss how this is supported mathematically.

Normalization

Like with any model created after 2015, networks need to normalize, with the basic idea being the same: network learns better when the distribution of inputs to each layer is constrained and well defined. Otherwise in a deep network, just after a few matrix multipliers we can reach really high values. Transformer contains:

Residual connections, famous for leading to explosion in activations.
Activation output depending on softmax where large inputs can lead to spiky outputs and hence poor cross token attention outputs.

Normalization is applied to $Attention^{output}$ and to either input or output of a transformer layer. For the latter, where normalization is applied matters and is known as Pre-Norm v/s Post-Norm.

Why pre-norm over post-norm?

Take a look at the output from a post-norm model:

Output_l = RMSNorm(I_l + F(I_l))

and compare this with the output from pre-norm model:

Output_l = I_l + F(RMSNorm(I_l))

Pre norm has two benefits:

The input residual is unobstructed and the gradients flow back nicely, without subjected to normalization pass.
This naturally leads to small perturbations from the latter layers as compared to initial layers. The input to a layer is normalized so $\frac{F(RMSNorm(I_l))}{I_l}$ gets smaller and smaller for deeper layers.

Multi Head Attention (MHA)

Increasing efficiency

Attention mechanism represents each token as three $d_{model}$ sized vectors. The tokens then figure out the relationship between each other. But what kind of relationship? Are those tokens going to tell what “it” refers to from the rest of the sentence? Or are they going to talk about the grammar of the sentence? Or maybe they are trying to understand a single word which was not in the dictionary and hence got split into multiple single character tokens.

Having a single large attention matrix makes it difficult for the model to reason about multiple things in parallel. Empirical evidence also shows that each attention is trying to reason about one specific stuff, so having such a large $d_{model}$ representation is a waste, where a low-rank representation will do.

So, rather than computing a single large attention, in MHA, attention is split into multiple smaller representations. Concretely, there are $h$ attention heads, each with internal representation of size $d_k = \frac{d_{model}}{h}$ . The output of each head is concatenated in the end to recover $d_{model}$ sized output vector of the attention layer.

Note how this further aids parallel execution of a transformer block, none of the attention heads are dependent upon each other.

In equations (where $i$ represents a specific attention head):

W_i^K, W_i^Q, W_i^V \in \mathbb{R}^{d_k \times d_{model}}

K_i = W^K_i \times I, Q_i = W^Q_i \times I, V_i = W^V_i \times I

Attention^{output}_i = V_i \times Softmax(\frac{K_i^T \times Q_i}{\sqrt{d_k}})

Attention^{output} = \begin{bmatrix} Attention^{output}_0 \\ Attention^{output}_1 \\ \vdots \\ Attention^{output}_{h-1} \end{bmatrix}

Output matrix

Notice that now the attention output is now poorly represented in the residual connection:

Intermediate = I + Attention^{output}

Intermediate = I + \begin{bmatrix} Attention^{output}_0 \\ Attention^{output}_1 \\ \vdots \\ Attention^{output}_{h-1} \end{bmatrix}

Each attention head can only affect a subset of the residual. And hence the need for an output matrix (denoted by $W^O$ ) which mixes up information from all attention heads.

Intermediate = I + W^O \times Attention^{output}

QWhy not keep the

V

still in

d_{model} \times T

That’s a valid representation, but that’s unnecessary heavy computation per head now and the whole point of MHA is to increase efficiency of each attention layer.

QWhy not merge

W^O

with each head's

W^V

as two transformations:

W^O_i \times W^V_i \times I

That’s also a valid representation, but now the number of matrix multiplication has increased from

1

h

and the storage requirement per attention head for

V

also increases from

d_k \times T

d_{model} \times T

. This low-rank representation and post output matrix multiplication helps with both storage and compute.

Finishing up transformer

Transformer output

The output from the stacked layers is just the last token’s final representation. The expectation is that this token now includes all the context the next stage needs. Using pre-norm requires that the output of the stacked transformer be normalized before it is fed further.

O_{transformer} = RMSNorm(O_{stacked}) \times \begin{bmatrix} 0 \\ \vdots \\ 0 \\ 1 \end{bmatrix}

Summed up

So the overall equation for one transformer layer looks like:

Intermediate_l = I_l + W^O \times Attention^{output}(RMSNorm(I_l))

Output_l = Intermediate_l + FFN(RMSNorm(Intermediate_l))

Where $I_1 = I$ and $I_{l+1} = Output_l$ . After stacking $L$ layers, the output of layer $L$ ( $Output_L$ ) is normalized to produce the final transformer vector:

O_{transformer} = RMSNorm(Output_L) \times \begin{bmatrix} 0 \\ \vdots \\ 0 \\ 1 \end{bmatrix}

MHA matmaul

In code, we are looping over all the attention heads and calculating one at a time, but we can also represent all the heads as a single matrix of size

h \times d_{model} \times d_{model}

Contemporay modifications

Grouped Query Attention (GQA)

Storing and retrieving KV cache is costly and we want to reduce that and make it tenable to serve the model on a single GPU. In GQA, multiple attention heads share the same KV matrices (called “KV heads”). The precursor to this was “Multi Query Attention” which forced all heads to share a single KV, but that decreased accuracy considerably. Whereas GQA maintained the accuracy while decreasing the memory requirements (savings depends on the ratio $\frac{KV\_heads}{Query\_heads}$ , but anywhere from 50% to 90%).

But why does this work? Each head is still asking a different query, but the high dimensional space of KV is equipped to answer those multiple queries. During training, the model learns to superimpose multiple concepts to same KV.

Rotary Positional Embeddings (RoPE)

Previously, I mentioned that position embeddings can be learned and simply added to the token embeddings. But this is suboptimal for multiple reasons:

That position additive term does not cancel itself in the $K^T \times Q$ product. In fact, it produces an additional term $P_i \cdot P_j$ which is dependent on the absolute location. Whereas attention’s requirement is to know the relative distance between tokens.
It has a fixed length, so positional embeddings are constrained by the training set size and does not generalize to longer sequences natively.
Storing positional embeddings for longer sequences waste quite a lot of memory (hence storage, bandwidth etc etc).

The original transformer paper uses sinusoidal position embeddings to solve some of these problems. That actually boils down to just a rotation in a constrained way (check "You could have designed state of the art positional encoding" for details).

The question boils down to “how to encode the absolute position for each token which when seen by attention layer acts as a relative distance between attending tokens”. One elegant way is the rotation in complex plane. Constraining in two-dimension, dot product of two complex numbers is represented by $\lVert \mathbf{v_1} \rVert \lVert \mathbf{v_2} \rVert cos(\theta)$ where $\theta$ represents the angular distance between the vectors. If each vector encodes a token’s absolute position, this naturally encodes the relative position of tokens in dot product as $cos({\theta}_1 - {\theta}_2)$ . Now choosing a fixed rotation frequency and $m$ and $n$ representing the absolute position of those vectors, that term becomes $cos((m - n) \theta)$ .

It is much easier to represent rotation in two dimensional space, than in higher dimensions where the freedom of rotation explodes. So to represent this rotation in embedding space, each pair of the embedding is thought to represent a complex number. Mathematically, embedding is transformed from a high dimensional space of $d_k$ real numbers into a high dimensional space of $\frac{d_k}{2}$ complex numbers.

For a 2D pair at position $m$ , the rotation by angle $m\theta$ is:

R_{\theta, m} \begin{bmatrix} x \\ y \end{bmatrix} = \begin{bmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{bmatrix} \begin{bmatrix} x \\ y \end{bmatrix}

Extending this to the full $d_k$ -dimensional embedding vector $\mathbf{x}$ , the block-diagonal rotation matrix $R_{\theta, m}$ is defined as:

R_{\theta, m} \mathbf{x} = \begin{bmatrix} \cos(m\theta) & -\sin(m\theta) & \dots & 0 & 0 \\ \sin(m\theta) & \cos(m\theta) & \dots & 0 & 0 \\ \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & \dots & \cos(m\theta) & -\sin(m\theta) \\ 0 & 0 & \dots & \sin(m\theta) & \cos(m\theta) \end{bmatrix} \begin{bmatrix} x_1 \\ x_2 \\ \vdots \\ x_{d_k-1} \\ x_{d_k} \end{bmatrix}

What should be the value of $\theta$ , i.e. what is the optimal rotation frequency?

If we choose a very fast moving frequency $\theta_{fast}$ , then the dot product will start repeating very soon. Like for $\theta = 45^\circ$ , $cos((10 - 2) \theta) = cos((2-2) \theta)$ , basically treating the token at 8 positions away as same as attending to self token.
If we choose a very slow moving frequency $\theta_{slow}$ , then the dot product will look almost the same. Like for $\theta = 0.0001^\circ$ , $cos((3 - 2) \theta) \approx cos((2 - 2) \theta)$ , basically failing to distinguish between nearby tokens.

A single $\theta$ choice fails to capture everything. So, RoPE uses pairwise $\theta$ encoded based on a decaying procedure:

\theta_i = b^{\frac{-(i - 1)}{d_k / 2}}

We use a separate $\theta$ for each complex number and $b$ is the base which controls how large our context window can be before the rotations start to repeat. The original transformer paper used $b = 10000$ but modern LLMs, with their million long context window, use much larger base.

Mathematically this looks like:

R_{\Theta, m} \mathbf{x} = \begin{bmatrix} \cos(m\theta_1) & -\sin(m\theta_1) & \dots & 0 & 0 \\ \sin(m\theta_1) & \cos(m\theta_1) & \dots & 0 & 0 \\ 0 & 0 & \dots & 0 & 0 \\ 0 & 0 & \dots & 0 & 0 \\ \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & \dots & \cos(m\theta_{d_k/2}) & -\sin(m\theta_{d_k/2}) \\ 0 & 0 & \dots & \sin(m\theta_{d_k/2}) & \cos(m\theta_{d_k/2}) \end{bmatrix} \begin{bmatrix} x_1 \\ x_2 \\ x_3 \\ x_4 \\ \vdots \\ x_{d_k-1} \\ x_{d_k} \end{bmatrix}

where the frequencies are $\theta_i = 10000^{-2(i-1)/d_k}$ .

Generalizing over the input sequence $I$ :

positional\_embedding = R_{\Theta,m} \times I

I also view RoPE acting as a forcing function to allow attention to learn important concepts of human language. If you look at a pronoun, its context is in nearby words only. Similarly this rotation acts as forcing function telling the attention layer if it is looking for such language concepts, it should force them to be in the earlier part of the embeddings.

For a more mathematically rigorous handling of RoPE, check “Rotary Embeddings: A Relative Revolution”.

Decoding

But where is the output of the model?

We expect that the final token of the last transformer layer encodes what comes next, let’s call it $O_{transformer}$ . A reverse embedding (called “LM Head”) inverts that to a logits over the dictionary. So the first token is:

T_1 = Sample(Softmax(O_{transformer} \times LM\_Head))

LM head can just be the inverse of the embedding as well. That will tie the encoding and decoding weights and saves parameters of the model in one of the heavier parts.

Now this procedure can be repeated as many times as needed. Append the latest generated token to the input, and go through the entire process again to generate the next token and so on and so on (this is what is happening when you talk to your AI bot, it is inherently a linear process of generating tokens.) Obviously, we are not going to calculate the $KQV$ for all the tokens that we have already processed. In fact, their $Q$ is not even required, just the $KV$ to calculate how the new token is impacted by those past tokens. So, we just store that $K$ and $V$ and that is the “KV cache”.

Implementing Qwen

Now that I have introduced all the building blocks of transformer and LLMs, let’s modify our implementation to specifics of “Qwen2.5-Coder-7B-Instruct” model and generate some text.

Introducing bias

Modern LLMs are actually bias free, because the inputs are normalized at all the stages and empirical evidence has shown that bias does not improve the speed of convergence of the model. But Qwen uses bias so we should add that to be able to run this model.

Split-half RoPE implementation

As I described, RoPE was implemented pairwise over $d_k$ . But Qwen uses split-half implementation, where the hidden state is not considered pairwise, rather split in the middle. The real part of the complex number belongs to the first half and the imaginary part to the second half. This does not change the math, just how the matrix is setup. The four components of the original rotation are split in four quadrants now when setting up the RoPE matrix.

R_{\Theta, m} \mathbf{x} = \begin{bmatrix} \cos(m\theta_1) & 0 & \dots & 0 & -\sin(m\theta_1) & 0 & \dots & 0 \\ 0 & \cos(m\theta_2) & \dots & 0 & 0 & -\sin(m\theta_2) & \dots & 0 \\ \vdots & \vdots & \ddots & \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \dots & \cos(m\theta_{d_k/2}) & 0 & 0 & \dots & -\sin(m\theta_{d_k/2}) \\ \sin(m\theta_1) & 0 & \dots & 0 & \cos(m\theta_1) & 0 & \dots & 0 \\ 0 & \sin(m\theta_2) & \dots & 0 & 0 & \cos(m\theta_2) & \dots & 0 \\ \vdots & \vdots & \ddots & \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \dots & \sin(m\theta_{d_k/2}) & 0 & 0 & \dots & \cos(m\theta_{d_k/2}) \end{bmatrix} \begin{bmatrix} x_1 \\ x_2 \\ \vdots \\ x_{d_k/2} \\ x_{d_k/2 + 1} \\ x_{d_k/2 + 2} \\ \vdots \\ x_{d_k} \end{bmatrix}

Plugging in the specific parameters

These are the specifics of our Qwen model:

Number of layers ( $L$ ): $28$ .
RoPE base ( $b$ ): $1,000,000$ .
Hidden size ( $d_{embed} = d_{model}$ ): $3584$ .
Attention heads ( $h$ ): $28$ .
KV heads: $4$ .
Head dimension ( $d_k$ ): $128$ .
FFN size ( $d_{ff}$ ): $18944$ .

Harness

A few more things to get Qwen to generate tokens:

Tokenization
Loading the right weights
Setting up prompt
Setting up end of stream token

Copy it and run it on your local - Qwen implementation in numpy

# /// script
# requires-python = ">=3.14"
# dependencies = [
#     "huggingface-hub>=1.27.0",
#     "jinja2>=3.1.6",
#     "numpy>=2.5.2",
#     "psutil>=7.2.2",
#     "transformers>=5.15.0",
# ]
# ///
import argparse
from dataclasses import dataclass
import json
import os
import struct
import time
import numpy as np
import numpy.typing as npt
import psutil
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
from typing import Annotated
EmbedxCtx = Annotated[npt.NDArray[np.float32], "shape=(d_embed, max_context)"]
EmbedxDict = Annotated[npt.NDArray[np.float32], "shape=(d_embed, dict_size)"]
ModelxEmbed = Annotated[npt.NDArray[np.float32], "shape=(d_model, d_embed)"]
EmbedxTokens = Annotated[npt.NDArray[np.float32], "shape=(d_embed, T)"]
ModelxTokens = Annotated[npt.NDArray[np.float32], "shape=(d_model, T)"]
ModelxModel = Annotated[npt.NDArray[np.float32], "shape=(d_model, d_model)"]
FFNxModel = Annotated[npt.NDArray[np.float32], "shape=(d_ff, d_model)"]
ModelxFFN = Annotated[npt.NDArray[np.float32], "shape=(d_model, d_ff)"]
KVxModel = Annotated[npt.NDArray[np.float32], "shape=(num_kv_heads * d_k, d_model)"]
def get_ram_mb() -> float:
  process = psutil.Process(os.getpid())
  return process.memory_info().rss / (1024 * 1024)
class WeightContext:
  def __init__(self, loader: "MemoryEfficientSafetensorsLoader"):
    self.loader = loader
    self.loaded_tensors: list[np.ndarray] = []
  def load(self, name: str) -> np.ndarray:
    arr = self.loader.load_tensor(name)
    self.loaded_tensors.append(arr)
    return arr
  def __enter__(self):
    return self
  def __exit__(self, exc_type, exc_val, exc_tb):
    self.loaded_tensors.clear()
class MemoryEfficientSafetensorsLoader:
  def __init__(self, repo_id: str = "Qwen/Qwen2.5-Coder-7B-Instruct"):
    print("Checking / downloading model weight shards from HuggingFace Hub...")
    self.model_dir = snapshot_download(
      repo_id=repo_id, allow_patterns=["*.safetensors", "*.json"]
    )
    index_path = os.path.join(self.model_dir, "model.safetensors.index.json")
    with open(index_path, "r") as f:
      index = json.load(f)
    self.weight_map = index["weight_map"]
    self.file_headers = {}
    shard_files = set(self.weight_map.values())
    for shard in shard_files:
      shard_path = os.path.join(self.model_dir, shard)
      with open(shard_path, "rb") as f:
        header_len = struct.unpack("<Q", f.read(8))[0]
        header_json = f.read(header_len).decode("utf-8")
        header = json.loads(header_json)
        self.file_headers[shard] = {"header_len": header_len, "header": header}
  def load_tensor(self, name: str) -> np.ndarray:
    shard = self.weight_map[name]
    shard_path = os.path.join(self.model_dir, shard)
    info = self.file_headers[shard]
    meta = info["header"][name]
    header_len = info["header_len"]
    start, end = meta["data_offsets"]
    shape = meta["shape"]
    dtype_str = meta["dtype"]
    with open(shard_path, "rb") as f:
      f.seek(8 + header_len + start)
      raw_bytes = f.read(end - start)
    if dtype_str == "BF16":
      u16 = np.frombuffer(raw_bytes, dtype=np.uint16)
      arr = (u16.astype(np.uint32) << 16).view(np.float32)
    elif dtype_str == "F32":
      arr = np.frombuffer(raw_bytes, dtype=np.float32)
    elif dtype_str == "F16":
      arr = np.frombuffer(raw_bytes, dtype=np.float16).astype(np.float32)
    else:
      raise ValueError(f"Unsupported tensor dtype: {dtype_str}")
    return arr.reshape(shape)
  def load_scope(self) -> WeightContext:
    return WeightContext(self)
@dataclass
class Qwen2_5_Coder_7B_Config:
  num_layers: int = 28
  hidden_size: int = 3584
  num_heads: int = 28
  num_kv_heads: int = 4
  head_dim: int = 128
  intermediate_size: int = 18944
  vocab_size: int = 152064
  rope_theta: float = 1000000.0
  rms_norm_eps: float = 1e-6
class Embedding:
  def __init__(self, w_embed: EmbedxDict):
    self.w_embed = w_embed
  def embed(self, token_ids: list[int]) -> EmbedxTokens:
    return self.w_embed[token_ids].T
class QwenRotaryEmbedding:
  def __init__(self, head_dim: int, base: float = 1000000.0):
    self.head_dim = head_dim
    self.base = base
    inv_freq = 1.0 / (base ** (np.arange(0, head_dim, 2, dtype=np.float32) / head_dim))
    self.inv_freq = inv_freq
  def apply(self, x: np.ndarray, positions: np.ndarray) -> np.ndarray:
    freqs = np.outer(self.inv_freq, positions)
    emb = np.concatenate([freqs, freqs], axis=0)
    cos = np.cos(emb)[None, :, :]
    sin = np.sin(emb)[None, :, :]
    half = self.head_dim // 2
    x1 = x[:, :half, :]
    x2 = x[:, half:, :]
    rotate_half = np.concatenate([-x2, x1], axis=1)
    return (x * cos) + (rotate_half * sin)
def softmax(x: np.ndarray, axis: int = 0) -> np.ndarray:
  x_max = np.max(x, axis=axis, keepdims=True)
  exp_x = np.exp(x - x_max)
  return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
class KVCache:
  def __init__(self):
    self.k_cache: dict[int, np.ndarray] = {}
    self.v_cache: dict[int, np.ndarray] = {}
  def update(
    self, layer_id: int, new_k: np.ndarray, new_v: np.ndarray
  ) -> tuple[np.ndarray, np.ndarray]:
    if layer_id not in self.k_cache:
      self.k_cache[layer_id] = new_k
      self.v_cache[layer_id] = new_v
    else:
      self.k_cache[layer_id] = np.concatenate([self.k_cache[layer_id], new_k], axis=-1)
      self.v_cache[layer_id] = np.concatenate([self.v_cache[layer_id], new_v], axis=-1)
    return self.k_cache[layer_id], self.v_cache[layer_id]
class GroupedQueryAttention:
  def __init__(
    self,
    num_heads: int,
    num_kv_heads: int,
    d_model: int,
    w_k: KVxModel,
    k_b: np.ndarray,
    w_q: ModelxModel,
    q_b: np.ndarray,
    w_v: KVxModel,
    v_b: np.ndarray,
    w_o: ModelxModel,
    rope: QwenRotaryEmbedding,
  ):
    self.num_heads = num_heads
    self.num_kv_heads = num_kv_heads
    self.d_model = d_model
    self.d_k = d_model // num_heads
    self.queries_per_kv = num_heads // num_kv_heads
    self.w_k, self.k_b = w_k, k_b[:, None]
    self.w_q, self.q_b = w_q, q_b[:, None]
    self.w_v, self.v_b = w_v, v_b[:, None]
    self.w_o = w_o
    self.rope = rope
  def forward(
    self, x: np.ndarray, positions: np.ndarray, kv_cache: KVCache, layer_id: int = 0
  ) -> np.ndarray:
    T = x.shape[1]
    q = self.w_q @ x + self.q_b
    k = self.w_k @ x + self.k_b
    v = self.w_v @ x + self.v_b
    q = q.reshape(self.num_heads, self.d_k, T)
    k = k.reshape(self.num_kv_heads, self.d_k, T)
    v = v.reshape(self.num_kv_heads, self.d_k, T)
    q = self.rope.apply(q, positions)
    k = self.rope.apply(k, positions)
    k, v = kv_cache.update(layer_id, k, v)
    T_total = k.shape[-1]
    k = np.repeat(k, self.queries_per_kv, axis=0)
    v = np.repeat(v, self.queries_per_kv, axis=0)
    # Note: Modifying the matrix multiplication here to take advantage of vector maths for speed. We'll talk more about it in the inference post.
    scores = (q.transpose(0, 2, 1) @ k) / np.sqrt(self.d_k)
    causal_mask = np.triu(np.ones((T, T_total), dtype=bool), k=T_total - T + 1)
    scores[:, causal_mask] = -1e9
    weights = softmax(scores, axis=-1)
    head_outputs = (v @ weights.transpose(0, 2, 1)).reshape(self.d_model, T)
    return self.w_o @ head_outputs
class SwiGLUFFN:
  def __init__(self, w_gate: FFNxModel, w_up: FFNxModel, w_down: ModelxFFN):
    self.w_gate = w_gate
    self.w_up = w_up
    self.w_down = w_down
  @staticmethod
  def swish(x: np.ndarray) -> np.ndarray:
    return x / (1.0 + np.exp(-x))
  def forward(self, x: ModelxTokens) -> ModelxTokens:
    gate = self.w_gate @ x
    up = self.w_up @ x
    gated_act = up * self.swish(gate)
    output = self.w_down @ gated_act
    return output
class RMSNorm:
  def __init__(self, weight: npt.NDArray[np.float32], eps: float = 1e-6):
    self.weight = weight[:, None]
    self.eps = eps
  def forward(self, x: np.ndarray) -> np.ndarray:
    variance = np.mean(x**2, axis=0, keepdims=True)
    return (x / np.sqrt(variance + self.eps)) * self.weight
class TransformerBlock:
  def __init__(
    self,
    attention: GroupedQueryAttention,
    ffn: SwiGLUFFN,
    norm1: RMSNorm,
    norm2: RMSNorm,
  ):
    self.attention = attention
    self.ffn = ffn
    self.norm1 = norm1
    self.norm2 = norm2
  def forward(
    self,
    x: ModelxTokens,
    positions: np.ndarray,
    kv_cache: KVCache,
    layer_id: int = 0,
  ) -> ModelxTokens:
    norm_x1 = self.norm1.forward(x)
    attn_out = self.attention.forward(norm_x1, positions, kv_cache, layer_id)
    intermediate = x + attn_out
    norm_x2 = self.norm2.forward(intermediate)
    ffn_out = self.ffn.forward(norm_x2)
    output = intermediate + ffn_out
    return output
class StackedTransformer:
  def __init__(
    self,
    loader: MemoryEfficientSafetensorsLoader,
    cfg: Qwen2_5_Coder_7B_Config,
    rope: QwenRotaryEmbedding,
  ):
    self.loader = loader
    self.cfg = cfg
    self.rope = rope
  def forward(
    self, x: ModelxTokens, positions: np.ndarray, kv_cache: KVCache
  ) -> ModelxTokens:
    for layer_idx in range(self.cfg.num_layers):
      prefix = f"model.layers.{layer_idx}."
      with self.loader.load_scope() as ctx:
        in_norm_w = ctx.load(prefix + "input_layernorm.weight")
        norm1 = RMSNorm(in_norm_w, eps=self.cfg.rms_norm_eps)
        q_w = ctx.load(prefix + "self_attn.q_proj.weight")
        q_b = ctx.load(prefix + "self_attn.q_proj.bias")
        k_w = ctx.load(prefix + "self_attn.k_proj.weight")
        k_b = ctx.load(prefix + "self_attn.k_proj.bias")
        v_w = ctx.load(prefix + "self_attn.v_proj.weight")
        v_b = ctx.load(prefix + "self_attn.v_proj.bias")
        o_w = ctx.load(prefix + "self_attn.o_proj.weight")
        attn = GroupedQueryAttention(
          self.cfg.num_heads,
          self.cfg.num_kv_heads,
          self.cfg.hidden_size,
          k_w,
          k_b,
          q_w,
          q_b,
          v_w,
          v_b,
          o_w,
          self.rope,
        )
        post_norm_w = ctx.load(prefix + "post_attention_layernorm.weight")
        norm2 = RMSNorm(post_norm_w, eps=self.cfg.rms_norm_eps)
        gate_w = ctx.load(prefix + "mlp.gate_proj.weight")
        up_w = ctx.load(prefix + "mlp.up_proj.weight")
        down_w = ctx.load(prefix + "mlp.down_proj.weight")
        ffn = SwiGLUFFN(gate_w, up_w, down_w)
        block = TransformerBlock(attn, ffn, norm1, norm2)
        x = block.forward(x, positions, kv_cache, layer_idx)
    with self.loader.load_scope() as ctx:
      final_norm_w = ctx.load("model.norm.weight")
      final_norm = RMSNorm(final_norm_w, eps=self.cfg.rms_norm_eps)
      output = final_norm.forward(x)[:, -1:]
    return output
class LMHead:
  def __init__(self, w_head: npt.NDArray[np.float32]):
    self.w_head = w_head
  def forward(
    self, last_hidden_state: npt.NDArray[np.float32]
  ) -> npt.NDArray[np.float32]:
    return (self.w_head @ last_hidden_state).squeeze(-1)
class Decoder:
  def __init__(
    self, transformers: StackedTransformer, loader: MemoryEfficientSafetensorsLoader
  ):
    self.transformers = transformers
    self.loader = loader
  def step(
    self, x: np.ndarray, positions: np.ndarray, kv_cache: KVCache
  ) -> tuple[int, np.ndarray, np.ndarray]:
    last_hidden_state = self.transformers.forward(x, positions, kv_cache)
    with self.loader.load_scope() as ctx:
      lm_head_w = ctx.load("lm_head.weight")
      lm_head = LMHead(lm_head_w)
      logits = lm_head.forward(last_hidden_state)
    next_token_id = int(np.argmax(logits))
    return next_token_id, logits, last_hidden_state
def generate_tokens(
  max_new_tokens: int = 3, prompt: str = "Give me the quicksort algorithm in python"
):
  print("=== Qwen2.5-Coder-7B-Instruct pure NumPy Inference ===")
  print(f"Prompt: {prompt!r}")
  print(f"Max New Tokens: {max_new_tokens}")
  print(f"Initial Process memory consumption: {get_ram_mb():.1f} MB\n")
  t_start = time.time()
  tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")
  eos_token_ids = {
    t_id
    for t_id in [
      tokenizer.eos_token_id,
      tokenizer.convert_tokens_to_ids("<|im_end|>"),
      tokenizer.convert_tokens_to_ids("<|endoftext|>"),
    ]
    if t_id is not None
  }
  messages = [
    {
      "role": "system",
      "content": "You are a helpful assistant specializing in coding.",
    },
    {"role": "user", "content": prompt},
  ]
  text = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
  )
  tokens = tokenizer.encode(text)
  seq_len = len(tokens)
  print(f"Formatted Chat Prompt:\n{text}")
  print(f"Tokenized Sequence Length: {seq_len} tokens.\n")
  loader = MemoryEfficientSafetensorsLoader("Qwen/Qwen2.5-Coder-7B-Instruct")
  cfg = Qwen2_5_Coder_7B_Config()
  rope = QwenRotaryEmbedding(head_dim=cfg.head_dim, base=cfg.rope_theta)
  transformers = StackedTransformer(loader, cfg, rope)
  decoder = Decoder(transformers, loader)
  generated_tokens = []
  cache = KVCache()
  for step in range(max_new_tokens):
    step_start = time.time()
    curr_seq_len = len(tokens)
    if max_new_tokens > 1:
      print(
        f"\n--- Generating Token {step + 1}/{max_new_tokens} (seq_len={curr_seq_len}) ---"
      )
    with loader.load_scope() as ctx:
      embed_weight = ctx.load("model.embed_tokens.weight")
      embedding = Embedding(embed_weight)
      if step == 0:
        x = embedding.embed(tokens)
        positions = np.arange(curr_seq_len, dtype=np.float32)
      else:
        x = embedding.embed([tokens[-1]])
        positions = np.array([curr_seq_len - 1], dtype=np.float32)
    next_token_id, logits, output = decoder.step(x, positions, cache)
    next_token_str = tokenizer.decode([next_token_id])
    top5_indices = np.argsort(logits)[-5:][::-1]
    top5_logits = logits[top5_indices]
    top5_tokens = [tokenizer.decode([idx]) for idx in top5_indices]
    step_end = time.time()
    tokens.append(next_token_id)
    generated_tokens.append(next_token_id)
    print(f"Generated Token {step + 1} ID  : {next_token_id}")
    print(f"Generated Token {step + 1} Text: {next_token_str!r}")
    print(f"Top 5 Logits            : {top5_logits.tolist()}")
    print(f"Top 5 Tokens            : {top5_tokens}")
    print(f"Step Time               : {step_end - step_start:.2f} s")
    print(f"Process memory          : {get_ram_mb():.1f} MB")
    if next_token_id in eos_token_ids:
      print(f"Reached EOS token ({next_token_str!r}, ID: {next_token_id}).")
      break
  t_end = time.time()
  full_generated_text = tokenizer.decode(generated_tokens)
  print("\n================ FINAL GENERATION RESULT ================")
  print(f"Total Tokens Generated  : {len(generated_tokens)}")
  print(f"Generated Tokens List   : {generated_tokens}")
  print(f"Generated Text          : {full_generated_text!r}")
  print(f"Total Computation Time  : {t_end - t_start:.2f} s")
  print(f"Peak Process memory: {get_ram_mb():.1f} MB")
  print("=========================================================\n")
if __name__ == "__main__":
  parser = argparse.ArgumentParser(
    description="Qwen2.5-Coder-7B-Instruct pure NumPy inference"
  )
  parser.add_argument(
    "num_tokens",
    type=int,
    nargs="?",
    default=None,
    help="Number of tokens to generate (positional)",
  )
  parser.add_argument(
    "--max-new-tokens",
    "-n",
    type=int,
    default=None,
    help="Number of tokens to generate",
  )
  parser.add_argument(
    "--prompt",
    "-p",
    type=str,
    default="Give me the quicksort algorithm in python",
    help="Input prompt",
  )
  args = parser.parse_args()
  max_new_tokens = (
    args.max_new_tokens
    if args.max_new_tokens is not None
    else (args.num_tokens if args.num_tokens is not None else 3)
  )
  generate_tokens(max_new_tokens=max_new_tokens, prompt=args.prompt)
---
[dependencies]
anyhow = "1.0"
hf-hub = "0.4"
rayon = "1.10"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
sysinfo = "0.33"
tokenizers = "0.21"
[profile.dev]
opt-level = 3
[profile.release]
opt-level = 3
---
#![allow(dead_code, unused_variables, deprecated, non_camel_case_types)]
use anyhow::{Context, Result};
use hf_hub::api::sync::Api;
use rayon::prelude::*;
use serde::Deserialize;
use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::PathBuf;
use sysinfo::{Pid, System};
use tokenizers::Tokenizer;
// --- MEMORY MONITORING ---
pub fn get_ram_mb() -> f64 {
  let mut sys = System::new();
  let pid = Pid::from_u32(std::process::id());
  sys.refresh_processes_specifics(
    sysinfo::ProcessesToUpdate::Some(&[pid]),
    true,
    sysinfo::ProcessRefreshKind::nothing().with_memory(),
  );
  if let Some(proc_) = sys.process(pid) {
    proc_.memory() as f64 / (1024.0 * 1024.0)
  } else {
    0.0
  }
}
// --- MEMORY EFFICIENT SAFETENSORS LOADER ---
#[derive(Deserialize)]
struct IndexJson {
  weight_map: HashMap<String, String>,
}
#[derive(Deserialize)]
struct TensorInfo {
  dtype: String,
  shape: Vec<usize>,
  data_offsets: (usize, usize),
}
struct ShardHeader {
  header_len: u64,
  tensors: HashMap<String, serde_json::Value>,
}
pub struct MemoryEfficientSafetensorsLoader {
  model_dir: PathBuf,
  weight_map: HashMap<String, String>,
  shard_headers: HashMap<String, ShardHeader>,
}
impl MemoryEfficientSafetensorsLoader {
  pub fn new(repo_id: &str) -> Result<Self> {
    println!("Checking / downloading model weight shards from HuggingFace Hub...");
    let api = Api::new().context("Failed to initialize HuggingFace API client")?;
    let repo = api.model(repo_id.to_string());
    let index_path = repo
      .get("model.safetensors.index.json")
      .context("Failed to get model.safetensors.index.json")?;
    let model_dir = index_path
      .parent()
      .context("Failed to get parent dir of index file")?
      .to_path_buf();
    let index_file = File::open(&index_path)?;
    let index: IndexJson = serde_json::from_reader(index_file)?;
    let weight_map = index.weight_map;
    let unique_shards: HashSet<String> = weight_map.values().cloned().collect();
    let mut shard_headers = HashMap::new();
    for shard in unique_shards {
      let shard_path = repo
        .get(&shard)
        .with_context(|| format!("Failed to fetch shard {}", shard))?;
      let mut file = File::open(&shard_path)?;
      let mut header_len_bytes = [0u8; 8];
      file.read_exact(&mut header_len_bytes)?;
      let header_len = u64::from_le_bytes(header_len_bytes);
      let mut header_bytes = vec![0u8; header_len as usize];
      file.read_exact(&mut header_bytes)?;
      let tensors: HashMap<String, serde_json::Value> = serde_json::from_slice(&header_bytes)?;
      shard_headers.insert(
        shard,
        ShardHeader {
          header_len,
          tensors,
        },
      );
    }
    Ok(Self {
      model_dir,
      weight_map,
      shard_headers,
    })
  }
  pub fn load_tensor(&self, name: &str) -> Result<(Vec<f32>, Vec<usize>)> {
    let shard = self
      .weight_map
      .get(name)
      .with_context(|| format!("Tensor '{}' not found in weight map", name))?;
    let shard_path = self.model_dir.join(shard);
    let header_info = &self.shard_headers[shard];
    let raw_val = header_info.tensors.get(name).with_context(|| {
      format!(
        "Tensor '{}' not found in header for shard '{}'",
        name, shard
      )
    })?;
    let tensor_meta: TensorInfo = serde_json::from_value(raw_val.clone())?;
    let header_len = header_info.header_len;
    let (start, end) = tensor_meta.data_offsets;
    let byte_count = end - start;
    let mut file = File::open(&shard_path)?;
    file.seek(SeekFrom::Start(8 + header_len + start as u64))?;
    let mut raw_bytes = vec![0u8; byte_count];
    file.read_exact(&mut raw_bytes)?;
    let shape = tensor_meta.shape.clone();
    let num_elements: usize = shape.iter().product();
    let data = match tensor_meta.dtype.as_str() {
      "BF16" => {
        let mut arr = Vec::with_capacity(num_elements);
        for chunk in raw_bytes.chunks_exact(2) {
          let u16_val = u16::from_le_bytes([chunk[0], chunk[1]]);
          let f32_val = f32::from_bits((u16_val as u32) << 16);
          arr.push(f32_val);
        }
        arr
      }
      "F32" => {
        let mut arr = Vec::with_capacity(num_elements);
        for chunk in raw_bytes.chunks_exact(4) {
          let f32_val = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
          arr.push(f32_val);
        }
        arr
      }
      "F16" => {
        let mut arr = Vec::with_capacity(num_elements);
        for chunk in raw_bytes.chunks_exact(2) {
          let u16_val = u16::from_le_bytes([chunk[0], chunk[1]]);
          let f32_val = f16_to_f32(u16_val);
          arr.push(f32_val);
        }
        arr
      }
      other => anyhow::bail!("Unsupported tensor dtype: {}", other),
    };
    Ok((data, shape))
  }
}
fn f16_to_f32(h: u16) -> f32 {
  let sign = ((h >> 15) as u32) << 31;
  let exp = ((h >> 10) & 0x1f) as u32;
  let mant = (h & 0x03ff) as u32;
  if exp == 0 {
    if mant == 0 {
      f32::from_bits(sign)
    } else {
      let mut m = mant << 1;
      let mut e = 0;
      while (m & 0x0400) == 0 {
        m <<= 1;
        e += 1;
      }
      let exp32 = (127 - 15 - e) << 23;
      let mant32 = (m & 0x03ff) << 13;
      f32::from_bits(sign | exp32 | mant32)
    }
  } else if exp == 31 {
    let exp32 = 0xff << 23;
    let mant32 = mant << 13;
    f32::from_bits(sign | exp32 | mant32)
  } else {
    let exp32 = (exp + (127 - 15)) << 23;
    let mant32 = mant << 13;
    f32::from_bits(sign | exp32 | mant32)
  }
}
// --- OPTIMIZED PARALLEL MATRIX MULTIPLICATION WITH CHUNKING ---
pub fn mat_mul(w: &[f32], m: usize, k: usize, x: &[f32], n: usize) -> Vec<f32> {
  let mut y = vec![0.0f32; m * n];
  let chunk_size = (m / (rayon::current_num_threads() * 4)).max(16);
  y.par_chunks_mut(chunk_size * n)
    .enumerate()
    .for_each(|(chunk_idx, chunk)| {
      let start_row = chunk_idx * chunk_size;
      let rows_in_chunk = chunk.len() / n;
      for r in 0..rows_in_chunk {
        let j = start_row + r;
        let w_row = &w[j * k..(j + 1) * k];
        let y_row = &mut chunk[r * n..(r + 1) * n];
        for i in 0..n {
          let mut sum = 0.0f32;
          for idx in 0..k {
            sum += w_row[idx] * x[idx * n + i];
          }
          y_row[i] = sum;
        }
      }
    });
  y
}
pub fn mat_mul_bias(w: &[f32], m: usize, k: usize, x: &[f32], n: usize, bias: &[f32]) -> Vec<f32> {
  let mut y = vec![0.0f32; m * n];
  let chunk_size = (m / (rayon::current_num_threads() * 4)).max(16);
  y.par_chunks_mut(chunk_size * n)
    .enumerate()
    .for_each(|(chunk_idx, chunk)| {
      let start_row = chunk_idx * chunk_size;
      let rows_in_chunk = chunk.len() / n;
      for r in 0..rows_in_chunk {
        let j = start_row + r;
        let w_row = &w[j * k..(j + 1) * k];
        let b = bias[j];
        let y_row = &mut chunk[r * n..(r + 1) * n];
        for i in 0..n {
          let mut sum = 0.0f32;
          for idx in 0..k {
            sum += w_row[idx] * x[idx * n + i];
          }
          y_row[i] = sum + b;
        }
      }
    });
  y
}
// --- KV CACHE FOR DECODING ---
#[derive(Default)]
pub struct KVCache {
  // layer_idx -> (k_cache, v_cache)
  // k_cache: (num_kv_heads * head_dim, total_seq_len)
  pub cache: HashMap<usize, (Vec<f32>, Vec<f32>)>,
}
impl KVCache {
  pub fn update(
    &mut self,
    layer_idx: usize,
    new_k: &[f32],
    new_v: &[f32],
    num_kv_heads: usize,
    head_dim: usize,
    seq_len: usize,
  ) -> (&[f32], &[f32], usize) {
    let kv_dim = num_kv_heads * head_dim;
    let entry = self
      .cache
      .entry(layer_idx)
      .or_insert_with(|| (Vec::new(), Vec::new()));
    let prev_len = if entry.0.is_empty() {
      0
    } else {
      entry.0.len() / kv_dim
    };
    let new_len = prev_len + seq_len;
    if prev_len == 0 {
      entry.0 = new_k.to_vec();
      entry.1 = new_v.to_vec();
    } else {
      let mut updated_k = vec![0.0f32; kv_dim * new_len];
      let mut updated_v = vec![0.0f32; kv_dim * new_len];
      for d in 0..kv_dim {
        let src_old_k = &entry.0[d * prev_len..(d + 1) * prev_len];
        let src_new_k = &new_k[d * seq_len..(d + 1) * seq_len];
        let dst_k = &mut updated_k[d * new_len..(d + 1) * new_len];
        dst_k[..prev_len].copy_from_slice(src_old_k);
        dst_k[prev_len..].copy_from_slice(src_new_k);
        let src_old_v = &entry.1[d * prev_len..(d + 1) * prev_len];
        let src_new_v = &new_v[d * seq_len..(d + 1) * seq_len];
        let dst_v = &mut updated_v[d * new_len..(d + 1) * new_len];
        dst_v[..prev_len].copy_from_slice(src_old_v);
        dst_v[prev_len..].copy_from_slice(src_new_v);
      }
      entry.0 = updated_k;
      entry.1 = updated_v;
    }
    (&entry.0, &entry.1, new_len)
  }
}
pub fn rms_norm(
  x: &[f32],
  hidden_size: usize,
  seq_len: usize,
  weight: &[f32],
  eps: f32,
) -> Vec<f32> {
  let mut out = vec![0.0f32; hidden_size * seq_len];
  let inv_rms: Vec<f32> = (0..seq_len)
    .into_par_iter()
    .map(|i| {
      let mut sum_sq = 0.0f32;
      for r in 0..hidden_size {
        let val = x[r * seq_len + i];
        sum_sq += val * val;
      }
      let mean_sq = sum_sq / (hidden_size as f32);
      1.0 / (mean_sq + eps).sqrt()
    })
    .collect();
  out
    .par_chunks_mut(seq_len)
    .enumerate()
    .for_each(|(r, row)| {
      let w = weight[r];
      let x_row = &x[r * seq_len..(r + 1) * seq_len];
      for i in 0..seq_len {
        row[i] = x_row[i] * inv_rms[i] * w;
      }
    });
  out
}
pub fn apply_rotary_pos_emb(
  q: &mut [f32],
  k: &mut [f32],
  seq_len: usize,
  start_pos: usize,
  num_heads: usize,
  num_kv_heads: usize,
  head_dim: usize,
  inv_freq: &[f32],
) {
  let half_dim = inv_freq.len();
  for step_i in 0..seq_len {
    let pos = (start_pos + step_i) as f32;
    let mut cos = vec![0.0f32; half_dim];
    let mut sin = vec![0.0f32; half_dim];
    for i in 0..half_dim {
      let freq = pos * inv_freq[i];
      cos[i] = freq.cos();
      sin[i] = freq.sin();
    }
    for h in 0..num_heads {
      let head_base = h * head_dim;
      for i in 0..half_dim {
        let idx1 = (head_base + i) * seq_len + step_i;
        let idx2 = (head_base + i + half_dim) * seq_len + step_i;
        let q1 = q[idx1];
        let q2 = q[idx2];
        q[idx1] = q1 * cos[i] - q2 * sin[i];
        q[idx2] = q2 * cos[i] + q1 * sin[i];
      }
    }
    for h in 0..num_kv_heads {
      let head_base = h * head_dim;
      for i in 0..half_dim {
        let idx1 = (head_base + i) * seq_len + step_i;
        let idx2 = (head_base + i + half_dim) * seq_len + step_i;
        let k1 = k[idx1];
        let k2 = k[idx2];
        k[idx1] = k1 * cos[i] - k2 * sin[i];
        k[idx2] = k2 * cos[i] + k1 * sin[i];
      }
    }
  }
}
pub fn softmax_inplace(x: &mut [f32]) {
  let max_val = x.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
  let mut sum = 0.0f32;
  for val in x.iter_mut() {
    *val = (*val - max_val).exp();
    sum += *val;
  }
  let inv_sum = 1.0 / sum;
  for val in x.iter_mut() {
    *val *= inv_sum;
  }
}
pub fn compute_gqa_attention(
  q: &[f32],
  k: &[f32],
  v: &[f32],
  step_seq_len: usize,
  total_seq_len: usize,
  num_heads: usize,
  num_kv_heads: usize,
  head_dim: usize,
) -> Vec<f32> {
  let num_kv_groups = num_heads / num_kv_heads;
  let scale = 1.0 / (head_dim as f32).sqrt();
  let hidden_size = num_heads * head_dim;
  let mut attn_out = vec![0.0f32; hidden_size * step_seq_len];
  let head_outputs: Vec<(usize, Vec<f32>)> = (0..num_heads)
    .into_par_iter()
    .map(|h| {
      let kv_h = h / num_kv_groups;
      let mut h_out = vec![0.0f32; head_dim * step_seq_len];
      let q_head_base = h * head_dim;
      let kv_head_base = kv_h * head_dim;
      for i in 0..step_seq_len {
        let current_pos = total_seq_len - step_seq_len + i;
        let mut scores = vec![0.0f32; total_seq_len];
        for j in 0..total_seq_len {
          if j > current_pos {
            scores[j] = -1e9;
          } else {
            let mut dot = 0.0f32;
            for d in 0..head_dim {
              let q_val = q[(q_head_base + d) * step_seq_len + i];
              let k_val = k[(kv_head_base + d) * total_seq_len + j];
              dot += q_val * k_val;
            }
            scores[j] = dot * scale;
          }
        }
        softmax_inplace(&mut scores);
        for j in 0..total_seq_len {
          let weight = scores[j];
          if weight > 0.0 {
            for d in 0..head_dim {
              let v_val = v[(kv_head_base + d) * total_seq_len + j];
              h_out[d * step_seq_len + i] += weight * v_val;
            }
          }
        }
      }
      (h, h_out)
    })
    .collect();
  for (h, h_out) in head_outputs {
    let q_head_base = h * head_dim;
    for d in 0..head_dim {
      let src_row = &h_out[d * step_seq_len..(d + 1) * step_seq_len];
      let dst_row =
        &mut attn_out[(q_head_base + d) * step_seq_len..(q_head_base + d + 1) * step_seq_len];
      dst_row.copy_from_slice(src_row);
    }
  }
  attn_out
}
pub fn swiglu(gate: &[f32], up: &[f32]) -> Vec<f32> {
  let mut out = vec![0.0f32; gate.len()];
  out
    .par_iter_mut()
    .zip(gate.par_iter())
    .zip(up.par_iter())
    .for_each(|((o, &g), &u)| {
      let silu = g / (1.0 + (-g).exp());
      *o = silu * u;
    });
  out
}
// --- MODEL FORWARD PASS ---
pub fn generate_tokens(max_new_tokens: usize, prompt: &str) {
  println!("=== Qwen2.5-Coder-7B-Instruct pure Rust Inference ===");
  println!("Prompt: {:?}", prompt);
  println!("Max New Tokens: {}", max_new_tokens);
  println!(
    "Initial Process memory consumption: {:.1} MB\n",
    get_ram_mb()
  );
  let t_start = std::time::Instant::now();
  let repo_id = "Qwen/Qwen2.5-Coder-7B-Instruct";
  let api = Api::new().unwrap();
  let repo = api.model(repo_id.to_string());
  let tokenizer_path = repo.get("tokenizer.json").unwrap();
  let tokenizer = Tokenizer::from_file(&tokenizer_path).unwrap();
  let loader = MemoryEfficientSafetensorsLoader::new(repo_id).unwrap();
  let text = format!(
        "<|im_start|>system\nYou are a helpful assistant specializing in coding.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n",
        prompt
    );
  let encoding = tokenizer.encode(text.as_str(), true).unwrap();
  let mut tokens: Vec<u32> = encoding.get_ids().to_vec();
  let initial_seq_len = tokens.len();
  println!("Formatted Chat Prompt:\n{}", text);
  println!("Tokenized Sequence Length: {} tokens.\n", initial_seq_len);
  let num_layers = 28;
  let hidden_size = 3584;
  let num_heads = 28;
  let num_kv_heads = 4;
  let head_dim = 128;
  let rope_theta = 1000000.0f32;
  let eps = 1e-6f32;
  // Precompute RoPE inverse frequencies once (mirrors QwenRotaryEmbedding).
  let half_dim = head_dim / 2;
  let mut inv_freq = vec![0.0f32; half_dim];
  for i in 0..half_dim {
    inv_freq[i] = 1.0 / rope_theta.powf((2 * i) as f32 / head_dim as f32);
  }
  let eos_token_ids: HashSet<u32> = [
    tokenizer.token_to_id("<|im_end|>"),
    tokenizer.token_to_id("<|endoftext|>"),
  ]
  .iter()
  .filter_map(|&id| id)
  .collect();
  let mut generated_tokens = Vec::new();
  let mut kv_cache = KVCache::default();
  for step in 0..max_new_tokens {
    let step_start = std::time::Instant::now();
    let total_seq_len = tokens.len();
    if max_new_tokens > 1 {
      println!(
        "\n--- Generating Token {}/{} (seq_len={}) ---",
        step + 1,
        max_new_tokens,
        total_seq_len
      );
    }
    let (step_tokens, start_pos) = if step == 0 {
      (tokens.as_slice(), 0)
    } else {
      (&tokens[tokens.len() - 1..], total_seq_len - 1)
    };
    let step_seq_len = step_tokens.len();
    // 1. Embedding Lookup
    let (embed_weight, _) = loader.load_tensor("model.embed_tokens.weight").unwrap();
    let mut x = vec![0.0f32; hidden_size * step_seq_len];
    for (i, &tok_id) in step_tokens.iter().enumerate() {
      let tok_idx = tok_id as usize;
      let src = &embed_weight[tok_idx * hidden_size..(tok_idx + 1) * hidden_size];
      for r in 0..hidden_size {
        x[r * step_seq_len + i] = src[r];
      }
    }
    drop(embed_weight);
    // 2. Sequential Layer Forward Pass
    for layer_idx in 0..num_layers {
      let pfx = format!("model.layers.{}.", layer_idx);
      // Pre-Attention RMSNorm
      let (in_norm_w, _) = loader
        .load_tensor(&format!("{}input_layernorm.weight", pfx))
        .unwrap();
      let norm_x = rms_norm(&x, hidden_size, step_seq_len, &in_norm_w, eps);
      drop(in_norm_w);
      // Projections
      let (q_w, q_shape) = loader
        .load_tensor(&format!("{}self_attn.q_proj.weight", pfx))
        .unwrap();
      let (q_b, _) = loader
        .load_tensor(&format!("{}self_attn.q_proj.bias", pfx))
        .unwrap();
      let mut q = mat_mul_bias(&q_w, q_shape[0], q_shape[1], &norm_x, step_seq_len, &q_b);
      drop(q_w);
      drop(q_b);
      let (k_w, k_shape) = loader
        .load_tensor(&format!("{}self_attn.k_proj.weight", pfx))
        .unwrap();
      let (k_b, _) = loader
        .load_tensor(&format!("{}self_attn.k_proj.bias", pfx))
        .unwrap();
      let mut k = mat_mul_bias(&k_w, k_shape[0], k_shape[1], &norm_x, step_seq_len, &k_b);
      drop(k_w);
      drop(k_b);
      let (v_w, v_shape) = loader
        .load_tensor(&format!("{}self_attn.v_proj.weight", pfx))
        .unwrap();
      let (v_b, _) = loader
        .load_tensor(&format!("{}self_attn.v_proj.bias", pfx))
        .unwrap();
      let v = mat_mul_bias(&v_w, v_shape[0], v_shape[1], &norm_x, step_seq_len, &v_b);
      drop(v_w);
      drop(v_b);
      // RoPE
      apply_rotary_pos_emb(
        &mut q,
        &mut k,
        step_seq_len,
        start_pos,
        num_heads,
        num_kv_heads,
        head_dim,
        &inv_freq,
      );
      // Update KV Cache
      let (full_k, full_v, cached_seq_len) =
        kv_cache.update(layer_idx, &k, &v, num_kv_heads, head_dim, step_seq_len);
      // GQA Attention
      let attn_out = compute_gqa_attention(
        &q,
        full_k,
        full_v,
        step_seq_len,
        cached_seq_len,
        num_heads,
        num_kv_heads,
        head_dim,
      );
      drop(q);
      drop(k);
      drop(v);
      // Output projection
      let (o_w, o_shape) = loader
        .load_tensor(&format!("{}self_attn.o_proj.weight", pfx))
        .unwrap();
      let attn_proj = mat_mul(&o_w, o_shape[0], o_shape[1], &attn_out, step_seq_len);
      drop(o_w);
      drop(attn_out);
      // Residual 1
      for i in 0..x.len() {
        x[i] += attn_proj[i];
      }
      drop(attn_proj);
      // Pre-MLP RMSNorm & SwiGLU MLP
      let (post_norm_w, _) = loader
        .load_tensor(&format!("{}post_attention_layernorm.weight", pfx))
        .unwrap();
      let norm_x2 = rms_norm(&x, hidden_size, step_seq_len, &post_norm_w, eps);
      drop(post_norm_w);
      let (gate_w, gate_shape) = loader
        .load_tensor(&format!("{}mlp.gate_proj.weight", pfx))
        .unwrap();
      let gate = mat_mul(
        &gate_w,
        gate_shape[0],
        gate_shape[1],
        &norm_x2,
        step_seq_len,
      );
      drop(gate_w);
      let (up_w, up_shape) = loader
        .load_tensor(&format!("{}mlp.up_proj.weight", pfx))
        .unwrap();
      let up = mat_mul(&up_w, up_shape[0], up_shape[1], &norm_x2, step_seq_len);
      drop(up_w);
      drop(norm_x2);
      let mlp_act = swiglu(&gate, &up);
      drop(gate);
      drop(up);
      let (down_w, down_shape) = loader
        .load_tensor(&format!("{}mlp.down_proj.weight", pfx))
        .unwrap();
      let mlp_out = mat_mul(
        &down_w,
        down_shape[0],
        down_shape[1],
        &mlp_act,
        step_seq_len,
      );
      drop(down_w);
      drop(mlp_act);
      // Residual 2
      for i in 0..x.len() {
        x[i] += mlp_out[i];
      }
      drop(mlp_out);
      if max_new_tokens == 1 {
        if (layer_idx + 1) % 7 == 0 || layer_idx == 0 || layer_idx == num_layers - 1 {
          println!(
            "Layer {:2}/{num_layers} completed. Process RAM: {:.1} MB",
            layer_idx + 1,
            get_ram_mb()
          );
        }
      } else {
        if (layer_idx + 1) % 14 == 0 || layer_idx == num_layers - 1 {
          println!(
            "Token {}/{} - Layer {:2}/{num_layers} completed. Process RAM: {:.1} MB",
            step + 1,
            max_new_tokens,
            layer_idx + 1,
            get_ram_mb()
          );
        }
      }
    }
    // 3. Final RMSNorm & LM Head Projection
    if max_new_tokens == 1 {
      println!("\nProcessing final norm and LM head...");
    }
    let (final_norm_w, _) = loader.load_tensor("model.norm.weight").unwrap();
    let x_norm = rms_norm(&x, hidden_size, step_seq_len, &final_norm_w, eps);
    drop(final_norm_w);
    drop(x);
    let mut last_token_hidden = vec![0.0f32; hidden_size];
    for r in 0..hidden_size {
      last_token_hidden[r] = x_norm[r * step_seq_len + (step_seq_len - 1)];
    }
    let (lm_head_w, lm_shape) = loader.load_tensor("lm_head.weight").unwrap();
    let vocab_size = lm_shape[0];
    let logits = mat_mul(&lm_head_w, vocab_size, hidden_size, &last_token_hidden, 1);
    drop(lm_head_w);
    let mut indexed_logits: Vec<(u32, f32)> = logits
      .iter()
      .copied()
      .enumerate()
      .map(|(idx, logit)| (idx as u32, logit))
      .collect();
    indexed_logits
      .sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    let top5 = &indexed_logits[..5];
    let next_token_id = top5[0].0;
    let next_token_str = tokenizer
      .decode(&[next_token_id], false)
      .unwrap_or_default();
    let top5_logits: Vec<f32> = top5.iter().map(|&(_, l)| l).collect();
    let top5_tokens: Vec<String> = top5
      .iter()
      .map(|&(id, _)| tokenizer.decode(&[id], false).unwrap_or_default())
      .collect();
    let step_end = std::time::Instant::now();
    let step_duration = step_end.duration_since(step_start).as_secs_f64();
    tokens.push(next_token_id);
    generated_tokens.push(next_token_id);
    if max_new_tokens == 1 {
      let total_duration = step_end.duration_since(t_start).as_secs_f64();
      println!("\n================ FINAL GENERATION RESULT ================");
      println!("Generated Next Token ID  : {}", next_token_id);
      println!("Generated Next Token Text: {:?}", next_token_str);
      println!("Top 5 Logits            : {:?}", top5_logits);
      println!("Top 5 Tokens            : {:?}", top5_tokens);
      println!("Total Computation Time  : {:.2} s", total_duration);
      println!("Peak Process RAM Usage  : {:.1} MB", get_ram_mb());
      println!("=========================================================\n");
    } else {
      println!("Generated Token {} ID  : {}", step + 1, next_token_id);
      println!("Generated Token {} Text: {:?}", step + 1, next_token_str);
      println!("Top 5 Logits            : {:?}", top5_logits);
      println!("Top 5 Tokens            : {:?}", top5_tokens);
      println!("Step Time               : {:.2} s", step_duration);
    }
    if eos_token_ids.contains(&next_token_id) {
      println!(
        "Reached EOS token ({:?}, ID: {}).",
        next_token_str, next_token_id
      );
      break;
    }
  }
  let t_end = std::time::Instant::now();
  if max_new_tokens > 1 {
    let full_generated_text = tokenizer
      .decode(&generated_tokens, false)
      .unwrap_or_default();
    let total_duration = t_end.duration_since(t_start).as_secs_f64();
    println!("\n================ FINAL GENERATION RESULT ================");
    println!("Total Tokens Generated  : {}", generated_tokens.len());
    println!("Generated Tokens List   : {:?}", generated_tokens);
    println!("Generated Text          : {:?}", full_generated_text);
    println!("Total Computation Time  : {:.2} s", total_duration);
    println!("Peak Process RAM Usage  : {:.1} MB", get_ram_mb());
    println!("=========================================================\n");
  }
}
fn main() {
  let args: Vec<String> = std::env::args().collect();
  let mut max_new_tokens: Option<usize> = None;
  let mut prompt: String = "Give me the quicksort algorithm in python".to_string();
  let mut i = 1;
  while i < args.len() {
    let arg = &args[i];
    if arg == "-n" || arg == "--max-new-tokens" {
      if i + 1 < args.len() {
        max_new_tokens = args[i + 1].parse().ok();
        i += 1;
      }
    } else if arg == "-p" || arg == "--prompt" {
      if i + 1 < args.len() {
        prompt = args[i + 1].clone();
        i += 1;
      }
    } else if !arg.starts_with('-') && max_new_tokens.is_none() {
      max_new_tokens = arg.parse().ok();
    }
    i += 1;
  }
  let max_new_tokens = max_new_tokens.unwrap_or(3);
  generate_tokens(max_new_tokens, &prompt);
}

This is the output when I prompt with “Quicksort in python, just the code no preamble please”:

```python
def quicksort(arr):
    if len(arr) <= 1:
        return arr
    pivot = arr[len(arr) // 2]
    left = [x for x in arr if x < pivot]
    middle = [x for x in arr if x == pivot]
 right = [x for x in arr if x > pivot]
    return quicksort(left) + middle + quicksort(right)
```<|im_end|>

Miscellaneous

Mixture of Experts (MoE)

How do we scale even more?

In all the model architectures, the FFN is used for every single token in every single layer. By increasing the size of FFN more information can be stored but that will also incur higher compute costs during decoding (as if decoding isn’t costly enough already).

The idea behind MOE is to split the FFN into many many small FFN networks (called experts) and only a few of those experts are activated per layer. The experts to activate is a fixed number $k$ which is controlled by another small “gating network” (or gating function) and then we take the weighted sum of the output of those activated FFN experts. In equation that looks like:

\text{Intermediate} = I + \text{Attention}_{\text{output}} \times W^O

\text{Output} = \text{Intermediate} + \text{MoE}(\text{RMSNorm}(\text{Intermediate}))

\text{MoE}(X) = \sum_{i \in \text{TopK}} G_i(X) \times \text{FFN}_i(X)

G_i(X) = \text{Softmax}\left(\text{TopK}\left(X \times W_{\text{gating}}, k\right)\right)_i

Shared MoE

Some experts are needed for all cases (like to understand name, grammar etc) which will now get duplicated over the experts. Another variation is having shared experts which are always active for all and use MoE for the rest of the FFNs. The output of shared experts is summed with the output of the gated MoE. Generally in shared MoE we decrease the parameters provided to individual expert which allows us to scale the number of experts and activate multiple of them.

One problem with MoE is breakdown of experts and network relying on only a few of them (as early on grammar might be learned by a few experts and the network routing everything to them). To account for this, an additional auxiliary loss is used during training, which I’ll talk about in the training blog.

RoPE and runtime context scaling (YARN)

One of the most important benefits of RoPE technique is that it allows a model to scale context window beyond the limits of its training (either zero-shot or very small fine-tuning step). If the base $b$ of RoPE is kept the same, but the context window is increased, then the model ends up in unknown territory where the new values are never seen by the model during training. But if base is modified so that even larger attentions are represented by the same distribution as in the training, then the model can perform just as well.

But scaling just the base leads to model scaling up frequency for the earlier pair in the embeddings as well, where the model is supposed to attend to nearby tokens. This leads to “blurriness” where the model fails to figure out the relationship between nearby tokens (as the model was trained for specific frequencies). To fix this, either the model is fine-tunned over a small corpus, or NTK-aware scaling is used where the faster frequencies are largely left alone and just the slower frequencies are modified to accommodate for larger contexts. The current SoTA for NTK-aware scaling is YARN (again Google around, its pretty straightforward).

Multi-head Latent Attention (MLA)

Another technique to save memory introduced by DeepSeek

The basic idea is: If each attention head is already learning some low-rank representation (as shown by empirical evidence), why not use a single KV per layer, rather than per attention head, which during decoding can be transformed to represent each head?

Hold your questions about compute costs.

So, what we have are 3 matrices:

A $W^{DKV}$ to project to some low-rank dimension $d_c$ which serves as a base for KV, generally called $c^{KV}$ . This matrix is of size $[d_c, d_{model}]$ signifying the transformation from embedding size to $d_c$ size.
A $W^{UK}_h$ to project $c^{KV}$ back to $d_{head}$ for key per head. This matrix is of size $[d_{head}, d_c]$ .
A $W^{UV}_h$ to project $c^{KV}$ back to $d_{head}$ for value per head. This matrix is of size $[d_{head}, d_c]$ .

W^{DKV}

is per layer, whereas

W^{UK}_h

and

W^{UV}_h

are per attention head (hence the

h

in subscript). We are just adding one additional matrix for down projection, up projection ones replace the older

W^K

and

W^V

For each token, its low-rank representation of size $c^{KV}$ is stored in the cache. Whereas previously, two vectors per token per KV head of size $d_{model}$ (and generally $d_{model} \gg d_c$ like $4096$ v/s $128$ ) was stored in the cache.

QBut what about the compute?

You might think now we have to reconstruct KV for each token when decoding. Ah no. And that’s where the beauty of this technique comes. Let’s see the equation we’ll have to apply during decoding of token $t$ w.r.t. past tokens (say we have seen $N$ tokens and $C^{KV}$ is the new low-rank kv cache of size $[d_c, N]$ ) for one head (omitting $h$ for clarity):

Attention_{score}^t = K^T \times Q_t = (W^{UK} \times C^{KV})^T \times Q_t = (C^{KV})^T \times ((W^{UK})^T \times Q_t)

using matrix associativity

Notice the brilliance here, each $K$ does not have to be computed separately, rather the attention head constant matrix $W^{UK}$ can be absorbed in to the newly computed query vector $Q_t$ for each head, denoted by $\tilde{Q_t}$ . Or thought from query’s perspective, it is projecting itself down to the low-rank space based on that attention head’s specific keys, asking the question in the key space. Similar behavior with absorbing the $Attention_{score}$ directly with $W^{UV}$ .

Empirical evidence shows that this model outperforms both GQA in both accuracy and memory savings, while maintaining the same speed (just one additional matmul, but reduces the dimension of other matmuls).

Encoder architecture

What I have described is what is commonly known as a “decoder only” architecture (as opposed to encoder-decoder or encoder only architecture). The only difference is that in an encoder all tokens attend to all tokens past and future. Decoder then attends to its own generated tokens’s KV, called self-attention, but also to the encoder’s KV, called cross-attention.

Each architecture has its own use based on the problem statement:

Encoder-only: Used generally for understanding tasks like sentiment analysis, image categorization etc.
Decoder-only: Natural language generation.
Encoder-decoder: Sequence to sequence tasks where target output sequence depends on a structured input like describe an image, or translation etc.

QK normalization

Remember the $\sqrt{d_k}$ when computing the attention score? Turns out that is not enough.

The goal of a softmax loss to make sure that the output of softmax is $1$ only for the right logit and $0$ every where else (the cross-entropy loss). For that to happen the gradient flowing back into the weights are continuously pushing them to be higher and higher positive values which can eventually break down the network. The assumption behind $\sqrt{d_k}$ is that the mean stays at $0$ and variance at $1$ but that isn’t true.
The residual network even with pre-norm keeps increasing the variance a bit which for deep layers add up to push $QK$ high enough to saturate softmax.
Deeper layers are also developing complex “thoughts” that they have to distinguish to figure out what token to attend to when all the tokens are saying they are the most important one. This will naturally push the weights to be higher, so that the softmax output is spread enough and it can isolate the important tokens.

The math is pretty simple actually, we just apply RMSNorm to the $QK$ vectors of each token.

Logit softcapping

QK normalization solves the problem for attention, but who solves the similar problem for LM Head’s softmax? That’s logit softcapping. The core idea is that we want to make sure that the raw value of the softmax output does not go beyond certain limits ( $30$ in Gemma 2 --- which BTW also used logit softcapping in attention layers as well rather than QK normalization with the value $50$ ). So, a hard cap is employed for softmax like $max(-50, min(50, value))$ .

But what about the gradients? Once a value goes beyond the limit its game over for training. To solve this, the same LSTM solution of $tanh$ is used, which has linear behavior within the limits and asymptotically approaches the limits. The formula used for LM head in Gemma 2:

Attention_{weights} = Softmax ( 50.0 \cdot tanh(\frac{QK^T}{50.0\cdot\sqrt{d_k}}))

Output_{logits} = 30.0 \cdot tanh(\frac{O_{transformer} \times W_{LM\_Head}}{30.0})

Sliding window attention

Do not attend to all the past tokens, attend to a fixed window and forget about the past before that. That’s it!

Multi Token Prediction (MTP)

Another one from DeepSeek

The idea is to have a smaller MTP model at the end of the main model which predicts the next token immediately rather than going through the entire autoregressive process of the main model. Take the hidden state, and then embedding of just predicted token, concat and pass through this MTP model. You can configure how many times this MTP model runs to generate multiple tokens (it was configured to $1$ for DeepSeek 3). Note that you still share the LM Head. It forces the model to learn a much richer representation because now it has to predict multiple future tokens from a single hidden state and it provides native speculative decoding.

Conclusion

The space of LLMs is crowded with many other ideas which are not based on transformer architecture (like SSMs). But transformers and the many concepts I have discussed here form the backbone of all major LLMs today.

In the next iteration in this series, I’ll build a highly efficient inference engine, targeting one processor (likely my M4 pro) and one single GPU node. I am hoping to be able to reach parity with some SoTA engine for one specific model.

Hacker News Discussion

Checking Hacker News for discussions...

Submit to Hacker News